parquet_variant_compute/variant_get/
mod.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.
17use arrow::{
18    array::{Array, ArrayRef},
19    compute::CastOptions,
20    error::Result,
21};
22use arrow_schema::{ArrowError, FieldRef};
23use parquet_variant::VariantPath;
24
25use crate::variant_array::ShreddingState;
26use crate::variant_get::output::instantiate_output_builder;
27use crate::VariantArray;
28
29mod output;
30
31/// Returns an array with the specified path extracted from the variant values.
32///
33/// The return array type depends on the `as_type` field of the options parameter
34/// 1. `as_type: None`: a VariantArray is returned. The values in this new VariantArray will point
35///    to the specified path.
36/// 2. `as_type: Some(<specific field>)`: an array of the specified type is returned.
37pub fn variant_get(input: &ArrayRef, options: GetOptions) -> Result<ArrayRef> {
38    let variant_array: &VariantArray = input.as_any().downcast_ref().ok_or_else(|| {
39        ArrowError::InvalidArgumentError(
40            "expected a VariantArray as the input for variant_get".to_owned(),
41        )
42    })?;
43
44    // Create the output writer based on the specified output options
45    let output_builder = instantiate_output_builder(options.clone())?;
46
47    // Dispatch based on the shredding state of the input variant array
48    match variant_array.shredding_state() {
49        ShreddingState::PartiallyShredded {
50            metadata,
51            value,
52            typed_value,
53        } => output_builder.partially_shredded(variant_array, metadata, value, typed_value),
54        ShreddingState::Typed {
55            metadata,
56            typed_value,
57        } => output_builder.typed(variant_array, metadata, typed_value),
58        ShreddingState::Unshredded { metadata, value } => {
59            output_builder.unshredded(variant_array, metadata, value)
60        }
61        ShreddingState::AllNull { metadata } => output_builder.all_null(variant_array, metadata),
62    }
63}
64
65/// Controls the action of the variant_get kernel.
66#[derive(Debug, Clone, Default)]
67pub struct GetOptions<'a> {
68    /// What path to extract
69    pub path: VariantPath<'a>,
70    /// if `as_type` is None, the returned array will itself be a VariantArray.
71    ///
72    /// if `as_type` is `Some(type)` the field is returned as the specified type.
73    pub as_type: Option<FieldRef>,
74    /// Controls the casting behavior (e.g. error vs substituting null on cast error).
75    pub cast_options: CastOptions<'a>,
76}
77
78impl<'a> GetOptions<'a> {
79    /// Construct default options to get the specified path as a variant.
80    pub fn new() -> Self {
81        Default::default()
82    }
83
84    /// Construct options to get the specified path as a variant.
85    pub fn new_with_path(path: VariantPath<'a>) -> Self {
86        Self {
87            path,
88            as_type: None,
89            cast_options: Default::default(),
90        }
91    }
92
93    /// Specify the type to return.
94    pub fn with_as_type(mut self, as_type: Option<FieldRef>) -> Self {
95        self.as_type = as_type;
96        self
97    }
98
99    /// Specify the cast options to use when casting to the specified type.
100    pub fn with_cast_options(mut self, cast_options: CastOptions<'a>) -> Self {
101        self.cast_options = cast_options;
102        self
103    }
104}
105
106#[cfg(test)]
107mod test {
108    use std::sync::Arc;
109
110    use arrow::array::{
111        Array, ArrayRef, BinaryViewArray, Float16Array, Float32Array, Float64Array, Int16Array,
112        Int32Array, Int64Array, Int8Array, StringArray, StructArray, UInt16Array, UInt32Array,
113        UInt64Array, UInt8Array,
114    };
115    use arrow::buffer::NullBuffer;
116    use arrow::compute::CastOptions;
117    use arrow_schema::{DataType, Field, FieldRef, Fields};
118    use parquet_variant::{Variant, VariantPath};
119
120    use crate::json_to_variant;
121    use crate::VariantArray;
122
123    use super::{variant_get, GetOptions};
124
125    fn single_variant_get_test(input_json: &str, path: VariantPath, expected_json: &str) {
126        // Create input array from JSON string
127        let input_array_ref: ArrayRef = Arc::new(StringArray::from(vec![Some(input_json)]));
128        let input_variant_array_ref: ArrayRef =
129            Arc::new(json_to_variant(&input_array_ref).unwrap());
130
131        let result =
132            variant_get(&input_variant_array_ref, GetOptions::new_with_path(path)).unwrap();
133
134        // Create expected array from JSON string
135        let expected_array_ref: ArrayRef = Arc::new(StringArray::from(vec![Some(expected_json)]));
136        let expected_variant_array = json_to_variant(&expected_array_ref).unwrap();
137
138        let result_array: &VariantArray = result.as_any().downcast_ref().unwrap();
139        assert_eq!(
140            result_array.len(),
141            1,
142            "Expected result array to have length 1"
143        );
144        assert!(
145            result_array.nulls().is_none(),
146            "Expected no nulls in result array"
147        );
148        let result_variant = result_array.value(0);
149        let expected_variant = expected_variant_array.value(0);
150        assert_eq!(
151            result_variant, expected_variant,
152            "Result variant does not match expected variant"
153        );
154    }
155
156    #[test]
157    fn get_primitive_variant_field() {
158        single_variant_get_test(
159            r#"{"some_field": 1234}"#,
160            VariantPath::from("some_field"),
161            "1234",
162        );
163    }
164
165    #[test]
166    fn get_primitive_variant_list_index() {
167        single_variant_get_test("[1234, 5678]", VariantPath::from(0), "1234");
168    }
169
170    #[test]
171    fn get_primitive_variant_inside_object_of_object() {
172        single_variant_get_test(
173            r#"{"top_level_field": {"inner_field": 1234}}"#,
174            VariantPath::from("top_level_field").join("inner_field"),
175            "1234",
176        );
177    }
178
179    #[test]
180    fn get_primitive_variant_inside_list_of_object() {
181        single_variant_get_test(
182            r#"[{"some_field": 1234}]"#,
183            VariantPath::from(0).join("some_field"),
184            "1234",
185        );
186    }
187
188    #[test]
189    fn get_primitive_variant_inside_object_of_list() {
190        single_variant_get_test(
191            r#"{"some_field": [1234]}"#,
192            VariantPath::from("some_field").join(0),
193            "1234",
194        );
195    }
196
197    #[test]
198    fn get_complex_variant() {
199        single_variant_get_test(
200            r#"{"top_level_field": {"inner_field": 1234}}"#,
201            VariantPath::from("top_level_field"),
202            r#"{"inner_field": 1234}"#,
203        );
204    }
205
206    /// Partial Shredding: extract a value as a VariantArray
207    macro_rules! numeric_partially_shredded_test {
208        ($primitive_type:ty, $data_fn:ident) => {
209            let array = $data_fn();
210            let options = GetOptions::new();
211            let result = variant_get(&array, options).unwrap();
212
213            // expect the result is a VariantArray
214            let result: &VariantArray = result.as_any().downcast_ref().unwrap();
215            assert_eq!(result.len(), 4);
216
217            // Expect the values are the same as the original values
218            assert_eq!(
219                result.value(0),
220                Variant::from(<$primitive_type>::try_from(34u8).unwrap())
221            );
222            assert!(!result.is_valid(1));
223            assert_eq!(result.value(2), Variant::from("n/a"));
224            assert_eq!(
225                result.value(3),
226                Variant::from(<$primitive_type>::try_from(100u8).unwrap())
227            );
228        };
229    }
230
231    #[test]
232    fn get_variant_partially_shredded_int8_as_variant() {
233        numeric_partially_shredded_test!(i8, partially_shredded_int8_variant_array);
234    }
235
236    #[test]
237    fn get_variant_partially_shredded_int16_as_variant() {
238        numeric_partially_shredded_test!(i16, partially_shredded_int16_variant_array);
239    }
240
241    #[test]
242    fn get_variant_partially_shredded_int32_as_variant() {
243        numeric_partially_shredded_test!(i32, partially_shredded_int32_variant_array);
244    }
245
246    #[test]
247    fn get_variant_partially_shredded_int64_as_variant() {
248        numeric_partially_shredded_test!(i64, partially_shredded_int64_variant_array);
249    }
250
251    #[test]
252    fn get_variant_partially_shredded_uint8_as_variant() {
253        numeric_partially_shredded_test!(u8, partially_shredded_uint8_variant_array);
254    }
255
256    #[test]
257    fn get_variant_partially_shredded_uint16_as_variant() {
258        numeric_partially_shredded_test!(u16, partially_shredded_uint16_variant_array);
259    }
260
261    #[test]
262    fn get_variant_partially_shredded_uint32_as_variant() {
263        numeric_partially_shredded_test!(u32, partially_shredded_uint32_variant_array);
264    }
265
266    #[test]
267    fn get_variant_partially_shredded_uint64_as_variant() {
268        numeric_partially_shredded_test!(u64, partially_shredded_uint64_variant_array);
269    }
270
271    #[test]
272    fn get_variant_partially_shredded_float16_as_variant() {
273        numeric_partially_shredded_test!(half::f16, partially_shredded_float16_variant_array);
274    }
275
276    #[test]
277    fn get_variant_partially_shredded_float32_as_variant() {
278        numeric_partially_shredded_test!(f32, partially_shredded_float32_variant_array);
279    }
280
281    #[test]
282    fn get_variant_partially_shredded_float64_as_variant() {
283        numeric_partially_shredded_test!(f64, partially_shredded_float64_variant_array);
284    }
285
286    /// Shredding: extract a value as an Int32Array
287    #[test]
288    fn get_variant_shredded_int32_as_int32_safe_cast() {
289        // Extract the typed value as Int32Array
290        let array = partially_shredded_int32_variant_array();
291        // specify we want the typed value as Int32
292        let field = Field::new("typed_value", DataType::Int32, true);
293        let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
294        let result = variant_get(&array, options).unwrap();
295        let expected: ArrayRef = Arc::new(Int32Array::from(vec![
296            Some(34),
297            None,
298            None, // "n/a" is not an Int32 so converted to null
299            Some(100),
300        ]));
301        assert_eq!(&result, &expected)
302    }
303
304    /// Shredding: extract a value as an Int32Array, unsafe cast (should error on "n/a")
305
306    #[test]
307    fn get_variant_shredded_int32_as_int32_unsafe_cast() {
308        // Extract the typed value as Int32Array
309        let array = partially_shredded_int32_variant_array();
310        let field = Field::new("typed_value", DataType::Int32, true);
311        let cast_options = CastOptions {
312            safe: false, // unsafe cast
313            ..Default::default()
314        };
315        let options = GetOptions::new()
316            .with_as_type(Some(FieldRef::from(field)))
317            .with_cast_options(cast_options);
318
319        let err = variant_get(&array, options).unwrap_err();
320        // TODO make this error message nicer (not Debug format)
321        assert_eq!(err.to_string(), "Cast error: Failed to extract primitive of type Int32 from variant ShortString(ShortString(\"n/a\")) at path VariantPath([])");
322    }
323
324    /// Perfect Shredding: extract the typed value as a VariantArray
325    macro_rules! numeric_perfectly_shredded_test {
326        ($primitive_type:ty, $data_fn:ident) => {
327            let array = $data_fn();
328            let options = GetOptions::new();
329            let result = variant_get(&array, options).unwrap();
330
331            // expect the result is a VariantArray
332            let result: &VariantArray = result.as_any().downcast_ref().unwrap();
333            assert_eq!(result.len(), 3);
334
335            // Expect the values are the same as the original values
336            assert_eq!(
337                result.value(0),
338                Variant::from(<$primitive_type>::try_from(1u8).unwrap())
339            );
340            assert_eq!(
341                result.value(1),
342                Variant::from(<$primitive_type>::try_from(2u8).unwrap())
343            );
344            assert_eq!(
345                result.value(2),
346                Variant::from(<$primitive_type>::try_from(3u8).unwrap())
347            );
348        };
349    }
350
351    #[test]
352    fn get_variant_perfectly_shredded_int8_as_variant() {
353        numeric_perfectly_shredded_test!(i8, perfectly_shredded_int8_variant_array);
354    }
355
356    #[test]
357    fn get_variant_perfectly_shredded_int16_as_variant() {
358        numeric_perfectly_shredded_test!(i16, perfectly_shredded_int16_variant_array);
359    }
360
361    #[test]
362    fn get_variant_perfectly_shredded_int32_as_variant() {
363        numeric_perfectly_shredded_test!(i32, perfectly_shredded_int32_variant_array);
364    }
365
366    #[test]
367    fn get_variant_perfectly_shredded_int64_as_variant() {
368        numeric_perfectly_shredded_test!(i64, perfectly_shredded_int64_variant_array);
369    }
370
371    #[test]
372    fn get_variant_perfectly_shredded_uint8_as_variant() {
373        numeric_perfectly_shredded_test!(u8, perfectly_shredded_uint8_variant_array);
374    }
375
376    #[test]
377    fn get_variant_perfectly_shredded_uint16_as_variant() {
378        numeric_perfectly_shredded_test!(u16, perfectly_shredded_uint16_variant_array);
379    }
380
381    #[test]
382    fn get_variant_perfectly_shredded_uint32_as_variant() {
383        numeric_perfectly_shredded_test!(u32, perfectly_shredded_uint32_variant_array);
384    }
385
386    #[test]
387    fn get_variant_perfectly_shredded_uint64_as_variant() {
388        numeric_perfectly_shredded_test!(u64, perfectly_shredded_uint64_variant_array);
389    }
390
391    #[test]
392    fn get_variant_perfectly_shredded_float16_as_variant() {
393        numeric_perfectly_shredded_test!(half::f16, perfectly_shredded_float16_variant_array);
394    }
395
396    #[test]
397    fn get_variant_perfectly_shredded_float32_as_variant() {
398        numeric_perfectly_shredded_test!(f32, perfectly_shredded_float32_variant_array);
399    }
400
401    #[test]
402    fn get_variant_perfectly_shredded_float64_as_variant() {
403        numeric_perfectly_shredded_test!(f64, perfectly_shredded_float64_variant_array);
404    }
405
406    /// Shredding: Extract the typed value as Int32Array
407    #[test]
408    fn get_variant_perfectly_shredded_int32_as_int32() {
409        // Extract the typed value as Int32Array
410        let array = perfectly_shredded_int32_variant_array();
411        // specify we want the typed value as Int32
412        let field = Field::new("typed_value", DataType::Int32, true);
413        let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
414        let result = variant_get(&array, options).unwrap();
415        let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]));
416        assert_eq!(&result, &expected)
417    }
418
419    /// AllNull: extract a value as a VariantArray
420    #[test]
421    fn get_variant_all_null_as_variant() {
422        let array = all_null_variant_array();
423        let options = GetOptions::new();
424        let result = variant_get(&array, options).unwrap();
425
426        // expect the result is a VariantArray
427        let result: &VariantArray = result.as_any().downcast_ref().unwrap();
428        assert_eq!(result.len(), 3);
429
430        // All values should be null
431        assert!(!result.is_valid(0));
432        assert!(!result.is_valid(1));
433        assert!(!result.is_valid(2));
434    }
435
436    /// AllNull: extract a value as an Int32Array
437    #[test]
438    fn get_variant_all_null_as_int32() {
439        let array = all_null_variant_array();
440        // specify we want the typed value as Int32
441        let field = Field::new("typed_value", DataType::Int32, true);
442        let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
443        let result = variant_get(&array, options).unwrap();
444
445        let expected: ArrayRef = Arc::new(Int32Array::from(vec![
446            Option::<i32>::None,
447            Option::<i32>::None,
448            Option::<i32>::None,
449        ]));
450        assert_eq!(&result, &expected)
451    }
452
453    #[test]
454    fn get_variant_perfectly_shredded_int16_as_int16() {
455        // Extract the typed value as Int16Array
456        let array = perfectly_shredded_int16_variant_array();
457        // specify we want the typed value as Int16
458        let field = Field::new("typed_value", DataType::Int16, true);
459        let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
460        let result = variant_get(&array, options).unwrap();
461        let expected: ArrayRef = Arc::new(Int16Array::from(vec![Some(1), Some(2), Some(3)]));
462        assert_eq!(&result, &expected)
463    }
464
465    /// Return a VariantArray that represents a perfectly "shredded" variant
466    /// for the given typed value.
467    ///
468    /// The schema of the corresponding `StructArray` would look like this:
469    ///
470    /// ```text
471    /// StructArray {
472    ///   metadata: BinaryViewArray,
473    ///   typed_value: Int32Array,
474    /// }
475    /// ```
476    macro_rules! numeric_perfectly_shredded_variant_array_fn {
477        ($func:ident, $array_type:ident, $primitive_type:ty) => {
478            fn $func() -> ArrayRef {
479                // At the time of writing, the `VariantArrayBuilder` does not support shredding.
480                // so we must construct the array manually.  see https://github.com/apache/arrow-rs/issues/7895
481                let (metadata, _value) = { parquet_variant::VariantBuilder::new().finish() };
482                let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 3));
483                let typed_value = $array_type::from(vec![
484                    Some(<$primitive_type>::try_from(1u8).unwrap()),
485                    Some(<$primitive_type>::try_from(2u8).unwrap()),
486                    Some(<$primitive_type>::try_from(3u8).unwrap()),
487                ]);
488
489                let struct_array = StructArrayBuilder::new()
490                    .with_field("metadata", Arc::new(metadata))
491                    .with_field("typed_value", Arc::new(typed_value))
492                    .build();
493
494                Arc::new(
495                    VariantArray::try_new(Arc::new(struct_array))
496                        .expect("should create variant array"),
497                )
498            }
499        };
500    }
501
502    numeric_perfectly_shredded_variant_array_fn!(
503        perfectly_shredded_int8_variant_array,
504        Int8Array,
505        i8
506    );
507    numeric_perfectly_shredded_variant_array_fn!(
508        perfectly_shredded_int16_variant_array,
509        Int16Array,
510        i16
511    );
512    numeric_perfectly_shredded_variant_array_fn!(
513        perfectly_shredded_int32_variant_array,
514        Int32Array,
515        i32
516    );
517    numeric_perfectly_shredded_variant_array_fn!(
518        perfectly_shredded_int64_variant_array,
519        Int64Array,
520        i64
521    );
522    numeric_perfectly_shredded_variant_array_fn!(
523        perfectly_shredded_uint8_variant_array,
524        UInt8Array,
525        u8
526    );
527    numeric_perfectly_shredded_variant_array_fn!(
528        perfectly_shredded_uint16_variant_array,
529        UInt16Array,
530        u16
531    );
532    numeric_perfectly_shredded_variant_array_fn!(
533        perfectly_shredded_uint32_variant_array,
534        UInt32Array,
535        u32
536    );
537    numeric_perfectly_shredded_variant_array_fn!(
538        perfectly_shredded_uint64_variant_array,
539        UInt64Array,
540        u64
541    );
542    numeric_perfectly_shredded_variant_array_fn!(
543        perfectly_shredded_float16_variant_array,
544        Float16Array,
545        half::f16
546    );
547    numeric_perfectly_shredded_variant_array_fn!(
548        perfectly_shredded_float32_variant_array,
549        Float32Array,
550        f32
551    );
552    numeric_perfectly_shredded_variant_array_fn!(
553        perfectly_shredded_float64_variant_array,
554        Float64Array,
555        f64
556    );
557
558    /// Return a VariantArray that represents a normal "shredded" variant
559    /// for the following example
560    ///
561    /// Based on the example from [the doc]
562    ///
563    /// [the doc]: https://docs.google.com/document/d/1pw0AWoMQY3SjD7R4LgbPvMjG_xSCtXp3rZHkVp9jpZ4/edit?tab=t.0
564    ///
565    /// ```text
566    /// 34
567    /// null (an Arrow NULL, not a Variant::Null)
568    /// "n/a" (a string)
569    /// 100
570    /// ```
571    ///
572    /// The schema of the corresponding `StructArray` would look like this:
573    ///
574    /// ```text
575    /// StructArray {
576    ///   metadata: BinaryViewArray,
577    ///   value: BinaryViewArray,
578    ///   typed_value: Int32Array,
579    /// }
580    /// ```
581    macro_rules! numeric_partially_shredded_variant_array_fn {
582        ($func:ident, $array_type:ident, $primitive_type:ty) => {
583            fn $func() -> ArrayRef {
584                // At the time of writing, the `VariantArrayBuilder` does not support shredding.
585                // so we must construct the array manually.  see https://github.com/apache/arrow-rs/issues/7895
586                let (metadata, string_value) = {
587                    let mut builder = parquet_variant::VariantBuilder::new();
588                    builder.append_value("n/a");
589                    builder.finish()
590                };
591
592                let nulls = NullBuffer::from(vec![
593                    true,  // row 0 non null
594                    false, // row 1 is null
595                    true,  // row 2 non null
596                    true,  // row 3 non null
597                ]);
598
599                // metadata is the same for all rows
600                let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 4));
601
602                // See https://docs.google.com/document/d/1pw0AWoMQY3SjD7R4LgbPvMjG_xSCtXp3rZHkVp9jpZ4/edit?disco=AAABml8WQrY
603                // about why row1 is an empty but non null, value.
604                let values = BinaryViewArray::from(vec![
605                    None,                // row 0 is shredded, so no value
606                    Some(b"" as &[u8]),  // row 1 is null, so empty value (why?)
607                    Some(&string_value), // copy the string value "N/A"
608                    None,                // row 3 is shredded, so no value
609                ]);
610
611                let typed_value = $array_type::from(vec![
612                    Some(<$primitive_type>::try_from(34u8).unwrap()), // row 0 is shredded, so it has a value
613                    None,                                             // row 1 is null, so no value
614                    None, // row 2 is a string, so no typed value
615                    Some(<$primitive_type>::try_from(100u8).unwrap()), // row 3 is shredded, so it has a value
616                ]);
617
618                let struct_array = StructArrayBuilder::new()
619                    .with_field("metadata", Arc::new(metadata))
620                    .with_field("typed_value", Arc::new(typed_value))
621                    .with_field("value", Arc::new(values))
622                    .with_nulls(nulls)
623                    .build();
624
625                Arc::new(
626                    VariantArray::try_new(Arc::new(struct_array))
627                        .expect("should create variant array"),
628                )
629            }
630        };
631    }
632
633    numeric_partially_shredded_variant_array_fn!(
634        partially_shredded_int8_variant_array,
635        Int8Array,
636        i8
637    );
638    numeric_partially_shredded_variant_array_fn!(
639        partially_shredded_int16_variant_array,
640        Int16Array,
641        i16
642    );
643    numeric_partially_shredded_variant_array_fn!(
644        partially_shredded_int32_variant_array,
645        Int32Array,
646        i32
647    );
648    numeric_partially_shredded_variant_array_fn!(
649        partially_shredded_int64_variant_array,
650        Int64Array,
651        i64
652    );
653    numeric_partially_shredded_variant_array_fn!(
654        partially_shredded_uint8_variant_array,
655        UInt8Array,
656        u8
657    );
658    numeric_partially_shredded_variant_array_fn!(
659        partially_shredded_uint16_variant_array,
660        UInt16Array,
661        u16
662    );
663    numeric_partially_shredded_variant_array_fn!(
664        partially_shredded_uint32_variant_array,
665        UInt32Array,
666        u32
667    );
668    numeric_partially_shredded_variant_array_fn!(
669        partially_shredded_uint64_variant_array,
670        UInt64Array,
671        u64
672    );
673    numeric_partially_shredded_variant_array_fn!(
674        partially_shredded_float16_variant_array,
675        Float16Array,
676        half::f16
677    );
678    numeric_partially_shredded_variant_array_fn!(
679        partially_shredded_float32_variant_array,
680        Float32Array,
681        f32
682    );
683    numeric_partially_shredded_variant_array_fn!(
684        partially_shredded_float64_variant_array,
685        Float64Array,
686        f64
687    );
688
689    /// Builds struct arrays from component fields
690    ///
691    /// TODO: move to arrow crate
692    #[derive(Debug, Default, Clone)]
693    struct StructArrayBuilder {
694        fields: Vec<FieldRef>,
695        arrays: Vec<ArrayRef>,
696        nulls: Option<NullBuffer>,
697    }
698
699    impl StructArrayBuilder {
700        fn new() -> Self {
701            Default::default()
702        }
703
704        /// Add an array to this struct array as a field with the specified name.
705        fn with_field(mut self, field_name: &str, array: ArrayRef) -> Self {
706            let field = Field::new(field_name, array.data_type().clone(), true);
707            self.fields.push(Arc::new(field));
708            self.arrays.push(array);
709            self
710        }
711
712        /// Set the null buffer for this struct array.
713        fn with_nulls(mut self, nulls: NullBuffer) -> Self {
714            self.nulls = Some(nulls);
715            self
716        }
717
718        pub fn build(self) -> StructArray {
719            let Self {
720                fields,
721                arrays,
722                nulls,
723            } = self;
724            StructArray::new(Fields::from(fields), arrays, nulls)
725        }
726    }
727
728    /// Return a VariantArray that represents an "all null" variant
729    /// for the following example (3 null values):
730    ///
731    /// ```text
732    /// null
733    /// null
734    /// null
735    /// ```
736    ///
737    /// The schema of the corresponding `StructArray` would look like this:
738    ///
739    /// ```text
740    /// StructArray {
741    ///   metadata: BinaryViewArray,
742    /// }
743    /// ```
744    fn all_null_variant_array() -> ArrayRef {
745        let (metadata, _value) = { parquet_variant::VariantBuilder::new().finish() };
746
747        let nulls = NullBuffer::from(vec![
748            false, // row 0 is null
749            false, // row 1 is null
750            false, // row 2 is null
751        ]);
752
753        // metadata is the same for all rows (though they're all null)
754        let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 3));
755
756        let struct_array = StructArrayBuilder::new()
757            .with_field("metadata", Arc::new(metadata))
758            .with_nulls(nulls)
759            .build();
760
761        Arc::new(
762            VariantArray::try_new(Arc::new(struct_array)).expect("should create variant array"),
763        )
764    }
765}
parquet_variant_compute/variant_get/mod.rs

parquet_variant_compute/variant_get/
mod.rs
