From 179b6ae1eaf4adf51c1c724de1b796ea0d352942 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 12 Jul 2026 09:28:59 -0600 Subject: [PATCH] perf: optimize escape_string in datafusion-comet-spark-expr --- 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(-) create mode 100644 native/spark-expr/benches/to_json.rs 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> = (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 = 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 = 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);