Skip to main content

datafusion_functions/math/
ceil.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 std::sync::Arc;
19
20use arrow::array::{ArrayRef, AsArray};
21use arrow::datatypes::{
22    DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Float32Type,
23    Float64Type,
24};
25use datafusion_common::{Result, ScalarValue, exec_err};
26use datafusion_expr::interval_arithmetic::Interval;
27use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
28use datafusion_expr::{
29    Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
30    TypeSignature, TypeSignatureClass, Volatility,
31};
32use datafusion_macros::user_doc;
33
34use super::decimal::{apply_decimal_op, ceil_decimal_value};
35
36#[user_doc(
37    doc_section(label = "Math Functions"),
38    description = "Returns the nearest integer greater than or equal to a number.",
39    syntax_example = "ceil(numeric_expression)",
40    standard_argument(name = "numeric_expression", prefix = "Numeric"),
41    sql_example = r#"```sql
42> SELECT ceil(3.14);
43+------------+
44| ceil(3.14) |
45+------------+
46| 4.0        |
47+------------+
48```"#
49)]
50#[derive(Debug, PartialEq, Eq, Hash)]
51pub struct CeilFunc {
52    signature: Signature,
53}
54
55impl Default for CeilFunc {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61impl CeilFunc {
62    pub fn new() -> Self {
63        let decimal_sig = Coercion::new_exact(TypeSignatureClass::Decimal);
64        Self {
65            signature: Signature::one_of(
66                vec![
67                    TypeSignature::Coercible(vec![decimal_sig]),
68                    TypeSignature::Uniform(1, vec![DataType::Float64, DataType::Float32]),
69                ],
70                Volatility::Immutable,
71            ),
72        }
73    }
74}
75
76impl ScalarUDFImpl for CeilFunc {
77    fn name(&self) -> &str {
78        "ceil"
79    }
80
81    fn signature(&self) -> &Signature {
82        &self.signature
83    }
84
85    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
86        match &arg_types[0] {
87            DataType::Null => Ok(DataType::Float64),
88            other => Ok(other.clone()),
89        }
90    }
91
92    fn is_strict(&self) -> bool {
93        true
94    }
95
96    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
97        let arg = &args.args[0];
98
99        // Scalar fast path for float types - avoid array conversion overhead entirely
100        if let ColumnarValue::Scalar(scalar) = arg {
101            match scalar {
102                ScalarValue::Float64(v) => {
103                    return Ok(ColumnarValue::Scalar(ScalarValue::Float64(
104                        v.map(f64::ceil),
105                    )));
106                }
107                ScalarValue::Float32(v) => {
108                    return Ok(ColumnarValue::Scalar(ScalarValue::Float32(
109                        v.map(f32::ceil),
110                    )));
111                }
112                ScalarValue::Null => {
113                    return Ok(ColumnarValue::Scalar(ScalarValue::Float64(None)));
114                }
115                // For decimals: convert to array of size 1, process, then extract scalar
116                // This ensures we don't expand the array while reusing overflow validation
117                _ => {}
118            }
119        }
120
121        // Track if input was a scalar to convert back at the end
122        let is_scalar = matches!(arg, ColumnarValue::Scalar(_));
123
124        // Array path (also handles decimal scalars converted to size-1 arrays)
125        let value = arg.to_array(args.number_rows)?;
126
127        let result: ArrayRef = match value.data_type() {
128            DataType::Float64 => Arc::new(
129                value
130                    .as_primitive::<Float64Type>()
131                    .unary::<_, Float64Type>(f64::ceil),
132            ),
133            DataType::Float32 => Arc::new(
134                value
135                    .as_primitive::<Float32Type>()
136                    .unary::<_, Float32Type>(f32::ceil),
137            ),
138            DataType::Null => {
139                return Ok(ColumnarValue::Scalar(ScalarValue::Float64(None)));
140            }
141            DataType::Decimal32(precision, scale) => {
142                apply_decimal_op::<Decimal32Type, _>(
143                    &value,
144                    *precision,
145                    *scale,
146                    self.name(),
147                    ceil_decimal_value,
148                )?
149            }
150            DataType::Decimal64(precision, scale) => {
151                apply_decimal_op::<Decimal64Type, _>(
152                    &value,
153                    *precision,
154                    *scale,
155                    self.name(),
156                    ceil_decimal_value,
157                )?
158            }
159            DataType::Decimal128(precision, scale) => {
160                apply_decimal_op::<Decimal128Type, _>(
161                    &value,
162                    *precision,
163                    *scale,
164                    self.name(),
165                    ceil_decimal_value,
166                )?
167            }
168            DataType::Decimal256(precision, scale) => {
169                apply_decimal_op::<Decimal256Type, _>(
170                    &value,
171                    *precision,
172                    *scale,
173                    self.name(),
174                    ceil_decimal_value,
175                )?
176            }
177            other => {
178                return exec_err!(
179                    "Unsupported data type {other:?} for function {}",
180                    self.name()
181                );
182            }
183        };
184
185        // If input was a scalar, convert result back to scalar
186        if is_scalar {
187            ScalarValue::try_from_array(&result, 0).map(ColumnarValue::Scalar)
188        } else {
189            Ok(ColumnarValue::Array(result))
190        }
191    }
192
193    fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
194        Ok(input[0].sort_properties)
195    }
196
197    fn evaluate_bounds(&self, inputs: &[&Interval]) -> Result<Interval> {
198        let data_type = inputs[0].data_type();
199        Interval::make_unbounded(&data_type)
200    }
201
202    fn documentation(&self) -> Option<&Documentation> {
203        self.doc()
204    }
205}