datafusion_functions_nested/
array_add.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_add function.
19
20use crate::utils::{coerce_array_math_arg_types, make_scalar_function};
21use arrow::array::{
22    Array, ArrayRef, Float64Array, GenericListArray, NullBufferBuilder, OffsetSizeTrait,
23};
24use arrow::buffer::{NullBuffer, OffsetBuffer};
25use arrow::datatypes::{
26    DataType,
27    DataType::{LargeList, List},
28    Field,
29};
30use datafusion_common::cast::{as_float64_array, as_generic_list_array};
31use datafusion_common::{Result, exec_err, utils::take_function_args};
32use datafusion_expr::{
33    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
34    Volatility,
35};
36use datafusion_macros::user_doc;
37use std::sync::Arc;
38
39make_udf_expr_and_func!(
40    ArrayAdd,
41    array_add,
42    array1 array2,
43    "returns the element-wise sum of two numeric arrays.",
44    array_add_udf
45);
46
47#[user_doc(
48    doc_section(label = "Array Functions"),
49    description = "Returns the element-wise sum of two numeric arrays of equal length, computed as `array1[i] + array2[i]` per position. NULL is propagated per element: if either input element at position `i` is NULL, the corresponding output element is NULL (positions are preserved). Returns NULL if either entire input array is NULL. Errors if the per-row lengths differ. Returns an empty array if both inputs are empty.",
50    syntax_example = "array_add(array1, array2)",
51    sql_example = r#"```sql
52> select array_add([1.0, 2.0, 3.0], [10.0, 20.0, 30.0]);
53+---------------------------------------------------------+
54| array_add(List([1.0,2.0,3.0]),List([10.0,20.0,30.0]))   |
55+---------------------------------------------------------+
56| [11.0, 22.0, 33.0]                                      |
57+---------------------------------------------------------+
58```"#,
59    argument(
60        name = "array1",
61        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
62    ),
63    argument(
64        name = "array2",
65        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
66    )
67)]
68#[derive(Debug, PartialEq, Eq, Hash)]
69pub struct ArrayAdd {
70    signature: Signature,
71    aliases: Vec<String>,
72}
73
74impl Default for ArrayAdd {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80impl ArrayAdd {
81    pub fn new() -> Self {
82        Self {
83            signature: Signature::user_defined(Volatility::Immutable),
84            aliases: vec!["list_add".to_string()],
85        }
86    }
87}
88
89impl ScalarUDFImpl for ArrayAdd {
90    fn name(&self) -> &str {
91        "array_add"
92    }
93
94    fn signature(&self) -> &Signature {
95        &self.signature
96    }
97
98    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
99        // After `coerce_types`, both args share the same List/LargeList<Float64> shape.
100        Ok(arg_types[0].clone())
101    }
102
103    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
104        let [_, _] = take_function_args(self.name(), arg_types)?;
105        coerce_array_math_arg_types(self.name(), arg_types)
106    }
107
108    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
109        make_scalar_function(array_add_inner)(&args.args)
110    }
111
112    fn aliases(&self) -> &[String] {
113        &self.aliases
114    }
115
116    fn documentation(&self) -> Option<&Documentation> {
117        self.doc()
118    }
119}
120
121fn array_add_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
122    let [array1, array2] = take_function_args("array_add", args)?;
123    match (array1.data_type(), array2.data_type()) {
124        (List(_), List(_)) => general_array_add::<i32>(array1, array2),
125        (LargeList(_), LargeList(_)) => general_array_add::<i64>(array1, array2),
126        (arg_type1, arg_type2) => exec_err!(
127            "array_add received unexpected types after coercion: {arg_type1} and {arg_type2}"
128        ),
129    }
130}
131
132fn general_array_add<O: OffsetSizeTrait>(
133    lhs: &ArrayRef,
134    rhs: &ArrayRef,
135) -> Result<ArrayRef> {
136    let lhs = as_generic_list_array::<O>(lhs)?;
137    let rhs = as_generic_list_array::<O>(rhs)?;
138
139    let lhs_values = as_float64_array(lhs.values())?;
140    let rhs_values = as_float64_array(rhs.values())?;
141    let lhs_offsets = lhs.value_offsets();
142    let rhs_offsets = rhs.value_offsets();
143
144    // Row-level validity: a row is valid iff both sides are valid at that row.
145    let row_nulls = NullBuffer::union(lhs.nulls(), rhs.nulls());
146
147    let mut out_values: Vec<f64> = Vec::with_capacity(lhs_values.len());
148    let mut out_inner_nulls = NullBufferBuilder::new(lhs_values.len());
149    let mut out_offsets = Vec::<O>::with_capacity(lhs.len() + 1);
150    out_offsets.push(O::zero());
151
152    for row in 0..lhs.len() {
153        // Whole-row NULL on either side -> NULL output row, no elements.
154        if row_nulls.as_ref().is_some_and(|nb| nb.is_null(row)) {
155            out_offsets.push(out_offsets[row]);
156            continue;
157        }
158
159        let start1 = lhs_offsets[row].as_usize();
160        let len1 = lhs.value_length(row).as_usize();
161        let start2 = rhs_offsets[row].as_usize();
162        let len2 = rhs.value_length(row).as_usize();
163
164        if len1 != len2 {
165            return exec_err!(
166                "array_add requires both list inputs to have the same length per row, got {len1} and {len2} at row {row}"
167            );
168        }
169
170        let l_slice = lhs_values.slice(start1, len1);
171        let r_slice = rhs_values.slice(start2, len2);
172
173        let l_vals = l_slice.values();
174        let r_vals = r_slice.values();
175
176        for i in 0..len1 {
177            out_values.push(l_vals[i] + r_vals[i]);
178        }
179
180        // Per-element validity: position `i` is valid iff both lhs[i] and rhs[i]
181        // are valid. `NullBuffer::union` returns `None` when both sides are
182        // entirely valid.
183        match NullBuffer::union(l_slice.nulls(), r_slice.nulls()) {
184            Some(nb) => out_inner_nulls.append_buffer(&nb),
185            None => out_inner_nulls.append_n_non_nulls(len1),
186        }
187
188        out_offsets.push(out_offsets[row] + O::usize_as(len1));
189    }
190
191    let values_array = Arc::new(Float64Array::new(
192        out_values.into(),
193        out_inner_nulls.finish(),
194    ));
195    let field = Arc::new(Field::new_list_field(DataType::Float64, true));
196
197    Ok(Arc::new(GenericListArray::<O>::try_new(
198        field,
199        OffsetBuffer::new(out_offsets.into()),
200        values_array,
201        row_nulls,
202    )?))
203}
datafusion_functions_nested/array_add.rs

datafusion_functions_nested/
array_add.rs
