You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/07/16 21:34:06 UTC
(datafusion-comet) branch main updated: perf: optimize `spark_get_json_object` (4x faster) (#4907)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new eb24be133b perf: optimize `spark_get_json_object` (4x faster) (#4907)
eb24be133b is described below
commit eb24be133b6c381a9d38bbad467e50e97d54094f
Author: Andy Grove <ag...@apache.org>
AuthorDate: Thu Jul 16 15:34:02 2026 -0600
perf: optimize `spark_get_json_object` (4x faster) (#4907)
* perf: optimize spark_get_json_object in datafusion-comet-spark-expr
* 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)
---
native/spark-expr/Cargo.toml | 4 +
native/spark-expr/benches/get_json_object.rs | 68 +++++++
.../spark-expr/src/string_funcs/get_json_object.rs | 217 ++++++++++++++++++---
3 files changed, 259 insertions(+), 30 deletions(-)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 586a69353e..bf6d43ee58 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -136,6 +136,10 @@ harness = false
name = "parse_url"
harness = false
+[[bench]]
+name = "get_json_object"
+harness = false
+
[[bench]]
name = "arrays_overlap"
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<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);
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..b61a3172fb 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<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();
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org