diff --git a/native/spark-expr/benches/cast_from_string.rs b/native/spark-expr/benches/cast_from_string.rs index 1fbb7c535b..480cb9416e 100644 --- a/native/spark-expr/benches/cast_from_string.rs +++ b/native/spark-expr/benches/cast_from_string.rs @@ -20,6 +20,8 @@ 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::sync::Arc; fn criterion_benchmark(c: &mut Criterion) { @@ -113,22 +115,22 @@ fn criterion_benchmark(c: &mut Criterion) { (EvalMode::Try, "try"), ] { let spark_cast_options = SparkCastOptions::new(mode, "", false); - let cast_to_decimal_38_10 = Cast::new( - expr.clone(), - DataType::Decimal128(38, 10), - spark_cast_options, - None, - None, - ); - let mut group = c.benchmark_group(format!("cast_string_to_decimal/{}", mode_name)); - group.bench_function("decimal_38_10", |b| { - b.iter(|| { - cast_to_decimal_38_10 - .evaluate(&decimal_string_batch) - .unwrap() + 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, + spark_cast_options.clone(), + None, + None, + ); + group.bench_function(name, |b| { + b.iter(|| cast.evaluate(&decimal_string_batch).unwrap()); }); - }); + } group.finish(); } } @@ -184,8 +186,9 @@ fn create_decimal_string_batch() -> RecordBatch { /// Create batch with decimal strings for string-to-decimal cast perf evaluation fn create_decimal_cast_string_batch() -> 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..1000 { + for i in 0..8192 { if i % 10 == 0 { b.append_null(); } else { @@ -193,30 +196,32 @@ fn create_decimal_cast_string_batch() -> RecordBatch { match i % 5 { 0 => { // gen simple decimals (ex : "123.45" - let int_part: u32 = rand::random::() % 1000000; - let dec_part: u32 = rand::random::() % 100000; + let int_part = rng.random_range(0..1_000_000u32); + let dec_part = rng.random_range(0..100_000u32); b.append_value(format!("{}.{}", int_part, dec_part)); } 1 => { // gen scientific notation like "123e5" - let mantissa: u32 = rand::random::() % 1000; - let exp: i8 = (rand::random::() % 10).abs(); - b.append_value(format!("{}.{}E{}", mantissa / 100, mantissa % 100, exp)); + b.append_value(format!( + "{}.{}E{}", + rng.random_range(0..10u32), + rng.random_range(0..100u32), + rng.random_range(0..10u32) + )); } 2 => { // Negative numbers - let int_part: u32 = rand::random::() % 1000000; - let dec_part: u32 = rand::random::() % 100000; + let int_part = rng.random_range(0..1_000_000u32); + let dec_part = rng.random_range(0..100_000u32); b.append_value(format!("-{}.{}", int_part, dec_part)); } 3 => { // Ints only - let val: i32 = rand::random::() % 1000000; - b.append_value(format!("{}", val)); + b.append_value(format!("{}", rng.random_range(-1_000_000..1_000_000i32))); } _ => { // Small decimals (ex : 0.001) - let dec_part: u32 = rand::random::() % 100000; + let dec_part = rng.random_range(0..100_000u32); b.append_value(format!("0.{:05}", dec_part)); } } diff --git a/native/spark-expr/src/conversion_funcs/string.rs b/native/spark-expr/src/conversion_funcs/string.rs index adfd6e2390..8a80e3a204 100644 --- a/native/spark-expr/src/conversion_funcs/string.rs +++ b/native/spark-expr/src/conversion_funcs/string.rs @@ -469,6 +469,68 @@ 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`, or `None` when the exponent overflows an `i128` (exp >= 39). +#[inline] +fn pow10_i128(exp: u32) -> Option { + POW10_I128.get(exp as usize).copied() +} + +/// 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 +556,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 +593,8 @@ 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)) + // Bounded above, so pow10_i128 always returns Some. + mantissa.checked_mul(pow10_i128(scale_adjustment as u32).unwrap()) } else { // Need to divide (decrease scale) let abs_scale_adjustment = (-scale_adjustment) as u32; @@ -546,7 +602,8 @@ 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); + // Bounded above, so pow10_i128 always returns Some. + let divisor = pow10_i128(abs_scale_adjustment).unwrap(); let quotient_opt = mantissa.checked_div(divisor); // Check if divisor is 0 if quotient_opt.is_none() { @@ -609,79 +666,75 @@ 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; } - (&mantissa_str[..dot_pos], &mantissa_str[dot_pos + 1..]) + b'e' | b'E' => { + exp_pos = Some(pos); + break; + } + _ => 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, + }; + + // 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 !fractional_part.is_empty() && !fractional_part.bytes().all(|b| b.is_ascii_digit()) { + 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)) + // Combine: value = integral * 10^fractional_scale + fractional. + // A fractional_scale beyond 38 cannot fit in an i128, so pow10_i128 returns None and this + // maps to the invalid-decimal error path instead of panicking. + let mantissa = pow10_i128(fractional_scale as u32) + .and_then(|p| integral_value.checked_mul(p)) .and_then(|v| v.checked_add(fractional_value)) .ok_or_else(|| invalid_decimal_cast(original_str, precision, scale))?; @@ -1989,6 +2042,37 @@ mod tests { } } + #[test] + fn test_digits_to_i128_boundary() { + // 38 nines is under i128::MAX (~1.7e38 vs 9.99e37); the head-only path parses it. + let d38 = "9".repeat(38); + assert_eq!( + digits_to_i128(d38.as_bytes()), + Some(99_999_999_999_999_999_999_999_999_999_999_999_999_i128) + ); + // 39 nines overflows i128 in the tail-checked step. + let d39 = "9".repeat(39); + assert_eq!(digits_to_i128(d39.as_bytes()), None); + // 40 characters that reduce to a small value once the leading zero(s) are seen: the + // head-only accumulator must not lose bits on 38 zeros followed by two digits. + let z38_plus = format!("{}42", "0".repeat(38)); + assert_eq!(digits_to_i128(z38_plus.as_bytes()), Some(42)); + } + + #[test] + fn test_parse_string_to_decimal_boundary() { + // 38-digit integral parses (fits i128). + let s38 = "9".repeat(38); + assert!(parse_string_to_decimal(&s38, 38, 0).unwrap().is_some()); + // 39-digit integral overflows i128, so returns the invalid_decimal_cast error. + let s39 = "9".repeat(39); + assert!(parse_string_to_decimal(&s39, 38, 0).is_err()); + // Very long fractional part now returns Err via the invalid_decimal_cast path instead + // of panicking on 10_i128.pow(fractional_scale). + let over_long = format!("0.{}", "0".repeat(40)); + assert!(parse_string_to_decimal(&over_long, 38, 10).is_err()); + } + #[test] #[cfg_attr(miri, ignore)] // test takes too long with miri fn test_cast_string_to_timestamp() {