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

parquet_variant_compute/
unshred_variant.rs
