Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: allow array_slice to take an optional stride parameter #10469

Merged
merged 5 commits into from
May 13, 2024
Merged

feat: allow array_slice to take an optional stride parameter #10469

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ac5c5f1
feat: allow array_slice to take an optional stride parameter
jonahgao May 12, 2024
0746583
Use ScalarUDF::call
jonahgao May 12, 2024
e981d28
Use create_function and add test
jonahgao May 13, 2024
debe6f5
format
jonahgao May 13, 2024
2df4fe0
fix cargo doc
jonahgao May 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions datafusion/functions-array/src/array_has.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/cardinality.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
6 changes: 3 additions & 3 deletions datafusion/functions-array/src/concat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -96,7 +96,7 @@ impl ScalarUDFImpl for ArrayAppend {
}
}

make_udf_function!(
make_udf_expr_and_func!(
ArrayPrepend,
array_prepend,
element array,
Expand Down Expand Up @@ -156,7 +156,7 @@ impl ScalarUDFImpl for ArrayPrepend {
}
}

make_udf_function!(
make_udf_expr_and_func!(
ArrayConcat,
array_concat,
"Concatenates arrays.",
Expand Down
4 changes: 2 additions & 2 deletions datafusion/functions-array/src/dimension.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -88,7 +88,7 @@ impl ScalarUDFImpl for ArrayDims {
}
}

make_udf_function!(
make_udf_expr_and_func!(
ArrayNdims,
array_ndims,
array,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/empty.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/except.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
23 changes: 13 additions & 10 deletions datafusion/functions-array/src/extract.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/flatten.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/length.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
44 changes: 21 additions & 23 deletions datafusion/functions-array/src/macros.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand All @@ -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
Expand All @@ -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) => {
Expand All @@ -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) => {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would recommend to pull out create_function like https://github.com/apache/datafusion/blob/main/datafusion/functions-aggregate/src/macros.rs.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Applied. This eliminates some duplication of macro codes. Thank you @jayzhan211 .

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)]
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/make_array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand Down
4 changes: 2 additions & 2 deletions datafusion/functions-array/src/position.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions datafusion/functions-array/src/range.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -106,7 +106,7 @@ impl ScalarUDFImpl for Range {
}
}

make_udf_function!(
make_udf_expr_and_func!(
GenSeries,
gen_series,
start stop step,
Expand Down
6 changes: 3 additions & 3 deletions datafusion/functions-array/src/remove.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -81,7 +81,7 @@ impl ScalarUDFImpl for ArrayRemove {
}
}

make_udf_function!(
make_udf_expr_and_func!(
ArrayRemoveN,
array_remove_n,
array element max,
Expand Down Expand Up @@ -130,7 +130,7 @@ impl ScalarUDFImpl for ArrayRemoveN {
}
}

make_udf_function!(
make_udf_expr_and_func!(
ArrayRemoveAll,
array_remove_all,
array element,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/repeat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions datafusion/functions-array/src/replace.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/resize.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/reverse.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-array/src/rewrite.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
};
Expand Down
6 changes: 3 additions & 3 deletions datafusion/functions-array/src/set_ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading