datafusion_functions_nested/
replace.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_replace, array_replace_n and array_replace_all functions.
19
20use arrow::array::{
21    Array, ArrayRef, AsArray, Capacities, GenericListArray, MutableArrayData,
22    NullBufferBuilder, OffsetSizeTrait, Scalar, new_null_array,
23};
24use arrow::buffer::OffsetBuffer;
25use arrow::datatypes::{DataType, Field, FieldRef};
26use datafusion_common::cast::as_int64_array;
27use datafusion_common::utils::ListCoercion;
28use datafusion_common::{
29    Result, ScalarValue, exec_err, internal_err, utils::take_function_args,
30};
31use datafusion_expr::{
32    ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation,
33    ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
34    Volatility,
35};
36use datafusion_macros::user_doc;
37
38use crate::utils::{compare_element_to_list, list_inner_field, list_type_with_element};
39
40use std::sync::Arc;
41
42// Create static instances of ScalarUDFs for each function
43make_udf_expr_and_func!(ArrayReplace,
44    array_replace,
45    array from to,
46    "replaces the first occurrence of the specified element with another specified element.",
47    array_replace_udf
48);
49make_udf_expr_and_func!(ArrayReplaceN,
50    array_replace_n,
51    array from to max,
52    "replaces the first `max` occurrences of the specified element with another specified element.",
53    array_replace_n_udf
54);
55make_udf_expr_and_func!(ArrayReplaceAll,
56    array_replace_all,
57    array from to,
58    "replaces all occurrences of the specified element with another specified element.",
59    array_replace_all_udf
60);
61
62#[user_doc(
63    doc_section(label = "Array Functions"),
64    description = "Replaces the first occurrence of the specified element with another specified element.",
65    syntax_example = "array_replace(array, from, to)",
66    sql_example = r#"```sql
67> select array_replace([1, 2, 2, 3, 2, 1, 4], 2, 5);
68+--------------------------------------------------------+
69| array_replace(List([1,2,2,3,2,1,4]),Int64(2),Int64(5)) |
70+--------------------------------------------------------+
71| [1, 5, 2, 3, 2, 1, 4]                                  |
72+--------------------------------------------------------+
73```"#,
74    argument(
75        name = "array",
76        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
77    ),
78    argument(name = "from", description = "Initial element."),
79    argument(name = "to", description = "Final element.")
80)]
81#[derive(Debug, PartialEq, Eq, Hash)]
82pub struct ArrayReplace {
83    signature: Signature,
84    aliases: Vec<String>,
85}
86
87impl Default for ArrayReplace {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93impl ArrayReplace {
94    pub fn new() -> Self {
95        Self {
96            signature: Signature {
97                type_signature: TypeSignature::ArraySignature(
98                    ArrayFunctionSignature::Array {
99                        arguments: vec![
100                            ArrayFunctionArgument::Array,
101                            ArrayFunctionArgument::Element,
102                            ArrayFunctionArgument::Element,
103                        ],
104                        array_coercion: Some(ListCoercion::FixedSizedListToList),
105                    },
106                ),
107                volatility: Volatility::Immutable,
108                parameter_names: None,
109            },
110            aliases: vec![String::from("list_replace")],
111        }
112    }
113}
114
115impl ScalarUDFImpl for ArrayReplace {
116    fn name(&self) -> &str {
117        "array_replace"
118    }
119
120    fn signature(&self) -> &Signature {
121        &self.signature
122    }
123
124    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
125        internal_err!("return_field_from_args should be used instead")
126    }
127
128    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
129        replace_return_field(self.name(), args.arg_fields)
130    }
131
132    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
133        let return_type = args.return_field.data_type().clone();
134        let [list_arg, from_arg, to_arg] = take_function_args(self.name(), &args.args)?;
135        let num_rows = args.number_rows;
136        let list_array = list_arg.to_array(num_rows)?;
137        match (from_arg, to_arg) {
138            (ColumnarValue::Scalar(scalar_from), ColumnarValue::Scalar(scalar_to)) => {
139                let result = array_replace_with_scalar_args(
140                    self.name(),
141                    &list_array,
142                    scalar_from,
143                    scalar_to,
144                    1i64,
145                    &return_type,
146                )?;
147                Ok(ColumnarValue::Array(result))
148            }
149            (from_arg, to_arg) => {
150                let from_array = from_arg.to_array(num_rows)?;
151                let to_array = to_arg.to_array(num_rows)?;
152                let result = array_replace_internal(
153                    self.name(),
154                    &list_array,
155                    &from_array,
156                    &to_array,
157                    &[Some(1)],
158                    &return_type,
159                )?;
160                Ok(ColumnarValue::Array(result))
161            }
162        }
163    }
164
165    fn aliases(&self) -> &[String] {
166        &self.aliases
167    }
168
169    fn documentation(&self) -> Option<&Documentation> {
170        self.doc()
171    }
172}
173
174#[user_doc(
175    doc_section(label = "Array Functions"),
176    description = "Replaces the first `max` occurrences of the specified element with another specified element.",
177    syntax_example = "array_replace_n(array, from, to, max)",
178    sql_example = r#"```sql
179> select array_replace_n([1, 2, 2, 3, 2, 1, 4], 2, 5, 2);
180+-------------------------------------------------------------------+
181| array_replace_n(List([1,2,2,3,2,1,4]),Int64(2),Int64(5),Int64(2)) |
182+-------------------------------------------------------------------+
183| [1, 5, 5, 3, 2, 1, 4]                                             |
184+-------------------------------------------------------------------+
185```"#,
186    argument(
187        name = "array",
188        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
189    ),
190    argument(name = "from", description = "Initial element."),
191    argument(name = "to", description = "Final element."),
192    argument(name = "max", description = "Number of first occurrences to replace.")
193)]
194#[derive(Debug, PartialEq, Eq, Hash)]
195pub(super) struct ArrayReplaceN {
196    signature: Signature,
197    aliases: Vec<String>,
198}
199
200impl ArrayReplaceN {
201    pub fn new() -> Self {
202        Self {
203            signature: Signature {
204                type_signature: TypeSignature::ArraySignature(
205                    ArrayFunctionSignature::Array {
206                        arguments: vec![
207                            ArrayFunctionArgument::Array,
208                            ArrayFunctionArgument::Element,
209                            ArrayFunctionArgument::Element,
210                            ArrayFunctionArgument::Index,
211                        ],
212                        array_coercion: Some(ListCoercion::FixedSizedListToList),
213                    },
214                ),
215                volatility: Volatility::Immutable,
216                parameter_names: None,
217            },
218            aliases: vec![String::from("list_replace_n")],
219        }
220    }
221}
222
223impl ScalarUDFImpl for ArrayReplaceN {
224    fn name(&self) -> &str {
225        "array_replace_n"
226    }
227
228    fn signature(&self) -> &Signature {
229        &self.signature
230    }
231
232    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
233        internal_err!("return_field_from_args should be used instead")
234    }
235
236    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
237        replace_return_field(self.name(), args.arg_fields)
238    }
239
240    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
241        let return_type = args.return_field.data_type().clone();
242        let [list_arg, from_arg, to_arg, max_arg] =
243            take_function_args(self.name(), &args.args)?;
244        let num_rows = args.number_rows;
245        let list_array = list_arg.to_array(num_rows)?;
246        match (from_arg, to_arg, max_arg) {
247            (
248                ColumnarValue::Scalar(scalar_from),
249                ColumnarValue::Scalar(scalar_to),
250                ColumnarValue::Scalar(scalar_max),
251            ) => {
252                let ScalarValue::Int64(Some(n)) = scalar_max else {
253                    return Ok(ColumnarValue::Array(new_null_array(
254                        &return_type,
255                        num_rows,
256                    )));
257                };
258                let result = array_replace_with_scalar_args(
259                    self.name(),
260                    &list_array,
261                    scalar_from,
262                    scalar_to,
263                    *n,
264                    &return_type,
265                )?;
266                Ok(ColumnarValue::Array(result))
267            }
268            (from_arg, to_arg, max_arg) => {
269                let from_array = from_arg.to_array(num_rows)?;
270                let to_array = to_arg.to_array(num_rows)?;
271                let max_array = max_arg.to_array(num_rows)?;
272                let result = array_replace_n_inner(
273                    self.name(),
274                    &list_array,
275                    &from_array,
276                    &to_array,
277                    &max_array,
278                    &return_type,
279                )?;
280                Ok(ColumnarValue::Array(result))
281            }
282        }
283    }
284
285    fn aliases(&self) -> &[String] {
286        &self.aliases
287    }
288
289    fn documentation(&self) -> Option<&Documentation> {
290        self.doc()
291    }
292}
293
294#[user_doc(
295    doc_section(label = "Array Functions"),
296    description = "Replaces all occurrences of the specified element with another specified element.",
297    syntax_example = "array_replace_all(array, from, to)",
298    sql_example = r#"```sql
299> select array_replace_all([1, 2, 2, 3, 2, 1, 4], 2, 5);
300+------------------------------------------------------------+
301| array_replace_all(List([1,2,2,3,2,1,4]),Int64(2),Int64(5)) |
302+------------------------------------------------------------+
303| [1, 5, 5, 3, 5, 1, 4]                                      |
304+------------------------------------------------------------+
305```"#,
306    argument(
307        name = "array",
308        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
309    ),
310    argument(name = "from", description = "Initial element."),
311    argument(name = "to", description = "Final element.")
312)]
313#[derive(Debug, PartialEq, Eq, Hash)]
314pub(super) struct ArrayReplaceAll {
315    signature: Signature,
316    aliases: Vec<String>,
317}
318
319impl ArrayReplaceAll {
320    pub fn new() -> Self {
321        Self {
322            signature: Signature {
323                type_signature: TypeSignature::ArraySignature(
324                    ArrayFunctionSignature::Array {
325                        arguments: vec![
326                            ArrayFunctionArgument::Array,
327                            ArrayFunctionArgument::Element,
328                            ArrayFunctionArgument::Element,
329                        ],
330                        array_coercion: Some(ListCoercion::FixedSizedListToList),
331                    },
332                ),
333                volatility: Volatility::Immutable,
334                parameter_names: None,
335            },
336            aliases: vec![String::from("list_replace_all")],
337        }
338    }
339}
340
341impl ScalarUDFImpl for ArrayReplaceAll {
342    fn name(&self) -> &str {
343        "array_replace_all"
344    }
345
346    fn signature(&self) -> &Signature {
347        &self.signature
348    }
349
350    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
351        internal_err!("return_field_from_args should be used instead")
352    }
353
354    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
355        replace_return_field(self.name(), args.arg_fields)
356    }
357
358    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
359        let return_type = args.return_field.data_type().clone();
360        let [list_arg, from_arg, to_arg] = take_function_args(self.name(), &args.args)?;
361        let num_rows = args.number_rows;
362        let list_array = list_arg.to_array(num_rows)?;
363        match (from_arg, to_arg) {
364            (ColumnarValue::Scalar(scalar_from), ColumnarValue::Scalar(scalar_to)) => {
365                let result = array_replace_with_scalar_args(
366                    self.name(),
367                    &list_array,
368                    scalar_from,
369                    scalar_to,
370                    i64::MAX,
371                    &return_type,
372                )?;
373                Ok(ColumnarValue::Array(result))
374            }
375            (from_arg, to_arg) => {
376                let from_array = from_arg.to_array(num_rows)?;
377                let to_array = to_arg.to_array(num_rows)?;
378                let result = array_replace_internal(
379                    self.name(),
380                    &list_array,
381                    &from_array,
382                    &to_array,
383                    &[Some(i64::MAX)],
384                    &return_type,
385                )?;
386                Ok(ColumnarValue::Array(result))
387            }
388        }
389    }
390
391    fn aliases(&self) -> &[String] {
392        &self.aliases
393    }
394
395    fn documentation(&self) -> Option<&Documentation> {
396        self.doc()
397    }
398}
399
400/// Return field shared by `array_replace`, `array_replace_n` and
401/// `array_replace_all`: the input list type, except that its inner field is
402/// nullable whenever the replacement element may be null.
403fn replace_return_field(name: &str, arg_fields: &[FieldRef]) -> Result<FieldRef> {
404    // `array` is at index 0 and `to` at index 2 for all three functions.
405    // `from` never contributes values to the output, so `to` is the only
406    // argument besides `array` that can affect the output's type.
407    let [array_field, _from_field, to_field, ..] = arg_fields else {
408        return exec_err!(
409            "{name} expects at least 3 arguments, got {}",
410            arg_fields.len()
411        );
412    };
413    let data_type =
414        list_type_with_element(array_field.data_type(), to_field.is_nullable());
415    Ok(Arc::new(Field::new(name, data_type, true)))
416}
417
418/// For each element of `list_array[i]`, replaces up to `arr_n[i]`  occurrences
419/// of `from_array[i]`, `to_array[i]`.
420///
421/// The type of each **element** in `list_array` must be the same as the type of
422/// `from_array` and `to_array`. This function also handles nested arrays
423/// (\[`ListArray`\] of \[`ListArray`\]s)
424///
425/// For example, when called to replace a list array (where each element is a
426/// list of int32s, the second and third argument are int32 arrays, and the
427/// fourth argument is the number of occurrences to replace
428///
429/// ```text
430/// general_replace(
431///   [1, 2, 3, 2], 2, 10, 1    ==> [1, 10, 3, 2]   (only the first 2 is replaced)
432///   [4, 5, 6, 5], 5, 20, 2    ==> [4, 20, 6, 20]  (both 5s are replaced)
433/// )
434/// ```
435fn general_replace<O: OffsetSizeTrait>(
436    list_array: &GenericListArray<O>,
437    from_array: &ArrayRef,
438    to_array: &ArrayRef,
439    arr_n: &[Option<i64>],
440    field: FieldRef,
441) -> Result<ArrayRef> {
442    // Build up the offsets for the final output array
443    let mut offsets: Vec<O> = Vec::with_capacity(list_array.len() + 1);
444    offsets.push(O::usize_as(0));
445    let values = list_array.values();
446    let original_data = values.to_data();
447    let to_data = to_array.to_data();
448    let capacity = Capacities::Array(original_data.len());
449
450    // First array is the original array, second array is the element to replace with.
451    let mut mutable = MutableArrayData::with_capacities(
452        vec![&original_data, &to_data],
453        false,
454        capacity,
455    );
456
457    let mut valid = NullBufferBuilder::new(list_array.len());
458
459    for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() {
460        if list_array.is_null(row_index) {
461            offsets.push(offsets[row_index]);
462            valid.append_null();
463            continue;
464        }
465
466        let n = if arr_n.len() == 1 {
467            arr_n[0]
468        } else {
469            arr_n[row_index]
470        };
471        let Some(n) = n else {
472            offsets.push(offsets[row_index]);
473            valid.append_null();
474            continue;
475        };
476
477        let start = offset_window[0];
478        let end = offset_window[1];
479
480        let list_array_row = list_array.value(row_index);
481
482        // Compute all positions in list_row_array (that is itself an
483        // array) that are equal to `from_array_row`
484        let eq_array =
485            compare_element_to_list(&list_array_row, &from_array, row_index, true)?;
486
487        let original_idx = O::usize_as(0);
488        let replace_idx = O::usize_as(1);
489        let mut counter = 0;
490
491        // All elements are false, no need to replace, just copy original data
492        if n <= 0 || !eq_array.has_true() {
493            mutable.try_extend(
494                original_idx.to_usize().unwrap(),
495                start.to_usize().unwrap(),
496                end.to_usize().unwrap(),
497            )?;
498            offsets.push(offsets[row_index] + (end - start));
499            valid.append_non_null();
500            continue;
501        }
502
503        let mut pending_retain: Option<O> = None;
504        for (i, to_replace) in eq_array.iter().enumerate() {
505            let i = O::usize_as(i);
506            if to_replace == Some(true) && counter < n {
507                // Flush any pending retain run before emitting the replacement.
508                if let Some(rs) = pending_retain.take() {
509                    mutable.try_extend(
510                        original_idx.to_usize().unwrap(),
511                        (start + rs).to_usize().unwrap(),
512                        (start + i).to_usize().unwrap(),
513                    )?;
514                }
515                mutable.try_extend(
516                    replace_idx.to_usize().unwrap(),
517                    row_index,
518                    row_index + 1,
519                )?;
520                counter += 1;
521                if counter == n {
522                    // copy original data for any matches past n
523                    mutable.try_extend(
524                        original_idx.to_usize().unwrap(),
525                        (start + i).to_usize().unwrap() + 1,
526                        end.to_usize().unwrap(),
527                    )?;
528                    break;
529                }
530            } else if pending_retain.is_none() {
531                pending_retain = Some(i);
532            }
533        }
534
535        // Flush trailing retain run when we exited the loop without ever
536        // hitting `counter == n` (i.e. fewer than `n` matches in this row).
537        if counter < n
538            && let Some(rs) = pending_retain
539        {
540            mutable.try_extend(
541                original_idx.to_usize().unwrap(),
542                (start + rs).to_usize().unwrap(),
543                end.to_usize().unwrap(),
544            )?;
545        }
546
547        offsets.push(offsets[row_index] + (end - start));
548        valid.append_non_null();
549    }
550
551    let data = mutable.freeze();
552
553    Ok(Arc::new(GenericListArray::<O>::try_new(
554        field,
555        OffsetBuffer::<O>::new(offsets.into()),
556        arrow::array::make_array(data),
557        valid.finish(),
558    )?))
559}
560
561/// Replaces up to `max_replacements` occurrences of `needle` with the single
562/// element in `to_array` for each row in `list_array`.
563///
564/// This is a specialized fast path for the all-scalar case that uses a single
565/// bulk `not_distinct` comparison over only the visible values range, then
566/// iterates match positions via `set_indices` instead of scanning every bit.
567fn general_replace_with_scalar<O: OffsetSizeTrait>(
568    list_array: &GenericListArray<O>,
569    needle: &Scalar<ArrayRef>,
570    scalar_to: &ScalarValue,
571    max_replacements: i64,
572    field: FieldRef,
573) -> Result<ArrayRef> {
574    // No replacement needed, but the output still has to carry the promised
575    // field, which may be more nullable than the input's.
576    if max_replacements <= 0 {
577        return Ok(Arc::new(GenericListArray::<O>::try_new(
578            field,
579            list_array.offsets().clone(),
580            Arc::clone(list_array.values()),
581            list_array.nulls().cloned(),
582        )?));
583    }
584
585    let first_offset = list_array.offsets()[0].to_usize().unwrap();
586    let last_offset = list_array.offsets()[list_array.len()].to_usize().unwrap();
587    let visible_values = list_array
588        .values()
589        .slice(first_offset, last_offset - first_offset);
590
591    let to_array = scalar_to.to_array_of_size(1)?;
592    let original_data = visible_values.to_data();
593    let to_data = to_array.to_data();
594    let capacity = Capacities::Array(original_data.len());
595
596    let mut mutable = MutableArrayData::with_capacities(
597        vec![&original_data, &to_data],
598        false,
599        capacity,
600    );
601
602    let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
603    offsets.push(O::zero());
604
605    // Single bulk comparison over the visible values only.
606    let match_bitmap = arrow_ord::cmp::not_distinct(&visible_values, needle)?;
607    let match_bits = match_bitmap.values();
608
609    for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() {
610        // Offsets relative to visible_values (subtract first_offset).
611        let start = offset_window[0].to_usize().unwrap() - first_offset;
612        let end = offset_window[1].to_usize().unwrap() - first_offset;
613        let row_len = end - start;
614
615        if list_array.is_null(row_index) {
616            offsets.push(offsets[row_index]);
617            continue;
618        }
619
620        // Slice the match bits to this row and iterate only over true positions.
621        let row_bits = match_bits.slice(start, row_len);
622        let mut match_positions = row_bits
623            .set_indices()
624            .take(max_replacements as usize)
625            .peekable();
626        if match_positions.peek().is_none() {
627            mutable.try_extend(0, start, end)?;
628            offsets.push(offsets[row_index] + O::usize_as(row_len));
629            continue;
630        }
631
632        // Iterate only over the positions that match using set_indices,
633        // which is more efficient than scanning every bit because the number
634        // of matches is typically much smaller than the total array size.
635        let mut prev_end = 0usize;
636        for match_pos in match_positions {
637            // Retain elements before this match.
638            if match_pos > prev_end {
639                mutable.try_extend(0, start + prev_end, start + match_pos)?;
640            }
641            // Emit the replacement element.
642            mutable.try_extend(1, 0, 1)?;
643            prev_end = match_pos + 1;
644        }
645
646        // Copy remaining elements after the last replacement.
647        if prev_end < row_len {
648            mutable.try_extend(0, start + prev_end, end)?;
649        }
650
651        offsets.push(offsets[row_index] + O::usize_as(row_len));
652    }
653
654    let data = mutable.freeze();
655
656    Ok(Arc::new(GenericListArray::<O>::try_new(
657        field,
658        OffsetBuffer::new(offsets.into()),
659        arrow::array::make_array(data),
660        list_array.nulls().cloned(),
661    )?))
662}
663
664/// Fast path for `array_replace` when all arguments are scalars.
665///
666/// Uses a single bulk `not_distinct` comparison instead of per-row comparisons.
667fn array_replace_with_scalar_args(
668    name: &str,
669    list_array: &ArrayRef,
670    scalar_from: &ScalarValue,
671    scalar_to: &ScalarValue,
672    max_replacements: i64,
673    return_type: &DataType,
674) -> Result<ArrayRef> {
675    // `not_distinct` doesn't support nested types, fall back to the generic array path.
676    if scalar_from.data_type().is_nested() {
677        let num_rows = list_array.len();
678        let from_array = scalar_from.to_array_of_size(num_rows)?;
679        let to_array = scalar_to.to_array_of_size(num_rows)?;
680        return array_replace_internal(
681            name,
682            list_array,
683            &from_array,
684            &to_array,
685            &vec![Some(max_replacements); num_rows],
686            return_type,
687        );
688    }
689
690    let needle = Scalar::new(scalar_from.to_array_of_size(1)?);
691    match list_array.data_type() {
692        DataType::List(_) => general_replace_with_scalar::<i32>(
693            list_array.as_list::<i32>(),
694            &needle,
695            scalar_to,
696            max_replacements,
697            list_inner_field(name, return_type)?,
698        ),
699        DataType::LargeList(_) => general_replace_with_scalar::<i64>(
700            list_array.as_list::<i64>(),
701            &needle,
702            scalar_to,
703            max_replacements,
704            list_inner_field(name, return_type)?,
705        ),
706        DataType::Null => Ok(new_null_array(return_type, list_array.len())),
707        array_type => exec_err!("{name} does not support type '{array_type}'."),
708    }
709}
710
711fn array_replace_internal(
712    name: &str,
713    array: &ArrayRef,
714    from: &ArrayRef,
715    to: &ArrayRef,
716    arr_n: &[Option<i64>],
717    return_type: &DataType,
718) -> Result<ArrayRef> {
719    match array.data_type() {
720        DataType::List(_) => general_replace::<i32>(
721            array.as_list::<i32>(),
722            from,
723            to,
724            arr_n,
725            list_inner_field(name, return_type)?,
726        ),
727        DataType::LargeList(_) => general_replace::<i64>(
728            array.as_list::<i64>(),
729            from,
730            to,
731            arr_n,
732            list_inner_field(name, return_type)?,
733        ),
734        DataType::Null => Ok(new_null_array(return_type, array.len())),
735        array_type => exec_err!("{name} does not support type '{array_type}'."),
736    }
737}
738
739fn array_replace_n_inner(
740    name: &str,
741    array: &ArrayRef,
742    from: &ArrayRef,
743    to: &ArrayRef,
744    max: &ArrayRef,
745    return_type: &DataType,
746) -> Result<ArrayRef> {
747    let arr_n = as_int64_array(max)?.iter().collect::<Vec<_>>();
748    array_replace_internal(name, array, from, to, &arr_n, return_type)
749}
750
751#[cfg(test)]
752mod tests {
753    use super::{ArrayReplaceN, array_replace_n_inner};
754    use arrow::array::{ArrayRef, AsArray, Int32Array, Int64Array, ListArray};
755    use arrow::buffer::{NullBuffer, ScalarBuffer};
756    use arrow::datatypes::{DataType, Field, Int32Type};
757    use datafusion_common::{Result, ScalarValue, config::ConfigOptions};
758    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
759    use std::sync::Arc;
760
761    #[test]
762    fn test_array_replace_n_null_max_returns_null() -> Result<()> {
763        let array: ArrayRef =
764            Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
765                Some(vec![Some(1), Some(2), Some(3)]),
766                Some(vec![Some(4), Some(2)]),
767            ]));
768        let from: ArrayRef = Arc::new(Int32Array::from(vec![2, 2]));
769        let to: ArrayRef = Arc::new(Int32Array::from(vec![9, 9]));
770        let max: ArrayRef = Arc::new(Int64Array::new(
771            ScalarBuffer::from(vec![1, 1]),
772            Some(NullBuffer::from(vec![true, false])),
773        ));
774
775        let result = array_replace_n_inner(
776            "array_replace_n",
777            &array,
778            &from,
779            &to,
780            &max,
781            array.data_type(),
782        )?;
783        let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
784            Some(vec![Some(1), Some(9), Some(3)]),
785            None,
786        ]);
787
788        assert_eq!(result.as_list::<i32>(), &expected);
789
790        Ok(())
791    }
792
793    #[test]
794    fn test_array_replace_n_scalar_null_max_returns_null() -> Result<()> {
795        let array: ArrayRef =
796            Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
797                Some(vec![Some(1), Some(2), Some(3)]),
798                Some(vec![Some(4), Some(2)]),
799            ]));
800        let array_field = Arc::new(Field::new("array", array.data_type().clone(), true));
801
802        let result = ArrayReplaceN::new().invoke_with_args(ScalarFunctionArgs {
803            args: vec![
804                ColumnarValue::Array(Arc::clone(&array)),
805                ColumnarValue::Scalar(ScalarValue::Int32(Some(2))),
806                ColumnarValue::Scalar(ScalarValue::Int32(Some(9))),
807                ColumnarValue::Scalar(ScalarValue::Int64(None)),
808            ],
809            arg_fields: vec![
810                Arc::clone(&array_field),
811                Arc::new(Field::new("from", DataType::Int32, false)),
812                Arc::new(Field::new("to", DataType::Int32, false)),
813                Arc::new(Field::new("max", DataType::Int64, true)),
814            ],
815            number_rows: array.len(),
816            return_field: Arc::clone(&array_field),
817            config_options: Arc::new(ConfigOptions::default()),
818        })?;
819
820        let result = result.into_array(array.len())?;
821        let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
822            Option::<Vec<Option<i32>>>::None,
823            Option::<Vec<Option<i32>>>::None,
824        ]);
825
826        assert_eq!(result.as_list::<i32>(), &expected);
827
828        Ok(())
829    }
830}
datafusion_functions_nested/replace.rs

datafusion_functions_nested/
replace.rs
