Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 30 additions & 25 deletions native/spark-expr/benches/cast_from_string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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();
}
}
Expand Down Expand Up @@ -184,39 +186,42 @@ 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 {
// Generate various decimal formats
match i % 5 {
0 => {
// gen simple decimals (ex : "123.45"
let int_part: u32 = rand::random::<u32>() % 1000000;
let dec_part: u32 = rand::random::<u32>() % 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::<u32>() % 1000;
let exp: i8 = (rand::random::<i8>() % 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::<u32>() % 1000000;
let dec_part: u32 = rand::random::<u32>() % 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::<i32>() % 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::<u32>() % 100000;
let dec_part = rng.random_range(0..100_000u32);
b.append_value(format!("0.{:05}", dec_part));
}
}
Expand Down
228 changes: 156 additions & 72 deletions native/spark-expr/src/conversion_funcs/string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<i128> {
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<i128> {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

digits_to_i128 at native/spark-expr/src/conversion_funcs/string.rs (new helper) splits the digit slice at 38 and skips overflow checks on the first 38 digits, relying on the fact that 38 nines (about 9.99e37) fit under i128::MAX (about 1.7e38). The reasoning is correct, and leading zeros are handled correctly because zeros in the head accumulate to 0 and the tail uses checked_mul/checked_add. But this boundary is exactly where a future edit could regress, and the current Scala fuzz tests (CometCastSuite.scala:958-961) cap the generated digit count at 38, so they never exercise a 39-plus-digit integral part.

Suggested change: add a Rust unit test in the #[cfg(test)] block of string.rs asserting parse results for a 38-digit value (parses), a 39-digit value (overflows to the invalid_decimal_cast error / NULL in non-ANSI), and a 40-digit value with leading zeros that reduce to a small value (parses). This locks the boundary that the optimization now depends on.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added test_digits_to_i128_boundary: 38 nines parses, 39 nines returns None, 38 leading zeros plus 42 reduces to 42. Added test_parse_string_to_decimal_boundary at the higher level too.

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
Expand All @@ -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);
}

Expand All @@ -538,15 +593,17 @@ 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;
if abs_scale_adjustment > 38 {
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() {
Expand Down Expand Up @@ -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))?;

Expand Down Expand Up @@ -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() {
Expand Down
Loading