datafusion_variant/
variant_object_construct.rs

1use std::sync::Arc;
2
3use arrow::array::StructArray;
4use arrow_schema::{DataType, Field, Fields};
5use datafusion::{
6    common::exec_datafusion_err,
7    error::{DataFusionError, Result},
8    logical_expr::{
9        ColumnarValue, ReturnFieldArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
10    },
11    scalar::ScalarValue,
12};
13use parquet_variant::{Variant, VariantBuilder};
14use parquet_variant_compute::{VariantArray, VariantType};
15
16use crate::shared::{ensure, try_parse_string_scalar, try_parse_variant_scalar};
17
18#[derive(Debug, Hash, PartialEq, Eq)]
19pub struct VariantObjectConstruct {
20    signature: Signature,
21}
22
23impl Default for VariantObjectConstruct {
24    fn default() -> Self {
25        Self {
26            signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable),
27        }
28    }
29}
30
31impl ScalarUDFImpl for VariantObjectConstruct {
32    fn as_any(&self) -> &dyn std::any::Any {
33        self
34    }
35
36    fn name(&self) -> &str {
37        "variant_object_construct"
38    }
39
40    fn signature(&self) -> &Signature {
41        &self.signature
42    }
43
44    fn return_type(&self, _arg_types: &[arrow_schema::DataType]) -> Result<arrow_schema::DataType> {
45        Err(DataFusionError::Internal(
46            "implemented return_field_from_args instead".into(),
47        ))
48    }
49
50    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Arc<Field>> {
51        let data_type = DataType::Struct(Fields::from(vec![
52            Field::new("metadata", DataType::BinaryView, false),
53            Field::new("value", DataType::BinaryView, false),
54        ]));
55
56        Ok(Arc::new(
57            Field::new(self.name(), data_type, true).with_extension_type(VariantType),
58        ))
59    }
60
61    fn invoke_with_args(
62        &self,
63        args: datafusion::logical_expr::ScalarFunctionArgs,
64    ) -> Result<ColumnarValue> {
65        // validate arguments
66        let argument_fields = args.arg_fields;
67        let argument_values = args.args;
68
69        ensure(
70            argument_fields.len() == argument_values.len(),
71            "argument fields and values must be of same length",
72        )?;
73
74        ensure(
75            argument_fields.len() & 1 == 0,
76            "list of arguments must be (key, value) pair",
77        )?;
78
79        let (_key_fields, value_fields): (Vec<_>, Vec<_>) = argument_fields
80            .into_iter()
81            .enumerate()
82            .partition(|(i, _)| i & 1 == 0);
83
84        let all_value_fields_have_variant_ext = value_fields
85            .iter()
86            .all(|(_, v)| matches!(v.extension_type(), VariantType));
87
88        ensure(
89            all_value_fields_have_variant_ext,
90            "expected all values in (key, value) to have a Variant ext type",
91        )?;
92
93        // sometimes field metadata is super redundant. Why check if the key fields are String
94        // when you can just try to parse out the ScalarValues?
95
96        let all_arguments_scalar = argument_values
97            .iter()
98            .all(|v| matches!(v, ColumnarValue::Scalar(_)));
99
100        // for now, let's just handle the scalar case
101        ensure(
102            all_arguments_scalar,
103            "all arguments must be scalar, todo: how do array arguments look like?",
104        )?;
105
106        let (key_values, variant_values): (Vec<_>, Vec<_>) = argument_values
107            .into_iter()
108            .enumerate()
109            .partition(|(i, _)| i & 1 == 0);
110
111        let object_keys = key_values
112            .into_iter()
113            .map(|(_, v)| {
114                let ColumnarValue::Scalar(sv) = v else {
115                    unreachable!()
116                };
117
118                try_parse_string_scalar(&sv).and_then(|opt| {
119                    opt.cloned()
120                        .ok_or_else(|| exec_datafusion_err!("expected non null string"))
121                })
122            })
123            .collect::<Result<Vec<_>, _>>()?;
124
125        let object_values = variant_values
126            .into_iter()
127            .map(|(_, v)| {
128                let ColumnarValue::Scalar(sv) = v else {
129                    unreachable!()
130                };
131
132                try_parse_variant_scalar(&sv)
133            })
134            .collect::<Result<Vec<_>, _>>()?;
135
136        // note, should we have the ability to configure behavior for duplicate keys?
137
138        let mut v = VariantBuilder::new();
139        let mut o = v.new_object();
140
141        for (k, v) in object_keys.iter().zip(object_values) {
142            let v = v.value(0);
143            o.try_insert(k, v)?;
144        }
145
146        o.finish();
147
148        let (m, v) = v.finish();
149
150        let v = Variant::new(m.as_ref(), v.as_ref());
151
152        let out: StructArray = VariantArray::from_iter([v]).into();
153
154        Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(out))))
155    }
156}
datafusion_variant/variant_object_construct.rs

datafusion_variant/
variant_object_construct.rs
