parquet_variant_compute/
variant_array.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//! [`VariantArray`] implementation
19
20use crate::VariantArrayBuilder;
21use crate::type_conversion::{
22    generic_conversion_single_value, generic_conversion_single_value_with_result,
23    primitive_conversion_single_value,
24};
25use arrow::array::{Array, ArrayRef, AsArray, StructArray, new_null_array};
26use arrow::buffer::NullBuffer;
27use arrow::compute::cast;
28use arrow::datatypes::{
29    Date32Type, Decimal32Type, Decimal64Type, Decimal128Type, Float16Type, Float32Type,
30    Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Time64MicrosecondType,
31    TimestampMicrosecondType, TimestampNanosecondType,
32};
33use arrow::error::Result;
34use arrow_schema::extension::{ExtensionType, Uuid as UuidExtension};
35use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields, TimeUnit};
36use chrono::{DateTime, NaiveTime};
37use parquet_variant::{
38    Uuid, Variant, VariantDecimal4, VariantDecimal8, VariantDecimal16, VariantDecimalType as _,
39};
40
41use std::borrow::Cow;
42use std::sync::Arc;
43
44/// Returns the raw bytes at the given index from a binary-like array, return `None` if the array isn't binary-like.
45pub(crate) fn binary_array_value(array: &dyn Array, index: usize) -> Option<&[u8]> {
46    match array.data_type() {
47        DataType::Binary => Some(array.as_binary::<i32>().value(index)),
48        DataType::LargeBinary => Some(array.as_binary::<i64>().value(index)),
49        DataType::BinaryView => Some(array.as_binary_view().value(index)),
50        _ => None,
51    }
52}
53
54/// Returns a [`Variant`] from a `metadata` and `value` byte arrays, returns `None`
55/// if one of them is of invalid type.
56pub(crate) fn variant_from_arrays_at<'m, 'v>(
57    metadata: &'m dyn Array,
58    value: &'v dyn Array,
59    index: usize,
60) -> Option<Variant<'m, 'v>> {
61    let metadata = binary_array_value(metadata, index)?;
62    let value = binary_array_value(value, index)?;
63    Some(Variant::new(metadata, value))
64}
65
66/// Returns an all-null binary `value` column of the given length.
67///
68/// The shredding spec requires the `value` column to always be present in the
69/// schema, so producers that have no unshredded values to store must still
70/// emit an all-null column. See <https://github.com/apache/arrow-rs/issues/10306>.
71pub(crate) fn all_null_value_column(len: usize) -> ArrayRef {
72    new_null_array(&DataType::BinaryView, len)
73}
74
75/// Validates that an array has a binary-like data type.
76pub(crate) fn validate_binary_array(array: &dyn Array, field_name: &str) -> Result<()> {
77    match array.data_type() {
78        DataType::Binary | DataType::LargeBinary | DataType::BinaryView => Ok(()),
79        _ => Err(ArrowError::InvalidArgumentError(format!(
80            "VariantArray '{field_name}' field must be Binary, LargeBinary, or BinaryView, got {}",
81            array.data_type()
82        ))),
83    }
84}
85
86/// Arrow Variant [`ExtensionType`].
87///
88/// Represents the canonical Arrow Extension Type for storing variants.
89/// See [`VariantArray`] for more examples of using this extension type.
90pub struct VariantType;
91
92impl ExtensionType for VariantType {
93    const NAME: &'static str = "arrow.parquet.variant";
94
95    // Variants extension metadata is an empty string
96    // <https://github.com/apache/arrow/blob/d803afcc43f5d132506318fd9e162d33b2c3d4cd/docs/source/format/CanonicalExtensions.rst?plain=1#L473>
97    type Metadata = &'static str;
98
99    fn metadata(&self) -> &Self::Metadata {
100        &""
101    }
102
103    fn serialize_metadata(&self) -> Option<String> {
104        Some(String::new())
105    }
106
107    fn deserialize_metadata(_metadata: Option<&str>) -> Result<Self::Metadata> {
108        Ok("")
109    }
110
111    fn supports_data_type(&self, data_type: &DataType) -> Result<()> {
112        if matches!(data_type, DataType::Struct(_)) {
113            Ok(())
114        } else {
115            Err(ArrowError::InvalidArgumentError(format!(
116                "VariantType only supports StructArray, got {data_type}"
117            )))
118        }
119    }
120
121    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self> {
122        Self.supports_data_type(data_type)?;
123        Ok(Self)
124    }
125
126    fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<()> {
127        Self.supports_data_type(data_type)
128    }
129}
130
131/// An array of Parquet [`Variant`] values
132///
133/// A [`VariantArray`] wraps an Arrow [`StructArray`] that stores the underlying
134/// `metadata` and `value` fields, and adds convenience methods to access
135/// the [`Variant`]s.
136///
137/// See [`VariantArrayBuilder`] for constructing `VariantArray` row by row.
138///
139/// See the examples below from converting between `VariantArray` and
140/// `StructArray`.
141///
142/// [`VariantArrayBuilder`]: crate::VariantArrayBuilder
143///
144/// # Documentation
145///
146/// Variant is documented as a canonical Arrow extension type in the
147/// [Parquet Variant] section of the [official list of extension types] on
148/// the Apache Arrow website.
149///
150/// [Parquet Variant]: https://arrow.apache.org/docs/format/CanonicalExtensions.html#parquet-variant
151/// [official list of extension types]: https://arrow.apache.org/docs/format/CanonicalExtensions.html
152///
153/// # Example: Check if a [`StructArray`] has the [`VariantType`] extension
154///
155/// Arrow Arrays only provide [`DataType`], but the extension type information
156/// is stored on a [`Field`]. Thus, you must have access to the [`Schema`] or
157/// [`Field`] to check for the extension type.
158///
159/// [`Schema`]: arrow_schema::Schema
160/// ```
161/// # use arrow::array::StructArray;
162/// # use arrow_schema::{Schema, Field, DataType};
163/// # use parquet_variant::Variant;
164/// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray, VariantType};
165/// # fn get_variant_array() -> VariantArray {
166/// #   let mut builder = VariantArrayBuilder::new(10);
167/// #   builder.append_variant(Variant::from("such wow"));
168/// #   builder.build()
169/// # }
170/// # fn get_schema() -> Schema {
171/// #   Schema::new(vec![
172/// #     Field::new("id", DataType::Int32, false),
173/// #     get_variant_array().field("var"),
174/// #   ])
175/// # }
176/// let schema = get_schema();
177/// assert_eq!(schema.fields().len(), 2);
178/// // first field is not a Variant
179/// assert!(!schema.field(0).has_valid_extension_type::<VariantType>());
180/// // second field is a Variant
181/// assert!(schema.field(1).has_valid_extension_type::<VariantType>());
182/// ```
183///
184/// # Example: Constructing the correct [`Field`] for a [`VariantArray`]
185///
186/// You can construct the correct [`Field`] for a [`VariantArray`] using the
187/// [`VariantArray::field`] method.
188///
189/// ```
190/// # use arrow_schema::{Schema, Field, DataType};
191/// # use parquet_variant::Variant;
192/// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray, VariantType};
193/// # fn get_variant_array() -> VariantArray {
194/// #   let mut builder = VariantArrayBuilder::new(10);
195/// #   builder.append_variant(Variant::from("such wow"));
196/// #   builder.build()
197/// # }
198/// let variant_array = get_variant_array();
199/// // First field is an integer id, second field is a variant
200/// let schema = Schema::new(vec![
201///   Field::new("id", DataType::Int32, false),
202///   // call VariantArray::field to get the correct Field
203///   variant_array.field("var"),
204/// ]);
205/// ```
206///
207/// You can also construct the [`Field`] using [`VariantType`] directly
208///
209/// ```
210/// # use arrow_schema::{Schema, Field, DataType};
211/// # use parquet_variant::Variant;
212/// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray, VariantType};
213/// # fn get_variant_array() -> VariantArray {
214/// #   let mut builder = VariantArrayBuilder::new(10);
215/// #   builder.append_variant(Variant::from("such wow"));
216/// #   builder.build()
217/// # }
218/// # let variant_array = get_variant_array();
219/// // The DataType of a VariantArray varies depending on how it is shredded
220/// let data_type = variant_array.data_type().clone();
221/// // First field is an integer id, second field is a variant
222/// let schema = Schema::new(vec![
223///   Field::new("id", DataType::Int32, false),
224///   Field::new("var", data_type, false)
225///     // Add extension metadata to the field using `VariantType`
226///     .with_extension_type(VariantType),
227/// ]);
228/// ```
229///
230/// # Example: Converting a [`VariantArray`] to a [`StructArray`]
231///
232/// ```
233/// # use arrow::array::StructArray;
234/// # use parquet_variant::Variant;
235/// # use parquet_variant_compute::VariantArrayBuilder;
236/// // Create Variant Array
237/// let mut builder = VariantArrayBuilder::new(10);
238/// builder.append_variant(Variant::from("such wow"));
239/// let variant_array = builder.build();
240/// // convert to StructArray
241/// let struct_array: StructArray = variant_array.into();
242/// ```
243///
244/// # Example: Converting a [`StructArray`] to a [`VariantArray`]
245///
246/// ```
247/// # use arrow::array::StructArray;
248/// # use parquet_variant::Variant;
249/// # use parquet_variant_compute::{VariantArrayBuilder, VariantArray};
250/// # fn get_struct_array() -> StructArray {
251/// #   let mut builder = VariantArrayBuilder::new(10);
252/// #   builder.append_variant(Variant::from("such wow"));
253/// #   builder.build().into()
254/// # }
255/// let struct_array: StructArray = get_struct_array();
256/// // try and create a VariantArray from it
257/// let variant_array = VariantArray::try_new(&struct_array).unwrap();
258/// assert_eq!(variant_array.value(0), Variant::from("such wow"));
259/// ```
260///
261#[derive(Debug, Clone)]
262pub struct VariantArray {
263    /// Reference to the underlying StructArray
264    inner: StructArray,
265
266    /// The metadata column of this variant (Binary, LargeBinary, or BinaryView)
267    metadata: ArrayRef,
268
269    /// how is this variant array shredded?
270    shredding_state: ShreddingState,
271}
272
273impl VariantArray {
274    /// Creates a new `VariantArray` from a [`StructArray`].
275    ///
276    /// # Arguments
277    /// - `inner` - The underlying [`StructArray`] that contains the variant data.
278    ///
279    /// # Returns
280    /// - A new instance of `VariantArray`.
281    ///
282    /// # Errors:
283    /// - If the `StructArray` does not contain the required fields
284    ///
285    /// # Requirements of the `StructArray`
286    ///
287    /// 1. A required field named `metadata` which is binary, large_binary, or
288    ///    binary_view
289    ///
290    /// 2. A required field named `value` that is binary, large_binary, or
291    ///    binary_view
292    ///
293    /// 3. An optional field named `typed_value` which can be any primitive type
294    ///    or be a list, large_list, list_view or struct
295    ///
296    /// NOTE: It is also permissible for the metadata field to be
297    /// Dictionary-Encoded, preferably (but not required) with an index type of
298    /// int8.
299    ///
300    pub fn try_new(inner: &dyn Array) -> Result<Self> {
301        // Canonicalize shredded typed_value fields (e.g. decimal narrowing)
302        let inner = canonicalize_shredded_types(inner)?;
303
304        let Some(inner) = inner.as_struct_opt() else {
305            return Err(ArrowError::InvalidArgumentError(
306                "Invalid VariantArray: requires StructArray as input".to_string(),
307            ));
308        };
309
310        // Note the specification allows for any order so we must search by name
311
312        // Ensure the StructArray has a metadata field that is a binary type
313        let Some(metadata_col) = inner.column_by_name("metadata") else {
314            return Err(ArrowError::InvalidArgumentError(
315                "Invalid VariantArray: StructArray must contain a 'metadata' field".to_string(),
316            ));
317        };
318        validate_binary_array(metadata_col.as_ref(), "metadata")?;
319
320        let shredding_state = ShreddingState::try_from(inner)?;
321
322        // `try_from` synthesizes an all-null `value` when the input omits it. Rebuild the inner
323        // struct so `inner()` and write-back carry the column too.
324        if inner.column_by_name("value").is_none() {
325            return Ok(Self::from_parts(
326                metadata_col.clone(),
327                shredding_state.value_column().clone(),
328                shredding_state.typed_value_column().cloned(),
329                inner.nulls().cloned(),
330            ));
331        }
332
333        // Note these clones are cheap, they just bump the ref count
334        Ok(Self {
335            inner: inner.clone(),
336            metadata: metadata_col.clone(),
337            shredding_state,
338        })
339    }
340
341    /// Note: annotates `value` as nullable, which the spec only permits for shredded
342    /// variants. It is also needed by `variant_get`'s unshredded intermediates, whose
343    /// `value` column can contain unmasked nulls. Unshredded producers should use
344    /// [`Self::from_parts_unshredded`] instead.
345    pub(crate) fn from_parts(
346        metadata: ArrayRef,
347        value: ArrayRef,
348        typed_value: Option<ArrayRef>,
349        nulls: Option<NullBuffer>,
350    ) -> Self {
351        Self::from_parts_with_nullable_value(metadata, value, typed_value, nulls, true)
352    }
353
354    /// Construct an unshredded `VariantArray`, annotating `value` as non-nullable as the
355    /// spec requires when there is no `typed_value` column.
356    ///
357    /// # Panics
358    /// If `value` contains nulls not masked by `nulls`.
359    pub(crate) fn from_parts_unshredded(
360        metadata: ArrayRef,
361        value: ArrayRef,
362        nulls: Option<NullBuffer>,
363    ) -> Self {
364        Self::from_parts_with_nullable_value(metadata, value, None, nulls, false)
365    }
366
367    fn from_parts_with_nullable_value(
368        metadata: ArrayRef,
369        value: ArrayRef,
370        typed_value: Option<ArrayRef>,
371        nulls: Option<NullBuffer>,
372        value_nullable: bool,
373    ) -> Self {
374        let mut builder = StructArrayBuilder::new()
375            .with_field("metadata", metadata.clone(), false)
376            .with_field("value", value.clone(), value_nullable);
377        if let Some(typed_value) = typed_value.clone() {
378            builder = builder.with_field_ref(typed_value_field(&typed_value), typed_value);
379        }
380        if let Some(nulls) = nulls {
381            builder = builder.with_nulls(nulls);
382        }
383
384        Self {
385            inner: builder.build(),
386            metadata,
387            shredding_state: ShreddingState::new(value, typed_value),
388        }
389    }
390
391    /// Returns a reference to the underlying [`StructArray`].
392    pub fn inner(&self) -> &StructArray {
393        &self.inner
394    }
395
396    /// Returns the inner [`StructArray`], consuming self
397    pub fn into_inner(self) -> StructArray {
398        self.inner
399    }
400
401    /// Return the shredding state of this `VariantArray`
402    pub fn shredding_state(&self) -> &ShreddingState {
403        &self.shredding_state
404    }
405
406    /// Return the [`Variant`] instance stored at the given row
407    ///
408    /// This is a convenience wrapper that calls [`VariantArray::try_value`] and unwraps the `Result`.
409    /// Use `try_value` if you need to handle conversion errors gracefully.
410    ///
411    /// # Panics
412    /// * if the index is out of bounds
413    /// * if the array value is null
414    /// * if `try_value` returns an error.
415    pub fn value(&self, index: usize) -> Variant<'_, '_> {
416        self.try_value(index).unwrap()
417    }
418
419    /// Return the [`Variant`] instance stored at the given row
420    ///
421    /// Note: This method does not check for nulls and the value is arbitrary
422    /// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
423    ///
424    /// # Panics
425    ///
426    /// Panics if
427    /// * the index is out of bounds
428    /// * the array value is null
429    ///
430    /// # Errors
431    ///
432    /// Errors if
433    /// - the data in `typed_value` cannot be interpreted as a valid `Variant`
434    ///
435    /// If this is a shredded variant but has no value at the shredded location, it
436    /// will return [`Variant::Null`].
437    ///
438    ///
439    /// # Performance Note
440    ///
441    /// This is certainly not the most efficient way to access values in a
442    /// `VariantArray`, but it is useful for testing and debugging.
443    ///
444    /// Note: Does not do deep validation of the [`Variant`], so it is up to the
445    /// caller to ensure that the metadata and value were constructed correctly.
446    pub fn try_value(&self, index: usize) -> Result<Variant<'_, '_>> {
447        let value = self.value_column();
448        match self.typed_value_column() {
449            // Always prefer typed_value, if available
450            Some(typed_value) if typed_value.is_valid(index) => {
451                if !matches!(typed_value.data_type(), DataType::Struct(_)) && value.is_valid(index) {
452                    // Only a partially shredded struct is allowed to have values for both columns
453                    panic!("Invalid variant, conflicting value and typed_value");
454                }
455                typed_value_to_variant(typed_value, index)
456            }
457            // Otherwise fall back to value, if available
458            _ if value.is_valid(index) => variant_from_arrays_at(&self.metadata, value, index)
459                .ok_or_else(|| {
460                    ArrowError::InvalidArgumentError(format!(
461                        "metadata and value fields must be binary-like arrays, instead got {} and {}",
462                        self.metadata.data_type(),
463                        value.data_type()
464                    ))
465                }),
466            // It is technically invalid for both value and typed_value to be null,
467            // but the spec specifically requires readers to return Variant::Null in this case.
468            _ => Ok(Variant::Null),
469        }
470    }
471
472    /// Return a reference to the `metadata` column of the [`StructArray`]
473    pub fn metadata_column(&self) -> &ArrayRef {
474        &self.metadata
475    }
476
477    /// Return a reference to the `value` column of the [`StructArray`]
478    pub fn value_column(&self) -> &ArrayRef {
479        self.shredding_state.value_column()
480    }
481
482    /// Return a reference to the `typed_value` column of the [`StructArray`], if present
483    pub fn typed_value_column(&self) -> Option<&ArrayRef> {
484        self.shredding_state.typed_value_column()
485    }
486
487    /// Return a field to represent this VariantArray in a `Schema` with
488    /// a particular name
489    pub fn field(&self, name: impl Into<String>) -> Field {
490        Field::new(
491            name.into(),
492            self.data_type().clone(),
493            self.inner.is_nullable(),
494        )
495        .with_extension_type(VariantType)
496    }
497
498    /// Returns a new DataType representing this VariantArray's inner type
499    pub fn data_type(&self) -> &DataType {
500        self.inner.data_type()
501    }
502
503    pub fn slice(&self, offset: usize, length: usize) -> Self {
504        let inner = self.inner.slice(offset, length);
505        let metadata = self.metadata.slice(offset, length);
506        let shredding_state = self.shredding_state.slice(offset, length);
507        Self {
508            inner,
509            metadata,
510            shredding_state,
511        }
512    }
513
514    pub fn len(&self) -> usize {
515        self.inner.len()
516    }
517
518    pub fn is_empty(&self) -> bool {
519        self.inner.is_empty()
520    }
521
522    pub fn nulls(&self) -> Option<&NullBuffer> {
523        self.inner.nulls()
524    }
525
526    /// Is the element at index null?
527    pub fn is_null(&self, index: usize) -> bool {
528        self.nulls().is_some_and(|n| n.is_null(index))
529    }
530
531    /// Is the element at index valid (not null)?
532    pub fn is_valid(&self, index: usize) -> bool {
533        !self.is_null(index)
534    }
535
536    /// Returns an iterator over the values in this array
537    pub fn iter(&self) -> VariantArrayIter<'_> {
538        VariantArrayIter::new(self)
539    }
540}
541
542impl PartialEq for VariantArray {
543    fn eq(&self, other: &Self) -> bool {
544        self.inner == other.inner
545    }
546}
547
548impl From<VariantArray> for StructArray {
549    fn from(variant_array: VariantArray) -> Self {
550        variant_array.into_inner()
551    }
552}
553
554impl From<VariantArray> for ArrayRef {
555    fn from(variant_array: VariantArray) -> Self {
556        Arc::new(variant_array.into_inner())
557    }
558}
559
560impl<'m, 'v> FromIterator<Option<Variant<'m, 'v>>> for VariantArray {
561    fn from_iter<T: IntoIterator<Item = Option<Variant<'m, 'v>>>>(iter: T) -> Self {
562        let iter = iter.into_iter();
563
564        let mut b = VariantArrayBuilder::new(iter.size_hint().0);
565        b.extend(iter);
566        b.build()
567    }
568}
569
570impl<'m, 'v> FromIterator<Variant<'m, 'v>> for VariantArray {
571    fn from_iter<T: IntoIterator<Item = Variant<'m, 'v>>>(iter: T) -> Self {
572        Self::from_iter(iter.into_iter().map(Some))
573    }
574}
575
576/// An iterator over [`VariantArray`]
577///
578/// This iterator returns `Option<Option<Variant<'a, 'a>>>` where:
579/// - `None` indicates the end of iteration
580/// - `Some(None)` indicates a null value at this position
581/// - `Some(Some(variant))` indicates a valid variant value
582///
583/// # Example
584///
585/// ```
586/// # use parquet_variant::Variant;
587/// # use parquet_variant_compute::VariantArrayBuilder;
588/// let mut builder = VariantArrayBuilder::new(10);
589/// builder.append_variant(Variant::from(42));
590/// builder.append_null();
591/// builder.append_variant(Variant::from("hello"));
592/// let array = builder.build();
593///
594/// let values = array.iter().collect::<Vec<_>>();
595/// assert_eq!(values.len(), 3);
596/// assert_eq!(values[0], Some(Variant::from(42)));
597/// assert_eq!(values[1], None);
598/// assert_eq!(values[2], Some(Variant::from("hello")));
599/// ```
600#[derive(Debug)]
601pub struct VariantArrayIter<'a> {
602    array: &'a VariantArray,
603    head_i: usize,
604    tail_i: usize,
605}
606
607impl<'a> VariantArrayIter<'a> {
608    /// Creates a new iterator over the given [`VariantArray`]
609    pub fn new(array: &'a VariantArray) -> Self {
610        Self {
611            array,
612            head_i: 0,
613            tail_i: array.len(),
614        }
615    }
616
617    fn value_opt(&self, i: usize) -> Option<Variant<'a, 'a>> {
618        self.array.is_valid(i).then(|| self.array.value(i))
619    }
620}
621
622impl<'a> Iterator for VariantArrayIter<'a> {
623    type Item = Option<Variant<'a, 'a>>;
624
625    #[inline]
626    fn next(&mut self) -> Option<Self::Item> {
627        if self.head_i == self.tail_i {
628            return None;
629        }
630
631        let out = self.value_opt(self.head_i);
632
633        self.head_i += 1;
634
635        Some(out)
636    }
637
638    fn size_hint(&self) -> (usize, Option<usize>) {
639        let remainder = self.tail_i - self.head_i;
640
641        (remainder, Some(remainder))
642    }
643}
644
645impl<'a> DoubleEndedIterator for VariantArrayIter<'a> {
646    fn next_back(&mut self) -> Option<Self::Item> {
647        if self.head_i == self.tail_i {
648            return None;
649        }
650
651        self.tail_i -= 1;
652
653        Some(self.value_opt(self.tail_i))
654    }
655}
656
657impl<'a> ExactSizeIterator for VariantArrayIter<'a> {}
658
659/// One shredded field of a partially or perfectly shredded variant. For example, suppose the
660/// shredding schema for variant `v` treats it as an object with a single field `a`, where `a` is
661/// itself a struct with the single field `b` of type INT. Then the physical layout of the column
662/// is:
663///
664/// ```text
665/// v: VARIANT {
666///     metadata: BINARY,
667///     value: BINARY,
668///     typed_value: STRUCT {
669///         a: SHREDDED_VARIANT_FIELD {
670///             value: BINARY,
671///             typed_value: STRUCT {
672///                 a: SHREDDED_VARIANT_FIELD {
673///                     value: BINARY,
674///                     typed_value: INT,
675///                 },
676///             },
677///         },
678///     },
679/// }
680/// ```
681///
682/// In the above, each row of `v.value` is either a variant value (shredding failed, `v` was not an
683/// object at all) or a variant object (partial shredding, `v` was an object but included unexpected
684/// fields other than `a`), or is NULL (perfect shredding, `v` was an object containing only the
685/// single expected field `a`).
686///
687/// A similar story unfolds for each `v.typed_value.a.value` -- a variant value if shredding failed
688/// (`v:a` was not an object at all), or a variant object (`v:a` was an object with unexpected
689/// additional fields), or NULL (`v:a` was an object containing only the single expected field `b`).
690///
691/// Finally, `v.typed_value.a.typed_value.b.value` is either NULL (`v:a.b` was an integer) or else a
692/// variant value (which could be `Variant::Null`).
693#[derive(Debug)]
694pub struct ShreddedVariantFieldArray {
695    /// Reference to the underlying StructArray
696    inner: StructArray,
697    shredding_state: ShreddingState,
698}
699
700impl ShreddedVariantFieldArray {
701    /// Creates a new `ShreddedVariantFieldArray` from a [`StructArray`].
702    ///
703    /// # Arguments
704    /// - `inner` - The underlying [`StructArray`] that contains the variant data.
705    ///
706    /// # Returns
707    /// - A new instance of `ShreddedVariantFieldArray`.
708    ///
709    /// # Errors:
710    /// - If the `StructArray` does not contain the required fields
711    ///
712    /// # Requirements of the `StructArray`
713    ///
714    /// 1. A required field named `value` that is binary, large_binary, or
715    ///    binary_view
716    ///
717    /// 2. An optional field named `typed_value` which can be any primitive type
718    ///    or be a list, large_list, list_view or struct
719    ///
720    pub fn try_new(inner: &dyn Array) -> Result<Self> {
721        let Some(inner_struct) = inner.as_struct_opt() else {
722            return Err(ArrowError::InvalidArgumentError(
723                "Invalid ShreddedVariantFieldArray: requires StructArray as input".to_string(),
724            ));
725        };
726
727        let shredding_state = ShreddingState::try_from(inner_struct)?;
728
729        // `try_from` synthesizes an all-null `value` when the input omits it. Rebuild the inner
730        // struct so `inner()` and write-back carry the column too (see `VariantArray::try_new`).
731        if inner_struct.column_by_name("value").is_none() {
732            return Ok(Self::from_parts(
733                shredding_state.value_column().clone(),
734                shredding_state.typed_value_column().cloned(),
735                inner_struct.nulls().cloned(),
736            ));
737        }
738
739        // Note this clone is cheap, it just bumps the ref count
740        Ok(Self {
741            inner: inner_struct.clone(),
742            shredding_state,
743        })
744    }
745
746    /// Return the shredding state of this `VariantArray`
747    pub fn shredding_state(&self) -> &ShreddingState {
748        &self.shredding_state
749    }
750
751    /// Return a reference to the `value` column of the [`StructArray`]
752    pub fn value_column(&self) -> &ArrayRef {
753        self.shredding_state.value_column()
754    }
755
756    /// Return a reference to the `typed_value` column of the [`StructArray`], if present
757    pub fn typed_value_column(&self) -> Option<&ArrayRef> {
758        self.shredding_state.typed_value_column()
759    }
760
761    /// Returns a reference to the underlying [`StructArray`].
762    pub fn inner(&self) -> &StructArray {
763        &self.inner
764    }
765
766    pub(crate) fn from_parts(
767        value: ArrayRef,
768        typed_value: Option<ArrayRef>,
769        nulls: Option<NullBuffer>,
770    ) -> Self {
771        let mut builder = StructArrayBuilder::new().with_field("value", value.clone(), true);
772        if let Some(typed_value) = typed_value.clone() {
773            builder = builder.with_field_ref(typed_value_field(&typed_value), typed_value);
774        }
775        if let Some(nulls) = nulls {
776            builder = builder.with_nulls(nulls);
777        }
778
779        Self {
780            inner: builder.build(),
781            shredding_state: ShreddingState::new(value, typed_value),
782        }
783    }
784
785    /// Returns the inner [`StructArray`], consuming self
786    pub fn into_inner(self) -> StructArray {
787        self.inner
788    }
789
790    pub fn data_type(&self) -> &DataType {
791        self.inner.data_type()
792    }
793
794    pub fn len(&self) -> usize {
795        self.inner.len()
796    }
797
798    pub fn is_empty(&self) -> bool {
799        self.inner.is_empty()
800    }
801
802    pub fn offset(&self) -> usize {
803        self.inner.offset()
804    }
805
806    pub fn nulls(&self) -> Option<&NullBuffer> {
807        // According to the shredding spec, ShreddedVariantFieldArray should be
808        // physically non-nullable - SQL NULL is inferred by both value and
809        // typed_value being physically NULL
810        None
811    }
812    /// Is the element at index null?
813    pub fn is_null(&self, index: usize) -> bool {
814        self.nulls().is_some_and(|n| n.is_null(index))
815    }
816
817    /// Is the element at index valid (not null)?
818    pub fn is_valid(&self, index: usize) -> bool {
819        !self.is_null(index)
820    }
821}
822
823impl From<ShreddedVariantFieldArray> for ArrayRef {
824    fn from(array: ShreddedVariantFieldArray) -> Self {
825        Arc::new(array.into_inner())
826    }
827}
828
829impl From<ShreddedVariantFieldArray> for StructArray {
830    fn from(array: ShreddedVariantFieldArray) -> Self {
831        array.into_inner()
832    }
833}
834
835/// Represents the shredding state of a [`VariantArray`]
836///
837/// [`VariantArray`]s can be shredded according to the [Parquet Variant
838/// Shredding Spec]. Shredding means that the actual value is stored in a typed
839/// `typed_field` instead of the generic `value` field.
840///
841/// The `value` column is always present (the spec requires writers to emit
842/// it); `typed_value` is optional. Values in the two columns must be
843/// interpreted according to the following table (see [Parquet Variant
844/// Shredding Spec] for more details):
845///
846/// | value    | typed_value  | Meaning |
847/// |----------|--------------|---------|
848/// | NULL     | NULL         | The value is missing; only valid for shredded object fields |
849/// | non-NULL | NULL         | The value is present and may be any type, including [`Variant::Null`] |
850/// | NULL     | non-NULL     | The value is present and is the shredded type |
851/// | non-NULL | non-NULL     | The value is present and is a partially shredded object |
852///
853///
854/// Applying the above rules to entire columns, we obtain the following:
855///
856/// | value  | typed_value  | Meaning |
857/// |--------|-------------|---------|
858/// | exists | --          | **Unshredded**: If present, the value may be any type, including [`Variant::Null`]
859/// | exists | exists      | **Shredded**: perfectly if `value` is all-null, otherwise imperfectly |
860///
861/// Note the spec requires the `value` column to always be present in the
862/// schema; structs without one are rejected
863/// (see <https://github.com/apache/arrow-rs/issues/10306>).
864///
865/// NOTE: Partial shredding is a row-wise situation that can arise under imperfect shredding (a
866/// column-wise situation): When both columns exist (imperfect shredding) and the typed_value column
867/// is a struct, then both columns can be non-NULL for the same row if value is a variant object
868/// (partial shredding).
869///
870/// [Parquet Variant Shredding Spec]: https://github.com/apache/parquet-format/blob/master/VariantShredding.md#value-shredding
871#[derive(Debug, Clone)]
872pub struct ShreddingState {
873    value: ArrayRef,
874    typed_value: Option<ArrayRef>,
875}
876
877impl ShreddingState {
878    /// Create a new `ShreddingState` from the given `value` and `typed_value` fields
879    ///
880    /// Note you can create a `ShreddingState` from a &[`StructArray`] using
881    /// `ShreddingState::try_from(&struct_array)`, for example:
882    ///
883    /// ```no_run
884    /// # use arrow::array::StructArray;
885    /// # use parquet_variant_compute::ShreddingState;
886    /// # fn get_struct_array() -> StructArray {
887    /// #   unimplemented!()
888    /// # }
889    /// let struct_array: StructArray = get_struct_array();
890    /// let shredding_state = ShreddingState::try_from(&struct_array).unwrap();
891    /// ```
892    pub fn new(value: ArrayRef, typed_value: Option<ArrayRef>) -> Self {
893        Self { value, typed_value }
894    }
895
896    /// Return a reference to the `value` column
897    pub fn value_column(&self) -> &ArrayRef {
898        &self.value
899    }
900
901    /// Return a reference to the `typed_value` column, if present
902    pub fn typed_value_column(&self) -> Option<&ArrayRef> {
903        self.typed_value.as_ref()
904    }
905
906    /// Slice all the underlying arrays
907    pub fn slice(&self, offset: usize, length: usize) -> Self {
908        Self {
909            value: self.value.slice(offset, length),
910            typed_value: self.typed_value.as_ref().map(|tv| tv.slice(offset, length)),
911        }
912    }
913}
914
915impl TryFrom<&StructArray> for ShreddingState {
916    type Error = ArrowError;
917
918    fn try_from(inner_struct: &StructArray) -> Result<Self> {
919        let typed_value = inner_struct.column_by_name("typed_value").cloned();
920        let value = match inner_struct.column_by_name("value") {
921            Some(value) => {
922                validate_binary_array(value.as_ref(), "value")?;
923                value.clone()
924            }
925            // Lenient read: a shredded group may omit the spec-required `value`. Synthesize an
926            // all-null one so the model always has a `value`.
927            None if typed_value.is_some() => all_null_value_column(inner_struct.len()),
928            None => {
929                return Err(ArrowError::InvalidArgumentError(
930                    "Invalid VariantArray: StructArray must contain a 'value' field".to_string(),
931                ));
932            }
933        };
934        Ok(ShreddingState::new(value, typed_value))
935    }
936}
937
938/// Build the `typed_value` [`FieldRef`] for a shredded column.
939///
940/// The Variant spec maps `FixedSizeBinary(16)` exclusively to UUID, so any
941/// shredded column of that type must carry the canonical [`UuidExtension`]
942/// extension metadata on its field.
943fn typed_value_field(array: &ArrayRef) -> FieldRef {
944    let mut field = Field::new("typed_value", array.data_type().clone(), true);
945    if matches!(array.data_type(), DataType::FixedSizeBinary(16)) {
946        field = field.with_extension_type(UuidExtension);
947    }
948    Arc::new(field)
949}
950
951/// Builds struct arrays from component fields
952///
953/// TODO: move to arrow crate
954#[derive(Debug, Default, Clone)]
955pub(crate) struct StructArrayBuilder {
956    fields: Vec<FieldRef>,
957    arrays: Vec<ArrayRef>,
958    nulls: Option<NullBuffer>,
959}
960
961impl StructArrayBuilder {
962    pub fn new() -> Self {
963        Default::default()
964    }
965
966    /// Add an array to this struct array as a field with the specified name.
967    pub fn with_field(mut self, field_name: &str, array: ArrayRef, nullable: bool) -> Self {
968        let field = Field::new(field_name, array.data_type().clone(), nullable);
969        self.fields.push(Arc::new(field));
970        self.arrays.push(array);
971        self
972    }
973
974    /// Add an array to this struct array using a caller-supplied [`FieldRef`].
975    ///
976    /// Use this when the field carries metadata (e.g. an extension type) that
977    /// would be lost if the field were synthesized from the array's data type alone.
978    pub fn with_field_ref(mut self, field: FieldRef, array: ArrayRef) -> Self {
979        self.fields.push(field);
980        self.arrays.push(array);
981        self
982    }
983
984    /// Set the null buffer for this struct array.
985    pub fn with_nulls(mut self, nulls: NullBuffer) -> Self {
986        self.nulls = Some(nulls);
987        self
988    }
989
990    pub fn build(self) -> StructArray {
991        let Self {
992            fields,
993            arrays,
994            nulls,
995        } = self;
996        StructArray::new(Fields::from(fields), arrays, nulls)
997    }
998}
999
1000/// returns the non-null element at index as a Variant
1001fn typed_value_to_variant(typed_value: &ArrayRef, index: usize) -> Result<Variant<'_, '_>> {
1002    let data_type = typed_value.data_type();
1003    match data_type {
1004        DataType::Null => Ok(Variant::Null),
1005        DataType::Boolean => {
1006            let boolean_array = typed_value.as_boolean();
1007            let value = boolean_array.value(index);
1008            Ok(Variant::from(value))
1009        }
1010        // 16-byte FixedSizeBinary alway corresponds to a UUID; all other sizes are illegal.
1011        DataType::FixedSizeBinary(16) => {
1012            let array = typed_value.as_fixed_size_binary();
1013            let value = array.value(index);
1014            Ok(Uuid::from_slice(value).unwrap().into()) // unwrap is safe: slice is always 16 bytes
1015        }
1016        DataType::Binary => {
1017            let array = typed_value.as_binary::<i32>();
1018            let value = array.value(index);
1019            Ok(Variant::from(value))
1020        }
1021        DataType::LargeBinary => {
1022            let array = typed_value.as_binary::<i64>();
1023            let value = array.value(index);
1024            Ok(Variant::from(value))
1025        }
1026        DataType::BinaryView => {
1027            let array = typed_value.as_binary_view();
1028            let value = array.value(index);
1029            Ok(Variant::from(value))
1030        }
1031        DataType::Utf8 => {
1032            let array = typed_value.as_string::<i32>();
1033            let value = array.value(index);
1034            Ok(Variant::from(value))
1035        }
1036        DataType::LargeUtf8 => {
1037            let array = typed_value.as_string::<i64>();
1038            let value = array.value(index);
1039            Ok(Variant::from(value))
1040        }
1041        DataType::Utf8View => {
1042            let array = typed_value.as_string_view();
1043            let value = array.value(index);
1044            Ok(Variant::from(value))
1045        }
1046        DataType::Int8 => {
1047            primitive_conversion_single_value!(Int8Type, typed_value, index)
1048        }
1049        DataType::Int16 => {
1050            primitive_conversion_single_value!(Int16Type, typed_value, index)
1051        }
1052        DataType::Int32 => {
1053            primitive_conversion_single_value!(Int32Type, typed_value, index)
1054        }
1055        DataType::Int64 => {
1056            primitive_conversion_single_value!(Int64Type, typed_value, index)
1057        }
1058        DataType::Float16 => {
1059            primitive_conversion_single_value!(Float16Type, typed_value, index)
1060        }
1061        DataType::Float32 => {
1062            primitive_conversion_single_value!(Float32Type, typed_value, index)
1063        }
1064        DataType::Float64 => {
1065            primitive_conversion_single_value!(Float64Type, typed_value, index)
1066        }
1067        DataType::Decimal32(_, s) => {
1068            generic_conversion_single_value_with_result!(
1069                Decimal32Type,
1070                as_primitive,
1071                |v| VariantDecimal4::try_new(v, *s as u8),
1072                typed_value,
1073                index
1074            )
1075        }
1076        DataType::Decimal64(_, s) => {
1077            generic_conversion_single_value_with_result!(
1078                Decimal64Type,
1079                as_primitive,
1080                |v| VariantDecimal8::try_new(v, *s as u8),
1081                typed_value,
1082                index
1083            )
1084        }
1085        DataType::Decimal128(_, s) => {
1086            generic_conversion_single_value_with_result!(
1087                Decimal128Type,
1088                as_primitive,
1089                |v| VariantDecimal16::try_new(v, *s as u8),
1090                typed_value,
1091                index
1092            )
1093        }
1094        DataType::Date32 => {
1095            generic_conversion_single_value!(
1096                Date32Type,
1097                as_primitive,
1098                |v| Date32Type::to_naive_date_opt(v).unwrap(),
1099                typed_value,
1100                index
1101            )
1102        }
1103        DataType::Time64(TimeUnit::Microsecond) => {
1104            generic_conversion_single_value_with_result!(
1105                Time64MicrosecondType,
1106                as_primitive,
1107                |v| NaiveTime::from_num_seconds_from_midnight_opt(
1108                    (v / 1_000_000) as u32,
1109                    (v % 1_000_000) as u32 * 1000
1110                )
1111                .ok_or_else(|| format!("Invalid microsecond from midnight: {}", v)),
1112                typed_value,
1113                index
1114            )
1115        }
1116        DataType::Timestamp(TimeUnit::Microsecond, Some(_)) => {
1117            generic_conversion_single_value!(
1118                TimestampMicrosecondType,
1119                as_primitive,
1120                |v| DateTime::from_timestamp_micros(v).unwrap(),
1121                typed_value,
1122                index
1123            )
1124        }
1125        DataType::Timestamp(TimeUnit::Microsecond, None) => {
1126            generic_conversion_single_value!(
1127                TimestampMicrosecondType,
1128                as_primitive,
1129                |v| DateTime::from_timestamp_micros(v).unwrap().naive_utc(),
1130                typed_value,
1131                index
1132            )
1133        }
1134        DataType::Timestamp(TimeUnit::Nanosecond, Some(_)) => {
1135            generic_conversion_single_value!(
1136                TimestampNanosecondType,
1137                as_primitive,
1138                DateTime::from_timestamp_nanos,
1139                typed_value,
1140                index
1141            )
1142        }
1143        DataType::Timestamp(TimeUnit::Nanosecond, None) => {
1144            generic_conversion_single_value!(
1145                TimestampNanosecondType,
1146                as_primitive,
1147                |v| DateTime::from_timestamp_nanos(v).naive_utc(),
1148                typed_value,
1149                index
1150            )
1151        }
1152        // todo other types here (note this is very similar to cast_to_variant.rs)
1153        // so it would be great to figure out how to share this code
1154        //
1155        // Composite shredded values may require combining `value` and
1156        // `typed_value` and allocating new encoded bytes. `try_value` returns
1157        // borrowed Variant, so callers must unshred the array first.
1158        _ => Err(ArrowError::NotYetImplemented(format!(
1159            "VariantArray::try_value cannot materialize typed_value of type {} \
1160             as a borrowed Variant; call unshred_variant first",
1161            typed_value.data_type()
1162        ))),
1163    }
1164}
1165
1166/// Canonicalize shredded typed_value fields (e.g. decimal narrowing) and
1167/// verify that all data types in the struct are legal for a variant array.
1168fn canonicalize_shredded_types(array: &dyn Array) -> Result<ArrayRef> {
1169    let new_type = canonicalize_and_verify_data_type(array.data_type())?;
1170    if let Cow::Borrowed(_) = new_type
1171        && let Some(array) = array.as_struct_opt()
1172    {
1173        return Ok(Arc::new(array.clone())); // bypass the unnecessary cast
1174    }
1175    cast(array, new_type.as_ref())
1176}
1177
1178/// Recursively visits a data type, ensuring that it only contains data types that can legally
1179/// appear in a (possibly shredded) variant array. It also narrows decimal types to the smallest
1180/// valid precision (e.g. Decimal128 -> Decimal32 when the precision fits).
1181fn canonicalize_and_verify_data_type(data_type: &DataType) -> Result<Cow<'_, DataType>> {
1182    use DataType::*;
1183
1184    // helper macros
1185    macro_rules! fail {
1186        () => {
1187            return Err(ArrowError::InvalidArgumentError(format!(
1188                "Illegal shredded value type: {data_type}"
1189            )))
1190        };
1191    }
1192    macro_rules! borrow {
1193        () => {
1194            Cow::Borrowed(data_type)
1195        };
1196    }
1197
1198    let new_data_type = match data_type {
1199        // Primitive arrow types that have a direct variant counterpart are allowed
1200        Null | Boolean => borrow!(),
1201        Int8 | Int16 | Int32 | Int64 | Float32 | Float64 => borrow!(),
1202
1203        // Unsigned integers and half-float are not allowed
1204        UInt8 | UInt16 | UInt32 | UInt64 | Float16 => fail!(),
1205
1206        // Most decimal types are allowed, with restrictions on precision and scale
1207        //
1208        // NOTE: arrow-parquet reads widens 32- and 64-bit decimals to 128-bit, but the variant spec
1209        // requires using the narrowest decimal type for a given precision. Fix those up first.
1210        Decimal64(p, s) | Decimal128(p, s)
1211            if VariantDecimal4::is_valid_precision_and_scale(p, s) =>
1212        {
1213            Cow::Owned(Decimal32(*p, *s))
1214        }
1215        Decimal128(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => {
1216            Cow::Owned(Decimal64(*p, *s))
1217        }
1218        Decimal32(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => borrow!(),
1219        Decimal64(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => borrow!(),
1220        Decimal128(p, s) if VariantDecimal16::is_valid_precision_and_scale(p, s) => borrow!(),
1221        Decimal32(..) | Decimal64(..) | Decimal128(..) | Decimal256(..) => fail!(),
1222
1223        // Only micro and nano timestamps are allowed
1224        Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, _) => borrow!(),
1225        Timestamp(TimeUnit::Millisecond | TimeUnit::Second, _) => fail!(),
1226
1227        // Only 32-bit dates and 64-bit microsecond time are allowed.
1228        Date32 | Time64(TimeUnit::Microsecond) => borrow!(),
1229        Date64 | Time32(_) | Time64(_) | Duration(_) | Interval(_) => fail!(),
1230
1231        // Binary, string, and their view counterparts are allowed.
1232        Binary | LargeBinary | BinaryView | Utf8 | LargeUtf8 | Utf8View => borrow!(),
1233
1234        // UUID maps to 16-byte fixed-size binary; no other width is allowed
1235        FixedSizeBinary(16) => borrow!(),
1236        FixedSizeBinary(_) | FixedSizeList(..) => fail!(),
1237
1238        // List-like containers and struct are allowed, maps and unions are not
1239        List(field) => match canonicalize_and_verify_field(field)? {
1240            Cow::Borrowed(_) => borrow!(),
1241            Cow::Owned(new_field) => Cow::Owned(DataType::List(new_field)),
1242        },
1243        LargeList(field) => match canonicalize_and_verify_field(field)? {
1244            Cow::Borrowed(_) => borrow!(),
1245            Cow::Owned(new_field) => Cow::Owned(DataType::LargeList(new_field)),
1246        },
1247        ListView(field) => match canonicalize_and_verify_field(field)? {
1248            Cow::Borrowed(_) => borrow!(),
1249            Cow::Owned(new_field) => Cow::Owned(DataType::ListView(new_field)),
1250        },
1251        LargeListView(field) => match canonicalize_and_verify_field(field)? {
1252            Cow::Borrowed(_) => borrow!(),
1253            Cow::Owned(new_field) => Cow::Owned(DataType::LargeListView(new_field)),
1254        },
1255        // Struct is used by the internal layout, and can also represent a shredded variant object.
1256        Struct(fields) => {
1257            // Avoid allocation unless at least one field changes, to avoid unnecessary deep cloning
1258            // of the data type. Even if some fields change, the others are shallow arc clones.
1259            let mut new_fields = std::collections::HashMap::new();
1260            for (i, field) in fields.iter().enumerate() {
1261                if let Cow::Owned(new_field) = canonicalize_and_verify_field(field)? {
1262                    new_fields.insert(i, new_field);
1263                }
1264            }
1265
1266            if new_fields.is_empty() {
1267                borrow!()
1268            } else {
1269                let new_fields = fields
1270                    .iter()
1271                    .enumerate()
1272                    .map(|(i, field)| new_fields.remove(&i).unwrap_or_else(|| field.clone()));
1273                Cow::Owned(DataType::Struct(new_fields.collect()))
1274            }
1275        }
1276        Map(..) | Union(..) => fail!(),
1277
1278        // We can _possibly_ support (some of) these some day?
1279        Dictionary(..) | RunEndEncoded(..) => fail!(),
1280    };
1281    Ok(new_data_type)
1282}
1283
1284fn canonicalize_and_verify_field(field: &Arc<Field>) -> Result<Cow<'_, Arc<Field>>> {
1285    let new_data_type = canonicalize_and_verify_data_type(field.data_type())?;
1286
1287    // A shredded FixedSizeBinary(16) column is always a UUID. Tag it with the UUID extension type
1288    // on read, as a safety net against writers that emit the column without the extension metadata.
1289    // Canonicalization never rewrites FixedSizeBinary(16), so the type is already correct here.
1290    if matches!(new_data_type.as_ref(), DataType::FixedSizeBinary(16))
1291        && !field.has_valid_extension_type::<UuidExtension>()
1292    {
1293        let new_field = field.as_ref().clone().with_extension_type(UuidExtension);
1294        return Ok(Cow::Owned(Arc::new(new_field)));
1295    }
1296
1297    let Cow::Owned(new_data_type) = new_data_type else {
1298        return Ok(Cow::Borrowed(field));
1299    };
1300    let new_field = field.as_ref().clone().with_data_type(new_data_type);
1301    Ok(Cow::Owned(Arc::new(new_field)))
1302}
1303
1304/// Test-only constructors for perfectly shredded arrays, where every value lives in
1305/// `typed_value` and the required `value` column is all-null.
1306#[cfg(test)]
1307impl VariantArray {
1308    pub(crate) fn perfectly_shredded(
1309        metadata: ArrayRef,
1310        typed_value: ArrayRef,
1311        nulls: Option<NullBuffer>,
1312    ) -> Self {
1313        let value = all_null_value_column(typed_value.len());
1314        Self::from_parts(metadata, value, Some(typed_value), nulls)
1315    }
1316}
1317
1318#[cfg(test)]
1319impl ShreddedVariantFieldArray {
1320    pub(crate) fn perfectly_shredded(typed_value: ArrayRef) -> Self {
1321        let value = all_null_value_column(typed_value.len());
1322        Self::from_parts(value, Some(typed_value), None)
1323    }
1324}
1325
1326#[cfg(test)]
1327mod test {
1328    use crate::VariantArrayBuilder;
1329    use std::str::FromStr;
1330
1331    use super::*;
1332    use arrow::array::{
1333        BinaryArray, BinaryViewArray, Decimal32Array, Decimal64Array, Decimal128Array,
1334        FixedSizeBinaryArray, Int32Array, Int64Array, LargeBinaryArray, LargeListArray,
1335        LargeListViewArray, ListArray, ListViewArray, Time64MicrosecondArray,
1336    };
1337    use arrow::buffer::{OffsetBuffer, ScalarBuffer};
1338    use arrow_schema::{Field, Fields};
1339    use parquet_variant::{EMPTY_VARIANT_METADATA_BYTES, ShortString};
1340
1341    #[test]
1342    fn invalid_not_a_struct_array() {
1343        let array = make_binary_view_array();
1344        // Should fail because the input is not a StructArray
1345        let err = VariantArray::try_new(&array);
1346        assert_eq!(
1347            err.unwrap_err().to_string(),
1348            "Invalid argument error: Invalid VariantArray: requires StructArray as input"
1349        );
1350    }
1351
1352    #[test]
1353    fn invalid_missing_metadata() {
1354        let fields = Fields::from(vec![Field::new("value", DataType::BinaryView, true)]);
1355        let array = StructArray::new(fields, vec![make_binary_view_array()], None);
1356        // Should fail because the StructArray does not contain a 'metadata' field
1357        let err = VariantArray::try_new(&array);
1358        assert_eq!(
1359            err.unwrap_err().to_string(),
1360            "Invalid argument error: Invalid VariantArray: StructArray must contain a 'metadata' field"
1361        );
1362    }
1363
1364    #[test]
1365    fn read_missing_value_column() {
1366        // With `typed_value` present, a missing `value` is read leniently: an all-null `value` is
1367        // synthesized and normalized into the inner struct (so a write-back emits it), and values
1368        // still decode from `typed_value`.
1369        let typed_value = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef;
1370        let metadata =
1371            BinaryViewArray::from_iter_values(std::iter::repeat_n(EMPTY_VARIANT_METADATA_BYTES, 3));
1372        let struct_array = StructArrayBuilder::new()
1373            .with_field("metadata", Arc::new(metadata), false)
1374            .with_field("typed_value", typed_value, true)
1375            .build();
1376        assert!(struct_array.column_by_name("value").is_none());
1377
1378        let variant_array = VariantArray::try_new(&struct_array).unwrap();
1379        assert_eq!(variant_array.value_column().len(), 3);
1380        assert_eq!(variant_array.value_column().null_count(), 3);
1381        assert!(variant_array.inner().column_by_name("value").is_some());
1382        assert!(variant_array.typed_value_column().is_some());
1383        assert_eq!(variant_array.value(0), Variant::from(1i64));
1384        assert_eq!(variant_array.value(1), Variant::Null);
1385        assert_eq!(variant_array.value(2), Variant::from(3i64));
1386
1387        // With no `typed_value` either, there is nothing to read, so `value` stays required.
1388        let fields = Fields::from(vec![Field::new("metadata", DataType::BinaryView, false)]);
1389        let metadata_only = StructArray::new(fields, vec![make_binary_view_array()], None);
1390        let err = VariantArray::try_new(&metadata_only);
1391        assert_eq!(
1392            err.unwrap_err().to_string(),
1393            "Invalid argument error: Invalid VariantArray: StructArray must contain a 'value' field"
1394        );
1395    }
1396
1397    #[test]
1398    fn invalid_metadata_field_type() {
1399        let fields = Fields::from(vec![
1400            Field::new("metadata", DataType::Int32, true), // not supported
1401            Field::new("value", DataType::BinaryView, true),
1402        ]);
1403        let array = StructArray::new(
1404            fields,
1405            vec![make_int32_array(), make_binary_view_array()],
1406            None,
1407        );
1408        let err = VariantArray::try_new(&array);
1409        assert_eq!(
1410            err.unwrap_err().to_string(),
1411            "Invalid argument error: VariantArray 'metadata' field must be Binary, LargeBinary, or BinaryView, got Int32"
1412        );
1413    }
1414
1415    #[test]
1416    fn invalid_value_field_type() {
1417        let fields = Fields::from(vec![
1418            Field::new("metadata", DataType::BinaryView, true),
1419            Field::new("value", DataType::Int32, true),
1420        ]);
1421        let array = StructArray::new(
1422            fields,
1423            vec![make_binary_view_array(), make_int32_array()],
1424            None,
1425        );
1426        let err = VariantArray::try_new(&array);
1427        assert_eq!(
1428            err.unwrap_err().to_string(),
1429            "Invalid argument error: VariantArray 'value' field must be Binary, LargeBinary, or BinaryView, got Int32"
1430        );
1431    }
1432
1433    fn make_binary_view_array() -> ArrayRef {
1434        Arc::new(BinaryViewArray::from(vec![b"test" as &[u8]]))
1435    }
1436
1437    fn make_int32_array() -> ArrayRef {
1438        Arc::new(Int32Array::from(vec![1]))
1439    }
1440
1441    fn make_variant_struct_with_typed_value(typed_value: ArrayRef) -> StructArray {
1442        let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(
1443            EMPTY_VARIANT_METADATA_BYTES,
1444            typed_value.len(),
1445        ));
1446        let value = new_null_array(&DataType::BinaryView, typed_value.len());
1447        StructArrayBuilder::new()
1448            .with_field("metadata", Arc::new(metadata), false)
1449            .with_field("value", value, true)
1450            .with_field("typed_value", typed_value, true)
1451            .build()
1452    }
1453
1454    #[test]
1455    fn try_new_tags_untagged_uuid_on_read() {
1456        // Simulate a foreign writer that shredded a UUID column as bare FixedSizeBinary(16),
1457        // omitting the UUID extension type.
1458        let typed_value = FixedSizeBinaryArray::try_from_iter(std::iter::repeat_n([0u8; 16], 2));
1459        let input = make_variant_struct_with_typed_value(Arc::new(typed_value.unwrap()));
1460
1461        // try_new canonicalizes on the read path and attaches the extension.
1462        let variant_array = VariantArray::try_new(&input).unwrap();
1463        let typed_value = variant_array.inner().field_by_name("typed_value").unwrap();
1464        assert_eq!(typed_value.data_type(), &DataType::FixedSizeBinary(16));
1465        assert!(typed_value.has_valid_extension_type::<UuidExtension>());
1466    }
1467
1468    #[test]
1469    fn try_new_tags_untagged_nested_uuid_on_read() {
1470        // A shredded object { id: { typed_value: FixedSizeBinary(16) } } whose inner UUID leaf
1471        // carries no extension type; canonicalization must reach it recursively.
1472        let leaf = FixedSizeBinaryArray::try_from_iter(std::iter::repeat_n([0u8; 16], 1)).unwrap();
1473        let inner = StructArrayBuilder::new()
1474            .with_field("typed_value", Arc::new(leaf), true)
1475            .build();
1476        let object = StructArrayBuilder::new()
1477            .with_field("id", Arc::new(inner), false)
1478            .build();
1479        let input = make_variant_struct_with_typed_value(Arc::new(object));
1480
1481        // typed_value (struct) -> id (struct) -> typed_value (the FixedSizeBinary(16) UUID leaf).
1482        let variant_array = VariantArray::try_new(&input).unwrap();
1483        let object = variant_array.typed_value_column().unwrap().as_struct();
1484        let id = object.column_by_name("id").unwrap().as_struct();
1485        let uuid_leaf = id.field_by_name("typed_value").unwrap();
1486        assert!(uuid_leaf.has_valid_extension_type::<UuidExtension>());
1487    }
1488
1489    #[test]
1490    fn all_null_value_column_is_valid_and_unshredded() {
1491        // An all-null `value` column is accepted (only a *missing* column is
1492        // rejected) and the array is unshredded (no typed_value column)
1493        let metadata = BinaryViewArray::from(vec![b"test" as &[u8]; 3]);
1494        let value = new_null_array(&DataType::BinaryView, 3);
1495
1496        let fields = Fields::from(vec![
1497            Field::new("metadata", DataType::BinaryView, false),
1498            Field::new("value", DataType::BinaryView, true),
1499        ]);
1500        let struct_array = StructArray::new(fields, vec![Arc::new(metadata), value], None);
1501
1502        let variant_array = VariantArray::try_new(&struct_array).unwrap();
1503        assert!(variant_array.typed_value_column().is_none());
1504
1505        // The rows are valid but have neither value nor typed_value; the spec
1506        // requires readers to return Variant::Null in this case
1507        for i in 0..variant_array.len() {
1508            assert!(variant_array.is_valid(i));
1509            assert_eq!(variant_array.value(i), Variant::Null);
1510        }
1511    }
1512
1513    #[test]
1514    fn canonicalize_and_verify_list_like_data_types() {
1515        // `parquet/tests/variant_integration.rs` validates Parquet shredded-variant fixtures that
1516        // use Parquet LIST encoding, but those fixtures do not cover Arrow-specific list container
1517        // variants (`LargeList`, `ListView`, `LargeListView`) accepted by `VariantArray::try_new`.
1518        let make_item_binary = || Arc::new(Field::new("item", DataType::Binary, true));
1519        let make_large_binary = || Arc::new(Field::new("item", DataType::LargeBinary, true));
1520        let make_item_binary_view = || Arc::new(Field::new("item", DataType::BinaryView, true));
1521
1522        let cases = vec![
1523            // Binary item
1524            DataType::LargeList(make_item_binary()),
1525            DataType::ListView(make_item_binary()),
1526            DataType::LargeListView(make_item_binary()),
1527            // Large binary item
1528            DataType::LargeList(make_large_binary()),
1529            DataType::ListView(make_large_binary()),
1530            DataType::LargeListView(make_large_binary()),
1531            // Binary view item
1532            DataType::LargeList(make_item_binary_view()),
1533            DataType::ListView(make_item_binary_view()),
1534            DataType::LargeListView(make_item_binary_view()),
1535        ];
1536
1537        for input in cases {
1538            assert_eq!(
1539                canonicalize_and_verify_data_type(&input).unwrap().as_ref(),
1540                &input
1541            );
1542        }
1543    }
1544
1545    #[test]
1546    fn variant_array_try_new_supports_list_like_typed_value() {
1547        let item_field = Arc::new(Field::new("item", DataType::Int64, true));
1548        let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]));
1549
1550        let typed_values = vec![
1551            Arc::new(ListArray::new(
1552                item_field.clone(),
1553                OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])),
1554                values.clone(),
1555                None,
1556            )) as ArrayRef,
1557            Arc::new(LargeListArray::new(
1558                item_field.clone(),
1559                OffsetBuffer::new(ScalarBuffer::from(vec![0_i64, 2, 3])),
1560                values.clone(),
1561                None,
1562            )) as ArrayRef,
1563            Arc::new(ListViewArray::new(
1564                item_field.clone(),
1565                ScalarBuffer::from(vec![0, 2]),
1566                ScalarBuffer::from(vec![2, 1]),
1567                values.clone(),
1568                None,
1569            )) as ArrayRef,
1570            Arc::new(LargeListViewArray::new(
1571                item_field,
1572                ScalarBuffer::from(vec![0_i64, 2]),
1573                ScalarBuffer::from(vec![2_i64, 1]),
1574                values,
1575                None,
1576            )) as ArrayRef,
1577        ];
1578
1579        for typed_value in typed_values {
1580            let input = make_variant_struct_with_typed_value(typed_value.clone());
1581            let variant_array = VariantArray::try_new(&input).unwrap();
1582            assert_eq!(
1583                variant_array.typed_value_column().unwrap().data_type(),
1584                typed_value.data_type(),
1585            );
1586        }
1587    }
1588
1589    #[test]
1590    fn test_variant_array_iterable() {
1591        let mut b = VariantArrayBuilder::new(6);
1592
1593        b.append_null();
1594        b.append_variant(Variant::from(1_i8));
1595        b.append_variant(Variant::Null);
1596        b.append_variant(Variant::from(2_i32));
1597        b.append_variant(Variant::from(3_i64));
1598        b.append_null();
1599
1600        let v = b.build();
1601
1602        let variants = v.iter().collect::<Vec<_>>();
1603
1604        assert_eq!(
1605            variants,
1606            vec![
1607                None,
1608                Some(Variant::Int8(1)),
1609                Some(Variant::Null),
1610                Some(Variant::Int32(2)),
1611                Some(Variant::Int64(3)),
1612                None,
1613            ]
1614        );
1615    }
1616
1617    #[test]
1618    fn test_variant_array_iter_double_ended() {
1619        let mut b = VariantArrayBuilder::new(5);
1620
1621        b.append_variant(Variant::from(0_i32));
1622        b.append_null();
1623        b.append_variant(Variant::from(2_i32));
1624        b.append_null();
1625        b.append_variant(Variant::from(4_i32));
1626
1627        let array = b.build();
1628        let mut iter = array.iter();
1629
1630        assert_eq!(iter.next(), Some(Some(Variant::from(0_i32))));
1631        assert_eq!(iter.next(), Some(None));
1632
1633        assert_eq!(iter.next_back(), Some(Some(Variant::from(4_i32))));
1634        assert_eq!(iter.next_back(), Some(None));
1635        assert_eq!(iter.next_back(), Some(Some(Variant::from(2_i32))));
1636
1637        assert_eq!(iter.next_back(), None);
1638        assert_eq!(iter.next(), None);
1639    }
1640
1641    #[test]
1642    fn test_variant_array_iter_reverse() {
1643        let mut b = VariantArrayBuilder::new(5);
1644
1645        b.append_variant(Variant::from("a"));
1646        b.append_null();
1647        b.append_variant(Variant::from("aaa"));
1648        b.append_null();
1649        b.append_variant(Variant::from("aaaaa"));
1650
1651        let array = b.build();
1652
1653        let result: Vec<_> = array.iter().rev().collect();
1654        assert_eq!(
1655            result,
1656            vec![
1657                Some(Variant::from("aaaaa")),
1658                None,
1659                Some(Variant::from("aaa")),
1660                None,
1661                Some(Variant::from("a")),
1662            ]
1663        );
1664    }
1665
1666    #[test]
1667    fn test_variant_array_iter_empty() {
1668        let v = VariantArrayBuilder::new(0).build();
1669        let mut i = v.iter();
1670        assert!(i.next().is_none());
1671        assert!(i.next_back().is_none());
1672    }
1673
1674    #[test]
1675    fn test_from_variant_opts_into_variant_array() {
1676        let v = vec![None, Some(Variant::Null), Some(Variant::BooleanFalse), None];
1677
1678        let variant_array = VariantArray::from_iter(v);
1679
1680        assert_eq!(variant_array.len(), 4);
1681
1682        assert!(variant_array.is_null(0));
1683
1684        assert!(!variant_array.is_null(1));
1685        assert_eq!(variant_array.value(1), Variant::Null);
1686
1687        assert!(!variant_array.is_null(2));
1688        assert_eq!(variant_array.value(2), Variant::BooleanFalse);
1689
1690        assert!(variant_array.is_null(3));
1691    }
1692
1693    #[test]
1694    fn test_from_variants_into_variant_array() {
1695        let v = vec![
1696            Variant::Null,
1697            Variant::BooleanFalse,
1698            Variant::ShortString(ShortString::try_new("norm").unwrap()),
1699        ];
1700
1701        let variant_array = VariantArray::from_iter(v);
1702
1703        assert_eq!(variant_array.len(), 3);
1704
1705        assert!(!variant_array.is_null(0));
1706        assert_eq!(variant_array.value(0), Variant::Null);
1707
1708        assert!(!variant_array.is_null(1));
1709        assert_eq!(variant_array.value(1), Variant::BooleanFalse);
1710
1711        assert!(!variant_array.is_null(2));
1712        assert_eq!(
1713            variant_array.value(2),
1714            Variant::ShortString(ShortString::try_new("norm").unwrap())
1715        );
1716    }
1717
1718    #[test]
1719    fn test_variant_equality() {
1720        let v_iter = [None, Some(Variant::BooleanFalse), Some(Variant::Null), None];
1721        let v = VariantArray::from_iter(v_iter.clone());
1722
1723        {
1724            let v_copy = v.clone();
1725            assert_eq!(v, v_copy);
1726        }
1727
1728        {
1729            let v_iter_reversed = v_iter.iter().cloned().rev();
1730            let v_reversed = VariantArray::from_iter(v_iter_reversed);
1731
1732            assert_ne!(v, v_reversed);
1733        }
1734
1735        {
1736            let v_sliced = v.slice(0, 1);
1737            assert_ne!(v, v_sliced);
1738        }
1739    }
1740
1741    #[test]
1742    fn binary_typed_value_roundtrips() {
1743        // Verify that a shredded variant with Binary typed_value can be read back
1744        let typed_value: ArrayRef = Arc::new(BinaryArray::from(vec![b"hello" as &[u8]]));
1745        let struct_array = make_variant_struct_with_typed_value(typed_value);
1746
1747        let variant_array = VariantArray::try_new(&struct_array).unwrap();
1748        assert_eq!(variant_array.value(0), Variant::from(b"hello" as &[u8]));
1749    }
1750
1751    #[test]
1752    fn large_binary_typed_value_roundtrips() {
1753        // Verify that a shredded variant with LargeBinary typed_value can be read back
1754        let typed_value: ArrayRef = Arc::new(LargeBinaryArray::from(vec![b"world" as &[u8]]));
1755        let struct_array = make_variant_struct_with_typed_value(typed_value);
1756
1757        let variant_array = VariantArray::try_new(&struct_array).unwrap();
1758        assert_eq!(variant_array.value(0), Variant::from(b"world" as &[u8]));
1759    }
1760
1761    macro_rules! invalid_variant_array_test {
1762        ($fn_name: ident, $invalid_typed_value: expr, $error_msg: literal) => {
1763            #[test]
1764            fn $fn_name() {
1765                let invalid_typed_value = $invalid_typed_value;
1766
1767                let struct_array =
1768                    make_variant_struct_with_typed_value(Arc::new(invalid_typed_value));
1769
1770                let array: VariantArray = VariantArray::try_new(&struct_array)
1771                    .expect("should create variant array")
1772                    .into();
1773
1774                let result = array.try_value(0);
1775                assert!(result.is_err());
1776                let error = result.unwrap_err();
1777                assert!(matches!(error, ArrowError::CastError(_)));
1778
1779                let expected: &str = $error_msg;
1780                assert!(
1781                    error.to_string().contains($error_msg),
1782                    "error `{}` did not contain `{}`",
1783                    error,
1784                    expected
1785                )
1786            }
1787        };
1788    }
1789
1790    invalid_variant_array_test!(
1791        test_variant_array_invalide_time,
1792        Time64MicrosecondArray::from(vec![Some(86401000000)]),
1793        "Cast error: Cast failed at index 0 (array type: Time64(µs)): Invalid microsecond from midnight: 86401000000"
1794    );
1795
1796    invalid_variant_array_test!(
1797        test_variant_array_invalid_decimal32,
1798        Decimal32Array::from(vec![Some(1234567890)]),
1799        "Cast error: Cast failed at index 0 (array type: Decimal32(9, 2)): Invalid argument error: 1234567890 is wider than max precision 9"
1800    );
1801
1802    invalid_variant_array_test!(
1803        test_variant_array_invalid_decimal64,
1804        Decimal64Array::from(vec![Some(1234567890123456789)]),
1805        "Cast error: Cast failed at index 0 (array type: Decimal64(18, 6)): Invalid argument error: 1234567890123456789 is wider than max precision 18"
1806    );
1807
1808    invalid_variant_array_test!(
1809        test_variant_array_invalid_decimal128,
1810        Decimal128Array::from(vec![Some(
1811            i128::from_str("123456789012345678901234567890123456789").unwrap()
1812        ),]),
1813        "Cast error: Cast failed at index 0 (array type: Decimal128(38, 10)): Invalid argument error: 123456789012345678901234567890123456789 is wider than max precision 38"
1814    );
1815    #[test]
1816    fn try_value_errors_on_unimplemented_typed_value_type() {
1817        use crate::{json_to_variant, shred_variant};
1818        use arrow::array::StringArray;
1819
1820        let json: ArrayRef = Arc::new(StringArray::from(vec![r#"{"qty": 3}"#]));
1821        let variant = json_to_variant(&json).unwrap();
1822        let shred_type = DataType::Struct(vec![Field::new("qty", DataType::Int64, true)].into());
1823        let shredded = shred_variant(&variant, &shred_type).unwrap();
1824        // Object-shredded typed_value is not yet implemented: reading it must
1825        // error, never silently return Variant::Null
1826        // TODO: https://github.com/apache/arrow-rs/issues/10620
1827        let err = shredded.try_value(0).unwrap_err();
1828        assert!(
1829            err.to_string().starts_with(
1830                "Not yet implemented: VariantArray::try_value cannot materialize typed_value"
1831            ),
1832            "unexpected error: {err}"
1833        );
1834    }
1835}
parquet_variant_compute/variant_array.rs

parquet_variant_compute/
variant_array.rs
