[Variant] Support array shredding into List/LargeList/ListView/LargeListView

parquet-variant-compute/src/arrow_to_variant.rs

@@ -17,8 +17,8 @@
 
 use crate::type_conversion::CastOptions;
 use arrow::array::{
-    Array, AsArray, FixedSizeListArray, GenericBinaryArray, GenericListArray, GenericListViewArray,
-    GenericStringArray, OffsetSizeTrait, PrimitiveArray,
+    Array, ArrayRef, AsArray, FixedSizeListArray, GenericBinaryArray, GenericListArray,
+    GenericListViewArray, GenericStringArray, OffsetSizeTrait, PrimitiveArray,
 };
 use arrow::compute::kernels::cast;
 use arrow::datatypes::{
@@ -553,20 +553,42 @@ impl<'a, L: ListLikeArray> ListArrowToVariantBuilder<'a, L> {
     }
 }
 
-/// Trait for list-like arrays that can provide element ranges
+/// Trait for list-like arrays that can provide common helpers
 pub(crate) trait ListLikeArray: Array {
+    type OffsetSize: OffsetSizeTrait;
+
     /// Get the values array
-    fn values(&self) -> &dyn Array;
+    fn values(&self) -> &ArrayRef;
+
+    /// Get the offsets backing the list values
+    #[cfg(test)]
+    fn value_offsets(&self) -> Option<&[Self::OffsetSize]>;
+
+    /// Size (number of values) for the element at `index`
+    #[cfg(test)]
+    fn value_size(&self, index: usize) -> Self::OffsetSize;
 
     /// Get the start and end indices for a list element
     fn element_range(&self, index: usize) -> Range<usize>;
 }
 
 impl<O: OffsetSizeTrait> ListLikeArray for GenericListArray<O> {
-    fn values(&self) -> &dyn Array {
+    type OffsetSize = O;
+
+    fn values(&self) -> &ArrayRef {
         self.values()
     }
 
+    #[cfg(test)]
+    fn value_offsets(&self) -> Option<&[Self::OffsetSize]> {
+        Some(GenericListArray::value_offsets(self))
+    }
+
+    #[cfg(test)]
+    fn value_size(&self, index: usize) -> Self::OffsetSize {
+        GenericListArray::value_length(self, index)
+    }
+
     fn element_range(&self, index: usize) -> Range<usize> {
         let offsets = self.offsets();
         let start = offsets[index].as_usize();
@@ -576,10 +598,22 @@ impl<O: OffsetSizeTrait> ListLikeArray for GenericListArray<O> {
 }
 
 impl<O: OffsetSizeTrait> ListLikeArray for GenericListViewArray<O> {
-    fn values(&self) -> &dyn Array {
+    type OffsetSize = O;
+
+    fn values(&self) -> &ArrayRef {
         self.values()
     }
 
+    #[cfg(test)]
+    fn value_offsets(&self) -> Option<&[Self::OffsetSize]> {
+        Some(GenericListViewArray::value_offsets(self))
+    }
+
+    #[cfg(test)]
+    fn value_size(&self, index: usize) -> Self::OffsetSize {
+        GenericListViewArray::value_size(self, index)
+    }
+
     fn element_range(&self, index: usize) -> Range<usize> {
         let offsets = self.value_offsets();
         let sizes = self.value_sizes();
@@ -590,10 +624,22 @@ impl<O: OffsetSizeTrait> ListLikeArray for GenericListViewArray<O> {
 }
 
 impl ListLikeArray for FixedSizeListArray {
-    fn values(&self) -> &dyn Array {
+    type OffsetSize = i32;
+
+    fn values(&self) -> &ArrayRef {
         self.values()
     }
 
+    #[cfg(test)]
+    fn value_offsets(&self) -> Option<&[Self::OffsetSize]> {
+        None
+    }
+
+    #[cfg(test)]
+    fn value_size(&self, _index: usize) -> Self::OffsetSize {
+        self.value_length()
+    }
+
     fn element_range(&self, index: usize) -> Range<usize> {
         let value_length = self.value_length().as_usize();
         let offset = index * value_length;