feat: allow array_slice to take an optional stride parameter (apach…

…e#10469) * feat: allow array_slice to take an optional stride parameter * Use ScalarUDF::call * Use create_function and add test * format * fix cargo doc
findepi · Jul 16, 2024 · 73f3d5d · 73f3d5d
1 parent 6f90e52
commit 73f3d5d
Show file tree

Hide file tree

Showing 23 changed files with 74 additions and 67 deletions.
diff --git a/datafusion/functions-array/src/array_has.rs b/datafusion/functions-array/src/array_has.rs
@@ -34,19 +34,19 @@ use std::any::Any;
 use std::sync::Arc;
 
 // Create static instances of ScalarUDFs for each function
-make_udf_function!(ArrayHas,
+make_udf_expr_and_func!(ArrayHas,
     array_has,
     first_array second_array, // arg name
     "returns true, if the element appears in the first array, otherwise false.", // doc
     array_has_udf // internal function name
 );
-make_udf_function!(ArrayHasAll,
+make_udf_expr_and_func!(ArrayHasAll,
     array_has_all,
     first_array second_array, // arg name
     "returns true if each element of the second array appears in the first array; otherwise, it returns false.", // doc
     array_has_all_udf // internal function name
 );
-make_udf_function!(ArrayHasAny,
+make_udf_expr_and_func!(ArrayHasAny,
     array_has_any,
     first_array second_array, // arg name
     "returns true if at least one element of the second array appears in the first array; otherwise, it returns false.", // doc

diff --git a/datafusion/functions-array/src/cardinality.rs b/datafusion/functions-array/src/cardinality.rs
@@ -29,7 +29,7 @@ use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     Cardinality,
     cardinality,
     array,

diff --git a/datafusion/functions-array/src/concat.rs b/datafusion/functions-array/src/concat.rs
@@ -36,7 +36,7 @@ use datafusion_expr::{
 
 use crate::utils::{align_array_dimensions, check_datatypes, make_scalar_function};
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayAppend,
     array_append,
     array element,                                // arg name
@@ -96,7 +96,7 @@ impl ScalarUDFImpl for ArrayAppend {
     }
 }
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayPrepend,
     array_prepend,
     element array,
@@ -156,7 +156,7 @@ impl ScalarUDFImpl for ArrayPrepend {
     }
 }
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayConcat,
     array_concat,
     "Concatenates arrays.",

diff --git a/datafusion/functions-array/src/dimension.rs b/datafusion/functions-array/src/dimension.rs
@@ -33,7 +33,7 @@ use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility};
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayDims,
     array_dims,
     array,
@@ -88,7 +88,7 @@ impl ScalarUDFImpl for ArrayDims {
     }
 }
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayNdims,
     array_ndims,
     array,

diff --git a/datafusion/functions-array/src/empty.rs b/datafusion/functions-array/src/empty.rs
@@ -28,7 +28,7 @@ use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayEmpty,
     array_empty,
     array,

diff --git a/datafusion/functions-array/src/except.rs b/datafusion/functions-array/src/except.rs
@@ -31,7 +31,7 @@ use std::any::Any;
 use std::collections::HashSet;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayExcept,
     array_except,
     first_array second_array,

diff --git a/datafusion/functions-array/src/extract.rs b/datafusion/functions-array/src/extract.rs
@@ -44,31 +44,25 @@ use std::sync::Arc;
 use crate::utils::make_scalar_function;
 
 // Create static instances of ScalarUDFs for each function
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayElement,
     array_element,
     array element,
     "extracts the element with the index n from the array.",
     array_element_udf
 );
 
-make_udf_function!(
-    ArraySlice,
-    array_slice,
-    array begin end stride,
-    "returns a slice of the array.",
-    array_slice_udf
-);
+create_func!(ArraySlice, array_slice_udf);
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayPopFront,
     array_pop_front,
     array,
     "returns the array without the first element.",
     array_pop_front_udf
 );
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayPopBack,
     array_pop_back,
     array,
@@ -224,6 +218,15 @@ where
     Ok(arrow::array::make_array(data))
 }
 
+#[doc = "returns a slice of the array."]
+pub fn array_slice(array: Expr, begin: Expr, end: Expr, stride: Option<Expr>) -> Expr {
+    let args = match stride {
+        Some(stride) => vec![array, begin, end, stride],
+        None => vec![array, begin, end],
+    };
+    array_slice_udf().call(args)
+}
+
 #[derive(Debug)]
 pub(super) struct ArraySlice {
     signature: Signature,

diff --git a/datafusion/functions-array/src/flatten.rs b/datafusion/functions-array/src/flatten.rs
@@ -31,7 +31,7 @@ use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     Flatten,
     flatten,
     array,

diff --git a/datafusion/functions-array/src/length.rs b/datafusion/functions-array/src/length.rs
@@ -32,7 +32,7 @@ use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayLength,
     array_length,
     array,

diff --git a/datafusion/functions-array/src/macros.rs b/datafusion/functions-array/src/macros.rs
@@ -19,8 +19,8 @@
 ///
 /// 1. Single `ScalarUDF` instance
 ///
-/// Creates a singleton `ScalarUDF` of the `$UDF` function named `$GNAME` and a
-/// function named `$NAME` which returns that function named $NAME.
+/// Creates a singleton `ScalarUDF` of the `$UDF` function named `STATIC_$(UDF)` and a
+/// function named `$SCALAR_UDF_FUNC` which returns that function named `STATIC_$(UDF)`.
 ///
 /// This is used to ensure creating the list of `ScalarUDF` only happens once.
 ///
@@ -41,10 +41,9 @@
 /// * `arg`: 0 or more named arguments for the function
 /// * `DOC`: documentation string for the function
 /// * `SCALAR_UDF_FUNC`: name of the function to create (just) the `ScalarUDF`
-/// * `GNAME`: name for the single static instance of the `ScalarUDF`
 ///
 /// [`ScalarUDFImpl`]: datafusion_expr::ScalarUDFImpl
-macro_rules! make_udf_function {
+macro_rules! make_udf_expr_and_func {
     ($UDF:ty, $EXPR_FN:ident, $($arg:ident)*, $DOC:expr , $SCALAR_UDF_FN:ident) => {
         paste::paste! {
             // "fluent expr_fn" style function
@@ -55,25 +54,7 @@ macro_rules! make_udf_function {
                     vec![$($arg),*],
                 ))
             }
-
-            /// Singleton instance of [`$UDF`], ensures the UDF is only created once
-            /// named STATIC_$(UDF). For example `STATIC_ArrayToString`
-            #[allow(non_upper_case_globals)]
-            static [< STATIC_ $UDF >]: std::sync::OnceLock<std::sync::Arc<datafusion_expr::ScalarUDF>> =
-                std::sync::OnceLock::new();
-
-            /// ScalarFunction that returns a [`ScalarUDF`] for [`$UDF`]
-            ///
-            /// [`ScalarUDF`]: datafusion_expr::ScalarUDF
-            pub fn $SCALAR_UDF_FN() -> std::sync::Arc<datafusion_expr::ScalarUDF> {
-                [< STATIC_ $UDF >]
-                    .get_or_init(|| {
-                        std::sync::Arc::new(datafusion_expr::ScalarUDF::new_from_impl(
-                            <$UDF>::new(),
-                        ))
-                    })
-                    .clone()
-            }
+            create_func!($UDF, $SCALAR_UDF_FN);
         }
     };
     ($UDF:ty, $EXPR_FN:ident, $DOC:expr , $SCALAR_UDF_FN:ident) => {
@@ -86,7 +67,24 @@ macro_rules! make_udf_function {
                     arg,
                 ))
             }
+            create_func!($UDF, $SCALAR_UDF_FN);
+        }
+    };
+}
 
+/// Creates a singleton `ScalarUDF` of the `$UDF` function named `STATIC_$(UDF)` and a
+/// function named `$SCALAR_UDF_FUNC` which returns that function named `STATIC_$(UDF)`.
+///
+/// This is used to ensure creating the list of `ScalarUDF` only happens once.
+///
+/// # Arguments
+/// * `UDF`: name of the [`ScalarUDFImpl`]
+/// * `SCALAR_UDF_FUNC`: name of the function to create (just) the `ScalarUDF`
+///
+/// [`ScalarUDFImpl`]: datafusion_expr::ScalarUDFImpl
+macro_rules! create_func {
+    ($UDF:ty, $SCALAR_UDF_FN:ident) => {
+        paste::paste! {
             /// Singleton instance of [`$UDF`], ensures the UDF is only created once
             /// named STATIC_$(UDF). For example `STATIC_ArrayToString`
             #[allow(non_upper_case_globals)]

diff --git a/datafusion/functions-array/src/make_array.rs b/datafusion/functions-array/src/make_array.rs
@@ -35,7 +35,7 @@ use datafusion_expr::{Expr, TypeSignature};
 
 use crate::utils::make_scalar_function;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     MakeArray,
     make_array,
     "Returns an Arrow array using the specified input expressions.",

diff --git a/datafusion/functions-array/src/position.rs b/datafusion/functions-array/src/position.rs
@@ -37,7 +37,7 @@ use itertools::Itertools;
 
 use crate::utils::{compare_element_to_list, make_scalar_function};
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayPosition,
     array_position,
     array element index,
@@ -168,7 +168,7 @@ fn generic_position<OffsetSize: OffsetSizeTrait>(
     Ok(Arc::new(UInt64Array::from(data)))
 }
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayPositions,
     array_positions,
     array element, // arg name

diff --git a/datafusion/functions-array/src/range.rs b/datafusion/functions-array/src/range.rs
@@ -35,7 +35,7 @@ use datafusion_expr::{
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     Range,
     range,
     start stop step,
@@ -106,7 +106,7 @@ impl ScalarUDFImpl for Range {
     }
 }
 
-make_udf_function!(
+make_udf_expr_and_func!(
     GenSeries,
     gen_series,
     start stop step,

diff --git a/datafusion/functions-array/src/remove.rs b/datafusion/functions-array/src/remove.rs
@@ -32,7 +32,7 @@ use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayRemove,
     array_remove,
     array element,
@@ -81,7 +81,7 @@ impl ScalarUDFImpl for ArrayRemove {
     }
 }
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayRemoveN,
     array_remove_n,
     array element max,
@@ -130,7 +130,7 @@ impl ScalarUDFImpl for ArrayRemoveN {
     }
 }
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayRemoveAll,
     array_remove_all,
     array element,

diff --git a/datafusion/functions-array/src/repeat.rs b/datafusion/functions-array/src/repeat.rs
@@ -34,7 +34,7 @@ use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayRepeat,
     array_repeat,
     element count, // arg name

diff --git a/datafusion/functions-array/src/replace.rs b/datafusion/functions-array/src/replace.rs
@@ -38,19 +38,19 @@ use std::any::Any;
 use std::sync::Arc;
 
 // Create static instances of ScalarUDFs for each function
-make_udf_function!(ArrayReplace,
+make_udf_expr_and_func!(ArrayReplace,
     array_replace,
     array from to,
     "replaces the first occurrence of the specified element with another specified element.",
     array_replace_udf
 );
-make_udf_function!(ArrayReplaceN,
+make_udf_expr_and_func!(ArrayReplaceN,
     array_replace_n,
     array from to max,
     "replaces the first `max` occurrences of the specified element with another specified element.",
     array_replace_n_udf
 );
-make_udf_function!(ArrayReplaceAll,
+make_udf_expr_and_func!(ArrayReplaceAll,
     array_replace_all,
     array from to,
     "replaces all occurrences of the specified element with another specified element.",

diff --git a/datafusion/functions-array/src/resize.rs b/datafusion/functions-array/src/resize.rs
@@ -30,7 +30,7 @@ use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayResize,
     array_resize,
     array size value,

diff --git a/datafusion/functions-array/src/reverse.rs b/datafusion/functions-array/src/reverse.rs
@@ -30,7 +30,7 @@ use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}
 use std::any::Any;
 use std::sync::Arc;
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayReverse,
     array_reverse,
     array,

diff --git a/datafusion/functions-array/src/rewrite.rs b/datafusion/functions-array/src/rewrite.rs
@@ -171,7 +171,7 @@ impl FunctionRewrite for ArrayFunctionRewriter {
                         stop,
                         stride,
                     },
-            }) => Transformed::yes(array_slice(*expr, *start, *stop, *stride)),
+            }) => Transformed::yes(array_slice(*expr, *start, *stop, Some(*stride))),
 
             _ => Transformed::no(expr),
         };

diff --git a/datafusion/functions-array/src/set_ops.rs b/datafusion/functions-array/src/set_ops.rs
@@ -37,23 +37,23 @@ use std::fmt::{Display, Formatter};
 use std::sync::Arc;
 
 // Create static instances of ScalarUDFs for each function
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayUnion,
     array_union,
     array1 array2,
     "returns an array of the elements in the union of array1 and array2 without duplicates.",
     array_union_udf
 );
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayIntersect,
     array_intersect,
     first_array second_array,
     "returns an array of the elements in the intersection of array1 and array2.",
     array_intersect_udf
 );
 
-make_udf_function!(
+make_udf_expr_and_func!(
     ArrayDistinct,
     array_distinct,
     array,