datafusion_functions/macros.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/// macro that exports a list of function names as:
19/// 1. individual functions in an `expr_fn` module
20/// 2. a single function that returns a list of all functions
21///
22/// Equivalent to
23/// ```text
24/// pub mod expr_fn {
25/// use super::*;
26/// /// Return encode(arg)
27/// pub fn encode(args: Vec<Expr>) -> Expr {
28/// super::encode().call(args)
29/// }
30/// ...
31/// /// Return a list of all functions in this package
32/// pub(crate) fn functions() -> Vec<Arc<ScalarUDF>> {
33/// vec![
34/// encode(),
35/// decode()
36/// ]
37/// }
38/// ```
39///
40/// Exported functions accept:
41/// - `Vec<Expr>` argument (single argument followed by a comma)
42/// - Variable number of `Expr` arguments (zero or more arguments, must be without commas)
43/// - Functions that require config (marked with `@config` prefix)
44///
45/// Note on configuration construction paths:
46/// - The convenience wrappers generated for `@config` functions call the inner
47/// constructor with `ConfigOptions::default()`. These wrappers are intended
48/// primarily for programmatic `Expr` construction and convenience usage.
49/// - When functions are registered in a session, DataFusion will call
50/// `with_updated_config()` to create a `ScalarUDF` instance using the session's
51/// actual `ConfigOptions`. This also happens when configuration changes at runtime
52/// (e.g., via `SET` statements). In short: the macro uses the default config for
53/// convenience constructors; the session config is applied when functions are
54/// registered or when configuration is updated.
55#[macro_export]
56macro_rules! export_functions {
57 ($(($FUNC:ident, $DOC:expr, $($arg:tt)*)),*) => {
58 $(
59 // switch to single-function cases below
60 $crate::export_functions!(single $FUNC, $DOC, $($arg)*);
61 )*
62 };
63
64 // function that requires config (marked with @config)
65 (single $FUNC:ident, $DOC:expr, @config) => {
66 #[doc = $DOC]
67 pub fn $FUNC() -> datafusion_expr::Expr {
68 use datafusion_common::config::ConfigOptions;
69 super::$FUNC(&ConfigOptions::default()).call(vec![])
70 }
71 };
72
73 // function that requires config and takes a vector argument
74 (single $FUNC:ident, $DOC:expr, @config $arg:ident,) => {
75 #[doc = $DOC]
76 pub fn $FUNC($arg: Vec<datafusion_expr::Expr>) -> datafusion_expr::Expr {
77 use datafusion_common::config::ConfigOptions;
78 super::$FUNC(&ConfigOptions::default()).call($arg)
79 }
80 };
81
82 // function that requires config and variadic arguments
83 (single $FUNC:ident, $DOC:expr, @config $($arg:ident)*) => {
84 #[doc = $DOC]
85 pub fn $FUNC($($arg: datafusion_expr::Expr),*) -> datafusion_expr::Expr {
86 use datafusion_common::config::ConfigOptions;
87 super::$FUNC(&ConfigOptions::default()).call(vec![$($arg),*])
88 }
89 };
90
91 // single vector argument (a single argument followed by a comma)
92 (single $FUNC:ident, $DOC:expr, $arg:ident,) => {
93 #[doc = $DOC]
94 pub fn $FUNC($arg: Vec<datafusion_expr::Expr>) -> datafusion_expr::Expr {
95 super::$FUNC().call($arg)
96 }
97 };
98
99 // variadic arguments (zero or more arguments, without commas)
100 (single $FUNC:ident, $DOC:expr, $($arg:ident)*) => {
101 #[doc = $DOC]
102 pub fn $FUNC($($arg: datafusion_expr::Expr),*) -> datafusion_expr::Expr {
103 super::$FUNC().call(vec![$($arg),*])
104 }
105 };
106}
107
108/// Creates a singleton `ScalarUDF` of the `$UDF` function and a function
109/// named `$NAME` which returns that singleton. Optionally use a custom constructor
110/// `$CTOR` which defaults to `$UDF::new()` if not specified.
111///
112/// This is used to ensure creating the list of `ScalarUDF` only happens once.
113#[macro_export]
114macro_rules! make_udf_function {
115 ($UDF:ty, $NAME:ident, $CTOR:expr) => {
116 #[doc = concat!("Return a [`ScalarUDF`](datafusion_expr::ScalarUDF) implementation of ", stringify!($NAME))]
117 pub fn $NAME() -> std::sync::Arc<datafusion_expr::ScalarUDF> {
118 // Singleton instance of the function
119 static INSTANCE: std::sync::LazyLock<
120 std::sync::Arc<datafusion_expr::ScalarUDF>,
121 > = std::sync::LazyLock::new(|| {
122 std::sync::Arc::new(datafusion_expr::ScalarUDF::new_from_impl(
123 ($CTOR)(),
124 ))
125 });
126 std::sync::Arc::clone(&INSTANCE)
127 }
128 };
129 ($UDF:ty, $NAME:ident) => {
130 make_udf_function!($UDF, $NAME, <$UDF>::new);
131 };
132}
133
134/// Creates a singleton `ScalarUDF` of the `$UDF` function and a function
135/// named `$NAME` which returns that singleton. The function takes a
136/// configuration argument of type `$CONFIG_TYPE` to create the UDF.
137#[macro_export]
138macro_rules! make_udf_function_with_config {
139 ($UDF:ty, $NAME:ident) => {
140 #[doc = concat!("Return a [`ScalarUDF`](datafusion_expr::ScalarUDF) implementation of ", stringify!($NAME))]
141 pub fn $NAME(config: &datafusion_common::config::ConfigOptions) -> std::sync::Arc<datafusion_expr::ScalarUDF> {
142 std::sync::Arc::new(datafusion_expr::ScalarUDF::new_from_impl(
143 <$UDF>::new_with_config(&config),
144 ))
145 }
146 };
147}
148
149/// Macro creates a sub module if the feature is not enabled
150///
151/// The rationale for providing stub functions is to help users to configure datafusion
152/// properly (so they get an error telling them why a function is not available)
153/// instead of getting a cryptic "no function found" message at runtime.
154macro_rules! make_stub_package {
155 ($name:ident, $feature:literal) => {
156 #[cfg(not(feature = $feature))]
157 #[doc = concat!("Disabled. Enable via feature flag `", $feature, "`")]
158 pub mod $name {
159 use datafusion_expr::ScalarUDF;
160 use log::debug;
161 use std::sync::Arc;
162
163 /// Returns an empty list of functions when the feature is not enabled
164 pub fn functions() -> Vec<Arc<ScalarUDF>> {
165 debug!("{} functions disabled", stringify!($name));
166 vec![]
167 }
168 }
169 };
170}
171
172/// Downcast a named argument to a specific array type, returning an internal error
173/// if the cast fails
174///
175/// $ARG: ArrayRef
176/// $NAME: name of the argument (for error messages)
177/// $ARRAY_TYPE: the type of array to cast the argument to
178#[macro_export]
179macro_rules! downcast_named_arg {
180 ($ARG:expr, $NAME:expr, $ARRAY_TYPE:ident) => {{
181 $ARG.as_any().downcast_ref::<$ARRAY_TYPE>().ok_or_else(|| {
182 datafusion_common::internal_datafusion_err!(
183 "could not cast {} to {}",
184 $NAME,
185 std::any::type_name::<$ARRAY_TYPE>()
186 )
187 })?
188 }};
189}
190
191/// Downcast an argument to a specific array type, returning an internal error
192/// if the cast fails
193///
194/// $ARG: ArrayRef
195/// $ARRAY_TYPE: the type of array to cast the argument to
196#[macro_export]
197macro_rules! downcast_arg {
198 ($ARG:expr, $ARRAY_TYPE:ident) => {{ $crate::downcast_named_arg!($ARG, "", $ARRAY_TYPE) }};
199}
200
201/// Macro to create a unary math UDF.
202///
203/// A unary math function takes an argument of type Float32 or Float64,
204/// applies a unary floating function to the argument, and returns a value of the same type.
205///
206/// $UDF: the name of the UDF struct that implements `ScalarUDFImpl`
207/// $NAME: the name of the function
208/// $UNARY_FUNC: the unary function to apply to the argument
209/// $OUTPUT_ORDERING: the output ordering calculation method of the function
210/// $STRICT: whether the function returns NULL when any argument is NULL
211/// $GET_DOC: the function to get the documentation of the UDF
212macro_rules! make_math_unary_udf {
213 ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $STRICT:expr, $GET_DOC:expr) => {
214 make_math_unary_udf!(
215 $UDF,
216 $NAME,
217 $UNARY_FUNC,
218 $OUTPUT_ORDERING,
219 $EVALUATE_BOUNDS,
220 $STRICT,
221 $GET_DOC,
222 None::<fn(f64) -> Result<()>>
223 );
224 };
225 ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $STRICT:expr, $GET_DOC:expr, $VALIDATOR:expr) => {
226 $crate::make_udf_function!($NAME::$UDF, $NAME);
227
228 mod $NAME {
229
230 use std::sync::Arc;
231
232 use arrow::array::{ArrayRef, AsArray};
233 use arrow::datatypes::{DataType, Float32Type, Float64Type};
234 use arrow::error::ArrowError;
235 use datafusion_common::{Result, exec_err};
236 use datafusion_expr::interval_arithmetic::Interval;
237 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
238 use datafusion_expr::{
239 ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl,
240 Signature, Volatility,
241 };
242
243 #[derive(Debug, PartialEq, Eq, Hash)]
244 pub struct $UDF {
245 signature: Signature,
246 }
247
248 impl $UDF {
249 pub fn new() -> Self {
250 Self {
251 signature: Signature::uniform(
252 1,
253 vec![DataType::Float64, DataType::Float32],
254 Volatility::Immutable,
255 ),
256 }
257 }
258 }
259
260 impl ScalarUDFImpl for $UDF {
261 fn name(&self) -> &str {
262 stringify!($NAME)
263 }
264
265 fn signature(&self) -> &Signature {
266 &self.signature
267 }
268
269 fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
270 let arg_type = &arg_types[0];
271
272 match arg_type {
273 DataType::Float32 => Ok(DataType::Float32),
274 _ => Ok(DataType::Float64),
275 }
276 }
277
278 fn is_strict(&self) -> bool {
279 $STRICT
280 }
281
282 fn output_ordering(
283 &self,
284 input: &[ExprProperties],
285 ) -> Result<SortProperties> {
286 $OUTPUT_ORDERING(input)
287 }
288
289 fn evaluate_bounds(&self, inputs: &[&Interval]) -> Result<Interval> {
290 $EVALUATE_BOUNDS(inputs)
291 }
292
293 fn invoke_with_args(
294 &self,
295 args: ScalarFunctionArgs,
296 ) -> Result<ColumnarValue> {
297 let args = ColumnarValue::values_to_arrays(&args.args)?;
298 let arr: ArrayRef = match args[0].data_type() {
299 DataType::Float64 => {
300 let values = args[0]
301 .as_primitive::<Float64Type>()
302 .try_unary::<_, Float64Type, _>(
303 |x: f64| -> std::result::Result<f64, ArrowError> {
304 if let Some(validate) = $VALIDATOR {
305 validate(x).map_err(|error| {
306 ArrowError::ComputeError(error.to_string())
307 })?;
308 }
309
310 Ok(f64::$UNARY_FUNC(x))
311 },
312 )?;
313 Arc::new(values) as ArrayRef
314 }
315 DataType::Float32 => {
316 let values = args[0]
317 .as_primitive::<Float32Type>()
318 .try_unary::<_, Float32Type, _>(
319 |x: f32| -> std::result::Result<f32, ArrowError> {
320 if let Some(validate) = $VALIDATOR {
321 validate(x as f64).map_err(|error| {
322 ArrowError::ComputeError(error.to_string())
323 })?;
324 }
325
326 Ok(f32::$UNARY_FUNC(x))
327 },
328 )?;
329 Arc::new(values) as ArrayRef
330 }
331 other => {
332 return exec_err!(
333 "Unsupported data type {other:?} for function {}",
334 self.name()
335 );
336 }
337 };
338
339 Ok(ColumnarValue::Array(arr))
340 }
341
342 fn documentation(&self) -> Option<&Documentation> {
343 Some($GET_DOC())
344 }
345 }
346 }
347 };
348}
349
350/// Macro to create a binary math UDF.
351///
352/// A binary math function takes two numeric arguments. When both arguments are
353/// Float32 the function is evaluated in single precision and returns Float32.
354/// Any other combination of numeric (or null) argument types is coerced to
355/// Float64 and returns Float64; in particular integers are widened to Float64
356/// rather than Float32 so that values needing more than 24 bits of mantissa are
357/// not silently rounded.
358///
359/// $UDF: the name of the UDF struct that implements `ScalarUDFImpl`
360/// $NAME: the name of the function
361/// $BINARY_FUNC: the binary function to apply to the argument
362/// $OUTPUT_ORDERING: the output ordering calculation method of the function
363/// $STRICT: whether the function returns NULL when any argument is NULL
364/// $GET_DOC: the function to get the documentation of the UDF
365macro_rules! make_math_binary_udf {
366 ($UDF:ident, $NAME:ident, $BINARY_FUNC:ident, $OUTPUT_ORDERING:expr, $STRICT:expr, $GET_DOC:expr) => {
367 $crate::make_udf_function!($NAME::$UDF, $NAME);
368
369 mod $NAME {
370
371 use std::sync::Arc;
372
373 use arrow::array::{ArrayRef, AsArray};
374 use arrow::datatypes::{DataType, Float32Type, Float64Type};
375 use datafusion_common::utils::take_function_args;
376 use datafusion_common::{Result, ScalarValue, internal_err};
377 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
378 use datafusion_expr::{
379 ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl,
380 Signature, Volatility,
381 };
382
383 #[derive(Debug, PartialEq, Eq, Hash)]
384 pub struct $UDF {
385 signature: Signature,
386 }
387
388 impl $UDF {
389 pub fn new() -> Self {
390 Self {
391 // Float64 is listed first so that integer (and other
392 // non-float) arguments coerce to Float64 rather than
393 // Float32; genuine Float32 arguments still match
394 // exactly and stay in single precision. Coercing
395 // integers to Float64 matters for correctness: Float32
396 // has only a 24-bit mantissa, so widening a large
397 // integer to Float32 would round it before the function
398 // is ever applied.
399 signature: Signature::uniform(
400 2,
401 vec![DataType::Float64, DataType::Float32],
402 Volatility::Immutable,
403 ),
404 }
405 }
406 }
407
408 impl ScalarUDFImpl for $UDF {
409 fn name(&self) -> &str {
410 stringify!($NAME)
411 }
412
413 fn signature(&self) -> &Signature {
414 &self.signature
415 }
416
417 fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
418 match (&arg_types[0], &arg_types[1]) {
419 (DataType::Float32, DataType::Float32) => Ok(DataType::Float32),
420 _ => Ok(DataType::Float64),
421 }
422 }
423
424 fn is_strict(&self) -> bool {
425 $STRICT
426 }
427
428 fn output_ordering(
429 &self,
430 input: &[ExprProperties],
431 ) -> Result<SortProperties> {
432 $OUTPUT_ORDERING(input)
433 }
434
435 fn invoke_with_args(
436 &self,
437 args: ScalarFunctionArgs,
438 ) -> Result<ColumnarValue> {
439 let ScalarFunctionArgs {
440 args, return_field, ..
441 } = args;
442 let return_type = return_field.data_type();
443 let [y, x] = take_function_args(self.name(), args)?;
444
445 match (y, x) {
446 (
447 ColumnarValue::Scalar(y_scalar),
448 ColumnarValue::Scalar(x_scalar),
449 ) => match (&y_scalar, &x_scalar) {
450 (y, x) if y.is_null() || x.is_null() => {
451 ColumnarValue::Scalar(ScalarValue::Null)
452 .cast_to(return_type, None)
453 }
454 (
455 ScalarValue::Float64(Some(yv)),
456 ScalarValue::Float64(Some(xv)),
457 ) => Ok(ColumnarValue::Scalar(ScalarValue::Float64(Some(
458 f64::$BINARY_FUNC(*yv, *xv),
459 )))),
460 (
461 ScalarValue::Float32(Some(yv)),
462 ScalarValue::Float32(Some(xv)),
463 ) => Ok(ColumnarValue::Scalar(ScalarValue::Float32(Some(
464 f32::$BINARY_FUNC(*yv, *xv),
465 )))),
466 _ => internal_err!(
467 "Unexpected scalar types for function {}: {:?}, {:?}",
468 self.name(),
469 y_scalar.data_type(),
470 x_scalar.data_type()
471 ),
472 },
473 (y, x) => {
474 let args = ColumnarValue::values_to_arrays(&[y, x])?;
475 let arr: ArrayRef = match args[0].data_type() {
476 DataType::Float64 => {
477 let y = args[0].as_primitive::<Float64Type>();
478 let x = args[1].as_primitive::<Float64Type>();
479 let result =
480 arrow::compute::binary::<_, _, _, Float64Type>(
481 y,
482 x,
483 |y, x| f64::$BINARY_FUNC(y, x),
484 )?;
485 Arc::new(result) as _
486 }
487 DataType::Float32 => {
488 let y = args[0].as_primitive::<Float32Type>();
489 let x = args[1].as_primitive::<Float32Type>();
490 let result =
491 arrow::compute::binary::<_, _, _, Float32Type>(
492 y,
493 x,
494 |y, x| f32::$BINARY_FUNC(y, x),
495 )?;
496 Arc::new(result) as _
497 }
498 other => {
499 return internal_err!(
500 "Unsupported data type {other:?} for function {}",
501 self.name()
502 );
503 }
504 };
505
506 Ok(ColumnarValue::Array(arr))
507 }
508 }
509 }
510
511 fn documentation(&self) -> Option<&Documentation> {
512 Some($GET_DOC())
513 }
514 }
515 }
516 };
517}