datafusion_functions_nested/
except.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`] definition for array_except function.
19
20use crate::utils::{check_datatypes, make_scalar_function};
21use arrow::array::new_null_array;
22use arrow::array::{
23    Array, ArrayRef, GenericListArray, OffsetSizeTrait, UInt32Array, UInt64Array,
24    cast::AsArray,
25};
26use arrow::buffer::{NullBuffer, OffsetBuffer};
27use arrow::compute::take;
28use arrow::datatypes::{DataType, FieldRef};
29use arrow::row::{RowConverter, SortField};
30use datafusion_common::utils::{ListCoercion, normalize_float_zero, take_function_args};
31use datafusion_common::{HashSet, Result, internal_err};
32use datafusion_expr::{
33    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
34    Volatility,
35};
36use datafusion_macros::user_doc;
37use itertools::Itertools;
38use std::sync::Arc;
39
40make_udf_expr_and_func!(
41    ArrayExcept,
42    array_except,
43    first_array second_array,
44    "returns an array of the elements that appear in the first array but not in the second.",
45    array_except_udf
46);
47
48#[user_doc(
49    doc_section(label = "Array Functions"),
50    description = "Returns an array of the elements that appear in the first array but not in the second.",
51    syntax_example = "array_except(array1, array2)",
52    sql_example = r#"```sql
53> select array_except([1, 2, 3, 4], [5, 6, 3, 4]);
54+----------------------------------------------------+
55| array_except([1, 2, 3, 4], [5, 6, 3, 4]);           |
56+----------------------------------------------------+
57| [1, 2]                                              |
58+----------------------------------------------------+
59> select array_except([1, 2, 3, 4], [3, 4, 5, 6]);
60+----------------------------------------------------+
61| array_except([1, 2, 3, 4], [3, 4, 5, 6]);           |
62+----------------------------------------------------+
63| [1, 2]                                              |
64+----------------------------------------------------+
65```"#,
66    argument(
67        name = "array1",
68        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
69    ),
70    argument(
71        name = "array2",
72        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
73    )
74)]
75#[derive(Debug, PartialEq, Eq, Hash)]
76pub struct ArrayExcept {
77    signature: Signature,
78    aliases: Vec<String>,
79}
80
81impl Default for ArrayExcept {
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87impl ArrayExcept {
88    pub fn new() -> Self {
89        Self {
90            signature: Signature::arrays(
91                2,
92                Some(ListCoercion::FixedSizedListToList),
93                Volatility::Immutable,
94            ),
95            aliases: vec!["list_except".to_string()],
96        }
97    }
98}
99
100impl ScalarUDFImpl for ArrayExcept {
101    fn name(&self) -> &str {
102        "array_except"
103    }
104
105    fn signature(&self) -> &Signature {
106        &self.signature
107    }
108
109    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
110        match (&arg_types[0], &arg_types[1]) {
111            (DataType::Null, DataType::Null) => {
112                Ok(DataType::new_list(DataType::Null, true))
113            }
114            (DataType::Null, dt) | (dt, DataType::Null) => Ok(dt.clone()),
115            (dt, _) => Ok(dt.clone()),
116        }
117    }
118
119    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
120        make_scalar_function(array_except_inner)(&args.args)
121    }
122
123    fn aliases(&self) -> &[String] {
124        &self.aliases
125    }
126
127    fn documentation(&self) -> Option<&Documentation> {
128        self.doc()
129    }
130}
131
132fn array_except_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
133    let [array1, array2] = take_function_args("array_except", args)?;
134
135    let len = array1.len();
136    match (array1.data_type(), array2.data_type()) {
137        (DataType::Null, DataType::Null) => Ok(new_null_array(
138            &DataType::new_list(DataType::Null, true),
139            len,
140        )),
141        (DataType::Null, dt @ DataType::List(_))
142        | (DataType::Null, dt @ DataType::LargeList(_))
143        | (dt @ DataType::List(_), DataType::Null)
144        | (dt @ DataType::LargeList(_), DataType::Null) => Ok(new_null_array(dt, len)),
145        (DataType::List(field), DataType::List(_)) => {
146            check_datatypes("array_except", &[array1, array2])?;
147            let list1 = array1.as_list::<i32>();
148            let list2 = array2.as_list::<i32>();
149            let result = general_except::<i32>(list1, list2, field)?;
150            Ok(Arc::new(result))
151        }
152        (DataType::LargeList(field), DataType::LargeList(_)) => {
153            check_datatypes("array_except", &[array1, array2])?;
154            let list1 = array1.as_list::<i64>();
155            let list2 = array2.as_list::<i64>();
156            let result = general_except::<i64>(list1, list2, field)?;
157            Ok(Arc::new(result))
158        }
159        (dt1, dt2) => {
160            internal_err!("array_except got unexpected types: {dt1:?} and {dt2:?}")
161        }
162    }
163}
164
165fn general_except<OffsetSize: OffsetSizeTrait>(
166    l: &GenericListArray<OffsetSize>,
167    r: &GenericListArray<OffsetSize>,
168    field: &FieldRef,
169) -> Result<GenericListArray<OffsetSize>> {
170    let converter = RowConverter::new(vec![SortField::new(l.value_type())])?;
171
172    // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) groups
173    // ±0 together for both the rhs lookup set and the lhs probe.
174    let l_values_norm = normalize_float_zero(l.values());
175    let r_values_norm = normalize_float_zero(r.values());
176
177    // Only convert the visible portion of the values array. For sliced
178    // ListArrays, values() returns the full underlying array but only
179    // elements between the first and last offset are referenced.
180    let l_first = l.offsets()[0].as_usize();
181    let l_len = l.offsets()[l.len()].as_usize() - l_first;
182    let l_values = converter.convert_columns(&[l_values_norm.slice(l_first, l_len)])?;
183
184    let r_first = r.offsets()[0].as_usize();
185    let r_len = r.offsets()[r.len()].as_usize() - r_first;
186    let r_values = converter.convert_columns(&[r_values_norm.slice(r_first, r_len)])?;
187
188    let mut offsets = Vec::<OffsetSize>::with_capacity(l.len() + 1);
189    offsets.push(OffsetSize::usize_as(0));
190
191    let mut indices: Vec<usize> = Vec::with_capacity(l_values.num_rows());
192    let mut dedup = HashSet::new();
193
194    let nulls = NullBuffer::union(l.nulls(), r.nulls());
195
196    let l_offsets_iter = l.offsets().iter().tuple_windows();
197    let r_offsets_iter = r.offsets().iter().tuple_windows();
198    for (list_index, ((l_start, l_end), (r_start, r_end))) in
199        l_offsets_iter.zip(r_offsets_iter).enumerate()
200    {
201        if nulls
202            .as_ref()
203            .is_some_and(|nulls| nulls.is_null(list_index))
204        {
205            offsets.push(OffsetSize::usize_as(indices.len()));
206            continue;
207        }
208
209        for element_index in r_start.as_usize() - r_first..r_end.as_usize() - r_first {
210            let right_row = r_values.row(element_index);
211            dedup.insert(right_row);
212        }
213        for element_index in l_start.as_usize() - l_first..l_end.as_usize() - l_first {
214            let left_row = l_values.row(element_index);
215            if dedup.insert(left_row) {
216                indices.push(element_index + l_first);
217            }
218        }
219
220        offsets.push(OffsetSize::usize_as(indices.len()));
221        dedup.clear();
222    }
223
224    // Gather distinct left-side values by index.
225    // Use UInt64Array for LargeList to support values arrays exceeding u32::MAX.
226    let values = if indices.is_empty() {
227        arrow::array::new_empty_array(&l.value_type())
228    } else if OffsetSize::IS_LARGE {
229        let indices =
230            UInt64Array::from(indices.into_iter().map(|i| i as u64).collect::<Vec<_>>());
231        take(l_values_norm.as_ref(), &indices, None)?
232    } else {
233        let indices =
234            UInt32Array::from(indices.into_iter().map(|i| i as u32).collect::<Vec<_>>());
235        take(l_values_norm.as_ref(), &indices, None)?
236    };
237
238    Ok(GenericListArray::<OffsetSize>::new(
239        field.to_owned(),
240        OffsetBuffer::new(offsets.into()),
241        values,
242        nulls,
243    ))
244}
245
246#[cfg(test)]
247mod tests {
248    use super::ArrayExcept;
249    use arrow::array::{Array, AsArray, Int32Array, ListArray};
250    use arrow::datatypes::{Field, Int32Type};
251    use datafusion_common::{Result, config::ConfigOptions};
252    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
253    use std::sync::Arc;
254
255    #[test]
256    fn test_array_except_sliced_lists() -> Result<()> {
257        // l: [[1,2], [3,4], [5,6], [7,8]]  →  slice(1,2)  →  [[3,4], [5,6]]
258        // r: [[3],   [5],   [6],   [8]]    →  slice(1,2)  →  [[5],   [6]]
259        // except(l, r) should be [[3,4], [5]]
260        let l_full = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
261            Some(vec![Some(1), Some(2)]),
262            Some(vec![Some(3), Some(4)]),
263            Some(vec![Some(5), Some(6)]),
264            Some(vec![Some(7), Some(8)]),
265        ]);
266        let r_full = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
267            Some(vec![Some(3)]),
268            Some(vec![Some(5)]),
269            Some(vec![Some(6)]),
270            Some(vec![Some(8)]),
271        ]);
272        let l_sliced = l_full.slice(1, 2);
273        let r_sliced = r_full.slice(1, 2);
274
275        let list_field = Arc::new(Field::new("item", l_sliced.data_type().clone(), true));
276        let return_field =
277            Arc::new(Field::new("return", l_sliced.data_type().clone(), true));
278
279        let result = ArrayExcept::new().invoke_with_args(ScalarFunctionArgs {
280            args: vec![
281                ColumnarValue::Array(Arc::new(l_sliced)),
282                ColumnarValue::Array(Arc::new(r_sliced)),
283            ],
284            arg_fields: vec![Arc::clone(&list_field), Arc::clone(&list_field)],
285            number_rows: 2,
286            return_field,
287            config_options: Arc::new(ConfigOptions::default()),
288        })?;
289
290        let output = result.into_array(2)?;
291        let output = output.as_list::<i32>();
292
293        // Row 0: [3,4] except [5] = [3,4]
294        let row0 = output.value(0);
295        let row0 = row0.as_any().downcast_ref::<Int32Array>().unwrap();
296        assert_eq!(row0.values().as_ref(), &[3, 4]);
297
298        // Row 1: [5,6] except [6] = [5]
299        let row1 = output.value(1);
300        let row1 = row1.as_any().downcast_ref::<Int32Array>().unwrap();
301        assert_eq!(row1.values().as_ref(), &[5]);
302
303        Ok(())
304    }
305}
datafusion_functions_nested/except.rs

datafusion_functions_nested/
except.rs
