datafusion_functions_nested/
remove.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//! [`ScalarUDFImpl`] definitions for array_remove, array_remove_n, array_remove_all functions.
19
20use crate::utils;
21use arrow::array::{
22    Array, ArrayRef, Capacities, GenericListArray, MutableArrayData, NullBufferBuilder,
23    OffsetSizeTrait, Scalar, cast::AsArray, make_array, new_null_array,
24};
25use arrow::buffer::OffsetBuffer;
26use arrow::datatypes::{DataType, FieldRef};
27use datafusion_common::cast::as_int64_array;
28use datafusion_common::utils::ListCoercion;
29use datafusion_common::{
30    Result, ScalarValue, exec_err, internal_err, utils::take_function_args,
31};
32use datafusion_expr::{
33    ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation,
34    ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
35};
36use datafusion_macros::user_doc;
37use std::sync::Arc;
38
39make_udf_expr_and_func!(
40    ArrayRemove,
41    array_remove,
42    array element,
43    "removes the first element from the array equal to the given value. NULL elements already in the array are preserved when removing a non-NULL value. If `element` evaluates to NULL, the result is NULL rather than removing NULL entries.",
44    array_remove_udf
45);
46
47#[user_doc(
48    doc_section(label = "Array Functions"),
49    description = "Removes the first element from the array equal to the given value. NULL elements already in the array are preserved when removing a non-NULL value. If `element` evaluates to NULL, the result is NULL rather than removing NULL entries.",
50    syntax_example = "array_remove(array, element)",
51    sql_example = r#"```sql
52> select array_remove([1, 2, 2, 3, 2, 1, 4], 2);
53+----------------------------------------------+
54| array_remove(List([1,2,2,3,2,1,4]),Int64(2)) |
55+----------------------------------------------+
56| [1, 2, 3, 2, 1, 4]                           |
57+----------------------------------------------+
58
59> select array_remove([1, 2, NULL, 2, 4], 2);
60+---------------------------------------------------+
61| array_remove(List([1,2,NULL,2,4]),Int64(2)) |
62+---------------------------------------------------+
63| [1, NULL, 2, 4]                              |
64+---------------------------------------------------+
65```"#,
66    argument(
67        name = "array",
68        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
69    ),
70    argument(
71        name = "element",
72        description = "Element to be removed from the array."
73    )
74)]
75#[derive(Debug, PartialEq, Eq, Hash)]
76pub struct ArrayRemove {
77    signature: Signature,
78    aliases: Vec<String>,
79}
80
81impl Default for ArrayRemove {
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87impl ArrayRemove {
88    pub fn new() -> Self {
89        Self {
90            signature: Signature::array_and_element(Volatility::Immutable),
91            aliases: vec!["list_remove".to_string()],
92        }
93    }
94}
95
96impl ScalarUDFImpl for ArrayRemove {
97    fn name(&self) -> &str {
98        "array_remove"
99    }
100
101    fn signature(&self) -> &Signature {
102        &self.signature
103    }
104
105    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
106        internal_err!("return_field_from_args should be used instead")
107    }
108
109    fn return_field_from_args(
110        &self,
111        args: datafusion_expr::ReturnFieldArgs,
112    ) -> Result<FieldRef> {
113        let array_field = args.arg_fields[0].as_ref().clone();
114        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
115        Ok(Arc::new(array_field.with_nullable(nullable)))
116    }
117
118    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
119        let [list_arg, element_arg] = take_function_args(self.name(), &args.args)?;
120        let num_rows = args.number_rows;
121        let list_array = list_arg.to_array(num_rows)?;
122        match element_arg {
123            ColumnarValue::Scalar(scalar_element)
124                if !scalar_element.is_null()
125                    && !scalar_element.data_type().is_nested() =>
126            {
127                let result =
128                    array_remove_with_scalar_args(&list_array, scalar_element, 1i64)?;
129                Ok(ColumnarValue::Array(result))
130            }
131            element_arg => {
132                let element_array = element_arg.to_array(num_rows)?;
133                let result =
134                    array_remove_internal(&list_array, &element_array, &[Some(1)])?;
135                Ok(ColumnarValue::Array(result))
136            }
137        }
138    }
139
140    fn aliases(&self) -> &[String] {
141        &self.aliases
142    }
143
144    fn documentation(&self) -> Option<&Documentation> {
145        self.doc()
146    }
147}
148
149make_udf_expr_and_func!(
150    ArrayRemoveN,
151    array_remove_n,
152    array element max,
153    "removes the first `max` elements from the array equal to the given value. NULL elements already in the array are preserved when removing a non-NULL value. If `element` evaluates to NULL, the result is NULL rather than removing NULL entries.",
154    array_remove_n_udf
155);
156
157#[user_doc(
158    doc_section(label = "Array Functions"),
159    description = "Removes the first `max` elements from the array equal to the given value. NULL elements already in the array are preserved when removing a non-NULL value. If `element` evaluates to NULL, the result is NULL rather than removing NULL entries.",
160    syntax_example = "array_remove_n(array, element, max)",
161    sql_example = r#"```sql
162> select array_remove_n([1, 2, 2, 3, 2, 1, 4], 2, 2);
163+---------------------------------------------------------+
164| array_remove_n(List([1,2,2,3,2,1,4]),Int64(2),Int64(2)) |
165+---------------------------------------------------------+
166| [1, 3, 2, 1, 4]                                         |
167+---------------------------------------------------------+
168
169> select array_remove_n([1, 2, NULL, 2, 4], 2, 2);
170+----------------------------------------------------------+
171| array_remove_n(List([1,2,NULL,2,4]),Int64(2),Int64(2)) |
172+----------------------------------------------------------+
173| [1, NULL, 4]                                            |
174+----------------------------------------------------------+
175```"#,
176    argument(
177        name = "array",
178        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
179    ),
180    argument(
181        name = "element",
182        description = "Element to be removed from the array."
183    ),
184    argument(name = "max", description = "Number of first occurrences to remove.")
185)]
186#[derive(Debug, PartialEq, Eq, Hash)]
187pub struct ArrayRemoveN {
188    signature: Signature,
189    aliases: Vec<String>,
190}
191
192impl Default for ArrayRemoveN {
193    fn default() -> Self {
194        Self::new()
195    }
196}
197
198impl ArrayRemoveN {
199    pub fn new() -> Self {
200        Self {
201            signature: Signature::new(
202                TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
203                    arguments: vec![
204                        ArrayFunctionArgument::Array,
205                        ArrayFunctionArgument::Element,
206                        ArrayFunctionArgument::Index,
207                    ],
208                    array_coercion: Some(ListCoercion::FixedSizedListToList),
209                }),
210                Volatility::Immutable,
211            ),
212            aliases: vec!["list_remove_n".to_string()],
213        }
214    }
215}
216
217impl ScalarUDFImpl for ArrayRemoveN {
218    fn name(&self) -> &str {
219        "array_remove_n"
220    }
221
222    fn signature(&self) -> &Signature {
223        &self.signature
224    }
225
226    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
227        internal_err!("return_field_from_args should be used instead")
228    }
229
230    fn return_field_from_args(
231        &self,
232        args: datafusion_expr::ReturnFieldArgs,
233    ) -> Result<FieldRef> {
234        let array_field = args.arg_fields[0].as_ref().clone();
235        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
236        Ok(Arc::new(array_field.with_nullable(nullable)))
237    }
238
239    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
240        let [list_arg, element_arg, max_arg] =
241            take_function_args(self.name(), &args.args)?;
242        let num_rows = args.number_rows;
243        let list_array = list_arg.to_array(num_rows)?;
244        match (element_arg, max_arg) {
245            (
246                ColumnarValue::Scalar(scalar_element),
247                ColumnarValue::Scalar(scalar_max),
248            ) if !scalar_element.is_null() && !scalar_element.data_type().is_nested() => {
249                let ScalarValue::Int64(Some(n)) = scalar_max else {
250                    return Ok(ColumnarValue::Array(new_null_array(
251                        list_array.data_type(),
252                        num_rows,
253                    )));
254                };
255                let result =
256                    array_remove_with_scalar_args(&list_array, scalar_element, *n)?;
257                Ok(ColumnarValue::Array(result))
258            }
259            (element_arg, max_arg) => {
260                let element_array = element_arg.to_array(num_rows)?;
261                let max_array = max_arg.to_array(num_rows)?;
262                let arr_n = as_int64_array(&max_array)?.iter().collect::<Vec<_>>();
263                let result = array_remove_internal(&list_array, &element_array, &arr_n)?;
264                Ok(ColumnarValue::Array(result))
265            }
266        }
267    }
268
269    fn aliases(&self) -> &[String] {
270        &self.aliases
271    }
272
273    fn documentation(&self) -> Option<&Documentation> {
274        self.doc()
275    }
276}
277
278make_udf_expr_and_func!(
279    ArrayRemoveAll,
280    array_remove_all,
281    array element,
282    "removes all elements from the array equal to the given value. NULL elements already in the array are preserved when removing a non-NULL value. If `element` evaluates to NULL, the result is NULL rather than removing NULL entries.",
283    array_remove_all_udf
284);
285
286#[user_doc(
287    doc_section(label = "Array Functions"),
288    description = "Removes all elements from the array equal to the given value. NULL elements already in the array are preserved when removing a non-NULL value. If `element` evaluates to NULL, the result is NULL rather than removing NULL entries.",
289    syntax_example = "array_remove_all(array, element)",
290    sql_example = r#"```sql
291> select array_remove_all([1, 2, 2, 3, 2, 1, 4], 2);
292+--------------------------------------------------+
293| array_remove_all(List([1,2,2,3,2,1,4]),Int64(2)) |
294+--------------------------------------------------+
295| [1, 3, 1, 4]                                     |
296+--------------------------------------------------+
297
298> select array_remove_all([1, 2, NULL, 2, 4], 2);
299+-----------------------------------------------------+
300| array_remove_all(List([1,2,NULL,2,4]),Int64(2)) |
301+-----------------------------------------------------+
302| [1, NULL, 4]                                     |
303+-----------------------------------------------------+
304```"#,
305    argument(
306        name = "array",
307        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
308    ),
309    argument(
310        name = "element",
311        description = "Element to be removed from the array."
312    )
313)]
314#[derive(Debug, PartialEq, Eq, Hash)]
315pub struct ArrayRemoveAll {
316    signature: Signature,
317    aliases: Vec<String>,
318}
319
320impl Default for ArrayRemoveAll {
321    fn default() -> Self {
322        Self::new()
323    }
324}
325
326impl ArrayRemoveAll {
327    pub fn new() -> Self {
328        Self {
329            signature: Signature::array_and_element(Volatility::Immutable),
330            aliases: vec!["list_remove_all".to_string()],
331        }
332    }
333}
334
335impl ScalarUDFImpl for ArrayRemoveAll {
336    fn name(&self) -> &str {
337        "array_remove_all"
338    }
339
340    fn signature(&self) -> &Signature {
341        &self.signature
342    }
343
344    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
345        internal_err!("return_field_from_args should be used instead")
346    }
347
348    fn return_field_from_args(
349        &self,
350        args: datafusion_expr::ReturnFieldArgs,
351    ) -> Result<FieldRef> {
352        let array_field = args.arg_fields[0].as_ref().clone();
353        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
354        Ok(Arc::new(array_field.with_nullable(nullable)))
355    }
356
357    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
358        let [list_arg, element_arg] = take_function_args(self.name(), &args.args)?;
359        let num_rows = args.number_rows;
360        let list_array = list_arg.to_array(num_rows)?;
361        match element_arg {
362            ColumnarValue::Scalar(scalar_element)
363                if !scalar_element.is_null()
364                    && !scalar_element.data_type().is_nested() =>
365            {
366                let result =
367                    array_remove_with_scalar_args(&list_array, scalar_element, i64::MAX)?;
368                Ok(ColumnarValue::Array(result))
369            }
370            element_arg => {
371                let element_array = element_arg.to_array(num_rows)?;
372                let result = array_remove_internal(
373                    &list_array,
374                    &element_array,
375                    &[Some(i64::MAX)],
376                )?;
377                Ok(ColumnarValue::Array(result))
378            }
379        }
380    }
381
382    fn aliases(&self) -> &[String] {
383        &self.aliases
384    }
385
386    fn documentation(&self) -> Option<&Documentation> {
387        self.doc()
388    }
389}
390
391fn array_remove_internal(
392    array: &ArrayRef,
393    element_array: &ArrayRef,
394    arr_n: &[Option<i64>],
395) -> Result<ArrayRef> {
396    match array.data_type() {
397        DataType::List(_) => {
398            let list_array = array.as_list::<i32>();
399            general_remove::<i32>(list_array, element_array, arr_n)
400        }
401        DataType::LargeList(_) => {
402            let list_array = array.as_list::<i64>();
403            general_remove::<i64>(list_array, element_array, arr_n)
404        }
405        DataType::Null => Ok(new_null_array(array.data_type(), array.len())),
406        array_type => {
407            exec_err!("array_remove_all does not support type '{array_type}'.")
408        }
409    }
410}
411
412/// Fast path for `array_remove` when the needle is a non-null, non-nested scalar.
413/// Dispatches to the bulk `not_distinct` comparison kernel.
414fn array_remove_with_scalar_args(
415    array: &ArrayRef,
416    scalar_needle: &ScalarValue,
417    max_removals: i64,
418) -> Result<ArrayRef> {
419    match array.data_type() {
420        DataType::List(_) => {
421            let list_array = array.as_list::<i32>();
422            general_remove_with_scalar::<i32>(list_array, scalar_needle, max_removals)
423        }
424        DataType::LargeList(_) => {
425            let list_array = array.as_list::<i64>();
426            general_remove_with_scalar::<i64>(list_array, scalar_needle, max_removals)
427        }
428        DataType::Null => Ok(new_null_array(array.data_type(), array.len())),
429        array_type => exec_err!(
430            "array_remove/array_remove_n/array_remove_all does not support type '{array_type}'."
431        ),
432    }
433}
434
435/// For each element of `list_array[i]`, removed up to `arr_n[i]` occurrences
436/// of `element_array[i]`.
437///
438/// The type of each **element** in `list_array` must be the same as the type of
439/// `element_array`. This function also handles nested arrays
440/// ([`arrow::array::ListArray`] of [`arrow::array::ListArray`]s)
441///
442/// For example, when called to remove a list array (where each element is a
443/// list of int32s, the second argument are int32 arrays, and the
444/// third argument is the number of occurrences to remove
445///
446/// ```text
447/// general_remove(
448///   [1, 2, 3, 2], 2, 1    ==> [1, 3, 2]   (only the first 2 is removed)
449///   [4, 5, 6, 5], 5, 2    ==> [4, 6]  (both 5s are removed)
450/// )
451/// ```
452fn general_remove<OffsetSize: OffsetSizeTrait>(
453    list_array: &GenericListArray<OffsetSize>,
454    element_array: &ArrayRef,
455    arr_n: &[Option<i64>],
456) -> Result<ArrayRef> {
457    let list_field = match list_array.data_type() {
458        DataType::List(field) | DataType::LargeList(field) => field,
459        _ => {
460            return exec_err!(
461                "Expected List or LargeList data type, got {:?}",
462                list_array.data_type()
463            );
464        }
465    };
466    let original_data = list_array.values().to_data();
467    // Build up the offsets for the final output array
468    let mut offsets = Vec::<OffsetSize>::with_capacity(list_array.len() + 1);
469    offsets.push(OffsetSize::zero());
470
471    let mut mutable = MutableArrayData::with_capacities(
472        vec![&original_data],
473        false,
474        Capacities::Array(original_data.len()),
475    );
476    let mut valid = NullBufferBuilder::new(list_array.len());
477
478    for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() {
479        if list_array.is_null(row_index) || element_array.is_null(row_index) {
480            offsets.push(offsets[row_index]);
481            valid.append_null();
482            continue;
483        }
484
485        let n = if arr_n.len() == 1 {
486            arr_n[0]
487        } else {
488            arr_n[row_index]
489        };
490        let Some(n) = n else {
491            offsets.push(offsets[row_index]);
492            valid.append_null();
493            continue;
494        };
495
496        let start = offset_window[0].to_usize().unwrap();
497        let end = offset_window[1].to_usize().unwrap();
498
499        // compare each element in the list, `false` means the element matches and should be removed
500        let eq_array = utils::compare_element_to_list(
501            &list_array.value(row_index),
502            element_array,
503            row_index,
504            false,
505        )?;
506
507        let num_to_remove = eq_array.false_count();
508
509        // Fast path: no elements to remove, copy entire row
510        if num_to_remove == 0 {
511            mutable.try_extend(0, start, end)?;
512            offsets.push(offsets[row_index] + OffsetSize::usize_as(end - start));
513            valid.append_non_null();
514            continue;
515        }
516
517        // Remove at most `n` matching elements
518        let max_removals = n.min(num_to_remove as i64);
519        let mut removed = 0i64;
520        let mut copied = 0usize;
521        // marks the beginning of a range of elements pending to be copied.
522        let mut pending_batch_to_retain: Option<usize> = None;
523        for (i, keep) in eq_array.iter().enumerate() {
524            if keep == Some(false) && removed < max_removals {
525                // Flush pending batch before skipping this element
526                if let Some(bs) = pending_batch_to_retain {
527                    mutable.try_extend(0, start + bs, start + i)?;
528                    copied += i - bs;
529                    pending_batch_to_retain = None;
530                }
531                removed += 1;
532            } else if pending_batch_to_retain.is_none() {
533                pending_batch_to_retain = Some(i);
534            }
535        }
536
537        // Flush remaining batch
538        if let Some(bs) = pending_batch_to_retain {
539            mutable.try_extend(0, start + bs, start + eq_array.len())?;
540            copied += eq_array.len() - bs;
541        }
542
543        offsets.push(offsets[row_index] + OffsetSize::usize_as(copied));
544        valid.append_non_null();
545    }
546
547    let new_values = make_array(mutable.freeze());
548    Ok(Arc::new(GenericListArray::<OffsetSize>::try_new(
549        Arc::clone(list_field),
550        OffsetBuffer::new(offsets.into()),
551        new_values,
552        valid.finish(),
553    )?))
554}
555
556/// For each element of `list_array[i]`, removes up to `max_removals` occurrences
557/// of the scalar needle.
558///
559/// This is a specialized version of `general_remove` for scalar elements that
560/// uses bulk comparison for better performance.
561fn general_remove_with_scalar<OffsetSize: OffsetSizeTrait>(
562    list_array: &GenericListArray<OffsetSize>,
563    scalar_needle: &ScalarValue,
564    max_removals: i64,
565) -> Result<ArrayRef> {
566    if max_removals <= 0 {
567        return Ok(Arc::new(list_array.clone()));
568    }
569
570    let list_field = match list_array.data_type() {
571        DataType::List(field) | DataType::LargeList(field) => field,
572        _ => {
573            return exec_err!(
574                "Expected List or LargeList data type, got {:?}",
575                list_array.data_type()
576            );
577        }
578    };
579
580    let list_offsets = list_array.offsets();
581    let first_offset = list_offsets[0].to_usize().unwrap();
582    let last_offset = list_offsets[list_offsets.len() - 1].to_usize().unwrap();
583    let values_range_len = last_offset - first_offset;
584    let values_slice = list_array.values().slice(first_offset, values_range_len);
585    let original_data = values_slice.to_data();
586    let mut offsets = Vec::<OffsetSize>::with_capacity(list_array.len() + 1);
587    offsets.push(OffsetSize::zero());
588
589    let mut mutable = MutableArrayData::with_capacities(
590        vec![&original_data],
591        false,
592        Capacities::Array(original_data.len()),
593    );
594    let nulls = list_array.nulls().cloned();
595    let needle = scalar_needle.to_array_of_size(1)?;
596    let remove_mask = arrow_ord::cmp::not_distinct(&values_slice, &Scalar::new(needle))?;
597    let remove_bits = remove_mask.values();
598
599    for (row_index, offset_window) in list_offsets.windows(2).enumerate() {
600        if nulls.as_ref().is_some_and(|nulls| nulls.is_null(row_index)) {
601            offsets.push(offsets[row_index]);
602            continue;
603        }
604
605        let start = offset_window[0].to_usize().unwrap() - first_offset;
606        let end = offset_window[1].to_usize().unwrap() - first_offset;
607        let row_len = end - start;
608
609        let row_remove_bits = remove_bits.slice(start, row_len);
610        let num_to_remove = row_remove_bits.count_set_bits();
611
612        if num_to_remove == 0 {
613            mutable.try_extend(0, start, end)?;
614            offsets.push(offsets[row_index] + OffsetSize::usize_as(row_len));
615            continue;
616        }
617
618        let removals_to_apply = max_removals.min(num_to_remove as i64) as usize;
619
620        // Iterate only over the removal positions via set_indices. This is
621        // efficient when the number of removals is small relative to the row
622        // length (common case), since it skips over retained elements.
623        let mut removed = 0usize;
624        let mut copied = 0usize;
625        let mut prev_end = start;
626        for remove_pos in row_remove_bits.set_indices() {
627            let abs_pos = start + remove_pos;
628            if abs_pos > prev_end {
629                mutable.try_extend(0, prev_end, abs_pos)?;
630                copied += abs_pos - prev_end;
631            }
632            prev_end = abs_pos + 1;
633            removed += 1;
634            if removed == removals_to_apply {
635                break;
636            }
637        }
638        // Copy the remaining tail after the last removal
639        if prev_end < end {
640            mutable.try_extend(0, prev_end, end)?;
641            copied += end - prev_end;
642        }
643
644        offsets.push(offsets[row_index] + OffsetSize::usize_as(copied));
645    }
646
647    let new_values = make_array(mutable.freeze());
648    Ok(Arc::new(GenericListArray::<OffsetSize>::try_new(
649        Arc::clone(list_field),
650        OffsetBuffer::new(offsets.into()),
651        new_values,
652        nulls,
653    )?))
654}
655
656#[cfg(test)]
657mod tests {
658    use crate::remove::{ArrayRemove, ArrayRemoveAll, ArrayRemoveN};
659    use arrow::array::{
660        Array, ArrayRef, AsArray, GenericListArray, Int32Array, Int64Array, ListArray,
661        OffsetSizeTrait,
662    };
663    use arrow::buffer::{NullBuffer, ScalarBuffer};
664    use arrow::datatypes::{DataType, Field, Int32Type};
665    use datafusion_common::ScalarValue;
666    use datafusion_expr::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl};
667    use datafusion_expr_common::columnar_value::ColumnarValue;
668    use std::ops::Deref;
669    use std::sync::Arc;
670
671    #[test]
672    fn test_array_remove_nullability() {
673        for nullability in [true, false] {
674            for item_nullability in [true, false] {
675                for element_nullability in [true, false] {
676                    let input_field = Arc::new(Field::new(
677                        "num",
678                        DataType::new_list(DataType::Int32, item_nullability),
679                        nullability,
680                    ));
681                    let args_fields = vec![
682                        Arc::clone(&input_field),
683                        Arc::new(Field::new("a", DataType::Int32, element_nullability)),
684                    ];
685                    let scalar_args = vec![None, Some(&ScalarValue::Int32(Some(1)))];
686
687                    let result = ArrayRemove::new()
688                        .return_field_from_args(ReturnFieldArgs {
689                            arg_fields: &args_fields,
690                            scalar_arguments: &scalar_args,
691                        })
692                        .unwrap();
693
694                    let expected = Arc::new(
695                        input_field
696                            .as_ref()
697                            .clone()
698                            .with_nullable(nullability || element_nullability),
699                    );
700
701                    assert_eq!(result, expected);
702                }
703            }
704        }
705    }
706
707    #[test]
708    fn test_array_remove_n_nullability() {
709        for nullability in [true, false] {
710            for item_nullability in [true, false] {
711                for element_nullability in [true, false] {
712                    for count_nullability in [true, false] {
713                        let input_field = Arc::new(Field::new(
714                            "num",
715                            DataType::new_list(DataType::Int32, item_nullability),
716                            nullability,
717                        ));
718                        let args_fields = vec![
719                            Arc::clone(&input_field),
720                            Arc::new(Field::new(
721                                "a",
722                                DataType::Int32,
723                                element_nullability,
724                            )),
725                            Arc::new(Field::new("b", DataType::Int64, count_nullability)),
726                        ];
727                        let scalar_args = vec![
728                            None,
729                            Some(&ScalarValue::Int32(Some(1))),
730                            Some(&ScalarValue::Int64(Some(1))),
731                        ];
732
733                        let result = ArrayRemoveN::new()
734                            .return_field_from_args(ReturnFieldArgs {
735                                arg_fields: &args_fields,
736                                scalar_arguments: &scalar_args,
737                            })
738                            .unwrap();
739
740                        let expected_nullable =
741                            nullability || element_nullability || count_nullability;
742                        let expected = Arc::new(
743                            input_field
744                                .as_ref()
745                                .clone()
746                                .with_nullable(expected_nullable),
747                        );
748
749                        assert_eq!(result, expected);
750                    }
751                }
752            }
753        }
754    }
755
756    #[test]
757    fn test_array_remove_all_nullability() {
758        for nullability in [true, false] {
759            for item_nullability in [true, false] {
760                for element_nullability in [true, false] {
761                    let input_field = Arc::new(Field::new(
762                        "num",
763                        DataType::new_list(DataType::Int32, item_nullability),
764                        nullability,
765                    ));
766                    let args_fields = vec![
767                        Arc::clone(&input_field),
768                        Arc::new(Field::new("a", DataType::Int32, element_nullability)),
769                    ];
770                    let scalar_args = vec![None, Some(&ScalarValue::Int32(Some(1)))];
771                    let result = ArrayRemoveAll::new()
772                        .return_field_from_args(ReturnFieldArgs {
773                            arg_fields: &args_fields,
774                            scalar_arguments: &scalar_args,
775                        })
776                        .unwrap();
777
778                    let expected = Arc::new(
779                        input_field
780                            .as_ref()
781                            .clone()
782                            .with_nullable(nullability || element_nullability),
783                    );
784
785                    assert_eq!(result, expected);
786                }
787            }
788        }
789    }
790
791    fn ensure_field_nullability<O: OffsetSizeTrait>(
792        field_nullable: bool,
793        list: GenericListArray<O>,
794    ) -> GenericListArray<O> {
795        let (field, offsets, values, nulls) = list.into_parts();
796
797        if field.is_nullable() == field_nullable {
798            return GenericListArray::new(field, offsets, values, nulls);
799        }
800        if !field_nullable {
801            assert_eq!(nulls, None);
802        }
803
804        let field = Arc::new(field.deref().clone().with_nullable(field_nullable));
805
806        GenericListArray::new(field, offsets, values, nulls)
807    }
808
809    #[test]
810    fn test_array_remove_non_nullable() {
811        let input_list = Arc::new(ensure_field_nullability(
812            false,
813            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
814                Some(([1, 2, 2, 3, 2, 1, 4]).iter().copied().map(Some)),
815                Some(([42, 2, 55, 63, 2]).iter().copied().map(Some)),
816            ]),
817        ));
818        let expected_list = ensure_field_nullability(
819            false,
820            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
821                Some(([1, 2, 3, 2, 1, 4]).iter().copied().map(Some)),
822                Some(([42, 55, 63, 2]).iter().copied().map(Some)),
823            ]),
824        );
825
826        let element_to_remove = ScalarValue::Int32(Some(2));
827
828        assert_array_remove(input_list, expected_list, element_to_remove);
829    }
830
831    #[test]
832    fn test_array_remove_nullable() {
833        let input_list = Arc::new(ensure_field_nullability(
834            true,
835            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
836                Some(vec![
837                    Some(1),
838                    Some(2),
839                    Some(2),
840                    Some(3),
841                    None,
842                    Some(1),
843                    Some(4),
844                ]),
845                Some(vec![Some(42), Some(2), None, Some(63), Some(2)]),
846            ]),
847        ));
848        let expected_list = ensure_field_nullability(
849            true,
850            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
851                Some(vec![Some(1), Some(2), Some(3), None, Some(1), Some(4)]),
852                Some(vec![Some(42), None, Some(63), Some(2)]),
853            ]),
854        );
855
856        let element_to_remove = ScalarValue::Int32(Some(2));
857
858        assert_array_remove(input_list, expected_list, element_to_remove);
859    }
860
861    fn assert_array_remove(
862        input_list: ArrayRef,
863        expected_list: GenericListArray<i32>,
864        element_to_remove: ScalarValue,
865    ) {
866        assert_eq!(input_list.data_type(), expected_list.data_type());
867        assert_eq!(expected_list.value_type(), element_to_remove.data_type());
868        let input_list_len = input_list.len();
869        let input_list_data_type = input_list.data_type().clone();
870
871        let udf = ArrayRemove::new();
872        let args_fields = vec![
873            Arc::new(Field::new("num", input_list.data_type().clone(), false)),
874            Arc::new(Field::new(
875                "el",
876                element_to_remove.data_type(),
877                element_to_remove.is_null(),
878            )),
879        ];
880        let scalar_args = vec![None, Some(&element_to_remove)];
881
882        let return_field = udf
883            .return_field_from_args(ReturnFieldArgs {
884                arg_fields: &args_fields,
885                scalar_arguments: &scalar_args,
886            })
887            .unwrap();
888
889        let result = udf
890            .invoke_with_args(ScalarFunctionArgs {
891                args: vec![
892                    ColumnarValue::Array(input_list),
893                    ColumnarValue::Scalar(element_to_remove),
894                ],
895                arg_fields: args_fields,
896                number_rows: input_list_len,
897                return_field,
898                config_options: Arc::new(Default::default()),
899            })
900            .unwrap();
901
902        assert_eq!(result.data_type(), input_list_data_type);
903        match result {
904            ColumnarValue::Array(array) => {
905                let result_list = array.as_list::<i32>();
906                assert_eq!(result_list, &expected_list);
907            }
908            _ => panic!("Expected ColumnarValue::Array"),
909        }
910    }
911
912    #[test]
913    fn test_array_remove_n_non_nullable() {
914        let input_list = Arc::new(ensure_field_nullability(
915            false,
916            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
917                Some(([1, 2, 2, 3, 2, 1, 4]).iter().copied().map(Some)),
918                Some(([42, 2, 55, 63, 2]).iter().copied().map(Some)),
919            ]),
920        ));
921        let expected_list = ensure_field_nullability(
922            false,
923            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
924                Some(([1, 3, 2, 1, 4]).iter().copied().map(Some)),
925                Some(([42, 55, 63]).iter().copied().map(Some)),
926            ]),
927        );
928
929        let element_to_remove = ScalarValue::Int32(Some(2));
930
931        assert_array_remove_n(input_list, expected_list, element_to_remove, 2);
932    }
933
934    #[test]
935    fn test_array_remove_n_nullable() {
936        let input_list = Arc::new(ensure_field_nullability(
937            true,
938            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
939                Some(vec![
940                    Some(1),
941                    Some(2),
942                    Some(2),
943                    Some(3),
944                    None,
945                    Some(1),
946                    Some(4),
947                ]),
948                Some(vec![Some(42), Some(2), None, Some(63), Some(2)]),
949            ]),
950        ));
951        let expected_list = ensure_field_nullability(
952            true,
953            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
954                Some(vec![Some(1), Some(3), None, Some(1), Some(4)]),
955                Some(vec![Some(42), None, Some(63)]),
956            ]),
957        );
958
959        let element_to_remove = ScalarValue::Int32(Some(2));
960
961        assert_array_remove_n(input_list, expected_list, element_to_remove, 2);
962    }
963
964    #[test]
965    fn test_array_remove_n_null_count_returns_null() {
966        let array: ArrayRef =
967            Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
968                Some(vec![Some(1), Some(2), Some(2)]),
969                Some(vec![Some(4), Some(2)]),
970            ]));
971        let element: ArrayRef = Arc::new(Int32Array::from(vec![2, 2]));
972        let max: ArrayRef = Arc::new(Int64Array::new(
973            ScalarBuffer::from(vec![1, 1]),
974            Some(NullBuffer::from(vec![true, false])),
975        ));
976
977        let udf = ArrayRemoveN::new();
978        let args_fields = vec![
979            Arc::new(Field::new("num", array.data_type().clone(), false)),
980            Arc::new(Field::new("el", DataType::Int32, false)),
981            Arc::new(Field::new("count", DataType::Int64, true)),
982        ];
983        let scalar_args = vec![None, None, None];
984        let return_field = udf
985            .return_field_from_args(ReturnFieldArgs {
986                arg_fields: &args_fields,
987                scalar_arguments: &scalar_args,
988            })
989            .unwrap();
990        let result = udf
991            .invoke_with_args(ScalarFunctionArgs {
992                args: vec![
993                    ColumnarValue::Array(array),
994                    ColumnarValue::Array(element),
995                    ColumnarValue::Array(max),
996                ],
997                arg_fields: args_fields,
998                number_rows: 2,
999                return_field,
1000                config_options: Arc::new(Default::default()),
1001            })
1002            .unwrap();
1003        let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1004            Some(vec![Some(1), Some(2)]),
1005            None,
1006        ]);
1007
1008        match result {
1009            ColumnarValue::Array(array) => {
1010                assert_eq!(array.as_list::<i32>(), &expected);
1011            }
1012            _ => panic!("Expected ColumnarValue::Array"),
1013        }
1014    }
1015
1016    fn assert_array_remove_n(
1017        input_list: ArrayRef,
1018        expected_list: GenericListArray<i32>,
1019        element_to_remove: ScalarValue,
1020        n: i64,
1021    ) {
1022        assert_eq!(input_list.data_type(), expected_list.data_type());
1023        assert_eq!(expected_list.value_type(), element_to_remove.data_type());
1024        let input_list_len = input_list.len();
1025        let input_list_data_type = input_list.data_type().clone();
1026
1027        let count_scalar = ScalarValue::Int64(Some(n));
1028
1029        let udf = ArrayRemoveN::new();
1030        let args_fields = vec![
1031            Arc::new(Field::new("num", input_list.data_type().clone(), false)),
1032            Arc::new(Field::new(
1033                "el",
1034                element_to_remove.data_type(),
1035                element_to_remove.is_null(),
1036            )),
1037            Arc::new(Field::new("count", DataType::Int64, false)),
1038        ];
1039        let scalar_args = vec![None, Some(&element_to_remove), Some(&count_scalar)];
1040
1041        let return_field = udf
1042            .return_field_from_args(ReturnFieldArgs {
1043                arg_fields: &args_fields,
1044                scalar_arguments: &scalar_args,
1045            })
1046            .unwrap();
1047
1048        let result = udf
1049            .invoke_with_args(ScalarFunctionArgs {
1050                args: vec![
1051                    ColumnarValue::Array(input_list),
1052                    ColumnarValue::Scalar(element_to_remove),
1053                    ColumnarValue::Scalar(count_scalar),
1054                ],
1055                arg_fields: args_fields,
1056                number_rows: input_list_len,
1057                return_field,
1058                config_options: Arc::new(Default::default()),
1059            })
1060            .unwrap();
1061
1062        assert_eq!(result.data_type(), input_list_data_type);
1063        match result {
1064            ColumnarValue::Array(array) => {
1065                let result_list = array.as_list::<i32>();
1066                assert_eq!(result_list, &expected_list);
1067            }
1068            _ => panic!("Expected ColumnarValue::Array"),
1069        }
1070    }
1071
1072    #[test]
1073    fn test_array_remove_all_non_nullable() {
1074        let input_list = Arc::new(ensure_field_nullability(
1075            false,
1076            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1077                Some(([1, 2, 2, 3, 2, 1, 4]).iter().copied().map(Some)),
1078                Some(([42, 2, 55, 63, 2]).iter().copied().map(Some)),
1079            ]),
1080        ));
1081        let expected_list = ensure_field_nullability(
1082            false,
1083            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1084                Some(([1, 3, 1, 4]).iter().copied().map(Some)),
1085                Some(([42, 55, 63]).iter().copied().map(Some)),
1086            ]),
1087        );
1088
1089        let element_to_remove = ScalarValue::Int32(Some(2));
1090
1091        assert_array_remove_all(input_list, expected_list, element_to_remove);
1092    }
1093
1094    #[test]
1095    fn test_array_remove_all_nullable() {
1096        let input_list = Arc::new(ensure_field_nullability(
1097            true,
1098            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1099                Some(vec![
1100                    Some(1),
1101                    Some(2),
1102                    Some(2),
1103                    Some(3),
1104                    None,
1105                    Some(1),
1106                    Some(4),
1107                ]),
1108                Some(vec![Some(42), Some(2), None, Some(63), Some(2)]),
1109            ]),
1110        ));
1111        let expected_list = ensure_field_nullability(
1112            true,
1113            ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1114                Some(vec![Some(1), Some(3), None, Some(1), Some(4)]),
1115                Some(vec![Some(42), None, Some(63)]),
1116            ]),
1117        );
1118
1119        let element_to_remove = ScalarValue::Int32(Some(2));
1120
1121        assert_array_remove_all(input_list, expected_list, element_to_remove);
1122    }
1123
1124    fn assert_array_remove_all(
1125        input_list: ArrayRef,
1126        expected_list: GenericListArray<i32>,
1127        element_to_remove: ScalarValue,
1128    ) {
1129        assert_eq!(input_list.data_type(), expected_list.data_type());
1130        assert_eq!(expected_list.value_type(), element_to_remove.data_type());
1131        let input_list_len = input_list.len();
1132        let input_list_data_type = input_list.data_type().clone();
1133
1134        let udf = ArrayRemoveAll::new();
1135        let args_fields = vec![
1136            Arc::new(Field::new("num", input_list.data_type().clone(), false)),
1137            Arc::new(Field::new(
1138                "el",
1139                element_to_remove.data_type(),
1140                element_to_remove.is_null(),
1141            )),
1142        ];
1143        let scalar_args = vec![None, Some(&element_to_remove)];
1144
1145        let return_field = udf
1146            .return_field_from_args(ReturnFieldArgs {
1147                arg_fields: &args_fields,
1148                scalar_arguments: &scalar_args,
1149            })
1150            .unwrap();
1151
1152        let result = udf
1153            .invoke_with_args(ScalarFunctionArgs {
1154                args: vec![
1155                    ColumnarValue::Array(input_list),
1156                    ColumnarValue::Scalar(element_to_remove),
1157                ],
1158                arg_fields: args_fields,
1159                number_rows: input_list_len,
1160                return_field,
1161                config_options: Arc::new(Default::default()),
1162            })
1163            .unwrap();
1164
1165        assert_eq!(result.data_type(), input_list_data_type);
1166        match result {
1167            ColumnarValue::Array(array) => {
1168                let result_list = array.as_list::<i32>();
1169                assert_eq!(result_list, &expected_list);
1170            }
1171            _ => panic!("Expected ColumnarValue::Array"),
1172        }
1173    }
1174}
datafusion_functions_nested/remove.rs

datafusion_functions_nested/
remove.rs
