From 8fddeb26cd84efb1daaefab0c5b801cde4366a2b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 15 Jul 2026 10:24:57 -0600 Subject: [PATCH 1/2] perf: vectorize float/decimal to narrow-integer casts The four macros for float-to-int and decimal-to-int narrowing casts (cast_float_to_int16_down, cast_float_to_int32_up, cast_decimal_to_int16_down, cast_decimal_to_int32_up) built the output with a per-element iterator-collect over Option/Result. Replace that with Arrow's unary (legacy) and try_unary (ANSI) kernels, which map the values buffer in one pass and carry the null buffer over, following the same pattern used by cast_int_to_int_macro. The decimal macros also hoist the constant scale divisor out of the per-element loop. Overflow, NaN, saturation, and wrap-through-Int semantics are preserved unchanged; only the iteration mechanism changes. Add Rust unit tests for the float-to-Byte and decimal-to-Int/Byte legacy wrap paths and the ANSI overflow error paths (previously covered only by Scala tests), plus a benchmark. Non-overflow casts are 49-91% faster with no regression. Part of #4936. --- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/cast_narrowing.rs | 100 +++++ .../src/conversion_funcs/numeric.rs | 381 ++++++++++-------- 3 files changed, 319 insertions(+), 166 deletions(-) create mode 100644 native/spark-expr/benches/cast_narrowing.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..4860f7233c 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_narrowing" +harness = false diff --git a/native/spark-expr/benches/cast_narrowing.rs b/native/spark-expr/benches/cast_narrowing.rs new file mode 100644 index 0000000000..af6aab9997 --- /dev/null +++ b/native/spark-expr/benches/cast_narrowing.rs @@ -0,0 +1,100 @@ +// 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::{Array, Decimal128Array, Float64Array, 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::hint::black_box; +use std::sync::Arc; + +fn f64_batch(size: usize) -> RecordBatch { + // Small in-range values so narrowing to i8 does not overflow. + let a: Float64Array = (0..size) + .map(|i| { + if i % 10 == 0 { + None + } else { + Some((i % 100) as f64) + } + }) + .collect(); + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)])); + RecordBatch::try_new(schema, vec![Arc::new(a)]).unwrap() +} + +fn dec_batch(size: usize) -> RecordBatch { + let a: Decimal128Array = (0..size) + .map(|i| { + if i % 10 == 0 { + None + } else { + Some((i % 100) as i128 * 100) + } + }) + .collect::() + .with_precision_and_scale(10, 2) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + a.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![Arc::new(a)]).unwrap() +} + +fn cast(to: DataType, mode: EvalMode) -> Cast { + Cast::new( + Arc::new(Column::new("a", 0)), + to, + SparkCastOptions::new_without_timezone(mode, false), + None, + None, + ) +} + +fn criterion_benchmark(c: &mut Criterion) { + let size = 8192; + let f = f64_batch(size); + let d = dec_batch(size); + + let f_i8 = cast(DataType::Int8, EvalMode::Legacy); + let f_i32 = cast(DataType::Int32, EvalMode::Legacy); + let f_i32_ansi = cast(DataType::Int32, EvalMode::Ansi); + let d_i8 = cast(DataType::Int8, EvalMode::Legacy); + let d_i32 = cast(DataType::Int32, EvalMode::Legacy); + + c.bench_function("cast_narrowing: f64 -> i8", |b| { + b.iter(|| black_box(f_i8.evaluate(black_box(&f)).unwrap())) + }); + c.bench_function("cast_narrowing: f64 -> i32", |b| { + b.iter(|| black_box(f_i32.evaluate(black_box(&f)).unwrap())) + }); + c.bench_function("cast_narrowing: f64 -> i32 ansi", |b| { + b.iter(|| black_box(f_i32_ansi.evaluate(black_box(&f)).unwrap())) + }); + c.bench_function("cast_narrowing: dec -> i8", |b| { + b.iter(|| black_box(d_i8.evaluate(black_box(&d)).unwrap())) + }); + c.bench_function("cast_narrowing: dec -> i32", |b| { + b.iter(|| black_box(d_i32.evaluate(black_box(&d)).unwrap())) + }); +} + +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 ce2fe515b6..0ec8abb684 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -301,6 +301,7 @@ macro_rules! cast_float_to_int16_down { $rust_dest_type:ty, $src_type_str:expr, $dest_type_str:expr, + $dest_arrow_type:ty, $format_str:expr ) => {{ let cast_array = $array @@ -308,45 +309,29 @@ macro_rules! cast_float_to_int16_down { .downcast_ref::<$src_array_type>() .expect(concat!("Expected a ", stringify!($src_array_type))); - let output_array = match $eval_mode { - EvalMode::Ansi => cast_array - .iter() - .map(|value| match value { - Some(value) => { - let is_overflow = value.is_nan() || value.abs() as i32 == i32::MAX; - if is_overflow { - return Err(cast_overflow( - &format!($format_str, value).replace("e", "E"), - $src_type_str, - $dest_type_str, - )); - } - let i32_value = value as i32; - <$rust_dest_type>::try_from(i32_value) - .map_err(|_| { - cast_overflow( - &format!($format_str, value).replace("e", "E"), - $src_type_str, - $dest_type_str, - ) - }) - .map(Some) - } - None => Ok(None), - }) - .collect::>()?, - _ => cast_array - .iter() - .map(|value| match value { - Some(value) => { - let i32_value = value as i32; - Ok::, SparkError>(Some( - i32_value as $rust_dest_type, - )) - } - None => Ok(None), + // Spark casts float -> Byte/Short by going through Int first (with its own overflow), + // then narrowing. `unary`/`try_unary` map the values buffer in one pass and carry the + // null buffer over, replacing the per-element iterator-collect. + let output_array: $dest_array_type = match $eval_mode { + EvalMode::Ansi => cast_array.try_unary::<_, $dest_arrow_type, SparkError>(|value| { + let is_overflow = value.is_nan() || value.abs() as i32 == i32::MAX; + if is_overflow { + return Err(cast_overflow( + &format!($format_str, value).replace("e", "E"), + $src_type_str, + $dest_type_str, + )); + } + let i32_value = value as i32; + <$rust_dest_type>::try_from(i32_value).map_err(|_| { + cast_overflow( + &format!($format_str, value).replace("e", "E"), + $src_type_str, + $dest_type_str, + ) }) - .collect::>()?, + })?, + _ => cast_array.unary::<_, $dest_arrow_type>(|value| (value as i32) as $rust_dest_type), }; Ok(Arc::new(output_array) as ArrayRef) }}; @@ -363,6 +348,7 @@ macro_rules! cast_float_to_int32_up { $src_type_str:expr, $dest_type_str:expr, $max_dest_val:expr, + $dest_arrow_type:ty, $format_str:expr ) => {{ let cast_array = $array @@ -370,34 +356,21 @@ macro_rules! cast_float_to_int32_up { .downcast_ref::<$src_array_type>() .expect(concat!("Expected a ", stringify!($src_array_type))); - let output_array = match $eval_mode { - EvalMode::Ansi => cast_array - .iter() - .map(|value| match value { - Some(value) => { - let is_overflow = - value.is_nan() || value.abs() as $rust_dest_type == $max_dest_val; - if is_overflow { - return Err(cast_overflow( - &format!($format_str, value).replace("e", "E"), - $src_type_str, - $dest_type_str, - )); - } - Ok(Some(value as $rust_dest_type)) - } - None => Ok(None), - }) - .collect::>()?, - _ => cast_array - .iter() - .map(|value| match value { - Some(value) => { - Ok::, SparkError>(Some(value as $rust_dest_type)) - } - None => Ok(None), - }) - .collect::>()?, + // `unary`/`try_unary` map the values buffer in one pass and carry the null buffer over, + // replacing the per-element iterator-collect. + let output_array: $dest_array_type = match $eval_mode { + EvalMode::Ansi => cast_array.try_unary::<_, $dest_arrow_type, SparkError>(|value| { + let is_overflow = value.is_nan() || value.abs() as $rust_dest_type == $max_dest_val; + if is_overflow { + return Err(cast_overflow( + &format!($format_str, value).replace("e", "E"), + $src_type_str, + $dest_type_str, + )); + } + Ok(value as $rust_dest_type) + })?, + _ => cast_array.unary::<_, $dest_arrow_type>(|value| value as $rust_dest_type), }; Ok(Arc::new(output_array) as ArrayRef) }}; @@ -414,69 +387,47 @@ macro_rules! cast_decimal_to_int16_down { $rust_dest_type:ty, $dest_type_str:expr, $precision:expr, - $scale:expr + $scale:expr, + $dest_arrow_type:ty ) => {{ let cast_array = $array .as_any() .downcast_ref::() .expect("Expected a Decimal128ArrayType"); - let output_array = match $eval_mode { - EvalMode::Ansi => cast_array - .iter() - .map(|value| match value { - Some(value) => { - let divisor = 10_i128.pow($scale as u32); - let truncated = value / divisor; - let is_overflow = truncated.abs() > i32::MAX.into(); - if is_overflow { - return Err(cast_overflow( - &format!( - "{}BD", - format_decimal_str( - &value.to_string(), - $precision as usize, - $scale - ) - ), - &format!("DECIMAL({},{})", $precision, $scale), - $dest_type_str, - )); - } - let i32_value = truncated as i32; - <$rust_dest_type>::try_from(i32_value) - .map_err(|_| { - cast_overflow( - &format!( - "{}BD", - format_decimal_str( - &value.to_string(), - $precision as usize, - $scale - ) - ), - &format!("DECIMAL({},{})", $precision, $scale), - $dest_type_str, - ) - }) - .map(Some) - } - None => Ok(None), - }) - .collect::>()?, - _ => cast_array - .iter() - .map(|value| match value { - Some(value) => { - let divisor = 10_i128.pow($scale as u32); - let i32_value = (value / divisor) as i32; - Ok::, SparkError>(Some( - i32_value as $rust_dest_type, - )) - } - None => Ok(None), + // The scale divisor is constant across the batch, so hoist it out of the per-element + // loop. `unary`/`try_unary` then map the values buffer in one pass, carrying the null + // buffer over, instead of the per-element iterator-collect. + let divisor = 10_i128.pow($scale as u32); + let output_array: $dest_array_type = match $eval_mode { + EvalMode::Ansi => cast_array.try_unary::<_, $dest_arrow_type, SparkError>(|value| { + let truncated = value / divisor; + let is_overflow = truncated.abs() > i32::MAX.into(); + if is_overflow { + return Err(cast_overflow( + &format!( + "{}BD", + format_decimal_str(&value.to_string(), $precision as usize, $scale) + ), + &format!("DECIMAL({},{})", $precision, $scale), + $dest_type_str, + )); + } + let i32_value = truncated as i32; + <$rust_dest_type>::try_from(i32_value).map_err(|_| { + cast_overflow( + &format!( + "{}BD", + format_decimal_str(&value.to_string(), $precision as usize, $scale) + ), + &format!("DECIMAL({},{})", $precision, $scale), + $dest_type_str, + ) }) - .collect::>()?, + })?, + _ => cast_array.unary::<_, $dest_arrow_type>(|value| { + ((value / divisor) as i32) as $rust_dest_type + }), }; Ok(Arc::new(output_array) as ArrayRef) }}; @@ -491,53 +442,36 @@ macro_rules! cast_decimal_to_int32_up { $dest_type_str:expr, $max_dest_val:expr, $precision:expr, - $scale:expr + $scale:expr, + $dest_arrow_type:ty ) => {{ let cast_array = $array .as_any() .downcast_ref::() .expect("Expected a Decimal128ArrayType"); - let output_array = match $eval_mode { - EvalMode::Ansi => cast_array - .iter() - .map(|value| match value { - Some(value) => { - let divisor = 10_i128.pow($scale as u32); - let truncated = value / divisor; - let is_overflow = truncated.abs() > $max_dest_val.into(); - if is_overflow { - return Err(cast_overflow( - &format!( - "{}BD", - format_decimal_str( - &value.to_string(), - $precision as usize, - $scale - ) - ), - &format!("DECIMAL({},{})", $precision, $scale), - $dest_type_str, - )); - } - Ok(Some(truncated as $rust_dest_type)) - } - None => Ok(None), - }) - .collect::>()?, + // The scale divisor is constant across the batch, so hoist it out of the per-element + // loop. `unary`/`try_unary` then map the values buffer in one pass, carrying the null + // buffer over, instead of the per-element iterator-collect. + let divisor = 10_i128.pow($scale as u32); + let output_array: $dest_array_type = match $eval_mode { + EvalMode::Ansi => cast_array.try_unary::<_, $dest_arrow_type, SparkError>(|value| { + let truncated = value / divisor; + let is_overflow = truncated.abs() > $max_dest_val.into(); + if is_overflow { + return Err(cast_overflow( + &format!( + "{}BD", + format_decimal_str(&value.to_string(), $precision as usize, $scale) + ), + &format!("DECIMAL({},{})", $precision, $scale), + $dest_type_str, + )); + } + Ok(truncated as $rust_dest_type) + })?, _ => cast_array - .iter() - .map(|value| match value { - Some(value) => { - let divisor = 10_i128.pow($scale as u32); - let truncated = value / divisor; - Ok::, SparkError>(Some( - truncated as $rust_dest_type, - )) - } - None => Ok(None), - }) - .collect::>()?, + .unary::<_, $dest_arrow_type>(|value| (value / divisor) as $rust_dest_type), }; Ok(Arc::new(output_array) as ArrayRef) }}; @@ -880,6 +814,7 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( i8, "FLOAT", "TINYINT", + Int8Type, "{:e}" ), (DataType::Float32, DataType::Int16) => cast_float_to_int16_down!( @@ -891,6 +826,7 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( i16, "FLOAT", "SMALLINT", + Int16Type, "{:e}" ), (DataType::Float32, DataType::Int32) => cast_float_to_int32_up!( @@ -903,6 +839,7 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( "FLOAT", "INT", i32::MAX, + Int32Type, "{:e}" ), (DataType::Float32, DataType::Int64) => cast_float_to_int32_up!( @@ -915,6 +852,7 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( "FLOAT", "BIGINT", i64::MAX, + Int64Type, "{:e}" ), (DataType::Float64, DataType::Int8) => cast_float_to_int16_down!( @@ -926,6 +864,7 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( i8, "DOUBLE", "TINYINT", + Int8Type, "{:e}D" ), (DataType::Float64, DataType::Int16) => cast_float_to_int16_down!( @@ -937,6 +876,7 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( i16, "DOUBLE", "SMALLINT", + Int16Type, "{:e}D" ), (DataType::Float64, DataType::Int32) => cast_float_to_int32_up!( @@ -949,6 +889,7 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( "DOUBLE", "INT", i32::MAX, + Int32Type, "{:e}D" ), (DataType::Float64, DataType::Int64) => cast_float_to_int32_up!( @@ -961,16 +902,17 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( "DOUBLE", "BIGINT", i64::MAX, + Int64Type, "{:e}D" ), (DataType::Decimal128(precision, scale), DataType::Int8) => { cast_decimal_to_int16_down!( - array, eval_mode, Int8Array, i8, "TINYINT", *precision, *scale + array, eval_mode, Int8Array, i8, "TINYINT", *precision, *scale, Int8Type ) } (DataType::Decimal128(precision, scale), DataType::Int16) => { cast_decimal_to_int16_down!( - array, eval_mode, Int16Array, i16, "SMALLINT", *precision, *scale + array, eval_mode, Int16Array, i16, "SMALLINT", *precision, *scale, Int16Type ) } (DataType::Decimal128(precision, scale), DataType::Int32) => { @@ -982,7 +924,8 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( "INT", i32::MAX, *precision, - *scale + *scale, + Int32Type ) } (DataType::Decimal128(precision, scale), DataType::Int64) => { @@ -994,7 +937,8 @@ pub(crate) fn spark_cast_nonintegral_numeric_to_integral( "BIGINT", i64::MAX, *precision, - *scale + *scale, + Int64Type ) } _ => unreachable!( @@ -1160,6 +1104,111 @@ mod tests { assert_eq!(decimal_array.value(1), -10000); // -100 * 10^2 assert!(decimal_array.is_null(2)); } + + #[test] + fn test_cast_float64_to_int8_legacy_wraps() { + // Spark narrows float -> Int (truncate) -> Byte (wrap). 300.7 -> 300 -> 44; -1.9 -> -1. + let a: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(300.7), + Some(-1.9), + None, + Some(42.0), + ])); + let r = spark_cast_nonintegral_numeric_to_integral( + &a, + EvalMode::Legacy, + &DataType::Float64, + &DataType::Int8, + ) + .unwrap(); + let d = r.as_primitive::(); + assert_eq!(d.value(0), 44); // 300 wraps to 44 in i8 + assert_eq!(d.value(1), -1); + assert!(d.is_null(2)); + assert_eq!(d.value(3), 42); + } + + #[test] + fn test_cast_float64_to_int32_ansi_ok_and_overflow() { + let ok: ArrayRef = Arc::new(Float64Array::from(vec![Some(42.0), None, Some(-5.0)])); + let r = spark_cast_nonintegral_numeric_to_integral( + &ok, + EvalMode::Ansi, + &DataType::Float64, + &DataType::Int32, + ) + .unwrap(); + let d = r.as_primitive::(); + assert_eq!(d.value(0), 42); + assert!(d.is_null(1)); + assert_eq!(d.value(2), -5); + + let of: ArrayRef = Arc::new(Float64Array::from(vec![Some(1e30)])); + let e = spark_cast_nonintegral_numeric_to_integral( + &of, + EvalMode::Ansi, + &DataType::Float64, + &DataType::Int32, + ); + assert!(e.is_err()); + } + + #[test] + fn test_cast_decimal_to_int32_legacy() { + // Decimal128(10,2): 123.45 truncates to 123; -1.00 to -1. + let a: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(12345), None, Some(-100)]) + .with_precision_and_scale(10, 2) + .unwrap(), + ); + let r = spark_cast_nonintegral_numeric_to_integral( + &a, + EvalMode::Legacy, + &DataType::Decimal128(10, 2), + &DataType::Int32, + ) + .unwrap(); + let d = r.as_primitive::(); + assert_eq!(d.value(0), 123); + assert!(d.is_null(1)); + assert_eq!(d.value(2), -1); + } + + #[test] + fn test_cast_decimal_to_int8_legacy_wraps() { + // 300.00 truncates to 300 (Int), which wraps to 44 in Byte. + let a: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(30000)]) + .with_precision_and_scale(10, 2) + .unwrap(), + ); + let r = spark_cast_nonintegral_numeric_to_integral( + &a, + EvalMode::Legacy, + &DataType::Decimal128(10, 2), + &DataType::Int8, + ) + .unwrap(); + assert_eq!(r.as_primitive::().value(0), 44); + } + + #[test] + fn test_cast_decimal_to_int32_ansi_overflow_errors() { + // 10_000_000_000 (scale 0) exceeds i32::MAX -> ANSI error. + let a: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(10_000_000_000_i128)]) + .with_precision_and_scale(20, 0) + .unwrap(), + ); + let e = spark_cast_nonintegral_numeric_to_integral( + &a, + EvalMode::Ansi, + &DataType::Decimal128(20, 0), + &DataType::Int32, + ); + assert!(e.is_err()); + } + #[test] fn test_cast_int_to_timestamp() { let timezones: [Option>; 6] = [ From 5a2a9313cee538521325261bf331fea9d23ccf63 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 15 Jul 2026 10:25:40 -0600 Subject: [PATCH 2/2] docs: record narrowing float/decimal to int cast performance audit entry --- .../contributor-guide/expression-audits/conversion_funcs.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/contributor-guide/expression-audits/conversion_funcs.md b/docs/source/contributor-guide/expression-audits/conversion_funcs.md index 45b6a8f03b..91c0b77e29 100644 --- a/docs/source/contributor-guide/expression-audits/conversion_funcs.md +++ b/docs/source/contributor-guide/expression-audits/conversion_funcs.md @@ -35,5 +35,6 @@ - `spark.sql.legacy.castComplexTypesToString.enabled=true` is not honoured by Comet (https://github.com/apache/datafusion-comet/issues/4492). - `CAST( AS DECIMAL)` rounding may differ from Spark (`Incompatible`, gated by `spark.comet.expression.Cast.allowIncompatible`, tracked at https://github.com/apache/datafusion-comet/issues/1371). - Spark registers the type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `decimal`, `double`, `float`, `int`, `smallint`, `string`, `timestamp`, `tinyint`) as cast aliases. Each lowers to the same `Cast` node, so Comet handles it via the `cast` implementation with the same compatibility profile. +- Performance (tuned 2026-07-15, PR #4941): float-to-int and decimal-to-int narrowing casts (`cast_float_to_int16_down`/`cast_float_to_int32_up`/`cast_decimal_to_int16_down`/`cast_decimal_to_int32_up`) now map the values buffer with Arrow `unary` (legacy) / `try_unary` (ANSI) instead of a per-element `Option`/`Result` iterator-collect, and the decimal macros hoist the constant `10^scale` divisor out of the loop. 49-91% faster with no regression; overflow/NaN/wrap semantics unchanged. Benchmark: `benches/cast_narrowing.rs`. [Spark Expression Support]: ../../user-guide/latest/expressions.md