datafusion_comet_spark_expr/math_funcs/
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 crate::downcast_compute_op;
19use crate::math_funcs::utils::{get_precision_scale, make_decimal_array, make_decimal_scalar};
20use arrow::array::{Float32Array, Float64Array, Int64Array};
21use arrow_array::{Array, ArrowNativeTypeOp};
22use arrow_schema::DataType;
23use datafusion::physical_plan::ColumnarValue;
24use datafusion_common::{DataFusionError, ScalarValue};
25use num::integer::div_ceil;
26use std::sync::Arc;
27
28/// `ceil` function that simulates Spark `ceil` expression
29pub fn spark_ceil(
30    args: &[ColumnarValue],
31    data_type: &DataType,
32) -> Result<ColumnarValue, DataFusionError> {
33    let value = &args[0];
34    match value {
35        ColumnarValue::Array(array) => match array.data_type() {
36            DataType::Float32 => {
37                let result = downcast_compute_op!(array, "ceil", ceil, Float32Array, Int64Array);
38                Ok(ColumnarValue::Array(result?))
39            }
40            DataType::Float64 => {
41                let result = downcast_compute_op!(array, "ceil", ceil, Float64Array, Int64Array);
42                Ok(ColumnarValue::Array(result?))
43            }
44            DataType::Int64 => {
45                let result = array.as_any().downcast_ref::<Int64Array>().unwrap();
46                Ok(ColumnarValue::Array(Arc::new(result.clone())))
47            }
48            DataType::Decimal128(_, scale) if *scale > 0 => {
49                let f = decimal_ceil_f(scale);
50                let (precision, scale) = get_precision_scale(data_type);
51                make_decimal_array(array, precision, scale, &f)
52            }
53            other => Err(DataFusionError::Internal(format!(
54                "Unsupported data type {:?} for function ceil",
55                other,
56            ))),
57        },
58        ColumnarValue::Scalar(a) => match a {
59            ScalarValue::Float32(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(
60                a.map(|x| x.ceil() as i64),
61            ))),
62            ScalarValue::Float64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(
63                a.map(|x| x.ceil() as i64),
64            ))),
65            ScalarValue::Int64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(a.map(|x| x)))),
66            ScalarValue::Decimal128(a, _, scale) if *scale > 0 => {
67                let f = decimal_ceil_f(scale);
68                let (precision, scale) = get_precision_scale(data_type);
69                make_decimal_scalar(a, precision, scale, &f)
70            }
71            _ => Err(DataFusionError::Internal(format!(
72                "Unsupported data type {:?} for function ceil",
73                value.data_type(),
74            ))),
75        },
76    }
77}
78
79#[inline]
80fn decimal_ceil_f(scale: &i8) -> impl Fn(i128) -> i128 {
81    let div = 10_i128.pow_wrapping(*scale as u32);
82    move |x: i128| div_ceil(x, div)
83}
datafusion_comet_spark_expr/math_funcs/ceil.rs

datafusion_comet_spark_expr/math_funcs/
ceil.rs
