From a3c712111425bfd757bbdfdc0037caf1643ac3ef Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 05:13:41 -0600 Subject: [PATCH] perf: optimize spark_cast_float64_to_utf8 in datafusion-comet-spark-expr --- native/spark-expr/Cargo.toml | 4 + .../benches/cast_float_to_string.rs | 101 +++++++++++++++++ .../src/conversion_funcs/numeric.rs | 102 ++++++++++-------- 3 files changed, 163 insertions(+), 44 deletions(-) create mode 100644 native/spark-expr/benches/cast_float_to_string.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..186f2b5575 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -139,3 +139,7 @@ harness = false [[bench]] name = "to_json" harness = false + +[[bench]] +name = "cast_float_to_string" +harness = false diff --git a/native/spark-expr/benches/cast_float_to_string.rs b/native/spark-expr/benches/cast_float_to_string.rs new file mode 100644 index 0000000000..deb24a3e04 --- /dev/null +++ b/native/spark-expr/benches/cast_float_to_string.rs @@ -0,0 +1,101 @@ +// 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::builder::{Float32Builder, Float64Builder}; +use arrow::array::ArrayRef; +use arrow::datatypes::DataType; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::{spark_cast, EvalMode, SparkCastOptions}; +use std::hint::black_box; +use std::sync::Arc; + +/// Mix of values that exercise both the plain and the scientific notation paths. +fn create_f64_array(size: usize) -> ArrayRef { + let mut builder = Float64Builder::with_capacity(size); + for i in 0..size { + match i % 10 { + 0 => builder.append_null(), + 1 => builder.append_value(0.0), + 2 => builder.append_value(i as f64), + 3 => builder.append_value(i as f64 + 0.125), + 4 => builder.append_value(-(i as f64) * 1.5), + 5 => builder.append_value(1e20 * i as f64), + 6 => builder.append_value(1e-12 * i as f64), + 7 => builder.append_value(f64::NAN), + 8 => builder.append_value(1234.5678), + _ => builder.append_value(-0.001234), + } + } + Arc::new(builder.finish()) +} + +fn create_f32_array(size: usize) -> ArrayRef { + let mut builder = Float32Builder::with_capacity(size); + for i in 0..size { + match i % 10 { + 0 => builder.append_null(), + 1 => builder.append_value(0.0), + 2 => builder.append_value(i as f32), + 3 => builder.append_value(i as f32 + 0.125), + 4 => builder.append_value(-(i as f32) * 1.5), + 5 => builder.append_value(1e20 * i as f32), + 6 => builder.append_value(1e-12 * i as f32), + 7 => builder.append_value(f32::NAN), + 8 => builder.append_value(1234.5678), + _ => builder.append_value(-0.001234), + } + } + Arc::new(builder.finish()) +} + +fn criterion_benchmark(c: &mut Criterion) { + let size = 8192; + let f64_array = create_f64_array(size); + let f32_array = create_f32_array(size); + let cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false); + + let mut group = c.benchmark_group("cast_float_to_string"); + group.bench_function("cast_f64_to_utf8", |b| { + b.iter(|| { + black_box( + spark_cast( + ColumnarValue::Array(Arc::clone(&f64_array)), + &DataType::Utf8, + &cast_options, + ) + .unwrap(), + ) + }) + }); + group.bench_function("cast_f32_to_utf8", |b| { + b.iter(|| { + black_box( + spark_cast( + ColumnarValue::Array(Arc::clone(&f32_array)), + &DataType::Utf8, + &cast_options, + ) + .unwrap(), + ) + }) + }); + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs index bba34d0368..b0e8712312 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -20,7 +20,7 @@ use crate::conversion_funcs::utils::MICROS_PER_SECOND; use crate::{EvalMode, SparkError, SparkResult}; use arrow::array::{ Array, ArrayRef, AsArray, BooleanBuilder, Decimal128Array, Decimal128Builder, Float32Array, - Float64Array, GenericStringArray, Int16Array, Int32Array, Int64Array, Int8Array, + Float64Array, GenericStringBuilder, Int16Array, Int32Array, Int64Array, Int8Array, OffsetSizeTrait, PrimitiveArray, StringBuilder, TimestampMicrosecondBuilder, }; use arrow::datatypes::{ @@ -142,6 +142,8 @@ macro_rules! cast_float_to_string { ) -> SparkResult where OffsetSize: OffsetSizeTrait, { + use std::fmt::Write; + let array = from.as_any().downcast_ref::<$output_type>().unwrap(); // If the absolute number is less than 10,000,000 and greater or equal than 0.001, the @@ -155,53 +157,65 @@ macro_rules! cast_float_to_string { const LOWER_SCIENTIFIC_BOUND: $type = 0.001; const UPPER_SCIENTIFIC_BOUND: $type = 10000000.0; - let output_array = array - .iter() - .map(|value| match value { - Some(value) if value == <$type>::INFINITY => Ok(Some("Infinity".to_string())), - Some(value) if value == <$type>::NEG_INFINITY => Ok(Some("-Infinity".to_string())), - Some(value) - if (value.abs() < UPPER_SCIENTIFIC_BOUND - && value.abs() >= LOWER_SCIENTIFIC_BOUND) - || value.abs() == 0.0 => - { - let trailing_zero = if value.fract() == 0.0 { ".0" } else { "" }; - - Ok(Some(format!("{value}{trailing_zero}"))) + // Values are formatted straight into the builder, so no intermediate String + // is allocated per row. + let mut builder = GenericStringBuilder::::with_capacity( + array.len(), + array.len() * 8, + ); + // Reused across rows by the scientific-notation path, which has to inspect + // the formatted text before emitting it. + let mut scratch = String::with_capacity(32); + + for value in array.iter() { + let Some(value) = value else { + builder.append_null(); + continue; + }; + let abs = value.abs(); + if (LOWER_SCIENTIFIC_BOUND..UPPER_SCIENTIFIC_BOUND).contains(&abs) + || abs == 0.0 + { + let _ = write!(builder, "{value}"); + if value.fract() == 0.0 { + // Spark always renders a fractional digit; Rust omits it. + let _ = builder.write_str(".0"); } - Some(value) - if value.abs() >= UPPER_SCIENTIFIC_BOUND - || value.abs() < LOWER_SCIENTIFIC_BOUND => - { - // Spark uses Java's Float.MIN_VALUE / Double.MIN_VALUE strings for - // the smallest subnormal values; Rust's formatter rounds them more. - if value.abs().to_bits() == 1 { - let sign = if value.is_sign_negative() { "-" } else { "" }; - Ok(Some(format!("{sign}{}", $min_value))) - } else { - let formatted = format!("{value:E}"); - - if formatted.contains(".") { - Ok(Some(formatted)) - } else { - // `formatted` is already in scientific notation and can be split up by E - // in order to add the missing trailing 0 which gets removed for numbers with a fraction of 0.0 - let prepare_number: Vec<&str> = formatted.split("E").collect(); - - let coefficient = prepare_number[0]; - - let exponent = prepare_number[1]; - - Ok(Some(format!("{coefficient}.0E{exponent}"))) - } + builder.append_value(""); + } else if !value.is_finite() { + // NaN and the infinities are excluded by the range check above. + builder.append_value(if value.is_nan() { + "NaN" + } else if value.is_sign_positive() { + "Infinity" + } else { + "-Infinity" + }); + } else if abs.to_bits() == 1 { + // Java's Double.toString / Float.toString are not shortest-roundtrip + // and render the smallest subnormals with more digits than Rust does. + builder.append_value(if value.is_sign_negative() { + concat!("-", $min_value) + } else { + $min_value + }); + } else { + scratch.clear(); + let _ = write!(scratch, "{value:E}"); + match scratch.split_once('E') { + Some((coefficient, exponent)) if !coefficient.contains('.') => { + // Spark keeps the fractional digit Rust drops from a whole + // coefficient. + let _ = builder.write_str(coefficient); + let _ = builder.write_str(".0E"); + builder.append_value(exponent); } + _ => builder.append_value(&scratch), } - Some(value) => Ok(Some(value.to_string())), - _ => Ok(None), - }) - .collect::, SparkError>>()?; + } + } - Ok(Arc::new(output_array)) + Ok(Arc::new(builder.finish())) } cast::<$offset_type>($from, $eval_mode)