Skip to main content

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