Support reverse for ListView

Author: vegarsti

datafusion/common/src/cast.rs

@@ -24,8 +24,8 @@ use crate::{downcast_value, Result};
 use arrow::array::{
     BinaryViewArray, Decimal32Array, Decimal64Array, DurationMicrosecondArray,
     DurationMillisecondArray, DurationNanosecondArray, DurationSecondArray, Float16Array,
-    Int16Array, Int8Array, LargeBinaryArray, LargeStringArray, StringViewArray,
-    UInt16Array,
+    Int16Array, Int8Array, LargeBinaryArray, LargeListViewArray, LargeStringArray,
+    ListViewArray, StringViewArray, UInt16Array,
 };
 use arrow::{
     array::{
@@ -324,3 +324,13 @@ pub fn as_generic_string_array<T: OffsetSizeTrait>(
 ) -> Result<&GenericStringArray<T>> {
     Ok(downcast_value!(array, GenericStringArray, T))
 }
+
+// Downcast Array to ListViewArray
+pub fn as_list_view_array(array: &dyn Array) -> Result<&ListViewArray> {
+    Ok(downcast_value!(array, ListViewArray))
+}
+
+// Downcast Array to LargeListViewArray
+pub fn as_large_list_view_array(array: &dyn Array) -> Result<&LargeListViewArray> {
+    Ok(downcast_value!(array, LargeListViewArray))
+}

datafusion/functions-nested/src/reverse.rs

@@ -19,21 +19,26 @@
 
 use crate::utils::make_scalar_function;
 use arrow::array::{
-    Array, ArrayRef, Capacities, FixedSizeListArray, GenericListArray, MutableArrayData,
-    OffsetSizeTrait,
+    Array, ArrayRef, Capacities, FixedSizeListArray, GenericListArray,
+    GenericListViewArray, MutableArrayData, OffsetSizeTrait, UInt32Array,
+};
+use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
+use arrow::compute::take;
+use arrow::datatypes::DataType::{
+    FixedSizeList, LargeList, LargeListView, List, ListView, Null,
 };
-use arrow::buffer::OffsetBuffer;
-use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null};
 use arrow::datatypes::{DataType, FieldRef};
 use datafusion_common::cast::{
-    as_fixed_size_list_array, as_large_list_array, as_list_array,
+    as_fixed_size_list_array, as_large_list_array, as_large_list_view_array,
+    as_list_array, as_list_view_array,
 };
 use datafusion_common::{exec_err, utils::take_function_args, Result};
 use datafusion_expr::{
     ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
 };
 use datafusion_macros::user_doc;
 use std::any::Any;
+use std::mem::size_of;
 use std::sync::Arc;
 
 make_udf_expr_and_func!(
@@ -133,6 +138,14 @@ pub fn array_reverse_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
             fixed_size_array_reverse(array, field)
         }
         Null => Ok(Arc::clone(input_array)),
+        ListView(field) => {
+            let array = as_list_view_array(input_array)?;
+            list_view_reverse::<i32>(array, field)
+        }
+        LargeListView(field) => {
+            let array = as_large_list_view_array(input_array)?;
+            list_view_reverse::<i64>(array, field)
+        }
         array_type => exec_err!("array_reverse does not support type '{array_type}'."),
     }
 }
@@ -183,6 +196,75 @@ fn general_array_reverse<O: OffsetSizeTrait + TryFrom<i64>>(
     )?))
 }
 
+fn list_view_reverse<O: OffsetSizeTrait + TryFrom<i64>>(
+    array: &GenericListViewArray<O>,
+    field: &FieldRef,
+) -> Result<ArrayRef> {
+    let (_, offsets, sizes, values, nulls) = array.clone().into_parts();
+
+    // Construct indices, sizes and offsets for the reversed array by iterating over
+    // the list view array in the logical order, and reversing the order of the elements.
+    // We end up with a list view array where the elements are in order,
+    // even if the original array had elements out of order.
+    let mut indices: Vec<O> = Vec::with_capacity(values.len());
+    let mut new_sizes = Vec::with_capacity(sizes.len());
+    let mut new_offsets: Vec<O> = Vec::with_capacity(offsets.len());
+    let mut new_nulls =
+        Vec::with_capacity(nulls.clone().map(|nulls| nulls.len()).unwrap_or(0));
+    new_offsets.push(O::zero());
+    let has_nulls = nulls.is_some();
+    for (i, offset) in offsets.iter().enumerate().take(offsets.len()) {
+        // If this array is null, we set the new array to null with size 0 and continue
+        if let Some(ref nulls) = nulls {
+            if nulls.is_null(i) {
+                new_nulls.push(false); // null
+                new_sizes.push(O::zero());
+                new_offsets.push(new_offsets[i]);
+                continue;
+            } else {
+                new_nulls.push(true); // valid
+            }
+        }
+
+        // Each array is located at [offset, offset + size), so we collect indices in the reverse order
+        let array_start = offset.as_usize();
+        let array_end = array_start + sizes[i].as_usize();
+        for idx in (array_start..array_end).rev() {
+            indices.push(O::usize_as(idx));
+        }
+        new_sizes.push(sizes[i]);
+        if i < sizes.len() - 1 {
+            new_offsets.push(new_offsets[i] + sizes[i]);
+        }
+    }
+
+    // Materialize values from underlying array with take
+    let indices_array: ArrayRef = match size_of::<O>() {
+        4 => Arc::new(UInt32Array::from(
+            indices
+                .iter()
+                .map(|i| i.as_usize() as u32)
+                .collect::<Vec<_>>(),
+        )),
+        8 => Arc::new(arrow::array::UInt64Array::from(
+            indices
+                .iter()
+                .map(|i| i.as_usize() as u64)
+                .collect::<Vec<_>>(),
+        )),
+        _ => exec_err!("OffsetSizeTrait only supports i32 and i64")?,
+    };
+    let values_reversed = take(&values, &indices_array, None)?;
+
+    Ok(Arc::new(GenericListViewArray::<O>::try_new(
+        Arc::clone(field),
+        ScalarBuffer::from(new_offsets),
+        ScalarBuffer::from(new_sizes),
+        values_reversed,
+        has_nulls.then_some(NullBuffer::from(new_nulls)),
+    )?))
+}
+
 fn fixed_size_array_reverse(
     array: &FixedSizeListArray,
     field: &FieldRef,
@@ -219,3 +301,129 @@ fn fixed_size_array_reverse(
         Some(nulls.into()),
     )?))
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::reverse::list_view_reverse;
+    use arrow::{
+        array::{AsArray, Int32Array, LargeListViewArray, ListViewArray},
+        buffer::{NullBuffer, ScalarBuffer},
+        datatypes::{DataType, Field, Int32Type},
+    };
+    use std::sync::Arc;
+
+    #[test]
+    fn test_reverse_list_view_and_large_list_view() {
+        // ListView
+        let list_view = ListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![0, 1, 6, 6]),
+            ScalarBuffer::from(vec![1, 5, 0, 3]),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            Some(NullBuffer::from(vec![true, true, false, true])),
+        );
+        let result = list_view_reverse(
+            &list_view,
+            &Arc::new(Field::new("test", DataType::Int32, true)),
+        )
+        .unwrap();
+        let list_view_reversed: Vec<_> = result
+            .as_list_view::<i32>()
+            .iter()
+            .map(|x| x.map(|x| x.as_primitive::<Int32Type>().values().to_vec()))
+            .collect();
+
+        // LargeListView
+        let large_list_view = LargeListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![0, 1, 6, 6]),
+            ScalarBuffer::from(vec![1, 5, 0, 3]),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            Some(NullBuffer::from(vec![true, true, false, true])),
+        );
+        let result = list_view_reverse(
+            &large_list_view,
+            &Arc::new(Field::new("test", DataType::Int32, true)),
+        )
+        .unwrap();
+        let large_list_view_reversed: Vec<_> = result
+            .as_list_view::<i64>()
+            .iter()
+            .map(|x| x.map(|x| x.as_primitive::<Int32Type>().values().to_vec()))
+            .collect();
+
+        // Check results
+        let expected = vec![
+            Some(vec![1]),
+            Some(vec![6, 5, 4, 3, 2]),
+            None,
+            Some(vec![9, 8, 7]),
+        ];
+        assert_eq!(expected, list_view_reversed);
+        assert_eq!(expected, large_list_view_reversed);
+    }
+
+    #[test]
+    fn test_reverse_list_view_out_of_order() {
+        let list_view = ListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![6, 1, 6, 0]), // out of order
+            ScalarBuffer::from(vec![3, 5, 0, 1]),
+            Arc::new(Int32Array::from(vec![
+                1, // fourth array: offset 0, size 1
+                2, 3, 4, 5, 6, // second array: offset 1, size 5
+                // third array null but size 0
+                7, 8, 9, // first array: offset 6, size 3
+            ])),
+            Some(NullBuffer::from(vec![true, true, false, true])),
+        );
+        let list_view_reversed: Vec<_> = list_view_reverse(
+            &list_view,
+            &Arc::new(Field::new("test", DataType::Int32, true)),
+        )
+        .unwrap()
+        .as_list_view::<i32>()
+        .iter()
+        .map(|x| x.map(|x| x.as_primitive::<Int32Type>().values().to_vec()))
+        .collect();
+        let expected = vec![
+            Some(vec![9, 8, 7]),
+            Some(vec![6, 5, 4, 3, 2]),
+            None,
+            Some(vec![1]),
+        ];
+        assert_eq!(expected, list_view_reversed);
+    }
+
+    #[test]
+    fn test_reverse_list_view_with_nulls() {
+        let list_view = ListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![16, 1, 6, 0]), // out of order
+            ScalarBuffer::from(vec![3, 5, 10, 1]),
+            Arc::new(Int32Array::from(vec![
+                1, // fourth array: offset 0, size 1
+                2, 3, 4, 5, 6, // second array: offset 1, size 5
+                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, // third array: offset 6, size 10
+                7, 8, 9, // first array: offset 6, size 3
+            ])),
+            Some(NullBuffer::from(vec![true, true, false, true])),
+        );
+        let list_view_reversed: Vec<_> = list_view_reverse(
+            &list_view,
+            &Arc::new(Field::new("test", DataType::Int32, true)),
+        )
+        .unwrap()
+        .as_list_view::<i32>()
+        .iter()
+        .map(|x| x.map(|x| x.as_primitive::<Int32Type>().values().to_vec()))
+        .collect();
+        let expected = vec![
+            Some(vec![9, 8, 7]),
+            Some(vec![6, 5, 4, 3, 2]),
+            None,
+            Some(vec![1]),
+        ];
+        assert_eq!(expected, list_view_reversed);
+    }
+}