diff --git a/native/spark-expr/src/string_funcs/get_json_object.rs b/native/spark-expr/src/string_funcs/get_json_object.rs index b61a3172fb..a127b51e81 100644 --- a/native/spark-expr/src/string_funcs/get_json_object.rs +++ b/native/spark-expr/src/string_funcs/get_json_object.rs @@ -255,56 +255,75 @@ fn parse_json_path(path: &str) -> Option { /// Evaluate a parsed JSONPath against a JSON string. /// Returns the result as a string, or None if no match. fn evaluate_path(json_str: &str, path: &ParsedPath) -> Option { - if !path.has_wildcard { - return value_into_string(extract_no_wildcard(json_str, &path.segments)?); + let mut result = extract_path(json_str, &path.segments)?; + if !result.matched { + return None; } - let value: Value = serde_json::from_str(json_str).ok()?; - - // Wildcard path: may return multiple results - let results = evaluate_with_wildcard(&value, &path.segments); + if !path.has_wildcard { + return value_into_string(result.values.pop()?); + } - match results.len() { + match result.values.len() { 0 => None, - 1 => { - // Single wildcard match: Spark preserves JSON serialization format - // (strings keep their quotes, numbers don't) - if results[0].is_null() { - None - } else { - serde_json::to_string(results[0]).ok() - } - } - // Multiple results: wrap in JSON array. A slice of `&Value` serializes - // as a JSON array, so no clone into an owned `Value::Array` is needed. - _ => serde_json::to_string(&results).ok(), + // A wildcard match preserves JSON serialization format, including + // quotes around strings and the JSON text `null`. + 1 => serde_json::to_string(&result.values[0]).ok(), + _ => serde_json::to_string(&result.values).ok(), } } -/// Evaluation for paths without wildcards. -/// /// Descends into the document while it is being parsed, so only the matched -/// subtree is materialized as a `Value`; everything else is skipped by the -/// parser without allocating. The whole document is still consumed, so -/// malformed JSON anywhere in the input yields no match, as a full parse would. -fn extract_no_wildcard(json_str: &str, segments: &[PathSegment]) -> Option { +/// subtrees are materialized as `Value`s; everything else is skipped by the +/// parser without allocating. The whole document is still consumed, so malformed +/// JSON anywhere in the input yields no match, as a full parse would. +fn extract_path(json_str: &str, segments: &[PathSegment]) -> Option { let mut de = serde_json::Deserializer::from_str(json_str); - let found = PathSeed { segments }.deserialize(&mut de).ok()?; + let found = PathSeed { + segments, + reject_direct_null: false, + } + .deserialize(&mut de) + .ok()?; de.end().ok()?; - found + Some(found) } /// Deserializes the value at `segments`, discarding everything else. struct PathSeed<'a> { segments: &'a [PathSegment], + /// A JSON null directly below a named field is not a match in Spark. Nulls + /// reached through array traversal are matches and serialize as `null`. + reject_direct_null: bool, +} + +#[derive(Default)] +struct PathResult { + values: Vec, + matched: bool, +} + +impl PathResult { + fn matched(value: Value) -> Self { + Self { + values: vec![value], + matched: true, + } + } } impl<'de> DeserializeSeed<'de> for PathSeed<'_> { - type Value = Option; + type Value = PathResult; fn deserialize>(self, deserializer: D) -> Result { if self.segments.is_empty() { - return Value::deserialize(deserializer).map(Some); + return Value::deserialize(deserializer).map(|value| { + if self.reject_direct_null && value.is_null() { + PathResult::default() + } else { + PathResult::matched(value) + } + }); } deserializer.deserialize_any(SegmentVisitor { segments: self.segments, @@ -320,50 +339,52 @@ struct SegmentVisitor<'a> { } impl<'de> Visitor<'de> for SegmentVisitor<'_> { - type Value = Option; + type Value = PathResult; fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str("a JSON value") } fn visit_bool(self, _: bool) -> Result { - Ok(None) + Ok(PathResult::default()) } fn visit_i64(self, _: i64) -> Result { - Ok(None) + Ok(PathResult::default()) } fn visit_u64(self, _: u64) -> Result { - Ok(None) + Ok(PathResult::default()) } fn visit_f64(self, _: f64) -> Result { - Ok(None) + Ok(PathResult::default()) } fn visit_str(self, _: &str) -> Result { - Ok(None) + Ok(PathResult::default()) } fn visit_unit(self) -> Result { - Ok(None) + Ok(PathResult::default()) } fn visit_map>(self, mut map: A) -> Result { let PathSegment::Field(name) = &self.segments[0] else { IgnoredAny.visit_map(map)?; - return Ok(None); + return Ok(PathResult::default()); }; - let mut found = None; - // Every entry is visited so that a duplicated key resolves to its last - // occurrence, as it would in a parsed object. + let mut found = PathResult::default(); while let Some(matched) = map.next_key_seed(KeySeed(name))? { - if matched { - found = map.next_value_seed(PathSeed { + if matched && !found.matched { + let candidate = map.next_value_seed(PathSeed { segments: &self.segments[1..], + reject_direct_null: true, })?; + if candidate.matched { + found = candidate; + } } else { map.next_value::()?; } @@ -372,25 +393,42 @@ impl<'de> Visitor<'de> for SegmentVisitor<'_> { } fn visit_seq>(self, mut seq: A) -> Result { - let PathSegment::Index(idx) = &self.segments[0] else { - IgnoredAny.visit_seq(seq)?; - return Ok(None); - }; - - for _ in 0..*idx { - if seq.next_element::()?.is_none() { - return Ok(None); + match &self.segments[0] { + PathSegment::Index(idx) => { + for _ in 0..*idx { + if seq.next_element::()?.is_none() { + return Ok(PathResult::default()); + } + } + let found = seq + .next_element_seed(PathSeed { + segments: &self.segments[1..], + reject_direct_null: false, + })? + .unwrap_or_default(); + // The remaining elements are still visited, so that a malformed element + // after the match yields no match, as a full parse would. + IgnoredAny.visit_seq(seq)?; + Ok(found) + } + PathSegment::Wildcard => { + let mut found = PathResult::default(); + while let Some(mut result) = seq.next_element_seed(PathSeed { + segments: &self.segments[1..], + reject_direct_null: false, + })? { + if result.matched { + found.matched = true; + found.values.append(&mut result.values); + } + } + Ok(found) + } + PathSegment::Field(_) => { + IgnoredAny.visit_seq(seq)?; + Ok(PathResult::default()) } } - let found = seq - .next_element_seed(PathSeed { - segments: &self.segments[1..], - })? - .flatten(); - // The remaining elements are still visited, so that a malformed element - // after the match yields no match, as a full parse would. - IgnoredAny.visit_seq(seq)?; - Ok(found) } } @@ -417,47 +455,12 @@ impl<'de> Visitor<'de> for KeySeed<'_> { } } -/// Evaluation for paths containing wildcards. -/// Returns references to all matching values. -fn evaluate_with_wildcard<'a>(value: &'a Value, segments: &[PathSegment]) -> Vec<&'a Value> { - if segments.is_empty() { - return vec![value]; - } - - let rest = &segments[1..]; - - match &segments[0] { - PathSegment::Field(name) => match value { - Value::Object(map) => match map.get(name) { - Some(v) => evaluate_with_wildcard(v, rest), - None => vec![], - }, - _ => vec![], - }, - PathSegment::Index(idx) => match value { - Value::Array(arr) => match arr.get(*idx) { - Some(v) => evaluate_with_wildcard(v, rest), - None => vec![], - }, - _ => vec![], - }, - PathSegment::Wildcard => match value { - Value::Array(arr) => arr - .iter() - .flat_map(|v| evaluate_with_wildcard(v, rest)) - .collect(), - _ => vec![], - }, - } -} - /// Convert a JSON value to its string representation matching Spark behavior. /// - Strings are returned without quotes -/// - null returns None +/// - null reached through array traversal is returned as JSON text /// - Numbers, booleans, objects, arrays are serialized as JSON fn value_into_string(value: Value) -> Option { match value { - Value::Null => None, Value::String(s) => Some(s), other => Some(other.to_string()), } @@ -581,14 +584,86 @@ mod tests { } #[test] - fn test_duplicate_key_last_wins() { - // Every entry is visited, so a duplicated key resolves to its last - // occurrence, matching serde_json's preserve_order overwrite behavior. + fn test_duplicate_key_first_wins() { + // The first occurrence of a duplicated key wins, matching Spark. let path = parse_json_path("$.a").unwrap(); assert_eq!( evaluate_path(r#"{"a":1,"a":2}"#, &path), + Some("1".to_string()) + ); + } + + #[test] + fn test_duplicate_key_first_wins_nested() { + // First-wins also applies when recursing into the matched value. + let path = parse_json_path("$.a.b").unwrap(); + assert_eq!( + evaluate_path(r#"{"a":{"b":1},"a":{"b":2}}"#, &path), + Some("1".to_string()) + ); + } + + #[test] + fn test_duplicate_key_first_successful_match_wins() { + // A duplicate field is only locked in after the remaining path + // produces a non-null result. Spark continues to later occurrences + // when an earlier occurrence is null or misses the remaining path. + let path = parse_json_path("$.a.b").unwrap(); + assert_eq!( + evaluate_path(r#"{"a":{"x":1},"a":{"b":2}}"#, &path), Some("2".to_string()) ); + assert_eq!( + evaluate_path(r#"{"a":null,"a":{"b":2}}"#, &path), + Some("2".to_string()) + ); + + let path = parse_json_path("$.a").unwrap(); + assert_eq!( + evaluate_path(r#"{"a":null,"a":2}"#, &path), + Some("2".to_string()) + ); + + let path = parse_json_path("$.a.b").unwrap(); + assert_eq!( + evaluate_path(r#"{"a":{"b":null,"b":2}}"#, &path), + Some("2".to_string()) + ); + } + + #[test] + fn test_duplicate_key_first_successful_match_wins_with_wildcard() { + let path = parse_json_path("$.a[*].b").unwrap(); + assert_eq!( + evaluate_path( + r#"{"a":[{"x":1}],"a":[{"b":2},{"b":3}],"a":[{"b":4}]}"#, + &path + ), + Some("[2,3]".to_string()) + ); + + let path = parse_json_path("$.a[*]").unwrap(); + assert_eq!( + evaluate_path(r#"{"a":[1],"a":[2]}"#, &path), + Some("1".to_string()) + ); + assert_eq!( + evaluate_path(r#"{"a":[],"a":[2]}"#, &path), + Some("2".to_string()) + ); + } + + #[test] + fn test_duplicate_key_null_reached_through_array_locks_match() { + let json = r#"{"a":[null],"a":[2]}"#; + + // Unlike a null directly under a named field, a null reached through + // array traversal is emitted as JSON text and locks that field match. + let path = parse_json_path("$.a[0]").unwrap(); + assert_eq!(evaluate_path(json, &path), Some("null".to_string())); + + let path = parse_json_path("$.a[*]").unwrap(); + assert_eq!(evaluate_path(json, &path), Some("null".to_string())); } #[test]