diff --git a/datafusion/spark/src/function/math/ceil.rs b/datafusion/spark/src/function/math/ceil.rs index 6042a69cd6543..cddec2e13e671 100644 --- a/datafusion/spark/src/function/math/ceil.rs +++ b/datafusion/spark/src/function/math/ceil.rs @@ -5,11 +5,14 @@ use std::sync::Arc; use arrow::array::{ArrowNativeTypeOp, AsArray, Decimal128Array}; -use arrow::datatypes::{DataType, Decimal128Type, Float32Type, Float64Type, Int64Type}; -use datafusion_common::utils::take_function_args; -use datafusion_common::{Result, ScalarValue, exec_err}; +use arrow::datatypes::{ + DataType, Decimal128Type, Field, FieldRef, Float32Type, Float64Type, Int64Type, +}; +use datafusion_common::types::{NativeType, logical_int32}; +use datafusion_common::{Result, ScalarValue, exec_err, plan_err}; use datafusion_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, + Coercion, ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, + Signature, TypeSignature, TypeSignatureClass, Volatility, }; /// Spark-compatible `ceil` expression @@ -23,8 +26,9 @@ use datafusion_expr::{ /// - Spark only supports Decimal128; DataFusion also supports Decimal32/64/256 /// - Spark does not check for decimal overflow; DataFusion errors on overflow /// -/// 2-argument ceil(value, scale) is not yet implemented -/// +/// Two-arg form `ceil(decimal, scale)` matches Spark `RoundCeil`: decimal value, constant +/// Int32 scale, result type from Spark `RoundBase::dataType`. Extreme scales that would +/// wrap or exceed the result precision return an error (Spark: NUMERIC_VALUE_OUT_OF_RANGE). #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkCeil { signature: Signature, @@ -39,8 +43,24 @@ impl Default for SparkCeil { impl SparkCeil { pub fn new() -> Self { + let scale = Coercion::new_implicit( + TypeSignatureClass::Native(logical_int32()), + vec![TypeSignatureClass::Integer], + NativeType::Int32, + ); Self { - signature: Signature::numeric(1, Volatility::Immutable), + signature: Signature::one_of( + vec![ + // ceil(value) + TypeSignature::Numeric(1), + // ceil(decimal, scale) + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Decimal), + scale, + ]), + ], + Volatility::Immutable, + ), aliases: vec!["ceiling".to_string()], } } @@ -55,6 +75,43 @@ impl ScalarUDFImpl for SparkCeil { &self.signature } + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); + let value_type = args.arg_fields[0].data_type(); + + let data_type = if args.arg_fields.len() == 1 { + self.return_type(std::slice::from_ref(value_type))? + } else { + // Spark requires a foldable scale: a non-constant one is a plan-time error there + // (NON_FOLDABLE_INPUT) rather than something evaluated per row. + let scale = match args.scalar_arguments.get(1).copied().flatten() { + Some(ScalarValue::Int32(Some(scale))) => *scale, + Some(ScalarValue::Int32(None)) => { + return plan_err!("Function ceil requires a non-null scale argument"); + } + _ => { + return plan_err!( + "Function ceil requires a constant integer scale argument" + ); + } + }; + + match value_type { + DataType::Decimal128(p, s) => DataType::Decimal128( + ceil_scaled_precision(*p, *s, scale), + ceil_scaled_scale(*s, scale), + ), + other => { + return plan_err!( + "Function ceil does not support a scale argument for {other:?}" + ); + } + } + }; + + Ok(Arc::new(Field::new("ceil", data_type, nullable))) + } + fn return_type(&self, arg_types: &[DataType]) -> Result { match &arg_types[0] { DataType::Decimal128(p, s) => { @@ -85,15 +142,147 @@ impl ScalarUDFImpl for SparkCeil { } fn spark_ceil(args: &[ColumnarValue]) -> Result { - let [input] = take_function_args("ceil", args)?; + match args { + [input] => match input { + ColumnarValue::Scalar(value) => spark_ceil_scalar(value), + ColumnarValue::Array(input) => spark_ceil_array(input), + }, + [input, scale] => spark_ceil_with_scale(input, scale), + _ => exec_err!("Function ceil expects one or two arguments"), + } +} + +/// Result scale for `ceil(decimal(p, s), scale)`, per Spark's `RoundBase::dataType`. +#[inline] +fn ceil_scaled_scale(input_scale: i8, scale: i32) -> i8 { + if scale < 0 { + 0 + } else { + (input_scale as i32).min(scale) as i8 + } +} + +/// Result precision for `ceil(decimal(p, s), scale)`, per Spark's `RoundBase::dataType`. +/// +/// Rounding up can carry into a new integral digit (`ceil(9.9, 0)` is `10`), so the integral +/// part always gets one digit more than the input had. +#[inline] +fn ceil_scaled_precision(precision: u8, input_scale: i8, scale: i32) -> u8 { + let integral_least_num_digits = precision as i32 - input_scale as i32 + 1; + let new_precision = if scale < 0 { + integral_least_num_digits.max(-scale + 1) + } else { + integral_least_num_digits + (input_scale as i32).min(scale) + }; + new_precision.clamp(1, 38) as u8 +} + +/// True when `|value|` fits in `Decimal128(precision, _)` (unscaled absolute value < 10^p). +#[inline] +fn decimal128_fits_precision(value: i128, precision: u8) -> bool { + match 10_i128.checked_pow(precision as u32) { + Some(limit) => value > -limit && value < limit, + // precision == 38: 10^38 does not fit in i128, but every i128 fits Decimal128(38, _). + None => true, + } +} + +/// Round `value` up to `target_scale` decimal places, where `value` is held at `input_scale`. +/// +/// Errors on arithmetic overflow or when the result does not fit the Spark result precision, +/// matching Spark's `NUMERIC_VALUE_OUT_OF_RANGE` for extreme negative scales. +fn decimal128_ceil_to_scale( + value: i128, + input_scale: i8, + input_precision: u8, + target_scale: i32, +) -> Result { + let out_p = ceil_scaled_precision(input_precision, input_scale, target_scale); + + let result = if target_scale >= input_scale as i32 { + // Already at least as fine as the target; keep the unscaled digits. + // Spark still caps the *declared* result scale at min(s, scale). + value + } else { + let dropped = (input_scale as i32 - target_scale) as u32; + let rounded = decimal128_ceil_checked(value, dropped)?; + + if target_scale < 0 { + // Result is stored at scale 0; restore the trailing zeros for the places + // below the decimal point that were rounded away into the integer part. + let Some(factor) = 10_i128.checked_pow((-target_scale) as u32) else { + return exec_err!( + "Decimal overflow in ceil: 10^{} exceeds Decimal128 range", + -target_scale + ); + }; + rounded.checked_mul(factor).ok_or_else(|| { + (exec_err!("Decimal overflow in ceil while rescaling to scale 0") + as Result<()>) + .unwrap_err() + })? + } else { + rounded + } + }; + + if !decimal128_fits_precision(result, out_p) { + return exec_err!( + "Decimal overflow in ceil: result does not fit Decimal128({out_p}, {})", + ceil_scaled_scale(input_scale, target_scale) + ); + } + Ok(result) +} + +fn spark_ceil_with_scale( + input: &ColumnarValue, + scale: &ColumnarValue, +) -> Result { + let ColumnarValue::Scalar(scale) = scale else { + return exec_err!("Function ceil requires a constant integer scale argument"); + }; + let ScalarValue::Int32(Some(scale)) = scale else { + return exec_err!("Function ceil requires a non-null integer scale argument"); + }; + let scale = *scale; match input { - ColumnarValue::Scalar(value) => spark_ceil_scalar(value), - ColumnarValue::Array(input) => spark_ceil_array(input), + ColumnarValue::Scalar(ScalarValue::Decimal128(value, p, s)) => { + let out_p = ceil_scaled_precision(*p, *s, scale); + let out_s = ceil_scaled_scale(*s, scale); + let result = match value { + Some(x) => Some(decimal128_ceil_to_scale(*x, *s, *p, scale)?), + None => None, + }; + Ok(ColumnarValue::Scalar(ScalarValue::Decimal128( + result, out_p, out_s, + ))) + } + ColumnarValue::Array(array) => match array.data_type() { + DataType::Decimal128(p, s) => { + let (p, s) = (*p, *s); + let out_p = ceil_scaled_precision(p, s, scale); + let out_s = ceil_scaled_scale(s, scale); + let result: Decimal128Array = array + .as_primitive::() + .try_unary(|x| decimal128_ceil_to_scale(x, s, p, scale))?; + Ok(ColumnarValue::Array(Arc::new( + result.with_data_type(DataType::Decimal128(out_p, out_s)), + ))) + } + other => { + exec_err!("Function ceil does not support a scale argument for {other:?}") + } + }, + other => exec_err!( + "Function ceil does not support a scale argument for {:?}", + other.data_type() + ), } } -/// Compute ceil for a single decimal128 value with the given scale. +/// Compute ceil for a single decimal128 value with the given scale (1-arg path). #[inline] fn decimal128_ceil(value: i128, scale: u32) -> i128 { let div = 10_i128.pow_wrapping(scale); @@ -102,6 +291,25 @@ fn decimal128_ceil(value: i128, scale: u32) -> i128 { if r > 0 { d + 1 } else { d } } +/// Checked variant used by `ceil(decimal, scale)`. Errors instead of wrapping when +/// `10^scale` itself overflows i128. +fn decimal128_ceil_checked(value: i128, scale: u32) -> Result { + let Some(div) = 10_i128.checked_pow(scale) else { + return exec_err!( + "Decimal overflow in ceil: 10^{scale} exceeds Decimal128 range" + ); + }; + let d = value / div; + let r = value % div; + if r > 0 { + d.checked_add(1).ok_or_else(|| { + (exec_err!("Decimal overflow in ceil") as Result<()>).unwrap_err() + }) + } else { + Ok(d) + } +} + /// Compute the return precision for a decimal128 ceil result. #[inline] fn decimal128_ceil_precision(precision: u8, scale: i8) -> u8 { @@ -155,6 +363,64 @@ fn spark_ceil_array(input: &Arc) -> Result Decimal('3.142'), decimal(5,3) + let value = 31411_i128; // 3.1411 at scale 4 + assert_eq!(decimal128_ceil_to_scale(value, 4, 5, 3).unwrap(), 3142); + assert_eq!(ceil_scaled_precision(5, 4, 3), 5); + assert_eq!(ceil_scaled_scale(4, 3), 3); + + // ceil(3.1411, -3) -> Decimal('1000'), decimal(4,0) + assert_eq!(decimal128_ceil_to_scale(value, 4, 5, -3).unwrap(), 1000); + assert_eq!(ceil_scaled_precision(5, 4, -3), 4); + assert_eq!(ceil_scaled_scale(4, -3), 0); + } + + #[test] + fn test_scale_rounds_toward_positive_infinity() { + assert_eq!(decimal128_ceil_to_scale(-31411, 4, 5, 3).unwrap(), -3141); + assert_eq!(decimal128_ceil_to_scale(-31411, 4, 5, 0).unwrap(), -3); + } + + #[test] + fn test_scale_noop_when_request_ge_input_scale() { + assert_eq!(decimal128_ceil_to_scale(31411, 4, 5, 4).unwrap(), 31411); + assert_eq!(decimal128_ceil_to_scale(31411, 4, 5, 9).unwrap(), 31411); + assert_eq!(ceil_scaled_scale(4, 9), 4); + assert_eq!(ceil_scaled_precision(5, 4, 9), 6); + } + + #[test] + fn test_scale_exact_value_and_carry() { + assert_eq!(decimal128_ceil_to_scale(31400, 4, 5, 2).unwrap(), 314); + assert_eq!(decimal128_ceil_to_scale(99, 1, 2, 0).unwrap(), 10); + assert_eq!(ceil_scaled_precision(2, 1, 0), 2); + } + + #[test] + fn test_scaled_precision_bounds() { + assert_eq!(ceil_scaled_precision(38, 0, -40), 38); + assert_eq!(ceil_scaled_precision(38, 2, 30), 38); + assert_eq!(ceil_scaled_precision(1, 1, 0), 1); + } + + #[test] + fn test_scale_errors_on_extreme_negative_scale() { + // Spark raises NUMERIC_VALUE_OUT_OF_RANGE for these; do not wrap. + assert!(decimal128_ceil_to_scale(31411, 4, 5, -40).is_err()); + assert!(decimal128_ceil_to_scale(31411, 4, 5, -38).is_err()); + // ceil(10^38 - 1 at scale 0, -1) overflows Decimal128(38, 0) + let max38 = 10_i128.pow(38) - 1; + assert!(decimal128_ceil_to_scale(max38, 0, 38, -1).is_err()); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/sqllogictest/test_files/spark/math/ceil.slt b/datafusion/sqllogictest/test_files/spark/math/ceil.slt index cf8c72ae037d8..d109bec4250e7 100644 --- a/datafusion/sqllogictest/test_files/spark/math/ceil.slt +++ b/datafusion/sqllogictest/test_files/spark/math/ceil.slt @@ -28,15 +28,24 @@ SELECT ceil(-0.1::decimal(1,1)), arrow_typeof(ceil(-0.1::decimal(1,1))); ## Original Query: SELECT ceil(3.1411, -3); ## PySpark 3.5.5 Result: {'ceil(3.1411, -3)': Decimal('1000'), 'typeof(ceil(3.1411, -3))': 'decimal(4,0)', 'typeof(3.1411)': 'decimal(5,4)', 'typeof(-3)': 'int'} -## TODO: 2-argument ceil(value, scale) is not yet implemented -#query -#SELECT ceil(3.1411::decimal(5,4), -3::int); +query RT +SELECT ceil(3.1411::decimal(5,4), -3), arrow_typeof(ceil(3.1411::decimal(5,4), -3)); +---- +1000 Decimal128(4, 0) ## Original Query: SELECT ceil(3.1411, 3); ## PySpark 3.5.5 Result: {'ceil(3.1411, 3)': Decimal('3.142'), 'typeof(ceil(3.1411, 3))': 'decimal(5,3)', 'typeof(3.1411)': 'decimal(5,4)', 'typeof(3)': 'int'} -## TODO: 2-argument ceil(value, scale) is not yet implemented -#query -#SELECT ceil(3.1411::decimal(5,4), 3::int); +# +# Spelled `ceiling` rather than `ceil` on purpose. sqlparser recognises +# `CEIL(, )` as the SQL-standard scale form and hands it to the +# planner as `CeilFloorKind::Scale`, which datafusion/sql rejects with "CEIL with scale is +# not supported" before any function is resolved. A negative scale does not match that +# production, which is why the case above can use `ceil` directly. The `ceiling` alias is +# parsed as an ordinary function call in both cases. +query RT +SELECT ceiling(3.1411::decimal(5,4), 3), arrow_typeof(ceiling(3.1411::decimal(5,4), 3)); +---- +3.142 Decimal128(5, 3) ## Original Query: SELECT ceil(5); ## PySpark 3.5.5 Result: {'CEIL(5)': 5, 'typeof(CEIL(5))': 'bigint', 'typeof(5)': 'int'}