[Variant] Support primitive variant to arrow row for boolean

Author: klion26

parquet-variant-compute/src/variant_get.rs

@@ -307,6 +307,7 @@ mod test {
     use arrow::buffer::NullBuffer;
     use arrow::compute::CastOptions;
     use arrow::datatypes::DataType::{Int16, Int32, Int64};
+    use arrow_schema::DataType::{Boolean, Float32, Float64, Int8};
     use arrow_schema::{DataType, Field, FieldRef, Fields};
     use chrono::DateTime;
     use parquet_variant::{Variant, VariantPath, EMPTY_VARIANT_METADATA_BYTES};
@@ -711,6 +712,13 @@ mod test {
         };
     }
 
+    perfectly_shredded_to_arrow_primitive_test!(
+        get_variant_perfectly_shredded_int18_as_int8,
+        Int8,
+        perfectly_shredded_int8_variant_array,
+        Int8Array::from(vec![Some(1), Some(2), Some(3)])
+    );
+
     perfectly_shredded_to_arrow_primitive_test!(
         get_variant_perfectly_shredded_int16_as_int16,
         Int16,
@@ -732,31 +740,37 @@ mod test {
         Int64Array::from(vec![Some(1), Some(2), Some(3)])
     );
 
-    /// Return a VariantArray that represents a perfectly "shredded" variant
-    /// for the given typed value.
-    ///
-    /// The schema of the corresponding `StructArray` would look like this:
-    ///
-    /// ```text
-    /// StructArray {
-    ///   metadata: BinaryViewArray,
-    ///   typed_value: Int32Array,
-    /// }
-    /// ```
-    macro_rules! numeric_perfectly_shredded_variant_array_fn {
-        ($func:ident, $array_type:ident, $primitive_type:ty) => {
+    perfectly_shredded_to_arrow_primitive_test!(
+        get_variant_perfectly_shredded_float32_as_float32,
+        Float32,
+        perfectly_shredded_float32_variant_array,
+        Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])
+    );
+
+    perfectly_shredded_to_arrow_primitive_test!(
+        get_variant_perfectly_shredded_float64_as_float64,
+        Float64,
+        perfectly_shredded_float64_variant_array,
+        Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])
+    );
+
+    perfectly_shredded_to_arrow_primitive_test!(
+        get_variant_perfectly_shredded_boolean_as_boolean,
+        Boolean,
+        perfectly_shredded_bool_variant_array,
+        BooleanArray::from(vec![Some(true), Some(false), Some(true)])
+    );
+
+    macro_rules! perfectly_shredded_variant_array_fn {
+        ($func:ident, $typed_value_gen:expr) => {
             fn $func() -> ArrayRef {
                 // At the time of writing, the `VariantArrayBuilder` does not support shredding.
                 // so we must construct the array manually.  see https://github.com/apache/arrow-rs/issues/7895
                 let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(
                     EMPTY_VARIANT_METADATA_BYTES,
                     3,
                 ));
-                let typed_value = $array_type::from(vec![
-                    Some(<$primitive_type>::try_from(1u8).unwrap()),
-                    Some(<$primitive_type>::try_from(2u8).unwrap()),
-                    Some(<$primitive_type>::try_from(3u8).unwrap()),
-                ]);
+                let typed_value = $typed_value_gen();
 
                 let struct_array = StructArrayBuilder::new()
                     .with_field("metadata", Arc::new(metadata), false)
@@ -770,6 +784,33 @@ mod test {
         };
     }
 
+    perfectly_shredded_variant_array_fn!(perfectly_shredded_bool_variant_array, || {
+        BooleanArray::from(vec![Some(true), Some(false), Some(true)])
+    });
+
+    /// Return a VariantArray that represents a perfectly "shredded" variant
+    /// for the given typed value.
+    ///
+    /// The schema of the corresponding `StructArray` would look like this:
+    ///
+    /// ```text
+    /// StructArray {
+    ///   metadata: BinaryViewArray,
+    ///   typed_value: Int32Array,
+    /// }
+    /// ```
+    macro_rules! numeric_perfectly_shredded_variant_array_fn {
+        ($func:ident, $array_type:ident, $primitive_type:ty) => {
+            perfectly_shredded_variant_array_fn!($func, || {
+                $array_type::from(vec![
+                    Some(<$primitive_type>::try_from(1u8).unwrap()),
+                    Some(<$primitive_type>::try_from(2u8).unwrap()),
+                    Some(<$primitive_type>::try_from(3u8).unwrap()),
+                ])
+            });
+        };
+    }
+
     numeric_perfectly_shredded_variant_array_fn!(
         perfectly_shredded_int8_variant_array,
         Int8Array,

parquet-variant-compute/src/variant_to_arrow.rs

@@ -15,7 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{ArrayRef, BinaryViewArray, NullBufferBuilder, PrimitiveBuilder};
+use arrow::array::{
+    Array, ArrayRef, BinaryViewArray, NullBufferBuilder, PrimitiveArray, PrimitiveBuilder,
+};
 use arrow::compute::CastOptions;
 use arrow::datatypes::{self, ArrowPrimitiveType, DataType};
 use arrow::error::{ArrowError, Result};
@@ -30,6 +32,7 @@ use std::sync::Arc;
 /// `VariantToArrowRowBuilder` (below) and `VariantToShreddedPrimitiveVariantRowBuilder` (in
 /// `shred_variant.rs`).
 pub(crate) enum PrimitiveVariantToArrowRowBuilder<'a> {
+    Boolean(VariantToBooleanArrowRowBuilder<'a>),
     Int8(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int8Type>),
     Int16(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int16Type>),
     Int32(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int32Type>),
@@ -41,6 +44,7 @@ pub(crate) enum PrimitiveVariantToArrowRowBuilder<'a> {
     Float16(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Float16Type>),
     Float32(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Float32Type>),
     Float64(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Float64Type>),
+    Boolean(VariantToBooleanArrowRowBuilder<'a>),
 }
 
 /// Builder for converting variant values into strongly typed Arrow arrays.
@@ -59,6 +63,7 @@ impl<'a> PrimitiveVariantToArrowRowBuilder<'a> {
     pub fn append_null(&mut self) -> Result<()> {
         use PrimitiveVariantToArrowRowBuilder::*;
         match self {
+            Boolean(b) => b.append_null(),
             Int8(b) => b.append_null(),
             Int16(b) => b.append_null(),
             Int32(b) => b.append_null(),
@@ -70,6 +75,8 @@ impl<'a> PrimitiveVariantToArrowRowBuilder<'a> {
             Float16(b) => b.append_null(),
             Float32(b) => b.append_null(),
             Float64(b) => b.append_null(),
+            TimestampMicro(b) => b.append_null(),
+            TimestampNano(b) => b.append_null(),
         }
     }
 
@@ -87,12 +94,14 @@ impl<'a> PrimitiveVariantToArrowRowBuilder<'a> {
             Float16(b) => b.append_value(value),
             Float32(b) => b.append_value(value),
             Float64(b) => b.append_value(value),
+            Boolean(b) => b.append_value(value),
         }
     }
 
     pub fn finish(self) -> Result<ArrayRef> {
         use PrimitiveVariantToArrowRowBuilder::*;
         match self {
+            Boolean(b) => b.finish(),
             Int8(b) => b.finish(),
             Int16(b) => b.finish(),
             Int32(b) => b.finish(),
@@ -104,6 +113,7 @@ impl<'a> PrimitiveVariantToArrowRowBuilder<'a> {
             Float16(b) => b.finish(),
             Float32(b) => b.finish(),
             Float64(b) => b.finish(),
+            Boolean(b) => b.finish(),
         }
     }
 }
@@ -146,6 +156,7 @@ pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>(
     use PrimitiveVariantToArrowRowBuilder::*;
 
     let builder = match data_type {
+        DataType::Boolean => Boolean(VariantToBooleanArrowRowBuilder::new(cast_options, capacity)),
         DataType::Int8 => Int8(VariantToPrimitiveArrowRowBuilder::new(
             cast_options,
             capacity,
@@ -190,6 +201,7 @@ pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>(
             cast_options,
             capacity,
         )),
+        DataType::Boolean => Boolean(VariantToBooleanArrowRowBuilder::new(cast_options, capacity)),
         _ if data_type.is_primitive() => {
             return Err(ArrowError::NotYetImplemented(format!(
                 "Primitive data_type {data_type:?} not yet implemented"
@@ -297,6 +309,49 @@ fn get_type_name<T: ArrowPrimitiveType>() -> &'static str {
     }
 }
 
+/// Builder for converting variant values to boolean values
+/// Boolean is not primitive types in Arrow, so we need a separate builder
+pub(crate) struct VariantToBooleanArrowRowBuilder<'a> {
+    builder: arrow::array::BooleanBuilder,
+    cast_options: &'a CastOptions<'a>,
+}
+
+impl<'a> VariantToBooleanArrowRowBuilder<'a> {
+    fn new(cast_options: &'a CastOptions<'a>, capacity: usize) -> Self {
+        Self {
+            builder: arrow::array::BooleanBuilder::with_capacity(capacity),
+            cast_options,
+        }
+    }
+
+    fn append_null(&mut self) -> Result<()> {
+        self.builder.append_null();
+        Ok(())
+    }
+
+    fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
+        if let Some(v) = value.as_boolean() {
+            self.builder.append_value(v);
+            Ok(true)
+        } else {
+            if !self.cast_options.safe {
+                // Unsafe casting: return error on conversion failure
+                return Err(ArrowError::CastError(format!(
+                    "Failed to extract boolean from variant {:?} at path VariantPath([])",
+                    value
+                )));
+            }
+            // Safe casting: append null on conversion failure
+            self.builder.append_null();
+            Ok(false)
+        }
+    }
+
+    fn finish(mut self) -> Result<ArrayRef> {
+        Ok(Arc::new(self.builder.finish()))
+    }
+}
+
 /// Builder for converting variant values to primitive values
 pub(crate) struct VariantToPrimitiveArrowRowBuilder<'a, T: ArrowPrimitiveType> {
     builder: arrow::array::PrimitiveBuilder<T>,