Skip to content

Commit ef51f11

Browse files
authored
[Variant] Reverse VariantAsPrimitive trait to PrimitiveFromVariant (#8519)
# Which issue does this PR close? - Relates to #8515 - Relates to #8516 # Rationale for this change The existing `VariantAsPrimitive` trait is kind of "backward" and requires very complex type bounds to use, e.g.: ```rust impl<'a, T> VariantToPrimitiveArrowRowBuilder<'a, T> where T: ArrowPrimitiveType, for<'m, 'v> Variant<'m, 'v>: VariantAsPrimitive<T>, ``` This is a code smell. # What changes are included in this PR? "Reverse" the trait -- instead of extending `Variant`, the trait extends `T: ArrowPrimitiveType`. The resulting type bounds are much more intuitive: ```rust impl<'a, T: PrimitiveFromVariant> VariantToPrimitiveArrowRowBuilder<'a, T> ``` # Are these changes tested? Existing unit tests cover this refactor. # Are there any user-facing changes? No.
1 parent d7a871f commit ef51f11

File tree

2 files changed

+28
-69
lines changed

2 files changed

+28
-69
lines changed

parquet-variant-compute/src/type_conversion.rs

Lines changed: 23 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -33,70 +33,33 @@ impl Default for CastOptions {
3333
}
3434
}
3535

36-
/// Helper trait for converting `Variant` values to arrow primitive values.
37-
pub(crate) trait VariantAsPrimitive<T: ArrowPrimitiveType> {
38-
fn as_primitive(&self) -> Option<T::Native>;
36+
/// Extension trait for Arrow primitive types that can extract their native value from a Variant
37+
pub(crate) trait PrimitiveFromVariant: ArrowPrimitiveType {
38+
fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native>;
3939
}
4040

41-
impl VariantAsPrimitive<datatypes::Int32Type> for Variant<'_, '_> {
42-
fn as_primitive(&self) -> Option<i32> {
43-
self.as_int32()
44-
}
45-
}
46-
impl VariantAsPrimitive<datatypes::Int16Type> for Variant<'_, '_> {
47-
fn as_primitive(&self) -> Option<i16> {
48-
self.as_int16()
49-
}
50-
}
51-
impl VariantAsPrimitive<datatypes::Int8Type> for Variant<'_, '_> {
52-
fn as_primitive(&self) -> Option<i8> {
53-
self.as_int8()
54-
}
55-
}
56-
impl VariantAsPrimitive<datatypes::Int64Type> for Variant<'_, '_> {
57-
fn as_primitive(&self) -> Option<i64> {
58-
self.as_int64()
59-
}
60-
}
61-
impl VariantAsPrimitive<datatypes::Float16Type> for Variant<'_, '_> {
62-
fn as_primitive(&self) -> Option<half::f16> {
63-
self.as_f16()
64-
}
65-
}
66-
impl VariantAsPrimitive<datatypes::Float32Type> for Variant<'_, '_> {
67-
fn as_primitive(&self) -> Option<f32> {
68-
self.as_f32()
69-
}
70-
}
71-
impl VariantAsPrimitive<datatypes::Float64Type> for Variant<'_, '_> {
72-
fn as_primitive(&self) -> Option<f64> {
73-
self.as_f64()
74-
}
75-
}
76-
77-
impl VariantAsPrimitive<datatypes::UInt8Type> for Variant<'_, '_> {
78-
fn as_primitive(&self) -> Option<u8> {
79-
self.as_u8()
80-
}
81-
}
82-
83-
impl VariantAsPrimitive<datatypes::UInt16Type> for Variant<'_, '_> {
84-
fn as_primitive(&self) -> Option<u16> {
85-
self.as_u16()
86-
}
87-
}
88-
89-
impl VariantAsPrimitive<datatypes::UInt32Type> for Variant<'_, '_> {
90-
fn as_primitive(&self) -> Option<u32> {
91-
self.as_u32()
92-
}
41+
/// Macro to generate PrimitiveFromVariant implementations for Arrow primitive types
42+
macro_rules! impl_primitive_from_variant {
43+
($arrow_type:ty, $variant_method:ident) => {
44+
impl PrimitiveFromVariant for $arrow_type {
45+
fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native> {
46+
variant.$variant_method()
47+
}
48+
}
49+
};
9350
}
9451

95-
impl VariantAsPrimitive<datatypes::UInt64Type> for Variant<'_, '_> {
96-
fn as_primitive(&self) -> Option<u64> {
97-
self.as_u64()
98-
}
99-
}
52+
impl_primitive_from_variant!(datatypes::Int32Type, as_int32);
53+
impl_primitive_from_variant!(datatypes::Int16Type, as_int16);
54+
impl_primitive_from_variant!(datatypes::Int8Type, as_int8);
55+
impl_primitive_from_variant!(datatypes::Int64Type, as_int64);
56+
impl_primitive_from_variant!(datatypes::UInt8Type, as_u8);
57+
impl_primitive_from_variant!(datatypes::UInt16Type, as_u16);
58+
impl_primitive_from_variant!(datatypes::UInt32Type, as_u32);
59+
impl_primitive_from_variant!(datatypes::UInt64Type, as_u64);
60+
impl_primitive_from_variant!(datatypes::Float16Type, as_f16);
61+
impl_primitive_from_variant!(datatypes::Float32Type, as_f32);
62+
impl_primitive_from_variant!(datatypes::Float64Type, as_f64);
10063

10164
/// Convert the value at a specific index in the given array into a `Variant`.
10265
macro_rules! non_generic_conversion_single_value {

parquet-variant-compute/src/variant_to_arrow.rs

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use arrow::datatypes::{self, ArrowPrimitiveType, DataType};
2121
use arrow::error::{ArrowError, Result};
2222
use parquet_variant::{Variant, VariantPath};
2323

24-
use crate::type_conversion::VariantAsPrimitive;
24+
use crate::type_conversion::PrimitiveFromVariant;
2525
use crate::{VariantArray, VariantValueArrayBuilder};
2626

2727
use std::sync::Arc;
@@ -298,12 +298,12 @@ fn get_type_name<T: ArrowPrimitiveType>() -> &'static str {
298298
}
299299

300300
/// Builder for converting variant values to primitive values
301-
pub(crate) struct VariantToPrimitiveArrowRowBuilder<'a, T: ArrowPrimitiveType> {
301+
pub(crate) struct VariantToPrimitiveArrowRowBuilder<'a, T: PrimitiveFromVariant> {
302302
builder: arrow::array::PrimitiveBuilder<T>,
303303
cast_options: &'a CastOptions<'a>,
304304
}
305305

306-
impl<'a, T: ArrowPrimitiveType> VariantToPrimitiveArrowRowBuilder<'a, T> {
306+
impl<'a, T: PrimitiveFromVariant> VariantToPrimitiveArrowRowBuilder<'a, T> {
307307
fn new(cast_options: &'a CastOptions<'a>, capacity: usize) -> Self {
308308
Self {
309309
builder: PrimitiveBuilder::<T>::with_capacity(capacity),
@@ -312,18 +312,14 @@ impl<'a, T: ArrowPrimitiveType> VariantToPrimitiveArrowRowBuilder<'a, T> {
312312
}
313313
}
314314

315-
impl<'a, T> VariantToPrimitiveArrowRowBuilder<'a, T>
316-
where
317-
T: ArrowPrimitiveType,
318-
for<'m, 'v> Variant<'m, 'v>: VariantAsPrimitive<T>,
319-
{
315+
impl<'a, T: PrimitiveFromVariant> VariantToPrimitiveArrowRowBuilder<'a, T> {
320316
fn append_null(&mut self) -> Result<()> {
321317
self.builder.append_null();
322318
Ok(())
323319
}
324320

325321
fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
326-
if let Some(v) = value.as_primitive() {
322+
if let Some(v) = T::from_variant(value) {
327323
self.builder.append_value(v);
328324
Ok(true)
329325
} else {

0 commit comments

Comments
 (0)