This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 81867eb69e [Variant] Implement `VariantArray::value` for shredded
variants (#8105)
81867eb69e is described below
commit 81867eb69ebc52ecd91731e63d3dc943469f24eb
Author: Congxian Qiu <qcx978132...@gmail.com>
AuthorDate: Sat Aug 23 18:53:08 2025 +0800
[Variant] Implement `VariantArray::value` for shredded variants (#8105)
# Which issue does this PR close?
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
- Closes #8091 .
# Rationale for this change
Implement `VariantArray::value` for some more shredded variants(eg.
primitive_conversion/generic_conversion/non_generic_conversion).
# What changes are included in this PR?
- Extract all `macroRules` to a separate module `type_conversion.rs`
- Add a macro for `variant value`
# Are these changes tested?
Covered by the existing test
# Are there any user-facing changes?
No
---
parquet-variant-compute/src/cast_to_variant.rs | 174 ++++-----------------
parquet-variant-compute/src/lib.rs | 1 +
parquet-variant-compute/src/type_conversion.rs | 125 +++++++++++++++
parquet-variant-compute/src/variant_array.rs | 10 +-
parquet-variant-compute/src/variant_get/mod.rs | 58 +++++--
.../src/variant_get/output/mod.rs | 7 +-
.../src/variant_get/output/primitive.rs | 8 +-
.../src/variant_get/output/variant.rs | 18 +--
8 files changed, 231 insertions(+), 170 deletions(-)
diff --git a/parquet-variant-compute/src/cast_to_variant.rs
b/parquet-variant-compute/src/cast_to_variant.rs
index 8841ced27c..3850579946 100644
--- a/parquet-variant-compute/src/cast_to_variant.rs
+++ b/parquet-variant-compute/src/cast_to_variant.rs
@@ -17,6 +17,10 @@
use std::sync::Arc;
+use crate::type_conversion::{
+ decimal_to_variant_decimal, generic_conversion_array,
non_generic_conversion_array,
+ primitive_conversion_array,
+};
use crate::{VariantArray, VariantArrayBuilder};
use arrow::array::{
Array, AsArray, TimestampMicrosecondArray, TimestampMillisecondArray,
TimestampNanosecondArray,
@@ -37,60 +41,10 @@ use arrow::temporal_conversions::{
};
use arrow_schema::{ArrowError, DataType, TimeUnit};
use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc};
-use half::f16;
use parquet_variant::{
Variant, VariantBuilder, VariantDecimal16, VariantDecimal4,
VariantDecimal8,
};
-/// Convert the input array of a specific primitive type to a `VariantArray`
-/// row by row
-macro_rules! primitive_conversion {
- ($t:ty, $input:expr, $builder:expr) => {{
- let array = $input.as_primitive::<$t>();
- for i in 0..array.len() {
- if array.is_null(i) {
- $builder.append_null();
- continue;
- }
- $builder.append_variant(Variant::from(array.value(i)));
- }
- }};
-}
-
-/// Convert the input array to a `VariantArray` row by row, using `method`
-/// requiring a generic type to downcast the generic array to a specific
-/// array type and `cast_fn` to transform each element to a type compatible
with Variant
-macro_rules! generic_conversion {
- ($t:ty, $method:ident, $cast_fn:expr, $input:expr, $builder:expr) => {{
- let array = $input.$method::<$t>();
- for i in 0..array.len() {
- if array.is_null(i) {
- $builder.append_null();
- continue;
- }
- let cast_value = $cast_fn(array.value(i));
- $builder.append_variant(Variant::from(cast_value));
- }
- }};
-}
-
-/// Convert the input array to a `VariantArray` row by row, using `method`
-/// not requiring a generic type to downcast the generic array to a specific
-/// array type and `cast_fn` to transform each element to a type compatible
with Variant
-macro_rules! non_generic_conversion {
- ($method:ident, $cast_fn:expr, $input:expr, $builder:expr) => {{
- let array = $input.$method();
- for i in 0..array.len() {
- if array.is_null(i) {
- $builder.append_null();
- continue;
- }
- let cast_value = $cast_fn(array.value(i));
- $builder.append_variant(Variant::from(cast_value));
- }
- }};
-}
-
fn convert_timestamp(
time_unit: &TimeUnit,
time_zone: &Option<Arc<str>>,
@@ -159,61 +113,6 @@ fn convert_timestamp(
}
}
-/// Convert a decimal value to a `VariantDecimal`
-macro_rules! decimal_to_variant_decimal {
- ($v:ident, $scale:expr, $value_type:ty, $variant_type:ty) => {
- if *$scale < 0 {
- // For negative scale, we need to multiply the value by 10^|scale|
- // For example: 123 with scale -2 becomes 12300
- let multiplier = (10 as $value_type).pow((-*$scale) as u32);
- // Check for overflow
- if $v > 0 && $v > <$value_type>::MAX / multiplier {
- return Variant::Null;
- }
- if $v < 0 && $v < <$value_type>::MIN / multiplier {
- return Variant::Null;
- }
- <$variant_type>::try_new($v * multiplier, 0)
- .map(|v| v.into())
- .unwrap_or(Variant::Null)
- } else {
- <$variant_type>::try_new($v, *$scale as u8)
- .map(|v| v.into())
- .unwrap_or(Variant::Null)
- }
- };
-}
-
-/// Convert arrays that don't need generic type parameters
-macro_rules! cast_conversion_nongeneric {
- ($method:ident, $cast_fn:expr, $input:expr, $builder:expr) => {{
- let array = $input.$method();
- for i in 0..array.len() {
- if array.is_null(i) {
- $builder.append_null();
- continue;
- }
- let cast_value = $cast_fn(array.value(i));
- $builder.append_variant(Variant::from(cast_value));
- }
- }};
-}
-
-/// Convert string arrays using the offset size as the type parameter
-macro_rules! cast_conversion_string {
- ($offset_type:ty, $method:ident, $cast_fn:expr, $input:expr,
$builder:expr) => {{
- let array = $input.$method::<$offset_type>();
- for i in 0..array.len() {
- if array.is_null(i) {
- $builder.append_null();
- continue;
- }
- let cast_value = $cast_fn(array.value(i));
- $builder.append_variant(Variant::from(cast_value));
- }
- }};
-}
-
/// Casts a typed arrow [`Array`] to a [`VariantArray`]. This is useful when
you
/// need to convert a specific data type
///
@@ -250,58 +149,52 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
// todo: handle other types like Boolean, Date, Timestamp, etc.
match input_type {
DataType::Boolean => {
- non_generic_conversion!(as_boolean, |v| v, input, builder);
+ non_generic_conversion_array!(input.as_boolean(), |v| v, builder);
}
DataType::Binary => {
- generic_conversion!(BinaryType, as_bytes, |v| v, input, builder);
+ generic_conversion_array!(BinaryType, as_bytes, |v| v, input,
builder);
}
DataType::LargeBinary => {
- generic_conversion!(LargeBinaryType, as_bytes, |v| v, input,
builder);
+ generic_conversion_array!(LargeBinaryType, as_bytes, |v| v, input,
builder);
}
DataType::BinaryView => {
- generic_conversion!(BinaryViewType, as_byte_view, |v| v, input,
builder);
+ generic_conversion_array!(BinaryViewType, as_byte_view, |v| v,
input, builder);
}
DataType::Int8 => {
- primitive_conversion!(Int8Type, input, builder);
+ primitive_conversion_array!(Int8Type, input, builder);
}
DataType::Int16 => {
- primitive_conversion!(Int16Type, input, builder);
+ primitive_conversion_array!(Int16Type, input, builder);
}
DataType::Int32 => {
- primitive_conversion!(Int32Type, input, builder);
+ primitive_conversion_array!(Int32Type, input, builder);
}
DataType::Int64 => {
- primitive_conversion!(Int64Type, input, builder);
+ primitive_conversion_array!(Int64Type, input, builder);
}
DataType::UInt8 => {
- primitive_conversion!(UInt8Type, input, builder);
+ primitive_conversion_array!(UInt8Type, input, builder);
}
DataType::UInt16 => {
- primitive_conversion!(UInt16Type, input, builder);
+ primitive_conversion_array!(UInt16Type, input, builder);
}
DataType::UInt32 => {
- primitive_conversion!(UInt32Type, input, builder);
+ primitive_conversion_array!(UInt32Type, input, builder);
}
DataType::UInt64 => {
- primitive_conversion!(UInt64Type, input, builder);
+ primitive_conversion_array!(UInt64Type, input, builder);
}
DataType::Float16 => {
- generic_conversion!(
- Float16Type,
- as_primitive,
- |v: f16| -> f32 { v.into() },
- input,
- builder
- );
+ generic_conversion_array!(Float16Type, as_primitive, f32::from,
input, builder);
}
DataType::Float32 => {
- primitive_conversion!(Float32Type, input, builder);
+ primitive_conversion_array!(Float32Type, input, builder);
}
DataType::Float64 => {
- primitive_conversion!(Float64Type, input, builder);
+ primitive_conversion_array!(Float64Type, input, builder);
}
DataType::Decimal32(_, scale) => {
- generic_conversion!(
+ generic_conversion_array!(
Decimal32Type,
as_primitive,
|v| decimal_to_variant_decimal!(v, scale, i32,
VariantDecimal4),
@@ -310,7 +203,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
);
}
DataType::Decimal64(_, scale) => {
- generic_conversion!(
+ generic_conversion_array!(
Decimal64Type,
as_primitive,
|v| decimal_to_variant_decimal!(v, scale, i64,
VariantDecimal8),
@@ -319,7 +212,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
);
}
DataType::Decimal128(_, scale) => {
- generic_conversion!(
+ generic_conversion_array!(
Decimal128Type,
as_primitive,
|v| decimal_to_variant_decimal!(v, scale, i128,
VariantDecimal16),
@@ -328,7 +221,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
);
}
DataType::Decimal256(_, scale) => {
- generic_conversion!(
+ generic_conversion_array!(
Decimal256Type,
as_primitive,
|v: i256| {
@@ -346,7 +239,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
);
}
DataType::FixedSizeBinary(_) => {
- non_generic_conversion!(as_fixed_size_binary, |v| v, input,
builder);
+ non_generic_conversion_array!(input.as_fixed_size_binary(), |v| v,
builder);
}
DataType::Null => {
for _ in 0..input.len() {
@@ -359,7 +252,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
DataType::Time32(unit) => {
match *unit {
TimeUnit::Second => {
- generic_conversion!(
+ generic_conversion_array!(
Time32SecondType,
as_primitive,
// nano second are always 0
@@ -369,7 +262,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
);
}
TimeUnit::Millisecond => {
- generic_conversion!(
+ generic_conversion_array!(
Time32MillisecondType,
as_primitive,
|v| NaiveTime::from_num_seconds_from_midnight_opt(
@@ -392,7 +285,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
DataType::Time64(unit) => {
match *unit {
TimeUnit::Microsecond => {
- generic_conversion!(
+ generic_conversion_array!(
Time64MicrosecondType,
as_primitive,
|v| NaiveTime::from_num_seconds_from_midnight_opt(
@@ -405,7 +298,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
);
}
TimeUnit::Nanosecond => {
- generic_conversion!(
+ generic_conversion_array!(
Time64NanosecondType,
as_primitive,
|v| NaiveTime::from_num_seconds_from_midnight_opt(
@@ -433,13 +326,13 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
));
}
DataType::Utf8 => {
- cast_conversion_string!(i32, as_string, |v| v, input, builder);
+ generic_conversion_array!(i32, as_string, |v| v, input, builder);
}
DataType::LargeUtf8 => {
- cast_conversion_string!(i64, as_string, |v| v, input, builder);
+ generic_conversion_array!(i64, as_string, |v| v, input, builder);
}
DataType::Utf8View => {
- cast_conversion_nongeneric!(as_string_view, |v| v, input, builder);
+ non_generic_conversion_array!(input.as_string_view(), |v| v,
builder);
}
DataType::Struct(_) => {
let struct_array = input.as_struct();
@@ -487,7 +380,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
}
}
DataType::Date32 => {
- generic_conversion!(
+ generic_conversion_array!(
Date32Type,
as_primitive,
|v: i32| -> NaiveDate { Date32Type::to_naive_date(v) },
@@ -496,7 +389,7 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
);
}
DataType::Date64 => {
- generic_conversion!(
+ generic_conversion_array!(
Date64Type,
as_primitive,
|v: i64| { Date64Type::to_naive_date_opt(v).unwrap() },
@@ -723,6 +616,7 @@ mod tests {
use arrow_schema::{
DECIMAL128_MAX_PRECISION, DECIMAL32_MAX_PRECISION,
DECIMAL64_MAX_PRECISION,
};
+ use half::f16;
use parquet_variant::{Variant, VariantDecimal16};
use std::{sync::Arc, vec};
diff --git a/parquet-variant-compute/src/lib.rs
b/parquet-variant-compute/src/lib.rs
index 245e344488..ef674d9614 100644
--- a/parquet-variant-compute/src/lib.rs
+++ b/parquet-variant-compute/src/lib.rs
@@ -38,6 +38,7 @@
pub mod cast_to_variant;
mod from_json;
mod to_json;
+mod type_conversion;
mod variant_array;
mod variant_array_builder;
pub mod variant_get;
diff --git a/parquet-variant-compute/src/type_conversion.rs
b/parquet-variant-compute/src/type_conversion.rs
new file mode 100644
index 0000000000..647d2c705f
--- /dev/null
+++ b/parquet-variant-compute/src/type_conversion.rs
@@ -0,0 +1,125 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Module for transforming a typed arrow `Array` to `VariantArray`.
+
+/// Convert the input array to a `VariantArray` row by row, using `method`
+/// not requiring a generic type to downcast the generic array to a specific
+/// array type and `cast_fn` to transform each element to a type compatible
with Variant
+macro_rules! non_generic_conversion_array {
+ ($array:expr, $cast_fn:expr, $builder:expr) => {{
+ let array = $array;
+ for i in 0..array.len() {
+ if array.is_null(i) {
+ $builder.append_null();
+ continue;
+ }
+ let cast_value = $cast_fn(array.value(i));
+ $builder.append_variant(Variant::from(cast_value));
+ }
+ }};
+}
+pub(crate) use non_generic_conversion_array;
+
+/// Convert the value at a specific index in the given array into a `Variant`.
+macro_rules! non_generic_conversion_single_value {
+ ($array:expr, $cast_fn:expr, $index:expr) => {{
+ let array = $array;
+ if array.is_null($index) {
+ Variant::Null
+ } else {
+ let cast_value = $cast_fn(array.value($index));
+ Variant::from(cast_value)
+ }
+ }};
+}
+pub(crate) use non_generic_conversion_single_value;
+
+/// Convert the input array to a `VariantArray` row by row, using `method`
+/// requiring a generic type to downcast the generic array to a specific
+/// array type and `cast_fn` to transform each element to a type compatible
with Variant
+macro_rules! generic_conversion_array {
+ ($t:ty, $method:ident, $cast_fn:expr, $input:expr, $builder:expr) => {{
+ $crate::type_conversion::non_generic_conversion_array!(
+ $input.$method::<$t>(),
+ $cast_fn,
+ $builder
+ )
+ }};
+}
+pub(crate) use generic_conversion_array;
+
+/// Convert the value at a specific index in the given array into a `Variant`,
+/// using `method` requiring a generic type to downcast the generic array
+/// to a specific array type and `cast_fn` to transform the element.
+macro_rules! generic_conversion_single_value {
+ ($t:ty, $method:ident, $cast_fn:expr, $input:expr, $index:expr) => {{
+ $crate::type_conversion::non_generic_conversion_single_value!(
+ $input.$method::<$t>(),
+ $cast_fn,
+ $index
+ )
+ }};
+}
+pub(crate) use generic_conversion_single_value;
+
+/// Convert the input array of a specific primitive type to a `VariantArray`
+/// row by row
+macro_rules! primitive_conversion_array {
+ ($t:ty, $input:expr, $builder:expr) => {{
+ $crate::type_conversion::generic_conversion_array!(
+ $t,
+ as_primitive,
+ |v| v,
+ $input,
+ $builder
+ )
+ }};
+}
+pub(crate) use primitive_conversion_array;
+
+/// Convert the value at a specific index in the given array into a `Variant`.
+macro_rules! primitive_conversion_single_value {
+ ($t:ty, $input:expr, $index:expr) => {{
+ $crate::type_conversion::generic_conversion_single_value!(
+ $t,
+ as_primitive,
+ |v| v,
+ $input,
+ $index
+ )
+ }};
+}
+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)
+ };
+
+ v.and_then(|v| <$variant_type>::try_new(v, scale).ok())
+ .map_or(Variant::Null, Variant::from)
+ }};
+}
+pub(crate) use decimal_to_variant_decimal;
diff --git a/parquet-variant-compute/src/variant_array.rs
b/parquet-variant-compute/src/variant_array.rs
index c541258942..10fb5f67ee 100644
--- a/parquet-variant-compute/src/variant_array.rs
+++ b/parquet-variant-compute/src/variant_array.rs
@@ -19,12 +19,14 @@
use arrow::array::{Array, ArrayData, ArrayRef, AsArray, BinaryViewArray,
StructArray};
use arrow::buffer::NullBuffer;
-use arrow::datatypes::Int32Type;
+use arrow::datatypes::{Int16Type, Int32Type};
use arrow_schema::{ArrowError, DataType};
use parquet_variant::Variant;
use std::any::Any;
use std::sync::Arc;
+use crate::type_conversion::primitive_conversion_single_value;
+
/// An array of Parquet [`Variant`] values
///
/// A [`VariantArray`] wraps an Arrow [`StructArray`] that stores the
underlying
@@ -350,8 +352,10 @@ impl ShreddingState {
fn typed_value_to_variant(typed_value: &ArrayRef, index: usize) -> Variant<'_,
'_> {
match typed_value.data_type() {
DataType::Int32 => {
- let typed_value = typed_value.as_primitive::<Int32Type>();
- Variant::from(typed_value.value(index))
+ primitive_conversion_single_value!(Int32Type, typed_value, index)
+ }
+ DataType::Int16 => {
+ primitive_conversion_single_value!(Int16Type, typed_value, index)
}
// todo other types here (note this is very similar to
cast_to_variant.rs)
// so it would be great to figure out how to share this code
diff --git a/parquet-variant-compute/src/variant_get/mod.rs
b/parquet-variant-compute/src/variant_get/mod.rs
index 0c9d2686c0..4460705cba 100644
--- a/parquet-variant-compute/src/variant_get/mod.rs
+++ b/parquet-variant-compute/src/variant_get/mod.rs
@@ -107,7 +107,10 @@ impl<'a> GetOptions<'a> {
mod test {
use std::sync::Arc;
- use arrow::array::{Array, ArrayRef, BinaryViewArray, Int32Array,
StringArray, StructArray};
+ use arrow::array::{
+ Array, ArrayRef, BinaryViewArray, Int16Array, Int32Array,
PrimitiveArray, StringArray,
+ StructArray,
+ };
use arrow::buffer::NullBuffer;
use arrow::compute::CastOptions;
use arrow_schema::{DataType, Field, FieldRef, Fields};
@@ -258,7 +261,8 @@ mod test {
/// Perfect Shredding: extract the typed value as a VariantArray
#[test]
fn get_variant_perfectly_shredded_int32_as_variant() {
- let array = perfectly_shredded_int32_variant_array();
+ let array =
+ perfectly_shredded_variant_array(Int32Array::from(vec![Some(1),
Some(2), Some(3)]));
let options = GetOptions::new();
let result = variant_get(&array, options).unwrap();
@@ -276,7 +280,8 @@ mod test {
#[test]
fn get_variant_perfectly_shredded_int32_as_int32() {
// Extract the typed value as Int32Array
- let array = perfectly_shredded_int32_variant_array();
+ let array =
+ perfectly_shredded_variant_array(Int32Array::from(vec![Some(1),
Some(2), Some(3)]));
// specify we want the typed value as Int32
let field = Field::new("typed_value", DataType::Int32, true);
let options =
GetOptions::new().with_as_type(Some(FieldRef::from(field)));
@@ -319,14 +324,38 @@ mod test {
assert_eq!(&result, &expected)
}
+ #[test]
+ fn get_variant_perfectly_shredded_int16_as_variant() {
+ let array =
+ perfectly_shredded_variant_array(Int16Array::from(vec![Some(1),
Some(2), Some(3)]));
+ let options = GetOptions::new();
+ let result = variant_get(&array, options).unwrap();
+
+ // expect the result is a VariantArray
+ let result: &VariantArray = result.as_any().downcast_ref().unwrap();
+ assert_eq!(result.len(), 3);
+
+ // Expect the values are the same as the original values
+ assert_eq!(result.value(0), Variant::Int16(1));
+ assert_eq!(result.value(1), Variant::Int16(2));
+ assert_eq!(result.value(2), Variant::Int16(3));
+ }
+
+ #[test]
+ fn get_variant_perfectly_shredded_int16_as_int16() {
+ // Extract the typed value as Int16Array
+ let array =
+ perfectly_shredded_variant_array(Int16Array::from(vec![Some(1),
Some(2), Some(3)]));
+ // specify we want the typed value as Int16
+ let field = Field::new("typed_value", DataType::Int16, true);
+ let options =
GetOptions::new().with_as_type(Some(FieldRef::from(field)));
+ let result = variant_get(&array, options).unwrap();
+ let expected: ArrayRef = Arc::new(Int16Array::from(vec![Some(1),
Some(2), Some(3)]));
+ assert_eq!(&result, &expected)
+ }
+
/// Return a VariantArray that represents a perfectly "shredded" variant
- /// for the following example (3 Variant::Int32 values):
- ///
- /// ```text
- /// 1
- /// 2
- /// 3
- /// ```
+ /// for the given typed value.
///
/// The schema of the corresponding `StructArray` would look like this:
///
@@ -336,13 +365,16 @@ mod test {
/// typed_value: Int32Array,
/// }
/// ```
- fn perfectly_shredded_int32_variant_array() -> ArrayRef {
+ fn perfectly_shredded_variant_array<T>(typed_value: PrimitiveArray<T>) ->
ArrayRef
+ where
+ T: arrow::datatypes::ArrowPrimitiveType,
+ {
// At the time of writing, the `VariantArrayBuilder` does not support
shredding.
// so we must construct the array manually. see
https://github.com/apache/arrow-rs/issues/7895
let (metadata, _value) = {
parquet_variant::VariantBuilder::new().finish() };
- let metadata =
BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 3));
- let typed_value = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
+ let metadata =
+ BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata,
typed_value.len()));
let struct_array = StructArrayBuilder::new()
.with_field("metadata", Arc::new(metadata))
diff --git a/parquet-variant-compute/src/variant_get/output/mod.rs
b/parquet-variant-compute/src/variant_get/output/mod.rs
index 52a8f5bc02..3ca21d482f 100644
--- a/parquet-variant-compute/src/variant_get/output/mod.rs
+++ b/parquet-variant-compute/src/variant_get/output/mod.rs
@@ -23,7 +23,7 @@ use crate::variant_get::output::variant::VariantOutputBuilder;
use crate::variant_get::GetOptions;
use crate::VariantArray;
use arrow::array::{ArrayRef, BinaryViewArray};
-use arrow::datatypes::Int32Type;
+use arrow::datatypes::{Int16Type, Int32Type};
use arrow::error::Result;
use arrow_schema::{ArrowError, DataType};
@@ -87,6 +87,11 @@ pub(crate) fn instantiate_output_builder<'a>(
as_type,
cast_options,
))),
+ DataType::Int16 =>
Ok(Box::new(PrimitiveOutputBuilder::<Int16Type>::new(
+ path,
+ as_type,
+ cast_options,
+ ))),
dt => Err(ArrowError::NotYetImplemented(format!(
"variant_get with as_type={dt} is not implemented yet",
))),
diff --git a/parquet-variant-compute/src/variant_get/output/primitive.rs
b/parquet-variant-compute/src/variant_get/output/primitive.rs
index aabc9827a7..ff3e58c3c3 100644
--- a/parquet-variant-compute/src/variant_get/output/primitive.rs
+++ b/parquet-variant-compute/src/variant_get/output/primitive.rs
@@ -24,7 +24,7 @@ use arrow::array::{
NullBufferBuilder, PrimitiveArray,
};
use arrow::compute::{cast_with_options, CastOptions};
-use arrow::datatypes::Int32Type;
+use arrow::datatypes::{Int16Type, Int32Type};
use arrow_schema::{ArrowError, FieldRef};
use parquet_variant::{Variant, VariantPath};
use std::marker::PhantomData;
@@ -176,3 +176,9 @@ impl ArrowPrimitiveVariant for Int32Type {
variant.as_int32()
}
}
+
+impl ArrowPrimitiveVariant for Int16Type {
+ fn from_variant(variant: &Variant) -> Option<Self::Native> {
+ variant.as_int16()
+ }
+}
diff --git a/parquet-variant-compute/src/variant_get/output/variant.rs
b/parquet-variant-compute/src/variant_get/output/variant.rs
index 7c8b4da2f5..203fab233b 100644
--- a/parquet-variant-compute/src/variant_get/output/variant.rs
+++ b/parquet-variant-compute/src/variant_get/output/variant.rs
@@ -16,9 +16,9 @@
// under the License.
use crate::variant_get::output::OutputBuilder;
-use crate::{VariantArray, VariantArrayBuilder};
+use crate::{type_conversion::primitive_conversion_array, VariantArray,
VariantArrayBuilder};
use arrow::array::{Array, ArrayRef, AsArray, BinaryViewArray};
-use arrow::datatypes::Int32Type;
+use arrow::datatypes::{Int16Type, Int32Type};
use arrow_schema::{ArrowError, DataType};
use parquet_variant::{Variant, VariantPath};
use std::sync::Arc;
@@ -93,16 +93,10 @@ impl OutputBuilder for VariantOutputBuilder<'_> {
let mut array_builder = VariantArrayBuilder::new(variant_array.len());
match typed_value.data_type() {
DataType::Int32 => {
- let primitive_array = typed_value.as_primitive::<Int32Type>();
- for i in 0..variant_array.len() {
- if primitive_array.is_null(i) {
- array_builder.append_null();
- continue;
- }
-
- let int_value = primitive_array.value(i);
- array_builder.append_variant(Variant::from(int_value));
- }
+ primitive_conversion_array!(Int32Type, typed_value,
array_builder);
+ }
+ DataType::Int16 => {
+ primitive_conversion_array!(Int16Type, typed_value,
array_builder);
}
dt => {
// https://github.com/apache/arrow-rs/issues/8087
