Create TestListLikeArray instead of updating ListLikeArray

Author: liamzwbao

parquet-variant-compute/src/arrow_to_variant.rs

@@ -553,42 +553,20 @@ impl<'a, L: ListLikeArray> ListArrowToVariantBuilder<'a, L> {
     }
 }
 
-/// Trait for list-like arrays that can provide common helpers
+/// Trait for list-like arrays that can provide element ranges
 pub(crate) trait ListLikeArray: Array {
-    type OffsetSize: OffsetSizeTrait;
-
     /// Get the values 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> {
-    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();
@@ -598,22 +576,10 @@ impl<O: OffsetSizeTrait> ListLikeArray for GenericListArray<O> {
 }
 
 impl<O: OffsetSizeTrait> ListLikeArray for GenericListViewArray<O> {
-    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();
@@ -624,22 +590,10 @@ impl<O: OffsetSizeTrait> ListLikeArray for GenericListViewArray<O> {
 }
 
 impl ListLikeArray for FixedSizeListArray {
-    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;

parquet-variant-compute/src/shred_variant.rs

@@ -928,9 +928,39 @@ mod tests {
         builder.build()
     }
 
+    trait TestListLikeArray: ListLikeArray {
+        type OffsetSize: OffsetSizeTrait;
+        fn value_offsets(&self) -> Option<&[Self::OffsetSize]>;
+        fn value_size(&self, index: usize) -> Self::OffsetSize;
+    }
+
+    impl<O: OffsetSizeTrait> TestListLikeArray for GenericListArray<O> {
+        type OffsetSize = O;
+
+        fn value_offsets(&self) -> Option<&[Self::OffsetSize]> {
+            Some(GenericListArray::value_offsets(self))
+        }
+
+        fn value_size(&self, index: usize) -> Self::OffsetSize {
+            GenericListArray::value_length(self, index)
+        }
+    }
+
+    impl<O: OffsetSizeTrait> TestListLikeArray for GenericListViewArray<O> {
+        type OffsetSize = O;
+
+        fn value_offsets(&self) -> Option<&[Self::OffsetSize]> {
+            Some(GenericListViewArray::value_offsets(self))
+        }
+
+        fn value_size(&self, index: usize) -> Self::OffsetSize {
+            GenericListViewArray::value_size(self, index)
+        }
+    }
+
     fn downcast_list_like_array<O: OffsetSizeTrait>(
         array: &VariantArray,
-    ) -> &dyn ListLikeArray<OffsetSize = O> {
+    ) -> &dyn TestListLikeArray<OffsetSize = O> {
         let typed_value = array.typed_value_field().unwrap();
         if let Some(list) = typed_value.as_any().downcast_ref::<GenericListArray<O>>() {
             list
@@ -2456,21 +2486,18 @@ mod tests {
 
     #[test]
     fn test_variant_schema_builder_with_shred_variant() {
-        let mut builder = VariantArrayBuilder::new(3);
-        builder
-            .new_object()
-            .with_field("time", 1234567890i64)
-            .with_field("hostname", "server1")
-            .with_field("extra", 42)
-            .finish();
-        builder
-            .new_object()
-            .with_field("time", 9876543210i64)
-            .with_field("hostname", "server2")
-            .finish();
-        builder.append_null();
-
-        let input = builder.build();
+        let input = build_variant_array(vec![
+            VariantRow::Object(vec![
+                ("time", VariantValue::from(1234567890i64)),
+                ("hostname", VariantValue::from("server1")),
+                ("extra", VariantValue::from(42)),
+            ]),
+            VariantRow::Object(vec![
+                ("time", VariantValue::from(9876543210i64)),
+                ("hostname", VariantValue::from("server2")),
+            ]),
+            VariantRow::Null,
+        ]);
 
         let shredding_type = ShreddedSchemaBuilder::default()
             .with_path("time", &DataType::Int64)