[Variant] VariantArrayBuilder uses MetadataBuilder and ValueBuilder (#8206)

# Which issue does this PR close?

- Closes #8205

# Rationale for this change

`VariantArrayBuilder` had a very complex choreography with the
`VariantBuilder` API, that required lots of manual drop glue to deal
with ownership transfers between it and the `VariantArrayVariantBuilder`
it delegates the actual work to. Rework the whole thing to use a
(now-reusable) `MetadataBuilder` and `ValueBuilder`, with rollbacks
largely handled by `ParentState` -- just like the other builders in the
parquet-variant crate.

# What changes are included in this PR?

Five changes (curated as five commits that reviewers may want to examine
individually):
1. Make a bunch of parquet-variant builder infrastructure public, so
that `VariantArrayBuilder` can access it from the
parquet-variant-compute crate.
2. Make `MetadataBuilder` reusable. Its `finish` method appends the
bytes of a new serialized metadata dictionary to the underlying buffer
and resets the remaining builder state. The builder is thus ready to
create a brand new metadata dictionary whose serialized bytes will also
be appended to the underlying buffer once finished.
3. Rework `VariantArrayBuilder` to use `MetadataBuilder` and
`ValueBuilder`, coordinated via `ParentState`. This is the main feature
of the PR and also the most complicated/subtle.
4. Delete now-unused code that had been added previously in order to
support the old implementation of `VariantArrayBuilder`.
5. Add missing doc comments for now-public types and methods


# Are these changes tested?

Existing variant array builder tests cover the change.

# Are there any user-facing changes?

A lot of builder-related types and methods from the parquet-variant
crate are now public.

Author: scovich

parquet-variant-compute/src/variant_array_builder.rs

@@ -20,7 +20,8 @@
 use crate::VariantArray;
 use arrow::array::{ArrayRef, BinaryViewArray, BinaryViewBuilder, NullBufferBuilder, StructArray};
 use arrow_schema::{ArrowError, DataType, Field, Fields};
-use parquet_variant::{ListBuilder, ObjectBuilder, Variant, VariantBuilder, VariantBuilderExt};
+use parquet_variant::{ListBuilder, ObjectBuilder, Variant, VariantBuilderExt};
+use parquet_variant::{MetadataBuilder, ParentState, ValueBuilder};
 use std::sync::Arc;
 
 /// A builder for [`VariantArray`]
@@ -72,12 +73,12 @@ use std::sync::Arc;
 pub struct VariantArrayBuilder {
     /// Nulls
     nulls: NullBufferBuilder,
-    /// buffer for all the metadata
-    metadata_buffer: Vec<u8>,
+    /// builder for all the metadata
+    metadata_builder: MetadataBuilder,
     /// ending offset for each serialized metadata dictionary in the buffer
     metadata_offsets: Vec<usize>,
-    /// buffer for values
-    value_buffer: Vec<u8>,
+    /// builder for values
+    value_builder: ValueBuilder,
     /// ending offset for each serialized variant value in the buffer
     value_offsets: Vec<usize>,
     /// The fields of the final `StructArray`
@@ -95,9 +96,9 @@ impl VariantArrayBuilder {
 
         Self {
             nulls: NullBufferBuilder::new(row_capacity),
-            metadata_buffer: Vec::new(), // todo allocation capacity
+            metadata_builder: MetadataBuilder::default(),
             metadata_offsets: Vec::with_capacity(row_capacity),
-            value_buffer: Vec::new(),
+            value_builder: ValueBuilder::new(),
             value_offsets: Vec::with_capacity(row_capacity),
             fields: Fields::from(vec![metadata_field, value_field]),
         }
@@ -107,15 +108,17 @@ impl VariantArrayBuilder {
     pub fn build(self) -> VariantArray {
         let Self {
             mut nulls,
-            metadata_buffer,
+            metadata_builder,
             metadata_offsets,
-            value_buffer,
+            value_builder,
             value_offsets,
             fields,
         } = self;
 
+        let metadata_buffer = metadata_builder.into_inner();
         let metadata_array = binary_view_array_from_buffers(metadata_buffer, metadata_offsets);
 
+        let value_buffer = value_builder.into_inner();
         let value_array = binary_view_array_from_buffers(value_buffer, value_offsets);
 
         // The build the final struct array
@@ -136,14 +139,14 @@ impl VariantArrayBuilder {
     pub fn append_null(&mut self) {
         self.nulls.append_null();
         // The subfields are expected to be non-nullable according to the parquet variant spec.
-        self.metadata_offsets.push(self.metadata_buffer.len());
-        self.value_offsets.push(self.value_buffer.len());
+        self.metadata_offsets.push(self.metadata_builder.offset());
+        self.value_offsets.push(self.value_builder.offset());
     }
 
     /// Append the [`Variant`] to the builder as the next row
     pub fn append_variant(&mut self, variant: Variant) {
         let mut direct_builder = self.variant_builder();
-        direct_builder.variant_builder.append_value(variant);
+        direct_builder.append_value(variant);
         direct_builder.finish()
     }
 
@@ -194,32 +197,23 @@ impl VariantArrayBuilder {
 ///
 /// See [`VariantArrayBuilder::variant_builder`] for an example
 pub struct VariantArrayVariantBuilder<'a> {
-    /// was finish called?
-    finished: bool,
-    /// starting offset in the variant_builder's `metadata` buffer
-    metadata_offset: usize,
-    /// starting offset in the variant_builder's `value` buffer
-    value_offset: usize,
-    /// Parent array builder that this variant builder writes to. Buffers
-    /// have been moved into the variant builder, and must be returned on
-    /// drop
-    array_builder: &'a mut VariantArrayBuilder,
-    /// Builder for the in progress variant value, temporarily owns the buffers
-    /// from `array_builder`
-    variant_builder: VariantBuilder,
+    parent_state: ParentState<'a>,
+    metadata_offsets: &'a mut Vec<usize>,
+    value_offsets: &'a mut Vec<usize>,
+    nulls: &'a mut NullBufferBuilder,
 }
 
 impl VariantBuilderExt for VariantArrayVariantBuilder<'_> {
     fn append_value<'m, 'v>(&mut self, value: impl Into<Variant<'m, 'v>>) {
-        self.variant_builder.append_value(value);
+        ValueBuilder::append_variant(self.parent_state(), value.into());
     }
 
     fn try_new_list(&mut self) -> Result<ListBuilder<'_>, ArrowError> {
-        Ok(self.variant_builder.new_list())
+        Ok(ListBuilder::new(self.parent_state(), false))
     }
 
     fn try_new_object(&mut self) -> Result<ObjectBuilder<'_>, ArrowError> {
-        Ok(self.variant_builder.new_object())
+        Ok(ObjectBuilder::new(self.parent_state(), false))
     }
 }
 
@@ -228,103 +222,40 @@ impl<'a> VariantArrayVariantBuilder<'a> {
     ///
     /// Note this is not public as this is a structure that is logically
     /// part of the [`VariantArrayBuilder`] and relies on its internal structure
-    fn new(array_builder: &'a mut VariantArrayBuilder) -> Self {
-        // append directly into the metadata and value buffers
-        let metadata_buffer = std::mem::take(&mut array_builder.metadata_buffer);
-        let value_buffer = std::mem::take(&mut array_builder.value_buffer);
-        let metadata_offset = metadata_buffer.len();
-        let value_offset = value_buffer.len();
+    fn new(builder: &'a mut VariantArrayBuilder) -> Self {
+        let parent_state =
+            ParentState::variant(&mut builder.value_builder, &mut builder.metadata_builder);
         VariantArrayVariantBuilder {
-            finished: false,
-            metadata_offset,
-            value_offset,
-            variant_builder: VariantBuilder::new_with_buffers(metadata_buffer, value_buffer),
-            array_builder,
+            parent_state,
+            metadata_offsets: &mut builder.metadata_offsets,
+            value_offsets: &mut builder.value_offsets,
+            nulls: &mut builder.nulls,
         }
     }
 
-    /// Return a reference to the underlying `VariantBuilder`
-    pub fn inner(&self) -> &VariantBuilder {
-        &self.variant_builder
-    }
-
-    /// Return a mutable reference to the underlying `VariantBuilder`
-    pub fn inner_mut(&mut self) -> &mut VariantBuilder {
-        &mut self.variant_builder
-    }
-
     /// Called to finish the in progress variant and write it to the underlying
     /// buffers
     ///
     /// Note if you do not call finish, on drop any changes made to the
     /// underlying buffers will be rolled back.
     pub fn finish(mut self) {
-        self.finished = true;
-
-        let metadata_offset = self.metadata_offset;
-        let value_offset = self.value_offset;
-        // get the buffers back from the variant builder
-        let (metadata_buffer, value_buffer) = std::mem::take(&mut self.variant_builder).finish();
-
-        // Sanity Check: if the buffers got smaller, something went wrong (previous data was lost)
-        assert!(
-            metadata_offset <= metadata_buffer.len(),
-            "metadata length decreased unexpectedly"
-        );
-        assert!(
-            value_offset <= value_buffer.len(),
-            "value length decreased unexpectedly"
-        );
-
-        // commit the changes by putting the
-        // ending offsets into the parent array builder.
-        let builder = &mut self.array_builder;
-        builder.metadata_offsets.push(metadata_buffer.len());
-        builder.value_offsets.push(value_buffer.len());
-        builder.nulls.append_non_null();
+        // Record the ending offsets after finishing metadata and finish the parent state.
+        let (value_builder, metadata_builder) = self.parent_state.value_and_metadata_builders();
+        self.metadata_offsets.push(metadata_builder.finish());
+        self.value_offsets.push(value_builder.offset());
+        self.nulls.append_non_null();
+        self.parent_state.finish();
+    }
 
-        // put the buffers back into the array builder
-        builder.metadata_buffer = metadata_buffer;
-        builder.value_buffer = value_buffer;
+    fn parent_state(&mut self) -> ParentState<'_> {
+        let (value_builder, metadata_builder) = self.parent_state.value_and_metadata_builders();
+        ParentState::variant(value_builder, metadata_builder)
     }
 }
 
+// Empty Drop to help with borrow checking - warns users if they forget to call finish()
 impl Drop for VariantArrayVariantBuilder<'_> {
-    /// If the builder was not finished, roll back any changes made to the
-    /// underlying buffers (by truncating them)
-    fn drop(&mut self) {
-        if self.finished {
-            return;
-        }
-
-        // if the object was not finished, need to rollback any changes by
-        // truncating the buffers to the original offsets
-        let metadata_offset = self.metadata_offset;
-        let value_offset = self.value_offset;
-
-        // get the buffers back from the variant builder
-        let (mut metadata_buffer, mut value_buffer) =
-            std::mem::take(&mut self.variant_builder).into_buffers();
-
-        // Sanity Check: if the buffers got smaller, something went wrong (previous data was lost) so panic immediately
-        metadata_buffer
-            .len()
-            .checked_sub(metadata_offset)
-            .expect("metadata length decreased unexpectedly");
-        value_buffer
-            .len()
-            .checked_sub(value_offset)
-            .expect("value length decreased unexpectedly");
-
-        // Note this truncate is fast because truncate doesn't free any memory:
-        // it just has to drop elements (and u8 doesn't have a destructor)
-        metadata_buffer.truncate(metadata_offset);
-        value_buffer.truncate(value_offset);
-
-        // put the buffers back into the array builder
-        self.array_builder.metadata_buffer = metadata_buffer;
-        self.array_builder.value_buffer = value_buffer;
-    }
+    fn drop(&mut self) {}
 }
 
 fn binary_view_array_from_buffers(buffer: Vec<u8>, offsets: Vec<usize>) -> BinaryViewArray {
@@ -457,12 +388,18 @@ mod test {
         assert_eq!(variant_array.len(), 2);
         assert!(!variant_array.is_null(0));
         let variant = variant_array.value(0);
-        let variant = variant.as_object().expect("variant to be an object");
-        assert_eq!(variant.get("foo").unwrap(), Variant::from(1i32));
+        assert_eq!(
+            variant.get_object_field("foo"),
+            Some(Variant::from(1i32)),
+            "Expected an object with field \"foo\", got: {variant:?}"
+        );
 
         assert!(!variant_array.is_null(1));
         let variant = variant_array.value(1);
-        let variant = variant.as_object().expect("variant to be an object");
-        assert_eq!(variant.get("baz").unwrap(), Variant::from(3i32));
+        assert_eq!(
+            variant.get_object_field("baz"),
+            Some(Variant::from(3i32)),
+            "Expected an object with field \"baz\", got: {variant:?}"
+        );
     }
 }