Initial plumbing into variant_array and builder

Author: friendlymatthew

parquet-variant-compute/src/from_json.rs

@@ -78,7 +78,7 @@ mod test {
         let variant_array = batch_json_string_to_variant(&array_ref).unwrap();
 
         let metadata_array = variant_array.metadata_field().as_binary_view();
-        let value_array = variant_array.value_field().as_binary_view();
+        let value_array = variant_array.value_field().unwrap().as_binary_view();
 
         // Compare row 0
         assert!(!variant_array.is_null(0));

parquet-variant-compute/src/shredding.rs

@@ -22,7 +22,7 @@ pub const METADATA: &str = "metadata";
 pub const VALUE: &str = "value";
 pub const TYPED_VALUE: &str = "typed_value";
 
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Clone)]
 pub enum ValueSchema {
     MissingValue,
     Value(usize),
@@ -33,7 +33,7 @@ pub enum ValueSchema {
     },
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct VariantSchema {
     inner: Fields,
 
@@ -54,7 +54,7 @@ impl VariantSchema {
         let (metadata_idx, metadata_field) = fields
             .iter()
             .enumerate()
-            .find(|(i, f)| f.name() == METADATA)
+            .find(|(_, f)| f.name() == METADATA)
             .ok_or_else(|| {
                 ArrowError::InvalidArgumentError(
                     "Invalid VariantArray: StructArray must contain a 'metadata' field".to_string(),
@@ -169,29 +169,42 @@ impl VariantSchema {
         self.inner
     }
 
+    pub fn metadata_idx(&self) -> usize {
+        self.metadata_idx
+    }
+
     pub fn metadata(&self) -> &FieldRef {
         self.inner.get(self.metadata_idx).unwrap()
     }
 
-    pub fn value(&self) -> Option<&FieldRef> {
+    pub fn value_idx(&self) -> Option<usize> {
         match self.value_schema {
             ValueSchema::MissingValue => None,
             ValueSchema::ShreddedValue(_) => None,
-            ValueSchema::Value(value_idx) => self.inner.get(value_idx),
-            ValueSchema::PartiallyShredded { value_idx, .. } => self.inner.get(value_idx),
+            ValueSchema::Value(value_idx) => Some(value_idx),
+            ValueSchema::PartiallyShredded { value_idx, .. } => Some(value_idx),
         }
     }
 
-    pub fn shredded_value(&self) -> Option<&FieldRef> {
+    pub fn value(&self) -> Option<&FieldRef> {
+        self.value_idx().map(|i| self.inner.get(i).unwrap())
+    }
+
+    pub fn shredded_value_idx(&self) -> Option<usize> {
         match self.value_schema {
             ValueSchema::MissingValue => None,
             ValueSchema::Value(_) => None,
-            ValueSchema::ShreddedValue(shredded_idx) => self.inner.get(shredded_idx),
+            ValueSchema::ShreddedValue(shredded_idx) => Some(shredded_idx),
             ValueSchema::PartiallyShredded {
                 shredded_value_idx, ..
-            } => self.inner.get(shredded_value_idx),
+            } => Some(shredded_value_idx),
         }
     }
+
+    pub fn shredded_value(&self) -> Option<&FieldRef> {
+        self.shredded_value_idx()
+            .map(|i| self.inner.get(i).unwrap())
+    }
 }
 
 #[cfg(test)]

parquet-variant-compute/src/variant_array.rs

@@ -24,6 +24,8 @@ use parquet_variant::Variant;
 use std::any::Any;
 use std::sync::Arc;
 
+use crate::shredding::VariantSchema;
+
 /// An array of Parquet [`Variant`] values
 ///
 /// A [`VariantArray`] wraps an Arrow [`StructArray`] that stores the underlying
@@ -60,11 +62,7 @@ pub struct VariantArray {
     /// int8.
     inner: StructArray,
 
-    /// Reference to the metadata column of inner
-    metadata_ref: ArrayRef,
-
-    /// Reference to the value column of inner
-    value_ref: ArrayRef,
+    variant_schema: VariantSchema,
 }
 
 impl VariantArray {
@@ -94,23 +92,11 @@ impl VariantArray {
             ));
         };
 
-        // todo, remove this since we already do it in validate_shredded_schema
-        let Some(metadata_field) = VariantArray::find_metadata_field(inner) else {
-            return Err(ArrowError::InvalidArgumentError(
-                "Invalid VariantArray: StructArray must contain a 'metadata' field".to_string(),
-            ));
-        };
-
-        let Some(value_field) = VariantArray::find_value_field(inner) else {
-            return Err(ArrowError::InvalidArgumentError(
-                "Invalid VariantArray: StructArray must contain a 'value' field".to_string(),
-            ));
-        };
+        let variant_schema = VariantSchema::try_new(inner.fields().clone())?;
 
         Ok(Self {
             inner: inner.clone(),
-            metadata_ref: metadata_field,
-            value_ref: value_field,
+            variant_schema,
         })
     }
 
@@ -126,34 +112,41 @@ impl VariantArray {
 
     /// Return the [`Variant`] instance stored at the given row
     ///
-    /// Panics if the index is out of bounds.
+    /// Panics if the index is out of bounds or value array does not exist.
     ///
     /// Note: Does not do deep validation of the [`Variant`], so it is up to the
     /// caller to ensure that the metadata and value were constructed correctly.
+    ///
+    /// Todo: reconstruct partially shredded or shredded variants
     pub fn value(&self, index: usize) -> Variant {
         let metadata = self.metadata_field().as_binary_view().value(index);
-        let value = self.value_field().as_binary_view().value(index);
+        let value = self
+            .value_field()
+            .expect("value field does not exist")
+            .as_binary_view()
+            .value(index);
         Variant::new(metadata, value)
     }
 
-    fn find_metadata_field(array: &StructArray) -> Option<ArrayRef> {
-        array.column_by_name("metadata").cloned()
-    }
-
-    fn find_value_field(array: &StructArray) -> Option<ArrayRef> {
-        array.column_by_name("value").cloned()
-    }
-
     /// Return a reference to the metadata field of the [`StructArray`]
     pub fn metadata_field(&self) -> &ArrayRef {
-        // spec says fields order is not guaranteed, so we search by name
-        &self.metadata_ref
+        let metadata_idx = self.variant_schema.metadata_idx();
+
+        self.inner.column(metadata_idx)
     }
 
     /// Return a reference to the value field of the `StructArray`
-    pub fn value_field(&self) -> &ArrayRef {
-        // spec says fields order is not guaranteed, so we search by name
-        &self.value_ref
+    pub fn value_field(&self) -> Option<&ArrayRef> {
+        self.variant_schema
+            .value_idx()
+            .map(|i| self.inner.column(i))
+    }
+
+    /// Return a reference to the shredded value field of the `StructArray`
+    pub fn shredded_value_field(&self) -> Option<&ArrayRef> {
+        self.variant_schema
+            .shredded_value_idx()
+            .map(|i| self.inner.column(i))
     }
 }
 
@@ -176,12 +169,9 @@ impl Array for VariantArray {
 
     fn slice(&self, offset: usize, length: usize) -> ArrayRef {
         let slice = self.inner.slice(offset, length);
-        let met = self.metadata_ref.slice(offset, length);
-        let val = self.value_ref.slice(offset, length);
         Arc::new(Self {
             inner: slice,
-            metadata_ref: met,
-            value_ref: val,
+            variant_schema: self.variant_schema.clone(),
         })
     }
 

parquet-variant-compute/src/variant_array_builder.rs

@@ -19,7 +19,7 @@
 
 use crate::{shredding::VariantSchema, VariantArray};
 use arrow::array::{ArrayRef, BinaryViewArray, BinaryViewBuilder, NullBufferBuilder, StructArray};
-use arrow_schema::{ArrowError, DataType, Field, Fields};
+use arrow_schema::{ArrowError, Fields};
 use parquet_variant::{Variant, VariantBuilder};
 use std::sync::Arc;
 
@@ -191,6 +191,7 @@ fn binary_view_array_from_buffers(
 mod test {
     use super::*;
     use arrow::array::Array;
+    use arrow_schema::{DataType, Field};
 
     /// Test that both the metadata and value buffers are non nullable
     #[test]
@@ -212,7 +213,7 @@ mod test {
 
         // the metadata and value fields of non shredded variants should not be null
         assert!(variant_array.metadata_field().nulls().is_none());
-        assert!(variant_array.value_field().nulls().is_none());
+        assert!(variant_array.value_field().unwrap().nulls().is_none());
         let DataType::Struct(fields) = variant_array.data_type() else {
             panic!("Expected VariantArray to have Struct data type");
         };