From 5c544bcd6f5a77163fd63292933bb903cc6bf8ad Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 04:29:11 -0600 Subject: [PATCH] perf: optimize parse_string_to_decimal in datafusion-comet-spark-expr --- native/spark-expr/Cargo.toml | 4 + .../benches/cast_string_to_decimal.rs | 93 +++++++++ .../spark-expr/src/conversion_funcs/string.rs | 192 +++++++++++------- 3 files changed, 219 insertions(+), 70 deletions(-) create mode 100644 native/spark-expr/benches/cast_string_to_decimal.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..dd8d3ddeee 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_string_to_decimal" +harness = false diff --git a/native/spark-expr/benches/cast_string_to_decimal.rs b/native/spark-expr/benches/cast_string_to_decimal.rs new file mode 100644 index 0000000000..09ab5efe3d --- /dev/null +++ b/native/spark-expr/benches/cast_string_to_decimal.rs @@ -0,0 +1,93 @@ +// 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::StringBuilder, 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 rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; +use std::hint::black_box; +use std::sync::Arc; + +/// A batch of decimal strings covering the shapes Spark sees in practice: plain +/// integers, fixed-point values, negatives, and scientific notation. +fn create_decimal_string_batch(size: usize) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)])); + let mut rng = StdRng::seed_from_u64(42); + let mut b = StringBuilder::new(); + for i in 0..size { + if i % 10 == 0 { + b.append_null(); + } else { + match i % 5 { + 0 => b.append_value(format!( + "{}.{}", + rng.random_range(0..1_000_000u32), + rng.random_range(0..100_000u32) + )), + 1 => b.append_value(format!( + "{}.{}E{}", + rng.random_range(0..10u32), + rng.random_range(0..100u32), + rng.random_range(0..10u32) + )), + 2 => b.append_value(format!( + "-{}.{}", + rng.random_range(0..1_000_000u32), + rng.random_range(0..100_000u32) + )), + 3 => b.append_value(format!("{}", rng.random_range(-1_000_000..1_000_000i32))), + _ => b.append_value(format!("0.{:05}", rng.random_range(0..100_000u32))), + } + } + } + RecordBatch::try_new(schema, vec![Arc::new(b.finish())]).unwrap() +} + +fn criterion_benchmark(c: &mut Criterion) { + let batch = create_decimal_string_batch(8192); + let expr = Arc::new(Column::new("a", 0)); + + for (mode, mode_name) in [ + (EvalMode::Legacy, "legacy"), + (EvalMode::Ansi, "ansi"), + (EvalMode::Try, "try"), + ] { + let mut group = c.benchmark_group(format!("cast_string_to_decimal/{mode_name}")); + for (data_type, name) in [ + (DataType::Decimal128(38, 10), "decimal_38_10"), + (DataType::Decimal128(18, 2), "decimal_18_2"), + ] { + let cast = Cast::new( + expr.clone(), + data_type, + SparkCastOptions::new(mode, "", false), + None, + None, + ); + group.bench_function(name, |b| { + b.iter(|| black_box(cast.evaluate(&batch).unwrap())); + }); + } + group.finish(); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/conversion_funcs/string.rs b/native/spark-expr/src/conversion_funcs/string.rs index adfd6e2390..e136a064d7 100644 --- a/native/spark-expr/src/conversion_funcs/string.rs +++ b/native/spark-expr/src/conversion_funcs/string.rs @@ -469,6 +469,71 @@ fn normalize_fullwidth_digits(s: &str) -> String { unsafe { String::from_utf8_unchecked(out) } } +/// Powers of ten that fit in an `i128` (`10^0` through `10^38`). +const POW10_I128: [i128; 39] = { + let mut table = [1i128; 39]; + let mut i = 1; + while i < 39 { + table[i] = table[i - 1] * 10; + i += 1; + } + table +}; + +/// `10^exp`, using the precomputed table for the range that fits in an `i128`. +#[inline] +fn pow10_i128(exp: u32) -> i128 { + match POW10_I128.get(exp as usize) { + Some(v) => *v, + None => 10_i128.pow(exp), + } +} + +/// Accumulate an ASCII-digit slice into an `i128`, returning `None` on overflow. +/// +/// The first 38 digits always fit (`i128::MAX` is ~1.7e38), so only the digits past +/// them need the per-digit overflow checks. +#[inline] +fn digits_to_i128(digits: &[u8]) -> Option { + let (head, tail) = digits.split_at(digits.len().min(38)); + let mut value: i128 = 0; + for &d in head { + value = value * 10 + (d - b'0') as i128; + } + for &d in tail { + value = value.checked_mul(10)?.checked_add((d - b'0') as i128)?; + } + Some(value) +} + +/// Values that Spark parses as NULL rather than as a decimal, matched case-insensitively. +const SPECIAL_DECIMAL_VALUES: [&str; 7] = [ + "inf", + "+inf", + "-inf", + "infinity", + "+infinity", + "-infinity", + "nan", +]; + +/// True if `trimmed` is one of [`SPECIAL_DECIMAL_VALUES`]. +#[inline] +fn is_special_value(trimmed: &str) -> bool { + // Every special value starts with `i`/`n`, or with a sign followed by `i`, so + // ordinary numeric input is ruled out after inspecting a single byte. + let bytes = trimmed.as_bytes(); + let plausible = match bytes.first() { + Some(b'i' | b'I' | b'n' | b'N') => true, + Some(b'+' | b'-') => matches!(bytes.get(1), Some(b'i' | b'I')), + _ => false, + }; + plausible + && SPECIAL_DECIMAL_VALUES + .iter() + .any(|v| trimmed.eq_ignore_ascii_case(v)) +} + /// Parse a decimal string into mantissa and scale /// e.g., "123.45" -> (12345, 2), "-0.001" -> (-1, 3) , 0e50 -> (0,50) etc /// Parse a string to decimal following Spark's behavior @@ -494,25 +559,18 @@ fn parse_string_to_decimal(input_str: &str, precision: u8, scale: i8) -> SparkRe // Normalize fullwidth digits to ASCII. Fast path skips the allocation for // pure-ASCII strings, which is the common case. let normalized; - let trimmed = if trimmed.bytes().any(|b| b > 0x7F) { + let trimmed = if trimmed.is_ascii() { + trimmed + } else { normalized = normalize_fullwidth_digits(trimmed); normalized.as_str() - } else { - trimmed }; if trimmed.is_empty() { return Ok(None); } // Handle special values (inf, nan, etc.) - if trimmed.eq_ignore_ascii_case("inf") - || trimmed.eq_ignore_ascii_case("+inf") - || trimmed.eq_ignore_ascii_case("infinity") - || trimmed.eq_ignore_ascii_case("+infinity") - || trimmed.eq_ignore_ascii_case("-inf") - || trimmed.eq_ignore_ascii_case("-infinity") - || trimmed.eq_ignore_ascii_case("nan") - { + if is_special_value(trimmed) { return Ok(None); } @@ -538,7 +596,7 @@ fn parse_string_to_decimal(input_str: &str, precision: u8, scale: i8) -> SparkRe if scale_adjustment > 38 { return Ok(None); } - mantissa.checked_mul(10_i128.pow(scale_adjustment as u32)) + mantissa.checked_mul(pow10_i128(scale_adjustment as u32)) } else { // Need to divide (decrease scale) let abs_scale_adjustment = (-scale_adjustment) as u32; @@ -546,7 +604,7 @@ fn parse_string_to_decimal(input_str: &str, precision: u8, scale: i8) -> SparkRe return Ok(Some(0)); } - let divisor = 10_i128.pow(abs_scale_adjustment); + let divisor = pow10_i128(abs_scale_adjustment); let quotient_opt = mantissa.checked_div(divisor); // Check if divisor is 0 if quotient_opt.is_none() { @@ -609,79 +667,73 @@ fn parse_decimal_str( precision: u8, scale: i8, ) -> SparkResult<(i128, i32)> { - if s.is_empty() { - return Err(invalid_decimal_cast(original_str, precision, scale)); - } - - let (mantissa_str, exponent) = if let Some(e_pos) = s.find(|c| ['e', 'E'].contains(&c)) { - let mantissa_part = &s[..e_pos]; - let exponent_part = &s[e_pos + 1..]; - // Parse exponent - let exp: i32 = exponent_part - .parse() - .map_err(|_| invalid_decimal_cast(original_str, precision, scale))?; - - (mantissa_part, exp) - } else { - (s, 0) - }; + let bytes = s.as_bytes(); - let negative = mantissa_str.starts_with('-'); - let mantissa_str = if negative || mantissa_str.starts_with('+') { - &mantissa_str[1..] - } else { - mantissa_str + let mut pos = 0; + let negative = match bytes.first() { + Some(b'-') => { + pos = 1; + true + } + Some(b'+') => { + pos = 1; + false + } + _ => false, }; - if mantissa_str.starts_with('+') || mantissa_str.starts_with('-') { - return Err(invalid_decimal_cast(original_str, precision, scale)); - } - - let (integral_part, fractional_part) = match mantissa_str.find('.') { - Some(dot_pos) => { - if mantissa_str[dot_pos + 1..].contains('.') { - return Err(invalid_decimal_cast(original_str, precision, scale)); + // Single validating pass over the mantissa: ASCII digits with at most one `.`, + // ending at the optional exponent marker. It also locates the integral/fractional + // split and the start of the exponent. `.`, `e` and `E` are ASCII, so scanning + // bytes can never land inside a multi-byte character and every index taken here is + // on a char boundary. + let digits_start = pos; + let mut dot_pos = None; + let mut exp_pos = None; + while pos < bytes.len() { + match bytes[pos] { + b'0'..=b'9' => pos += 1, + b'.' if dot_pos.is_none() => { + dot_pos = Some(pos); + pos += 1; + } + b'e' | b'E' => { + exp_pos = Some(pos); + break; } - (&mantissa_str[..dot_pos], &mantissa_str[dot_pos + 1..]) + _ => return Err(invalid_decimal_cast(original_str, precision, scale)), } - None => (mantissa_str, ""), - }; - - if integral_part.is_empty() && fractional_part.is_empty() { - return Err(invalid_decimal_cast(original_str, precision, scale)); } - if !integral_part.is_empty() && !integral_part.bytes().all(|b| b.is_ascii_digit()) { - return Err(invalid_decimal_cast(original_str, precision, scale)); - } + let exponent: i32 = match exp_pos { + Some(e_pos) => s[e_pos + 1..] + .parse() + .map_err(|_| invalid_decimal_cast(original_str, precision, scale))?, + None => 0, + }; - if !fractional_part.is_empty() && !fractional_part.bytes().all(|b| b.is_ascii_digit()) { + // An empty integral part is valid (e.g. ".5" or "-.7e9"), as is an empty + // fractional part, but they cannot both be empty. + let mantissa_end = exp_pos.unwrap_or(pos); + let (integral_part, fractional_part): (&[u8], &[u8]) = match dot_pos { + Some(dot) => (&bytes[digits_start..dot], &bytes[dot + 1..mantissa_end]), + None => (&bytes[digits_start..mantissa_end], &[]), + }; + + if integral_part.is_empty() && fractional_part.is_empty() { return Err(invalid_decimal_cast(original_str, precision, scale)); } - // Parse integral part - let integral_value: i128 = if integral_part.is_empty() { - // Empty integral part is valid (e.g., ".5" or "-.7e9") - 0 - } else { - integral_part - .parse() - .map_err(|_| invalid_decimal_cast(original_str, precision, scale))? - }; + let integral_value = digits_to_i128(integral_part) + .ok_or_else(|| invalid_decimal_cast(original_str, precision, scale))?; - // Parse fractional part let fractional_scale = fractional_part.len() as i32; - let fractional_value: i128 = if fractional_part.is_empty() { - 0 - } else { - fractional_part - .parse() - .map_err(|_| invalid_decimal_cast(original_str, precision, scale))? - }; + let fractional_value = digits_to_i128(fractional_part) + .ok_or_else(|| invalid_decimal_cast(original_str, precision, scale))?; // Combine: value = integral * 10^fractional_scale + fractional let mantissa = integral_value - .checked_mul(10_i128.pow(fractional_scale as u32)) + .checked_mul(pow10_i128(fractional_scale as u32)) .and_then(|v| v.checked_add(fractional_value)) .ok_or_else(|| invalid_decimal_cast(original_str, precision, scale))?;