perf: optimize to_json (2x faster)

native/spark-expr/Cargo.toml

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

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);

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);