datafusion_variant/
variant_object_delete.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 VariantObjectDelete {
20    signature: Signature,
21}
22
23impl Default for VariantObjectDelete {
24    fn default() -> Self {
25        Self {
26            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
27        }
28    }
29}
30
31impl ScalarUDFImpl for VariantObjectDelete {
32    fn as_any(&self) -> &dyn std::any::Any {
33        self
34    }
35
36    fn name(&self) -> &str {
37        "variant_object_delete"
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_to_delete] = argument_values.as_slice() else {
74            return exec_err!("expected 2 arguments");
75        };
76
77        ensure(
78            matches!(argument_fields[0].extension_type(), VariantType),
79            "expected extension type of VariantType for variant object argument",
80        )?;
81
82        let key = {
83            let ColumnarValue::Scalar(key) = key_to_delete else {
84                return exec_err!("expected scalar value for key");
85            };
86
87            try_parse_string_scalar(key)?
88                .ok_or_else(|| DataFusionError::Execution("expected non null string".into()))?
89        };
90
91        match variant_object_to_update {
92            ColumnarValue::Scalar(scalar_variant_object_to_update) => {
93                let variant_object = try_parse_variant_scalar(scalar_variant_object_to_update)?;
94                let variant_object = variant_object.value(0);
95                let Variant::Object(variant_object) = variant_object else {
96                    return exec_err!("expected variant object");
97                };
98
99                let mut v = VariantBuilder::new();
100                let mut o = v.new_object();
101
102                for (k, val) in variant_object.iter() {
103                    if k != key.as_str() {
104                        o.insert(k, val);
105                    }
106                }
107
108                o.finish();
109
110                let (m, v) = v.finish();
111                let v = Variant::try_new(m.as_ref(), v.as_ref())?;
112
113                let out: StructArray = VariantArray::from_iter([v]).into();
114
115                Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(out))))
116            }
117            ColumnarValue::Array(variant_objects_to_update) => {
118                let variant_objects = VariantArray::try_new(&variant_objects_to_update)?;
119
120                let buffers = variant_objects
121                    .iter()
122                    .map(|v_opt| {
123                        v_opt
124                            .map(|variant_object| {
125                                let Variant::Object(variant_object) = variant_object else {
126                                    return exec_err!("expected variant object");
127                                };
128
129                                let mut v = VariantBuilder::new();
130                                let mut o = v.new_object();
131
132                                for (k, val) in variant_object.iter() {
133                                    if k != key.as_str() {
134                                        o.insert(k, val);
135                                    }
136                                }
137
138                                o.finish();
139
140                                Ok(v.finish())
141                            })
142                            .transpose()
143                    })
144                    .collect::<Result<Vec<_>>>()?;
145
146                let variants = buffers
147                    .iter()
148                    .map(|opt| {
149                        opt.as_ref()
150                            .map(|(m, v)| Variant::new(m.as_ref(), v.as_ref()))
151                    })
152                    .collect::<Vec<_>>();
153
154                let out: StructArray = VariantArray::from_iter(variants).into();
155
156                Ok(ColumnarValue::Array(Arc::new(out)))
157            }
158        }
159    }
160}
datafusion_variant/variant_object_delete.rs

datafusion_variant/
variant_object_delete.rs
