datafusion_variant/
variant_list_construct.rs

1use std::sync::Arc;
2
3use arrow::array::StructArray;
4use arrow_schema::{DataType, Field, Fields};
5use datafusion::{
6    error::{DataFusionError, Result},
7    logical_expr::{
8        ColumnarValue, ReturnFieldArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
9    },
10    scalar::ScalarValue,
11};
12use parquet_variant::{Variant, VariantBuilder};
13use parquet_variant_compute::{VariantArray, VariantType};
14
15use crate::shared::{ensure, try_parse_variant_scalar};
16
17#[derive(Debug, Hash, PartialEq, Eq)]
18pub struct VariantListConstruct {
19    signature: Signature,
20}
21
22impl Default for VariantListConstruct {
23    fn default() -> Self {
24        Self {
25            signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable),
26        }
27    }
28}
29
30impl ScalarUDFImpl for VariantListConstruct {
31    fn as_any(&self) -> &dyn std::any::Any {
32        self
33    }
34
35    fn name(&self) -> &str {
36        "variant_list_construct"
37    }
38
39    fn signature(&self) -> &Signature {
40        &self.signature
41    }
42
43    fn return_type(&self, _arg_types: &[arrow_schema::DataType]) -> Result<arrow_schema::DataType> {
44        Err(DataFusionError::Internal(
45            "implemented return_field_from_args instead".into(),
46        ))
47    }
48
49    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Arc<Field>> {
50        let data_type = DataType::Struct(Fields::from(vec![
51            Field::new("metadata", DataType::BinaryView, false),
52            Field::new("value", DataType::BinaryView, false),
53        ]));
54
55        Ok(Arc::new(
56            Field::new(self.name(), data_type, true).with_extension_type(VariantType),
57        ))
58    }
59
60    fn invoke_with_args(
61        &self,
62        args: datafusion::logical_expr::ScalarFunctionArgs,
63    ) -> Result<ColumnarValue> {
64        // validate arguments
65        let argument_fields = args.arg_fields;
66        let argument_values = args.args;
67
68        // note: arguments must be nonempty
69        // this is required by the ListBuilder
70        ensure(
71            !argument_values.is_empty(),
72            "expected nonempty list of arguments",
73        )?;
74
75        ensure(
76            argument_fields.len() == argument_values.len(),
77            "argument fields and values must be of same length",
78        )?;
79
80        let all_variant_fields = argument_fields
81            .iter()
82            .all(|f| matches!(f.extension_type(), VariantType));
83
84        ensure(
85            all_variant_fields,
86            "argument fields must all have the VariantType extension",
87        )?;
88
89        let all_arguments_scalar = argument_values
90            .iter()
91            .all(|v| matches!(v, ColumnarValue::Scalar(_)));
92
93        ensure(
94            all_arguments_scalar,
95            "todo: how do you construct lists of lists?",
96        )?;
97
98        let mut v = VariantBuilder::new();
99        let mut l = v.new_list();
100
101        // note: it would be nice to reserve capacity
102        // somethign like: v.new_list_with_capacity(usize)
103        // or ListBuilder().with_capacity(usize)
104        // or l.reserve(usize)
105
106        for v in argument_values {
107            let ColumnarValue::Scalar(sv) = v else {
108                unreachable!()
109            };
110
111            let variant = try_parse_variant_scalar(&sv)?;
112            let v = variant.value(0);
113
114            l.append_value(v);
115        }
116
117        l.finish();
118
119        let (m, v) = v.finish();
120
121        let v = Variant::new(m.as_ref(), v.as_ref());
122
123        let out: StructArray = VariantArray::from_iter([v]).into();
124
125        Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(out))))
126    }
127}
datafusion_variant/variant_list_construct.rs

datafusion_variant/
variant_list_construct.rs
