Append complex variants

Author: friendlymatthew

parquet-variant/src/builder.rs

@@ -15,7 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 use crate::decoder::{VariantBasicType, VariantPrimitiveType};
-use crate::{ShortString, Variant, VariantDecimal16, VariantDecimal4, VariantDecimal8};
+use crate::{
+    ShortString, Variant, VariantDecimal16, VariantDecimal4, VariantDecimal8, VariantList,
+    VariantObject,
+};
 use arrow_schema::ArrowError;
 use indexmap::{IndexMap, IndexSet};
 use std::collections::HashSet;
@@ -192,8 +195,7 @@ impl ValueBuffer {
         self.0.len()
     }
 
-    fn append_non_nested_value<'m, 'd, T: Into<Variant<'m, 'd>>>(&mut self, value: T) {
-        let variant = value.into();
+    fn append_variant<'m, 'd>(&mut self, variant: Variant<'m, 'd>) {
         match variant {
             Variant::Null => self.append_null(),
             Variant::BooleanTrue => self.append_bool(true),
@@ -213,14 +215,14 @@ impl ValueBuffer {
             Variant::Binary(v) => self.append_binary(v),
             Variant::String(s) => self.append_string(s),
             Variant::ShortString(s) => self.append_short_string(s),
-            Variant::Object(_) | Variant::List(_) => {
-                unreachable!(
-                    "Nested values are handled specially by ObjectBuilder and ListBuilder"
-                );
-            }
+            _ => unreachable!("Objects and lists must be appended using VariantBuilder::append_object and VariantBuilder::append_list"),
         }
     }
 
+    fn append_non_nested_value<'m, 'd, T: Into<Variant<'m, 'd>>>(&mut self, value: T) {
+        self.append_variant(value.into());
+    }
+
     /// Writes out the header byte for a variant object or list
     fn append_header(&mut self, header_byte: u8, is_large: bool, num_items: usize) {
         let buf = self.inner_mut();
@@ -697,6 +699,73 @@ impl VariantBuilder {
         ObjectBuilder::new(parent_state, validate_unique_fields)
     }
 
+    /// Appends a [`VariantObject`] to the builder.
+    fn append_object<'m, 'v>(&mut self, object: VariantObject<'m, 'v>) -> Result<(), ArrowError> {
+        let (parent_state, validate_unique_fields) = self.parent_state();
+
+        let mut obj_builder = ObjectBuilder::new(parent_state, validate_unique_fields);
+
+        for (field_name, variant) in object.iter() {
+            obj_builder.insert(field_name, variant);
+        }
+
+        obj_builder.finish()?;
+
+        Ok(())
+    }
+
+    /// Appends a [`VariantObject`] to the builder with full validation during iteration.
+    ///
+    /// Recursively validates all nested variants in the object during iteration.
+    fn try_append_object<'m, 'v>(
+        &mut self,
+        object: VariantObject<'m, 'v>,
+    ) -> Result<(), ArrowError> {
+        let (parent_state, validate_unique_fields) = self.parent_state();
+
+        let mut obj_builder = ObjectBuilder::new(parent_state, validate_unique_fields);
+
+        for res in object.iter_try() {
+            let (field_name, variant) = res?;
+
+            obj_builder.insert(field_name, variant);
+        }
+
+        obj_builder.finish()?;
+
+        Ok(())
+    }
+
+    /// Appends a [`VariantList`] to the builder.
+    fn append_list<'m, 'v>(&mut self, list: VariantList<'m, 'v>) {
+        let (parent_state, validate_unique_fields) = self.parent_state();
+
+        let mut list_builder = ListBuilder::new(parent_state, validate_unique_fields);
+
+        for variant in list.iter() {
+            list_builder.append_value(variant);
+        }
+
+        list_builder.finish();
+    }
+
+    /// Appends a [`VariantList`] to the builder with full validation during iteration.
+    ///
+    /// Recursively validates all nested variants in the list during iteration.
+    fn try_append_list<'m, 'v>(&mut self, list: VariantList<'m, 'v>) -> Result<(), ArrowError> {
+        let (parent_state, validate_unique_fields) = self.parent_state();
+
+        let mut list_builder = ListBuilder::new(parent_state, validate_unique_fields);
+
+        for variant in list.iter_try() {
+            list_builder.append_value(variant?);
+        }
+
+        list_builder.finish();
+
+        Ok(())
+    }
+
     /// Append a non-nested value to the builder.
     ///
     /// # Example
@@ -707,7 +776,13 @@ impl VariantBuilder {
     /// builder.append_value(42i8);
     /// ```
     pub fn append_value<'m, 'd, T: Into<Variant<'m, 'd>>>(&mut self, value: T) {
-        self.buffer.append_non_nested_value(value);
+        let variant = value.into();
+
+        match variant {
+            Variant::Object(obj) => self.append_object(obj),
+            Variant::List(list) => self.append_list(list),
+            primitive => self.buffer.append_variant(primitive),
+        }
     }
 
     /// Finish the builder and return the metadata and value buffers.
@@ -2170,4 +2245,45 @@ mod tests {
         let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();
         assert_eq!(variant, Variant::Int8(2));
     }
+
+    #[test]
+    fn test_append_object() {
+        let (m1, v1) = make_object();
+        let variant = Variant::new(&m1, &v1);
+
+        let mut builder = VariantBuilder::new();
+        builder.append_value(variant.clone());
+        let (metadata, value) = builder.finish();
+        assert_eq!(variant, Variant::new(&metadata, &value));
+    }
+
+    /// make an object variant
+    fn make_object() -> (Vec<u8>, Vec<u8>) {
+        let mut builder = VariantBuilder::new();
+
+        let mut obj = builder.new_object();
+        obj.insert("a", true);
+        obj.finish().unwrap();
+        builder.finish()
+    }
+
+    #[test]
+    fn test_append_list() {
+        let (m1, v1) = make_list();
+        let variant = Variant::new(&m1, &v1);
+        let mut builder = VariantBuilder::new();
+        builder.append_value(variant.clone());
+        let (metadata, value) = builder.finish();
+        assert_eq!(variant, Variant::new(&metadata, &value));
+    }
+
+    /// make a simple List variant
+    fn make_list() -> (Vec<u8>, Vec<u8>) {
+        let mut builder = VariantBuilder::new();
+        let mut list = builder.new_list();
+        list.append_value(1234);
+        list.append_value("a string value");
+        list.finish();
+        builder.finish()
+    }
 }

parquet-variant/src/variant/metadata.rs

@@ -127,7 +127,7 @@ impl VariantMetadataHeader {
 /// [Variant Spec]: https://github.com/apache/parquet-format/blob/master/VariantEncoding.md#metadata-encoding
 #[derive(Debug, Clone, PartialEq)]
 pub struct VariantMetadata<'m> {
-    bytes: &'m [u8],
+    pub(crate) bytes: &'m [u8],
     header: VariantMetadataHeader,
     dictionary_size: u32,
     first_value_byte: u32,