diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..8a0de1d4d3 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_decimal_to_string" +harness = false diff --git a/native/spark-expr/benches/cast_decimal_to_string.rs b/native/spark-expr/benches/cast_decimal_to_string.rs new file mode 100644 index 0000000000..d0629300a9 --- /dev/null +++ b/native/spark-expr/benches/cast_decimal_to_string.rs @@ -0,0 +1,83 @@ +// 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::{Decimal128Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_expr::{expressions::Column, PhysicalExpr}; +use datafusion_comet_spark_expr::{Cast, EvalMode, SparkCastOptions}; +use std::sync::Arc; + +const NUM_ROWS: usize = 4096; + +fn create_decimal_batch(precision: u8, scale: i8, unscaled: impl Fn(usize) -> i128) -> RecordBatch { + let values: Vec> = (0..NUM_ROWS) + .map(|i| if i % 10 == 9 { None } else { Some(unscaled(i)) }) + .collect(); + let array = Decimal128Array::from(values) + .with_precision_and_scale(precision, scale) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Decimal128(precision, scale), + true, + )])); + RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap() +} + +fn cast_to_utf8() -> Cast { + Cast::new( + Arc::new(Column::new("a", 0)), + DataType::Utf8, + SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + None, + None, + ) +} + +fn criterion_benchmark(c: &mut Criterion) { + // plain notation, fraction shorter than the integer part + let small_scale = create_decimal_batch(18, 2, |i| (i as i128 * 7919) - 1_000_000); + // plain notation, leading zeroes in the fraction + let leading_zeroes = create_decimal_batch(18, 6, |i| (i as i128 % 97) + 1); + // wide coefficients spanning more than one 64-bit digit group + let wide = create_decimal_batch(38, 10, |i| { + i128::from(i as i64 + 1) * 1_234_567_890_123_456_789i128 * 1_000_000_007i128 + }); + // scientific notation (adjusted exponent < -6) + let scientific = create_decimal_batch(38, 30, |i| (i as i128 % 13) + 1); + + let cast = cast_to_utf8(); + + let mut group = c.benchmark_group("cast_decimal_to_string"); + group.bench_function("decimal128_scale_2", |b| { + b.iter(|| cast.evaluate(&small_scale).unwrap()); + }); + group.bench_function("decimal128_leading_zeroes", |b| { + b.iter(|| cast.evaluate(&leading_zeroes).unwrap()); + }); + group.bench_function("decimal128_wide", |b| { + b.iter(|| cast.evaluate(&wide).unwrap()); + }); + group.bench_function("decimal128_scientific", |b| { + b.iter(|| cast.evaluate(&scientific).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..6851c27b7c 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -596,14 +596,15 @@ pub(crate) fn cast_decimal128_to_utf8(array: &ArrayRef, scale: i8) -> SparkResul .downcast_ref::() .expect("Expected a Decimal128Array"); let mut builder = StringBuilder::with_capacity(decimal_array.len(), decimal_array.len() * 16); - // Reuse a single String buffer across rows to avoid one allocation per value. - let mut buf = String::with_capacity(40); + // Reuse the output and digit buffers across rows to keep the loop allocation-free. + let mut buf = String::with_capacity(48); + let mut digits = [0u8; MAX_COEFF_DIGITS]; for opt_val in decimal_array.iter() { match opt_val { None => builder.append_null(), Some(unscaled) => { buf.clear(); - decimal128_to_java_string(unscaled, scale, &mut buf); + decimal128_to_java_string(unscaled, scale, &mut digits, &mut buf); builder.append_value(&buf); } } @@ -611,42 +612,95 @@ pub(crate) fn cast_decimal128_to_utf8(array: &ArrayRef, scale: i8) -> SparkResul Ok(Arc::new(builder.finish())) } +/// The largest power of ten that fits in a `u64`. +const POW10_19: u128 = 10_000_000_000_000_000_000; + +/// Digits in the largest `u128`, which bounds the coefficient of any Decimal128. +const MAX_COEFF_DIGITS: usize = 39; + +/// Renders the base-10 digits of `value` into `buf`, returning them without leading zeroes +/// (except for `value == 0`, which renders as `"0"`). +fn render_digits(mut value: u128, buf: &mut [u8; MAX_COEFF_DIGITS]) -> &str { + let mut pos = buf.len(); + + // Peel off 19 digits at a time so that all but the final group is produced with 64-bit + // arithmetic, which is much cheaper than repeated 128-bit division. Every group taken here + // has a non-zero quotient above it, so emitting its leading zeroes is correct. + while value >= POW10_19 { + let mut group = (value % POW10_19) as u64; + value /= POW10_19; + for _ in 0..19 { + pos -= 1; + buf[pos] = b'0' + (group % 10) as u8; + group /= 10; + } + } + + let mut rest = value as u64; + loop { + pos -= 1; + buf[pos] = b'0' + (rest % 10) as u8; + rest /= 10; + if rest == 0 { + break; + } + } + + // Only ASCII digits were written above. + std::str::from_utf8(&buf[pos..]).expect("ascii digits") +} + /// Formats a Decimal128 unscaled value into `out` matching Java's BigDecimal.toString(): /// - Plain notation when scale >= 0 and adjusted_exponent >= -6 /// - Scientific notation otherwise /// /// adjusted_exponent = -scale + (numDigits - 1) -fn decimal128_to_java_string(unscaled: i128, scale: i8, out: &mut String) { +/// +/// `digits` is scratch space for the coefficient; the caller reuses one buffer across rows. +fn decimal128_to_java_string( + unscaled: i128, + scale: i8, + digits: &mut [u8; MAX_COEFF_DIGITS], + out: &mut String, +) { use std::fmt::Write; - let negative = unscaled < 0; - let sign = if negative { "-" } else { "" }; - let coeff = unscaled.unsigned_abs().to_string(); - let num_digits = coeff.len() as i64; - let adj_exp = -(scale as i64) + (num_digits - 1); + let coeff = render_digits(unscaled.unsigned_abs(), digits); + let num_digits = coeff.len(); + let adj_exp = -(scale as i64) + (num_digits as i64 - 1); + + if unscaled < 0 { + out.push('-'); + } if scale >= 0 && adj_exp >= -6 { - let scale_u = scale as usize; - let num_digits_u = num_digits as usize; - if scale_u == 0 { - write!(out, "{sign}{coeff}").unwrap(); - } else if num_digits_u > scale_u { - let (int_part, frac_part) = coeff.split_at(num_digits_u - scale_u); - write!(out, "{sign}{int_part}.{frac_part}").unwrap(); + let scale = scale as usize; + if scale == 0 { + out.push_str(coeff); + } else if num_digits > scale { + let (int_part, frac_part) = coeff.split_at(num_digits - scale); + out.push_str(int_part); + out.push('.'); + out.push_str(frac_part); } else { - let leading = scale_u - num_digits_u; - write!(out, "{sign}0.{}{coeff}", "0".repeat(leading)).unwrap(); + out.push_str("0."); + for _ in 0..scale - num_digits { + out.push('0'); + } + out.push_str(coeff); } } else { if num_digits > 1 { - write!(out, "{sign}{}.{}", &coeff[..1], &coeff[1..]).unwrap(); + out.push_str(&coeff[..1]); + out.push('.'); + out.push_str(&coeff[1..]); } else { - write!(out, "{sign}{coeff}").unwrap(); + out.push_str(coeff); } + out.push('E'); if adj_exp > 0 { - write!(out, "E+{adj_exp}").unwrap(); - } else { - write!(out, "E{adj_exp}").unwrap(); + out.push('+'); } + write!(out, "{adj_exp}").unwrap(); } } @@ -1417,7 +1471,7 @@ mod tests { fn test_decimal128_to_java_string() { fn fmt(unscaled: i128, scale: i8) -> String { let mut buf = String::new(); - decimal128_to_java_string(unscaled, scale, &mut buf); + decimal128_to_java_string(unscaled, scale, &mut [0u8; MAX_COEFF_DIGITS], &mut buf); buf } // scale >= 0, adj_exp >= -6 → plain notation @@ -1441,5 +1495,14 @@ mod tests { assert_eq!(fmt(1, -2), "1E+2"); assert_eq!(fmt(123, -2), "1.23E+4"); assert_eq!(fmt(-123, -2), "-1.23E+4"); + + // values needing more than one 64-bit digit group + assert_eq!(fmt(10_000_000_000_000_000_000, 0), "10000000000000000000"); + assert_eq!(fmt(i128::MAX, 0), i128::MAX.to_string()); + assert_eq!(fmt(i128::MIN, 0), i128::MIN.to_string()); + assert_eq!( + fmt(99_999_999_999_999_999_999_999_999_999_999_999_999, 10), + "9999999999999999999999999999.9999999999" + ); } }