Skip to content

Commit f813b08

Browse files
committed
address comment
1 parent b124132 commit f813b08

File tree

3 files changed

+30
-68
lines changed

3 files changed

+30
-68
lines changed

parquet-variant-compute/src/cast_to_variant.rs

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@
1818
use std::sync::Arc;
1919

2020
use crate::{
21-
cast_conversion_nongeneric, cast_conversion_string, decimal_to_variant_decimal,
22-
generic_conversion_array, non_generic_conversion_array, primitive_conversion,
21+
decimal_to_variant_decimal, generic_conversion_array, non_generic_conversion_array,
22+
primitive_conversion_array,
2323
};
2424
use crate::{VariantArray, VariantArrayBuilder};
2525
use arrow::array::{
@@ -161,28 +161,28 @@ pub fn cast_to_variant(input: &dyn Array) -> Result<VariantArray, ArrowError> {
161161
generic_conversion_array!(BinaryViewType, as_byte_view, |v| v, input, builder);
162162
}
163163
DataType::Int8 => {
164-
primitive_conversion!(Int8Type, input, builder);
164+
primitive_conversion_array!(Int8Type, input, builder);
165165
}
166166
DataType::Int16 => {
167-
primitive_conversion!(Int16Type, input, builder);
167+
primitive_conversion_array!(Int16Type, input, builder);
168168
}
169169
DataType::Int32 => {
170-
primitive_conversion!(Int32Type, input, builder);
170+
primitive_conversion_array!(Int32Type, input, builder);
171171
}
172172
DataType::Int64 => {
173-
primitive_conversion!(Int64Type, input, builder);
173+
primitive_conversion_array!(Int64Type, input, builder);
174174
}
175175
DataType::UInt8 => {
176-
primitive_conversion!(UInt8Type, input, builder);
176+
primitive_conversion_array!(UInt8Type, input, builder);
177177
}
178178
DataType::UInt16 => {
179-
primitive_conversion!(UInt16Type, input, builder);
179+
primitive_conversion_array!(UInt16Type, input, builder);
180180
}
181181
DataType::UInt32 => {
182-
primitive_conversion!(UInt32Type, input, builder);
182+
primitive_conversion_array!(UInt32Type, input, builder);
183183
}
184184
DataType::UInt64 => {
185-
primitive_conversion!(UInt64Type, input, builder);
185+
primitive_conversion_array!(UInt64Type, input, builder);
186186
}
187187
DataType::Float16 => {
188188
generic_conversion_array!(
@@ -194,10 +194,10 @@ pub fn cast_to_variant(input: &dyn Array) -> Result<VariantArray, ArrowError> {
194194
);
195195
}
196196
DataType::Float32 => {
197-
primitive_conversion!(Float32Type, input, builder);
197+
primitive_conversion_array!(Float32Type, input, builder);
198198
}
199199
DataType::Float64 => {
200-
primitive_conversion!(Float64Type, input, builder);
200+
primitive_conversion_array!(Float64Type, input, builder);
201201
}
202202
DataType::Decimal32(_, scale) => {
203203
generic_conversion_array!(
@@ -332,13 +332,13 @@ pub fn cast_to_variant(input: &dyn Array) -> Result<VariantArray, ArrowError> {
332332
));
333333
}
334334
DataType::Utf8 => {
335-
cast_conversion_string!(i32, as_string, |v| v, input, builder);
335+
generic_conversion_array!(i32, as_string, |v| v, input, builder);
336336
}
337337
DataType::LargeUtf8 => {
338-
cast_conversion_string!(i64, as_string, |v| v, input, builder);
338+
generic_conversion_array!(i64, as_string, |v| v, input, builder);
339339
}
340340
DataType::Utf8View => {
341-
cast_conversion_nongeneric!(as_string_view, |v| v, input, builder);
341+
non_generic_conversion_array!(as_string_view, |v| v, input, builder);
342342
}
343343
DataType::Struct(_) => {
344344
let struct_array = input.as_struct();

parquet-variant-compute/src/type_conversion.rs

Lines changed: 12 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
/// Convert the input array of a specific primitive type to a `VariantArray`
2121
/// row by row
2222
#[macro_export]
23-
macro_rules! primitive_conversion {
23+
macro_rules! primitive_conversion_array {
2424
($t:ty, $input:expr, $builder:expr) => {{
2525
let array = $input.as_primitive::<$t>();
2626
for i in 0..array.len() {
@@ -96,71 +96,33 @@ macro_rules! non_generic_conversion_array {
9696
}};
9797
}
9898

99+
/// Convert the value at a specific index in the given array into a `Variant`.
99100
#[macro_export]
100101
macro_rules! non_generic_conversion_single_value {
101102
($method:ident, $cast_fn:expr, $input:expr, $index:expr) => {{
102103
let array = $input.$method();
103104
if array.is_null($index) {
104105
return Variant::Null;
105106
}
107+
let cast_value = $cast_fn(array.value($index));
106108
Variant::from(cast_value)
107109
}};
108110
}
109111

110112
/// Convert a decimal value to a `VariantDecimal`
111113
#[macro_export]
112114
macro_rules! decimal_to_variant_decimal {
113-
($v:ident, $scale:expr, $value_type:ty, $variant_type:ty) => {
114-
if *$scale < 0 {
115+
($v:ident, $scale:expr, $value_type:ty, $variant_type:ty) => {{
116+
let (v, scale) = if *$scale < 0 {
115117
// For negative scale, we need to multiply the value by 10^|scale|
116-
// For example: 123 with scale -2 becomes 12300
117-
let multiplier = (10 as $value_type).pow((-*$scale) as u32);
118-
// Check for overflow
119-
if $v > 0 && $v > <$value_type>::MAX / multiplier {
120-
return Variant::Null;
121-
}
122-
if $v < 0 && $v < <$value_type>::MIN / multiplier {
123-
return Variant::Null;
124-
}
125-
<$variant_type>::try_new($v * multiplier, 0)
126-
.map(|v| v.into())
127-
.unwrap_or(Variant::Null)
118+
// For example: 123 with scale -2 becomes 12300 with scale 0
119+
let multiplier = <$value_type>::pow(10, (-*$scale) as u32);
120+
(<$value_type>::checked_mul($v, multiplier), 0u8)
128121
} else {
129-
<$variant_type>::try_new($v, *$scale as u8)
130-
.map(|v| v.into())
131-
.unwrap_or(Variant::Null)
132-
}
133-
};
134-
}
122+
(Some($v), *$scale as u8)
123+
};
135124

136-
/// Convert arrays that don't need generic type parameters
137-
#[macro_export]
138-
macro_rules! cast_conversion_nongeneric {
139-
($method:ident, $cast_fn:expr, $input:expr, $builder:expr) => {{
140-
let array = $input.$method();
141-
for i in 0..array.len() {
142-
if array.is_null(i) {
143-
$builder.append_null();
144-
continue;
145-
}
146-
let cast_value = $cast_fn(array.value(i));
147-
$builder.append_variant(Variant::from(cast_value));
148-
}
149-
}};
150-
}
151-
152-
/// Convert string arrays using the offset size as the type parameter
153-
#[macro_export]
154-
macro_rules! cast_conversion_string {
155-
($offset_type:ty, $method:ident, $cast_fn:expr, $input:expr, $builder:expr) => {{
156-
let array = $input.$method::<$offset_type>();
157-
for i in 0..array.len() {
158-
if array.is_null(i) {
159-
$builder.append_null();
160-
continue;
161-
}
162-
let cast_value = $cast_fn(array.value(i));
163-
$builder.append_variant(Variant::from(cast_value));
164-
}
125+
v.and_then(|v| <$variant_type>::try_new(v, scale).ok())
126+
.map_or(Variant::Null, Variant::from)
165127
}};
166128
}

parquet-variant-compute/src/variant_get/output/variant.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
// under the License.
1717

1818
use crate::variant_get::output::OutputBuilder;
19-
use crate::{primitive_conversion, VariantArray, VariantArrayBuilder};
19+
use crate::{primitive_conversion_array, VariantArray, VariantArrayBuilder};
2020
use arrow::array::{Array, ArrayRef, AsArray, BinaryViewArray};
2121
use arrow::datatypes::{Int16Type, Int32Type};
2222
use arrow_schema::{ArrowError, DataType};
@@ -93,10 +93,10 @@ impl OutputBuilder for VariantOutputBuilder<'_> {
9393
let mut array_builder = VariantArrayBuilder::new(variant_array.len());
9494
match typed_value.data_type() {
9595
DataType::Int32 => {
96-
primitive_conversion!(Int32Type, typed_value, array_builder);
96+
primitive_conversion_array!(Int32Type, typed_value, array_builder);
9797
}
9898
DataType::Int16 => {
99-
primitive_conversion!(Int16Type, typed_value, array_builder);
99+
primitive_conversion_array!(Int16Type, typed_value, array_builder);
100100
}
101101
dt => {
102102
// https://github.com/apache/arrow-rs/issues/8087

0 commit comments

Comments
 (0)