FreakingMind commented on code in PR #7534:
URL: https://github.com/apache/datafusion-comet/pull/7534#discussion_r4535678763
##########
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:
Fixed the precisions to` (33, 0)` and (20, 0), and added bigint,
`DECIMAL(38, 6)` and `DECIMAL(38, 19)` rows to both files, read from a table. I
also checked that ceil.sql fails if the wide path rounds the wrong way
--
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]
