Initial plumbing into variant_array and builder

Author: friendlymatthew

parquet-variant-compute/benches/variant_kernels.rs

@@ -17,6 +17,7 @@
 
 use arrow::array::{Array, ArrayRef, StringArray};
 use arrow::util::test_util::seedable_rng;
+use arrow_schema::{DataType, Field, Fields};
 use criterion::{criterion_group, criterion_main, Criterion};
 use parquet_variant::{Variant, VariantBuilder};
 use parquet_variant_compute::variant_get::{variant_get, GetOptions};
@@ -27,6 +28,14 @@ use rand::Rng;
 use rand::SeedableRng;
 use std::fmt::Write;
 use std::sync::Arc;
+
+fn unshredded_schema_fields() -> Fields {
+    let metadata_field = Field::new("metadata", DataType::BinaryView, false);
+    let value_field = Field::new("value", DataType::BinaryView, false);
+
+    Fields::from(vec![metadata_field, value_field])
+}
+
 fn benchmark_batch_json_string_to_variant(c: &mut Criterion) {
     let input_array = StringArray::from_iter_values(json_repeated_struct(8000));
     let array_ref: ArrayRef = Arc::new(input_array);
@@ -93,7 +102,7 @@ pub fn variant_get_bench(c: &mut Criterion) {
     };
 
     c.bench_function("variant_get_primitive", |b| {
-        b.iter(|| variant_get(&input.clone(), options.clone()))
+        b.iter(|| variant_get(&input.clone(), options.clone(), unshredded_schema_fields()))
     });
 }
 
@@ -108,7 +117,7 @@ criterion_main!(benches);
 fn create_primitive_variant_array(size: usize) -> VariantArray {
     let mut rng = StdRng::seed_from_u64(42);
 
-    let mut variant_builder = VariantArrayBuilder::new(1);
+    let mut variant_builder = VariantArrayBuilder::try_new(1, unshredded_schema_fields()).unwrap();
 
     for _ in 0..size {
         let mut builder = VariantBuilder::new();

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(),
@@ -165,33 +165,50 @@ impl VariantSchema {
         })
     }
 
-    pub fn inner(self) -> Fields {
+    pub fn inner(&self) -> &Fields {
+        &self.inner
+    }
+
+    pub fn into_inner(self) -> Fields {
         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");
         };

parquet-variant-compute/src/variant_get.rs

@@ -49,7 +49,7 @@ pub fn variant_get(
         )));
     }
 
-    let mut builder = VariantArrayBuilder::try_new(variant_array.len(), shredded_schema).unwrap();
+    let mut builder = VariantArrayBuilder::try_new(variant_array.len(), shredded_schema)?;
     for i in 0..variant_array.len() {
         let new_variant = variant_array.value(i);
         // TODO: perf?