datafusion_functions_nested/
concat.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_append`, `array_prepend` and `array_concat` functions.
19
20use std::sync::Arc;
21
22use crate::make_array::make_array_inner;
23use crate::utils::{
24    align_array_dimensions, check_datatypes, list_inner_field, list_type_with_element,
25    make_scalar_function,
26};
27use arrow::array::{
28    Array, ArrayData, ArrayRef, Capacities, GenericListArray, MutableArrayData,
29    OffsetSizeTrait,
30};
31use arrow::buffer::{NullBuffer, OffsetBuffer};
32use arrow::datatypes::{DataType, Field, FieldRef};
33use datafusion_common::Result;
34use datafusion_common::utils::{
35    ListCoercion, base_type, coerced_type_with_base_type_only,
36};
37use datafusion_common::{
38    cast::as_generic_list_array,
39    exec_err, internal_err, plan_err,
40    utils::{list_ndims, take_function_args},
41};
42use datafusion_expr::binary::type_union_resolution;
43use datafusion_expr::{
44    ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl,
45    Signature, Volatility,
46};
47use datafusion_macros::user_doc;
48use itertools::Itertools;
49
50make_udf_expr_and_func!(
51    ArrayAppend,
52    array_append,
53    array element,                                // arg name
54    "appends an element to the end of an array.", // doc
55    array_append_udf                              // internal function name
56);
57
58#[user_doc(
59    doc_section(label = "Array Functions"),
60    description = "Appends an element to the end of an array.",
61    syntax_example = "array_append(array, element)",
62    sql_example = r#"```sql
63> select array_append([1, 2, 3], 4);
64+--------------------------------------+
65| array_append(List([1,2,3]),Int64(4)) |
66+--------------------------------------+
67| [1, 2, 3, 4]                         |
68+--------------------------------------+
69```"#,
70    argument(
71        name = "array",
72        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
73    ),
74    argument(name = "element", description = "Element to append to the array.")
75)]
76#[derive(Debug, PartialEq, Eq, Hash)]
77pub struct ArrayAppend {
78    signature: Signature,
79    aliases: Vec<String>,
80}
81
82impl Default for ArrayAppend {
83    fn default() -> Self {
84        Self::new()
85    }
86}
87
88impl ArrayAppend {
89    pub fn new() -> Self {
90        Self {
91            signature: Signature::array_and_element(Volatility::Immutable),
92            aliases: vec![
93                String::from("list_append"),
94                String::from("array_push_back"),
95                String::from("list_push_back"),
96            ],
97        }
98    }
99}
100
101impl ScalarUDFImpl for ArrayAppend {
102    fn name(&self) -> &str {
103        "array_append"
104    }
105
106    fn signature(&self) -> &Signature {
107        &self.signature
108    }
109
110    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
111        internal_err!("return_field_from_args should be used instead")
112    }
113
114    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
115        let [array_field, element_field] =
116            take_function_args(self.name(), args.arg_fields)?;
117        let data_type = append_prepend_return_type(
118            array_field.data_type(),
119            element_field.data_type(),
120            element_field.is_nullable(),
121        );
122        Ok(Arc::new(Field::new(self.name(), data_type, true)))
123    }
124
125    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
126        let return_type = args.return_field.data_type().clone();
127        make_scalar_function(|args: &[ArrayRef]| array_append_inner(args, &return_type))(
128            &args.args,
129        )
130    }
131
132    fn aliases(&self) -> &[String] {
133        &self.aliases
134    }
135
136    fn documentation(&self) -> Option<&Documentation> {
137        self.doc()
138    }
139}
140
141make_udf_expr_and_func!(
142    ArrayPrepend,
143    array_prepend,
144    element array,
145    "Prepends an element to the beginning of an array.",
146    array_prepend_udf
147);
148
149#[user_doc(
150    doc_section(label = "Array Functions"),
151    description = "Prepends an element to the beginning of an array.",
152    syntax_example = "array_prepend(element, array)",
153    sql_example = r#"```sql
154> select array_prepend(1, [2, 3, 4]);
155+---------------------------------------+
156| array_prepend(Int64(1),List([2,3,4])) |
157+---------------------------------------+
158| [1, 2, 3, 4]                          |
159+---------------------------------------+
160```"#,
161    argument(
162        name = "array",
163        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
164    ),
165    argument(name = "element", description = "Element to prepend to the array.")
166)]
167#[derive(Debug, PartialEq, Eq, Hash)]
168pub struct ArrayPrepend {
169    signature: Signature,
170    aliases: Vec<String>,
171}
172
173impl Default for ArrayPrepend {
174    fn default() -> Self {
175        Self::new()
176    }
177}
178
179impl ArrayPrepend {
180    pub fn new() -> Self {
181        Self {
182            signature: Signature::element_and_array(Volatility::Immutable),
183            aliases: vec![
184                String::from("list_prepend"),
185                String::from("array_push_front"),
186                String::from("list_push_front"),
187            ],
188        }
189    }
190}
191
192impl ScalarUDFImpl for ArrayPrepend {
193    fn name(&self) -> &str {
194        "array_prepend"
195    }
196
197    fn signature(&self) -> &Signature {
198        &self.signature
199    }
200
201    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
202        internal_err!("return_field_from_args should be used instead")
203    }
204
205    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
206        let [element_field, array_field] =
207            take_function_args(self.name(), args.arg_fields)?;
208        let data_type = append_prepend_return_type(
209            array_field.data_type(),
210            element_field.data_type(),
211            element_field.is_nullable(),
212        );
213        Ok(Arc::new(Field::new(self.name(), data_type, true)))
214    }
215
216    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
217        let return_type = args.return_field.data_type().clone();
218        make_scalar_function(|args: &[ArrayRef]| array_prepend_inner(args, &return_type))(
219            &args.args,
220        )
221    }
222
223    fn aliases(&self) -> &[String] {
224        &self.aliases
225    }
226
227    fn documentation(&self) -> Option<&Documentation> {
228        self.doc()
229    }
230}
231
232make_udf_expr_and_func!(
233    ArrayConcat,
234    array_concat,
235    "Concatenates arrays.",
236    array_concat_udf
237);
238
239#[user_doc(
240    doc_section(label = "Array Functions"),
241    description = "Concatenates arrays.",
242    syntax_example = "array_concat(array[, ..., array_n])",
243    sql_example = r#"```sql
244> select array_concat([1, 2], [3, 4], [5, 6]);
245+---------------------------------------------------+
246| array_concat(List([1,2]),List([3,4]),List([5,6])) |
247+---------------------------------------------------+
248| [1, 2, 3, 4, 5, 6]                                |
249+---------------------------------------------------+
250```"#,
251    argument(
252        name = "array",
253        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
254    ),
255    argument(
256        name = "array_n",
257        description = "Subsequent array column or literal array to concatenate."
258    )
259)]
260#[derive(Debug, PartialEq, Eq, Hash)]
261pub struct ArrayConcat {
262    signature: Signature,
263    aliases: Vec<String>,
264}
265
266impl Default for ArrayConcat {
267    fn default() -> Self {
268        Self::new()
269    }
270}
271
272impl ArrayConcat {
273    pub fn new() -> Self {
274        Self {
275            signature: Signature::user_defined(Volatility::Immutable),
276            aliases: vec![
277                String::from("array_cat"),
278                String::from("list_concat"),
279                String::from("list_cat"),
280            ],
281        }
282    }
283}
284
285impl ScalarUDFImpl for ArrayConcat {
286    fn name(&self) -> &str {
287        "array_concat"
288    }
289
290    fn signature(&self) -> &Signature {
291        &self.signature
292    }
293
294    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
295        let mut max_dims = 0;
296        let mut large_list = false;
297        let mut element_types = Vec::with_capacity(arg_types.len());
298        for arg_type in arg_types {
299            match arg_type {
300                DataType::Null | DataType::List(_) | DataType::FixedSizeList(..) => (),
301                DataType::LargeList(_) => large_list = true,
302                arg_type => {
303                    return plan_err!("{} does not support type {arg_type}", self.name());
304                }
305            }
306
307            max_dims = max_dims.max(list_ndims(arg_type));
308            element_types.push(base_type(arg_type))
309        }
310
311        if max_dims == 0 {
312            Ok(DataType::Null)
313        } else if let Some(mut return_type) = type_union_resolution(&element_types) {
314            for _ in 1..max_dims {
315                return_type = DataType::new_list(return_type, true)
316            }
317
318            if large_list {
319                Ok(DataType::new_large_list(return_type, true))
320            } else {
321                Ok(DataType::new_list(return_type, true))
322            }
323        } else {
324            plan_err!(
325                "Failed to unify argument types of {}: [{}]",
326                self.name(),
327                arg_types.iter().join(", ")
328            )
329        }
330    }
331
332    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
333        make_scalar_function(array_concat_inner)(&args.args)
334    }
335
336    fn aliases(&self) -> &[String] {
337        &self.aliases
338    }
339
340    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
341        let return_type = self.return_type(arg_types)?;
342        let base_type = base_type(&return_type);
343        let coercion = Some(&ListCoercion::FixedSizedListToList);
344        // When the return type is a `LargeList`, the outer container of every
345        // input must be widened to `LargeList` as well. Otherwise
346        // `array_concat_inner` would later try to downcast a `List` argument
347        // to `GenericListArray<i64>` and fail.
348        let promote_to_large_list = matches!(return_type, DataType::LargeList(_));
349        let arg_types = arg_types.iter().map(|arg_type| {
350            let coerced =
351                coerced_type_with_base_type_only(arg_type, &base_type, coercion);
352            match coerced {
353                DataType::List(field) if promote_to_large_list => {
354                    DataType::LargeList(field)
355                }
356                other => other,
357            }
358        });
359
360        Ok(arg_types.collect())
361    }
362
363    fn documentation(&self) -> Option<&Documentation> {
364        self.doc()
365    }
366}
367
368pub fn array_concat_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
369    if args.is_empty() {
370        return exec_err!("array_concat expects at least one argument");
371    }
372
373    let mut all_null = true;
374    let mut large_list = false;
375    for arg in args {
376        match arg.data_type() {
377            DataType::Null => continue,
378            DataType::LargeList(_) => large_list = true,
379            _ => (),
380        }
381        if arg.null_count() < arg.len() {
382            all_null = false;
383        }
384    }
385
386    if all_null {
387        // Return a null array with the same type as the first non-null-type argument
388        let return_type = args
389            .iter()
390            .map(|arg| arg.data_type())
391            .find_or_first(|d| !d.is_null())
392            .unwrap(); // Safe because args is non-empty
393
394        Ok(arrow::array::make_array(ArrayData::new_null(
395            return_type,
396            args[0].len(),
397        )))
398    } else if large_list {
399        concat_internal::<i64>(args, None)
400    } else {
401        concat_internal::<i32>(args, None)
402    }
403}
404
405/// Return type shared by `array_append` and `array_prepend`: the input list
406/// type, except that its inner field is nullable whenever the appended or
407/// prepended element may be null.
408fn append_prepend_return_type(
409    array_type: &DataType,
410    element_type: &DataType,
411    element_nullable: bool,
412) -> DataType {
413    if array_type.is_null() {
414        DataType::new_list(element_type.clone(), true)
415    } else {
416        list_type_with_element(array_type, element_nullable)
417    }
418}
419
420/// Concatenates the list arrays in `args` row-wise.
421///
422/// `field` is the list field the output must carry. `array_concat` passes `None`
423/// because its `return_type` derives a fresh field from the unified element
424/// types, which is what deriving the field from the aligned inputs reproduces.
425/// `array_append` / `array_prepend` promise their input's field verbatim and so
426/// must pass it in explicitly.
427fn concat_internal<O: OffsetSizeTrait>(
428    args: &[ArrayRef],
429    field: Option<&FieldRef>,
430) -> Result<ArrayRef> {
431    let args = align_array_dimensions::<O>(args.to_vec())?;
432
433    let list_arrays = args
434        .iter()
435        .map(|arg| as_generic_list_array::<O>(arg))
436        .collect::<Result<Vec<_>>>()?;
437    let row_count = list_arrays[0].len();
438
439    // Extract underlying values ArrayData from each list array for MutableArrayData.
440    let values_data: Vec<ArrayData> =
441        list_arrays.iter().map(|la| la.values().to_data()).collect();
442    let values_data_refs: Vec<&ArrayData> = values_data.iter().collect();
443
444    // Estimate capacity as the sum of all values arrays' lengths.
445    let total_capacity: usize = values_data.iter().map(|d| d.len()).sum();
446
447    let mut mutable = MutableArrayData::with_capacities(
448        values_data_refs,
449        false,
450        Capacities::Array(total_capacity),
451    );
452    let mut offsets: Vec<O> = Vec::with_capacity(row_count + 1);
453    offsets.push(O::zero());
454
455    // Compute the output null buffer: a row is null only if null in ALL input
456    // arrays. This is the bitwise OR of validity bits (valid if valid in ANY
457    // input). If any array has no null buffer (all valid), no output row can be
458    // null.
459    let nulls = list_arrays
460        .iter()
461        .filter_map(|la| la.nulls())
462        .collect::<Vec<_>>();
463    let valid = if nulls.len() == list_arrays.len() {
464        nulls
465            .iter()
466            .map(|n| n.inner().clone())
467            .reduce(|a, b| &a | &b)
468            .map(NullBuffer::new)
469    } else {
470        None
471    };
472
473    for row_idx in 0..row_count {
474        for (arr_idx, list_array) in list_arrays.iter().enumerate() {
475            if list_array.is_null(row_idx) {
476                continue;
477            }
478            let start = list_array.offsets()[row_idx].to_usize().unwrap();
479            let end = list_array.offsets()[row_idx + 1].to_usize().unwrap();
480            if start < end {
481                mutable.try_extend(arr_idx, start, end)?;
482            }
483        }
484        offsets.push(O::usize_as(mutable.len()));
485    }
486
487    let field = match field {
488        Some(field) => Arc::clone(field),
489        None => Arc::new(Field::new_list_field(list_arrays[0].value_type(), true)),
490    };
491    let data = mutable.freeze();
492
493    Ok(Arc::new(GenericListArray::<O>::try_new(
494        field,
495        OffsetBuffer::new(offsets.into()),
496        arrow::array::make_array(data),
497        valid,
498    )?))
499}
500
501// Kernel functions
502
503fn array_append_inner(args: &[ArrayRef], return_type: &DataType) -> Result<ArrayRef> {
504    let [array, values] = take_function_args("array_append", args)?;
505    match array.data_type() {
506        DataType::Null => make_array_inner(&[Arc::clone(values)]),
507        DataType::List(_) => general_append_and_prepend::<i32>(args, true, return_type),
508        DataType::LargeList(_) => {
509            general_append_and_prepend::<i64>(args, true, return_type)
510        }
511        arg_type => exec_err!("array_append does not support type {arg_type}"),
512    }
513}
514
515fn array_prepend_inner(args: &[ArrayRef], return_type: &DataType) -> Result<ArrayRef> {
516    let [values, array] = take_function_args("array_prepend", args)?;
517    match array.data_type() {
518        DataType::Null => make_array_inner(&[Arc::clone(values)]),
519        DataType::List(_) => general_append_and_prepend::<i32>(args, false, return_type),
520        DataType::LargeList(_) => {
521            general_append_and_prepend::<i64>(args, false, return_type)
522        }
523        arg_type => exec_err!("array_prepend does not support type {arg_type}"),
524    }
525}
526
527fn general_append_and_prepend<O: OffsetSizeTrait>(
528    args: &[ArrayRef],
529    is_append: bool,
530    return_type: &DataType,
531) -> Result<ArrayRef>
532where
533    i64: TryInto<O>,
534{
535    let (list_array, element_array) = if is_append {
536        let list_array = as_generic_list_array::<O>(&args[0])?;
537        let element_array = &args[1];
538        check_datatypes("array_append", &[element_array, list_array.values()])?;
539        (list_array, element_array)
540    } else {
541        let list_array = as_generic_list_array::<O>(&args[1])?;
542        let element_array = &args[0];
543        check_datatypes("array_prepend", &[list_array.values(), element_array])?;
544        (list_array, element_array)
545    };
546
547    let name = if is_append {
548        "array_append"
549    } else {
550        "array_prepend"
551    };
552    let field = list_inner_field(name, return_type)?;
553
554    let res = match list_array.value_type() {
555        DataType::List(_) | DataType::LargeList(_) => {
556            concat_internal::<O>(args, Some(&field))?
557        }
558        _ => {
559            return generic_append_and_prepend::<O>(
560                list_array,
561                element_array,
562                field,
563                is_append,
564            );
565        }
566    };
567
568    Ok(res)
569}
570
571/// Appends or prepends elements to a ListArray.
572///
573/// This function takes a ListArray, an ArrayRef, a FieldRef, and a boolean flag
574/// indicating whether to append or prepend the elements. It returns a `Result<ArrayRef>`
575/// representing the resulting ListArray after the operation.
576///
577/// # Arguments
578///
579/// * `list_array` - A reference to the ListArray to which elements will be appended/prepended.
580/// * `element_array` - A reference to the Array containing elements to be appended/prepended.
581/// * `field` - The list field the output must carry, taken from the promised return type.
582/// * `is_append` - A boolean flag indicating whether to append (`true`) or prepend (`false`) elements.
583///
584/// # Examples
585///
586/// generic_append_and_prepend(
587///     [1, 2, 3], 4, append => [1, 2, 3, 4]
588///     5, [6, 7, 8], prepend => [5, 6, 7, 8]
589/// )
590fn generic_append_and_prepend<O: OffsetSizeTrait>(
591    list_array: &GenericListArray<O>,
592    element_array: &ArrayRef,
593    field: FieldRef,
594    is_append: bool,
595) -> Result<ArrayRef>
596where
597    i64: TryInto<O>,
598{
599    let mut offsets = vec![O::usize_as(0)];
600    let values = list_array.values();
601    let original_data = values.to_data();
602    let element_data = element_array.to_data();
603    let capacity = Capacities::Array(original_data.len() + element_data.len());
604
605    let mut mutable = MutableArrayData::with_capacities(
606        vec![&original_data, &element_data],
607        false,
608        capacity,
609    );
610
611    let values_index = 0;
612    let element_index = 1;
613
614    for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() {
615        let start = offset_window[0].to_usize().unwrap();
616        let end = offset_window[1].to_usize().unwrap();
617        if is_append {
618            mutable.try_extend(values_index, start, end)?;
619            mutable.try_extend(element_index, row_index, row_index + 1)?;
620        } else {
621            mutable.try_extend(element_index, row_index, row_index + 1)?;
622            mutable.try_extend(values_index, start, end)?;
623        }
624        offsets.push(offsets[row_index] + O::usize_as(end - start + 1));
625    }
626
627    let data = mutable.freeze();
628
629    Ok(Arc::new(GenericListArray::<O>::try_new(
630        field,
631        OffsetBuffer::new(offsets.into()),
632        arrow::array::make_array(data),
633        None,
634    )?))
635}
datafusion_functions_nested/concat.rs

datafusion_functions_nested/
concat.rs
