perf: optimize spark_get_json_object (4x faster)

native/spark-expr/Cargo.toml

@@ -135,3 +135,7 @@ harness = false
 [[bench]]
 name = "parse_url"
 harness = false
+
+[[bench]]
+name = "get_json_object"
+harness = false

native/spark-expr/benches/get_json_object.rs

@@ -0,0 +1,68 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::StringArray;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::common::ScalarValue;
+use datafusion::logical_expr::ColumnarValue;
+use datafusion_comet_spark_expr::spark_get_json_object;
+use std::hint::black_box;
+use std::sync::Arc;
+
+const N: usize = 8192;
+
+fn json_docs() -> ColumnarValue {
+    let values: Vec<String> = (0..N)
+        .map(|i| {
+            format!(
+                r#"{{"id":{i},"name":"user{i}","email":"user{i}@example.com","tags":["a","b","c"],"address":{{"city":"city{i}","zip":"{i:05}"}},"active":{},"score":{}.5,"notes":"some longer free text field for row {i}"}}"#,
+                i % 2 == 0,
+                i % 100
+            )
+        })
+        .collect();
+    ColumnarValue::Array(Arc::new(StringArray::from(values)))
+}
+
+fn path(p: &str) -> ColumnarValue {
+    ColumnarValue::Scalar(ScalarValue::Utf8(Some(p.to_string())))
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("get_json_object");
+    let docs = json_docs();
+
+    // The common Spark shape: extract one field from each document using a
+    // constant path.
+    for (name, p) in [
+        ("top_level_field", "$.name"),
+        ("nested_field", "$.address.city"),
+        ("array_index", "$.tags[1]"),
+    ] {
+        let path = path(p);
+        group.bench_function(name, |b| {
+            b.iter(|| {
+                black_box(spark_get_json_object(&[docs.clone(), path.clone()]).unwrap());
+            });
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

native/spark-expr/src/string_funcs/get_json_object.rs

@@ -20,7 +20,10 @@ use datafusion::common::{
     cast::as_generic_string_array, exec_err, Result as DataFusionResult, ScalarValue,
 };
 use datafusion::logical_expr::ColumnarValue;
+use serde::de::{DeserializeSeed, IgnoredAny, MapAccess, SeqAccess, Visitor};
+use serde::{Deserialize, Deserializer};
 use serde_json::Value;
+use std::fmt;
 use std::sync::Arc;
 
 /// Extracts a string from a ScalarValue, returning Ok(None) for null values.
@@ -69,7 +72,7 @@ pub fn spark_get_json_object(args: &[ColumnarValue]) -> DataFusionResult<Columna
             };
 
             let json_strings = as_generic_string_array::<i32>(json_array)?;
-            let mut builder = StringBuilder::new();
+            let mut builder = StringBuilder::with_capacity(json_strings.len(), 0);
 
             for i in 0..json_strings.len() {
                 if json_strings.is_null(i) {
@@ -108,7 +111,7 @@ pub fn spark_get_json_object(args: &[ColumnarValue]) -> DataFusionResult<Columna
         (ColumnarValue::Array(json_array), ColumnarValue::Array(path_array)) => {
             let json_strings = as_generic_string_array::<i32>(json_array)?;
             let path_strings = as_generic_string_array::<i32>(path_array)?;
-            let mut builder = StringBuilder::new();
+            let mut builder = StringBuilder::with_capacity(json_strings.len(), 0);
 
             for i in 0..json_strings.len() {
                 if json_strings.is_null(i) || path_strings.is_null(i) {
@@ -252,14 +255,12 @@ fn parse_json_path(path: &str) -> Option<ParsedPath> {
 /// 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<String> {
-    let value: Value = serde_json::from_str(json_str).ok()?;
-
     if !path.has_wildcard {
-        // Fast path: no wildcards, no Vec allocations
-        let result = evaluate_no_wildcard(&value, &path.segments)?;
-        return value_to_string(result);
+        return value_into_string(extract_no_wildcard(json_str, &path.segments)?);
     }
 
+    let value: Value = serde_json::from_str(json_str).ok()?;
+
     // Wildcard path: may return multiple results
     let results = evaluate_with_wildcard(&value, &path.segments);
 
@@ -268,38 +269,151 @@ fn evaluate_path(json_str: &str, path: &ParsedPath) -> Option<String> {
         1 => {
             // Single wildcard match: Spark preserves JSON serialization format
             // (strings keep their quotes, numbers don't)
-            let s = serde_json::to_string(results[0]).ok()?;
-            if s == "null" {
+            if results[0].is_null() {
                 None
             } else {
-                Some(s)
+                serde_json::to_string(results[0]).ok()
             }
         }
-        _ => {
-            // Multiple results: wrap in JSON array
-            let arr = Value::Array(results.into_iter().cloned().collect());
-            Some(arr.to_string())
+        // 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(),
+    }
+}
+
+/// 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<Value> {
+    let mut de = serde_json::Deserializer::from_str(json_str);
+    let found = PathSeed { segments }.deserialize(&mut de).ok()?;
+    de.end().ok()?;
+    found
+}
+
+/// Deserializes the value at `segments`, discarding everything else.
+struct PathSeed<'a> {
+    segments: &'a [PathSegment],
+}
+
+impl<'de> DeserializeSeed<'de> for PathSeed<'_> {
+    type Value = Option<Value>;
+
+    fn deserialize<D: Deserializer<'de>>(self, deserializer: D) -> Result<Self::Value, D::Error> {
+        if self.segments.is_empty() {
+            return Value::deserialize(deserializer).map(Some);
         }
+        deserializer.deserialize_any(SegmentVisitor {
+            segments: self.segments,
+        })
     }
 }
 
-/// Fast-path evaluation for paths without wildcards.
-/// Returns a reference to the matched value, or None if no match.
-fn evaluate_no_wildcard<'a>(value: &'a Value, segments: &[PathSegment]) -> Option<&'a Value> {
-    if segments.is_empty() {
-        return Some(value);
+/// Applies `segments[0]` to the value being visited. A value whose shape does
+/// not match the segment (an index into an object, say) is skipped and reported
+/// as no match rather than as an error, matching a lookup on a parsed document.
+struct SegmentVisitor<'a> {
+    segments: &'a [PathSegment],
+}
+
+impl<'de> Visitor<'de> for SegmentVisitor<'_> {
+    type Value = Option<Value>;
+
+    fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("a JSON value")
     }
 
-    match &segments[0] {
-        PathSegment::Field(name) => {
-            let child = value.as_object()?.get(name)?;
-            evaluate_no_wildcard(child, &segments[1..])
+    fn visit_bool<E>(self, _: bool) -> Result<Self::Value, E> {
+        Ok(None)
+    }
+
+    fn visit_i64<E>(self, _: i64) -> Result<Self::Value, E> {
+        Ok(None)
+    }
+
+    fn visit_u64<E>(self, _: u64) -> Result<Self::Value, E> {
+        Ok(None)
+    }
+
+    fn visit_f64<E>(self, _: f64) -> Result<Self::Value, E> {
+        Ok(None)
+    }
+
+    fn visit_str<E>(self, _: &str) -> Result<Self::Value, E> {
+        Ok(None)
+    }
+
+    fn visit_unit<E>(self) -> Result<Self::Value, E> {
+        Ok(None)
+    }
+
+    fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
+        let PathSegment::Field(name) = &self.segments[0] else {
+            IgnoredAny.visit_map(map)?;
+            return Ok(None);
+        };
+
+        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.
+        while let Some(matched) = map.next_key_seed(KeySeed(name))? {
+            if matched {
+                found = map.next_value_seed(PathSeed {
+                    segments: &self.segments[1..],
+                })?;
+            } else {
+                map.next_value::<IgnoredAny>()?;
+            }
         }
-        PathSegment::Index(idx) => {
-            let child = value.as_array()?.get(*idx)?;
-            evaluate_no_wildcard(child, &segments[1..])
+        Ok(found)
+    }
+
+    fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
+        let PathSegment::Index(idx) = &self.segments[0] else {
+            IgnoredAny.visit_seq(seq)?;
+            return Ok(None);
+        };
+
+        for _ in 0..*idx {
+            if seq.next_element::<IgnoredAny>()?.is_none() {
+                return Ok(None);
+            }
         }
-        PathSegment::Wildcard => unreachable!("wildcard in no-wildcard path"),
+        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)
+    }
+}
+
+/// Compares an object key against a field name without allocating it.
+struct KeySeed<'a>(&'a str);
+
+impl<'de> DeserializeSeed<'de> for KeySeed<'_> {
+    type Value = bool;
+
+    fn deserialize<D: Deserializer<'de>>(self, deserializer: D) -> Result<bool, D::Error> {
+        deserializer.deserialize_str(self)
+    }
+}
+
+impl<'de> Visitor<'de> for KeySeed<'_> {
+    type Value = bool;
+
+    fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("an object key")
+    }
+
+    fn visit_str<E>(self, key: &str) -> Result<bool, E> {
+        Ok(key == self.0)
     }
 }
 
@@ -341,11 +455,11 @@ fn evaluate_with_wildcard<'a>(value: &'a Value, segments: &[PathSegment]) -> Vec
 /// - Strings are returned without quotes
 /// - null returns None
 /// - Numbers, booleans, objects, arrays are serialized as JSON
-fn value_to_string(value: &Value) -> Option<String> {
+fn value_into_string(value: Value) -> Option<String> {
     match value {
         Value::Null => None,
-        Value::String(s) => Some(s.clone()),
-        _ => Some(value.to_string()),
+        Value::String(s) => Some(s),
+        other => Some(other.to_string()),
     }
 }
 
@@ -434,6 +548,49 @@ mod tests {
         assert_eq!(evaluate_path(r#"{"a":null}"#, &path), None);
     }
 
+    #[test]
+    fn test_null_midpath_is_null() {
+        // A null encountered before the path is exhausted has no child to
+        // recurse into, so the lookup yields no match (SegmentVisitor::visit_unit).
+        let path = parse_json_path("$.a.b").unwrap();
+        assert_eq!(evaluate_path(r#"{"a":null}"#, &path), None);
+    }
+
+    #[test]
+    fn test_match_then_trailing_garbage_is_null() {
+        // The match is found early, but trailing content makes the document
+        // malformed; the full document must still be validated (de.end()).
+        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() {
+        // "a" matches early, but the "b" sibling is malformed; visit_map must
+        // keep draining entries after a match so the parse still rejects this.
+        let path = parse_json_path("$.a").unwrap();
+        assert_eq!(evaluate_path(r#"{"a":1,"b":}"#, &path), None);
+    }
+
+    #[test]
+    fn test_array_match_then_malformed_element_is_null() {
+        // The element at index 0 matches, but a later element is malformed;
+        // visit_seq must keep draining elements after a match.
+        let path = parse_json_path("$[0]").unwrap();
+        assert_eq!(evaluate_path(r#"[1,,]"#, &path), None);
+    }
+
+    #[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.
+        let path = parse_json_path("$.a").unwrap();
+        assert_eq!(
+            evaluate_path(r#"{"a":1,"a":2}"#, &path),
+            Some("2".to_string())
+        );
+    }
+
     #[test]
     fn test_evaluate_missing_field() {
         let path = parse_json_path("$.c").unwrap();