From 63cbf21f013fe565aa59a2f4168cfdd9a77f7827 Mon Sep 17 00:00:00 2001
From: Ryan Johnson <scovich@gmail.com>
Date: Tue, 30 Sep 2025 10:45:33 -0700
Subject: [PATCH] [Variant] Reverse VariantAsPrimitive trait to
 PrimitiveFromVariant

---
 .../src/type_conversion.rs                    | 83 +++++--------------
 .../src/variant_to_arrow.rs                   | 14 ++--
 2 files changed, 28 insertions(+), 69 deletions(-)

diff --git a/parquet-variant-compute/src/type_conversion.rs b/parquet-variant-compute/src/type_conversion.rs
index ccecd510f6cf..27bd0f949029 100644
--- a/parquet-variant-compute/src/type_conversion.rs
+++ b/parquet-variant-compute/src/type_conversion.rs
@@ -33,70 +33,33 @@ impl Default for CastOptions {
     }
 }
 
-/// Helper trait for converting `Variant` values to arrow primitive values.
-pub(crate) trait VariantAsPrimitive<T: ArrowPrimitiveType> {
-    fn as_primitive(&self) -> Option<T::Native>;
+/// Extension trait for Arrow primitive types that can extract their native value from a Variant
+pub(crate) trait PrimitiveFromVariant: ArrowPrimitiveType {
+    fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native>;
 }
 
-impl VariantAsPrimitive<datatypes::Int32Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<i32> {
-        self.as_int32()
-    }
-}
-impl VariantAsPrimitive<datatypes::Int16Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<i16> {
-        self.as_int16()
-    }
-}
-impl VariantAsPrimitive<datatypes::Int8Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<i8> {
-        self.as_int8()
-    }
-}
-impl VariantAsPrimitive<datatypes::Int64Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<i64> {
-        self.as_int64()
-    }
-}
-impl VariantAsPrimitive<datatypes::Float16Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<half::f16> {
-        self.as_f16()
-    }
-}
-impl VariantAsPrimitive<datatypes::Float32Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<f32> {
-        self.as_f32()
-    }
-}
-impl VariantAsPrimitive<datatypes::Float64Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<f64> {
-        self.as_f64()
-    }
-}
-
-impl VariantAsPrimitive<datatypes::UInt8Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<u8> {
-        self.as_u8()
-    }
-}
-
-impl VariantAsPrimitive<datatypes::UInt16Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<u16> {
-        self.as_u16()
-    }
-}
-
-impl VariantAsPrimitive<datatypes::UInt32Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<u32> {
-        self.as_u32()
-    }
+/// Macro to generate PrimitiveFromVariant implementations for Arrow primitive types
+macro_rules! impl_primitive_from_variant {
+    ($arrow_type:ty, $variant_method:ident) => {
+        impl PrimitiveFromVariant for $arrow_type {
+            fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native> {
+                variant.$variant_method()
+            }
+        }
+    };
 }
 
-impl VariantAsPrimitive<datatypes::UInt64Type> for Variant<'_, '_> {
-    fn as_primitive(&self) -> Option<u64> {
-        self.as_u64()
-    }
-}
+impl_primitive_from_variant!(datatypes::Int32Type, as_int32);
+impl_primitive_from_variant!(datatypes::Int16Type, as_int16);
+impl_primitive_from_variant!(datatypes::Int8Type, as_int8);
+impl_primitive_from_variant!(datatypes::Int64Type, as_int64);
+impl_primitive_from_variant!(datatypes::UInt8Type, as_u8);
+impl_primitive_from_variant!(datatypes::UInt16Type, as_u16);
+impl_primitive_from_variant!(datatypes::UInt32Type, as_u32);
+impl_primitive_from_variant!(datatypes::UInt64Type, as_u64);
+impl_primitive_from_variant!(datatypes::Float16Type, as_f16);
+impl_primitive_from_variant!(datatypes::Float32Type, as_f32);
+impl_primitive_from_variant!(datatypes::Float64Type, as_f64);
 
 /// Convert the value at a specific index in the given array into a `Variant`.
 macro_rules! non_generic_conversion_single_value {
diff --git a/parquet-variant-compute/src/variant_to_arrow.rs b/parquet-variant-compute/src/variant_to_arrow.rs
index c1483b74bc5b..50249aa63d20 100644
--- a/parquet-variant-compute/src/variant_to_arrow.rs
+++ b/parquet-variant-compute/src/variant_to_arrow.rs
@@ -21,7 +21,7 @@ use arrow::datatypes::{self, ArrowPrimitiveType, DataType};
 use arrow::error::{ArrowError, Result};
 use parquet_variant::{Variant, VariantPath};
 
-use crate::type_conversion::VariantAsPrimitive;
+use crate::type_conversion::PrimitiveFromVariant;
 use crate::{VariantArray, VariantValueArrayBuilder};
 
 use std::sync::Arc;
@@ -298,12 +298,12 @@ fn get_type_name<T: ArrowPrimitiveType>() -> &'static str {
 }
 
 /// Builder for converting variant values to primitive values
-pub(crate) struct VariantToPrimitiveArrowRowBuilder<'a, T: ArrowPrimitiveType> {
+pub(crate) struct VariantToPrimitiveArrowRowBuilder<'a, T: PrimitiveFromVariant> {
     builder: arrow::array::PrimitiveBuilder<T>,
     cast_options: &'a CastOptions<'a>,
 }
 
-impl<'a, T: ArrowPrimitiveType> VariantToPrimitiveArrowRowBuilder<'a, T> {
+impl<'a, T: PrimitiveFromVariant> VariantToPrimitiveArrowRowBuilder<'a, T> {
     fn new(cast_options: &'a CastOptions<'a>, capacity: usize) -> Self {
         Self {
             builder: PrimitiveBuilder::<T>::with_capacity(capacity),
@@ -312,18 +312,14 @@ impl<'a, T: ArrowPrimitiveType> VariantToPrimitiveArrowRowBuilder<'a, T> {
     }
 }
 
-impl<'a, T> VariantToPrimitiveArrowRowBuilder<'a, T>
-where
-    T: ArrowPrimitiveType,
-    for<'m, 'v> Variant<'m, 'v>: VariantAsPrimitive<T>,
-{
+impl<'a, T: PrimitiveFromVariant> VariantToPrimitiveArrowRowBuilder<'a, T> {
     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_primitive() {
+        if let Some(v) = T::from_variant(value) {
             self.builder.append_value(v);
             Ok(true)
         } else {