datafusion_comet_spark_expr/math_funcs/
utils.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::{ArrayRef, Decimal128Array};
21use arrow::datatypes::DataType;
22use datafusion::common::{DataFusionError, ScalarValue};
23use datafusion::physical_plan::ColumnarValue;
24use std::sync::Arc;
25
26#[macro_export]
27macro_rules! downcast_compute_op {
28    ($ARRAY:expr, $NAME:expr, $FUNC:ident, $TYPE:ident, $RESULT:ident) => {{
29        let n = $ARRAY.as_any().downcast_ref::<$TYPE>();
30        match n {
31            Some(array) => {
32                let res: $RESULT =
33                    arrow::compute::kernels::arity::unary(array, |x| x.$FUNC() as i64);
34                Ok(Arc::new(res))
35            }
36            _ => Err(DataFusionError::Internal(format!(
37                "Invalid data type for {}",
38                $NAME
39            ))),
40        }
41    }};
42}
43
44#[inline]
45pub(crate) fn make_decimal_scalar(
46    a: &Option<i128>,
47    precision: u8,
48    scale: i8,
49    f: &dyn Fn(i128) -> i128,
50) -> Result<ColumnarValue, DataFusionError> {
51    let result = ScalarValue::Decimal128(a.map(f), precision, scale);
52    Ok(ColumnarValue::Scalar(result))
53}
54
55#[inline]
56pub(crate) fn make_decimal_array(
57    array: &ArrayRef,
58    precision: u8,
59    scale: i8,
60    f: &dyn Fn(i128) -> i128,
61) -> Result<ColumnarValue, DataFusionError> {
62    let array = array.as_primitive::<Decimal128Type>();
63    let result: Decimal128Array = arrow::compute::kernels::arity::unary(array, f);
64    let result = result.with_data_type(DataType::Decimal128(precision, scale));
65    Ok(ColumnarValue::Array(Arc::new(result)))
66}
67
68#[inline]
69pub(crate) fn get_precision_scale(data_type: &DataType) -> (u8, i8) {
70    let DataType::Decimal128(precision, scale) = data_type else {
71        unreachable!()
72    };
73    (*precision, *scale)
74}
datafusion_comet_spark_expr/math_funcs/utils.rs

datafusion_comet_spark_expr/math_funcs/
utils.rs
