[Variant] Define and use VariantDecimalType trait

parquet-variant-compute/src/arrow_to_variant.rs

@@ -15,27 +15,25 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::type_conversion::{CastOptions, decimal_to_variant_decimal};
+use crate::type_conversion::CastOptions;
 use arrow::array::{
     Array, AsArray, FixedSizeListArray, GenericBinaryArray, GenericListArray, GenericListViewArray,
     GenericStringArray, OffsetSizeTrait, PrimitiveArray,
 };
 use arrow::compute::kernels::cast;
 use arrow::datatypes::{
-    ArrowNativeType, ArrowPrimitiveType, ArrowTemporalType, ArrowTimestampType, Date32Type,
-    Date64Type, Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type,
-    RunEndIndexType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
-    Time64NanosecondType, TimestampMicrosecondType, TimestampMillisecondType,
-    TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
+    self as datatypes, ArrowNativeType, ArrowPrimitiveType, ArrowTemporalType, ArrowTimestampType,
+    DecimalType, RunEndIndexType,
 };
 use arrow::temporal_conversions::{as_date, as_datetime, as_time};
 use arrow_schema::{ArrowError, DataType, TimeUnit};
 use chrono::{DateTime, TimeZone, Utc};
 use parquet_variant::{
     ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal4, VariantDecimal8,
-    VariantDecimal16,
+    VariantDecimal16, VariantDecimalType,
 };
 use std::collections::HashMap;
+use std::marker::PhantomData;
 use std::ops::Range;
 
 // ============================================================================
@@ -46,31 +44,31 @@ use std::ops::Range;
 pub(crate) enum ArrowToVariantRowBuilder<'a> {
     Null(NullArrowToVariantBuilder),
     Boolean(BooleanArrowToVariantBuilder<'a>),
-    PrimitiveInt8(PrimitiveArrowToVariantBuilder<'a, Int8Type>),
-    PrimitiveInt16(PrimitiveArrowToVariantBuilder<'a, Int16Type>),
-    PrimitiveInt32(PrimitiveArrowToVariantBuilder<'a, Int32Type>),
-    PrimitiveInt64(PrimitiveArrowToVariantBuilder<'a, Int64Type>),
-    PrimitiveUInt8(PrimitiveArrowToVariantBuilder<'a, UInt8Type>),
-    PrimitiveUInt16(PrimitiveArrowToVariantBuilder<'a, UInt16Type>),
-    PrimitiveUInt32(PrimitiveArrowToVariantBuilder<'a, UInt32Type>),
-    PrimitiveUInt64(PrimitiveArrowToVariantBuilder<'a, UInt64Type>),
-    PrimitiveFloat16(PrimitiveArrowToVariantBuilder<'a, Float16Type>),
-    PrimitiveFloat32(PrimitiveArrowToVariantBuilder<'a, Float32Type>),
-    PrimitiveFloat64(PrimitiveArrowToVariantBuilder<'a, Float64Type>),
-    Decimal32(Decimal32ArrowToVariantBuilder<'a>),
-    Decimal64(Decimal64ArrowToVariantBuilder<'a>),
-    Decimal128(Decimal128ArrowToVariantBuilder<'a>),
+    PrimitiveInt8(PrimitiveArrowToVariantBuilder<'a, datatypes::Int8Type>),
+    PrimitiveInt16(PrimitiveArrowToVariantBuilder<'a, datatypes::Int16Type>),
+    PrimitiveInt32(PrimitiveArrowToVariantBuilder<'a, datatypes::Int32Type>),
+    PrimitiveInt64(PrimitiveArrowToVariantBuilder<'a, datatypes::Int64Type>),
+    PrimitiveUInt8(PrimitiveArrowToVariantBuilder<'a, datatypes::UInt8Type>),
+    PrimitiveUInt16(PrimitiveArrowToVariantBuilder<'a, datatypes::UInt16Type>),
+    PrimitiveUInt32(PrimitiveArrowToVariantBuilder<'a, datatypes::UInt32Type>),
+    PrimitiveUInt64(PrimitiveArrowToVariantBuilder<'a, datatypes::UInt64Type>),
+    PrimitiveFloat16(PrimitiveArrowToVariantBuilder<'a, datatypes::Float16Type>),
+    PrimitiveFloat32(PrimitiveArrowToVariantBuilder<'a, datatypes::Float32Type>),
+    PrimitiveFloat64(PrimitiveArrowToVariantBuilder<'a, datatypes::Float64Type>),
+    Decimal32(DecimalArrowToVariantBuilder<'a, datatypes::Decimal32Type, VariantDecimal4>),
+    Decimal64(DecimalArrowToVariantBuilder<'a, datatypes::Decimal64Type, VariantDecimal8>),
+    Decimal128(DecimalArrowToVariantBuilder<'a, datatypes::Decimal128Type, VariantDecimal16>),
     Decimal256(Decimal256ArrowToVariantBuilder<'a>),
-    TimestampSecond(TimestampArrowToVariantBuilder<'a, TimestampSecondType>),
-    TimestampMillisecond(TimestampArrowToVariantBuilder<'a, TimestampMillisecondType>),
-    TimestampMicrosecond(TimestampArrowToVariantBuilder<'a, TimestampMicrosecondType>),
-    TimestampNanosecond(TimestampArrowToVariantBuilder<'a, TimestampNanosecondType>),
-    Date32(DateArrowToVariantBuilder<'a, Date32Type>),
-    Date64(DateArrowToVariantBuilder<'a, Date64Type>),
-    Time32Second(TimeArrowToVariantBuilder<'a, Time32SecondType>),
-    Time32Millisecond(TimeArrowToVariantBuilder<'a, Time32MillisecondType>),
-    Time64Microsecond(TimeArrowToVariantBuilder<'a, Time64MicrosecondType>),
-    Time64Nanosecond(TimeArrowToVariantBuilder<'a, Time64NanosecondType>),
+    TimestampSecond(TimestampArrowToVariantBuilder<'a, datatypes::TimestampSecondType>),
+    TimestampMillisecond(TimestampArrowToVariantBuilder<'a, datatypes::TimestampMillisecondType>),
+    TimestampMicrosecond(TimestampArrowToVariantBuilder<'a, datatypes::TimestampMicrosecondType>),
+    TimestampNanosecond(TimestampArrowToVariantBuilder<'a, datatypes::TimestampNanosecondType>),
+    Date32(DateArrowToVariantBuilder<'a, datatypes::Date32Type>),
+    Date64(DateArrowToVariantBuilder<'a, datatypes::Date64Type>),
+    Time32Second(TimeArrowToVariantBuilder<'a, datatypes::Time32SecondType>),
+    Time32Millisecond(TimeArrowToVariantBuilder<'a, datatypes::Time32MillisecondType>),
+    Time64Microsecond(TimeArrowToVariantBuilder<'a, datatypes::Time64MicrosecondType>),
+    Time64Nanosecond(TimeArrowToVariantBuilder<'a, datatypes::Time64NanosecondType>),
     Binary(BinaryArrowToVariantBuilder<'a, i32>),
     LargeBinary(BinaryArrowToVariantBuilder<'a, i64>),
     BinaryView(BinaryViewArrowToVariantBuilder<'a>),
@@ -87,9 +85,9 @@ pub(crate) enum ArrowToVariantRowBuilder<'a> {
     Map(MapArrowToVariantBuilder<'a>),
     Union(UnionArrowToVariantBuilder<'a>),
     Dictionary(DictionaryArrowToVariantBuilder<'a>),
-    RunEndEncodedInt16(RunEndEncodedArrowToVariantBuilder<'a, Int16Type>),
-    RunEndEncodedInt32(RunEndEncodedArrowToVariantBuilder<'a, Int32Type>),
-    RunEndEncodedInt64(RunEndEncodedArrowToVariantBuilder<'a, Int64Type>),
+    RunEndEncodedInt16(RunEndEncodedArrowToVariantBuilder<'a, datatypes::Int16Type>),
+    RunEndEncodedInt32(RunEndEncodedArrowToVariantBuilder<'a, datatypes::Int32Type>),
+    RunEndEncodedInt64(RunEndEncodedArrowToVariantBuilder<'a, datatypes::Int64Type>),
 }
 
 impl<'a> ArrowToVariantRowBuilder<'a> {
@@ -173,17 +171,26 @@ pub(crate) fn make_arrow_to_variant_row_builder<'a>(
             DataType::Float16 => PrimitiveFloat16(PrimitiveArrowToVariantBuilder::new(array)),
             DataType::Float32 => PrimitiveFloat32(PrimitiveArrowToVariantBuilder::new(array)),
             DataType::Float64 => PrimitiveFloat64(PrimitiveArrowToVariantBuilder::new(array)),
-            DataType::Decimal32(_, scale) => {
-                Decimal32(Decimal32ArrowToVariantBuilder::new(array, options, *scale))
-            }
-            DataType::Decimal64(_, scale) => {
-                Decimal64(Decimal64ArrowToVariantBuilder::new(array, options, *scale))
-            }
-            DataType::Decimal128(_, scale) => {
-                Decimal128(Decimal128ArrowToVariantBuilder::new(array, options, *scale))
-            }
-            DataType::Decimal256(_, scale) => {
-                Decimal256(Decimal256ArrowToVariantBuilder::new(array, options, *scale))
+            DataType::Decimal32(_, s) => Decimal32(DecimalArrowToVariantBuilder::new(
+                array,
+                options,
+                *s,
+                PhantomData,
+            )),
+            DataType::Decimal64(_, s) => Decimal64(DecimalArrowToVariantBuilder::new(
+                array,
+                options,
+                *s,
+                PhantomData,
+            )),
+            DataType::Decimal128(_, s) => Decimal128(DecimalArrowToVariantBuilder::new(
+                array,
+                options,
+                *s,
+                PhantomData,
+            )),
+            DataType::Decimal256(_, s) => {
+                Decimal256(Decimal256ArrowToVariantBuilder::new(array, options, *s))
             }
             DataType::Timestamp(time_unit, time_zone) => {
                 match time_unit {
@@ -320,20 +327,20 @@ pub(crate) fn make_arrow_to_variant_row_builder<'a>(
 // worth the trouble, tho, because it makes for some pretty bulky and unwieldy macro expansions.
 macro_rules! define_row_builder {
     (
-        struct $name:ident<$lifetime:lifetime $(, $generic:ident: $bound:path )?>
+        struct $name:ident<$lifetime:lifetime $(, $generic:ident: $bound:path )*>
         $( where $where_path:path: $where_bound:path $(,)? )?
         $({ $($field:ident: $field_type:ty),+ $(,)? })?,
         |$array_param:ident| -> $array_type:ty { $init_expr:expr }
         $(, |$value:ident| $(-> Option<$option_ty:ty>)? $value_transform:expr)?
     ) => {
-        pub(crate) struct $name<$lifetime $(, $generic: $bound )?>
+        pub(crate) struct $name<$lifetime $(, $generic: $bound )*>
         $( where $where_path: $where_bound )?
         {
             array: &$lifetime $array_type,
             $( $( $field: $field_type, )+ )?
         }
 
-        impl<$lifetime $(, $generic: $bound+ )?> $name<$lifetime $(, $generic)?>
+        impl<$lifetime $(, $generic: $bound )*> $name<$lifetime $(, $generic)*>
         $( where $where_path: $where_bound )?
         {
             pub(crate) fn new($array_param: &$lifetime dyn Array $(, $( $field: $field_type ),+ )?) -> Self {
@@ -401,43 +408,28 @@ define_row_builder!(
 );
 
 define_row_builder!(
-    struct Decimal32ArrowToVariantBuilder<'a> {
-        options: &'a CastOptions,
-        scale: i8,
-    },
-    |array| -> arrow::array::Decimal32Array { array.as_primitive() },
-    |value| -> Option<_> { decimal_to_variant_decimal!(value, scale, i32, VariantDecimal4) }
-);
-
-define_row_builder!(
-    struct Decimal64ArrowToVariantBuilder<'a> {
-        options: &'a CastOptions,
-        scale: i8,
-    },
-    |array| -> arrow::array::Decimal64Array { array.as_primitive() },
-    |value| -> Option<_> { decimal_to_variant_decimal!(value, scale, i64, VariantDecimal8) }
-);
-
-define_row_builder!(
-    struct Decimal128ArrowToVariantBuilder<'a> {
+    struct DecimalArrowToVariantBuilder<'a, A: DecimalType, V: VariantDecimalType>
+    where
+        V::Native: From<A::Native>,
+    {
         options: &'a CastOptions,
         scale: i8,
+        _phantom: PhantomData<V>,
     },
-    |array| -> arrow::array::Decimal128Array { array.as_primitive() },
-    |value| -> Option<_> { decimal_to_variant_decimal!(value, scale, i128, VariantDecimal16) }
+    |array| -> PrimitiveArray<A> { array.as_primitive() },
+    |value| -> Option<_> { V::try_new_with_signed_scale(value.into(), *scale).ok() }
 );
 
+// Decimal256 needs a two-stage conversion via i128
 define_row_builder!(
     struct Decimal256ArrowToVariantBuilder<'a> {
         options: &'a CastOptions,
         scale: i8,
     },
     |array| -> arrow::array::Decimal256Array { array.as_primitive() },
     |value| -> Option<_> {
-        // Decimal256 needs special handling - convert to i128 if possible
-        value.to_i128().and_then(|i128_val| {
-            decimal_to_variant_decimal!(i128_val, scale, i128, VariantDecimal16)
-        })
+        let value = value.to_i128();
+        value.and_then(|v| VariantDecimal16::try_new_with_signed_scale(v, *scale).ok())
     }
 );
 
@@ -911,6 +903,7 @@ mod tests {
     use super::*;
     use crate::{VariantArray, VariantArrayBuilder};
     use arrow::array::{ArrayRef, BooleanArray, Int32Array, StringArray};
+    use arrow::datatypes::Int32Type;
     use std::sync::Arc;
 
     /// Builds a VariantArray from an Arrow array using the row builder.

parquet-variant-compute/src/type_conversion.rs

@@ -102,22 +102,3 @@ macro_rules! primitive_conversion_single_value {
     }};
 }
 pub(crate) use primitive_conversion_single_value;
-
-/// Convert a decimal value to a `VariantDecimal`
-macro_rules! decimal_to_variant_decimal {
-    ($v:ident, $scale:expr, $value_type:ty, $variant_type:ty) => {{
-        let (v, scale) = if *$scale < 0 {
-            // For negative scale, we need to multiply the value by 10^|scale|
-            // For example: 123 with scale -2 becomes 12300 with scale 0
-            let multiplier = <$value_type>::pow(10, (-*$scale) as u32);
-            (<$value_type>::checked_mul($v, multiplier), 0u8)
-        } else {
-            (Some($v), *$scale as u8)
-        };
-
-        // Return an Option to allow callers to decide whether to error (strict)
-        // or append null (non-strict) on conversion failure
-        v.and_then(|v| <$variant_type>::try_new(v, scale).ok())
-    }};
-}
-pub(crate) use decimal_to_variant_decimal;