From f98ac3e68f77e1a83c9078c1c65d7efc031e011f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 13:42:14 -0600 Subject: [PATCH] perf: optimize struct_to_json in datafusion-comet-spark-expr --- native/spark-expr/src/json_funcs/to_json.rs | 157 +++++++++++++------- 1 file changed, 100 insertions(+), 57 deletions(-) diff --git a/native/spark-expr/src/json_funcs/to_json.rs b/native/spark-expr/src/json_funcs/to_json.rs index caf87514fa..d4f4b32a69 100644 --- a/native/spark-expr/src/json_funcs/to_json.rs +++ b/native/spark-expr/src/json_funcs/to_json.rs @@ -29,7 +29,7 @@ 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::fmt::{Debug, Display, Formatter, Write}; use std::hash::Hash; use std::sync::Arc; @@ -187,87 +187,130 @@ fn escape_string(input: &str) -> Cow<'_, str> { Cow::Owned(escaped) } +/// How the rendered value of a column has to be written into the JSON output. +enum FieldStyle { + /// String column: always quoted. + Quoted, + /// The rendering may be `NaN` / `Infinity`, which Spark writes as a quoted string. + MaybeQuoted, + /// The rendering is always written verbatim. + Raw, +} + +/// Peels off a dictionary encoding to get at the type the values are rendered from. +fn value_type(data_type: &DataType) -> &DataType { + match data_type { + DataType::Dictionary(_, dt) => dt.as_ref(), + dt => dt, + } +} + +/// True for types whose string rendering can never be `NaN` or `Infinity`, so the +/// per-value check for those spellings can be skipped entirely. +fn never_renders_special_float(data_type: &DataType) -> bool { + let data_type = value_type(data_type); + data_type.is_null() + || data_type.is_integer() + || data_type.is_decimal() + || data_type.is_temporal() + || matches!(data_type, DataType::Boolean) +} + fn struct_to_json( array: &StructArray, timezone: &str, ignore_null_fields: bool, ) -> Result { - // get field names and escape any quotes - let field_names: Vec = array - .fields() - .iter() - .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 - .fields() - .iter() - .map(|f| match f.data_type() { - DataType::Utf8 | DataType::LargeUtf8 => true, - DataType::Dictionary(_, dt) => { - matches!(dt.as_ref(), DataType::Utf8 | DataType::LargeUtf8) - } - _ => false, - }) - .collect(); // create JSON string representation of each column let string_arrays: Vec = array .columns() .iter() .map(|arr| array_to_json_string(arr, timezone, ignore_null_fields)) .collect::>>()?; - let string_arrays: Vec<&StringArray> = string_arrays + + // pre-render `"field_name":` so that each field costs a single copy rather than one + // per punctuation character, and decide up front how its values have to be written + let fields: Vec<(&StringArray, String, FieldStyle)> = array + .fields() .iter() - .map(|arr| { - arr.as_any() + .zip(string_arrays.iter()) + .map(|(f, values)| { + let values = values + .as_any() .downcast_ref::() - .expect("string array") + .expect("string array"); + let prefix = format!("\"{}\":", escape_string(f.name())); + let style = match value_type(f.data_type()) { + DataType::Utf8 | DataType::LargeUtf8 => FieldStyle::Quoted, + _ if never_renders_special_float(f.data_type()) => FieldStyle::Raw, + _ => FieldStyle::MaybeQuoted, + }; + (values, prefix, style) }) .collect(); - // build the JSON string containing entries in the format `"field_name":field_value` - let mut builder = StringBuilder::with_capacity(array.len(), array.len() * 16); - let mut json = String::with_capacity(array.len() * 16); - for row_index in 0..array.len() { + + // size the output buffer from the rendered columns so that it does not have to grow: + // the braces of each row, plus every field's `,"name":` and its (possibly quoted) value + let num_rows = array.len(); + let data_capacity = 2 * num_rows + + fields + .iter() + .map(|(values, prefix, _)| values.value_data().len() + num_rows * (prefix.len() + 3)) + .sum::(); + + // build the JSON string containing entries in the format `"field_name":field_value`, + // writing straight into the builder rather than staging each row in a temporary + let mut builder = StringBuilder::with_capacity(num_rows, data_capacity); + for row_index in 0..num_rows { if array.is_null(row_index) { builder.append_null(); - } else { - json.clear(); - let mut any_fields_written = false; - json.push('{'); - for col_index in 0..string_arrays.len() { - let is_null = string_arrays[col_index].is_null(row_index); - if is_null && ignore_null_fields { - continue; - } - if any_fields_written { - json.push(','); + continue; + } + let mut any_fields_written = false; + append(&mut builder, "{"); + for (values, prefix, style) in &fields { + let is_null = values.is_null(row_index); + if is_null && ignore_null_fields { + continue; + } + if any_fields_written { + append(&mut builder, ","); + } + any_fields_written = true; + append(&mut builder, prefix); + if is_null { + append(&mut builder, "null"); + continue; + } + let string_value = values.value(row_index); + match style { + FieldStyle::Quoted => { + append(&mut builder, "\""); + append(&mut builder, escape_string(string_value).as_ref()); + append(&mut builder, "\""); } - // quoted field name - json.push('"'); - json.push_str(&field_names[col_index]); - json.push_str("\":"); - if is_null { - json.push_str("null"); - } else { - // value - 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).as_ref()); - json.push('"'); - } else { - json.push_str(string_value); - } + // the special float spellings contain nothing that needs escaping + FieldStyle::MaybeQuoted if is_infinity(string_value) || is_nan(string_value) => { + append(&mut builder, "\""); + append(&mut builder, string_value); + append(&mut builder, "\""); } - any_fields_written = true; + _ => append(&mut builder, string_value), } - json.push('}'); - builder.append_value(&json); } + append(&mut builder, "}"); + // finalizes the value accumulated by the `append` calls above + builder.append_value(""); } Ok(Arc::new(builder.finish())) } +/// Appends to the value currently being built. The builder's `Write` impl is infallible. +#[inline] +fn append(builder: &mut StringBuilder, s: &str) { + let _ = builder.write_str(s); +} + fn is_infinity(input: &str) -> bool { input == "Infinity" || input == "-Infinity" }