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 dd30f411d0 perf: optimize escape_string in datafusion-comet-spark-expr
(#4902)
dd30f411d0 is described below
commit dd30f411d0fc704c1c5f576299fd0a424aff4624
Author: Andy Grove <[email protected]>
AuthorDate: Mon Jul 13 15:24:14 2026 -0600
perf: optimize escape_string in datafusion-comet-spark-expr (#4902)
---
native/spark-expr/Cargo.toml | 4 ++
native/spark-expr/benches/to_json.rs | 77 +++++++++++++++++++++++++
native/spark-expr/src/json_funcs/to_json.rs | 87 ++++++++++++++---------------
3 files changed, 122 insertions(+), 46 deletions(-)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index ef00bef1ea..6f90181ee0 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 = "to_json"
+harness = false
diff --git a/native/spark-expr/benches/to_json.rs
b/native/spark-expr/benches/to_json.rs
new file mode 100644
index 0000000000..0c28cd4653
--- /dev/null
+++ b/native/spark-expr/benches/to_json.rs
@@ -0,0 +1,77 @@
+// 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::{Array, ArrayRef, StringArray, StructArray};
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_plan::expressions::Column;
+use datafusion_comet_spark_expr::ToJson;
+use std::hint::black_box;
+use std::sync::Arc;
+
+fn create_batch(array_size: usize) -> RecordBatch {
+ // A struct of string fields exercises the per-row escaping path, which is
the
+ // dominant per-element cost of to_json for string-heavy structs.
+ let samples = [
+ "",
+ "a",
+ "abc",
+ "hello world",
+ "some longer string value here",
+ ];
+ let build = |offset: usize| -> ArrayRef {
+ let vals: Vec<Option<String>> = (0..array_size)
+ .map(|i| {
+ if i % 10 == 0 {
+ None
+ } else {
+ Some(samples[(i + offset) % samples.len()].to_string())
+ }
+ })
+ .collect();
+ Arc::new(StringArray::from(vals))
+ };
+
+ let f1 = build(0);
+ let f2 = build(2);
+ let f3 = build(4);
+ let struct_array = StructArray::from(vec![
+ (Arc::new(Field::new("f1", DataType::Utf8, true)), f1),
+ (Arc::new(Field::new("f2", DataType::Utf8, true)), f2),
+ (Arc::new(Field::new("f3", DataType::Utf8, true)), f3),
+ ]);
+
+ let schema = Schema::new(vec![Field::new(
+ "s",
+ struct_array.data_type().clone(),
+ true,
+ )]);
+ RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(struct_array)]).unwrap()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let batch = create_batch(8192);
+ let expr = ToJson::new(Arc::new(Column::new("s", 0)), "UTC", false);
+ c.bench_function("to_json", |b| {
+ b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap()))
+ });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/src/json_funcs/to_json.rs
b/native/spark-expr/src/json_funcs/to_json.rs
index 016b1315ac..caf87514fa 100644
--- a/native/spark-expr/src/json_funcs/to_json.rs
+++ b/native/spark-expr/src/json_funcs/to_json.rs
@@ -28,6 +28,7 @@ use datafusion::common::Result;
use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_plan::ColumnarValue;
use std::any::Any;
+use std::borrow::Cow;
use std::fmt::{Debug, Display, Formatter};
use std::hash::Hash;
use std::sync::Arc;
@@ -144,52 +145,46 @@ fn array_to_json_string(
}
}
-fn escape_string(input: &str) -> String {
- let mut escaped_string = String::with_capacity(input.len());
- let mut is_escaped = false;
- for c in input.chars() {
- match c {
- '\"' | '\\' if !is_escaped => {
- escaped_string.push('\\');
- escaped_string.push(c);
- is_escaped = false;
- }
- '\t' => {
- escaped_string.push('\\');
- escaped_string.push('t');
- is_escaped = false;
- }
- '\r' => {
- escaped_string.push('\\');
- escaped_string.push('r');
- is_escaped = false;
- }
- '\n' => {
- escaped_string.push('\\');
- escaped_string.push('n');
- is_escaped = false;
- }
- '\x0C' => {
- escaped_string.push('\\');
- escaped_string.push('f');
- is_escaped = false;
- }
- '\x08' => {
- escaped_string.push('\\');
- escaped_string.push('b');
- is_escaped = false;
- }
- '\\' => {
- escaped_string.push('\\');
- is_escaped = true;
- }
- _ => {
- escaped_string.push(c);
- is_escaped = false;
- }
+/// Returns the JSON escape sequence for a byte that requires escaping, or
`None`
+/// otherwise. Every escaped character is ASCII, so scanning the input as
bytes is
+/// safe: a UTF-8 continuation or lead byte (>= 0x80) can never match here, so
run
+/// boundaries always fall on `char` boundaries.
+#[inline]
+fn escape_replacement(b: u8) -> Option<&'static str> {
+ match b {
+ b'"' => Some("\\\""),
+ b'\\' => Some("\\\\"),
+ b'\t' => Some("\\t"),
+ b'\r' => Some("\\r"),
+ b'\n' => Some("\\n"),
+ 0x0C => Some("\\f"),
+ 0x08 => Some("\\b"),
+ _ => None,
+ }
+}
+
+/// Escapes the characters that Spark's JSON writer escapes (`"`, `\`, and the
+/// `\t`/`\r`/`\n`/`\f`/`\b` control characters). Returns the input unchanged
+/// (without allocating) when nothing needs escaping, which is the common case,
+/// and otherwise copies unescaped runs in bulk rather than character by
character.
+fn escape_string(input: &str) -> Cow<'_, str> {
+ let bytes = input.as_bytes();
+ let first = match bytes.iter().position(|&b|
escape_replacement(b).is_some()) {
+ None => return Cow::Borrowed(input),
+ Some(pos) => pos,
+ };
+
+ let mut escaped = String::with_capacity(input.len() + 8);
+ let mut run_start = 0;
+ for i in first..bytes.len() {
+ if let Some(replacement) = escape_replacement(bytes[i]) {
+ escaped.push_str(&input[run_start..i]);
+ escaped.push_str(replacement);
+ run_start = i + 1;
}
}
- escaped_string
+ escaped.push_str(&input[run_start..]);
+ Cow::Owned(escaped)
}
fn struct_to_json(
@@ -201,7 +196,7 @@ fn struct_to_json(
let field_names: Vec<String> = array
.fields()
.iter()
- .map(|f| escape_string(f.name().as_str()))
+ .map(|f| escape_string(f.name().as_str()).into_owned())
.collect();
// determine which fields need to have their values quoted
let is_string: Vec<bool> = array
@@ -258,7 +253,7 @@ fn struct_to_json(
let string_value =
string_arrays[col_index].value(row_index);
if is_string[col_index] || is_infinity(string_value) ||
is_nan(string_value) {
json.push('"');
- json.push_str(&escape_string(string_value));
+ json.push_str(escape_string(string_value).as_ref());
json.push('"');
} else {
json.push_str(string_value);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
