datafusion_variant/
variant_list_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_variant_scalar};
17
18#[derive(Debug, Hash, PartialEq, Eq)]
19pub struct VariantListDelete {
20    signature: Signature,
21}
22
23impl Default for VariantListDelete {
24    fn default() -> Self {
25        Self {
26            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
27        }
28    }
29}
30
31// just a helper to parse index elements
32fn try_parse_index_scalar(scalar: &ScalarValue) -> Result<usize> {
33    match scalar {
34        ScalarValue::Int8(Some(v)) => Ok(*v as usize),
35        ScalarValue::Int16(Some(v)) => Ok(*v as usize),
36        ScalarValue::Int32(Some(v)) => Ok(*v as usize),
37        ScalarValue::Int64(Some(v)) => Ok(*v as usize),
38        _ => exec_err!("expected non-null integer scalar for index"),
39    }
40}
41
42fn delete_list_element(variant_list: Variant, index: usize) -> Result<(Vec<u8>, Vec<u8>)> {
43    let Variant::List(variant_list) = variant_list else {
44        return exec_err!("expected variant list");
45    };
46
47    if index >= variant_list.len() {
48        return exec_err!(
49            "index {} oob for list of length {}",
50            index,
51            variant_list.len()
52        );
53    }
54
55    let mut v = VariantBuilder::new();
56    let mut l = v.new_list();
57
58    for (i, val) in variant_list.iter().enumerate() {
59        if i != index {
60            l.append_value(val);
61        }
62    }
63
64    l.finish();
65
66    Ok(v.finish())
67}
68
69impl ScalarUDFImpl for VariantListDelete {
70    fn as_any(&self) -> &dyn std::any::Any {
71        self
72    }
73
74    fn name(&self) -> &str {
75        "variant_list_delete"
76    }
77
78    fn signature(&self) -> &Signature {
79        &self.signature
80    }
81
82    fn return_type(&self, _arg_types: &[arrow_schema::DataType]) -> Result<arrow_schema::DataType> {
83        Err(DataFusionError::Internal(
84            "implemented return_field_from_args instead".into(),
85        ))
86    }
87
88    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Arc<Field>> {
89        let data_type = DataType::Struct(Fields::from(vec![
90            Field::new("metadata", DataType::BinaryView, false),
91            Field::new("value", DataType::BinaryView, false),
92        ]));
93
94        Ok(Arc::new(
95            Field::new(self.name(), data_type, true).with_extension_type(VariantType),
96        ))
97    }
98
99    fn invoke_with_args(
100        &self,
101        args: datafusion::logical_expr::ScalarFunctionArgs,
102    ) -> Result<ColumnarValue> {
103        let argument_fields = args.arg_fields;
104        let argument_values = args.args;
105
106        ensure(
107            argument_fields.len() == argument_values.len(),
108            "argument fields and values must be of same length",
109        )?;
110
111        let [variant_list_to_update, index_to_delete] = argument_values.as_slice() else {
112            return exec_err!("expected 2 arguments");
113        };
114
115        ensure(
116            matches!(argument_fields[0].extension_type(), VariantType),
117            "expected extension type of VariantType for variant list argument",
118        )?;
119
120        let index = {
121            let ColumnarValue::Scalar(index) = index_to_delete else {
122                return exec_err!("expected scalar value for index");
123            };
124
125            try_parse_index_scalar(index)?
126        };
127
128        match variant_list_to_update {
129            ColumnarValue::Scalar(scalar_variant_list) => {
130                let variant_list = try_parse_variant_scalar(scalar_variant_list)?;
131                let variant_list = variant_list.value(0);
132
133                let (m, v) = delete_list_element(variant_list, index)?;
134                // safety: delete_list_element uses the list builder to construct a valid variant ahead of time
135                let v = Variant::new(m.as_ref(), v.as_ref());
136
137                let out: StructArray = VariantArray::from_iter([v]).into();
138
139                Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(out))))
140            }
141            ColumnarValue::Array(variant_lists) => {
142                let variant_array = VariantArray::try_new(&variant_lists)?;
143
144                let buffers = variant_array
145                    .iter()
146                    .map(|v_opt| {
147                        v_opt
148                            .map(|variant_list| delete_list_element(variant_list, index))
149                            .transpose()
150                    })
151                    .collect::<Result<Vec<_>>>()?;
152
153                let variants = buffers
154                    .iter()
155                    .map(|opt| {
156                        opt.as_ref()
157                            .map(|(m, v)| Variant::new(m.as_ref(), v.as_ref()))
158                    })
159                    .collect::<Vec<_>>();
160
161                let out: StructArray = VariantArray::from_iter(variants).into();
162
163                Ok(ColumnarValue::Array(Arc::new(out)))
164            }
165        }
166    }
167}
datafusion_variant/variant_list_delete.rs

datafusion_variant/
variant_list_delete.rs
