parquet_variant_compute/
unshred_variant.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Module for unshredding VariantArray by folding typed_value columns back into the value column.
6
7use crate::variant_array::{binary_array_value, validate_binary_array};
8use crate::{VariantArray, VariantValueArrayBuilder};
9use arrow::array::{
10    Array, ArrayRef, AsArray as _, BinaryArray, BinaryViewArray, BooleanArray,
11    FixedSizeBinaryArray, GenericListArray, GenericListViewArray, LargeBinaryArray,
12    LargeStringArray, ListLikeArray, PrimitiveArray, StringArray, StringViewArray, StructArray,
13};
14use arrow::buffer::NullBuffer;
15use arrow::datatypes::{
16    ArrowPrimitiveType, DataType, Date32Type, Decimal32Type, Decimal64Type, Decimal128Type,
17    DecimalType, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type,
18    Time64MicrosecondType, TimeUnit, TimestampMicrosecondType, TimestampNanosecondType,
19};
20use arrow::error::{ArrowError, Result};
21use arrow::temporal_conversions::time64us_to_time;
22use chrono::{DateTime, Utc};
23use indexmap::IndexMap;
24use parquet_variant::{
25    ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal4, VariantDecimal8,
26    VariantDecimal16, VariantDecimalType, VariantMetadata,
27};
28use std::marker::PhantomData;
29use std::sync::Arc;
30use uuid::Uuid;
31
32/// Removes all (nested) typed_value columns from a VariantArray by converting them back to binary
33/// variant and merging the resulting values back into the value column.
34///
35/// This function efficiently converts a shredded VariantArray back to an unshredded form where all
36/// data resides in the value column.
37///
38/// # Arguments
39/// * `array` - The VariantArray to unshred
40///
41/// # Returns
42/// A new VariantArray with all data in the value column and no typed_value column
43///
44/// # Errors
45/// - If the shredded data contains spec violations (e.g., field name conflicts)
46/// - If unsupported data types are encountered in typed_value columns
47pub fn unshred_variant(array: &VariantArray) -> Result<VariantArray> {
48    let nulls = array.nulls();
49    let metadata = array.metadata_column();
50    let value_col = array.value_column();
51    let typed_value_col = array.typed_value_column();
52
53    // Already unshredded: no data movement needed, but the output must annotate `value` as
54    // non-nullable per the spec. Inputs whose value-nulls are not all masked by the parent null
55    // buffer (spec-invalid "missing" rows) cannot be re-annotated and fall through to the row
56    // loop below, whose top-level builder materializes `Variant::Null` for such rows.
57    if typed_value_col.is_none() {
58        if value_field_is_non_nullable(array) {
59            return Ok(array.clone());
60        }
61        if value_nulls_are_masked(value_col, nulls) {
62            return Ok(VariantArray::from_parts_unshredded(
63                metadata.clone(),
64                value_col.clone(),
65                nulls.cloned(),
66            ));
67        }
68    }
69
70    // NOTE: None/None at top-level is technically invalid, but the shredding spec requires us to
71    // emit `Variant::Null` when a required value is missing.
72    let mut row_builder = UnshredVariantRowBuilder::try_new_opt(array.inner())?
73        .unwrap_or_else(UnshredVariantRowBuilder::null);
74
75    let mut value_builder = VariantValueArrayBuilder::new(array.len());
76    for i in 0..array.len() {
77        if array.is_null(i) {
78            value_builder.append_null();
79        } else {
80            let metadata_bytes = binary_array_value(metadata.as_ref(), i).ok_or_else(|| {
81                ArrowError::InvalidArgumentError(
82                    "metadata field must be a binary-like array".to_string(),
83                )
84            })?;
85            let metadata = VariantMetadata::try_new(metadata_bytes)?;
86            let mut value_builder = value_builder.builder_ext(&metadata);
87            row_builder.append_row(&mut value_builder, &metadata, i)?;
88        }
89    }
90
91    let value = value_builder.build()?;
92    Ok(VariantArray::from_parts_unshredded(
93        metadata.clone(),
94        Arc::new(value),
95        nulls.cloned(),
96    ))
97}
98
99fn value_field_is_non_nullable(array: &VariantArray) -> bool {
100    array
101        .inner()
102        .field_by_name("value")
103        .is_some_and(|field| !field.is_nullable())
104}
105
106/// Returns true if every null in `value` is masked by a parent null, i.e. the column may be
107/// annotated non-nullable.
108fn value_nulls_are_masked(value: &ArrayRef, parent_nulls: Option<&NullBuffer>) -> bool {
109    value.null_count() == 0
110        || parent_nulls
111            .zip(value.nulls())
112            .is_some_and(|(parent, value_nulls)| parent.contains(value_nulls))
113}
114
115/// Row builder for converting shredded VariantArray rows back to unshredded form
116enum UnshredVariantRowBuilder<'a> {
117    PrimitiveInt8(UnshredPrimitiveRowBuilder<'a, PrimitiveArray<Int8Type>>),
118    PrimitiveInt16(UnshredPrimitiveRowBuilder<'a, PrimitiveArray<Int16Type>>),
119    PrimitiveInt32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray<Int32Type>>),
120    PrimitiveInt64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray<Int64Type>>),
121    PrimitiveFloat32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray<Float32Type>>),
122    PrimitiveFloat64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray<Float64Type>>),
123    Decimal32(DecimalUnshredRowBuilder<'a, Decimal32Type, VariantDecimal4>),
124    Decimal64(DecimalUnshredRowBuilder<'a, Decimal64Type, VariantDecimal8>),
125    Decimal128(DecimalUnshredRowBuilder<'a, Decimal128Type, VariantDecimal16>),
126    PrimitiveDate32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray<Date32Type>>),
127    PrimitiveTime64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray<Time64MicrosecondType>>),
128    TimestampMicrosecond(TimestampUnshredRowBuilder<'a, TimestampMicrosecondType>),
129    TimestampNanosecond(TimestampUnshredRowBuilder<'a, TimestampNanosecondType>),
130    PrimitiveBoolean(UnshredPrimitiveRowBuilder<'a, BooleanArray>),
131    PrimitiveString(UnshredPrimitiveRowBuilder<'a, StringArray>),
132    PrimitiveStringView(UnshredPrimitiveRowBuilder<'a, StringViewArray>),
133    PrimitiveLargeString(UnshredPrimitiveRowBuilder<'a, LargeStringArray>),
134    PrimitiveBinary(UnshredPrimitiveRowBuilder<'a, BinaryArray>),
135    PrimitiveBinaryView(UnshredPrimitiveRowBuilder<'a, BinaryViewArray>),
136    PrimitiveLargeBinary(UnshredPrimitiveRowBuilder<'a, LargeBinaryArray>),
137    PrimitiveUuid(UnshredPrimitiveRowBuilder<'a, FixedSizeBinaryArray>),
138    List(ListUnshredVariantBuilder<'a, GenericListArray<i32>>),
139    LargeList(ListUnshredVariantBuilder<'a, GenericListArray<i64>>),
140    ListView(ListUnshredVariantBuilder<'a, GenericListViewArray<i32>>),
141    LargeListView(ListUnshredVariantBuilder<'a, GenericListViewArray<i64>>),
142    Struct(StructUnshredVariantBuilder<'a>),
143    ValueOnly(ValueOnlyUnshredVariantBuilder<'a>),
144    Null(NullUnshredVariantBuilder),
145}
146
147impl<'a> UnshredVariantRowBuilder<'a> {
148    /// Creates an all-null row builder.
149    fn null() -> Self {
150        Self::Null(NullUnshredVariantBuilder)
151    }
152
153    /// Appends a single row at the given value index to the supplied builder.
154    fn append_row(
155        &mut self,
156        builder: &mut impl VariantBuilderExt,
157        metadata: &VariantMetadata,
158        index: usize,
159    ) -> Result<()> {
160        match self {
161            Self::PrimitiveInt8(b) => b.append_row(builder, metadata, index),
162            Self::PrimitiveInt16(b) => b.append_row(builder, metadata, index),
163            Self::PrimitiveInt32(b) => b.append_row(builder, metadata, index),
164            Self::PrimitiveInt64(b) => b.append_row(builder, metadata, index),
165            Self::PrimitiveFloat32(b) => b.append_row(builder, metadata, index),
166            Self::PrimitiveFloat64(b) => b.append_row(builder, metadata, index),
167            Self::Decimal32(b) => b.append_row(builder, metadata, index),
168            Self::Decimal64(b) => b.append_row(builder, metadata, index),
169            Self::Decimal128(b) => b.append_row(builder, metadata, index),
170            Self::PrimitiveDate32(b) => b.append_row(builder, metadata, index),
171            Self::PrimitiveTime64(b) => b.append_row(builder, metadata, index),
172            Self::TimestampMicrosecond(b) => b.append_row(builder, metadata, index),
173            Self::TimestampNanosecond(b) => b.append_row(builder, metadata, index),
174            Self::PrimitiveBoolean(b) => b.append_row(builder, metadata, index),
175            Self::PrimitiveString(b) => b.append_row(builder, metadata, index),
176            Self::PrimitiveStringView(b) => b.append_row(builder, metadata, index),
177            Self::PrimitiveLargeString(b) => b.append_row(builder, metadata, index),
178            Self::PrimitiveBinary(b) => b.append_row(builder, metadata, index),
179            Self::PrimitiveBinaryView(b) => b.append_row(builder, metadata, index),
180            Self::PrimitiveLargeBinary(b) => b.append_row(builder, metadata, index),
181            Self::PrimitiveUuid(b) => b.append_row(builder, metadata, index),
182            Self::List(b) => b.append_row(builder, metadata, index),
183            Self::LargeList(b) => b.append_row(builder, metadata, index),
184            Self::ListView(b) => b.append_row(builder, metadata, index),
185            Self::LargeListView(b) => b.append_row(builder, metadata, index),
186            Self::Struct(b) => b.append_row(builder, metadata, index),
187            Self::ValueOnly(b) => b.append_row(builder, metadata, index),
188            Self::Null(b) => b.append_row(builder, metadata, index),
189        }
190    }
191
192    /// Creates a new UnshredVariantRowBuilder from the `(value, typed_value)` pair of a shredded
193    /// variant struct. Returns None for the None/None case - caller decides how to handle based on
194    /// context.
195    fn try_new_opt(inner_struct: &'a StructArray) -> Result<Option<Self>> {
196        let value = if let Some(value_col) = inner_struct.column_by_name("value") {
197            validate_binary_array(value_col.as_ref(), "value")?;
198            Some(value_col)
199        } else {
200            None
201        };
202        let Some(typed_value) = inner_struct.column_by_name("typed_value") else {
203            // Copy the value across directly, if present. Else caller decides what to do.
204            return Ok(value.map(|v| Self::ValueOnly(ValueOnlyUnshredVariantBuilder::new(v))));
205        };
206
207        // Has typed_value -> determine type and create appropriate builder
208        macro_rules! primitive_builder {
209            ($enum_variant:ident, $cast_fn:ident) => {
210                Self::$enum_variant(UnshredPrimitiveRowBuilder::new(
211                    value,
212                    typed_value.$cast_fn(),
213                ))
214            };
215        }
216
217        let builder = match typed_value.data_type() {
218            DataType::Int8 => primitive_builder!(PrimitiveInt8, as_primitive),
219            DataType::Int16 => primitive_builder!(PrimitiveInt16, as_primitive),
220            DataType::Int32 => primitive_builder!(PrimitiveInt32, as_primitive),
221            DataType::Int64 => primitive_builder!(PrimitiveInt64, as_primitive),
222            DataType::Float32 => primitive_builder!(PrimitiveFloat32, as_primitive),
223            DataType::Float64 => primitive_builder!(PrimitiveFloat64, as_primitive),
224            DataType::Decimal32(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => {
225                Self::Decimal32(DecimalUnshredRowBuilder::new(value, typed_value, *s))
226            }
227            DataType::Decimal64(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => {
228                Self::Decimal64(DecimalUnshredRowBuilder::new(value, typed_value, *s))
229            }
230            DataType::Decimal128(p, s) if VariantDecimal16::is_valid_precision_and_scale(p, s) => {
231                Self::Decimal128(DecimalUnshredRowBuilder::new(value, typed_value, *s))
232            }
233            DataType::Decimal32(_, _)
234            | DataType::Decimal64(_, _)
235            | DataType::Decimal128(_, _)
236            | DataType::Decimal256(_, _) => {
237                return Err(ArrowError::InvalidArgumentError(format!(
238                    "{} is not a valid variant shredding type",
239                    typed_value.data_type()
240                )));
241            }
242            DataType::Date32 => primitive_builder!(PrimitiveDate32, as_primitive),
243            DataType::Time64(TimeUnit::Microsecond) => {
244                primitive_builder!(PrimitiveTime64, as_primitive)
245            }
246            DataType::Time64(time_unit) => {
247                return Err(ArrowError::InvalidArgumentError(format!(
248                    "Time64({time_unit}) is not a valid variant shredding type",
249                )));
250            }
251            DataType::Timestamp(TimeUnit::Microsecond, timezone) => Self::TimestampMicrosecond(
252                TimestampUnshredRowBuilder::new(value, typed_value, timezone.is_some()),
253            ),
254            DataType::Timestamp(TimeUnit::Nanosecond, timezone) => Self::TimestampNanosecond(
255                TimestampUnshredRowBuilder::new(value, typed_value, timezone.is_some()),
256            ),
257            DataType::Timestamp(time_unit, _) => {
258                return Err(ArrowError::InvalidArgumentError(format!(
259                    "Timestamp({time_unit}) is not a valid variant shredding type",
260                )));
261            }
262            DataType::Boolean => primitive_builder!(PrimitiveBoolean, as_boolean),
263            DataType::Utf8 => primitive_builder!(PrimitiveString, as_string),
264            DataType::Utf8View => primitive_builder!(PrimitiveStringView, as_string_view),
265            DataType::LargeUtf8 => primitive_builder!(PrimitiveLargeString, as_string),
266            DataType::Binary => primitive_builder!(PrimitiveBinary, as_binary),
267            DataType::BinaryView => primitive_builder!(PrimitiveBinaryView, as_binary_view),
268            DataType::LargeBinary => primitive_builder!(PrimitiveLargeBinary, as_binary),
269            DataType::FixedSizeBinary(16) => {
270                primitive_builder!(PrimitiveUuid, as_fixed_size_binary)
271            }
272            DataType::FixedSizeBinary(size) => {
273                return Err(ArrowError::InvalidArgumentError(format!(
274                    "FixedSizeBinary({size}) is not a valid variant shredding type",
275                )));
276            }
277            DataType::Struct(_) => Self::Struct(StructUnshredVariantBuilder::try_new(
278                value,
279                typed_value.as_struct(),
280            )?),
281            DataType::List(_) => Self::List(ListUnshredVariantBuilder::try_new(
282                value,
283                typed_value.as_list(),
284            )?),
285            DataType::LargeList(_) => Self::LargeList(ListUnshredVariantBuilder::try_new(
286                value,
287                typed_value.as_list(),
288            )?),
289            DataType::ListView(_) => Self::ListView(ListUnshredVariantBuilder::try_new(
290                value,
291                typed_value.as_list_view(),
292            )?),
293            DataType::LargeListView(_) => Self::LargeListView(ListUnshredVariantBuilder::try_new(
294                value,
295                typed_value.as_list_view(),
296            )?),
297            _ => {
298                return Err(ArrowError::NotYetImplemented(format!(
299                    "Unshredding not yet supported for type: {}",
300                    typed_value.data_type()
301                )));
302            }
303        };
304        Ok(Some(builder))
305    }
306}
307
308/// Builder for arrays with neither typed_value nor value (all Variant::Null)
309struct NullUnshredVariantBuilder;
310
311impl NullUnshredVariantBuilder {
312    fn append_row(
313        &self,
314        builder: &mut impl VariantBuilderExt,
315        _metadata: &VariantMetadata,
316        _index: usize,
317    ) -> Result<()> {
318        builder.append_value(Variant::Null);
319        Ok(())
320    }
321}
322
323/// Builder for arrays that only have value column (already unshredded)
324struct ValueOnlyUnshredVariantBuilder<'a> {
325    value: &'a ArrayRef,
326}
327
328impl<'a> ValueOnlyUnshredVariantBuilder<'a> {
329    fn new(value: &'a ArrayRef) -> Self {
330        Self { value }
331    }
332
333    fn append_row(
334        &self,
335        builder: &mut impl VariantBuilderExt,
336        metadata: &VariantMetadata,
337        index: usize,
338    ) -> Result<()> {
339        if self.value.is_null(index) {
340            builder.append_null();
341        } else {
342            let value_bytes = binary_array_value(self.value.as_ref(), index).ok_or_else(|| {
343                ArrowError::InvalidArgumentError(
344                    "value field must be a binary-like array".to_string(),
345                )
346            })?;
347            let variant = Variant::try_new_with_metadata(metadata.clone(), value_bytes)?;
348            builder.append_value(variant);
349        }
350        Ok(())
351    }
352}
353
354/// Extension trait that directly adds row builder support for arrays that correspond to primitive
355/// variant types.
356trait AppendToVariantBuilder: Array {
357    fn append_to_variant_builder(
358        &self,
359        builder: &mut impl VariantBuilderExt,
360        index: usize,
361    ) -> Result<()>;
362}
363
364/// Macro that handles the unshredded case (typed_value is missing or NULL) and returns early if
365/// handled.  If not handled (shredded case), validates and returns the extracted value.
366macro_rules! handle_unshredded_case {
367    ($self:expr, $builder:expr, $metadata:expr, $index:expr, $partial_shredding:expr) => {{
368        let value = $self.value.as_ref().filter(|v| v.is_valid($index));
369        let value = value
370            .map(|v| {
371                let bytes = binary_array_value(v.as_ref(), $index).ok_or_else(|| {
372                    ArrowError::InvalidArgumentError(format!(
373                        "value field must be a binary-like array, instead got {}",
374                        v.data_type(),
375                    ))
376                })?;
377                Variant::try_new_with_metadata($metadata.clone(), bytes)
378            })
379            .transpose()?;
380
381        // If typed_value is null, handle unshredded case and return early
382        if $self.typed_value.is_null($index) {
383            match value {
384                Some(value) => $builder.append_value(value),
385                None => $builder.append_null(),
386            }
387            return Ok(());
388        }
389
390        // Only partial shredding allows value and typed_value to both be non-NULL
391        if !$partial_shredding && value.is_some() {
392            return Err(ArrowError::InvalidArgumentError(
393                "Invalid shredded variant: both value and typed_value are non-null".to_string(),
394            ));
395        }
396
397        // Return the extracted value for the partial shredded case
398        value
399    }};
400}
401
402/// Generic unshred builder that works with any Array implementing AppendToVariantBuilder
403struct UnshredPrimitiveRowBuilder<'a, T> {
404    value: Option<&'a ArrayRef>,
405    typed_value: &'a T,
406}
407
408impl<'a, T: AppendToVariantBuilder> UnshredPrimitiveRowBuilder<'a, T> {
409    fn new(value: Option<&'a ArrayRef>, typed_value: &'a T) -> Self {
410        Self { value, typed_value }
411    }
412
413    fn append_row(
414        &self,
415        builder: &mut impl VariantBuilderExt,
416        metadata: &VariantMetadata,
417        index: usize,
418    ) -> Result<()> {
419        handle_unshredded_case!(self, builder, metadata, index, false);
420
421        // If we get here, typed_value is valid and value is NULL
422        self.typed_value.append_to_variant_builder(builder, index)
423    }
424}
425
426// Macro to generate AppendToVariantBuilder implementations with optional value transformation
427macro_rules! impl_append_to_variant_builder {
428    ($array_type:ty $(, |$v:ident| $transform:expr)? ) => {
429        impl AppendToVariantBuilder for $array_type {
430            fn append_to_variant_builder(
431                &self,
432                builder: &mut impl VariantBuilderExt,
433                index: usize,
434            ) -> Result<()> {
435                let value = self.value(index);
436                $(
437                    let $v = value;
438                    let value = $transform;
439                )?
440                builder.append_value(value);
441                Ok(())
442            }
443        }
444    };
445}
446
447impl_append_to_variant_builder!(BooleanArray);
448impl_append_to_variant_builder!(StringArray);
449impl_append_to_variant_builder!(StringViewArray);
450impl_append_to_variant_builder!(LargeStringArray);
451impl_append_to_variant_builder!(BinaryArray);
452impl_append_to_variant_builder!(BinaryViewArray);
453impl_append_to_variant_builder!(LargeBinaryArray);
454impl_append_to_variant_builder!(PrimitiveArray<Int8Type>);
455impl_append_to_variant_builder!(PrimitiveArray<Int16Type>);
456impl_append_to_variant_builder!(PrimitiveArray<Int32Type>);
457impl_append_to_variant_builder!(PrimitiveArray<Int64Type>);
458impl_append_to_variant_builder!(PrimitiveArray<Float32Type>);
459impl_append_to_variant_builder!(PrimitiveArray<Float64Type>);
460
461impl_append_to_variant_builder!(PrimitiveArray<Date32Type>, |days_since_epoch| {
462    Date32Type::to_naive_date_opt(days_since_epoch).ok_or_else(|| {
463        ArrowError::InvalidArgumentError(format!("Invalid Date32 value: {days_since_epoch}"))
464    })?
465});
466
467impl_append_to_variant_builder!(
468    PrimitiveArray<Time64MicrosecondType>,
469    |micros_since_midnight| {
470        time64us_to_time(micros_since_midnight).ok_or_else(|| {
471            ArrowError::InvalidArgumentError(format!(
472                "Invalid Time64 microsecond value: {micros_since_midnight}"
473            ))
474        })?
475    }
476);
477
478// UUID from FixedSizeBinary(16)
479// NOTE: FixedSizeBinaryArray guarantees the byte length, so we can safely unwrap
480impl_append_to_variant_builder!(FixedSizeBinaryArray, |bytes| {
481    Uuid::from_slice(bytes).unwrap()
482});
483
484/// Trait for timestamp types to handle conversion to `DateTime<Utc>`
485trait TimestampType: ArrowPrimitiveType<Native = i64> {
486    fn to_datetime_utc(value: i64) -> Result<DateTime<Utc>>;
487}
488
489impl TimestampType for TimestampMicrosecondType {
490    fn to_datetime_utc(micros: i64) -> Result<DateTime<Utc>> {
491        DateTime::from_timestamp_micros(micros).ok_or_else(|| {
492            ArrowError::InvalidArgumentError(format!(
493                "Invalid timestamp microsecond value: {micros}"
494            ))
495        })
496    }
497}
498
499impl TimestampType for TimestampNanosecondType {
500    fn to_datetime_utc(nanos: i64) -> Result<DateTime<Utc>> {
501        Ok(DateTime::from_timestamp_nanos(nanos))
502    }
503}
504
505/// Generic builder for timestamp types that handles timezone-aware conversion
506struct TimestampUnshredRowBuilder<'a, T: TimestampType> {
507    value: Option<&'a ArrayRef>,
508    typed_value: &'a PrimitiveArray<T>,
509    has_timezone: bool,
510}
511
512impl<'a, T: TimestampType> TimestampUnshredRowBuilder<'a, T> {
513    fn new(value: Option<&'a ArrayRef>, typed_value: &'a dyn Array, has_timezone: bool) -> Self {
514        Self {
515            value,
516            typed_value: typed_value.as_primitive(),
517            has_timezone,
518        }
519    }
520
521    fn append_row(
522        &self,
523        builder: &mut impl VariantBuilderExt,
524        metadata: &VariantMetadata,
525        index: usize,
526    ) -> Result<()> {
527        handle_unshredded_case!(self, builder, metadata, index, false);
528
529        // If we get here, typed_value is valid and value is NULL
530        let timestamp_value = self.typed_value.value(index);
531        let dt = T::to_datetime_utc(timestamp_value)?;
532        if self.has_timezone {
533            builder.append_value(dt);
534        } else {
535            builder.append_value(dt.naive_utc());
536        }
537        Ok(())
538    }
539}
540
541/// Generic builder for decimal unshredding
542struct DecimalUnshredRowBuilder<'a, A: DecimalType, V>
543where
544    V: VariantDecimalType<Native = A::Native>,
545{
546    value: Option<&'a ArrayRef>,
547    typed_value: &'a PrimitiveArray<A>,
548    scale: i8,
549    _phantom: PhantomData<V>,
550}
551
552impl<'a, A: DecimalType, V> DecimalUnshredRowBuilder<'a, A, V>
553where
554    V: VariantDecimalType<Native = A::Native>,
555{
556    fn new(value: Option<&'a ArrayRef>, typed_value: &'a dyn Array, scale: i8) -> Self {
557        Self {
558            value,
559            typed_value: typed_value.as_primitive(),
560            scale,
561            _phantom: PhantomData,
562        }
563    }
564
565    fn append_row(
566        &self,
567        builder: &mut impl VariantBuilderExt,
568        metadata: &VariantMetadata,
569        index: usize,
570    ) -> Result<()> {
571        handle_unshredded_case!(self, builder, metadata, index, false);
572
573        let raw = self.typed_value.value(index);
574        let variant = V::try_new_with_signed_scale(raw, self.scale)?;
575        builder.append_value(variant);
576        Ok(())
577    }
578}
579
580/// Builder for unshredding struct/object types with nested fields
581struct StructUnshredVariantBuilder<'a> {
582    value: Option<&'a ArrayRef>,
583    typed_value: &'a arrow::array::StructArray,
584    field_unshredders: IndexMap<&'a str, Option<UnshredVariantRowBuilder<'a>>>,
585}
586
587impl<'a> StructUnshredVariantBuilder<'a> {
588    fn try_new(value: Option<&'a ArrayRef>, typed_value: &'a StructArray) -> Result<Self> {
589        // Create unshredders for each field in constructor
590        let mut field_unshredders = IndexMap::new();
591        for (field, field_array) in typed_value.fields().iter().zip(typed_value.columns()) {
592            // Factory returns None for None/None case -- these are missing fields we should skip
593            let Some(field_array) = field_array.as_struct_opt() else {
594                return Err(ArrowError::InvalidArgumentError(format!(
595                    "Invalid shredded variant object field: expected Struct, got {}",
596                    field_array.data_type()
597                )));
598            };
599            let field_unshredder = UnshredVariantRowBuilder::try_new_opt(field_array)?;
600            field_unshredders.insert(field.name().as_ref(), field_unshredder);
601        }
602
603        Ok(Self {
604            value,
605            typed_value,
606            field_unshredders,
607        })
608    }
609
610    fn append_row(
611        &mut self,
612        builder: &mut impl VariantBuilderExt,
613        metadata: &VariantMetadata,
614        index: usize,
615    ) -> Result<()> {
616        let value = handle_unshredded_case!(self, builder, metadata, index, true);
617
618        // If we get here, typed_value is valid and value may or may not be valid
619        let mut object_builder = builder.try_new_object()?;
620
621        // Process typed fields (skip empty builders that indicate missing fields)
622        for (field_name, field_unshredder_opt) in &mut self.field_unshredders {
623            if let Some(field_unshredder) = field_unshredder_opt {
624                let mut field_builder = ObjectFieldBuilder::new(field_name, &mut object_builder);
625                field_unshredder.append_row(&mut field_builder, metadata, index)?;
626            }
627        }
628
629        // Process any unshredded fields (partial shredding)
630        if let Some(value) = value {
631            let Variant::Object(object) = value else {
632                return Err(ArrowError::InvalidArgumentError(
633                    "Expected object in value field for partially shredded struct".to_string(),
634                ));
635            };
636
637            for entry in object.iter_try() {
638                let (field_name, field_value) = entry?;
639                if self.field_unshredders.contains_key(field_name) {
640                    return Err(ArrowError::InvalidArgumentError(format!(
641                        "Field '{field_name}' appears in both typed_value and value",
642                    )));
643                }
644                object_builder.insert_bytes(field_name, field_value);
645            }
646        }
647
648        object_builder.finish();
649        Ok(())
650    }
651}
652
653/// Builder for unshredding list/array types with recursive element processing
654struct ListUnshredVariantBuilder<'a, L: ListLikeArray> {
655    value: Option<&'a ArrayRef>,
656    typed_value: &'a L,
657    element_unshredder: Box<UnshredVariantRowBuilder<'a>>,
658}
659
660impl<'a, L: ListLikeArray> ListUnshredVariantBuilder<'a, L> {
661    fn try_new(value: Option<&'a ArrayRef>, typed_value: &'a L) -> Result<Self> {
662        // Create a recursive unshredder for the list elements
663        // The element type comes from the values array of the list
664        let element_values = typed_value.values();
665
666        // For shredded lists, each element would be a ShreddedVariantFieldArray (struct)
667        // Extract value/typed_value from the element struct
668        let Some(element_values) = element_values.as_struct_opt() else {
669            return Err(ArrowError::InvalidArgumentError(format!(
670                "Invalid shredded variant array element: expected Struct, got {}",
671                element_values.data_type()
672            )));
673        };
674
675        // Create recursive unshredder for elements
676        //
677        // NOTE: A None/None array element is technically invalid, but the shredding spec
678        // requires us to emit `Variant::Null` when a required value is missing.
679        let element_unshredder = UnshredVariantRowBuilder::try_new_opt(element_values)?
680            .unwrap_or_else(UnshredVariantRowBuilder::null);
681
682        Ok(Self {
683            value,
684            typed_value,
685            element_unshredder: Box::new(element_unshredder),
686        })
687    }
688
689    fn append_row(
690        &mut self,
691        builder: &mut impl VariantBuilderExt,
692        metadata: &VariantMetadata,
693        index: usize,
694    ) -> Result<()> {
695        handle_unshredded_case!(self, builder, metadata, index, false);
696
697        // If we get here, typed_value is valid and value is NULL -- process the list elements
698        let mut list_builder = builder.try_new_list()?;
699        for element_index in self.typed_value.element_range(index) {
700            self.element_unshredder
701                .append_row(&mut list_builder, metadata, element_index)?;
702        }
703
704        list_builder.finish();
705        Ok(())
706    }
707}
708
709#[cfg(test)]
710mod tests {
711    use crate::{VariantArray, VariantArrayBuilder, shred_variant};
712    use arrow::array::{
713        Array, ArrayRef, AsArray, BinaryArray, BinaryViewArray, Int64Array, LargeBinaryArray,
714        LargeStringArray, StringViewArray,
715    };
716    use arrow::datatypes::{DataType, Field, Fields, TimeUnit};
717    use parquet_variant::{EMPTY_VARIANT_METADATA_BYTES, Variant, VariantBuilder, VariantDecimal8};
718    use std::sync::Arc;
719
720    /// Returns the nullability annotation of the `value` field
721    fn value_field_is_nullable(array: &VariantArray) -> bool {
722        array.inner().field_by_name("value").unwrap().is_nullable()
723    }
724
725    #[test]
726    fn test_unshred_utf8view_typed_value() {
727        let metadata_bytes = EMPTY_VARIANT_METADATA_BYTES;
728        let metadata: ArrayRef =
729            Arc::new(BinaryViewArray::from_iter_values(vec![metadata_bytes; 3]));
730
731        let typed_value: ArrayRef = Arc::new(StringViewArray::from(vec![
732            Some("hello"),
733            Some("middle"),
734            Some("world"),
735        ]));
736
737        let variant_array = VariantArray::perfectly_shredded(metadata, typed_value, None);
738
739        let result = crate::unshred_variant(&variant_array).unwrap();
740
741        assert_eq!(result.len(), 3);
742        assert_eq!(result.value(0), Variant::from("hello"));
743        assert_eq!(result.value(1), Variant::from("middle"));
744        assert_eq!(result.value(2), Variant::from("world"));
745    }
746
747    #[test]
748    fn test_unshred_largeutf8_typed_value() {
749        let metadata_bytes = EMPTY_VARIANT_METADATA_BYTES;
750        let metadata: ArrayRef =
751            Arc::new(BinaryViewArray::from_iter_values(vec![metadata_bytes; 3]));
752
753        let typed_value: ArrayRef = Arc::new(LargeStringArray::from(vec![
754            Some("hello"),
755            Some("middle"),
756            Some("world"),
757        ]));
758
759        let variant_array = VariantArray::perfectly_shredded(metadata, typed_value, None);
760
761        let result = crate::unshred_variant(&variant_array).unwrap();
762
763        assert_eq!(result.len(), 3);
764        assert_eq!(result.value(0), Variant::from("hello"));
765        assert_eq!(result.value(1), Variant::from("middle"));
766        assert_eq!(result.value(2), Variant::from("world"));
767    }
768
769    #[test]
770    fn test_unshred_binary_typed_value() {
771        let metadata_bytes = EMPTY_VARIANT_METADATA_BYTES;
772        let metadata: ArrayRef =
773            Arc::new(BinaryViewArray::from_iter_values(vec![metadata_bytes; 3]));
774
775        let typed_value: ArrayRef = Arc::new(BinaryArray::from_iter_values(vec![
776            &b"\x00\x01\x02"[..],
777            &b"\xff\xaa"[..],
778            &b"\xde\xad\xbe\xef"[..],
779        ]));
780
781        let variant_array = VariantArray::perfectly_shredded(metadata, typed_value, None);
782
783        let result = crate::unshred_variant(&variant_array).unwrap();
784
785        assert_eq!(result.len(), 3);
786        assert_eq!(result.value(0), Variant::from(&b"\x00\x01\x02"[..]));
787        assert_eq!(result.value(1), Variant::from(&b"\xff\xaa"[..]));
788        assert_eq!(result.value(2), Variant::from(&b"\xde\xad\xbe\xef"[..]));
789    }
790
791    #[test]
792    fn test_unshred_largebinary_typed_value() {
793        let metadata_bytes = EMPTY_VARIANT_METADATA_BYTES;
794        let metadata: ArrayRef =
795            Arc::new(BinaryViewArray::from_iter_values(vec![metadata_bytes; 3]));
796
797        let typed_value: ArrayRef = Arc::new(LargeBinaryArray::from_iter_values(vec![
798            &b"\x00\x01\x02"[..],
799            &b"\xff\xaa"[..],
800            &b"\xde\xad\xbe\xef"[..],
801        ]));
802
803        let variant_array = VariantArray::perfectly_shredded(metadata, typed_value, None);
804
805        let result = crate::unshred_variant(&variant_array).unwrap();
806
807        assert_eq!(result.len(), 3);
808        assert_eq!(result.value(0), Variant::from(&b"\x00\x01\x02"[..]));
809        assert_eq!(result.value(1), Variant::from(&b"\xff\xaa"[..]));
810        assert_eq!(result.value(2), Variant::from(&b"\xde\xad\xbe\xef"[..]));
811    }
812
813    #[test]
814    fn test_shred_unshred_round_trip_annotates_value_non_nullable() {
815        let original = VariantArray::from_iter([Some(42_i64)]);
816        assert!(!value_field_is_nullable(&original));
817
818        let shredded = shred_variant(&original, &DataType::Int64).unwrap();
819        assert!(value_field_is_nullable(&shredded)); // legal: typed_value present
820
821        let unshredded = crate::unshred_variant(&shredded).unwrap();
822        assert!(!value_field_is_nullable(&unshredded));
823        assert_eq!(unshredded, original);
824    }
825
826    #[test]
827    fn test_unshred_with_nulls_annotates_value_non_nullable() {
828        // a null row plus an unshreddable row, so the shredded input exercises parent
829        // nulls and both value/typed_value columns
830        let mut builder = VariantArrayBuilder::new(3);
831        builder.append_variant(Variant::from(1i64));
832        builder.append_null();
833        builder.append_variant(Variant::from("s"));
834        let original = builder.build();
835
836        let shredded = shred_variant(&original, &DataType::Int64).unwrap();
837        let unshredded = crate::unshred_variant(&shredded).unwrap();
838
839        assert!(!value_field_is_nullable(&unshredded));
840        assert_eq!(unshredded.len(), 3);
841        assert_eq!(unshredded.value(0), Variant::from(1i64));
842        assert!(unshredded.is_null(1));
843        assert_eq!(unshredded.value(2), Variant::from("s"));
844    }
845
846    #[test]
847    fn test_unshred_already_unshredded_reannotates_nullable_value() {
848        let original = VariantArray::from_iter([Some(42_i64), None]);
849
850        // same data, but with the out-of-spec nullable `value` annotation
851        let nullable_input = VariantArray::from_parts(
852            original.metadata_column().clone(),
853            original.value_column().clone(),
854            None,
855            original.nulls().cloned(),
856        );
857        assert!(value_field_is_nullable(&nullable_input));
858
859        let unshredded = crate::unshred_variant(&nullable_input).unwrap();
860        assert!(!value_field_is_nullable(&unshredded));
861        assert_eq!(unshredded, original);
862    }
863
864    #[test]
865    fn test_unshred_missing_top_level_value_becomes_variant_null() {
866        let metadata_bytes = EMPTY_VARIANT_METADATA_BYTES;
867        let metadata: ArrayRef =
868            Arc::new(BinaryViewArray::from_iter_values(vec![metadata_bytes; 3]));
869
870        // row 1 is null in both value and typed_value with a valid parent row: spec-invalid
871        // "missing" value, tolerated as Variant::Null (like `VariantArray::try_value`)
872        let typed_value: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]));
873        let variant_array = VariantArray::perfectly_shredded(metadata, typed_value, None);
874
875        let result = crate::unshred_variant(&variant_array).unwrap();
876
877        assert!(!value_field_is_nullable(&result));
878        assert_eq!(result.inner().null_count(), 0);
879        assert_eq!(result.value(0), Variant::from(1i64));
880        assert_eq!(result.value(1), Variant::Null);
881        assert_eq!(result.value(2), Variant::from(3i64));
882    }
883
884    /// Shreds `original` to `as_type`, then drops the parent null buffer so the parent-null row
885    /// becomes a spec-invalid "missing" row (value and typed_value both NULL with a valid
886    /// parent), and asserts unshredding turns exactly that row into `Variant::Null`.
887    fn assert_missing_row_unshreds_to_variant_null(original: &VariantArray, as_type: &DataType) {
888        let shredded = shred_variant(original, as_type).unwrap();
889        // Row 0 must actually shred, so that its round trip below exercises the typed
890        // reconstruction path of this shape's row builder, not the value fallback.
891        assert!(shredded.typed_value_column().unwrap().is_valid(0));
892
893        // The parent-null row carries no metadata bytes, so give every row the metadata of row 0.
894        let metadata_bytes = shredded.metadata_column().as_binary_view().value(0);
895        let metadata: ArrayRef = Arc::new(BinaryViewArray::from_iter_values(std::iter::repeat_n(
896            metadata_bytes,
897            original.len(),
898        )));
899
900        let input = VariantArray::from_parts(
901            metadata,
902            shredded.value_column().clone(),
903            shredded.typed_value_column().cloned(),
904            None,
905        );
906
907        let result = crate::unshred_variant(&input).unwrap();
908        assert!(!value_field_is_nullable(&result));
909        assert_eq!(result.inner().null_count(), 0);
910        assert_eq!(result.value(0), original.value(0));
911        assert_eq!(result.value(1), Variant::Null);
912    }
913
914    /// Missing rows must become `Variant::Null` through every row-builder shape, since each
915    /// shape has its own expansion of `handle_unshredded_case`.
916    #[test]
917    fn test_unshred_missing_row_for_decimal_timestamp_object_list() {
918        let decimals =
919            VariantArray::from_iter([Some(VariantDecimal8::try_new(1234, 2).unwrap()), None]);
920        assert_missing_row_unshreds_to_variant_null(&decimals, &DataType::Decimal64(18, 2));
921
922        let timestamps =
923            VariantArray::from_iter([Some(chrono::DateTime::from_timestamp(1, 0).unwrap()), None]);
924        assert_missing_row_unshreds_to_variant_null(
925            &timestamps,
926            &DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
927        );
928
929        let mut variant_builder = VariantBuilder::new();
930        let mut object_builder = variant_builder.new_object();
931        object_builder.insert("a", 1i64);
932        object_builder.finish();
933        let (object_metadata, object_value) = variant_builder.finish();
934        let mut builder = VariantArrayBuilder::new(2);
935        builder.append_variant(Variant::try_new(&object_metadata, &object_value).unwrap());
936        builder.append_null();
937        assert_missing_row_unshreds_to_variant_null(
938            &builder.build(),
939            &DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int64, true)])),
940        );
941
942        let mut variant_builder = VariantBuilder::new();
943        let mut list_builder = variant_builder.new_list();
944        list_builder.append_value(1i64);
945        list_builder.append_value(2i64);
946        list_builder.finish();
947        let (list_metadata, list_value) = variant_builder.finish();
948        let mut builder = VariantArrayBuilder::new(2);
949        builder.append_variant(Variant::try_new(&list_metadata, &list_value).unwrap());
950        builder.append_null();
951        assert_missing_row_unshreds_to_variant_null(
952            &builder.build(),
953            &DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
954        );
955    }
956
957    #[test]
958    fn test_unshred_value_only_with_unmasked_nulls_materializes_variant_null() {
959        let single = VariantArray::from_iter([Some(42_i64)]);
960        let metadata_bytes = single.metadata_column().as_binary_view().value(0);
961        let value_bytes = single.value_column().as_binary_view().value(0);
962
963        // unshredded input whose `value` null is not masked by a parent null: cannot be
964        // re-annotated in place, so unshredding must materialize Variant::Null bytes
965        let metadata: ArrayRef =
966            Arc::new(BinaryViewArray::from_iter_values(vec![metadata_bytes; 2]));
967        let value: ArrayRef = Arc::new(BinaryViewArray::from(vec![Some(value_bytes), None]));
968        let input = VariantArray::from_parts(metadata, value, None, None);
969
970        let result = crate::unshred_variant(&input).unwrap();
971
972        assert!(!value_field_is_nullable(&result));
973        assert_eq!(result.inner().null_count(), 0);
974        assert_eq!(result.value(0), Variant::from(42i64));
975        assert_eq!(result.value(1), Variant::Null);
976    }
977
978    #[test]
979    fn test_unshred_returns_err_on_malformed_metadata() {
980        // empty metadata bytes fail VariantMetadata's header parse. before this fix the
981        // call inside unshred_variant used the panicking `VariantMetadata::new`, which
982        // crashed the thread instead of surfacing the spec violation through the
983        // documented `Result` return type.
984        let metadata: ArrayRef = Arc::new(BinaryViewArray::from_iter_values(vec![&b""[..]]));
985
986        let typed_value: ArrayRef = Arc::new(StringViewArray::from(vec![Some("hello")]));
987
988        let variant_array = VariantArray::perfectly_shredded(metadata, typed_value, None);
989
990        let result = crate::unshred_variant(&variant_array);
991
992        assert!(
993            result.is_err(),
994            "unshred_variant must return Err on malformed metadata, not panic",
995        );
996    }
997}
parquet_variant_compute/unshred_variant.rs

parquet_variant_compute/
unshred_variant.rs
