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::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
63pub(crate) fn instantiate_output_builder<'a>(
64    options: GetOptions<'a>,
65) -> Result<Box<dyn OutputBuilder + 'a>> {
66    let GetOptions {
67        as_type,
68        path,
69        cast_options,
70    } = options;
71
72    let Some(as_type) = as_type else {
73        return Ok(Box::new(VariantOutputBuilder::new(path)));
74    };
75
76    // handle typed output
77    match as_type.data_type() {
78        DataType::Int32 => Ok(Box::new(PrimitiveOutputBuilder::<Int32Type>::new(
79            path,
80            as_type,
81            cast_options,
82        ))),
83        dt => Err(ArrowError::NotYetImplemented(format!(
84            "variant_get with as_type={dt} is not implemented yet",
85        ))),
86    }
87}