Avoid a copy in VariantArrayBuilder when creating JSON

Author: alamb

parquet-variant-compute/src/from_json.rs

@@ -21,7 +21,6 @@
 use crate::{VariantArray, VariantArrayBuilder};
 use arrow::array::{Array, ArrayRef, StringArray};
 use arrow_schema::ArrowError;
-use parquet_variant::VariantBuilder;
 use parquet_variant_json::json_to_variant;
 
 /// Parse a batch of JSON strings into a batch of Variants represented as
@@ -41,10 +40,10 @@ pub fn batch_json_string_to_variant(input: &ArrayRef) -> Result<VariantArray, Ar
             // The subfields are expected to be non-nullable according to the parquet variant spec.
             variant_array_builder.append_null();
         } else {
-            let mut vb = VariantBuilder::new();
+            let mut vb = variant_array_builder.variant_builder();
+            // parse JSON directly to the variant builder
             json_to_variant(input_string_array.value(i), &mut vb)?;
-            let (metadata, value) = vb.finish();
-            variant_array_builder.append_variant_buffers(&metadata, &value);
+            vb.finish()
         }
     }
     Ok(variant_array_builder.build())

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`]
@@ -147,11 +147,9 @@ impl VariantArrayBuilder {
 
     /// Append the [`Variant`] to the builder as the next row
     pub fn append_variant(&mut self, variant: Variant) {
-        // TODO make this more efficient by avoiding the intermediate buffers
-        let mut variant_builder = VariantBuilder::new();
-        variant_builder.append_value(variant);
-        let (metadata, value) = variant_builder.finish();
-        self.append_variant_buffers(&metadata, &value);
+        let mut direct_builder = self.variant_builder();
+        direct_builder.variant_builder.append_value(variant);
+        direct_builder.finish()
     }
 
     /// Append a metadata and values buffer to the builder
@@ -168,7 +166,156 @@ impl VariantArrayBuilder {
         self.value_buffer.extend_from_slice(value);
     }
 
-    // TODO: Return a Variant builder that will write to the underlying buffers (TODO)
+    /// Return a `VariantArrayVariantBuilder` that writes directly to the
+    /// buffers of this builder.
+    ///
+    /// You must call [`VariantArrayVariantBuilder::finish`] to complete the 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();
+        VariantArrayVariantBuilder {
+            finished: false,
+            metadata_offset,
+            value_offset,
+            variant_builder: VariantBuilder::new_with_buffers(metadata_buffer, value_buffer),
+            array_builder: self,
+        }
+    }
+}
+
+/// A `VariantBuilder` that writes directly to the buffers of a `VariantArrayBuilder`.
+///
+/// Note this struct implements [`VariantBuilderExt`], so it can be used
+/// as a drop-in replacement for [`VariantBuilder`] in most cases.
+///
+/// If [`Self::finish`] is not called, any changes will be rolled back
+///
+/// See [`VariantArrayBuilder::variant_builder`] for an example
+pub struct VariantArrayVariantBuilder<'a> {
+    /// was finish called?
+    finished: bool,
+    /// starting metadata offset in the underlying buffers
+    metadata_offset: usize,
+    /// starting value offset
+    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 variant, temporarily owns the buffers from `array_builder`
+    variant_builder: VariantBuilder,
+}
+
+impl<'a> VariantBuilderExt for VariantArrayVariantBuilder<'a> {
+    fn append_value<'m, 'v>(&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
+    }
+
+    /// 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;
+        // Note: buffers are returned and replaced in the drop impl
+    }
+}
+
+impl<'a> Drop for VariantArrayVariantBuilder<'a> {
+    /// If the builder was not finished, roll back any changes made to the
+    /// underlying buffers (by truncating them)
+    fn drop(&mut self) {
+        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).finish();
+
+        // Sanity Check: if the buffers got smaller, something went wrong (previous data was lost)
+        let metadata_len = metadata_buffer
+            .len()
+            .checked_sub(metadata_offset)
+            .expect("metadata length decreased unexpectedly");
+        let value_len = value_buffer
+            .len()
+            .checked_sub(value_offset)
+            .expect("value length decreased unexpectedly");
+
+        if self.finished {
+            // if the object was finished, commit the changes by putting the
+            // offsets and lengths into the parent array builder.
+            self.array_builder
+                .metadata_locations
+                .push((metadata_offset, metadata_len));
+            self.array_builder
+                .value_locations
+                .push((value_offset, value_len));
+            self.array_builder.nulls.append_non_null();
+        } else {
+            // if the object was not finished, truncate the buffers to the
+            // original offsets to roll back any changes. Note this is fast
+            // because truncate doesn't free any memory: it just has to drop
+            // elements (but 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 binary_view_array_from_buffers(

parquet-variant-json/src/from_json.rs

@@ -18,7 +18,7 @@
 //! Module for parsing JSON strings as Variant
 
 use arrow_schema::ArrowError;
-use parquet_variant::{ListBuilder, ObjectBuilder, Variant, VariantBuilder, VariantBuilderExt};
+use parquet_variant::{ListBuilder, ObjectBuilder, Variant, VariantBuilderExt};
 use serde_json::{Number, Value};
 
 /// Converts a JSON string to Variant using [`VariantBuilder`]. The resulting `value` and `metadata`
@@ -62,15 +62,15 @@ use serde_json::{Number, Value};
 /// assert_eq!(json_result, serde_json::to_string(&json_value)?);
 /// # Ok::<(), Box<dyn std::error::Error>>(())
 /// ```
-pub fn json_to_variant(json: &str, builder: &mut VariantBuilder) -> Result<(), ArrowError> {
+pub fn json_to_variant(json: &str, builder: &mut impl VariantBuilderExt) -> Result<(), ArrowError> {
     let json: Value = serde_json::from_str(json)
         .map_err(|e| ArrowError::InvalidArgumentError(format!("JSON format error: {e}")))?;
 
     build_json(&json, builder)?;
     Ok(())
 }
 
-fn build_json(json: &Value, builder: &mut VariantBuilder) -> Result<(), ArrowError> {
+fn build_json(json: &Value, builder: &mut impl VariantBuilderExt) -> Result<(), ArrowError> {
     append_json(json, builder)?;
     Ok(())
 }
@@ -101,7 +101,7 @@ fn variant_from_number<'m, 'v>(n: &Number) -> Result<Variant<'m, 'v>, ArrowError
 
 fn append_json<'m, 'v>(
     json: &'v Value,
-    builder: &mut impl VariantBuilderExt<'m, 'v>,
+    builder: &mut impl VariantBuilderExt,
 ) -> Result<(), ArrowError> {
     match json {
         Value::Null => builder.append_value(Variant::Null),
@@ -137,8 +137,8 @@ struct ObjectFieldBuilder<'o, 'v, 's> {
     builder: &'o mut ObjectBuilder<'v>,
 }
 
-impl<'m, 'v> VariantBuilderExt<'m, 'v> for ObjectFieldBuilder<'_, '_, '_> {
-    fn append_value(&mut self, value: impl Into<Variant<'m, 'v>>) {
+impl VariantBuilderExt for ObjectFieldBuilder<'_, '_, '_> {
+    fn append_value<'m, 'v>(&mut self, value: impl Into<Variant<'m, 'v>>) {
         self.builder.insert(self.key, value);
     }
 

parquet-variant/src/builder.rs

@@ -1047,16 +1047,16 @@ impl Drop for ObjectBuilder<'_> {
 ///
 /// Allows users to append values to a [`VariantBuilder`], [`ListBuilder`] or
 /// [`ObjectBuilder`]. using the same interface.
-pub trait VariantBuilderExt<'m, 'v> {
-    fn append_value(&mut self, value: impl Into<Variant<'m, 'v>>);
+pub trait VariantBuilderExt {
+    fn append_value<'m, 'v>(&mut self, value: impl Into<Variant<'m, 'v>>);
 
     fn new_list(&mut self) -> ListBuilder;
 
     fn new_object(&mut self) -> ObjectBuilder;
 }
 
-impl<'m, 'v> VariantBuilderExt<'m, 'v> for ListBuilder<'_> {
-    fn append_value(&mut self, value: impl Into<Variant<'m, 'v>>) {
+impl VariantBuilderExt for ListBuilder<'_> {
+    fn append_value<'m, 'v>(&mut self, value: impl Into<Variant<'m, 'v>>) {
         self.append_value(value);
     }
 
@@ -1069,8 +1069,8 @@ impl<'m, 'v> VariantBuilderExt<'m, 'v> for ListBuilder<'_> {
     }
 }
 
-impl<'m, 'v> VariantBuilderExt<'m, 'v> for VariantBuilder {
-    fn append_value(&mut self, value: impl Into<Variant<'m, 'v>>) {
+impl VariantBuilderExt for VariantBuilder {
+    fn append_value<'m, 'v>(&mut self, value: impl Into<Variant<'m, 'v>>) {
         self.append_value(value);
     }