Add VariantArray and VariantArrayBuilder for constructing Arrow Arrays of Variants (#7905)

# Which issue does this PR close?

- Part of #6736
- Part of #7895


# Rationale for this change

As we begin to add operations on Variants stored in arrays, we need some
better abstractions of working with those arrays

This PR builds on the great work of @harshmotw-db in
#7884 to start adding t


# What changes are included in this PR?

1. Add `VariantArray` that wraps a `StructArray` and adds useful
accessors
2. Add `VariantArrayBuilder` as described in
#7895 to construct
`VariantArrays`
2. rework `batch_json_string_to_variant` to use the new builder and
array wrapper
Note while these APIs have no shredding support yet, I think shredding
can be added in a straightforward way

# Are these changes tested?

Yes, unit tests and doc examples are included


# Are there any user-facing changes?

New VariantArray and VariantArrayBuilder

Author: alamb

parquet-variant-compute/src/from_json.rs

@@ -18,109 +18,42 @@
 //! Module for transforming a batch of JSON strings into a batch of Variants represented as
 //! STRUCT<metadata: BINARY, value: BINARY>
 
-use std::sync::Arc;
-
-use arrow::array::{Array, ArrayRef, BinaryArray, BooleanBufferBuilder, StringArray, StructArray};
-use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
-use arrow::datatypes::{DataType, Field};
+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;
 
-fn variant_arrow_repr() -> DataType {
-    // The subfields are expected to be non-nullable according to the parquet variant spec.
-    let metadata_field = Field::new("metadata", DataType::Binary, false);
-    let value_field = Field::new("value", DataType::Binary, false);
-    let fields = vec![metadata_field, value_field];
-    DataType::Struct(fields.into())
-}
-
 /// Parse a batch of JSON strings into a batch of Variants represented as
 /// STRUCT<metadata: BINARY, value: BINARY> where nulls are preserved. The JSON strings in the input
 /// must be valid.
-pub fn batch_json_string_to_variant(input: &ArrayRef) -> Result<StructArray, ArrowError> {
+pub fn batch_json_string_to_variant(input: &ArrayRef) -> Result<VariantArray, ArrowError> {
     let input_string_array = match input.as_any().downcast_ref::<StringArray>() {
         Some(string_array) => Ok(string_array),
         None => Err(ArrowError::CastError(
             "Expected reference to StringArray as input".into(),
         )),
     }?;
 
-    // Zero-copy builders
-    let mut metadata_buffer: Vec<u8> = Vec::with_capacity(input.len() * 128);
-    let mut metadata_offsets: Vec<i32> = Vec::with_capacity(input.len() + 1);
-    let mut metadata_validity = BooleanBufferBuilder::new(input.len());
-    let mut metadata_current_offset: i32 = 0;
-    metadata_offsets.push(metadata_current_offset);
-
-    let mut value_buffer: Vec<u8> = Vec::with_capacity(input.len() * 128);
-    let mut value_offsets: Vec<i32> = Vec::with_capacity(input.len() + 1);
-    let mut value_validity = BooleanBufferBuilder::new(input.len());
-    let mut value_current_offset: i32 = 0;
-    value_offsets.push(value_current_offset);
-
-    let mut validity = BooleanBufferBuilder::new(input.len());
+    let mut variant_array_builder = VariantArrayBuilder::new(input_string_array.len());
     for i in 0..input.len() {
         if input.is_null(i) {
             // The subfields are expected to be non-nullable according to the parquet variant spec.
-            metadata_validity.append(true);
-            value_validity.append(true);
-            metadata_offsets.push(metadata_current_offset);
-            value_offsets.push(value_current_offset);
-            validity.append(false);
+            variant_array_builder.append_null();
         } else {
             let mut vb = VariantBuilder::new();
             json_to_variant(input_string_array.value(i), &mut vb)?;
             let (metadata, value) = vb.finish();
-            validity.append(true);
-
-            metadata_current_offset += metadata.len() as i32;
-            metadata_buffer.extend(metadata);
-            metadata_offsets.push(metadata_current_offset);
-            metadata_validity.append(true);
-
-            value_current_offset += value.len() as i32;
-            value_buffer.extend(value);
-            value_offsets.push(value_current_offset);
-            value_validity.append(true);
+            variant_array_builder.append_variant_buffers(&metadata, &value);
         }
     }
-    let metadata_offsets_buffer = OffsetBuffer::new(ScalarBuffer::from(metadata_offsets));
-    let metadata_data_buffer = Buffer::from_vec(metadata_buffer);
-    let metadata_null_buffer = NullBuffer::new(metadata_validity.finish());
-
-    let value_offsets_buffer = OffsetBuffer::new(ScalarBuffer::from(value_offsets));
-    let value_data_buffer = Buffer::from_vec(value_buffer);
-    let value_null_buffer = NullBuffer::new(value_validity.finish());
-
-    let metadata_array = BinaryArray::new(
-        metadata_offsets_buffer,
-        metadata_data_buffer,
-        Some(metadata_null_buffer),
-    );
-    let value_array = BinaryArray::new(
-        value_offsets_buffer,
-        value_data_buffer,
-        Some(value_null_buffer),
-    );
-
-    let struct_fields: Vec<ArrayRef> = vec![Arc::new(metadata_array), Arc::new(value_array)];
-    let variant_fields = match variant_arrow_repr() {
-        DataType::Struct(fields) => fields,
-        _ => unreachable!("variant_arrow_repr is hard-coded and must match the expected schema"),
-    };
-    let null_buffer = NullBuffer::new(validity.finish());
-    Ok(StructArray::new(
-        variant_fields,
-        struct_fields,
-        Some(null_buffer),
-    ))
+    Ok(variant_array_builder.build())
 }
 
 #[cfg(test)]
 mod test {
     use crate::batch_json_string_to_variant;
-    use arrow::array::{Array, ArrayRef, BinaryArray, StringArray};
+    use arrow::array::{Array, ArrayRef, AsArray, StringArray};
     use arrow_schema::ArrowError;
     use parquet_variant::{Variant, VariantBuilder};
     use std::sync::Arc;
@@ -135,43 +68,38 @@ mod test {
             None,
         ]);
         let array_ref: ArrayRef = Arc::new(input);
-        let output = batch_json_string_to_variant(&array_ref).unwrap();
+        let variant_array = batch_json_string_to_variant(&array_ref).unwrap();
 
-        let struct_array = &output;
-        let metadata_array = struct_array
-            .column(0)
-            .as_any()
-            .downcast_ref::<BinaryArray>()
-            .unwrap();
-        let value_array = struct_array
-            .column(1)
-            .as_any()
-            .downcast_ref::<BinaryArray>()
-            .unwrap();
+        let metadata_array = variant_array.metadata_field().as_binary_view();
+        let value_array = variant_array.value_field().as_binary_view();
 
-        assert!(!struct_array.is_null(0));
-        assert!(struct_array.is_null(1));
-        assert!(!struct_array.is_null(2));
-        assert!(!struct_array.is_null(3));
-        assert!(struct_array.is_null(4));
+        // Compare row 0
+        assert!(!variant_array.is_null(0));
+        assert_eq!(variant_array.value(0), Variant::Int8(1));
 
-        assert_eq!(metadata_array.value(0), &[1, 0, 0]);
-        assert_eq!(value_array.value(0), &[12, 1]);
+        // Compare row 1
+        assert!(variant_array.is_null(1));
 
+        // Compare row 2
+        assert!(!variant_array.is_null(2));
         {
             let mut vb = VariantBuilder::new();
             let mut ob = vb.new_object();
             ob.insert("a", Variant::Int8(32));
             ob.finish()?;
             let (object_metadata, object_value) = vb.finish();
-            assert_eq!(metadata_array.value(2), &object_metadata);
-            assert_eq!(value_array.value(2), &object_value);
+            let expected = Variant::new(&object_metadata, &object_value);
+            assert_eq!(variant_array.value(2), expected);
         }
 
-        assert_eq!(metadata_array.value(3), &[1, 0, 0]);
-        assert_eq!(value_array.value(3), &[0]);
+        // Compare row 3 (Note this is a variant NULL, not a null row)
+        assert!(!variant_array.is_null(3));
+        assert_eq!(variant_array.value(3), Variant::Null);
+
+        // Compare row 4
+        assert!(variant_array.is_null(4));
 
-        // Ensure that the subfields are not actually nullable
+        // Ensure that the subfields are not nullable
         assert!(!metadata_array.is_null(1));
         assert!(!value_array.is_null(1));
         assert!(!metadata_array.is_null(4));

parquet-variant-compute/src/lib.rs

@@ -17,6 +17,11 @@
 
 mod from_json;
 mod to_json;
+mod variant_array;
+mod variant_array_builder;
+
+pub use variant_array::VariantArray;
+pub use variant_array_builder::VariantArrayBuilder;
 
 pub use from_json::batch_json_string_to_variant;
 pub use to_json::batch_variant_to_json_string;