perf: optimize spark_get_json_object (4x faster)#4907
Conversation
andygrove
changed the title
spark_get_json_object (4x faster)
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
First pass, thanks @andygrove! Mostly tests.
| fn extract_no_wildcard(json_str: &str, segments: &[PathSegment]) -> Option<Value> { | ||
| let mut de = serde_json::Deserializer::from_str(json_str); | ||
| let found = PathSeed { segments }.deserialize(&mut de).ok()?; | ||
| de.end().ok()?; | ||
| found | ||
| } |
There was a problem hiding this comment.
native/spark-expr/src/string_funcs/get_json_object.rs:290-295 (extract_no_wildcard) relies on three things staying true to preserve parity: de.end() (line 293) rejecting trailing content, visit_map draining every entry after a match (lines 362-370), and visit_seq calling IgnoredAny.visit_seq(seq) after a match (line 392). These are the load-bearing lines of the whole PR, since they are what makes "streaming stop-early" still behave like "full parse". None of them is covered by a test. If a future refactor drops the drain or the de.end(), every existing test still passes because they all use well-formed single-value documents.
Add tests that only pass if the whole document is still validated after the match:
#[test]
fn test_match_then_trailing_garbage_is_null() {
let path = parse_json_path("$.a").unwrap();
assert_eq!(evaluate_path(r#"{"a":1} garbage"#, &path), None);
}
#[test]
fn test_match_then_malformed_sibling_is_null() {
let path = parse_json_path("$.a").unwrap();
// "a" matches early, but "b" is malformed; a full parse rejects this.
assert_eq!(evaluate_path(r#"{"a":1,"b":}"#, &path), None);
}
#[test]
fn test_array_match_then_malformed_element_is_null() {
let path = parse_json_path("$[0]").unwrap();
assert_eq!(evaluate_path(r#"[1,,]"#, &path), None);
}| // Every entry is visited so that a duplicated key resolves to its last | ||
| // occurrence, as it would in a parsed object. |
There was a problem hiding this comment.
native/spark-expr/src/string_funcs/get_json_object.rs:360-361 visits every entry so a duplicated key "resolves to its last occurrence, as it would in a parsed object." That is correct for the old serde_json + preserve_order behavior (IndexMap overwrites the value, so .get returns the last), so this PR is faithful to the old Comet code. It is worth pinning with a test because the visit-every-entry loop is exactly what makes it work, and a future "break after match" optimization would silently flip it to first-wins:
#[test]
fn test_duplicate_key_last_wins() {
let path = parse_json_path("$.a").unwrap();
assert_eq!(evaluate_path(r#"{"a":1,"a":2}"#, &path), Some("2".to_string()));
}Separately, note for the record that both old and new Comet diverge from Spark here. Spark's GetJsonObjectEvaluator.evaluatePath (sql/catalyst/.../json/JsonExpressionEvalUtils.scala:478-488) stops at the first matching field once dirty is set, so Spark returns the first occurrence, not the last. This is a pre-existing Comet bug, not something #4907 introduces, but the comment claiming parity with "a parsed object" is only true against serde, not against Spark. Consider filing a follow-up issue rather than expanding this PR.
| segments: &'a [PathSegment], | ||
| } | ||
|
|
||
| impl<'de> Visitor<'de> for SegmentVisitor<'_> { |
There was a problem hiding this comment.
native/spark-expr/src/string_funcs/get_json_object.rs SegmentVisitor::visit_unit returns Ok(None), which is what makes $.a.b on {"a":null} return null (recurse into null with a remaining segment yields no match). The old code got this from Null.as_object()? -> None. Parity holds, but there is no test for a null encountered mid-path (the existing test_evaluate_null_value only covers a terminal null). Add:
#[test]
fn test_null_midpath_is_null() {
let path = parse_json_path("$.a.b").unwrap();
assert_eq!(evaluate_path(r#"{"a":null}"#, &path), None);
}Add regression tests pinning the behaviors that keep the streaming DeserializeSeed path walk equivalent to a full parse: - trailing garbage after a match is rejected (de.end()) - a malformed sibling after a match is rejected (visit_map drain) - a malformed array element after a match is rejected (visit_seq drain) - a duplicated key resolves to its last occurrence (preserve_order parity) - a null encountered mid-path yields no match (visit_unit)
|
Thanks @mbutrovich . I added the tests and filed #4947 |
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
Approved pending CI, thanks @andygrove!
2 participants
Which issue does this PR close?
N/A
Rationale for this change
Optimize existing expression.
What changes are included in this PR?
Replaced the full serde_json::Value materialization of each row with a streaming DeserializeSeed path walk that skips non-matching subtrees via IgnoredAny and allocates only the matched subtree, eliminating per-key and per-value heap allocations for the whole document.
How are these changes tested?
Existing tests.
Benchmark (criterion):
Full criterion output: