Skip to main content

datafusion_expr/
function.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
18//! Function module contains typing and signature for built-in and user defined functions.
19
20use crate::ColumnarValue;
21use crate::{Expr, PartitionEvaluator};
22use arrow::datatypes::DataType;
23use datafusion_common::Result;
24use std::sync::Arc;
25
26pub use datafusion_functions_aggregate_common::accumulator::{
27    AccumulatorArgs, AccumulatorFactoryFunction, StateFieldsArgs,
28};
29
30use crate::expr::{AggregateFunction, WindowFunction};
31use crate::simplify::SimplifyContext;
32pub use datafusion_functions_window_common::expr::ExpressionArgs;
33pub use datafusion_functions_window_common::field::WindowUDFFieldArgs;
34pub use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
35
36#[derive(Debug, Clone, Copy)]
37pub enum Hint {
38    /// Indicates the argument needs to be padded if it is scalar
39    Pad,
40    /// Indicates the argument can be converted to an array of length 1
41    AcceptsSingular,
42}
43
44/// Scalar function
45///
46/// The Fn param is the wrapped function but be aware that the function will
47/// be passed with the slice / vec of columnar values (either scalar or array)
48/// with the exception of zero param function, where a singular element vec
49/// will be passed. In that case the single element is a null array to indicate
50/// the batch's row count (so that the generative zero-argument function can know
51/// the result array size).
52pub type ScalarFunctionImplementation =
53    Arc<dyn Fn(&[ColumnarValue]) -> Result<ColumnarValue> + Send + Sync>;
54
55/// Factory that returns the functions's return type given the input argument types
56pub type ReturnTypeFunction =
57    Arc<dyn Fn(&[DataType]) -> Result<Arc<DataType>> + Send + Sync>;
58
59/// Factory that creates a PartitionEvaluator for the given window
60/// function
61pub type PartitionEvaluatorFactory =
62    Arc<dyn Fn() -> Result<Box<dyn PartitionEvaluator>> + Send + Sync>;
63
64/// Factory that returns the types used by an aggregator to serialize
65/// its state, given its return datatype.
66pub type StateTypeFunction =
67    Arc<dyn Fn(&DataType) -> Result<Arc<Vec<DataType>>> + Send + Sync>;
68
69/// Type alias for [crate::udaf::AggregateUDFImpl::simplify].
70///
71/// This closure is invoked with:
72/// * `aggregate_function`: [AggregateFunction] with already simplified arguments
73/// * `info`: [SimplifyContext]
74///
75/// It returns a simplified [Expr] or an error.
76pub type AggregateFunctionSimplification =
77    Box<dyn Fn(AggregateFunction, &SimplifyContext) -> Result<Expr>>;
78
79/// Type alias for [crate::udwf::WindowUDFImpl::simplify].
80///
81/// This closure is invoked with:
82/// * `window_function`: [WindowFunction] with already simplified arguments
83/// * `info`: [SimplifyContext]
84///
85/// It returns a simplified [Expr] or an error.
86pub type WindowFunctionSimplification =
87    Box<dyn Fn(WindowFunction, &SimplifyContext) -> Result<Expr>>;