Impl Array for VariantArray

Author: friendlymatthew

parquet-variant-compute/benches/variant_kernels.rs

@@ -85,7 +85,7 @@ fn benchmark_batch_json_string_to_variant(c: &mut Criterion) {
 
 pub fn variant_get_bench(c: &mut Criterion) {
     let variant_array = create_primitive_variant_array(8192);
-    let input = ArrayRef::from(variant_array);
+    let input = Arc::new(variant_array) as ArrayRef;
 
     let options = GetOptions {
         path: vec![].into(),

parquet-variant-compute/src/shred_variant.rs

@@ -22,7 +22,7 @@ use crate::variant_to_arrow::{
     PrimitiveVariantToArrowRowBuilder, make_primitive_variant_to_arrow_row_builder,
 };
 use crate::{VariantArray, VariantValueArrayBuilder};
-use arrow::array::{ArrayRef, BinaryViewArray, NullBufferBuilder};
+use arrow::array::{Array, ArrayRef, BinaryViewArray, NullBufferBuilder};
 use arrow::buffer::NullBuffer;
 use arrow::compute::CastOptions;
 use arrow::datatypes::{DataType, Fields};

parquet-variant-compute/src/variant_array.rs

@@ -159,6 +159,7 @@ impl ExtensionType for VariantType {
 ///
 /// ```
 /// # use arrow_schema::{Schema, Field, DataType};
+/// # use arrow::array::Array;
 /// # use parquet_variant::Variant;
 /// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray, VariantType};
 /// # fn get_variant_array() -> VariantArray {
@@ -385,54 +386,74 @@ impl VariantArray {
         .with_extension_type(VariantType)
     }
 
-    /// Returns a new DataType representing this VariantArray's inner type
-    pub fn data_type(&self) -> &DataType {
+    /// Is the element at index null?
+    pub fn is_null(&self, index: usize) -> bool {
+        self.nulls().is_some_and(|n| n.is_null(index))
+    }
+
+    /// Is the element at index valid (not null)?
+    pub fn is_valid(&self, index: usize) -> bool {
+        !self.is_null(index)
+    }
+}
+
+impl From<VariantArray> for StructArray {
+    fn from(variant_array: VariantArray) -> Self {
+        variant_array.into_inner()
+    }
+}
+
+impl Array for VariantArray {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn to_data(&self) -> arrow::array::ArrayData {
+        self.inner.to_data()
+    }
+
+    fn into_data(self) -> arrow::array::ArrayData {
+        self.inner.into_data()
+    }
+
+    fn data_type(&self) -> &DataType {
         self.inner.data_type()
     }
 
-    pub fn slice(&self, offset: usize, length: usize) -> Self {
+    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
         let inner = self.inner.slice(offset, length);
         let metadata = self.metadata.slice(offset, length);
         let shredding_state = self.shredding_state.slice(offset, length);
-        Self {
+
+        Arc::new(Self {
             inner,
             metadata,
             shredding_state,
-        }
+        })
     }
 
-    pub fn len(&self) -> usize {
+    fn len(&self) -> usize {
         self.inner.len()
     }
 
-    pub fn is_empty(&self) -> bool {
+    fn is_empty(&self) -> bool {
         self.inner.is_empty()
     }
 
-    pub fn nulls(&self) -> Option<&NullBuffer> {
-        self.inner.nulls()
-    }
-
-    /// Is the element at index null?
-    pub fn is_null(&self, index: usize) -> bool {
-        self.nulls().is_some_and(|n| n.is_null(index))
+    fn offset(&self) -> usize {
+        self.inner.offset()
     }
 
-    /// Is the element at index valid (not null)?
-    pub fn is_valid(&self, index: usize) -> bool {
-        !self.is_null(index)
+    fn nulls(&self) -> Option<&NullBuffer> {
+        self.inner.nulls()
     }
-}
 
-impl From<VariantArray> for StructArray {
-    fn from(variant_array: VariantArray) -> Self {
-        variant_array.into_inner()
+    fn get_buffer_memory_size(&self) -> usize {
+        self.inner.get_buffer_memory_size()
     }
-}
 
-impl From<VariantArray> for ArrayRef {
-    fn from(variant_array: VariantArray) -> Self {
-        Arc::new(variant_array.into_inner())
+    fn get_array_memory_size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.inner.get_buffer_memory_size()
     }
 }
 

parquet-variant-compute/src/variant_get.rs

@@ -196,7 +196,7 @@ fn shredded_get_path(
 
     // If our caller did not request any specific type, we can just return whatever we landed on.
     let Some(as_field) = as_field else {
-        return Ok(ArrayRef::from(target));
+        return Ok(Arc::new(target) as ArrayRef);
     };
 
     // Structs are special. Recurse into each field separately, hoping to follow the shredding even
@@ -314,7 +314,8 @@ mod test {
     fn single_variant_get_test(input_json: &str, path: VariantPath, expected_json: &str) {
         // Create input array from JSON string
         let input_array_ref: ArrayRef = Arc::new(StringArray::from(vec![Some(input_json)]));
-        let input_variant_array_ref = ArrayRef::from(json_to_variant(&input_array_ref).unwrap());
+        let input_variant_array_ref =
+            Arc::new(json_to_variant(&input_array_ref).unwrap()) as ArrayRef;
 
         let result =
             variant_get(&input_variant_array_ref, GetOptions::new_with_path(path)).unwrap();
@@ -452,9 +453,9 @@ mod test {
                     .with_field("value", Arc::new(values), true)
                     .with_nulls(nulls)
                     .build();
-                ArrayRef::from(
-                    VariantArray::try_new(&struct_array).expect("should create variant array"),
-                )
+
+                Arc::new(VariantArray::try_new(&struct_array).expect("should create variant array"))
+                    as ArrayRef
             }
         };
     }
@@ -779,9 +780,10 @@ mod test {
                     .with_field("typed_value", Arc::new(typed_value), true)
                     .build();
 
-                VariantArray::try_new(&struct_array)
-                    .expect("should create variant array")
-                    .into()
+                let variant_array =
+                    VariantArray::try_new(&struct_array).expect("should create variant array");
+
+                Arc::new(variant_array) as ArrayRef
             }
         };
     }
@@ -1189,9 +1191,10 @@ mod test {
                     .with_nulls(nulls)
                     .build();
 
-                ArrayRef::from(
-                    VariantArray::try_new(&struct_array).expect("should create variant array"),
-                )
+                let variant_array =
+                    VariantArray::try_new(&struct_array).expect("should create variant array");
+
+                Arc::new(variant_array) as ArrayRef
             }
         };
     }
@@ -1642,7 +1645,8 @@ mod test {
             }
         }
 
-        ArrayRef::from(builder.build())
+        let variant_array = builder.build();
+        Arc::new(variant_array) as _
     }
 
     /// Create test data for depth 1 (single nested field)
@@ -1672,7 +1676,7 @@ mod test {
             }
         }
 
-        ArrayRef::from(builder.build())
+        Arc::new(builder.build()) as _
     }
 
     /// Create test data for depth 2 (double nested field)
@@ -1713,7 +1717,7 @@ mod test {
             }
         }
 
-        ArrayRef::from(builder.build())
+        Arc::new(builder.build()) as _
     }
 
     /// Create simple shredded test data for depth 0 using a simplified working pattern
@@ -2182,7 +2186,7 @@ mod test {
             cast_options: CastOptions::default(),
         };
 
-        let variant_array_ref = ArrayRef::from(variant_array);
+        let variant_array_ref = Arc::new(variant_array) as ArrayRef;
         let result = variant_get(&variant_array_ref, options).unwrap();
 
         // Verify the result is a StructArray
@@ -2260,7 +2264,7 @@ mod test {
             cast_options: CastOptions::default(),
         };
 
-        let variant_array_ref = ArrayRef::from(variant_array);
+        let variant_array_ref = Arc::new(variant_array) as ArrayRef;
         let result_nullable = variant_get(&variant_array_ref, options_nullable).unwrap();
 
         // Verify we get an Int32Array with nulls for cast failures
@@ -2315,7 +2319,7 @@ mod test {
 
         // Create variant array again since we moved it
         let variant_array_2 = json_to_variant(&string_array).unwrap();
-        let variant_array_ref_2 = ArrayRef::from(variant_array_2);
+        let variant_array_ref_2 = Arc::new(variant_array_2) as ArrayRef;
         let result_non_nullable = variant_get(&variant_array_ref_2, options_non_nullable).unwrap();
         let int32_result_2 = result_non_nullable
             .as_any()
@@ -2632,7 +2636,7 @@ mod test {
             cast_options: CastOptions::default(),
         };
 
-        let variant_array_ref = ArrayRef::from(variant_array);
+        let variant_array_ref = Arc::new(variant_array) as ArrayRef;
         let result = variant_get(&variant_array_ref, options);
 
         // Should fail with NotYetImplemented when the row builder tries to handle struct type

parquet-variant-compute/src/variant_to_arrow.rs

@@ -469,6 +469,6 @@ impl VariantToBinaryVariantArrowRowBuilder {
             self.nulls.finish(),
         );
 
-        Ok(ArrayRef::from(variant_array))
+        Ok(Arc::new(variant_array) as ArrayRef)
     }
 }

parquet/src/variant.rs

@@ -147,7 +147,7 @@ mod tests {
     use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
     use crate::file::reader::ChunkReader;
     use arrow::util::test_util::parquet_test_data;
-    use arrow_array::{ArrayRef, RecordBatch};
+    use arrow_array::{Array, ArrayRef, RecordBatch};
     use arrow_schema::Schema;
     use bytes::Bytes;
     use parquet_variant::{Variant, VariantBuilderExt};
@@ -237,7 +237,7 @@ mod tests {
     fn variant_array_to_batch(array: VariantArray) -> RecordBatch {
         let field = array.field("data");
         let schema = Schema::new(vec![field]);
-        RecordBatch::try_new(Arc::new(schema), vec![ArrayRef::from(array)]).unwrap()
+        RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array) as ArrayRef]).unwrap()
     }
 
     /// writes a RecordBatch to memory buffer and returns the buffer

parquet/tests/variant_integration.rs

@@ -23,6 +23,7 @@
 //!
 //! Inspired by the arrow-go implementation: <https://github.com/apache/arrow-go/pull/455/files>
 
+use arrow::array::Array;
 use arrow::util::test_util::parquet_test_data;
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use parquet_variant::{Variant, VariantMetadata};