diff --git a/native/spark-expr/src/math_funcs/ceil.rs b/native/spark-expr/src/math_funcs/ceil.rs deleted file mode 100644 index 292ef117e9a..00000000000 --- a/native/spark-expr/src/math_funcs/ceil.rs +++ /dev/null @@ -1,297 +0,0 @@ -// 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 crate::downcast_compute_op; -use crate::math_funcs::utils::{ - dispatch_pow10, get_precision_scale, make_decimal_array, make_decimal_scalar, -}; -use arrow::array::{Array, ArrowNativeTypeOp}; -use arrow::array::{Float32Array, Float64Array, Int64Array}; -use arrow::datatypes::DataType; -use datafusion::common::{DataFusionError, ScalarValue}; -use datafusion::physical_plan::ColumnarValue; -use num::integer::div_ceil; -use std::sync::Arc; - -/// `ceil` function that simulates Spark `ceil` expression -pub fn spark_ceil( - args: &[ColumnarValue], - data_type: &DataType, -) -> Result { - let value = &args[0]; - match value { - ColumnarValue::Array(array) => match array.data_type() { - DataType::Float32 => { - let result = downcast_compute_op!(array, "ceil", ceil, Float32Array, Int64Array); - Ok(ColumnarValue::Array(result?)) - } - DataType::Float64 => { - let result = downcast_compute_op!(array, "ceil", ceil, Float64Array, Int64Array); - Ok(ColumnarValue::Array(result?)) - } - DataType::Int64 => { - let result = array.as_any().downcast_ref::().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_ceil_pow10::), - make_decimal_array(array, precision, scale, decimal_ceil_f(*input_scale)) - ) - } - other => Err(DataFusionError::Internal(format!( - "Unsupported data type {other:?} for function ceil", - ))), - }, - ColumnarValue::Scalar(a) => match a { - ScalarValue::Float32(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64( - a.map(|x| x.ceil() as i64), - ))), - ScalarValue::Float64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64( - a.map(|x| x.ceil() 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_ceil_f(*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 ceil", - value.data_type(), - ))), - }, - } -} - -#[inline] -fn decimal_ceil_f(scale: i8) -> impl Fn(i128) -> i128 { - let div = 10_i128.pow_wrapping(scale as u32); - move |x: i128| div_ceil(x, div) -} - -/// Ceiling-divides an unscaled decimal by `10^EXP`. -/// -/// `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_ceil_pow10(x: i128) -> i128 { - match i64::try_from(x) { - Ok(x) => div_ceil(x, const { 10_i64.pow(EXP) }) as i128, - Err(_) => decimal_ceil_wide(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_ceil_pow10`] instantiation. -#[cold] -#[inline(never)] -fn decimal_ceil_wide(x: i128, div: i128) -> i128 { - div_ceil(x, div) -} - -#[cfg(test)] -mod test { - use crate::spark_ceil; - use arrow::array::{Decimal128Array, Float32Array, Float64Array, Int64Array}; - use arrow::datatypes::DataType; - use datafusion::common::cast::as_int64_array; - use datafusion::common::{Result, ScalarValue}; - use datafusion::physical_plan::ColumnarValue; - use std::sync::Arc; - - #[test] - fn test_ceil_f32_array() -> Result<()> { - let input = Float32Array::from(vec![ - Some(125.2345), - Some(15.0001), - Some(0.1), - Some(-0.9), - Some(-1.1), - Some(123.0), - None, - ]); - let args = vec![ColumnarValue::Array(Arc::new(input))]; - let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Float32)? else { - unreachable!() - }; - let actual = as_int64_array(&result)?; - let expected = Int64Array::from(vec![ - Some(126), - Some(16), - Some(1), - Some(0), - Some(-1), - Some(123), - None, - ]); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_ceil_f64_array() -> Result<()> { - let input = Float64Array::from(vec![ - Some(125.2345), - Some(15.0001), - Some(0.1), - Some(-0.9), - Some(-1.1), - Some(123.0), - None, - ]); - let args = vec![ColumnarValue::Array(Arc::new(input))]; - let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Float64)? else { - unreachable!() - }; - let actual = as_int64_array(&result)?; - let expected = Int64Array::from(vec![ - Some(126), - Some(16), - Some(1), - Some(0), - Some(-1), - Some(123), - None, - ]); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_ceil_i64_array() -> Result<()> { - let input = Int64Array::from(vec![Some(-1), Some(0), Some(1), None]); - let args = vec![ColumnarValue::Array(Arc::new(input))]; - let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Int64)? else { - unreachable!() - }; - let actual = as_int64_array(&result)?; - let expected = Int64Array::from(vec![Some(-1), Some(0), Some(1), None]); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_ceil_decimal128_array() -> Result<()> { - let array = Decimal128Array::from(vec![ - Some(12345), // 123.45 - Some(12500), // 125.00 - Some(-12999), // -129.99 - None, - ]) - .with_precision_and_scale(5, 2)?; - let args = vec![ColumnarValue::Array(Arc::new(array))]; - let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Decimal128(4, 0))? else { - unreachable!() - }; - let expected = Decimal128Array::from(vec![ - Some(124), // 124.00 - Some(125), // 125.00 - Some(-129), // -129.00 - None, - ]) - .with_precision_and_scale(4, 0)?; - let actual = result.as_any().downcast_ref::().unwrap(); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_ceil_decimal128_wide_array() -> Result<()> { - // Unscaled values that exceed i64::MAX (~9.22e18) exercise the wide fallback in - // decimal_ceil_pow10. Values chosen at scale 6 targeting Decimal128(38, 0): - // 20_000_000_000_000_000_000 / 10^6 = 20_000_000_000_000 (exact multiple) - // 20_000_000_000_000_000_001 / 10^6 -> ceil 20_000_000_000_001 (positive remainder) - // -20_000_000_000_000_000_001 / 10^6 -> ceil -20_000_000_000_000 (negative remainder) - let array = Decimal128Array::from(vec![ - 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, - ]) - .with_precision_and_scale(38, 6)?; - let args = vec![ColumnarValue::Array(Arc::new(array))]; - let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Decimal128(38, 0))? else { - unreachable!() - }; - let expected = Decimal128Array::from(vec![ - Some(20_000_000_000_000_i128), - Some(20_000_000_000_001_i128), - Some(-20_000_000_000_000_i128), - None, - ]) - .with_precision_and_scale(38, 0)?; - let actual = result.as_any().downcast_ref::().unwrap(); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_ceil_f32_scalar() -> Result<()> { - let args = vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(125.2345)))]; - let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) = - spark_ceil(&args, &DataType::Float32)? - else { - unreachable!() - }; - assert_eq!(result, 126); - Ok(()) - } - - #[test] - fn test_ceil_f64_scalar() -> Result<()> { - let args = vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(-1.1)))]; - let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) = - spark_ceil(&args, &DataType::Float64)? - else { - unreachable!() - }; - assert_eq!(result, -1); - Ok(()) - } - - #[test] - fn test_ceil_i64_scalar() -> Result<()> { - let args = vec![ColumnarValue::Scalar(ScalarValue::Int64(Some(48)))]; - let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) = - spark_ceil(&args, &DataType::Int64)? - else { - unreachable!() - }; - assert_eq!(result, 48); - Ok(()) - } - - #[test] - fn test_ceil_decimal128_scalar() -> Result<()> { - let args = vec![ColumnarValue::Scalar(ScalarValue::Decimal128( - Some(567), - 3, - 1, - ))]; // 56.7 - let ColumnarValue::Scalar(ScalarValue::Decimal128(Some(result), 3, 0)) = - spark_ceil(&args, &DataType::Decimal128(3, 0))? - else { - unreachable!() - }; - assert_eq!(result, 57); // 57.0 - Ok(()) - } -} diff --git a/native/spark-expr/src/math_funcs/ceil_floor.rs b/native/spark-expr/src/math_funcs/ceil_floor.rs new file mode 100644 index 00000000000..d2290882d8b --- /dev/null +++ b/native/spark-expr/src/math_funcs/ceil_floor.rs @@ -0,0 +1,568 @@ +// 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 crate::math_funcs::utils::{get_precision_scale, make_decimal_array, make_decimal_scalar}; +use arrow::array::cast::AsArray; +use arrow::array::types::{Float32Type, Float64Type}; +use arrow::array::{Array, ArrowNativeTypeOp, Int64Array}; +use arrow::datatypes::DataType; +use datafusion::common::{DataFusionError, ScalarValue}; +use datafusion::physical_plan::ColumnarValue; +use num::integer::{div_ceil, div_floor, Integer}; +use num::Float; +use std::sync::Arc; + +/// Evaluates `$body` with `$exp` bound as a compile-time constant equal to `$scale`, for every +/// decimal scale whose divisor `10^scale` fits in an `i64`. Larger scales evaluate `$fallback`, +/// which must derive the divisor at runtime. +/// +/// Specializing on the scale lets the divisor be folded into the generated code, turning the +/// division by `10^scale` into a multiply-and-shift. +macro_rules! dispatch_pow10 { + ($scale:expr, $exp:ident => $body:expr, $fallback:expr) => { + match $scale { + 1 => { + const $exp: u32 = 1; + $body + } + 2 => { + const $exp: u32 = 2; + $body + } + 3 => { + const $exp: u32 = 3; + $body + } + 4 => { + const $exp: u32 = 4; + $body + } + 5 => { + const $exp: u32 = 5; + $body + } + 6 => { + const $exp: u32 = 6; + $body + } + 7 => { + const $exp: u32 = 7; + $body + } + 8 => { + const $exp: u32 = 8; + $body + } + 9 => { + const $exp: u32 = 9; + $body + } + 10 => { + const $exp: u32 = 10; + $body + } + 11 => { + const $exp: u32 = 11; + $body + } + 12 => { + const $exp: u32 = 12; + $body + } + 13 => { + const $exp: u32 = 13; + $body + } + 14 => { + const $exp: u32 = 14; + $body + } + 15 => { + const $exp: u32 = 15; + $body + } + 16 => { + const $exp: u32 = 16; + $body + } + 17 => { + const $exp: u32 = 17; + $body + } + 18 => { + const $exp: u32 = 18; + $body + } + _ => $fallback, + } + }; +} + +/// `ceil` function that simulates Spark `ceil` expression +pub fn spark_ceil( + args: &[ColumnarValue], + data_type: &DataType, +) -> Result { + spark_ceil_or_floor::(args, data_type) +} + +/// `floor` function that simulates Spark `floor` expression +pub fn spark_floor( + args: &[ColumnarValue], + data_type: &DataType, +) -> Result { + spark_ceil_or_floor::(args, data_type) +} + +#[inline] +fn round_float(x: T) -> T { + if CEIL { + x.ceil() + } else { + x.floor() + } +} + +#[inline] +fn div_round(x: T, div: T) -> T { + if CEIL { + div_ceil(x, div) + } else { + div_floor(x, div) + } +} + +/// Shared implementation of Spark's `ceil` and `floor` expressions. +/// +/// `CEIL` selects the rounding direction. It is a const parameter so that each direction is +/// monomorphized into the element loop rather than branched on for every value. +fn spark_ceil_or_floor( + args: &[ColumnarValue], + data_type: &DataType, +) -> Result { + let name = if CEIL { "ceil" } else { "floor" }; + 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::(), + |x| round_float::(x) as i64, + ); + Ok(ColumnarValue::Array(Arc::new(result))) + } + DataType::Float64 => { + let result: Int64Array = arrow::compute::kernels::arity::unary( + array.as_primitive::(), + |x| round_float::(x) as i64, + ); + Ok(ColumnarValue::Array(Arc::new(result))) + } + DataType::Int64 => { + let result = array.as_any().downcast_ref::().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::), + make_decimal_array(array, precision, scale, decimal_round_f::(*input_scale)) + ) + } + other => Err(DataFusionError::Internal(format!( + "Unsupported data type {other:?} for function {}", + name, + ))), + }, + ColumnarValue::Scalar(a) => match a { + ScalarValue::Float32(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64( + a.map(|x| round_float::(x) as i64), + ))), + ScalarValue::Float64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64( + a.map(|x| round_float::(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::(*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(), + name, + ))), + }, + } +} + +#[inline] +fn decimal_round_f(scale: i8) -> impl Fn(i128) -> i128 { + let div = 10_i128.pow_wrapping(scale as u32); + move |x: i128| div_round::(x, div) +} + +/// Rounds an unscaled decimal divided by `10^EXP` up if `CEIL` is set, down otherwise. +/// +/// `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(x: i128) -> i128 { + match i64::try_from(x) { + Ok(x) => div_round::(x, const { 10_i64.pow(EXP) }) as i128, + Err(_) => decimal_round_wide::(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(x: i128, div: i128) -> i128 { + div_round::(x, div) +} + +#[cfg(test)] +mod test { + use super::*; + use arrow::array::{ArrayRef, Decimal128Array, Float32Array, Float64Array}; + use datafusion::common::cast::as_int64_array; + use datafusion::common::Result; + + type RoundFn = fn(&[ColumnarValue], &DataType) -> Result; + + fn check_int64_array( + round: RoundFn, + input: ArrayRef, + expected: Vec>, + ) -> 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], + (precision, scale): (u8, i8), + output_precision: u8, + expected: Vec>, + ) -> 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::().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; 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; 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), + 33, + 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), + 33, + 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), + 33, + 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), + 33, + 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<()> { + // dispatch_pow10! only specializes scales whose divisor fits in an i64, so scale 19 takes + // the runtime `decimal_round_f` fallback. Untested on the array path otherwise. + let input = [ + Some(20_000_000_000_000_000_000_000_i128), // 2000.0 + Some(20_000_000_000_000_000_000_001_i128), // 2000.0000000000000000001 + Some(-20_000_000_000_000_000_000_001_i128), // -2000.0000000000000000001 + None, + ]; + check_decimal_array( + spark_ceil, + &input, + (38, 19), + 20, + vec![Some(2000), Some(2001), Some(-2000), None], + )?; + check_decimal_array( + spark_floor, + &input, + (38, 19), + 20, + vec![Some(2000), Some(2000), Some(-2001), None], + ) + } + + #[test] + fn test_f32_scalar() -> Result<()> { + check_scalar( + spark_ceil, + ScalarValue::Float32(Some(125.2345)), + &DataType::Float32, + ScalarValue::Int64(Some(126)), + )?; + check_scalar( + spark_floor, + ScalarValue::Float32(Some(125.9345)), + &DataType::Float32, + ScalarValue::Int64(Some(125)), + ) + } + + #[test] + fn test_f64_scalar() -> Result<()> { + check_scalar( + spark_ceil, + ScalarValue::Float64(Some(-1.1)), + &DataType::Float64, + ScalarValue::Int64(Some(-1)), + )?; + check_scalar( + spark_floor, + ScalarValue::Float64(Some(-1.999)), + &DataType::Float64, + ScalarValue::Int64(Some(-2)), + ) + } + + #[test] + fn test_i64_scalar() -> Result<()> { + for round in [spark_ceil as RoundFn, spark_floor] { + check_scalar( + round, + ScalarValue::Int64(Some(48)), + &DataType::Int64, + ScalarValue::Int64(Some(48)), + )?; + } + Ok(()) + } + + #[test] + fn test_decimal128_scalar() -> Result<()> { + let input = ScalarValue::Decimal128(Some(567), 3, 1); // 56.7 + let output_type = DataType::Decimal128(3, 0); + check_scalar( + spark_ceil, + input.clone(), + &output_type, + ScalarValue::Decimal128(Some(57), 3, 0), + )?; + check_scalar( + spark_floor, + input, + &output_type, + ScalarValue::Decimal128(Some(56), 3, 0), + ) + } +} diff --git a/native/spark-expr/src/math_funcs/floor.rs b/native/spark-expr/src/math_funcs/floor.rs deleted file mode 100644 index f97c66e2c86..00000000000 --- a/native/spark-expr/src/math_funcs/floor.rs +++ /dev/null @@ -1,353 +0,0 @@ -// 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 crate::downcast_compute_op; -use crate::math_funcs::utils::{ - dispatch_pow10, get_precision_scale, make_decimal_array, make_decimal_scalar, -}; -use arrow::array::{Array, ArrowNativeTypeOp}; -use arrow::array::{Float32Array, Float64Array, Int64Array}; -use arrow::datatypes::DataType; -use datafusion::common::{DataFusionError, ScalarValue}; -use datafusion::physical_plan::ColumnarValue; -use num::integer::div_floor; -use std::sync::Arc; - -/// `floor` function that simulates Spark `floor` expression -pub fn spark_floor( - args: &[ColumnarValue], - data_type: &DataType, -) -> Result { - let value = &args[0]; - match value { - ColumnarValue::Array(array) => match array.data_type() { - DataType::Float32 => { - let result = downcast_compute_op!(array, "floor", floor, Float32Array, Int64Array); - Ok(ColumnarValue::Array(result?)) - } - DataType::Float64 => { - let result = downcast_compute_op!(array, "floor", floor, Float64Array, Int64Array); - Ok(ColumnarValue::Array(result?)) - } - DataType::Int64 => { - let result = array.as_any().downcast_ref::().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_floor_pow10::), - make_decimal_array(array, precision, scale, decimal_floor_f(*input_scale)) - ) - } - other => Err(DataFusionError::Internal(format!( - "Unsupported data type {other:?} for function floor", - ))), - }, - ColumnarValue::Scalar(a) => match a { - ScalarValue::Float32(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64( - a.map(|x| x.floor() as i64), - ))), - ScalarValue::Float64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64( - a.map(|x| x.floor() 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_floor_f(*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 floor", - value.data_type(), - ))), - }, - } -} - -#[inline] -fn decimal_floor_f(scale: i8) -> impl Fn(i128) -> i128 { - let div = 10_i128.pow_wrapping(scale as u32); - move |x: i128| div_floor(x, div) -} - -/// Floor-divides an unscaled decimal by `10^EXP`. -/// -/// `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_floor_pow10(x: i128) -> i128 { - match i64::try_from(x) { - Ok(x) => div_floor(x, const { 10_i64.pow(EXP) }) as i128, - Err(_) => decimal_floor_wide(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_floor_pow10`] instantiation. -#[cold] -#[inline(never)] -fn decimal_floor_wide(x: i128, div: i128) -> i128 { - div_floor(x, div) -} - -#[cfg(test)] -mod test { - use crate::spark_floor; - use arrow::array::{Decimal128Array, Float32Array, Float64Array, Int64Array}; - use arrow::datatypes::DataType; - use datafusion::common::cast::as_int64_array; - use datafusion::common::{Result, ScalarValue}; - use datafusion::physical_plan::ColumnarValue; - use std::sync::Arc; - - #[test] - fn test_floor_f32_array() -> Result<()> { - let input = Float32Array::from(vec![ - Some(125.9345), - Some(15.9999), - Some(0.9), - Some(-0.1), - Some(-1.999), - Some(123.0), - None, - ]); - let args = vec![ColumnarValue::Array(Arc::new(input))]; - let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Float32)? else { - unreachable!() - }; - let actual = as_int64_array(&result)?; - let expected = Int64Array::from(vec![ - Some(125), - Some(15), - Some(0), - Some(-1), - Some(-2), - Some(123), - None, - ]); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_floor_f64_array() -> Result<()> { - let input = Float64Array::from(vec![ - Some(125.9345), - Some(15.9999), - Some(0.9), - Some(-0.1), - Some(-1.999), - Some(123.0), - None, - ]); - let args = vec![ColumnarValue::Array(Arc::new(input))]; - let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Float64)? else { - unreachable!() - }; - let actual = as_int64_array(&result)?; - let expected = Int64Array::from(vec![ - Some(125), - Some(15), - Some(0), - Some(-1), - Some(-2), - Some(123), - None, - ]); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_floor_i64_array() -> Result<()> { - let input = Int64Array::from(vec![Some(-1), Some(0), Some(1), None]); - let args = vec![ColumnarValue::Array(Arc::new(input))]; - let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Int64)? else { - unreachable!() - }; - let actual = as_int64_array(&result)?; - let expected = Int64Array::from(vec![Some(-1), Some(0), Some(1), None]); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_floor_decimal128_array() -> Result<()> { - let array = Decimal128Array::from(vec![ - Some(12345), // 123.45 - Some(12500), // 125.00 - Some(-12999), // -129.99 - None, - ]) - .with_precision_and_scale(5, 2)?; - let args = vec![ColumnarValue::Array(Arc::new(array))]; - let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Decimal128(4, 0))? else { - unreachable!() - }; - let expected = Decimal128Array::from(vec![ - Some(123), // 123.00 - Some(125), // 125.00 - Some(-130), // -130.00 - None, - ]) - .with_precision_and_scale(4, 0)?; - let actual = result.as_any().downcast_ref::().unwrap(); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_floor_decimal128_wide_array() -> Result<()> { - // Unscaled values above i64::MAX (~9.22e18) take the wide fallback in decimal_floor_pow10. - // Scale 6 targeting Decimal128(38, 0): - // 20_000_000_000_000_000_000 / 10^6 = 20_000_000_000_000 (exact multiple) - // 20_000_000_000_000_000_001 / 10^6 -> floor 20_000_000_000_000 (toward zero) - // -20_000_000_000_000_000_001 / 10^6 -> floor -20_000_000_000_001 (toward -inf) - let array = Decimal128Array::from(vec![ - 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, - ]) - .with_precision_and_scale(38, 6)?; - let args = vec![ColumnarValue::Array(Arc::new(array))]; - let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Decimal128(38, 0))? else { - unreachable!() - }; - let expected = Decimal128Array::from(vec![ - Some(20_000_000_000_000_i128), - Some(20_000_000_000_000_i128), - Some(-20_000_000_000_001_i128), - None, - ]) - .with_precision_and_scale(38, 0)?; - let actual = result.as_any().downcast_ref::().unwrap(); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_floor_decimal128_i64_boundary_array() -> Result<()> { - // decimal_floor_pow10 switches from the i64 fast path to decimal_floor_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 floors to the same value, so a divergence - // between the 64-bit and 128-bit divisions would show up as a mismatch here. - let array = Decimal128Array::from(vec![ - 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 - ]) - .with_precision_and_scale(38, 6)?; - let args = vec![ColumnarValue::Array(Arc::new(array))]; - let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Decimal128(38, 0))? else { - unreachable!() - }; - let expected = Decimal128Array::from(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), - ]) - .with_precision_and_scale(38, 0)?; - let actual = result.as_any().downcast_ref::().unwrap(); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_floor_decimal128_large_scale_array() -> Result<()> { - // dispatch_pow10! only specializes scales whose divisor fits in an i64, so scale 19 takes - // the runtime `decimal_floor_f` fallback. Untested on the array path otherwise. - let array = Decimal128Array::from(vec![ - Some(20_000_000_000_000_000_000_000_i128), // 2000.0 - Some(20_000_000_000_000_000_000_001_i128), // 2000.0000000000000000001 - Some(-20_000_000_000_000_000_000_001_i128), // -2000.0000000000000000001 - None, - ]) - .with_precision_and_scale(38, 19)?; - let args = vec![ColumnarValue::Array(Arc::new(array))]; - let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Decimal128(38, 0))? else { - unreachable!() - }; - let expected = Decimal128Array::from(vec![ - Some(2000_i128), - Some(2000_i128), - Some(-2001_i128), - None, - ]) - .with_precision_and_scale(38, 0)?; - let actual = result.as_any().downcast_ref::().unwrap(); - assert_eq!(actual, &expected); - Ok(()) - } - - #[test] - fn test_floor_f32_scalar() -> Result<()> { - let args = vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(125.9345)))]; - let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) = - spark_floor(&args, &DataType::Float32)? - else { - unreachable!() - }; - assert_eq!(result, 125); - Ok(()) - } - - #[test] - fn test_floor_f64_scalar() -> Result<()> { - let args = vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(-1.999)))]; - let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) = - spark_floor(&args, &DataType::Float64)? - else { - unreachable!() - }; - assert_eq!(result, -2); - Ok(()) - } - - #[test] - fn test_floor_i64_scalar() -> Result<()> { - let args = vec![ColumnarValue::Scalar(ScalarValue::Int64(Some(48)))]; - let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) = - spark_floor(&args, &DataType::Int64)? - else { - unreachable!() - }; - assert_eq!(result, 48); - Ok(()) - } - - #[test] - fn test_floor_decimal128_scalar() -> Result<()> { - let args = vec![ColumnarValue::Scalar(ScalarValue::Decimal128( - Some(567), - 3, - 1, - ))]; // 56.7 - let ColumnarValue::Scalar(ScalarValue::Decimal128(Some(result), 3, 0)) = - spark_floor(&args, &DataType::Decimal128(3, 0))? - else { - unreachable!() - }; - assert_eq!(result, 56); // 56.0 - Ok(()) - } -} diff --git a/native/spark-expr/src/math_funcs/mod.rs b/native/spark-expr/src/math_funcs/mod.rs index b27f9811f1d..de19eed9323 100644 --- a/native/spark-expr/src/math_funcs/mod.rs +++ b/native/spark-expr/src/math_funcs/mod.rs @@ -16,11 +16,10 @@ // under the License. pub(crate) mod abs; -mod ceil; +mod ceil_floor; pub(crate) mod checked_arithmetic; mod checked_binary_expr; mod div; -mod floor; mod greatest_least; pub mod internal; pub(crate) mod log; @@ -35,12 +34,11 @@ mod utils; mod wide_decimal_binary_expr; pub use abs::abs; -pub use ceil::spark_ceil; +pub use ceil_floor::{spark_ceil, spark_floor}; pub use checked_arithmetic::{checked_add, checked_div, checked_mul, checked_sub}; pub use checked_binary_expr::CheckedBinaryExpr; pub use div::spark_decimal_div; pub use div::spark_decimal_integral_div; -pub use floor::spark_floor; pub use greatest_least::SparkGreatestLeast; pub use internal::*; pub use log::spark_log; diff --git a/native/spark-expr/src/math_funcs/utils.rs b/native/spark-expr/src/math_funcs/utils.rs index 23e7b336ef8..9c68295c576 100644 --- a/native/spark-expr/src/math_funcs/utils.rs +++ b/native/spark-expr/src/math_funcs/utils.rs @@ -23,112 +23,6 @@ use datafusion::common::{DataFusionError, ScalarValue}; use datafusion::physical_plan::ColumnarValue; use std::sync::Arc; -#[macro_export] -macro_rules! downcast_compute_op { - ($ARRAY:expr, $NAME:expr, $FUNC:ident, $TYPE:ident, $RESULT:ident) => {{ - let n = $ARRAY.as_any().downcast_ref::<$TYPE>(); - match n { - Some(array) => { - let res: $RESULT = - arrow::compute::kernels::arity::unary(array, |x| x.$FUNC() as i64); - Ok(Arc::new(res)) - } - _ => Err(DataFusionError::Internal(format!( - "Invalid data type for {}", - $NAME - ))), - } - }}; -} - -/// Evaluates `$body` with `$exp` bound as a compile-time constant equal to `$scale`, for every -/// decimal scale whose divisor `10^scale` fits in an `i64`. Larger scales evaluate `$fallback`, -/// which must derive the divisor at runtime. -/// -/// Specializing on the scale lets the divisor be folded into the generated code, turning the -/// division by `10^scale` into a multiply-and-shift. -macro_rules! dispatch_pow10 { - ($scale:expr, $exp:ident => $body:expr, $fallback:expr) => { - match $scale { - 1 => { - const $exp: u32 = 1; - $body - } - 2 => { - const $exp: u32 = 2; - $body - } - 3 => { - const $exp: u32 = 3; - $body - } - 4 => { - const $exp: u32 = 4; - $body - } - 5 => { - const $exp: u32 = 5; - $body - } - 6 => { - const $exp: u32 = 6; - $body - } - 7 => { - const $exp: u32 = 7; - $body - } - 8 => { - const $exp: u32 = 8; - $body - } - 9 => { - const $exp: u32 = 9; - $body - } - 10 => { - const $exp: u32 = 10; - $body - } - 11 => { - const $exp: u32 = 11; - $body - } - 12 => { - const $exp: u32 = 12; - $body - } - 13 => { - const $exp: u32 = 13; - $body - } - 14 => { - const $exp: u32 = 14; - $body - } - 15 => { - const $exp: u32 = 15; - $body - } - 16 => { - const $exp: u32 = 16; - $body - } - 17 => { - const $exp: u32 = 17; - $body - } - 18 => { - const $exp: u32 = 18; - $body - } - _ => $fallback, - } - }; -} - -pub(crate) use dispatch_pow10; - #[inline] pub(crate) fn make_decimal_scalar( a: &Option, diff --git a/spark/src/test/resources/sql-tests/expressions/math/ceil.sql b/spark/src/test/resources/sql-tests/expressions/math/ceil.sql index a0fa163361b..38cf86af109 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/ceil.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/ceil.sql @@ -29,3 +29,14 @@ SELECT ceil(f), ceil(d), ceil(dec) FROM test_ceil -- literal arguments query SELECT ceil(1.1), ceil(-1.1), ceil(0.0), ceil(NULL), ceil(cast(1.10 as DECIMAL(5, 2))), ceil(cast(-1.10 as DECIMAL(5, 2))) + +-- bigint and wide decimals. Unscaled DECIMAL(38, 6) values cross i64::MAX and i64::MIN, and +-- DECIMAL(38, 19) has a divisor that does not fit in an i64 +statement +CREATE TABLE test_ceil_wide(l bigint, d6 DECIMAL(38, 6), d19 DECIMAL(38, 19)) USING parquet + +statement +INSERT INTO test_ceil_wide VALUES (9223372036854775807, 9223372036854.775807, 2000.0000000000000000000), (-9223372036854775808, 9223372036854.775808, 2000.0000000000000000001), (0, -9223372036854.775808, -2000.0000000000000000001), (1, -9223372036854.775809, 9999999999999999999.9999999999999999999), (-1, 20000000000000.000001, -9999999999999999999.9999999999999999999), (NULL, -20000000000000.000001, 0.0000000000000000001), (42, 99999999999999999999999999999999.999999, -0.0000000000000000001), (-42, -99999999999999999999999999999999.999999, NULL) + +query +SELECT ceil(l), ceil(d6), ceil(d19) FROM test_ceil_wide diff --git a/spark/src/test/resources/sql-tests/expressions/math/floor.sql b/spark/src/test/resources/sql-tests/expressions/math/floor.sql index fd59dc6857d..776bca10b09 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/floor.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/floor.sql @@ -29,3 +29,14 @@ SELECT floor(f), floor(d), floor(dec) FROM test_floor -- literal arguments query SELECT floor(1.9), floor(-1.1), floor(0.0), floor(NULL), floor(cast(1.90 as DECIMAL(5, 2))), floor(cast(-1.10 as DECIMAL(5, 2))) + +-- bigint and wide decimals. Unscaled DECIMAL(38, 6) values cross i64::MAX and i64::MIN, and +-- DECIMAL(38, 19) has a divisor that does not fit in an i64 +statement +CREATE TABLE test_floor_wide(l bigint, d6 DECIMAL(38, 6), d19 DECIMAL(38, 19)) USING parquet + +statement +INSERT INTO test_floor_wide VALUES (9223372036854775807, 9223372036854.775807, 2000.0000000000000000000), (-9223372036854775808, 9223372036854.775808, 2000.0000000000000000001), (0, -9223372036854.775808, -2000.0000000000000000001), (1, -9223372036854.775809, 9999999999999999999.9999999999999999999), (-1, 20000000000000.000001, -9999999999999999999.9999999999999999999), (NULL, -20000000000000.000001, 0.0000000000000000001), (42, 99999999999999999999999999999999.999999, -0.0000000000000000001), (-42, -99999999999999999999999999999999.999999, NULL) + +query +SELECT floor(l), floor(d6), floor(d19) FROM test_floor_wide