parquet_variant_compute/
type_conversion.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Module for transforming a typed arrow `Array` to `VariantArray`.
19
20use arrow::datatypes::{self, ArrowPrimitiveType};
21use parquet_variant::Variant;
22
23/// Options for controlling the behavior of `cast_to_variant_with_options`.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct CastOptions {
26    /// If true, return error on conversion failure. If false, insert null for failed conversions.
27    pub strict: bool,
28}
29
30impl Default for CastOptions {
31    fn default() -> Self {
32        Self { strict: true }
33    }
34}
35
36/// Extension trait for Arrow primitive types that can extract their native value from a Variant
37pub(crate) trait PrimitiveFromVariant: ArrowPrimitiveType {
38    fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native>;
39}
40
41/// Macro to generate PrimitiveFromVariant implementations for Arrow primitive types
42macro_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    };
50}
51
52impl_primitive_from_variant!(datatypes::Int32Type, as_int32);
53impl_primitive_from_variant!(datatypes::Int16Type, as_int16);
54impl_primitive_from_variant!(datatypes::Int8Type, as_int8);
55impl_primitive_from_variant!(datatypes::Int64Type, as_int64);
56impl_primitive_from_variant!(datatypes::UInt8Type, as_u8);
57impl_primitive_from_variant!(datatypes::UInt16Type, as_u16);
58impl_primitive_from_variant!(datatypes::UInt32Type, as_u32);
59impl_primitive_from_variant!(datatypes::UInt64Type, as_u64);
60impl_primitive_from_variant!(datatypes::Float16Type, as_f16);
61impl_primitive_from_variant!(datatypes::Float32Type, as_f32);
62impl_primitive_from_variant!(datatypes::Float64Type, as_f64);
63
64/// Convert the value at a specific index in the given array into a `Variant`.
65macro_rules! non_generic_conversion_single_value {
66    ($array:expr, $cast_fn:expr, $index:expr) => {{
67        let array = $array;
68        if array.is_null($index) {
69            Variant::Null
70        } else {
71            let cast_value = $cast_fn(array.value($index));
72            Variant::from(cast_value)
73        }
74    }};
75}
76pub(crate) use non_generic_conversion_single_value;
77
78/// Convert the value at a specific index in the given array into a `Variant`,
79/// using `method` requiring a generic type to downcast the generic array
80/// to a specific array type and `cast_fn` to transform the element.
81macro_rules! generic_conversion_single_value {
82    ($t:ty, $method:ident, $cast_fn:expr, $input:expr, $index:expr) => {{
83        $crate::type_conversion::non_generic_conversion_single_value!(
84            $input.$method::<$t>(),
85            $cast_fn,
86            $index
87        )
88    }};
89}
90pub(crate) use generic_conversion_single_value;
91
92/// Convert the value at a specific index in the given array into a `Variant`.
93macro_rules! primitive_conversion_single_value {
94    ($t:ty, $input:expr, $index:expr) => {{
95        $crate::type_conversion::generic_conversion_single_value!(
96            $t,
97            as_primitive,
98            |v| v,
99            $input,
100            $index
101        )
102    }};
103}
104pub(crate) use primitive_conversion_single_value;
105
106/// Convert a decimal value to a `VariantDecimal`
107macro_rules! decimal_to_variant_decimal {
108    ($v:ident, $scale:expr, $value_type:ty, $variant_type:ty) => {{
109        let (v, scale) = if *$scale < 0 {
110            // For negative scale, we need to multiply the value by 10^|scale|
111            // For example: 123 with scale -2 becomes 12300 with scale 0
112            let multiplier = <$value_type>::pow(10, (-*$scale) as u32);
113            (<$value_type>::checked_mul($v, multiplier), 0u8)
114        } else {
115            (Some($v), *$scale as u8)
116        };
117
118        // Return an Option to allow callers to decide whether to error (strict)
119        // or append null (non-strict) on conversion failure
120        v.and_then(|v| <$variant_type>::try_new(v, scale).ok())
121    }};
122}
123pub(crate) use decimal_to_variant_decimal;