From fac067f35c030aa2ff65ce41d31e91e39375ef47 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 13 Jul 2026 10:03:51 -0600 Subject: [PATCH 1/2] perf: optimize spark_get_json_object in datafusion-comet-spark-expr --- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/get_json_object.rs | 68 +++++++ .../src/string_funcs/get_json_object.rs | 174 +++++++++++++++--- 3 files changed, 216 insertions(+), 30 deletions(-) create mode 100644 native/spark-expr/benches/get_json_object.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index ef00bef1ea..c1ab8eccb6 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -135,3 +135,7 @@ harness = false [[bench]] name = "parse_url" harness = false + +[[bench]] +name = "get_json_object" +harness = false diff --git a/native/spark-expr/benches/get_json_object.rs b/native/spark-expr/benches/get_json_object.rs new file mode 100644 index 0000000000..12511a025c --- /dev/null +++ b/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 = (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); 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 4d9acc95af..654bf8aabc 100644 --- a/native/spark-expr/src/string_funcs/get_json_object.rs +++ b/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(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 { let json_strings = as_generic_string_array::(json_array)?; let path_strings = as_generic_string_array::(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 { /// 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 { - 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 { 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 { + 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; + + fn deserialize>(self, deserializer: D) -> Result { + 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; + + 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(self, _: bool) -> Result { + Ok(None) + } + + fn visit_i64(self, _: i64) -> Result { + Ok(None) + } + + fn visit_u64(self, _: u64) -> Result { + Ok(None) + } + + fn visit_f64(self, _: f64) -> Result { + Ok(None) + } + + fn visit_str(self, _: &str) -> Result { + Ok(None) + } + + fn visit_unit(self) -> Result { + Ok(None) + } + + fn visit_map>(self, mut map: A) -> Result { + 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::()?; + } } - PathSegment::Index(idx) => { - let child = value.as_array()?.get(*idx)?; - evaluate_no_wildcard(child, &segments[1..]) + Ok(found) + } + + 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); + } } - 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>(self, deserializer: D) -> Result { + 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(self, key: &str) -> Result { + 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 { +fn value_into_string(value: Value) -> Option { 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()), } } From a9df12d09af54110318324ed31303dd036adb056 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 15 Jul 2026 16:22:38 -0600 Subject: [PATCH 2/2] test: cover streaming stop-early parity guards for get_json_object 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) --- .../src/string_funcs/get_json_object.rs | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) 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 654bf8aabc..b61a3172fb 100644 --- a/native/spark-expr/src/string_funcs/get_json_object.rs +++ b/native/spark-expr/src/string_funcs/get_json_object.rs @@ -548,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();