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

datafusion_functions_nested/
reverse.rs
