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