datafusion_variant/
is_variant_null.rs

1use std::sync::Arc;
2
3use arrow::array::{ArrayRef, BooleanArray};
4use arrow_schema::DataType;
5use datafusion::common::{exec_datafusion_err, exec_err};
6use datafusion::error::Result;
7use datafusion::logical_expr::{
8    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
9};
10use datafusion::scalar::ScalarValue;
11use parquet_variant::Variant;
12use parquet_variant_compute::VariantArray;
13
14use crate::shared::{try_field_as_variant_array, try_parse_variant_scalar};
15
16#[derive(Debug, Hash, PartialEq, Eq)]
17pub struct IsVariantNullUdf {
18    signature: Signature,
19}
20
21impl Default for IsVariantNullUdf {
22    fn default() -> Self {
23        Self {
24            signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable),
25        }
26    }
27}
28
29impl ScalarUDFImpl for IsVariantNullUdf {
30    fn as_any(&self) -> &dyn std::any::Any {
31        self
32    }
33
34    fn name(&self) -> &str {
35        "is_variant_null"
36    }
37
38    fn signature(&self) -> &Signature {
39        &self.signature
40    }
41
42    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
43        Ok(DataType::Boolean)
44    }
45
46    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
47        let variant_field = args
48            .arg_fields
49            .first()
50            .ok_or_else(|| exec_datafusion_err!("expected 1 argument field type"))?;
51
52        try_field_as_variant_array(variant_field.as_ref())?;
53
54        let [variant_arg] = args.args.as_slice() else {
55            return exec_err!("expected 1 argument");
56        };
57
58        let out = match variant_arg {
59            ColumnarValue::Scalar(scalar_variant) => {
60                let variant_array = try_parse_variant_scalar(scalar_variant)?;
61                let variant = variant_array.value(0);
62                let is_variant_null = variant == Variant::Null;
63
64                ColumnarValue::Scalar(ScalarValue::Boolean(Some(is_variant_null)))
65            }
66            ColumnarValue::Array(variant_array) => {
67                let variant_array = VariantArray::try_new(variant_array.as_ref())?;
68
69                let out: BooleanArray = variant_array
70                    .iter()
71                    .map(|v| v.map(|v| v == Variant::Null))
72                    .collect::<Vec<_>>()
73                    .into();
74
75                ColumnarValue::Array(Arc::new(out) as ArrayRef)
76            }
77        };
78
79        Ok(out)
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use arrow::array::StructArray;
86    use arrow_schema::{Field, Fields};
87    use parquet_variant_compute::VariantType;
88
89    use crate::shared::{build_variant_array_from_json, build_variant_array_from_json_array};
90
91    use super::*;
92
93    #[test]
94    fn test_scalar() {
95        let expected_json = serde_json::json!(null);
96        let input = build_variant_array_from_json(&expected_json);
97
98        let variant_input = ScalarValue::Struct(Arc::new(input.into()));
99
100        let udf = IsVariantNullUdf::default();
101        let return_field = Arc::new(Field::new("result", DataType::Utf8View, true));
102        let arg_field = Arc::new(
103            Field::new("input", DataType::Struct(Fields::empty()), true)
104                .with_extension_type(VariantType),
105        );
106
107        let args = ScalarFunctionArgs {
108            args: vec![ColumnarValue::Scalar(variant_input)],
109            return_field,
110            arg_fields: vec![arg_field],
111            number_rows: Default::default(),
112            config_options: Default::default(),
113        };
114
115        let result = udf.invoke_with_args(args).unwrap();
116
117        let ColumnarValue::Scalar(ScalarValue::Boolean(Some(b))) = result else {
118            panic!("expected valid bool")
119        };
120
121        assert!(b);
122    }
123
124    #[test]
125    fn test_columnar() {
126        let input = build_variant_array_from_json_array(&[
127            Some(serde_json::json!(null)),
128            Some(serde_json::json!(null)),
129            Some(serde_json::json!("null")), // this is a ShortString('null')
130            Some(serde_json::json!({
131                "name": "norm"
132            })),
133        ]);
134
135        let input: StructArray = input.into();
136
137        let variant_input = Arc::new(input) as ArrayRef;
138
139        let udf = IsVariantNullUdf::default();
140        let return_field = Arc::new(Field::new("result", DataType::Utf8View, true));
141        let arg_field = Arc::new(
142            Field::new("input", DataType::Struct(Fields::empty()), true)
143                .with_extension_type(VariantType),
144        );
145
146        let args = ScalarFunctionArgs {
147            args: vec![ColumnarValue::Array(variant_input)],
148            return_field,
149            arg_fields: vec![arg_field],
150            number_rows: Default::default(),
151            config_options: Default::default(),
152        };
153
154        let result = udf.invoke_with_args(args).unwrap();
155
156        let ColumnarValue::Array(boolean_array) = result else {
157            panic!("expected valid bool")
158        };
159
160        let bool_array = boolean_array
161            .as_any()
162            .downcast_ref::<BooleanArray>()
163            .unwrap();
164
165        let bool_array = bool_array.into_iter().collect::<Vec<_>>();
166
167        assert_eq!(
168            bool_array,
169            vec![Some(true), Some(true), Some(false), Some(false),]
170        );
171    }
172}
datafusion_variant/is_variant_null.rs

datafusion_variant/
is_variant_null.rs
