impl VariantBuilderExt for VariantArrayVariantBuilder

Author: alamb

parquet-variant-compute/src/variant_array_builder.rs

@@ -20,7 +20,7 @@
 use crate::VariantArray;
 use arrow::array::{ArrayRef, BinaryViewArray, BinaryViewBuilder, NullBufferBuilder, StructArray};
 use arrow_schema::{DataType, Field, Fields};
-use parquet_variant::{Variant, VariantBuilder};
+use parquet_variant::{ListBuilder, ObjectBuilder, Variant, VariantBuilder, VariantBuilderExt};
 use std::sync::Arc;
 
 /// A builder for [`VariantArray`]
@@ -166,14 +166,41 @@ impl VariantArrayBuilder {
         self.value_buffer.extend_from_slice(value);
     }
 
-    /// Return a DirectVariantBuilder that writes directly to the buffers of this builder.
-    pub fn variant_builder(&mut self) -> DirectVariantBuilder {
+    /// Return a VariantArrayVariantBuilder that writes directly to the buffers of this builder.
+    ///
+    /// # Example
+    /// ```
+    /// # use parquet_variant::{Variant, VariantBuilder, VariantBuilderExt};
+    /// # use parquet_variant_compute::{VariantArray, VariantArrayBuilder};
+    /// let mut array_builder = VariantArrayBuilder::new(10);
+    ///
+    /// // First row has a string
+    /// let mut variant_builder = array_builder.variant_builder();
+    /// variant_builder.append_value("Hello, World!");
+    /// // must call finish to write the variant to the buffers
+    /// variant_builder.finish();
+    ///
+    /// // Second row is an object
+    /// let mut variant_builder = array_builder.variant_builder();
+    /// let mut obj = variant_builder.new_object();
+    /// obj.insert("my_field", 42i64);
+    /// obj.finish().unwrap();
+    /// variant_builder.finish();
+    ///
+    /// // finalize the array
+    /// let variant_array: VariantArray = array_builder.build();
+    ///
+    /// // verify what we wrote is still there
+    /// assert_eq!(variant_array.value(0), Variant::from("Hello, World!"));
+    /// assert!(variant_array.value(1).as_object().is_some());
+    ///  ```
+    pub fn variant_builder(&mut self) -> VariantArrayVariantBuilder {
         // append directly into the metadata and value buffers
         let metadata_buffer = std::mem::take(&mut self.metadata_buffer);
         let value_buffer = std::mem::take(&mut self.value_buffer);
         let metadata_offset = metadata_buffer.len();
         let value_offset = value_buffer.len();
-        DirectVariantBuilder {
+        VariantArrayVariantBuilder {
             finished: false,
             metadata_offset,
             value_offset,
@@ -185,19 +212,36 @@ impl VariantArrayBuilder {
 
 /// A `VariantBuilder` that writes directly to the buffers of a `VariantArrayBuilder`.
 ///
-/// See [`VariantArray::variant_builder`] for an example
-pub struct DirectVariantBuilder<'a> {
+/// Note this struct implements [`VariantBuilderExt`], so it can be used
+/// as a drop-in replacement for [`VariantBuilder`] in most cases.
+///
+/// See [`VariantArrayBuilder::variant_builder`] for an example
+pub struct VariantArrayVariantBuilder<'a> {
     /// was finish called?
     finished: bool,
-    /// starting metadata offset
+    /// starting metadata offset in the underlying buffers
     metadata_offset: usize,
     /// starting value offset
     value_offset: usize,
     array_builder: &'a mut VariantArrayBuilder,
     variant_builder: VariantBuilder,
 }
 
-impl DirectVariantBuilder<'_> {
+impl<'a, 'm, 'v> VariantBuilderExt<'m, 'v> for VariantArrayVariantBuilder<'a> {
+    fn append_value(&mut self, value: impl Into<Variant<'m, 'v>>) {
+        self.variant_builder.append_value(value);
+    }
+
+    fn new_list(&mut self) -> ListBuilder {
+        self.variant_builder.new_list()
+    }
+
+    fn new_object(&mut self) -> ObjectBuilder {
+        self.variant_builder.new_object()
+    }
+}
+
+impl VariantArrayVariantBuilder<'_> {
     /// Return a reference to the underlying `VariantBuilder`
     pub fn inner(&self) -> &VariantBuilder {
         &self.variant_builder
@@ -244,14 +288,14 @@ impl DirectVariantBuilder<'_> {
     }
 }
 
-impl<'a> Drop for DirectVariantBuilder<'a> {
+impl<'a> Drop for VariantArrayVariantBuilder<'a> {
     fn drop(&mut self) {
         if self.finished {
             // if the object was finished, we do not need to do anything
             return;
         }
         // if the object was not finished, we need to reset any partial state put the buffers back
-        println!("DirectVariantBuilder::drop");
+        println!("VariantArrayVariantBuilder::drop");
         let variant_builder = std::mem::take(&mut self.variant_builder);
         let (mut metadata_buffer, mut value_buffer) = variant_builder.finish();
         assert!(