parquet_variant_compute/variant_get/output/
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.
17
18mod primitive;
19mod variant;
20
21use crate::variant_get::output::primitive::PrimitiveOutputBuilder;
22use crate::variant_get::output::variant::VariantOutputBuilder;
23use crate::variant_get::GetOptions;
24use crate::VariantArray;
25use arrow::array::{ArrayRef, BinaryViewArray};
26use arrow::datatypes::{Int16Type, Int32Type};
27use arrow::error::Result;
28use arrow_schema::{ArrowError, DataType};
29
30/// This trait represents something that gets the output of the variant_get kernel.
31///
32/// For example, there are specializations for writing the output as a VariantArray,
33/// or as a specific type (e.g. Int32Array).
34///
35/// See [`instantiate_output_builder`] to create an instance of this trait.
36pub(crate) trait OutputBuilder {
37    /// create output for a shredded variant array
38    fn partially_shredded(
39        &self,
40        variant_array: &VariantArray,
41        metadata: &BinaryViewArray,
42        value_field: &BinaryViewArray,
43        typed_value: &ArrayRef,
44    ) -> Result<ArrayRef>;
45
46    /// output for a perfectly shredded variant array
47    fn typed(
48        &self,
49        variant_array: &VariantArray,
50        metadata: &BinaryViewArray,
51        typed_value: &ArrayRef,
52    ) -> Result<ArrayRef>;
53
54    /// write out an unshredded variant array
55    fn unshredded(
56        &self,
57        variant_array: &VariantArray,
58        metadata: &BinaryViewArray,
59        value_field: &BinaryViewArray,
60    ) -> Result<ArrayRef>;
61
62    /// write out an all-null variant array
63    fn all_null(
64        &self,
65        variant_array: &VariantArray,
66        metadata: &BinaryViewArray,
67    ) -> Result<ArrayRef>;
68}
69
70pub(crate) fn instantiate_output_builder<'a>(
71    options: GetOptions<'a>,
72) -> Result<Box<dyn OutputBuilder + 'a>> {
73    let GetOptions {
74        as_type,
75        path,
76        cast_options,
77    } = options;
78
79    let Some(as_type) = as_type else {
80        return Ok(Box::new(VariantOutputBuilder::new(path)));
81    };
82
83    // handle typed output
84    match as_type.data_type() {
85        DataType::Int32 => Ok(Box::new(PrimitiveOutputBuilder::<Int32Type>::new(
86            path,
87            as_type,
88            cast_options,
89        ))),
90        DataType::Int16 => Ok(Box::new(PrimitiveOutputBuilder::<Int16Type>::new(
91            path,
92            as_type,
93            cast_options,
94        ))),
95        dt => Err(ArrowError::NotYetImplemented(format!(
96            "variant_get with as_type={dt} is not implemented yet",
97        ))),
98    }
99}
parquet_variant_compute/variant_get/output/mod.rs

parquet_variant_compute/variant_get/output/
mod.rs
