Add Decimal128 support to Ceil and Floor

datafusion/execution/src/config.rs

@@ -114,10 +114,10 @@ impl Default for SessionConfig {
 }
 
 /// A type map for storing extensions.
-/// 
+///
 /// Extensions are indexed by their type `T`. If multiple values of the same type are provided, only the last one
 /// will be kept.
-/// 
+///
 /// Extensions are opaque objects that are unknown to DataFusion itself but can be downcast by optimizer rules,
 /// execution plans, or other components that have access to the session config.
 /// They provide a flexible way to attach extra data or behavior to the session config.

datafusion/functions/src/math/ceil.rs

@@ -0,0 +1,174 @@
+// 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 std::any::Any;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, AsArray};
+use arrow::datatypes::{
+    DataType, Decimal128Type, Decimal256Type, Decimal32Type, Decimal64Type, Float32Type,
+    Float64Type,
+};
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::interval_arithmetic::Interval;
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_expr::{
+    Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+    TypeSignature, TypeSignatureClass, Volatility,
+};
+use datafusion_macros::user_doc;
+
+use super::decimal::{apply_decimal_op, ceil_decimal_value};
+
+#[user_doc(
+    doc_section(label = "Math Functions"),
+    description = "Returns the nearest integer greater than or equal to a number.",
+    syntax_example = "ceil(numeric_expression)",
+    standard_argument(name = "numeric_expression", prefix = "Numeric"),
+    sql_example = r#"```sql
+> SELECT ceil(3.14);
++------------+
+| ceil(3.14) |
++------------+
+| 4.0        |
++------------+
+```"#
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct CeilFunc {
+    signature: Signature,
+}
+
+impl Default for CeilFunc {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl CeilFunc {
+    pub fn new() -> Self {
+        let decimal_sig = Coercion::new_exact(TypeSignatureClass::Decimal);
+        Self {
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::Coercible(vec![decimal_sig]),
+                    TypeSignature::Uniform(1, vec![DataType::Float64, DataType::Float32]),
+                ],
+                Volatility::Immutable,
+            ),
+        }
+    }
+}
+
+impl ScalarUDFImpl for CeilFunc {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "ceil"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        match &arg_types[0] {
+            DataType::Null => Ok(DataType::Float64),
+            other => Ok(other.clone()),
+        }
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let args = ColumnarValue::values_to_arrays(&args.args)?;
+        let value = &args[0];
+
+        let result: ArrayRef = match value.data_type() {
+            DataType::Float64 => Arc::new(
+                value
+                    .as_primitive::<Float64Type>()
+                    .unary::<_, Float64Type>(f64::ceil),
+            ),
+            DataType::Float32 => Arc::new(
+                value
+                    .as_primitive::<Float32Type>()
+                    .unary::<_, Float32Type>(f32::ceil),
+            ),
+            DataType::Null => {
+                return Ok(ColumnarValue::Scalar(ScalarValue::Float64(None)))
+            }
+            DataType::Decimal32(precision, scale) => {
+                apply_decimal_op::<Decimal32Type, _>(
+                    value,
+                    *precision,
+                    *scale,
+                    self.name(),
+                    ceil_decimal_value,
+                )?
+            }
+            DataType::Decimal64(precision, scale) => {
+                apply_decimal_op::<Decimal64Type, _>(
+                    value,
+                    *precision,
+                    *scale,
+                    self.name(),
+                    ceil_decimal_value,
+                )?
+            }
+            DataType::Decimal128(precision, scale) => {
+                apply_decimal_op::<Decimal128Type, _>(
+                    value,
+                    *precision,
+                    *scale,
+                    self.name(),
+                    ceil_decimal_value,
+                )?
+            }
+            DataType::Decimal256(precision, scale) => {
+                apply_decimal_op::<Decimal256Type, _>(
+                    value,
+                    *precision,
+                    *scale,
+                    self.name(),
+                    ceil_decimal_value,
+                )?
+            }
+            other => {
+                return exec_err!(
+                    "Unsupported data type {other:?} for function {}",
+                    self.name()
+                )
+            }
+        };
+
+        Ok(ColumnarValue::Array(result))
+    }
+
+    fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
+        Ok(input[0].sort_properties)
+    }
+
+    fn evaluate_bounds(&self, inputs: &[&Interval]) -> Result<Interval> {
+        let data_type = inputs[0].data_type();
+        Interval::make_unbounded(&data_type)
+    }
+
+    fn documentation(&self) -> Option<&Documentation> {
+        self.doc()
+    }
+}

datafusion/functions/src/math/decimal.rs

@@ -0,0 +1,107 @@
+// 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 std::sync::Arc;
+
+use arrow::array::{ArrayRef, AsArray, PrimitiveArray};
+use arrow::datatypes::{ArrowNativeTypeOp, DecimalType};
+use arrow::error::ArrowError;
+use arrow_buffer::ArrowNativeType;
+use datafusion_common::{DataFusionError, Result};
+
+pub(super) fn apply_decimal_op<T, F>(
+    array: &ArrayRef,
+    precision: u8,
+    scale: i8,
+    fn_name: &str,
+    op: F,
+) -> Result<ArrayRef>
+where
+    T: DecimalType,
+    T::Native: ArrowNativeType + ArrowNativeTypeOp,
+    F: Fn(T::Native, T::Native) -> T::Native,
+{
+    if scale <= 0 {
+        return Ok(Arc::clone(array));
+    }
+
+    let factor = decimal_scale_factor::<T>(scale, fn_name)?;
+    let decimal = array.as_primitive::<T>();
+    let data_type = array.data_type().clone();
+
+    let result: PrimitiveArray<T> = decimal.try_unary(|value| {
+        let new_value = op(value, factor);
+        T::validate_decimal_precision(new_value, precision, scale).map_err(|_| {
+            ArrowError::ComputeError(format!("Decimal overflow while applying {fn_name}"))
+        })?;
+        Ok::<_, ArrowError>(new_value)
+    })?;
+
+    let result = result.with_data_type(data_type);
+
+    Ok(Arc::new(result))
+}
+
+fn decimal_scale_factor<T>(scale: i8, fn_name: &str) -> Result<T::Native>
+where
+    T: DecimalType,
+    T::Native: ArrowNativeType + ArrowNativeTypeOp,
+{
+    let base = <T::Native as ArrowNativeType>::from_usize(10).ok_or_else(|| {
+        DataFusionError::Execution(format!("Decimal overflow while applying {fn_name}"))
+    })?;
+
+    base.pow_checked(scale as u32).map_err(|_| {
+        DataFusionError::Execution(format!("Decimal overflow while applying {fn_name}"))
+    })
+}
+
+pub(super) fn ceil_decimal_value<T>(value: T, factor: T) -> T
+where
+    T: ArrowNativeTypeOp + std::ops::Rem<Output = T>,
+{
+    let remainder = value % factor;
+
+    if remainder == T::ZERO {
+        return value;
+    }
+
+    if value >= T::ZERO {
+        let increment = factor.sub_wrapping(remainder);
+        value.add_wrapping(increment)
+    } else {
+        value.sub_wrapping(remainder)
+    }
+}
+
+pub(super) fn floor_decimal_value<T>(value: T, factor: T) -> T
+where
+    T: ArrowNativeTypeOp + std::ops::Rem<Output = T>,
+{
+    let remainder = value % factor;
+
+    if remainder == T::ZERO {
+        return value;
+    }
+
+    if value >= T::ZERO {
+        value.sub_wrapping(remainder)
+    } else {
+        let adjustment = factor.add_wrapping(remainder);
+        value.sub_wrapping(adjustment)
+    }
+}