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}