datafusion_spark/function/math/
floor.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow::array::cast::AsArray;
19use arrow::array::types::Decimal128Type;
20use arrow::array::{ArrowNativeTypeOp, Decimal128Array, Int64Array};
21use arrow::compute::kernels::arity::unary;
22use arrow::datatypes::{DataType, Field, FieldRef};
23use datafusion_common::{DataFusionError, ScalarValue, exec_err, internal_err};
24use datafusion_expr::{
25    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature,
26    Volatility,
27};
28use std::sync::Arc;
29
30/// Spark-compatible `floor` function.
31///
32/// Differences from DataFusion's floor:
33/// - Returns Int64 for float and integer inputs (while DataFusion preserves input type)
34/// - For Decimal128(p, s), returns Decimal128(p-s+1, 0) with scale 0
35///   (DataFusion preserves original precision and scale)
36///
37/// <https://spark.apache.org/docs/latest/api/sql/index.html#floor>
38#[derive(Debug, PartialEq, Eq, Hash)]
39pub struct SparkFloor {
40    signature: Signature,
41}
42
43impl Default for SparkFloor {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl SparkFloor {
50    pub fn new() -> Self {
51        Self {
52            signature: Signature::numeric(1, Volatility::Immutable),
53        }
54    }
55}
56
57impl ScalarUDFImpl for SparkFloor {
58    fn name(&self) -> &str {
59        "floor"
60    }
61
62    fn signature(&self) -> &Signature {
63        &self.signature
64    }
65
66    fn return_type(
67        &self,
68        _arg_types: &[DataType],
69    ) -> datafusion_common::Result<DataType> {
70        internal_err!("return_field_from_args should be called instead")
71    }
72
73    fn return_field_from_args(
74        &self,
75        args: ReturnFieldArgs,
76    ) -> datafusion_common::Result<FieldRef> {
77        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
78        let return_type = match args.arg_fields[0].data_type() {
79            DataType::Decimal128(p, s) if *s > 0 => {
80                let new_p = (*p - *s as u8 + 1).clamp(1, 38);
81                DataType::Decimal128(new_p, 0)
82            }
83            DataType::Decimal128(p, s) => DataType::Decimal128(*p, *s),
84            DataType::Float32
85            | DataType::Float64
86            | DataType::Int8
87            | DataType::Int16
88            | DataType::Int32
89            | DataType::Int64 => DataType::Int64,
90            _ => exec_err!(
91                "found unsupported return type {:?}",
92                args.arg_fields[0].data_type()
93            )?,
94        };
95        Ok(Arc::new(Field::new(self.name(), return_type, nullable)))
96    }
97
98    fn invoke_with_args(
99        &self,
100        args: ScalarFunctionArgs,
101    ) -> datafusion_common::Result<ColumnarValue> {
102        spark_floor(&args.args, args.return_field.data_type())
103    }
104}
105
106macro_rules! apply_int64 {
107    ($value:expr, $arr_type:ty, $scalar_variant:path, $f:expr) => {
108        match $value {
109            ColumnarValue::Array(array) => {
110                let result: Int64Array = unary(array.as_primitive::<$arr_type>(), $f);
111                Ok(ColumnarValue::Array(Arc::new(result)))
112            }
113            ColumnarValue::Scalar($scalar_variant(v)) => {
114                Ok(ColumnarValue::Scalar(ScalarValue::Int64(v.map($f))))
115            }
116            other => internal_err!(
117                "floor: data type mismatch — expected scalar of type {} but got {:?}",
118                stringify!($scalar_variant),
119                other.data_type()
120            ),
121        }
122    };
123}
124
125fn spark_floor(
126    args: &[ColumnarValue],
127    return_type: &DataType,
128) -> Result<ColumnarValue, DataFusionError> {
129    let value = &args[0];
130    match value.data_type() {
131        DataType::Float32 => apply_int64!(
132            value,
133            arrow::datatypes::Float32Type,
134            ScalarValue::Float32,
135            |x| x.floor() as i64
136        ),
137        DataType::Float64 => apply_int64!(
138            value,
139            arrow::datatypes::Float64Type,
140            ScalarValue::Float64,
141            |x| x.floor() as i64
142        ),
143        DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
144            value.cast_to(&DataType::Int64, None)
145        }
146        DataType::Decimal128(_, scale) if scale > 0 => {
147            let divisor = 10_i128.pow_wrapping(scale as u32);
148            let floor_decimal = |x: i128| {
149                let (d, r) = (x / divisor, x % divisor);
150                if r < 0 { d - 1 } else { d }
151            };
152            match value {
153                ColumnarValue::Array(array) => {
154                    let result: Decimal128Array =
155                        unary(array.as_primitive::<Decimal128Type>(), floor_decimal);
156                    Ok(ColumnarValue::Array(Arc::new(
157                        result.with_data_type(return_type.clone()),
158                    )))
159                }
160                ColumnarValue::Scalar(ScalarValue::Decimal128(v, _, _)) => {
161                    let DataType::Decimal128(new_p, new_s) = return_type else {
162                        return internal_err!(
163                            "floor: data type mismatch — expected Decimal128 return type but got {:?}",
164                            return_type
165                        );
166                    };
167                    Ok(ColumnarValue::Scalar(ScalarValue::Decimal128(
168                        v.map(floor_decimal),
169                        *new_p,
170                        *new_s,
171                    )))
172                }
173                other => internal_err!(
174                    "floor: data type mismatch — expected Decimal128 scalar but got {:?}",
175                    other.data_type()
176                ),
177            }
178        }
179        DataType::Decimal128(_, _) => Ok(value.clone()),
180        other => exec_err!("Unsupported data type {other:?} for function floor"),
181    }
182}
datafusion_spark/function/math/floor.rs

datafusion_spark/function/math/
floor.rs
