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

parquet_variant_compute/
unshred_variant.rs
