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/lib.rs

@@ -22,7 +22,7 @@ mod variant_array_builder;
 pub mod variant_get;
 
 pub use variant_array::VariantArray;
-pub use variant_array_builder::VariantArrayBuilder;
+pub use variant_array_builder::{VariantArrayBuilder, VariantArrayVariantBuilder};
 
 pub use from_json::batch_json_string_to_variant;
 pub use to_json::batch_variant_to_json_string;

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`]
@@ -37,13 +37,13 @@ use std::sync::Arc;
 /// ## Example:
 /// ```
 /// # use arrow::array::Array;
-/// # use parquet_variant::{Variant, VariantBuilder};
+/// # use parquet_variant::{Variant, VariantBuilder, VariantBuilderExt};
 /// # use parquet_variant_compute::VariantArrayBuilder;
 /// // Create a new VariantArrayBuilder with a capacity of 100 rows
 /// let mut builder = VariantArrayBuilder::new(100);
 /// // append variant values
 /// builder.append_variant(Variant::from(42));
-/// // append a null row
+/// // append a null row (note not a Variant::Null)
 /// builder.append_null();
 /// // append a pre-constructed metadata and value buffers
 /// let (metadata, value) = {
@@ -55,9 +55,14 @@ use std::sync::Arc;
 /// };
 /// builder.append_variant_buffers(&metadata, &value);
 ///
+/// // Use `variant_builder` method to write values directly to the output array
+/// let mut vb = builder.variant_builder();
+/// vb.append_value("Hello, World!");
+/// vb.finish(); // Note: call finish to write the variant to the buffers
+///
 /// // create the final VariantArray
 /// let variant_array = builder.build();
-/// assert_eq!(variant_array.len(), 3);
+/// assert_eq!(variant_array.len(), 4);
 /// // // Access the values
 /// // row 1 is not null and is an integer
 /// assert!(!variant_array.is_null(0));
@@ -67,6 +72,9 @@ use std::sync::Arc;
 /// // row 2 is not null and is an object
 /// assert!(!variant_array.is_null(2));
 /// assert!(variant_array.value(2).as_object().is_some());
+/// // row 3 is a string
+/// assert!(!variant_array.is_null(3));
+/// assert_eq!(variant_array.value(3), Variant::from("Hello, World!"));
 /// ```
 #[derive(Debug)]
 pub struct VariantArrayBuilder {
@@ -147,11 +155,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 +174,169 @@ 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);
+        VariantArrayVariantBuilder::new(self, metadata_buffer, value_buffer)
+    }
+}
+
+/// 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 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,
+}
+
+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<'a> VariantArrayVariantBuilder<'a> {
+    /// Constructs a new VariantArrayVariantBuilder
+    ///
+    /// 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,
+        metadata_buffer: Vec<u8>,
+        value_buffer: Vec<u8>,
+    ) -> Self {
+        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,
+        }
+    }
+
+    /// 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 (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 binary_view_array_from_buffers(
@@ -220,4 +388,82 @@ mod test {
             );
         }
     }
+
+    /// Test using sub builders to append variants
+    #[test]
+    fn test_variant_array_builder_variant_builder() {
+        let mut builder = VariantArrayBuilder::new(10);
+        builder.append_null(); // should not panic
+        builder.append_variant(Variant::from(42i32));
+
+        // let's make a sub-object in the next row
+        let mut sub_builder = builder.variant_builder();
+        let mut obj = sub_builder.new_object();
+        obj.insert("foo", "bar");
+        obj.finish().unwrap();
+        sub_builder.finish(); // must call finish to write the variant to the buffers
+
+        // append a new list
+        let mut sub_builder = builder.variant_builder();
+        let mut list_builder = sub_builder.new_list();
+        list_builder.append_value(Variant::from(1i32));
+        list_builder.append_value(Variant::from(2i32));
+        list_builder.finish();
+        sub_builder.finish();
+        let variant_array = builder.build();
+
+        assert_eq!(variant_array.len(), 4);
+        assert!(variant_array.is_null(0));
+        assert!(!variant_array.is_null(1));
+        assert_eq!(variant_array.value(1), Variant::from(42i32));
+        assert!(!variant_array.is_null(2));
+        let variant = variant_array.value(2);
+        let variant = variant.as_object().expect("variant to be an object");
+        assert_eq!(variant.get("foo").unwrap(), Variant::from("bar"));
+        assert!(!variant_array.is_null(3));
+        let variant = variant_array.value(3);
+        let list = variant.as_list().expect("variant to be a list");
+        assert_eq!(list.len(), 2);
+    }
+
+    /// Test using non-finished sub builders to append variants
+    #[test]
+    fn test_variant_array_builder_variant_builder_reset() {
+        let mut builder = VariantArrayBuilder::new(10);
+
+        // make a sub-object in the first row
+        let mut sub_builder = builder.variant_builder();
+        let mut obj = sub_builder.new_object();
+        obj.insert("foo", 1i32);
+        obj.finish().unwrap();
+        sub_builder.finish(); // must call finish to write the variant to the buffers
+
+        // start appending an object but don't finish
+        let mut sub_builder = builder.variant_builder();
+        let mut obj = sub_builder.new_object();
+        obj.insert("bar", 2i32);
+        obj.finish().unwrap();
+        drop(sub_builder); // drop the sub builder without finishing it
+
+        // make a third sub-object (this should reset the previous unfinished object)
+        let mut sub_builder = builder.variant_builder();
+        let mut obj = sub_builder.new_object();
+        obj.insert("baz", 3i32);
+        obj.finish().unwrap();
+        sub_builder.finish(); // must call finish to write the variant to the buffers
+
+        let variant_array = builder.build();
+
+        // only the two finished objects should be present
+        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!(!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));
+    }
 }