datafusion_functions_nested/
reverse.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_reverse function.
19
20use crate::utils::make_scalar_function;
21use arrow::array::{
22    Array, ArrayRef, FixedSizeListArray, GenericListArray, GenericListViewArray,
23    OffsetSizeTrait, UInt32Array, UInt64Array,
24};
25use arrow::buffer::{OffsetBuffer, ScalarBuffer};
26use arrow::compute::take;
27use arrow::datatypes::DataType::{
28    FixedSizeList, LargeList, LargeListView, List, ListView, Null,
29};
30use arrow::datatypes::{DataType, FieldRef};
31use datafusion_common::cast::{
32    as_fixed_size_list_array, as_large_list_array, as_large_list_view_array,
33    as_list_array, as_list_view_array,
34};
35use datafusion_common::{Result, exec_err, utils::take_function_args};
36use datafusion_expr::{
37    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
38};
39use datafusion_macros::user_doc;
40use itertools::Itertools;
41use std::any::Any;
42use std::sync::Arc;
43
44make_udf_expr_and_func!(
45    ArrayReverse,
46    array_reverse,
47    array,
48    "reverses the order of elements in the array.",
49    array_reverse_udf
50);
51
52#[user_doc(
53    doc_section(label = "Array Functions"),
54    description = "Returns the array with the order of the elements reversed.",
55    syntax_example = "array_reverse(array)",
56    sql_example = r#"```sql
57> select array_reverse([1, 2, 3, 4]);
58+------------------------------------------------------------+
59| array_reverse(List([1, 2, 3, 4]))                          |
60+------------------------------------------------------------+
61| [4, 3, 2, 1]                                               |
62+------------------------------------------------------------+
63```"#,
64    argument(
65        name = "array",
66        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
67    )
68)]
69#[derive(Debug, PartialEq, Eq, Hash)]
70pub struct ArrayReverse {
71    signature: Signature,
72    aliases: Vec<String>,
73}
74
75impl Default for ArrayReverse {
76    fn default() -> Self {
77        Self::new()
78    }
79}
80
81impl ArrayReverse {
82    pub fn new() -> Self {
83        Self {
84            signature: Signature::array(Volatility::Immutable),
85            aliases: vec!["list_reverse".to_string()],
86        }
87    }
88}
89
90impl ScalarUDFImpl for ArrayReverse {
91    fn as_any(&self) -> &dyn Any {
92        self
93    }
94
95    fn name(&self) -> &str {
96        "array_reverse"
97    }
98
99    fn signature(&self) -> &Signature {
100        &self.signature
101    }
102
103    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
104        Ok(arg_types[0].clone())
105    }
106
107    fn invoke_with_args(
108        &self,
109        args: datafusion_expr::ScalarFunctionArgs,
110    ) -> Result<ColumnarValue> {
111        make_scalar_function(array_reverse_inner)(&args.args)
112    }
113
114    fn aliases(&self) -> &[String] {
115        &self.aliases
116    }
117
118    fn documentation(&self) -> Option<&Documentation> {
119        self.doc()
120    }
121}
122
123/// array_reverse SQL function
124pub fn array_reverse_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
125    let [input_array] = take_function_args("array_reverse", arg)?;
126
127    match &input_array.data_type() {
128        List(field) => {
129            let array = as_list_array(input_array)?;
130            general_array_reverse::<i32>(array, field)
131        }
132        LargeList(field) => {
133            let array = as_large_list_array(input_array)?;
134            general_array_reverse::<i64>(array, field)
135        }
136        FixedSizeList(field, _) => {
137            let array = as_fixed_size_list_array(input_array)?;
138            fixed_size_array_reverse(array, field)
139        }
140        Null => Ok(Arc::clone(input_array)),
141        ListView(field) => {
142            let array = as_list_view_array(input_array)?;
143            list_view_reverse::<i32>(array, field)
144        }
145        LargeListView(field) => {
146            let array = as_large_list_view_array(input_array)?;
147            list_view_reverse::<i64>(array, field)
148        }
149        array_type => exec_err!("array_reverse does not support type '{array_type}'."),
150    }
151}
152
153fn general_array_reverse<O: OffsetSizeTrait>(
154    array: &GenericListArray<O>,
155    field: &FieldRef,
156) -> Result<ArrayRef> {
157    let values = array.values();
158    let mut offsets = vec![O::usize_as(0)];
159    let mut indices: Vec<O> = Vec::with_capacity(values.len());
160
161    for (row_index, (&start, &end)) in array.offsets().iter().tuple_windows().enumerate()
162    {
163        // skip the null value
164        if array.is_null(row_index) {
165            offsets.push(offsets[row_index]);
166            continue;
167        }
168
169        let mut index = end - O::one();
170        while index >= start {
171            indices.push(index);
172            index = index - O::one();
173        }
174        let size = end - start;
175        offsets.push(offsets[row_index] + size);
176    }
177
178    // Materialize values from underlying array with take
179    let indices_array: ArrayRef = if O::IS_LARGE {
180        Arc::new(UInt64Array::from(
181            indices
182                .iter()
183                .map(|i| i.as_usize() as u64)
184                .collect::<Vec<_>>(),
185        ))
186    } else {
187        Arc::new(UInt32Array::from(
188            indices
189                .iter()
190                .map(|i| i.as_usize() as u32)
191                .collect::<Vec<_>>(),
192        ))
193    };
194    let values = take(&values, &indices_array, None)?;
195    Ok(Arc::new(GenericListArray::<O>::try_new(
196        Arc::clone(field),
197        OffsetBuffer::<O>::new(offsets.into()),
198        values,
199        array.nulls().cloned(),
200    )?))
201}
202
203/// Reverses a list view array.
204///
205/// Construct indices, sizes and offsets for the reversed array by iterating over
206/// the list view array in the logical order, and reversing the order of the elements.
207/// We end up with a list view array where the elements are in order,
208/// even if the original array had elements out of order.
209fn list_view_reverse<O: OffsetSizeTrait>(
210    array: &GenericListViewArray<O>,
211    field: &FieldRef,
212) -> Result<ArrayRef> {
213    let offsets = array.offsets();
214    let values = array.values();
215    let sizes = array.sizes();
216
217    let mut new_offsets: Vec<O> = Vec::with_capacity(offsets.len());
218    let mut indices: Vec<O> = Vec::with_capacity(values.len());
219    let mut new_sizes = Vec::with_capacity(sizes.len());
220
221    let mut current_offset = O::zero();
222    for (row_index, offset) in offsets.iter().enumerate() {
223        new_offsets.push(current_offset);
224
225        // If this array is null, we set its size to 0 and continue
226        if array.is_null(row_index) {
227            new_sizes.push(O::zero());
228            continue;
229        }
230        let size = sizes[row_index];
231        new_sizes.push(size);
232
233        // Each array is located at [offset, offset + size), collect indices in the reverse order
234        let array_start = *offset;
235        let array_end = array_start + size;
236        let mut idx = array_end - O::one();
237        while idx >= array_start {
238            indices.push(idx);
239            idx = idx - O::one();
240        }
241
242        current_offset += size;
243    }
244
245    // Materialize values from underlying array with take
246    let indices_array: ArrayRef = if O::IS_LARGE {
247        Arc::new(UInt64Array::from(
248            indices
249                .iter()
250                .map(|i| i.as_usize() as u64)
251                .collect::<Vec<_>>(),
252        ))
253    } else {
254        Arc::new(UInt32Array::from(
255            indices
256                .iter()
257                .map(|i| i.as_usize() as u32)
258                .collect::<Vec<_>>(),
259        ))
260    };
261    let values = take(&values, &indices_array, None)?;
262    Ok(Arc::new(GenericListViewArray::<O>::try_new(
263        Arc::clone(field),
264        ScalarBuffer::from(new_offsets),
265        ScalarBuffer::from(new_sizes),
266        values,
267        array.nulls().cloned(),
268    )?))
269}
270
271fn fixed_size_array_reverse(
272    array: &FixedSizeListArray,
273    field: &FieldRef,
274) -> Result<ArrayRef> {
275    let values: &Arc<dyn Array> = array.values();
276
277    // Since each fixed size list in the physical array is the same size and we keep the order
278    // of the fixed size lists, we can reverse the indices for each fixed size list.
279    let mut indices: Vec<u64> = (0..values.len() as u64).collect();
280    for chunk in indices.chunks_mut(array.value_length() as usize) {
281        chunk.reverse();
282    }
283
284    // Materialize values from underlying array with take
285    let indices_array: ArrayRef = Arc::new(UInt64Array::from(indices));
286    let values = take(&values, &indices_array, None)?;
287
288    Ok(Arc::new(FixedSizeListArray::try_new(
289        Arc::clone(field),
290        array.value_length(),
291        values,
292        array.nulls().cloned(),
293    )?))
294}
295
296#[cfg(test)]
297mod tests {
298    use crate::reverse::{fixed_size_array_reverse, list_view_reverse};
299    use arrow::{
300        array::{
301            AsArray, FixedSizeListArray, GenericListViewArray, Int32Array,
302            LargeListViewArray, ListViewArray, OffsetSizeTrait,
303        },
304        buffer::{NullBuffer, ScalarBuffer},
305        datatypes::{DataType, Field, Int32Type},
306    };
307    use datafusion_common::Result;
308    use std::sync::Arc;
309
310    fn list_view_values<O: OffsetSizeTrait>(
311        array: &GenericListViewArray<O>,
312    ) -> Vec<Option<Vec<i32>>> {
313        array
314            .iter()
315            .map(|x| x.map(|x| x.as_primitive::<Int32Type>().values().to_vec()))
316            .collect()
317    }
318
319    fn fixed_size_list_values(array: &FixedSizeListArray) -> Vec<Option<Vec<i32>>> {
320        array
321            .iter()
322            .map(|x| x.map(|x| x.as_primitive::<Int32Type>().values().to_vec()))
323            .collect()
324    }
325
326    #[test]
327    fn test_reverse_list_view() -> Result<()> {
328        let field = Arc::new(Field::new("a", DataType::Int32, false));
329        let offsets = ScalarBuffer::from(vec![0, 1, 6, 6]);
330        let sizes = ScalarBuffer::from(vec![1, 5, 0, 3]);
331        let values = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9]));
332        let nulls = Some(NullBuffer::from(vec![true, true, false, true]));
333        let list_view = ListViewArray::new(field, offsets, sizes, values, nulls);
334        let result = list_view_reverse(
335            &list_view,
336            &Arc::new(Field::new("test", DataType::Int32, true)),
337        )?;
338        let reversed = list_view_values(result.as_list_view::<i32>());
339        let expected = vec![
340            Some(vec![1]),
341            Some(vec![6, 5, 4, 3, 2]),
342            None,
343            Some(vec![9, 8, 7]),
344        ];
345        assert_eq!(expected, reversed);
346        Ok(())
347    }
348
349    #[test]
350    fn test_reverse_large_list_view() -> Result<()> {
351        let field = Arc::new(Field::new("a", DataType::Int32, false));
352        let offsets = ScalarBuffer::from(vec![0, 1, 6, 6]);
353        let sizes = ScalarBuffer::from(vec![1, 5, 0, 3]);
354        let values = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9]));
355        let nulls = Some(NullBuffer::from(vec![true, true, false, true]));
356        let list_view = LargeListViewArray::new(field, offsets, sizes, values, nulls);
357        let result = list_view_reverse(
358            &list_view,
359            &Arc::new(Field::new("test", DataType::Int32, true)),
360        )?;
361        let reversed = list_view_values(result.as_list_view::<i64>());
362        let expected = vec![
363            Some(vec![1]),
364            Some(vec![6, 5, 4, 3, 2]),
365            None,
366            Some(vec![9, 8, 7]),
367        ];
368        assert_eq!(expected, reversed);
369        Ok(())
370    }
371
372    #[test]
373    fn test_reverse_list_view_out_of_order() -> Result<()> {
374        let field = Arc::new(Field::new("a", DataType::Int32, false));
375        let offsets = ScalarBuffer::from(vec![6, 1, 6, 0]); // out of order
376        let sizes = ScalarBuffer::from(vec![3, 5, 0, 1]);
377        let values = Arc::new(Int32Array::from(vec![
378            1, // fourth array: offset 0, size 1
379            2, 3, 4, 5, 6, // second array: offset 1, size 5
380            // third array: offset 6, size 0 (and null)
381            7, 8, 9, // first array: offset 6, size 3
382        ]));
383        let nulls = Some(NullBuffer::from(vec![true, true, false, true]));
384        let list_view = ListViewArray::new(field, offsets, sizes, values, nulls);
385        let result = list_view_reverse(
386            &list_view,
387            &Arc::new(Field::new("test", DataType::Int32, true)),
388        )?;
389        let reversed = list_view_values(result.as_list_view::<i32>());
390        let expected = vec![
391            Some(vec![9, 8, 7]),
392            Some(vec![6, 5, 4, 3, 2]),
393            None,
394            Some(vec![1]),
395        ];
396        assert_eq!(expected, reversed);
397        Ok(())
398    }
399
400    #[test]
401    fn test_reverse_list_view_with_nulls() -> Result<()> {
402        let field = Arc::new(Field::new("a", DataType::Int32, false));
403        let offsets = ScalarBuffer::from(vec![16, 1, 6, 0]); // out of order
404        let sizes = ScalarBuffer::from(vec![3, 5, 10, 1]);
405        let values = Arc::new(Int32Array::from(vec![
406            1, // fourth array: offset 0, size 1
407            2, 3, 4, 5, 6, // second array: offset 1, size 5
408            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, // third array: offset 6, size 10
409            7, 8, 9, // first array: offset 6, size 3
410        ]));
411        let nulls = Some(NullBuffer::from(vec![true, true, false, true]));
412        let list_view = ListViewArray::new(field, offsets, sizes, values, nulls);
413        let result = list_view_reverse(
414            &list_view,
415            &Arc::new(Field::new("test", DataType::Int32, true)),
416        )?;
417        let reversed = list_view_values(result.as_list_view::<i32>());
418        let expected = vec![
419            Some(vec![9, 8, 7]),
420            Some(vec![6, 5, 4, 3, 2]),
421            None,
422            Some(vec![1]),
423        ];
424        assert_eq!(expected, reversed);
425        Ok(())
426    }
427
428    #[test]
429    fn test_reverse_list_view_empty() -> Result<()> {
430        let field = Arc::new(Field::new("a", DataType::Int32, false));
431        let offsets = ScalarBuffer::from(vec![]);
432        let sizes = ScalarBuffer::from(vec![]);
433        let empty_array: Vec<i32> = vec![];
434        let values = Arc::new(Int32Array::from(empty_array));
435        let nulls = None;
436        let list_view = ListViewArray::new(field, offsets, sizes, values, nulls);
437        let result = list_view_reverse(
438            &list_view,
439            &Arc::new(Field::new("test", DataType::Int32, true)),
440        )?;
441        let reversed = list_view_values(result.as_list_view::<i32>());
442        let expected: Vec<Option<Vec<i32>>> = vec![];
443        assert_eq!(expected, reversed);
444        Ok(())
445    }
446
447    #[test]
448    fn test_reverse_list_view_all_nulls() -> Result<()> {
449        let field = Arc::new(Field::new("a", DataType::Int32, false));
450        let offsets = ScalarBuffer::from(vec![0, 1, 2, 3]);
451        let sizes = ScalarBuffer::from(vec![0, 1, 1, 1]);
452        let values = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
453        let nulls = Some(NullBuffer::from(vec![false, false, false, false]));
454        let list_view = ListViewArray::new(field, offsets, sizes, values, nulls);
455        let result = list_view_reverse(
456            &list_view,
457            &Arc::new(Field::new("test", DataType::Int32, true)),
458        )?;
459        let reversed = list_view_values(result.as_list_view::<i32>());
460        let expected: Vec<Option<Vec<i32>>> = vec![None, None, None, None];
461        assert_eq!(expected, reversed);
462        Ok(())
463    }
464
465    #[test]
466    fn test_reverse_fixed_size_list() -> Result<()> {
467        let field = Arc::new(Field::new("a", DataType::Int32, false));
468        let values = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9]));
469        let result = fixed_size_array_reverse(
470            &FixedSizeListArray::new(
471                field,
472                3,
473                values,
474                Some(NullBuffer::from(vec![true, false, true])),
475            ),
476            &Arc::new(Field::new("test", DataType::Int32, true)),
477        )?;
478        let reversed = fixed_size_list_values(result.as_fixed_size_list());
479        let expected = vec![Some(vec![3, 2, 1]), None, Some(vec![9, 8, 7])];
480        assert_eq!(expected, reversed);
481        Ok(())
482    }
483
484    #[test]
485    fn test_reverse_fixed_size_list_empty() -> Result<()> {
486        let field = Arc::new(Field::new("a", DataType::Int32, false));
487        let empty_array: Vec<i32> = vec![];
488        let values = Arc::new(Int32Array::from(empty_array));
489        let nulls = None;
490        let fixed_size_list = FixedSizeListArray::new(field, 3, values, nulls);
491        let result = fixed_size_array_reverse(
492            &fixed_size_list,
493            &Arc::new(Field::new("test", DataType::Int32, true)),
494        )?;
495        let reversed = fixed_size_list_values(result.as_fixed_size_list());
496        let expected: Vec<Option<Vec<i32>>> = vec![];
497        assert_eq!(expected, reversed);
498        Ok(())
499    }
500}
datafusion_functions_nested/reverse.rs

datafusion_functions_nested/
reverse.rs
