Skip to content

Commit 45d9820

Browse files
irenjjirenjj
authored
chore: Migrate Array Functions to invoke_with_args (#14726)
* chore: Migrate Array Functions to invoke_with_args * fix clippy * fix issues * fix tests * fix
1 parent 873b5f7 commit 45d9820

25 files changed

+137
-172
lines changed

‎datafusion/functions-nested/src/array_has.rs

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use arrow::datatypes::DataType;
2525
use arrow::row::{RowConverter, Rows, SortField};
2626
use datafusion_common::cast::as_generic_list_array;
2727
use datafusion_common::utils::string_utils::string_array_to_vec;
28+
use datafusion_common::utils::take_function_args;
2829
use datafusion_common::{exec_err, Result, ScalarValue};
2930
use datafusion_expr::{
3031
ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
@@ -120,15 +121,15 @@ impl ScalarUDFImpl for ArrayHas {
120121
Ok(DataType::Boolean)
121122
}
122123

123-
fn invoke_batch(
124+
fn invoke_with_args(
124125
&self,
125-
args: &[ColumnarValue],
126-
_number_rows: usize,
126+
args: datafusion_expr::ScalarFunctionArgs,
127127
) -> Result<ColumnarValue> {
128-
match &args[1] {
128+
let [first_arg, second_arg] = take_function_args(self.name(), &args.args)?;
129+
match &second_arg {
129130
ColumnarValue::Array(array_needle) => {
130131
// the needle is already an array, convert the haystack to an array of the same length
131-
let haystack = args[0].to_array(array_needle.len())?;
132+
let haystack = first_arg.to_array(array_needle.len())?;
132133
let array = array_has_inner_for_array(&haystack, array_needle)?;
133134
Ok(ColumnarValue::Array(array))
134135
}
@@ -140,11 +141,11 @@ impl ScalarUDFImpl for ArrayHas {
140141
}
141142

142143
// since the needle is a scalar, convert it to an array of size 1
143-
let haystack = args[0].to_array(1)?;
144+
let haystack = first_arg.to_array(1)?;
144145
let needle = scalar_needle.to_array_of_size(1)?;
145146
let needle = Scalar::new(needle);
146147
let array = array_has_inner_for_scalar(&haystack, &needle)?;
147-
if let ColumnarValue::Scalar(_) = &args[0] {
148+
if let ColumnarValue::Scalar(_) = &first_arg {
148149
// If both inputs are scalar, keeps output as scalar
149150
let scalar_value = ScalarValue::try_from_array(&array, 0)?;
150151
Ok(ColumnarValue::Scalar(scalar_value))
@@ -332,12 +333,11 @@ impl ScalarUDFImpl for ArrayHasAll {
332333
Ok(DataType::Boolean)
333334
}
334335

335-
fn invoke_batch(
336+
fn invoke_with_args(
336337
&self,
337-
args: &[ColumnarValue],
338-
_number_rows: usize,
338+
args: datafusion_expr::ScalarFunctionArgs,
339339
) -> Result<ColumnarValue> {
340-
make_scalar_function(array_has_all_inner)(args)
340+
make_scalar_function(array_has_all_inner)(&args.args)
341341
}
342342

343343
fn aliases(&self) -> &[String] {
@@ -407,12 +407,11 @@ impl ScalarUDFImpl for ArrayHasAny {
407407
Ok(DataType::Boolean)
408408
}
409409

410-
fn invoke_batch(
410+
fn invoke_with_args(
411411
&self,
412-
args: &[ColumnarValue],
413-
_number_rows: usize,
412+
args: datafusion_expr::ScalarFunctionArgs,
414413
) -> Result<ColumnarValue> {
415-
make_scalar_function(array_has_any_inner)(args)
414+
make_scalar_function(array_has_any_inner)(&args.args)
416415
}
417416

418417
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/cardinality.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -112,12 +112,11 @@ impl ScalarUDFImpl for Cardinality {
112112
})
113113
}
114114

115-
fn invoke_batch(
115+
fn invoke_with_args(
116116
&self,
117-
args: &[ColumnarValue],
118-
_number_rows: usize,
117+
args: datafusion_expr::ScalarFunctionArgs,
119118
) -> Result<ColumnarValue> {
120-
make_scalar_function(cardinality_inner)(args)
119+
make_scalar_function(cardinality_inner)(&args.args)
121120
}
122121

123122
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/concat.rs

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -109,12 +109,11 @@ impl ScalarUDFImpl for ArrayAppend {
109109
Ok(arg_types[0].clone())
110110
}
111111

112-
fn invoke_batch(
112+
fn invoke_with_args(
113113
&self,
114-
args: &[ColumnarValue],
115-
_number_rows: usize,
114+
args: datafusion_expr::ScalarFunctionArgs,
116115
) -> Result<ColumnarValue> {
117-
make_scalar_function(array_append_inner)(args)
116+
make_scalar_function(array_append_inner)(&args.args)
118117
}
119118

120119
fn aliases(&self) -> &[String] {
@@ -205,12 +204,11 @@ impl ScalarUDFImpl for ArrayPrepend {
205204
Ok(arg_types[1].clone())
206205
}
207206

208-
fn invoke_batch(
207+
fn invoke_with_args(
209208
&self,
210-
args: &[ColumnarValue],
211-
_number_rows: usize,
209+
args: datafusion_expr::ScalarFunctionArgs,
212210
) -> Result<ColumnarValue> {
213-
make_scalar_function(array_prepend_inner)(args)
211+
make_scalar_function(array_prepend_inner)(&args.args)
214212
}
215213

216214
fn aliases(&self) -> &[String] {
@@ -324,12 +322,11 @@ impl ScalarUDFImpl for ArrayConcat {
324322
Ok(expr_type)
325323
}
326324

327-
fn invoke_batch(
325+
fn invoke_with_args(
328326
&self,
329-
args: &[ColumnarValue],
330-
_number_rows: usize,
327+
args: datafusion_expr::ScalarFunctionArgs,
331328
) -> Result<ColumnarValue> {
332-
make_scalar_function(array_concat_inner)(args)
329+
make_scalar_function(array_concat_inner)(&args.args)
333330
}
334331

335332
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/dimension.rs

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -106,12 +106,11 @@ impl ScalarUDFImpl for ArrayDims {
106106
})
107107
}
108108

109-
fn invoke_batch(
109+
fn invoke_with_args(
110110
&self,
111-
args: &[ColumnarValue],
112-
_number_rows: usize,
111+
args: datafusion_expr::ScalarFunctionArgs,
113112
) -> Result<ColumnarValue> {
114-
make_scalar_function(array_dims_inner)(args)
113+
make_scalar_function(array_dims_inner)(&args.args)
115114
}
116115

117116
fn aliases(&self) -> &[String] {
@@ -184,12 +183,11 @@ impl ScalarUDFImpl for ArrayNdims {
184183
})
185184
}
186185

187-
fn invoke_batch(
186+
fn invoke_with_args(
188187
&self,
189-
args: &[ColumnarValue],
190-
_number_rows: usize,
188+
args: datafusion_expr::ScalarFunctionArgs,
191189
) -> Result<ColumnarValue> {
192-
make_scalar_function(array_ndims_inner)(args)
190+
make_scalar_function(array_ndims_inner)(&args.args)
193191
}
194192

195193
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/distance.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -124,12 +124,11 @@ impl ScalarUDFImpl for ArrayDistance {
124124
Ok(result)
125125
}
126126

127-
fn invoke_batch(
127+
fn invoke_with_args(
128128
&self,
129-
args: &[ColumnarValue],
130-
_number_rows: usize,
129+
args: datafusion_expr::ScalarFunctionArgs,
131130
) -> Result<ColumnarValue> {
132-
make_scalar_function(array_distance_inner)(args)
131+
make_scalar_function(array_distance_inner)(&args.args)
133132
}
134133

135134
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/empty.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -98,12 +98,11 @@ impl ScalarUDFImpl for ArrayEmpty {
9898
})
9999
}
100100

101-
fn invoke_batch(
101+
fn invoke_with_args(
102102
&self,
103-
args: &[ColumnarValue],
104-
_number_rows: usize,
103+
args: datafusion_expr::ScalarFunctionArgs,
105104
) -> Result<ColumnarValue> {
106-
make_scalar_function(array_empty_inner)(args)
105+
make_scalar_function(array_empty_inner)(&args.args)
107106
}
108107

109108
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/except.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -106,12 +106,11 @@ impl ScalarUDFImpl for ArrayExcept {
106106
}
107107
}
108108

109-
fn invoke_batch(
109+
fn invoke_with_args(
110110
&self,
111-
args: &[ColumnarValue],
112-
_number_rows: usize,
111+
args: datafusion_expr::ScalarFunctionArgs,
113112
) -> Result<ColumnarValue> {
114-
make_scalar_function(array_except_inner)(args)
113+
make_scalar_function(array_except_inner)(&args.args)
115114
}
116115

117116
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/extract.rs

Lines changed: 16 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -172,12 +172,11 @@ impl ScalarUDFImpl for ArrayElement {
172172
}
173173
}
174174

175-
fn invoke_batch(
175+
fn invoke_with_args(
176176
&self,
177-
args: &[ColumnarValue],
178-
_number_rows: usize,
177+
args: datafusion_expr::ScalarFunctionArgs,
179178
) -> Result<ColumnarValue> {
180-
make_scalar_function(array_element_inner)(args)
179+
make_scalar_function(array_element_inner)(&args.args)
181180
}
182181

183182
fn aliases(&self) -> &[String] {
@@ -395,12 +394,11 @@ impl ScalarUDFImpl for ArraySlice {
395394
Ok(arg_types[0].clone())
396395
}
397396

398-
fn invoke_batch(
397+
fn invoke_with_args(
399398
&self,
400-
args: &[ColumnarValue],
401-
_number_rows: usize,
399+
args: datafusion_expr::ScalarFunctionArgs,
402400
) -> Result<ColumnarValue> {
403-
make_scalar_function(array_slice_inner)(args)
401+
make_scalar_function(array_slice_inner)(&args.args)
404402
}
405403

406404
fn aliases(&self) -> &[String] {
@@ -704,12 +702,11 @@ impl ScalarUDFImpl for ArrayPopFront {
704702
Ok(arg_types[0].clone())
705703
}
706704

707-
fn invoke_batch(
705+
fn invoke_with_args(
708706
&self,
709-
args: &[ColumnarValue],
710-
_number_rows: usize,
707+
args: datafusion_expr::ScalarFunctionArgs,
711708
) -> Result<ColumnarValue> {
712-
make_scalar_function(array_pop_front_inner)(args)
709+
make_scalar_function(array_pop_front_inner)(&args.args)
713710
}
714711

715712
fn aliases(&self) -> &[String] {
@@ -812,12 +809,11 @@ impl ScalarUDFImpl for ArrayPopBack {
812809
Ok(arg_types[0].clone())
813810
}
814811

815-
fn invoke_batch(
812+
fn invoke_with_args(
816813
&self,
817-
args: &[ColumnarValue],
818-
_number_rows: usize,
814+
args: datafusion_expr::ScalarFunctionArgs,
819815
) -> Result<ColumnarValue> {
820-
make_scalar_function(array_pop_back_inner)(args)
816+
make_scalar_function(array_pop_back_inner)(&args.args)
821817
}
822818

823819
fn aliases(&self) -> &[String] {
@@ -918,13 +914,13 @@ impl ScalarUDFImpl for ArrayAnyValue {
918914
}
919915
}
920916

921-
fn invoke_batch(
917+
fn invoke_with_args(
922918
&self,
923-
args: &[ColumnarValue],
924-
_number_rows: usize,
919+
args: datafusion_expr::ScalarFunctionArgs,
925920
) -> Result<ColumnarValue> {
926-
make_scalar_function(array_any_value_inner)(args)
921+
make_scalar_function(array_any_value_inner)(&args.args)
927922
}
923+
928924
fn aliases(&self) -> &[String] {
929925
&self.aliases
930926
}

‎datafusion/functions-nested/src/flatten.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -124,12 +124,11 @@ impl ScalarUDFImpl for Flatten {
124124
Ok(data_type)
125125
}
126126

127-
fn invoke_batch(
127+
fn invoke_with_args(
128128
&self,
129-
args: &[ColumnarValue],
130-
_number_rows: usize,
129+
args: datafusion_expr::ScalarFunctionArgs,
131130
) -> Result<ColumnarValue> {
132-
make_scalar_function(flatten_inner)(args)
131+
make_scalar_function(flatten_inner)(&args.args)
133132
}
134133

135134
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/length.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -103,12 +103,11 @@ impl ScalarUDFImpl for ArrayLength {
103103
})
104104
}
105105

106-
fn invoke_batch(
106+
fn invoke_with_args(
107107
&self,
108-
args: &[ColumnarValue],
109-
_number_rows: usize,
108+
args: datafusion_expr::ScalarFunctionArgs,
110109
) -> Result<ColumnarValue> {
111-
make_scalar_function(array_length_inner)(args)
110+
make_scalar_function(array_length_inner)(&args.args)
112111
}
113112

114113
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/make_array.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -117,12 +117,11 @@ impl ScalarUDFImpl for MakeArray {
117117
}
118118
}
119119

120-
fn invoke_batch(
120+
fn invoke_with_args(
121121
&self,
122-
args: &[ColumnarValue],
123-
_number_rows: usize,
122+
args: datafusion_expr::ScalarFunctionArgs,
124123
) -> Result<ColumnarValue> {
125-
make_scalar_function(make_array_inner)(args)
124+
make_scalar_function(make_array_inner)(&args.args)
126125
}
127126

128127
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/map.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -273,12 +273,11 @@ impl ScalarUDFImpl for MapFunc {
273273
))
274274
}
275275

276-
fn invoke_batch(
276+
fn invoke_with_args(
277277
&self,
278-
args: &[ColumnarValue],
279-
_number_rows: usize,
278+
args: datafusion_expr::ScalarFunctionArgs,
280279
) -> Result<ColumnarValue> {
281-
make_map_batch(args)
280+
make_map_batch(&args.args)
282281
}
283282

284283
fn documentation(&self) -> Option<&Documentation> {

‎datafusion/functions-nested/src/map_extract.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -110,12 +110,11 @@ impl ScalarUDFImpl for MapExtract {
110110
))))
111111
}
112112

113-
fn invoke_batch(
113+
fn invoke_with_args(
114114
&self,
115-
args: &[ColumnarValue],
116-
_number_rows: usize,
115+
args: datafusion_expr::ScalarFunctionArgs,
117116
) -> Result<ColumnarValue> {
118-
make_scalar_function(map_extract_inner)(args)
117+
make_scalar_function(map_extract_inner)(&args.args)
119118
}
120119

121120
fn aliases(&self) -> &[String] {

‎datafusion/functions-nested/src/map_keys.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -100,12 +100,11 @@ impl ScalarUDFImpl for MapKeysFunc {
100100
))))
101101
}
102102

103-
fn invoke_batch(
103+
fn invoke_with_args(
104104
&self,
105-
args: &[ColumnarValue],
106-
_number_rows: usize,
105+
args: datafusion_expr::ScalarFunctionArgs,
107106
) -> Result<ColumnarValue> {
108-
make_scalar_function(map_keys_inner)(args)
107+
make_scalar_function(map_keys_inner)(&args.args)
109108
}
110109

111110
fn documentation(&self) -> Option<&Documentation> {

0 commit comments

Comments
 (0)