parquet_variant_compute/
variant_to_arrow.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
18use crate::shred_variant::{
19    NullValue, VariantToShreddedVariantRowBuilder,
20    make_variant_to_shredded_variant_arrow_row_builder,
21};
22use crate::type_conversion::{
23    PrimitiveFromVariant, TimestampFromVariant, variant_cast_with_options,
24    variant_to_unscaled_decimal,
25};
26use crate::variant_array::ShreddedVariantFieldArray;
27use crate::{VariantArray, VariantValueArrayBuilder};
28use arrow::array::{
29    ArrayRef, ArrowNativeTypeOp, BinaryBuilder, BinaryLikeArrayBuilder, BinaryViewBuilder,
30    BooleanBuilder, FixedSizeBinaryBuilder, FixedSizeListArray, GenericListArray,
31    GenericListViewArray, LargeBinaryBuilder, LargeStringBuilder, MapArray, NullArray,
32    NullBufferBuilder, OffsetSizeTrait, PrimitiveBuilder, StringBuilder, StringLikeArrayBuilder,
33    StringViewBuilder, StructArray,
34};
35use arrow::buffer::{OffsetBuffer, ScalarBuffer};
36use arrow::compute::{CastOptions, DecimalCast, cast_with_options};
37use arrow::datatypes::{self, DataType, DecimalType};
38use arrow::error::{ArrowError, Result};
39use arrow_schema::{FieldRef, Fields, TimeUnit};
40use parquet_variant::{Variant, VariantPath};
41use std::sync::Arc;
42
43/// Builder for converting variant values into strongly typed Arrow arrays.
44///
45/// Useful for variant_get kernels that need to extract specific paths from variant values, possibly
46/// with casting of leaf values to specific types.
47pub(crate) enum VariantToArrowRowBuilder<'a> {
48    Primitive(PrimitiveVariantToArrowRowBuilder<'a>),
49    Array(ArrayVariantToArrowRowBuilder<'a>),
50    Struct(StructVariantToArrowRowBuilder<'a>),
51    Map(MapVariantToArrowRowBuilder<'a>),
52    Encoded(EncodedVariantToArrowRowBuilder<'a>),
53    BinaryVariant(VariantToBinaryVariantArrowRowBuilder),
54
55    // Path extraction wrapper - contains a boxed enum for any of the above
56    WithPath(VariantPathRowBuilder<'a>),
57}
58
59impl<'a> VariantToArrowRowBuilder<'a> {
60    pub fn append_null(&mut self) -> Result<()> {
61        use VariantToArrowRowBuilder::*;
62        match self {
63            Primitive(b) => b.append_null(),
64            Array(b) => b.append_null(),
65            Struct(b) => b.append_null(),
66            Map(b) => b.append_null(),
67            Encoded(b) => b.append_null(),
68            BinaryVariant(b) => b.append_null(),
69            WithPath(path_builder) => path_builder.append_null(),
70        }
71    }
72
73    pub fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
74        use VariantToArrowRowBuilder::*;
75        match self {
76            Primitive(b) => b.append_value(&value),
77            Array(b) => b.append_value(&value),
78            Struct(b) => b.append_value(&value),
79            Map(b) => b.append_value(&value),
80            Encoded(b) => b.append_value(value),
81            BinaryVariant(b) => b.append_value(value),
82            WithPath(path_builder) => path_builder.append_value(value),
83        }
84    }
85
86    pub fn finish(self) -> Result<ArrayRef> {
87        use VariantToArrowRowBuilder::*;
88        match self {
89            Primitive(b) => b.finish(),
90            Array(b) => b.finish(),
91            Struct(b) => b.finish(),
92            Map(b) => b.finish(),
93            Encoded(b) => b.finish(),
94            BinaryVariant(b) => b.finish(),
95            WithPath(path_builder) => path_builder.finish(),
96        }
97    }
98}
99
100fn make_typed_variant_to_arrow_row_builder<'a>(
101    data_type: &'a DataType,
102    cast_options: &'a CastOptions,
103    capacity: usize,
104) -> Result<VariantToArrowRowBuilder<'a>> {
105    use VariantToArrowRowBuilder::*;
106
107    match data_type {
108        DataType::Struct(fields) => {
109            let builder = StructVariantToArrowRowBuilder::try_new(fields, cast_options, capacity)?;
110            Ok(Struct(builder))
111        }
112        data_type @ (DataType::List(_)
113        | DataType::LargeList(_)
114        | DataType::ListView(_)
115        | DataType::LargeListView(_)
116        | DataType::FixedSizeList(..)) => {
117            let builder =
118                ArrayVariantToArrowRowBuilder::try_new(data_type, cast_options, capacity, false)?;
119            Ok(Array(builder))
120        }
121        DataType::Map(entries_field, ordered) => {
122            let builder = MapVariantToArrowRowBuilder::try_new(
123                entries_field,
124                *ordered,
125                cast_options,
126                capacity,
127            )?;
128            Ok(Map(builder))
129        }
130        DataType::Dictionary(_, value_type) => {
131            let builder = EncodedVariantToArrowRowBuilder::try_new(
132                data_type,
133                value_type.as_ref(),
134                cast_options,
135                capacity,
136            )?;
137            Ok(Encoded(builder))
138        }
139        DataType::RunEndEncoded(_, value_field) => {
140            let builder = EncodedVariantToArrowRowBuilder::try_new(
141                data_type,
142                value_field.data_type(),
143                cast_options,
144                capacity,
145            )?;
146            Ok(Encoded(builder))
147        }
148        data_type => {
149            let builder =
150                make_primitive_variant_to_arrow_row_builder(data_type, cast_options, capacity)?;
151            Ok(Primitive(builder))
152        }
153    }
154}
155
156pub(crate) fn make_variant_to_arrow_row_builder<'a>(
157    metadata: &ArrayRef,
158    path: VariantPath<'a>,
159    data_type: Option<&'a DataType>,
160    cast_options: &'a CastOptions,
161    capacity: usize,
162) -> Result<VariantToArrowRowBuilder<'a>> {
163    use VariantToArrowRowBuilder::*;
164
165    let mut builder = match data_type {
166        // If no data type was requested, build an unshredded VariantArray.
167        None => BinaryVariant(VariantToBinaryVariantArrowRowBuilder::new(
168            metadata.clone(),
169            capacity,
170        )),
171        Some(data_type) => {
172            make_typed_variant_to_arrow_row_builder(data_type, cast_options, capacity)?
173        }
174    };
175
176    // Wrap with path extraction if needed
177    if !path.is_empty() {
178        builder = WithPath(VariantPathRowBuilder {
179            builder: Box::new(builder),
180            path,
181        })
182    };
183
184    Ok(builder)
185}
186
187/// Builder for converting primitive variant values to Arrow arrays. It is used by both
188/// `VariantToArrowRowBuilder` (below) and `VariantToShreddedPrimitiveVariantRowBuilder` (in
189/// `shred_variant.rs`).
190pub(crate) enum PrimitiveVariantToArrowRowBuilder<'a> {
191    Null(VariantToNullArrowRowBuilder<'a>),
192    Boolean(VariantToBooleanArrowRowBuilder<'a>),
193    Int8(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int8Type>),
194    Int16(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int16Type>),
195    Int32(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int32Type>),
196    Int64(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int64Type>),
197    UInt8(VariantToPrimitiveArrowRowBuilder<'a, datatypes::UInt8Type>),
198    UInt16(VariantToPrimitiveArrowRowBuilder<'a, datatypes::UInt16Type>),
199    UInt32(VariantToPrimitiveArrowRowBuilder<'a, datatypes::UInt32Type>),
200    UInt64(VariantToPrimitiveArrowRowBuilder<'a, datatypes::UInt64Type>),
201    Float16(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Float16Type>),
202    Float32(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Float32Type>),
203    Float64(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Float64Type>),
204    Decimal32(VariantToDecimalArrowRowBuilder<'a, datatypes::Decimal32Type>),
205    Decimal64(VariantToDecimalArrowRowBuilder<'a, datatypes::Decimal64Type>),
206    Decimal128(VariantToDecimalArrowRowBuilder<'a, datatypes::Decimal128Type>),
207    Decimal256(VariantToDecimalArrowRowBuilder<'a, datatypes::Decimal256Type>),
208    TimestampSecond(VariantToTimestampArrowRowBuilder<'a, datatypes::TimestampSecondType>),
209    TimestampSecondNtz(VariantToTimestampNtzArrowRowBuilder<'a, datatypes::TimestampSecondType>),
210    TimestampMilli(VariantToTimestampArrowRowBuilder<'a, datatypes::TimestampMillisecondType>),
211    TimestampMilliNtz(
212        VariantToTimestampNtzArrowRowBuilder<'a, datatypes::TimestampMillisecondType>,
213    ),
214    TimestampMicro(VariantToTimestampArrowRowBuilder<'a, datatypes::TimestampMicrosecondType>),
215    TimestampMicroNtz(
216        VariantToTimestampNtzArrowRowBuilder<'a, datatypes::TimestampMicrosecondType>,
217    ),
218    TimestampNano(VariantToTimestampArrowRowBuilder<'a, datatypes::TimestampNanosecondType>),
219    TimestampNanoNtz(VariantToTimestampNtzArrowRowBuilder<'a, datatypes::TimestampNanosecondType>),
220    Time32Second(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Time32SecondType>),
221    Time32Milli(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Time32MillisecondType>),
222    Time64Micro(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Time64MicrosecondType>),
223    Time64Nano(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Time64NanosecondType>),
224    Date32(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Date32Type>),
225    Date64(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Date64Type>),
226    Uuid(VariantToUuidArrowRowBuilder<'a>),
227    String(VariantToStringArrowBuilder<'a, StringBuilder>),
228    LargeString(VariantToStringArrowBuilder<'a, LargeStringBuilder>),
229    StringView(VariantToStringArrowBuilder<'a, StringViewBuilder>),
230    Binary(VariantToBinaryArrowRowBuilder<'a, BinaryBuilder>),
231    LargeBinary(VariantToBinaryArrowRowBuilder<'a, LargeBinaryBuilder>),
232    BinaryView(VariantToBinaryArrowRowBuilder<'a, BinaryViewBuilder>),
233}
234
235impl<'a> PrimitiveVariantToArrowRowBuilder<'a> {
236    pub fn append_null(&mut self) -> Result<()> {
237        use PrimitiveVariantToArrowRowBuilder::*;
238        match self {
239            Null(b) => b.append_null(),
240            Boolean(b) => b.append_null(),
241            Int8(b) => b.append_null(),
242            Int16(b) => b.append_null(),
243            Int32(b) => b.append_null(),
244            Int64(b) => b.append_null(),
245            UInt8(b) => b.append_null(),
246            UInt16(b) => b.append_null(),
247            UInt32(b) => b.append_null(),
248            UInt64(b) => b.append_null(),
249            Float16(b) => b.append_null(),
250            Float32(b) => b.append_null(),
251            Float64(b) => b.append_null(),
252            Decimal32(b) => b.append_null(),
253            Decimal64(b) => b.append_null(),
254            Decimal128(b) => b.append_null(),
255            Decimal256(b) => b.append_null(),
256            TimestampSecond(b) => b.append_null(),
257            TimestampSecondNtz(b) => b.append_null(),
258            TimestampMilli(b) => b.append_null(),
259            TimestampMilliNtz(b) => b.append_null(),
260            TimestampMicro(b) => b.append_null(),
261            TimestampMicroNtz(b) => b.append_null(),
262            TimestampNano(b) => b.append_null(),
263            TimestampNanoNtz(b) => b.append_null(),
264            Time32Second(b) => b.append_null(),
265            Time32Milli(b) => b.append_null(),
266            Time64Micro(b) => b.append_null(),
267            Time64Nano(b) => b.append_null(),
268            Date32(b) => b.append_null(),
269            Date64(b) => b.append_null(),
270            Uuid(b) => b.append_null(),
271            String(b) => b.append_null(),
272            LargeString(b) => b.append_null(),
273            StringView(b) => b.append_null(),
274            Binary(b) => b.append_null(),
275            LargeBinary(b) => b.append_null(),
276            BinaryView(b) => b.append_null(),
277        }
278    }
279
280    pub fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
281        use PrimitiveVariantToArrowRowBuilder::*;
282        match self {
283            Null(b) => b.append_value(value),
284            Boolean(b) => b.append_value(value),
285            Int8(b) => b.append_value(value),
286            Int16(b) => b.append_value(value),
287            Int32(b) => b.append_value(value),
288            Int64(b) => b.append_value(value),
289            UInt8(b) => b.append_value(value),
290            UInt16(b) => b.append_value(value),
291            UInt32(b) => b.append_value(value),
292            UInt64(b) => b.append_value(value),
293            Float16(b) => b.append_value(value),
294            Float32(b) => b.append_value(value),
295            Float64(b) => b.append_value(value),
296            Decimal32(b) => b.append_value(value),
297            Decimal64(b) => b.append_value(value),
298            Decimal128(b) => b.append_value(value),
299            Decimal256(b) => b.append_value(value),
300            TimestampSecond(b) => b.append_value(value),
301            TimestampSecondNtz(b) => b.append_value(value),
302            TimestampMilli(b) => b.append_value(value),
303            TimestampMilliNtz(b) => b.append_value(value),
304            TimestampMicro(b) => b.append_value(value),
305            TimestampMicroNtz(b) => b.append_value(value),
306            TimestampNano(b) => b.append_value(value),
307            TimestampNanoNtz(b) => b.append_value(value),
308            Time32Second(b) => b.append_value(value),
309            Time32Milli(b) => b.append_value(value),
310            Time64Micro(b) => b.append_value(value),
311            Time64Nano(b) => b.append_value(value),
312            Date32(b) => b.append_value(value),
313            Date64(b) => b.append_value(value),
314            Uuid(b) => b.append_value(value),
315            String(b) => b.append_value(value),
316            LargeString(b) => b.append_value(value),
317            StringView(b) => b.append_value(value),
318            Binary(b) => b.append_value(value),
319            LargeBinary(b) => b.append_value(value),
320            BinaryView(b) => b.append_value(value),
321        }
322    }
323
324    pub fn finish(self) -> Result<ArrayRef> {
325        use PrimitiveVariantToArrowRowBuilder::*;
326        match self {
327            Null(b) => b.finish(),
328            Boolean(b) => b.finish(),
329            Int8(b) => b.finish(),
330            Int16(b) => b.finish(),
331            Int32(b) => b.finish(),
332            Int64(b) => b.finish(),
333            UInt8(b) => b.finish(),
334            UInt16(b) => b.finish(),
335            UInt32(b) => b.finish(),
336            UInt64(b) => b.finish(),
337            Float16(b) => b.finish(),
338            Float32(b) => b.finish(),
339            Float64(b) => b.finish(),
340            Decimal32(b) => b.finish(),
341            Decimal64(b) => b.finish(),
342            Decimal128(b) => b.finish(),
343            Decimal256(b) => b.finish(),
344            TimestampSecond(b) => b.finish(),
345            TimestampSecondNtz(b) => b.finish(),
346            TimestampMilli(b) => b.finish(),
347            TimestampMilliNtz(b) => b.finish(),
348            TimestampMicro(b) => b.finish(),
349            TimestampMicroNtz(b) => b.finish(),
350            TimestampNano(b) => b.finish(),
351            TimestampNanoNtz(b) => b.finish(),
352            Time32Second(b) => b.finish(),
353            Time32Milli(b) => b.finish(),
354            Time64Micro(b) => b.finish(),
355            Time64Nano(b) => b.finish(),
356            Date32(b) => b.finish(),
357            Date64(b) => b.finish(),
358            Uuid(b) => b.finish(),
359            String(b) => b.finish(),
360            LargeString(b) => b.finish(),
361            StringView(b) => b.finish(),
362            Binary(b) => b.finish(),
363            LargeBinary(b) => b.finish(),
364            BinaryView(b) => b.finish(),
365        }
366    }
367}
368
369pub(crate) struct EncodedVariantToArrowRowBuilder<'a> {
370    data_type: &'a DataType,
371    cast_options: &'a CastOptions<'a>,
372    values_builder: Box<VariantToArrowRowBuilder<'a>>,
373}
374
375impl<'a> EncodedVariantToArrowRowBuilder<'a> {
376    fn try_new(
377        data_type: &'a DataType,
378        value_type: &'a DataType,
379        cast_options: &'a CastOptions,
380        capacity: usize,
381    ) -> Result<Self> {
382        let values_builder = Box::new(make_typed_variant_to_arrow_row_builder(
383            value_type,
384            cast_options,
385            capacity,
386        )?);
387        Ok(Self {
388            data_type,
389            cast_options,
390            values_builder,
391        })
392    }
393
394    fn append_null(&mut self) -> Result<()> {
395        self.values_builder.append_null()
396    }
397
398    fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
399        self.values_builder.append_value(value)
400    }
401
402    fn finish(self) -> Result<ArrayRef> {
403        let values = self.values_builder.finish()?;
404        cast_with_options(values.as_ref(), self.data_type, self.cast_options)
405    }
406}
407
408/// Creates a row builder that converts primitive `Variant` values into the requested Arrow data type.
409pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>(
410    data_type: &'a DataType,
411    cast_options: &'a CastOptions,
412    capacity: usize,
413) -> Result<PrimitiveVariantToArrowRowBuilder<'a>> {
414    use PrimitiveVariantToArrowRowBuilder::*;
415
416    let builder =
417        match data_type {
418            DataType::Null => Null(VariantToNullArrowRowBuilder::new(cast_options, capacity)),
419            DataType::Boolean => {
420                Boolean(VariantToBooleanArrowRowBuilder::new(cast_options, capacity))
421            }
422            DataType::Int8 => Int8(VariantToPrimitiveArrowRowBuilder::new(
423                cast_options,
424                capacity,
425            )),
426            DataType::Int16 => Int16(VariantToPrimitiveArrowRowBuilder::new(
427                cast_options,
428                capacity,
429            )),
430            DataType::Int32 => Int32(VariantToPrimitiveArrowRowBuilder::new(
431                cast_options,
432                capacity,
433            )),
434            DataType::Int64 => Int64(VariantToPrimitiveArrowRowBuilder::new(
435                cast_options,
436                capacity,
437            )),
438            DataType::UInt8 => UInt8(VariantToPrimitiveArrowRowBuilder::new(
439                cast_options,
440                capacity,
441            )),
442            DataType::UInt16 => UInt16(VariantToPrimitiveArrowRowBuilder::new(
443                cast_options,
444                capacity,
445            )),
446            DataType::UInt32 => UInt32(VariantToPrimitiveArrowRowBuilder::new(
447                cast_options,
448                capacity,
449            )),
450            DataType::UInt64 => UInt64(VariantToPrimitiveArrowRowBuilder::new(
451                cast_options,
452                capacity,
453            )),
454            DataType::Float16 => Float16(VariantToPrimitiveArrowRowBuilder::new(
455                cast_options,
456                capacity,
457            )),
458            DataType::Float32 => Float32(VariantToPrimitiveArrowRowBuilder::new(
459                cast_options,
460                capacity,
461            )),
462            DataType::Float64 => Float64(VariantToPrimitiveArrowRowBuilder::new(
463                cast_options,
464                capacity,
465            )),
466            DataType::Decimal32(precision, scale) => Decimal32(
467                VariantToDecimalArrowRowBuilder::new(cast_options, capacity, *precision, *scale)?,
468            ),
469            DataType::Decimal64(precision, scale) => Decimal64(
470                VariantToDecimalArrowRowBuilder::new(cast_options, capacity, *precision, *scale)?,
471            ),
472            DataType::Decimal128(precision, scale) => Decimal128(
473                VariantToDecimalArrowRowBuilder::new(cast_options, capacity, *precision, *scale)?,
474            ),
475            DataType::Decimal256(precision, scale) => Decimal256(
476                VariantToDecimalArrowRowBuilder::new(cast_options, capacity, *precision, *scale)?,
477            ),
478            DataType::Date32 => Date32(VariantToPrimitiveArrowRowBuilder::new(
479                cast_options,
480                capacity,
481            )),
482            DataType::Date64 => Date64(VariantToPrimitiveArrowRowBuilder::new(
483                cast_options,
484                capacity,
485            )),
486            DataType::Time32(TimeUnit::Second) => Time32Second(
487                VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity),
488            ),
489            DataType::Time32(TimeUnit::Millisecond) => Time32Milli(
490                VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity),
491            ),
492            DataType::Time32(t) => {
493                return Err(ArrowError::InvalidArgumentError(format!(
494                    "The unit for Time32 must be second/millisecond, received {t:?}"
495                )));
496            }
497            DataType::Time64(TimeUnit::Microsecond) => Time64Micro(
498                VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity),
499            ),
500            DataType::Time64(TimeUnit::Nanosecond) => Time64Nano(
501                VariantToPrimitiveArrowRowBuilder::new(cast_options, capacity),
502            ),
503            DataType::Time64(t) => {
504                return Err(ArrowError::InvalidArgumentError(format!(
505                    "The unit for Time64 must be micro/nano seconds, received {t:?}"
506                )));
507            }
508            DataType::Timestamp(TimeUnit::Second, None) => TimestampSecondNtz(
509                VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity),
510            ),
511            DataType::Timestamp(TimeUnit::Second, tz) => TimestampSecond(
512                VariantToTimestampArrowRowBuilder::new(cast_options, capacity, tz.clone()),
513            ),
514            DataType::Timestamp(TimeUnit::Millisecond, None) => TimestampMilliNtz(
515                VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity),
516            ),
517            DataType::Timestamp(TimeUnit::Millisecond, tz) => TimestampMilli(
518                VariantToTimestampArrowRowBuilder::new(cast_options, capacity, tz.clone()),
519            ),
520            DataType::Timestamp(TimeUnit::Microsecond, None) => TimestampMicroNtz(
521                VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity),
522            ),
523            DataType::Timestamp(TimeUnit::Microsecond, tz) => TimestampMicro(
524                VariantToTimestampArrowRowBuilder::new(cast_options, capacity, tz.clone()),
525            ),
526            DataType::Timestamp(TimeUnit::Nanosecond, None) => TimestampNanoNtz(
527                VariantToTimestampNtzArrowRowBuilder::new(cast_options, capacity),
528            ),
529            DataType::Timestamp(TimeUnit::Nanosecond, tz) => TimestampNano(
530                VariantToTimestampArrowRowBuilder::new(cast_options, capacity, tz.clone()),
531            ),
532            DataType::Duration(_) | DataType::Interval(_) => {
533                return Err(ArrowError::InvalidArgumentError(
534                    "Casting Variant to duration/interval types is not supported. \
535                    The Variant format does not define duration/interval types."
536                        .to_string(),
537                ));
538            }
539            DataType::Binary => Binary(VariantToBinaryArrowRowBuilder::new(cast_options, capacity)),
540            DataType::LargeBinary => {
541                LargeBinary(VariantToBinaryArrowRowBuilder::new(cast_options, capacity))
542            }
543            DataType::BinaryView => {
544                BinaryView(VariantToBinaryArrowRowBuilder::new(cast_options, capacity))
545            }
546            DataType::FixedSizeBinary(16) => {
547                Uuid(VariantToUuidArrowRowBuilder::new(cast_options, capacity))
548            }
549            DataType::FixedSizeBinary(_) => {
550                return Err(ArrowError::NotYetImplemented(format!(
551                    "DataType {data_type:?} not yet implemented"
552                )));
553            }
554            DataType::Utf8 => String(VariantToStringArrowBuilder::new(cast_options, capacity)),
555            DataType::LargeUtf8 => {
556                LargeString(VariantToStringArrowBuilder::new(cast_options, capacity))
557            }
558            DataType::Utf8View => {
559                StringView(VariantToStringArrowBuilder::new(cast_options, capacity))
560            }
561            DataType::List(_)
562            | DataType::LargeList(_)
563            | DataType::ListView(_)
564            | DataType::LargeListView(_)
565            | DataType::FixedSizeList(..)
566            | DataType::Struct(_)
567            | DataType::Map(..)
568            | DataType::Union(..)
569            | DataType::Dictionary(..)
570            | DataType::RunEndEncoded(..) => {
571                return Err(ArrowError::InvalidArgumentError(format!(
572                    "Casting to {data_type:?} is not applicable for primitive Variant types"
573                )));
574            }
575        };
576    Ok(builder)
577}
578
579pub(crate) enum ArrayVariantToArrowRowBuilder<'a> {
580    List(VariantToListArrowRowBuilder<'a, i32, false>),
581    LargeList(VariantToListArrowRowBuilder<'a, i64, false>),
582    ListView(VariantToListArrowRowBuilder<'a, i32, true>),
583    LargeListView(VariantToListArrowRowBuilder<'a, i64, true>),
584    FixedSizeList(VariantToFixedSizeListArrowRowBuilder<'a>),
585}
586
587pub(crate) struct StructVariantToArrowRowBuilder<'a> {
588    fields: &'a Fields,
589    field_builders: Vec<VariantToArrowRowBuilder<'a>>,
590    nulls: NullBufferBuilder,
591    cast_options: &'a CastOptions<'a>,
592}
593
594impl<'a> StructVariantToArrowRowBuilder<'a> {
595    fn try_new(
596        fields: &'a Fields,
597        cast_options: &'a CastOptions<'a>,
598        capacity: usize,
599    ) -> Result<Self> {
600        let mut field_builders = Vec::with_capacity(fields.len());
601        for field in fields.iter() {
602            field_builders.push(make_typed_variant_to_arrow_row_builder(
603                field.data_type(),
604                cast_options,
605                capacity,
606            )?);
607        }
608        Ok(Self {
609            fields,
610            field_builders,
611            nulls: NullBufferBuilder::new(capacity),
612            cast_options,
613        })
614    }
615
616    fn append_null(&mut self) -> Result<()> {
617        for builder in &mut self.field_builders {
618            builder.append_null()?;
619        }
620        self.nulls.append_null();
621        Ok(())
622    }
623
624    fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
625        match variant_cast_with_options(value, self.cast_options, Variant::as_object) {
626            Ok(Some(obj)) => {
627                for (index, field) in self.fields.iter().enumerate() {
628                    match obj.get(field.name()) {
629                        Some(field_value) => {
630                            self.field_builders[index].append_value(field_value)?;
631                        }
632                        None => {
633                            self.field_builders[index].append_null()?;
634                        }
635                    }
636                }
637
638                self.nulls.append_non_null();
639                Ok(true)
640            }
641            Ok(None) => {
642                self.append_null()?;
643                Ok(false)
644            }
645            Err(_) => Err(ArrowError::CastError(format!(
646                "Failed to extract struct from variant {value:?}"
647            ))),
648        }
649    }
650
651    fn finish(mut self) -> Result<ArrayRef> {
652        let mut children = Vec::with_capacity(self.field_builders.len());
653        for builder in self.field_builders {
654            children.push(builder.finish()?);
655        }
656        Ok(Arc::new(StructArray::try_new(
657            self.fields.clone(),
658            children,
659            self.nulls.finish(),
660        )?))
661    }
662}
663
664/// Builder for converting variant objects into Arrow `MapArray`s.
665///
666/// Each variant object field becomes one map entry: the field name is the map key and the field
667/// value is converted to the requested value type. Variant objects store their fields in
668/// lexicographic key order, so maps with string keys come out sorted by key.
669pub(crate) struct MapVariantToArrowRowBuilder<'a> {
670    entries_field: &'a FieldRef,
671    key_field: &'a FieldRef,
672    value_field: &'a FieldRef,
673    ordered: bool,
674    key_builder: Box<VariantToArrowRowBuilder<'a>>,
675    value_builder: Box<VariantToArrowRowBuilder<'a>>,
676    offsets: Vec<i32>,
677    current_offset: i32,
678    nulls: NullBufferBuilder,
679    cast_options: &'a CastOptions<'a>,
680}
681
682impl<'a> MapVariantToArrowRowBuilder<'a> {
683    fn try_new(
684        entries_field: &'a FieldRef,
685        ordered: bool,
686        cast_options: &'a CastOptions,
687        capacity: usize,
688    ) -> Result<Self> {
689        let DataType::Struct(entry_fields) = entries_field.data_type() else {
690            return Err(ArrowError::InvalidArgumentError(format!(
691                "Map entries must be Struct, got {:?}",
692                entries_field.data_type()
693            )));
694        };
695        let (key_field, value_field) = match entry_fields.as_ref() {
696            [key_field, value_field] => (key_field, value_field),
697            fields => {
698                return Err(ArrowError::InvalidArgumentError(format!(
699                    "Map entries must have exactly two fields (key, value), got {}",
700                    fields.len()
701                )));
702            }
703        };
704        let key_builder = Box::new(make_typed_variant_to_arrow_row_builder(
705            key_field.data_type(),
706            cast_options,
707            capacity,
708        )?);
709        let value_builder = Box::new(make_typed_variant_to_arrow_row_builder(
710            value_field.data_type(),
711            cast_options,
712            capacity,
713        )?);
714        if capacity >= isize::MAX as usize {
715            return Err(ArrowError::ComputeError(
716                "Capacity exceeds isize::MAX when reserving map offsets".to_string(),
717            ));
718        }
719        let mut offsets = Vec::with_capacity(capacity + 1);
720        offsets.push(0);
721        Ok(Self {
722            entries_field,
723            key_field,
724            value_field,
725            ordered,
726            key_builder,
727            value_builder,
728            offsets,
729            current_offset: 0,
730            nulls: NullBufferBuilder::new(capacity),
731            cast_options,
732        })
733    }
734
735    fn append_null(&mut self) -> Result<()> {
736        self.offsets.push(self.current_offset);
737        self.nulls.append_null();
738        Ok(())
739    }
740
741    fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
742        match variant_cast_with_options(value, self.cast_options, Variant::as_object) {
743            Ok(Some(obj)) => {
744                for (key, field_value) in obj.iter() {
745                    // Map keys cannot be null, so a failed key conversion is an error even when
746                    // cast_options.safe is true.
747                    if !self.key_builder.append_value(Variant::from(key))? {
748                        return Err(ArrowError::CastError(format!(
749                            "Failed to cast map key {key:?} to {:?}",
750                            self.key_field.data_type()
751                        )));
752                    }
753                    self.value_builder.append_value(field_value)?;
754                    self.current_offset = self.current_offset.add_checked(1)?;
755                }
756                self.offsets.push(self.current_offset);
757                self.nulls.append_non_null();
758                Ok(true)
759            }
760            Ok(None) => {
761                self.append_null()?;
762                Ok(false)
763            }
764            Err(_) => Err(ArrowError::CastError(format!(
765                "Failed to extract object from variant {value:?}"
766            ))),
767        }
768    }
769
770    fn finish(mut self) -> Result<ArrayRef> {
771        let keys = self.key_builder.finish()?;
772        let values = self.value_builder.finish()?;
773
774        let entry_fields = Fields::from(vec![
775            self.key_field
776                .as_ref()
777                .clone()
778                .with_data_type(keys.data_type().clone()),
779            self.value_field
780                .as_ref()
781                .clone()
782                .with_data_type(values.data_type().clone()),
783        ]);
784        let entries = StructArray::try_new(entry_fields.clone(), vec![keys, values], None)?;
785        let entries_field = Arc::new(
786            self.entries_field
787                .as_ref()
788                .clone()
789                .with_data_type(DataType::Struct(entry_fields)),
790        );
791        let map_array = MapArray::try_new(
792            entries_field,
793            OffsetBuffer::new(ScalarBuffer::from(self.offsets)),
794            entries,
795            self.nulls.finish(),
796            self.ordered,
797        )?;
798        Ok(Arc::new(map_array))
799    }
800}
801
802impl<'a> ArrayVariantToArrowRowBuilder<'a> {
803    /// Creates a new list builder for the given data type.
804    ///
805    /// # Arguments
806    /// * `shredded` - If true, element builders produce shredded structs with `value`/`typed_value`
807    ///   fields (for [`crate::shred_variant()`]). If false, element builders produce strongly typed
808    ///   arrays directly (for [`crate::variant_get()`]).
809    pub(crate) fn try_new(
810        data_type: &'a DataType,
811        cast_options: &'a CastOptions,
812        capacity: usize,
813        shredded: bool,
814    ) -> Result<Self> {
815        use ArrayVariantToArrowRowBuilder::*;
816
817        // Make List/ListView builders without repeating the constructor boilerplate.
818        macro_rules! make_list_builder {
819            ($variant:ident, $offset:ty, $is_view:expr, $field:ident) => {
820                $variant(VariantToListArrowRowBuilder::<$offset, $is_view>::try_new(
821                    $field.clone(),
822                    $field.data_type(),
823                    cast_options,
824                    capacity,
825                    shredded,
826                )?)
827            };
828        }
829
830        let builder = match data_type {
831            DataType::List(field) => make_list_builder!(List, i32, false, field),
832            DataType::LargeList(field) => make_list_builder!(LargeList, i64, false, field),
833            DataType::ListView(field) => make_list_builder!(ListView, i32, true, field),
834            DataType::LargeListView(field) => make_list_builder!(LargeListView, i64, true, field),
835            DataType::FixedSizeList(field, size) => {
836                FixedSizeList(VariantToFixedSizeListArrowRowBuilder::try_new(
837                    field.clone(),
838                    field.data_type(),
839                    *size,
840                    cast_options,
841                    capacity,
842                    shredded,
843                )?)
844            }
845            other => {
846                return Err(ArrowError::InvalidArgumentError(format!(
847                    "Casting to {other:?} is not applicable for array Variant types"
848                )));
849            }
850        };
851        Ok(builder)
852    }
853
854    pub(crate) fn append_null(&mut self) -> Result<()> {
855        match self {
856            Self::List(builder) => builder.append_null(),
857            Self::LargeList(builder) => builder.append_null(),
858            Self::ListView(builder) => builder.append_null(),
859            Self::LargeListView(builder) => builder.append_null(),
860            Self::FixedSizeList(builder) => builder.append_null(),
861        }
862    }
863
864    pub(crate) fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
865        match self {
866            Self::List(builder) => builder.append_value(value),
867            Self::LargeList(builder) => builder.append_value(value),
868            Self::ListView(builder) => builder.append_value(value),
869            Self::LargeListView(builder) => builder.append_value(value),
870            Self::FixedSizeList(builder) => builder.append_value(value),
871        }
872    }
873
874    pub(crate) fn finish(self) -> Result<ArrayRef> {
875        match self {
876            Self::List(builder) => builder.finish(),
877            Self::LargeList(builder) => builder.finish(),
878            Self::ListView(builder) => builder.finish(),
879            Self::LargeListView(builder) => builder.finish(),
880            Self::FixedSizeList(builder) => builder.finish(),
881        }
882    }
883}
884
885/// A thin wrapper whose only job is to extract a specific path from a variant value and pass the
886/// result to a nested builder.
887pub(crate) struct VariantPathRowBuilder<'a> {
888    builder: Box<VariantToArrowRowBuilder<'a>>,
889    path: VariantPath<'a>,
890}
891
892impl<'a> VariantPathRowBuilder<'a> {
893    fn append_null(&mut self) -> Result<()> {
894        self.builder.append_null()
895    }
896
897    fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
898        if let Some(v) = value.get_path(&self.path) {
899            self.builder.append_value(v)
900        } else {
901            self.builder.append_null()?;
902            Ok(false)
903        }
904    }
905
906    fn finish(self) -> Result<ArrayRef> {
907        self.builder.finish()
908    }
909}
910
911macro_rules! define_variant_to_primitive_builder {
912    (struct $name:ident<$lifetime:lifetime $(, $generic:ident: $bound:path )?>
913    |$array_param:ident $(, $field:ident: $field_type:ty)?| -> $builder_name:ident $(< $array_type:ty >)? { $init_expr: expr },
914    |$value: ident| $value_transform:expr,
915    type_name: $type_name:expr) => {
916        pub(crate) struct $name<$lifetime $(, $generic : $bound )?>
917        {
918            builder: $builder_name $(<$array_type>)?,
919            cast_options: &$lifetime CastOptions<$lifetime>,
920        }
921
922        impl<$lifetime $(, $generic: $bound+ )?> $name<$lifetime $(, $generic )?> {
923            fn new(
924                cast_options: &$lifetime CastOptions<$lifetime>,
925                $array_param: usize,
926                // add this so that $init_expr can use it
927                $( $field: $field_type, )?
928            ) -> Self {
929                Self {
930                    builder: $init_expr,
931                    cast_options,
932                }
933            }
934
935            fn append_null(&mut self) -> Result<()> {
936                self.builder.append_null();
937                Ok(())
938            }
939
940            fn append_value(&mut self, $value: &Variant<'_, '_>) -> Result<bool> {
941                match variant_cast_with_options(
942                    $value,
943                    self.cast_options,
944                    |$value| $value_transform,
945                ) {
946                    Ok(Some(v)) => {
947                        self.builder.append_value(v);
948                        Ok(true)
949                    }
950                    Ok(None) => {
951                        self.builder.append_null();
952                        Ok(false)
953                    }
954                    Err(_) => Err(ArrowError::CastError(format!(
955                        "Failed to extract primitive of type {type_name} from variant {value:?} at path VariantPath([])",
956                        type_name = $type_name,
957                        value = $value
958                    ))),
959                }
960            }
961
962            // Add this to silence unused mut warning from macro-generated code
963            // This is mainly for `FakeNullBuilder`
964            #[allow(unused_mut)]
965            fn finish(mut self) -> Result<ArrayRef> {
966                // If the builder produces T: Array, the compiler infers `<Arc<T> as From<T>>::from`
967                // (which then coerces to ArrayRef). If the builder produces ArrayRef directly, the
968                // compiler infers `<ArrayRef as From<ArrayRef>>::from` (no-op, From blanket impl).
969                Ok(Arc::from(self.builder.finish()))
970            }
971        }
972    }
973}
974
975define_variant_to_primitive_builder!(
976    struct VariantToStringArrowBuilder<'a, B: StringLikeArrayBuilder>
977    |capacity| -> B { B::with_capacity(capacity) },
978    |value| value.as_string(),
979    type_name: B::type_name()
980);
981
982define_variant_to_primitive_builder!(
983    struct VariantToBooleanArrowRowBuilder<'a>
984    |capacity| -> BooleanBuilder { BooleanBuilder::with_capacity(capacity) },
985    |value| value.as_boolean(),
986    type_name: datatypes::BooleanType::DATA_TYPE
987);
988
989define_variant_to_primitive_builder!(
990    struct VariantToPrimitiveArrowRowBuilder<'a, T:PrimitiveFromVariant>
991    |capacity| -> PrimitiveBuilder<T> { PrimitiveBuilder::<T>::with_capacity(capacity) },
992    |value| T::from_variant(value),
993    type_name: T::DATA_TYPE
994);
995
996define_variant_to_primitive_builder!(
997    struct VariantToTimestampNtzArrowRowBuilder<'a, T:TimestampFromVariant<true>>
998    |capacity| -> PrimitiveBuilder<T> { PrimitiveBuilder::<T>::with_capacity(capacity) },
999    |value| T::from_variant(value),
1000    type_name: T::DATA_TYPE
1001);
1002
1003define_variant_to_primitive_builder!(
1004    struct VariantToTimestampArrowRowBuilder<'a, T:TimestampFromVariant<false>>
1005    |capacity, tz: Option<Arc<str>> | -> PrimitiveBuilder<T> {
1006        PrimitiveBuilder::<T>::with_capacity(capacity).with_timezone_opt(tz)
1007    },
1008    |value| T::from_variant(value),
1009    type_name: T::DATA_TYPE
1010);
1011
1012define_variant_to_primitive_builder!(
1013    struct VariantToBinaryArrowRowBuilder<'a, B: BinaryLikeArrayBuilder>
1014    |capacity| -> B { B::with_capacity(capacity) },
1015    |value| value.as_u8_slice(),
1016    type_name: B::type_name()
1017);
1018
1019/// Builder for converting variant values to arrow Decimal values
1020pub(crate) struct VariantToDecimalArrowRowBuilder<'a, T>
1021where
1022    T: DecimalType,
1023    T::Native: DecimalCast,
1024{
1025    builder: PrimitiveBuilder<T>,
1026    cast_options: &'a CastOptions<'a>,
1027    precision: u8,
1028    scale: i8,
1029}
1030
1031impl<'a, T> VariantToDecimalArrowRowBuilder<'a, T>
1032where
1033    T: DecimalType,
1034    T::Native: DecimalCast,
1035{
1036    fn new(
1037        cast_options: &'a CastOptions<'a>,
1038        capacity: usize,
1039        precision: u8,
1040        scale: i8,
1041    ) -> Result<Self> {
1042        let builder = PrimitiveBuilder::<T>::with_capacity(capacity)
1043            .with_precision_and_scale(precision, scale)?;
1044        Ok(Self {
1045            builder,
1046            cast_options,
1047            precision,
1048            scale,
1049        })
1050    }
1051
1052    fn append_null(&mut self) -> Result<()> {
1053        self.builder.append_null();
1054        Ok(())
1055    }
1056
1057    fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
1058        match variant_cast_with_options(value, self.cast_options, |value| {
1059            variant_to_unscaled_decimal::<T>(value, self.precision, self.scale)
1060        }) {
1061            Ok(Some(scaled)) => {
1062                self.builder.append_value(scaled);
1063                Ok(true)
1064            }
1065            Ok(None) => {
1066                self.builder.append_null();
1067                Ok(false)
1068            }
1069            Err(_) => Err(ArrowError::CastError(format!(
1070                "Failed to cast to {prefix}(precision={precision}, scale={scale}) from variant {value:?}",
1071                prefix = T::PREFIX,
1072                precision = self.precision,
1073                scale = self.scale
1074            ))),
1075        }
1076    }
1077
1078    fn finish(mut self) -> Result<ArrayRef> {
1079        Ok(Arc::new(self.builder.finish()))
1080    }
1081}
1082
1083/// Builder for converting variant values to FixedSizeBinary(16) for UUIDs
1084pub(crate) struct VariantToUuidArrowRowBuilder<'a> {
1085    builder: FixedSizeBinaryBuilder,
1086    cast_options: &'a CastOptions<'a>,
1087}
1088
1089impl<'a> VariantToUuidArrowRowBuilder<'a> {
1090    fn new(cast_options: &'a CastOptions<'a>, capacity: usize) -> Self {
1091        Self {
1092            builder: FixedSizeBinaryBuilder::with_capacity(capacity, 16),
1093            cast_options,
1094        }
1095    }
1096
1097    fn append_null(&mut self) -> Result<()> {
1098        self.builder.append_null();
1099        Ok(())
1100    }
1101
1102    fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
1103        match variant_cast_with_options(value, self.cast_options, Variant::as_uuid) {
1104            Ok(Some(uuid)) => {
1105                self.builder
1106                    .append_value(uuid.as_bytes())
1107                    .map_err(|e| ArrowError::ExternalError(Box::new(e)))?;
1108                Ok(true)
1109            }
1110            Ok(None) => {
1111                self.builder.append_null();
1112                Ok(false)
1113            }
1114            Err(_) => Err(ArrowError::CastError(format!(
1115                "Failed to extract UUID from variant {value:?}"
1116            ))),
1117        }
1118    }
1119
1120    fn finish(mut self) -> Result<ArrayRef> {
1121        Ok(Arc::new(self.builder.finish()))
1122    }
1123}
1124
1125/// Element builder for list variants, supporting both typed (for [`crate::variant_get()`])
1126/// and shredded (for [`crate::shred_variant()`]) output modes.
1127enum ListElementBuilder<'a> {
1128    /// Produces the target array type directly.
1129    Typed(Box<VariantToArrowRowBuilder<'a>>),
1130    /// Produces a shredded struct with `value` and `typed_value` fields.
1131    Shredded(Box<VariantToShreddedVariantRowBuilder<'a>>),
1132}
1133
1134impl<'a> ListElementBuilder<'a> {
1135    fn append_null(&mut self) -> Result<()> {
1136        match self {
1137            Self::Typed(b) => b.append_null(),
1138            Self::Shredded(b) => b.append_null(),
1139        }
1140    }
1141
1142    fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
1143        match self {
1144            Self::Typed(b) => b.append_value(value),
1145            Self::Shredded(b) => b.append_value(value),
1146        }
1147    }
1148
1149    fn finish(self) -> Result<ArrayRef> {
1150        match self {
1151            Self::Typed(b) => b.finish(),
1152            Self::Shredded(b) => {
1153                let (value, typed_value, nulls) = b.finish()?;
1154                Ok(ArrayRef::from(ShreddedVariantFieldArray::from_parts(
1155                    Arc::new(value),
1156                    Some(typed_value),
1157                    nulls,
1158                )))
1159            }
1160        }
1161    }
1162}
1163
1164pub(crate) struct VariantToListArrowRowBuilder<'a, O, const IS_VIEW: bool>
1165where
1166    O: OffsetSizeTrait + ArrowNativeTypeOp,
1167{
1168    field: FieldRef,
1169    offsets: Vec<O>,
1170    element_builder: ListElementBuilder<'a>,
1171    nulls: NullBufferBuilder,
1172    current_offset: O,
1173    cast_options: &'a CastOptions<'a>,
1174}
1175
1176impl<'a, O, const IS_VIEW: bool> VariantToListArrowRowBuilder<'a, O, IS_VIEW>
1177where
1178    O: OffsetSizeTrait + ArrowNativeTypeOp,
1179{
1180    fn try_new(
1181        field: FieldRef,
1182        element_data_type: &'a DataType,
1183        cast_options: &'a CastOptions,
1184        capacity: usize,
1185        shredded: bool,
1186    ) -> Result<Self> {
1187        if capacity >= isize::MAX as usize {
1188            return Err(ArrowError::ComputeError(
1189                "Capacity exceeds isize::MAX when reserving list offsets".to_string(),
1190            ));
1191        }
1192        let mut offsets = Vec::with_capacity(capacity + 1);
1193        offsets.push(O::ZERO);
1194        let element_builder = if shredded {
1195            let builder = make_variant_to_shredded_variant_arrow_row_builder(
1196                element_data_type,
1197                cast_options,
1198                capacity,
1199                NullValue::ArrayElement,
1200            )?;
1201            ListElementBuilder::Shredded(Box::new(builder))
1202        } else {
1203            let builder =
1204                make_typed_variant_to_arrow_row_builder(element_data_type, cast_options, capacity)?;
1205            ListElementBuilder::Typed(Box::new(builder))
1206        };
1207
1208        Ok(Self {
1209            field,
1210            offsets,
1211            element_builder,
1212            nulls: NullBufferBuilder::new(capacity),
1213            current_offset: O::ZERO,
1214            cast_options,
1215        })
1216    }
1217
1218    fn append_null(&mut self) -> Result<()> {
1219        self.offsets.push(self.current_offset);
1220        self.nulls.append_null();
1221        Ok(())
1222    }
1223
1224    fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
1225        match variant_cast_with_options(value, self.cast_options, Variant::as_list) {
1226            Ok(Some(list)) => {
1227                for element in list.iter() {
1228                    self.element_builder.append_value(element)?;
1229                    self.current_offset = self.current_offset.add_checked(O::ONE)?;
1230                }
1231                self.offsets.push(self.current_offset);
1232                self.nulls.append_non_null();
1233                Ok(true)
1234            }
1235            Ok(None) => {
1236                self.append_null()?;
1237                Ok(false)
1238            }
1239            Err(_) => Err(ArrowError::CastError(format!(
1240                "Failed to extract list from variant {value:?}"
1241            ))),
1242        }
1243    }
1244
1245    fn finish(mut self) -> Result<ArrayRef> {
1246        let element_array: ArrayRef = self.element_builder.finish()?;
1247        let field = Arc::new(
1248            self.field
1249                .as_ref()
1250                .clone()
1251                .with_data_type(element_array.data_type().clone()),
1252        );
1253
1254        if IS_VIEW {
1255            // NOTE: `offsets` is never empty (constructor pushes an entry)
1256            let mut sizes = Vec::with_capacity(self.offsets.len() - 1);
1257            for i in 1..self.offsets.len() {
1258                sizes.push(self.offsets[i] - self.offsets[i - 1]);
1259            }
1260            self.offsets.pop();
1261            let list_view_array = GenericListViewArray::<O>::new(
1262                field,
1263                ScalarBuffer::from(self.offsets),
1264                ScalarBuffer::from(sizes),
1265                element_array,
1266                self.nulls.finish(),
1267            );
1268            Ok(Arc::new(list_view_array))
1269        } else {
1270            let list_array = GenericListArray::<O>::new(
1271                field,
1272                OffsetBuffer::<O>::new(ScalarBuffer::from(self.offsets)),
1273                element_array,
1274                self.nulls.finish(),
1275            );
1276            Ok(Arc::new(list_array))
1277        }
1278    }
1279}
1280
1281pub(crate) struct VariantToFixedSizeListArrowRowBuilder<'a> {
1282    field: FieldRef,
1283    list_size: i32,
1284    element_builder: ListElementBuilder<'a>,
1285    nulls: NullBufferBuilder,
1286    cast_options: &'a CastOptions<'a>,
1287    shredded: bool,
1288}
1289
1290impl<'a> VariantToFixedSizeListArrowRowBuilder<'a> {
1291    fn try_new(
1292        field: FieldRef,
1293        element_data_type: &'a DataType,
1294        list_size: i32,
1295        cast_options: &'a CastOptions,
1296        capacity: usize,
1297        shredded: bool,
1298    ) -> Result<Self> {
1299        let element_builder = if shredded {
1300            let builder = make_variant_to_shredded_variant_arrow_row_builder(
1301                element_data_type,
1302                cast_options,
1303                capacity,
1304                NullValue::ArrayElement,
1305            )?;
1306            ListElementBuilder::Shredded(Box::new(builder))
1307        } else {
1308            let builder =
1309                make_typed_variant_to_arrow_row_builder(element_data_type, cast_options, capacity)?;
1310            ListElementBuilder::Typed(Box::new(builder))
1311        };
1312        Ok(Self {
1313            field,
1314            list_size,
1315            element_builder,
1316            nulls: NullBufferBuilder::new(capacity),
1317            cast_options,
1318            shredded,
1319        })
1320    }
1321
1322    fn append_null(&mut self) -> Result<()> {
1323        for _ in 0..self.list_size {
1324            self.element_builder.append_null()?;
1325        }
1326        self.nulls.append_null();
1327        Ok(())
1328    }
1329
1330    fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
1331        match variant_cast_with_options(value, self.cast_options, Variant::as_list) {
1332            Ok(Some(list)) => {
1333                let len = list.len();
1334                if len != self.list_size as usize {
1335                    if self.cast_options.safe && !self.shredded {
1336                        self.append_null()?;
1337                        return Ok(false);
1338                    }
1339                    return Err(ArrowError::CastError(format!(
1340                        "Expected fixed size list of size {}, got size {}",
1341                        self.list_size, len
1342                    )));
1343                }
1344                for element in list.iter() {
1345                    self.element_builder.append_value(element)?;
1346                }
1347                self.nulls.append_non_null();
1348                Ok(true)
1349            }
1350            Ok(None) => {
1351                self.append_null()?;
1352                Ok(false)
1353            }
1354            Err(_) => Err(ArrowError::CastError(format!(
1355                "Failed to extract list from variant {value:?}"
1356            ))),
1357        }
1358    }
1359
1360    fn finish(mut self) -> Result<ArrayRef> {
1361        let element_array: ArrayRef = self.element_builder.finish()?;
1362        let field = Arc::new(
1363            self.field
1364                .as_ref()
1365                .clone()
1366                .with_data_type(element_array.data_type().clone()),
1367        );
1368        let fixed_size_list_array =
1369            FixedSizeListArray::try_new(field, self.list_size, element_array, self.nulls.finish())?;
1370        Ok(Arc::new(fixed_size_list_array))
1371    }
1372}
1373
1374/// Builder for creating VariantArray output (for path extraction without type conversion)
1375pub(crate) struct VariantToBinaryVariantArrowRowBuilder {
1376    metadata: ArrayRef,
1377    builder: VariantValueArrayBuilder,
1378    nulls: NullBufferBuilder,
1379}
1380
1381impl VariantToBinaryVariantArrowRowBuilder {
1382    fn new(metadata: ArrayRef, capacity: usize) -> Self {
1383        Self {
1384            metadata,
1385            builder: VariantValueArrayBuilder::new(capacity),
1386            nulls: NullBufferBuilder::new(capacity),
1387        }
1388    }
1389}
1390
1391impl VariantToBinaryVariantArrowRowBuilder {
1392    fn append_null(&mut self) -> Result<()> {
1393        self.builder.append_null();
1394        self.nulls.append_null();
1395        Ok(())
1396    }
1397
1398    fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
1399        self.builder.append_value(value);
1400        self.nulls.append_non_null();
1401        Ok(true)
1402    }
1403
1404    fn finish(mut self) -> Result<ArrayRef> {
1405        let variant_array = VariantArray::from_parts(
1406            self.metadata,
1407            Arc::new(self.builder.build()?),
1408            None, // no typed_value column
1409            self.nulls.finish(),
1410        );
1411
1412        Ok(ArrayRef::from(variant_array))
1413    }
1414}
1415
1416#[derive(Default)]
1417struct FakeNullBuilder {
1418    item_count: usize,
1419}
1420
1421impl FakeNullBuilder {
1422    fn append_value(&mut self, _: ()) {
1423        self.item_count += 1;
1424    }
1425
1426    fn append_null(&mut self) {
1427        self.item_count += 1;
1428    }
1429
1430    fn finish(self) -> NullArray {
1431        NullArray::new(self.item_count)
1432    }
1433}
1434
1435define_variant_to_primitive_builder!(
1436    struct VariantToNullArrowRowBuilder<'a>
1437    |_capacity| -> FakeNullBuilder { FakeNullBuilder::default() },
1438    |value| value.as_null(),
1439    type_name: "Null"
1440);
1441
1442#[cfg(test)]
1443mod tests {
1444    use super::{
1445        make_primitive_variant_to_arrow_row_builder, make_typed_variant_to_arrow_row_builder,
1446    };
1447    use arrow::array::{
1448        Array, Decimal32Array, FixedSizeBinaryArray, Int32Array, ListArray, StructArray,
1449    };
1450    use arrow::compute::CastOptions;
1451    use arrow::datatypes::{DataType, Field, Fields, UnionFields, UnionMode};
1452    use arrow::error::ArrowError;
1453    use parquet_variant::{Variant, VariantDecimal4};
1454    use std::sync::Arc;
1455    use uuid::Uuid;
1456
1457    #[test]
1458    fn make_primitive_builder_rejects_non_primitive_types() {
1459        let cast_options = CastOptions::default();
1460        let item_field = Arc::new(Field::new("item", DataType::Int32, true));
1461        let struct_fields = Fields::from(vec![Field::new("child", DataType::Int32, true)]);
1462        let map_entries_field = Arc::new(Field::new(
1463            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
1464            DataType::Struct(Fields::from(vec![
1465                Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
1466                Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Float64, true),
1467            ])),
1468            true,
1469        ));
1470        let union_fields =
1471            UnionFields::try_new(vec![1], vec![Field::new("child", DataType::Int32, true)])
1472                .unwrap();
1473        let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int32, false));
1474        let ree_values_field = Arc::new(Field::new("values", DataType::Utf8, true));
1475
1476        let non_primitive_types = vec![
1477            DataType::List(item_field.clone()),
1478            DataType::LargeList(item_field.clone()),
1479            DataType::ListView(item_field.clone()),
1480            DataType::LargeListView(item_field.clone()),
1481            DataType::FixedSizeList(item_field.clone(), 2),
1482            DataType::Struct(struct_fields.clone()),
1483            DataType::Map(map_entries_field.clone(), false),
1484            DataType::Union(union_fields.clone(), UnionMode::Dense),
1485            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1486            DataType::RunEndEncoded(run_ends_field.clone(), ree_values_field.clone()),
1487        ];
1488
1489        for data_type in non_primitive_types {
1490            let err =
1491                match make_primitive_variant_to_arrow_row_builder(&data_type, &cast_options, 1) {
1492                    Ok(_) => panic!("non-primitive type {data_type:?} should be rejected"),
1493                    Err(err) => err,
1494                };
1495
1496            match err {
1497                ArrowError::InvalidArgumentError(msg) => {
1498                    assert!(msg.contains(&format!("{data_type:?}")));
1499                }
1500                other => panic!("expected InvalidArgumentError, got {other:?}"),
1501            }
1502        }
1503    }
1504
1505    #[test]
1506    fn strict_cast_allows_variant_null_for_primitive_builder() {
1507        let cast_options = CastOptions {
1508            safe: false,
1509            ..Default::default()
1510        };
1511        let mut builder =
1512            make_primitive_variant_to_arrow_row_builder(&DataType::Int32, &cast_options, 2)
1513                .unwrap();
1514
1515        assert!(!builder.append_value(&Variant::Null).unwrap());
1516        assert!(builder.append_value(&Variant::Int32(42)).unwrap());
1517
1518        let array = builder.finish().unwrap();
1519        let int_array = array.as_any().downcast_ref::<Int32Array>().unwrap();
1520        assert!(int_array.is_null(0));
1521        assert_eq!(int_array.value(1), 42);
1522    }
1523
1524    #[test]
1525    fn strict_cast_allows_variant_null_for_decimal_builder() {
1526        let cast_options = CastOptions {
1527            safe: false,
1528            ..Default::default()
1529        };
1530        let mut builder = make_primitive_variant_to_arrow_row_builder(
1531            &DataType::Decimal32(9, 2),
1532            &cast_options,
1533            2,
1534        )
1535        .unwrap();
1536        let decimal_variant: Variant<'_, '_> = VariantDecimal4::try_new(1234, 2).unwrap().into();
1537
1538        assert!(!builder.append_value(&Variant::Null).unwrap());
1539        assert!(builder.append_value(&decimal_variant).unwrap());
1540
1541        let array = builder.finish().unwrap();
1542        let decimal_array = array.as_any().downcast_ref::<Decimal32Array>().unwrap();
1543        assert!(decimal_array.is_null(0));
1544        assert_eq!(decimal_array.value(1), 1234);
1545    }
1546
1547    #[test]
1548    fn strict_cast_allows_variant_null_for_uuid_builder() {
1549        let cast_options = CastOptions {
1550            safe: false,
1551            ..Default::default()
1552        };
1553        let mut builder = make_primitive_variant_to_arrow_row_builder(
1554            &DataType::FixedSizeBinary(16),
1555            &cast_options,
1556            2,
1557        )
1558        .unwrap();
1559        let uuid = Uuid::nil();
1560
1561        assert!(!builder.append_value(&Variant::Null).unwrap());
1562        assert!(builder.append_value(&Variant::Uuid(uuid)).unwrap());
1563
1564        let array = builder.finish().unwrap();
1565        let uuid_array = array
1566            .as_any()
1567            .downcast_ref::<FixedSizeBinaryArray>()
1568            .unwrap();
1569        assert!(uuid_array.is_null(0));
1570        assert_eq!(uuid_array.value(1), uuid.as_bytes());
1571    }
1572
1573    #[test]
1574    fn strict_cast_allows_variant_null_for_list_and_struct_builders() {
1575        let cast_options = CastOptions {
1576            safe: false,
1577            ..Default::default()
1578        };
1579
1580        let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int64, true)));
1581        let mut list_builder =
1582            make_typed_variant_to_arrow_row_builder(&list_type, &cast_options, 1).unwrap();
1583        assert!(!list_builder.append_value(Variant::Null).unwrap());
1584        let list_array = list_builder.finish().unwrap();
1585        let list_array = list_array.as_any().downcast_ref::<ListArray>().unwrap();
1586        assert!(list_array.is_null(0));
1587
1588        let struct_type =
1589            DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)]));
1590        let mut struct_builder =
1591            make_typed_variant_to_arrow_row_builder(&struct_type, &cast_options, 1).unwrap();
1592        assert!(!struct_builder.append_value(Variant::Null).unwrap());
1593        let struct_array = struct_builder.finish().unwrap();
1594        let struct_array = struct_array.as_any().downcast_ref::<StructArray>().unwrap();
1595        assert!(struct_array.is_null(0));
1596    }
1597}
parquet_variant_compute/variant_to_arrow.rs

parquet_variant_compute/
variant_to_arrow.rs
