datafusion_variant/
variant_object_insert.rs

1use std::sync::Arc;
2
3use arrow::array::StructArray;
4use arrow_schema::{DataType, Field, Fields};
5use datafusion::{
6    common::exec_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 VariantObjectInsert {
20    signature: Signature,
21}
22
23impl Default for VariantObjectInsert {
24    fn default() -> Self {
25        Self {
26            signature: Signature::new(TypeSignature::Any(3), Volatility::Immutable),
27        }
28    }
29}
30
31impl ScalarUDFImpl for VariantObjectInsert {
32    fn as_any(&self) -> &dyn std::any::Any {
33        self
34    }
35
36    fn name(&self) -> &str {
37        "variant_object_insert"
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        let argument_fields = args.arg_fields;
66        let argument_values = args.args;
67
68        ensure(
69            argument_fields.len() == argument_values.len(),
70            "argument fields and values must be of same length",
71        )?;
72
73        let [variant_object_to_update, key, value] = argument_values.as_slice() else {
74            return exec_err!("expected 3 arguments");
75        };
76
77        {
78            let variant_object_to_update_ext =
79                matches!(argument_fields[0].extension_type(), VariantType);
80            let value_ext = matches!(argument_fields[2].extension_type(), VariantType);
81
82            ensure(
83                variant_object_to_update_ext && value_ext,
84                "expected extension type of VariantType",
85            )?;
86        }
87
88        let key = {
89            let ColumnarValue::Scalar(key) = key else {
90                return exec_err!("expected scalar value for key");
91            };
92
93            try_parse_string_scalar(key)?
94                .ok_or_else(|| DataFusionError::Execution("expected non null string".into()))?
95        };
96
97        let value_array = {
98            let ColumnarValue::Scalar(value) = value else {
99                return exec_err!("expected scalar value for value");
100            };
101
102            try_parse_variant_scalar(value)?
103        };
104        let value = value_array.value(0);
105
106        match variant_object_to_update {
107            ColumnarValue::Scalar(scalar_variant_object_to_update) => {
108                let variant_object = try_parse_variant_scalar(scalar_variant_object_to_update)?;
109                let variant_object = variant_object.value(0);
110                let Variant::Object(variant_object) = variant_object else {
111                    return exec_err!("expected variant object");
112                };
113
114                let mut v = VariantBuilder::new();
115                let mut o = v.new_object();
116
117                o.extend(variant_object.iter());
118                o.insert(key, value);
119
120                o.finish();
121
122                let (m, v) = v.finish();
123                let v = Variant::try_new(m.as_ref(), v.as_ref())?;
124
125                let out: StructArray = VariantArray::from_iter([v]).into();
126
127                Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(out))))
128            }
129            ColumnarValue::Array(variant_objects_to_update) => {
130                let variant_objects = VariantArray::try_new(&variant_objects_to_update)?;
131
132                // grr lifetimes...
133                let buffers = variant_objects
134                    .iter()
135                    .map(|v_opt| {
136                        v_opt
137                            .map(|variant_object| {
138                                let Variant::Object(variant_object) = variant_object else {
139                                    return exec_err!("expected variant object");
140                                };
141
142                                let mut v = VariantBuilder::new();
143                                let mut o = v.new_object();
144
145                                o.extend(variant_object.iter());
146                                o.insert(key, value.clone());
147
148                                o.finish();
149
150                                Ok(v.finish())
151                            })
152                            .transpose()
153                    })
154                    .collect::<Result<Vec<_>>>()?;
155
156                let variants = buffers
157                    .iter()
158                    .map(|opt| {
159                        opt.as_ref()
160                            .map(|(m, v)| Variant::new(m.as_ref(), v.as_ref()))
161                    })
162                    .collect::<Vec<_>>();
163
164                let out: StructArray = VariantArray::from_iter(variants).into();
165
166                Ok(ColumnarValue::Array(Arc::new(out)))
167            }
168        }
169    }
170}
datafusion_variant/variant_object_insert.rs

datafusion_variant/
variant_object_insert.rs
