comphead commented on code in PR #7534:
URL: https://github.com/apache/datafusion-comet/pull/7534#discussion_r4532980591
##########
native/spark-expr/src/math_funcs/utils.rs:
##########
@@ -162,3 +143,486 @@ pub(crate) fn get_precision_scale(data_type: &DataType)
-> (u8, i8) {
};
(*precision, *scale)
}
+
+/// The rounding direction that distinguishes `ceil` from `floor`.
+///
+/// Implemented by marker types so that each direction is monomorphized into
the element loop
+/// rather than dispatched for every value.
+pub(crate) trait RoundingMode {
+ const NAME: &'static str;
+
+ fn round_f32(x: f32) -> f32;
+ fn round_f64(x: f64) -> f64;
+ fn div_i64(x: i64, div: i64) -> i64;
+ fn div_i128(x: i128, div: i128) -> i128;
Review Comment:
`div_ceil` and `div_floor` from `num::integer` are already generic over
`Integer`, and `f32`/`f64` both have `ceil`/`floor` through `num::Float`. So
`div_i64`/`div_i128` and `round_f32`/`round_f64` could each be one generic
method, which halves the trait and both impls. A `const CEIL: bool` parameter
would also fold away at compile time and remove the marker types. I have not
compiled either variant.
##########
native/spark-expr/src/math_funcs/utils.rs:
##########
@@ -162,3 +143,486 @@ pub(crate) fn get_precision_scale(data_type: &DataType)
-> (u8, i8) {
};
(*precision, *scale)
}
+
+/// The rounding direction that distinguishes `ceil` from `floor`.
Review Comment:
`utils.rs` holds the decimal helpers that `div.rs`, `round.rs`,
`make_decimal.rs` and `wide_decimal_binary_expr.rs` share. The trait, the
markers, the kernel and its tests (about 480 lines) are used only by `ceil` and
`floor`, and so is `dispatch_pow10!`. Since `ceil.rs` and `floor.rs` are now
thin wrappers, could the shared code live in one `ceil_floor.rs` module that
exports `spark_ceil` and `spark_floor`, and keep `utils.rs` for the
cross-function helpers?
##########
native/spark-expr/src/math_funcs/utils.rs:
##########
@@ -162,3 +143,486 @@ pub(crate) fn get_precision_scale(data_type: &DataType)
-> (u8, i8) {
};
(*precision, *scale)
}
+
+/// The rounding direction that distinguishes `ceil` from `floor`.
+///
+/// Implemented by marker types so that each direction is monomorphized into
the element loop
+/// rather than dispatched for every value.
+pub(crate) trait RoundingMode {
+ const NAME: &'static str;
+
+ fn round_f32(x: f32) -> f32;
+ fn round_f64(x: f64) -> f64;
+ fn div_i64(x: i64, div: i64) -> i64;
+ fn div_i128(x: i128, div: i128) -> i128;
+}
+
+pub(crate) struct Ceil;
+
+impl RoundingMode for Ceil {
+ const NAME: &'static str = "ceil";
+
+ #[inline]
+ fn round_f32(x: f32) -> f32 {
+ x.ceil()
+ }
+
+ #[inline]
+ fn round_f64(x: f64) -> f64 {
+ x.ceil()
+ }
+
+ #[inline]
+ fn div_i64(x: i64, div: i64) -> i64 {
+ div_ceil(x, div)
+ }
+
+ #[inline]
+ fn div_i128(x: i128, div: i128) -> i128 {
+ div_ceil(x, div)
+ }
+}
+
+pub(crate) struct Floor;
+
+impl RoundingMode for Floor {
+ const NAME: &'static str = "floor";
+
+ #[inline]
+ fn round_f32(x: f32) -> f32 {
+ x.floor()
+ }
+
+ #[inline]
+ fn round_f64(x: f64) -> f64 {
+ x.floor()
+ }
+
+ #[inline]
+ fn div_i64(x: i64, div: i64) -> i64 {
+ div_floor(x, div)
+ }
+
+ #[inline]
+ fn div_i128(x: i128, div: i128) -> i128 {
+ div_floor(x, div)
+ }
+}
+
+/// Shared implementation of Spark's `ceil` and `floor` expressions.
+pub(crate) fn spark_ceil_or_floor<M: RoundingMode>(
+ args: &[ColumnarValue],
+ data_type: &DataType,
+) -> Result<ColumnarValue, DataFusionError> {
+ let value = &args[0];
+ match value {
+ ColumnarValue::Array(array) => match array.data_type() {
+ DataType::Float32 => {
+ let result: Int64Array = arrow::compute::kernels::arity::unary(
+ array.as_primitive::<Float32Type>(),
+ |x| M::round_f32(x) as i64,
+ );
+ Ok(ColumnarValue::Array(Arc::new(result)))
+ }
+ DataType::Float64 => {
+ let result: Int64Array = arrow::compute::kernels::arity::unary(
+ array.as_primitive::<Float64Type>(),
+ |x| M::round_f64(x) as i64,
+ );
+ Ok(ColumnarValue::Array(Arc::new(result)))
+ }
+ DataType::Int64 => {
+ let result =
array.as_any().downcast_ref::<Int64Array>().unwrap();
+ Ok(ColumnarValue::Array(Arc::new(result.clone())))
+ }
+ DataType::Decimal128(_, input_scale) if *input_scale > 0 => {
+ let (precision, scale) = get_precision_scale(data_type);
+ dispatch_pow10!(
+ *input_scale,
+ EXP => make_decimal_array(array, precision, scale,
decimal_round_pow10::<M, EXP>),
+ make_decimal_array(array, precision, scale,
decimal_round_f::<M>(*input_scale))
+ )
+ }
+ other => Err(DataFusionError::Internal(format!(
+ "Unsupported data type {other:?} for function {}",
+ M::NAME,
+ ))),
+ },
+ ColumnarValue::Scalar(a) => match a {
+ ScalarValue::Float32(a) =>
Ok(ColumnarValue::Scalar(ScalarValue::Int64(
+ a.map(|x| M::round_f32(x) as i64),
+ ))),
+ ScalarValue::Float64(a) =>
Ok(ColumnarValue::Scalar(ScalarValue::Int64(
+ a.map(|x| M::round_f64(x) as i64),
+ ))),
+ ScalarValue::Int64(a) =>
Ok(ColumnarValue::Scalar(ScalarValue::Int64(a.map(|x| x)))),
+ ScalarValue::Decimal128(a, _, input_scale) if *input_scale > 0 => {
+ let f = decimal_round_f::<M>(*input_scale);
+ let (precision, scale) = get_precision_scale(data_type);
+ make_decimal_scalar(a, precision, scale, f)
+ }
+ _ => Err(DataFusionError::Internal(format!(
+ "Unsupported data type {:?} for function {}",
+ value.data_type(),
+ M::NAME,
+ ))),
+ },
+ }
+}
+
+#[inline]
+fn decimal_round_f<M: RoundingMode>(scale: i8) -> impl Fn(i128) -> i128 {
+ let div = 10_i128.pow_wrapping(scale as u32);
+ move |x: i128| M::div_i128(x, div)
+}
+
+/// Rounds an unscaled decimal divided by `10^EXP` in the direction of `M`.
+///
+/// `EXP` is a compile-time constant so that the divisor is folded in and the
division lowered to a
+/// multiply-and-shift. A 128-bit division is always a libcall, even by a
constant, so values that
+/// fit in 64 bits take a 64-bit path; unscaled decimals rarely exceed that
range.
+#[inline]
+fn decimal_round_pow10<M: RoundingMode, const EXP: u32>(x: i128) -> i128 {
+ match i64::try_from(x) {
+ Ok(x) => M::div_i64(x, const { 10_i64.pow(EXP) }) as i128,
+ Err(_) => decimal_round_wide::<M>(x, const { 10_i128.pow(EXP) }),
+ }
+}
+
+/// Kept out of line so that the libcall and its stack frame stay out of the
loop body of every
+/// [`decimal_round_pow10`] instantiation.
+#[cold]
+#[inline(never)]
+fn decimal_round_wide<M: RoundingMode>(x: i128, div: i128) -> i128 {
+ M::div_i128(x, div)
+}
+
+#[cfg(test)]
+mod test {
+ use super::*;
+ use crate::{spark_ceil, spark_floor};
+ use arrow::array::{Float32Array, Float64Array};
+ use datafusion::common::cast::as_int64_array;
+ use datafusion::common::Result;
+
+ type RoundFn = fn(&[ColumnarValue], &DataType) -> Result<ColumnarValue>;
+
+ fn check_int64_array(
+ round: RoundFn,
+ input: ArrayRef,
+ expected: Vec<Option<i64>>,
+ ) -> Result<()> {
+ let data_type = input.data_type().clone();
+ let ColumnarValue::Array(result) =
round(&[ColumnarValue::Array(input)], &data_type)?
+ else {
+ unreachable!()
+ };
+ assert_eq!(as_int64_array(&result)?, &Int64Array::from(expected));
+ Ok(())
+ }
+
+ fn check_decimal_array(
+ round: RoundFn,
+ input: &[Option<i128>],
+ (precision, scale): (u8, i8),
+ output_precision: u8,
+ expected: Vec<Option<i128>>,
+ ) -> Result<()> {
+ let array =
+
Decimal128Array::from(input.to_vec()).with_precision_and_scale(precision,
scale)?;
+ let args = [ColumnarValue::Array(Arc::new(array))];
+ let ColumnarValue::Array(result) =
+ round(&args, &DataType::Decimal128(output_precision, 0))?
+ else {
+ unreachable!()
+ };
+ let expected =
+
Decimal128Array::from(expected).with_precision_and_scale(output_precision, 0)?;
+ let actual =
result.as_any().downcast_ref::<Decimal128Array>().unwrap();
+ assert_eq!(actual, &expected);
+ Ok(())
+ }
+
+ fn check_scalar(
+ round: RoundFn,
+ input: ScalarValue,
+ data_type: &DataType,
+ expected: ScalarValue,
+ ) -> Result<()> {
+ let ColumnarValue::Scalar(actual) =
round(&[ColumnarValue::Scalar(input)], data_type)?
+ else {
+ unreachable!()
+ };
+ assert_eq!(actual, expected);
+ Ok(())
+ }
+
+ #[test]
+ fn test_f32_array() -> Result<()> {
+ let ceil_input = [125.2345, 15.0001, 0.1, -0.9, -1.1, 123.0].map(Some);
+ let floor_input = [125.9345, 15.9999, 0.9, -0.1, -1.999,
123.0].map(Some);
+ let with_null = |values: [Option<f32>; 6]| {
+ let values: Vec<_> = values.into_iter().chain([None]).collect();
+ Arc::new(Float32Array::from(values)) as ArrayRef
+ };
+ check_int64_array(
+ spark_ceil,
+ with_null(ceil_input),
+ vec![
+ Some(126),
+ Some(16),
+ Some(1),
+ Some(0),
+ Some(-1),
+ Some(123),
+ None,
+ ],
+ )?;
+ check_int64_array(
+ spark_floor,
+ with_null(floor_input),
+ vec![
+ Some(125),
+ Some(15),
+ Some(0),
+ Some(-1),
+ Some(-2),
+ Some(123),
+ None,
+ ],
+ )
+ }
+
+ #[test]
+ fn test_f64_array() -> Result<()> {
+ let ceil_input = [125.2345, 15.0001, 0.1, -0.9, -1.1, 123.0].map(Some);
+ let floor_input = [125.9345, 15.9999, 0.9, -0.1, -1.999,
123.0].map(Some);
+ let with_null = |values: [Option<f64>; 6]| {
+ let values: Vec<_> = values.into_iter().chain([None]).collect();
+ Arc::new(Float64Array::from(values)) as ArrayRef
+ };
+ check_int64_array(
+ spark_ceil,
+ with_null(ceil_input),
+ vec![
+ Some(126),
+ Some(16),
+ Some(1),
+ Some(0),
+ Some(-1),
+ Some(123),
+ None,
+ ],
+ )?;
+ check_int64_array(
+ spark_floor,
+ with_null(floor_input),
+ vec![
+ Some(125),
+ Some(15),
+ Some(0),
+ Some(-1),
+ Some(-2),
+ Some(123),
+ None,
+ ],
+ )
+ }
+
+ #[test]
+ fn test_i64_array() -> Result<()> {
+ for round in [spark_ceil as RoundFn, spark_floor] {
+ let values = vec![Some(-1), Some(0), Some(1), None];
+ let input = Arc::new(Int64Array::from(values.clone())) as ArrayRef;
+ check_int64_array(round, input, values)?;
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn test_decimal128_array() -> Result<()> {
+ let input = [
+ Some(12345), // 123.45
+ Some(12500), // 125.00
+ Some(-12999), // -129.99
+ None,
+ ];
+ check_decimal_array(
+ spark_ceil,
+ &input,
+ (5, 2),
+ 4,
+ vec![Some(124), Some(125), Some(-129), None],
+ )?;
+ check_decimal_array(
+ spark_floor,
+ &input,
+ (5, 2),
+ 4,
+ vec![Some(123), Some(125), Some(-130), None],
+ )
+ }
+
+ #[test]
+ fn test_decimal128_wide_array() -> Result<()> {
+ // Unscaled values above i64::MAX (~9.22e18) take the wide fallback in
decimal_round_pow10.
+ // The first value is an exact multiple of 10^6, the other two have a
remainder.
+ let input = [
+ Some(20_000_000_000_000_000_000_i128),
+ Some(20_000_000_000_000_000_001_i128),
+ Some(-20_000_000_000_000_000_001_i128),
+ None,
+ ];
+ check_decimal_array(
+ spark_ceil,
+ &input,
+ (38, 6),
+ 38,
+ vec![
+ Some(20_000_000_000_000_i128),
+ Some(20_000_000_000_001_i128),
+ Some(-20_000_000_000_000_i128),
+ None,
+ ],
+ )?;
+ check_decimal_array(
+ spark_floor,
+ &input,
+ (38, 6),
+ 38,
+ vec![
+ Some(20_000_000_000_000_i128),
+ Some(20_000_000_000_000_i128),
+ Some(-20_000_000_000_001_i128),
+ None,
+ ],
+ )
+ }
+
+ #[test]
+ fn test_decimal128_i64_boundary_array() -> Result<()> {
+ // decimal_round_pow10 switches from the i64 fast path to
decimal_round_wide exactly at
+ // i64::MAX / i64::MIN. Straddling both edges pins that the two paths
agree: each pair
+ // below differs by one unscaled unit but rounds to the same value, so
a divergence
+ // between the 64-bit and 128-bit divisions would show up as a
mismatch here.
+ let input = [
+ Some(i64::MAX as i128), // 9223372036854.775807 -> fast path
+ Some(i64::MAX as i128 + 1), // 9223372036854.775808 -> wide path
+ Some(i64::MIN as i128), // -9223372036854.775808 -> fast path
+ Some(i64::MIN as i128 - 1), // -9223372036854.775809 -> wide path
+ ];
+ check_decimal_array(
+ spark_ceil,
+ &input,
+ (38, 6),
+ 38,
+ vec![
+ Some(9_223_372_036_855_i128),
+ Some(9_223_372_036_855_i128),
+ Some(-9_223_372_036_854_i128),
+ Some(-9_223_372_036_854_i128),
+ ],
+ )?;
+ check_decimal_array(
+ spark_floor,
+ &input,
+ (38, 6),
+ 38,
+ vec![
+ Some(9_223_372_036_854_i128),
+ Some(9_223_372_036_854_i128),
+ Some(-9_223_372_036_855_i128),
+ Some(-9_223_372_036_855_i128),
+ ],
+ )
+ }
+
+ #[test]
+ fn test_decimal128_large_scale_array() -> Result<()> {
Review Comment:
`ceil.sql` and `floor.sql` only cover `FLOAT`, `DOUBLE` and `DECIMAL(5, 2)`,
so a `BIGINT` column and a wide decimal column (`DECIMAL(38, 6)`, `DECIMAL(38,
19)`) with values past `i64::MAX` are never compared with Spark. These Rust
tests also pass `Decimal128(38, 0)` as the output type for a `(38, 6)` input,
where Spark plans `(33, 0)` (`p - s + 1`). Could a few such rows be added to
the two SQL files, read from a table so constant folding does not hide the
native path? The i64 boundary and scale 19 cases could then be checked there
against Spark too.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
