Add cast support for (Large)ListView <-> (Large)List

Author: vegarsti

arrow-array/src/array/list_view_array.rs

@@ -24,7 +24,10 @@ use std::sync::Arc;
 
 use crate::array::{make_array, print_long_array};
 use crate::iterator::GenericListViewArrayIter;
-use crate::{Array, ArrayAccessor, ArrayRef, FixedSizeListArray, OffsetSizeTrait, new_empty_array};
+use crate::{
+    Array, ArrayAccessor, ArrayRef, FixedSizeListArray, GenericListArray, OffsetSizeTrait,
+    new_empty_array,
+};
 
 /// A [`GenericListViewArray`] of variable size lists, storing offsets as `i32`.
 pub type ListViewArray = GenericListViewArray<i32>;
@@ -454,6 +457,36 @@ impl<OffsetSize: OffsetSizeTrait> std::fmt::Debug for GenericListViewArray<Offse
     }
 }
 
+impl<OffsetSize: OffsetSizeTrait> From<GenericListArray<OffsetSize>>
+    for GenericListViewArray<OffsetSize>
+{
+    fn from(value: GenericListArray<OffsetSize>) -> Self {
+        let field = match value.data_type() {
+            DataType::List(f) | DataType::LargeList(f) => f.clone(),
+            _ => panic!(
+                "Expected infallible creation of GenericListViewArray from GenericList failed"
+            ),
+        };
+
+        let offsets = value.value_offsets();
+        let len = offsets.len() - 1;
+        let mut sizes = Vec::with_capacity(len);
+        let mut view_offsets = Vec::with_capacity(len);
+        for (i, offset) in offsets.iter().enumerate().take(len) {
+            view_offsets.push(*offset);
+            sizes.push(value.value_length(i));
+        }
+
+        Self::new(
+            field,
+            ScalarBuffer::from(view_offsets),
+            ScalarBuffer::from(sizes),
+            value.values().clone(),
+            value.nulls().cloned(),
+        )
+    }
+}
+
 impl<OffsetSize: OffsetSizeTrait> From<GenericListViewArray<OffsetSize>> for ArrayData {
     fn from(array: GenericListViewArray<OffsetSize>) -> Self {
         let len = array.len();

arrow-cast/src/cast/list.rs

@@ -50,6 +50,15 @@ pub(crate) fn cast_single_element_fixed_size_list_to_values(
     cast_with_options(values, to, cast_options)
 }
 
+pub(crate) fn cast_list_to_list_view<OffsetSize>(array: &dyn Array) -> Result<ArrayRef, ArrowError>
+where
+    OffsetSize: OffsetSizeTrait,
+{
+    let list = array.as_list::<OffsetSize>();
+    let list_view: GenericListViewArray<OffsetSize> = list.clone().into();
+    Ok(Arc::new(list_view))
+}
+
 pub(crate) fn cast_fixed_size_list_to_list<OffsetSize>(
     array: &dyn Array,
 ) -> Result<ArrayRef, ArrowError>
@@ -160,6 +169,40 @@ pub(crate) fn cast_list_values<O: OffsetSizeTrait>(
     )?))
 }
 
+/// Helper function to cast the values in a list view to a list
+pub(crate) fn cast_list_view_values<O: OffsetSizeTrait>(
+    array: &dyn Array,
+    to: &FieldRef,
+    cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError> {
+    let list_view = array.as_list_view::<O>();
+    let list_view_offsets = list_view.offsets();
+    let sizes = list_view.sizes();
+    let source_values = list_view.values();
+    let mut indices = Vec::with_capacity(list_view.values().len());
+    let mut offsets = Vec::with_capacity(list_view.len() + 1);
+    offsets.push(O::usize_as(0));
+    for i in 0..list_view.len() {
+        let offset = list_view_offsets[i].as_usize();
+        let size = sizes[i].as_usize();
+        let end = offset + size;
+        for j in offset..end {
+            indices.push(j as i32);
+        }
+        offsets.push(O::usize_as(indices.len()));
+    }
+    let indices_array = Int32Array::from(indices);
+    let values = arrow_select::take::take(source_values, &indices_array, None)?;
+    let values = cast_with_options(&values, to.data_type(), cast_options)?;
+    let offsets = OffsetBuffer::new(offsets.into());
+    Ok(Arc::new(GenericListArray::<O>::try_new(
+        to.clone(),
+        offsets,
+        values,
+        list_view.nulls().cloned(),
+    )?))
+}
+
 /// Cast the container type of List/Largelist array along with the inner datatype
 pub(crate) fn cast_list<I: OffsetSizeTrait, O: OffsetSizeTrait>(
     array: &dyn Array,

arrow-cast/src/cast/mod.rs

@@ -153,7 +153,20 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
         (List(list_from) | LargeList(list_from), FixedSizeList(list_to, _)) => {
             can_cast_types(list_from.data_type(), list_to.data_type())
         }
+        (List(list_from), ListView(list_to)) => {
+            can_cast_types(list_from.data_type(), list_to.data_type())
+        }
         (List(_), _) => false,
+        (ListView(list_from), List(list_to)) => {
+            can_cast_types(list_from.data_type(), list_to.data_type())
+        }
+        (LargeListView(list_from), LargeList(list_to)) => {
+            can_cast_types(list_from.data_type(), list_to.data_type())
+        }
+        (LargeList(list_from), LargeListView(list_to)) => {
+            can_cast_types(list_from.data_type(), list_to.data_type())
+        }
+
         (FixedSizeList(list_from, _), List(list_to))
         | (FixedSizeList(list_from, _), LargeList(list_to)) => {
             can_cast_types(list_from.data_type(), list_to.data_type())
@@ -864,9 +877,40 @@ pub fn cast_with_options(
             let array = array.as_list::<i64>();
             cast_list_to_fixed_size_list::<i64>(array, field, *size, cast_options)
         }
-        (List(_) | LargeList(_), _) => match to_type {
+        (ListView(_), List(_)) => match to_type {
+            List(list_to) => cast_list_view_values::<i32>(array, list_to, cast_options),
+            _ => Err(ArrowError::CastError(
+                "Cannot cast list view to non-list data types".to_string(),
+            )),
+        },
+        (LargeListView(_), LargeList(_)) => match to_type {
+            LargeList(list_to) => cast_list_view_values::<i64>(array, list_to, cast_options),
+            _ => Err(ArrowError::CastError(
+                "Cannot cast list view to non-list data types".to_string(),
+            )),
+        },
+        (List(list_from) | LargeList(list_from), _) => match to_type {
             Utf8 => value_to_string::<i32>(array, cast_options),
             LargeUtf8 => value_to_string::<i64>(array, cast_options),
+            ListView(list_to) => {
+                if list_to.data_type() != list_from.data_type() {
+                    Err(ArrowError::CastError(
+                        "Cannot cast list to list view with different inner type".to_string(),
+                    ))
+                } else {
+                    cast_list_to_list_view::<i32>(array)
+                }
+            }
+            LargeListView(list_to) => {
+                if list_to.data_type() != list_from.data_type() {
+                    Err(ArrowError::CastError(
+                        "Cannot cast large list to large list view with different inner type"
+                            .to_string(),
+                    ))
+                } else {
+                    cast_list_to_list_view::<i64>(array)
+                }
+            }
             _ => Err(ArrowError::CastError(
                 "Cannot cast list to non-list data types".to_string(),
             )),
@@ -2677,8 +2721,8 @@ mod tests {
     use super::*;
     use DataType::*;
     use arrow_array::{Int64Array, RunArray, StringArray};
-    use arrow_buffer::i256;
     use arrow_buffer::{Buffer, IntervalDayTime, NullBuffer};
+    use arrow_buffer::{ScalarBuffer, i256};
     use arrow_schema::{DataType, Field};
     use chrono::NaiveDate;
     use half::f16;
@@ -11866,4 +11910,212 @@ mod tests {
         // Verify the run-ends were cast correctly (run ends at 3, 6, 9)
         assert_eq!(run_array.run_ends().values(), &[3i64, 6i64, 9i64]);
     }
+
+    #[test]
+    fn test_cast_list_view_to_list() {
+        let list_view = ListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![0, 3, 6]),
+            ScalarBuffer::from(vec![3, 3, 3]),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+        let cast_result = cast(
+            &list_view,
+            &DataType::List(Arc::new(Field::new("a", DataType::Int32, false))),
+        )
+        .unwrap();
+        let got_list = cast_result.as_any().downcast_ref::<ListArray>().unwrap();
+
+        let mut offsets = OffsetBufferBuilder::new(0);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        let expected_list = ListArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            offsets.finish(),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+        assert_eq!(got_list, &expected_list);
+    }
+
+    #[test]
+    fn test_cast_list_to_list_view() {
+        let mut offsets = OffsetBufferBuilder::new(3);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        let list = ListArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            offsets.finish(),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+
+        let expected_list_view = ListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![0, 3, 6]),
+            ScalarBuffer::from(vec![3, 3, 3]),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+        let cast_result = cast(
+            &list,
+            &DataType::ListView(Arc::new(Field::new("a", DataType::Int32, false))),
+        )
+        .unwrap();
+
+        let got_list_view = cast_result
+            .as_any()
+            .downcast_ref::<ListViewArray>()
+            .unwrap();
+        assert_eq!(got_list_view, &expected_list_view);
+    }
+
+    #[test]
+    fn test_cast_large_list_view_to_large_list() {
+        let list_view = LargeListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![0, 3, 6]),
+            ScalarBuffer::from(vec![3, 3, 3]),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+        let cast_result = cast(
+            &list_view,
+            &DataType::LargeList(Arc::new(Field::new("a", DataType::Int32, false))),
+        )
+        .unwrap();
+        let got_list = cast_result
+            .as_any()
+            .downcast_ref::<LargeListArray>()
+            .unwrap();
+
+        let mut offsets = OffsetBufferBuilder::new(0);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        let expected_list = LargeListArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            offsets.finish(),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+        assert_eq!(got_list, &expected_list);
+    }
+
+    #[test]
+    fn test_cast_large_list_to_large_list_view() {
+        let mut offsets = OffsetBufferBuilder::new(3);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        let list = LargeListArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            offsets.finish(),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+
+        let expected_list_view = LargeListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![0, 3, 6]),
+            ScalarBuffer::from(vec![3, 3, 3]),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+        let cast_result = cast(
+            &list,
+            &DataType::LargeListView(Arc::new(Field::new("a", DataType::Int32, false))),
+        )
+        .unwrap();
+
+        let got_list_view = cast_result
+            .as_any()
+            .downcast_ref::<LargeListViewArray>()
+            .unwrap();
+        assert_eq!(got_list_view, &expected_list_view);
+    }
+
+    #[test]
+    fn test_cast_list_view_to_list_out_of_order() {
+        let list_view = ListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![0, 6, 3]),
+            ScalarBuffer::from(vec![3, 3, 3]),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9])),
+            None,
+        );
+        let cast_result = cast(
+            &list_view,
+            &DataType::List(Arc::new(Field::new("a", DataType::Int32, false))),
+        )
+        .unwrap();
+        let got_list = cast_result.as_any().downcast_ref::<ListArray>().unwrap();
+        let mut offsets = OffsetBufferBuilder::new(0);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        offsets.push_length(3);
+        let expected_list = ListArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            offsets.finish(),
+            Arc::new(Int32Array::from(vec![1, 2, 3, 7, 8, 9, 4, 5, 6])),
+            None,
+        );
+        assert_eq!(got_list, &expected_list);
+    }
+
+    #[test]
+    fn test_cast_list_view_to_list_overlapping() {
+        let list_view = ListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![0, 0]),
+            ScalarBuffer::from(vec![1, 2]),
+            Arc::new(Int32Array::from(vec![1, 2])),
+            None,
+        );
+        let cast_result = cast(
+            &list_view,
+            &DataType::List(Arc::new(Field::new("a", DataType::Int32, false))),
+        )
+        .unwrap();
+        let got_list = cast_result.as_any().downcast_ref::<ListArray>().unwrap();
+        let mut offsets = OffsetBufferBuilder::new(0);
+        offsets.push_length(1);
+        offsets.push_length(2);
+        let expected_list = ListArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            offsets.finish(),
+            Arc::new(Int32Array::from(vec![1, 1, 2])),
+            None,
+        );
+        assert_eq!(got_list, &expected_list);
+    }
+
+    #[test]
+    fn test_cast_list_view_to_list_empty() {
+        let empty_array: Vec<i32> = vec![];
+        let list_view = ListViewArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            ScalarBuffer::from(vec![]),
+            ScalarBuffer::from(vec![]),
+            Arc::new(Int32Array::from(empty_array.clone())),
+            None,
+        );
+        let cast_result = cast(
+            &list_view,
+            &DataType::List(Arc::new(Field::new("a", DataType::Int32, false))),
+        )
+        .unwrap();
+        let got_list = cast_result.as_any().downcast_ref::<ListArray>().unwrap();
+        let offsets = OffsetBuffer::new_empty();
+        let expected_list = ListArray::new(
+            Arc::new(Field::new("a", DataType::Int32, false)),
+            offsets,
+            Arc::new(Int32Array::from(empty_array.clone())),
+            None,
+        );
+        assert_eq!(got_list, &expected_list);
+    }
 }