add_tests

Author: B Vadlamani

datafusion/spark/src/function/math/ceil.rs

@@ -22,10 +22,11 @@ use arrow::array::{
     Int32Array, Int64Array,
 };
 use arrow::compute::kernels::arity::unary;
-use arrow::datatypes::DataType;
-use datafusion_common::{DataFusionError, ScalarValue};
+use arrow::datatypes::{DataType, Field, FieldRef};
+use datafusion_common::{DataFusionError, ScalarValue, internal_err};
 use datafusion_expr::{
-    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+    Volatility,
 };
 use std::any::Any;
 use std::sync::Arc;
@@ -62,17 +63,32 @@ pub fn spark_ceil(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionErr
                 Ok(ColumnarValue::Array(result?))
             }
             DataType::Int8 => {
-                let input = array.as_any().downcast_ref::<Int8Array>().unwrap();
+                let input =
+                    array.as_any().downcast_ref::<Int8Array>().ok_or_else(|| {
+                        DataFusionError::Internal(
+                            "Failed to downcast to Int8Array".to_string(),
+                        )
+                    })?;
                 let result: Int64Array = unary(input, |x| x as i64);
                 Ok(ColumnarValue::Array(Arc::new(result)))
             }
             DataType::Int16 => {
-                let input = array.as_any().downcast_ref::<Int16Array>().unwrap();
+                let input =
+                    array.as_any().downcast_ref::<Int16Array>().ok_or_else(|| {
+                        DataFusionError::Internal(
+                            "Failed to downcast to Int16Array".to_string(),
+                        )
+                    })?;
                 let result: Int64Array = unary(input, |x| x as i64);
                 Ok(ColumnarValue::Array(Arc::new(result)))
             }
             DataType::Int32 => {
-                let input = array.as_any().downcast_ref::<Int32Array>().unwrap();
+                let input =
+                    array.as_any().downcast_ref::<Int32Array>().ok_or_else(|| {
+                        DataFusionError::Internal(
+                            "Failed to downcast to Int32Array".to_string(),
+                        )
+                    })?;
                 let result: Int64Array = unary(input, |x| x as i64);
                 Ok(ColumnarValue::Array(Arc::new(result)))
             }
@@ -147,26 +163,43 @@ fn decimal_ceil_f(scale: i8) -> impl Fn(i128) -> i128 {
     }
 }
 
+/// Spark-compatible `ceil` function.
+///
+/// Returns the smallest integer greater than or equal to the input.
+/// Unlike standard DataFusion ceil, this returns Int64 for float/integer
+/// inputs (matching Spark behavior).
+///
+/// # Supported types
+/// - Float32, Float64 → Int64
+/// - Int8, Int16, Int32, Int64 → Int64
+/// - Decimal128(p, s) → Decimal128(p, s) (preserves scale)
+///
+/// # Examples
+/// ```sql
+/// SELECT ceil(1.5);    -- Returns 2 (Int64)
+/// SELECT ceil(-1.5);   -- Returns -1 (Int64)
+/// SELECT ceil(5);      -- Returns 5 (Int64)
+/// ```
 #[derive(Debug, PartialEq, Eq, Hash)]
-pub struct SparkCiel {
+pub struct SparkCeil {
     signature: Signature,
 }
 
-impl Default for SparkCiel {
+impl Default for SparkCeil {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl SparkCiel {
+impl SparkCeil {
     pub fn new() -> Self {
         Self {
             signature: Signature::numeric(1, Volatility::Immutable),
         }
     }
 }
 
-impl ScalarUDFImpl for SparkCiel {
+impl ScalarUDFImpl for SparkCeil {
     fn as_any(&self) -> &dyn Any {
         self
     }
@@ -183,7 +216,19 @@ impl ScalarUDFImpl for SparkCiel {
         &self,
         _arg_types: &[DataType],
     ) -> datafusion_common::Result<DataType> {
-        Ok(DataType::Int64)
+        internal_err!("return_field_from_args should be called instead")
+    }
+
+    fn return_field_from_args(
+        &self,
+        args: ReturnFieldArgs,
+    ) -> datafusion_common::Result<FieldRef> {
+        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
+        let return_type = match args.arg_fields[0].data_type() {
+            DataType::Decimal128(p, s) => DataType::Decimal128(*p, *s),
+            _ => DataType::Int64,
+        };
+        Ok(Arc::new(Field::new(self.name(), return_type, nullable)))
     }
 
     fn invoke_with_args(

datafusion/spark/src/function/math/mod.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 pub mod abs;
-mod ceil;
+pub mod ceil;
 pub mod expm1;
 pub mod factorial;
 pub mod hex;
@@ -43,7 +43,7 @@ make_udf_function!(width_bucket::SparkWidthBucket, width_bucket);
 make_udf_function!(trigonometry::SparkCsc, csc);
 make_udf_function!(trigonometry::SparkSec, sec);
 make_udf_function!(negative::SparkNegative, negative);
-make_udf_function!(ceil::SparkCiel, ceil);
+make_udf_function!(ceil::SparkCeil, ceil);
 
 pub mod expr_fn {
     use datafusion_functions::export_functions;

datafusion/sqllogictest/test_files/spark/math/ceil.slt

@@ -144,7 +144,7 @@ SELECT ceil(127::tinyint);
 127
 
 query I
-SELECT ceil(-128::tinyint);
+SELECT ceil(CAST(-128 AS tinyint));
 ----
 -128
 
@@ -170,7 +170,7 @@ SELECT ceil(32767::smallint);
 32767
 
 query I
-SELECT ceil(-32768::smallint);
+SELECT ceil(CAST(-32768 AS smallint));
 ----
 -32768
 
@@ -212,7 +212,7 @@ SELECT ceil(9223372036854775807::bigint);
 9223372036854775807
 
 query I
-SELECT ceil(-9223372036854775808::bigint);
+SELECT ceil(CAST(-9223372036854775808 AS bigint));
 ----
 -9223372036854775808