| // 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::arithmetic_overflow_error; |
| use crate::math_funcs::utils::{get_precision_scale, make_decimal_array, make_decimal_scalar}; |
| use arrow::array::{Array, ArrowNativeTypeOp}; |
| use arrow::array::{Int16Array, Int32Array, Int64Array, Int8Array}; |
| use arrow::datatypes::DataType; |
| use arrow::error::ArrowError; |
| use datafusion::common::{exec_err, internal_err, DataFusionError, ScalarValue}; |
| use datafusion::physical_plan::ColumnarValue; |
| use std::{cmp::min, sync::Arc}; |
| |
| macro_rules! integer_round { |
| ($X:expr, $DIV:expr, $HALF:expr, $FAIL_ON_ERROR:expr) => {{ |
| let rem = $X % $DIV; |
| if rem <= -$HALF { |
| if $FAIL_ON_ERROR { |
| ($X - rem).sub_checked($DIV).map_err(|_| { |
| ArrowError::ComputeError(arithmetic_overflow_error("integer").to_string()) |
| }) |
| } else { |
| Ok(($X - rem).sub_wrapping($DIV)) |
| } |
| } else if rem >= $HALF { |
| if $FAIL_ON_ERROR { |
| ($X - rem).add_checked($DIV).map_err(|_| { |
| ArrowError::ComputeError(arithmetic_overflow_error("integer").to_string()) |
| }) |
| } else { |
| Ok(($X - rem).add_wrapping($DIV)) |
| } |
| } else { |
| if $FAIL_ON_ERROR { |
| $X.sub_checked(rem).map_err(|_| { |
| ArrowError::ComputeError(arithmetic_overflow_error("integer").to_string()) |
| }) |
| } else { |
| Ok($X.sub_wrapping(rem)) |
| } |
| } |
| }}; |
| } |
| |
| // Round a single native integer when `10^(-point)` does not fit in the native |
| // integer type but still fits in i128. The caller has already excluded the |
| // case where `div` fits the native type, so `|x| <= NATIVE::MAX < div`, which |
| // makes `x % div == x`: no division is needed here, and the result is either |
| // `0` (when `|x| < half`) or `sign(x) * div` — the latter always overflows the |
| // native type. Under ANSI we throw, under legacy we wrap by truncation |
| // (matches `BigDecimal.longValue`'s low-64-bit semantics). |
| macro_rules! integer_round_widened { |
| ($X:expr, $DIV:expr, $HALF:expr, $NATIVE:ty, $FAIL_ON_ERROR:expr) => {{ |
| let x128 = $X as i128; |
| debug_assert!( |
| x128 > -$DIV && x128 < $DIV, |
| "integer_round_widened! requires div to overflow the native type" |
| ); |
| if x128 > -$HALF && x128 < $HALF { |
| Ok(0 as $NATIVE) |
| } else if $FAIL_ON_ERROR { |
| Err(ArrowError::ComputeError( |
| arithmetic_overflow_error("integer").to_string(), |
| )) |
| } else { |
| Ok((if x128 >= $HALF { $DIV } else { -$DIV }) as $NATIVE) |
| } |
| }}; |
| } |
| |
| // Lift an `ArrowError` from the rounding macros into a `DataFusionError`, |
| // returning early from the enclosing function. Used by `round_integer_scalar!`. |
| macro_rules! round_scalar_result { |
| ($RESULT:expr) => { |
| match $RESULT { |
| Ok(v) => Some(v), |
| Err(e) => { |
| return Err(DataFusionError::ArrowError( |
| Box::from(e), |
| Some(DataFusionError::get_back_trace()), |
| )) |
| } |
| } |
| }; |
| } |
| |
| macro_rules! round_integer_array { |
| ($ARRAY:expr, $POINT:expr, $TYPE:ty, $NATIVE:ty, $FAIL_ON_ERROR:expr) => {{ |
| let array = $ARRAY.as_any().downcast_ref::<$TYPE>().unwrap(); |
| let ten: $NATIVE = 10; |
| let point_abs = (-(*$POINT)) as u32; |
| let result: $TYPE = if let Some(div) = ten.checked_pow(point_abs) { |
| let half = div / 2; |
| arrow::compute::kernels::arity::try_unary(array, |x| { |
| integer_round!(x, div, half, $FAIL_ON_ERROR) |
| })? |
| } else if let Some(div) = 10_i128.checked_pow(point_abs) { |
| let half = div / 2; |
| arrow::compute::kernels::arity::try_unary(array, |x| { |
| integer_round_widened!(x, div, half, $NATIVE, $FAIL_ON_ERROR) |
| })? |
| } else { |
| // Even i128 cannot hold 10^(-point); every bounded native |
| // integer rounds to 0. |
| arrow::compute::kernels::arity::try_unary(array, |_| Ok(0))? |
| }; |
| Ok(ColumnarValue::Array(Arc::new(result))) |
| }}; |
| } |
| |
| macro_rules! round_integer_scalar { |
| ($SCALAR:expr, $POINT:expr, $TYPE:expr, $NATIVE:ty, $FAIL_ON_ERROR:expr) => {{ |
| let ten: $NATIVE = 10; |
| let point_abs = (-(*$POINT)) as u32; |
| let scalar_opt = match $SCALAR { |
| None => None, |
| Some(x) => { |
| if let Some(div) = ten.checked_pow(point_abs) { |
| let half = div / 2; |
| round_scalar_result!(integer_round!(*x, div, half, $FAIL_ON_ERROR)) |
| } else if let Some(div) = 10_i128.checked_pow(point_abs) { |
| let half = div / 2; |
| round_scalar_result!(integer_round_widened!( |
| *x, |
| div, |
| half, |
| $NATIVE, |
| $FAIL_ON_ERROR |
| )) |
| } else { |
| // Even i128 cannot hold 10^(-point); every bounded native |
| // integer rounds to 0. |
| Some(0) |
| } |
| } |
| }; |
| Ok(ColumnarValue::Scalar($TYPE(scalar_opt))) |
| }}; |
| } |
| |
| /// `round` function that simulates Spark `round` expression |
| /// |
| /// Float and double are deliberately absent: Spark rounds them through a `BigDecimal` built from |
| /// `java.lang.Double.toString()`, which no native kernel reproduces, so `CometRound` reports those |
| /// inputs as `Unsupported` and routes them to the JVM codegen dispatcher instead. They never reach |
| /// this function, and fall through to the catch-all error below if they somehow do. |
| pub fn spark_round( |
| args: &[ColumnarValue], |
| data_type: &DataType, |
| fail_on_error: bool, |
| ) -> Result<ColumnarValue, DataFusionError> { |
| let value = &args[0]; |
| let point = &args[1]; |
| let ColumnarValue::Scalar(ScalarValue::Int64(Some(point))) = point else { |
| return internal_err!("Invalid point argument for Round(): {:#?}", point); |
| }; |
| match value { |
| ColumnarValue::Array(array) => match array.data_type() { |
| DataType::Int64 if *point < 0 => { |
| round_integer_array!(array, point, Int64Array, i64, fail_on_error) |
| } |
| DataType::Int32 if *point < 0 => { |
| round_integer_array!(array, point, Int32Array, i32, fail_on_error) |
| } |
| DataType::Int16 if *point < 0 => { |
| round_integer_array!(array, point, Int16Array, i16, fail_on_error) |
| } |
| DataType::Int8 if *point < 0 => { |
| round_integer_array!(array, point, Int8Array, i8, fail_on_error) |
| } |
| DataType::Decimal128(_, scale) if *scale >= 0 => { |
| let f = decimal_round_f(scale, point); |
| let (precision, scale) = get_precision_scale(data_type); |
| make_decimal_array(array, precision, scale, &f) |
| } |
| dt => exec_err!("Not supported datatype for ROUND: {dt}"), |
| }, |
| ColumnarValue::Scalar(a) => match a { |
| ScalarValue::Int64(a) if *point < 0 => { |
| round_integer_scalar!(a, point, ScalarValue::Int64, i64, fail_on_error) |
| } |
| ScalarValue::Int32(a) if *point < 0 => { |
| round_integer_scalar!(a, point, ScalarValue::Int32, i32, fail_on_error) |
| } |
| ScalarValue::Int16(a) if *point < 0 => { |
| round_integer_scalar!(a, point, ScalarValue::Int16, i16, fail_on_error) |
| } |
| ScalarValue::Int8(a) if *point < 0 => { |
| round_integer_scalar!(a, point, ScalarValue::Int8, i8, fail_on_error) |
| } |
| ScalarValue::Decimal128(a, _, scale) if *scale >= 0 => { |
| let f = decimal_round_f(scale, point); |
| let (precision, scale) = get_precision_scale(data_type); |
| make_decimal_scalar(a, precision, scale, &f) |
| } |
| dt => exec_err!("Not supported datatype for ROUND: {dt}"), |
| }, |
| } |
| } |
| |
| // Spark uses BigDecimal. See RoundBase implementation in Spark. Instead, we do the same by |
| // 1) add the half of divisor, 2) round down by division, 3) adjust precision by multiplication |
| #[inline] |
| fn decimal_round_f(scale: &i8, point: &i64) -> Box<dyn Fn(i128) -> i128> { |
| if *point < 0 { |
| if let Some(div) = 10_i128.checked_pow((-(*point) as u32) + (*scale as u32)) { |
| let half = div / 2; |
| let mul = 10_i128.pow_wrapping((-(*point)) as u32); |
| // i128 can hold 39 digits of a base 10 number, adding half will not cause overflow |
| Box::new(move |x: i128| (x + x.signum() * half) / div * mul) |
| } else { |
| Box::new(move |_: i128| 0) |
| } |
| } else { |
| let div = 10_i128.pow_wrapping((*scale as u32) - min(*scale as u32, *point as u32)); |
| let half = div / 2; |
| Box::new(move |x: i128| (x + x.signum() * half) / div) |
| } |
| } |
| |
| #[cfg(test)] |
| mod test { |
| use std::sync::Arc; |
| |
| use crate::spark_round; |
| |
| use arrow::array::Int64Array; |
| use arrow::datatypes::DataType; |
| use datafusion::common::cast::as_int64_array; |
| use datafusion::common::{Result, ScalarValue}; |
| use datafusion::physical_plan::ColumnarValue; |
| |
| // Regression tests for https://github.com/apache/datafusion-comet/issues/5070: |
| // round(Int64, scale) where `10^(-scale)` does not fit in i64. For scale=-19, |
| // values with |x| >= 5e18 round to sign(x)*1e19, which does not fit in a long: |
| // Spark throws under ANSI and wraps (low-order 64 bits) under legacy. |
| |
| // 1e19 truncated to the low 64 bits, matching `BigDecimal.longValue`. Note |
| // this value is *negative* (-8446744073709551616): 1e19 exceeds i64::MAX, so |
| // reinterpreting its low 64 bits as a signed long flips the sign. Rounding |
| // -5e18 down to -1e19 therefore wraps to a positive value. |
| const WRAPPED_1E19: i64 = 10_000_000_000_000_000_000u64 as i64; |
| const WRAPPED_MINUS_1E19: i64 = -WRAPPED_1E19; |
| |
| fn assert_round_int64_ansi_overflows(value: ColumnarValue) { |
| let args = vec![value, ColumnarValue::Scalar(ScalarValue::Int64(Some(-19)))]; |
| let err = spark_round(&args, &DataType::Int64, true).unwrap_err(); |
| assert!( |
| err.to_string().to_ascii_lowercase().contains("overflow"), |
| "expected arithmetic overflow error, got: {err}" |
| ); |
| } |
| |
| #[test] |
| fn test_round_int64_negative_scale_overflow_ansi() { |
| // ±5e18 rounds away from zero to ±1e19, which overflows i64. |
| for value in [5_000_000_000_000_000_000i64, -5_000_000_000_000_000_000i64] { |
| assert_round_int64_ansi_overflows(ColumnarValue::Array(Arc::new(Int64Array::from( |
| vec![value], |
| )))); |
| } |
| } |
| |
| #[test] |
| fn test_round_int64_negative_scale_overflow_ansi_scalar() { |
| for value in [5_000_000_000_000_000_000i64, -5_000_000_000_000_000_000i64] { |
| assert_round_int64_ansi_overflows(ColumnarValue::Scalar(ScalarValue::Int64(Some( |
| value, |
| )))); |
| } |
| } |
| |
| #[test] |
| fn test_round_int64_negative_scale_overflow_legacy() -> Result<()> { |
| // Under legacy mode, ±1e19 wraps to its low-order 64 bits. |
| let args = vec![ |
| ColumnarValue::Array(Arc::new(Int64Array::from(vec![ |
| 5_000_000_000_000_000_000i64, |
| -5_000_000_000_000_000_000i64, |
| 4_999_999_999_999_999_999i64, |
| 0i64, |
| i64::MAX, |
| i64::MIN, |
| ]))), |
| ColumnarValue::Scalar(ScalarValue::Int64(Some(-19))), |
| ]; |
| let ColumnarValue::Array(result) = spark_round(&args, &DataType::Int64, false)? else { |
| unreachable!() |
| }; |
| let longs = as_int64_array(&result)?; |
| let expected = Int64Array::from(vec![ |
| WRAPPED_1E19, |
| WRAPPED_MINUS_1E19, |
| 0i64, |
| 0i64, |
| WRAPPED_1E19, |
| WRAPPED_MINUS_1E19, |
| ]); |
| assert_eq!(longs, &expected); |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_round_int64_negative_scale_below_threshold() -> Result<()> { |
| // scale=-20: threshold is 5e19, which exceeds i64::MAX, so every long |
| // rounds to 0 under both ANSI and legacy. |
| let arr = Int64Array::from(vec![i64::MAX, i64::MIN, 0, 1_000_000_000_000_000_000]); |
| for fail_on_error in [false, true] { |
| let args = vec![ |
| ColumnarValue::Array(Arc::new(arr.clone())), |
| ColumnarValue::Scalar(ScalarValue::Int64(Some(-20))), |
| ]; |
| let ColumnarValue::Array(result) = spark_round(&args, &DataType::Int64, fail_on_error)? |
| else { |
| unreachable!() |
| }; |
| let longs = as_int64_array(&result)?; |
| assert_eq!(longs, &Int64Array::from(vec![0i64; arr.len()])); |
| } |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_round_int64_scale_below_i128_range() -> Result<()> { |
| // scale=-40: 10^40 does not fit in i128 either, so `checked_pow` returns |
| // None and the fallback returns 0 for every long, under both modes. |
| let arr = Int64Array::from(vec![i64::MAX, i64::MIN, 0, 5_000_000_000_000_000_000]); |
| for fail_on_error in [false, true] { |
| let args = vec![ |
| ColumnarValue::Array(Arc::new(arr.clone())), |
| ColumnarValue::Scalar(ScalarValue::Int64(Some(-40))), |
| ]; |
| let ColumnarValue::Array(result) = spark_round(&args, &DataType::Int64, fail_on_error)? |
| else { |
| unreachable!() |
| }; |
| let longs = as_int64_array(&result)?; |
| assert_eq!(longs, &Int64Array::from(vec![0i64; arr.len()])); |
| |
| let scalar_args = vec![ |
| ColumnarValue::Scalar(ScalarValue::Int64(Some(5_000_000_000_000_000_000i64))), |
| ColumnarValue::Scalar(ScalarValue::Int64(Some(-40))), |
| ]; |
| let ColumnarValue::Scalar(ScalarValue::Int64(result)) = |
| spark_round(&scalar_args, &DataType::Int64, fail_on_error)? |
| else { |
| unreachable!() |
| }; |
| assert_eq!(result, Some(0)); |
| } |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_round_int64_negative_scale_null_scalar() -> Result<()> { |
| // A null long stays null in every scale band: 10^(-scale) fits in i64 |
| // (-9), only in i128 (-19), or in neither (-40). |
| for point in [-9i64, -19, -40] { |
| for fail_on_error in [false, true] { |
| let args = vec![ |
| ColumnarValue::Scalar(ScalarValue::Int64(None)), |
| ColumnarValue::Scalar(ScalarValue::Int64(Some(point))), |
| ]; |
| let ColumnarValue::Scalar(ScalarValue::Int64(result)) = |
| spark_round(&args, &DataType::Int64, fail_on_error)? |
| else { |
| unreachable!() |
| }; |
| assert_eq!(result, None, "scale={point}, ansi={fail_on_error}"); |
| } |
| } |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_round_int64_negative_scale_legacy_scalar() -> Result<()> { |
| let args = vec![ |
| ColumnarValue::Scalar(ScalarValue::Int64(Some(5_000_000_000_000_000_000i64))), |
| ColumnarValue::Scalar(ScalarValue::Int64(Some(-19))), |
| ]; |
| let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) = |
| spark_round(&args, &DataType::Int64, false)? |
| else { |
| unreachable!() |
| }; |
| assert_eq!(result, WRAPPED_1E19); |
| Ok(()) |
| } |
| } |
