feat: implement variant_length UDF

src/cast_to_variant.rs

@@ -157,7 +157,6 @@ impl CastToVariantUdf {
 }
 
 impl ScalarUDFImpl for CastToVariantUdf {
-
     fn name(&self) -> &str {
         "cast_to_variant"
     }

src/is_variant_null.rs

@@ -27,7 +27,6 @@ impl Default for IsVariantNullUdf {
 }
 
 impl ScalarUDFImpl for IsVariantNullUdf {
-
     fn name(&self) -> &str {
         "is_variant_null"
     }

src/json_to_variant.rs

@@ -38,7 +38,6 @@ impl Default for JsonToVariantUdf {
 }
 
 impl ScalarUDFImpl for JsonToVariantUdf {
-
     fn name(&self) -> &str {
         "json_to_variant"
     }

src/lib.rs

@@ -8,6 +8,7 @@ mod is_variant_null;
 mod json_to_variant;
 mod variant_contains;
 mod variant_get;
+mod variant_length;
 mod variant_list_construct;
 mod variant_list_delete;
 mod variant_list_insert;
@@ -24,6 +25,7 @@ pub use is_variant_null::*;
 pub use json_to_variant::*;
 pub use variant_contains::*;
 pub use variant_get::*;
+pub use variant_length::*;
 pub use variant_list_construct::*;
 pub use variant_list_delete::*;
 pub use variant_list_insert::*;

src/variant_get.rs

@@ -335,7 +335,6 @@ impl Default for VariantGetUdf {
 }
 
 impl ScalarUDFImpl for VariantGetUdf {
-
     fn name(&self) -> &str {
         "variant_get"
     }
@@ -389,7 +388,6 @@ impl Default for VariantGetFieldUdf {
 }
 
 impl ScalarUDFImpl for VariantGetFieldUdf {
-
     fn name(&self) -> &str {
         "variant_get_field"
     }

src/variant_length.rs

@@ -0,0 +1,243 @@
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, UInt64Array};
+use arrow_schema::DataType;
+use datafusion::common::exec_datafusion_err;
+use datafusion::error::Result;
+use datafusion::logical_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
+};
+use datafusion::scalar::ScalarValue;
+use parquet_variant::Variant;
+use parquet_variant_compute::VariantArray;
+
+use crate::shared::{path_from_scalar, try_field_as_variant_array};
+
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct VariantLengthUdf {
+    signature: Signature,
+}
+
+impl Default for VariantLengthUdf {
+    fn default() -> Self {
+        Self {
+            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
+        }
+    }
+}
+
+/// Returns the number of elements in an Array or Object variant at the given path.
+/// Returns NULL if the variant at the path is not an Array or Object.
+fn variant_length(variant: Option<Variant<'_, '_>>) -> Option<u64> {
+    let variant = variant?;
+    match variant {
+        Variant::List(list) => Some(list.len() as u64),
+        Variant::Object(obj) => Some(obj.len() as u64),
+        _ => None,
+    }
+}
+
+impl ScalarUDFImpl for VariantLengthUdf {
+    fn name(&self) -> &str {
+        "variant_length"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::UInt64)
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let (variant_arg, path_arg) = match args.args.as_slice() {
+            [variant_arg, path_arg] => (variant_arg, path_arg),
+            _ => return datafusion::common::exec_err!("expected 2 arguments"),
+        };
+
+        let variant_field = args
+            .arg_fields
+            .first()
+            .ok_or_else(|| exec_datafusion_err!("expected argument field"))?;
+
+        try_field_as_variant_array(variant_field.as_ref())?;
+
+        match (variant_arg, path_arg) {
+            (ColumnarValue::Array(variant_array), ColumnarValue::Scalar(path_scalar)) => {
+                let path = path_from_scalar(path_scalar)?;
+                let variant_array = VariantArray::try_new(variant_array.as_ref())?;
+                let values = variant_array
+                    .iter()
+                    .map(|variant| {
+                        let at_path = variant.as_ref().and_then(|v| v.get_path(&path));
+                        variant_length(at_path)
+                    })
+                    .collect::<Vec<_>>();
+
+                Ok(ColumnarValue::Array(
+                    Arc::new(UInt64Array::from(values)) as ArrayRef
+                ))
+            }
+            (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Scalar(path_scalar)) => {
+                let ScalarValue::Struct(variant_array) = scalar_variant else {
+                    return datafusion::common::exec_err!("expected struct array");
+                };
+
+                let path = path_from_scalar(path_scalar)?;
+                let variant_array = VariantArray::try_new(variant_array.as_ref())?;
+                let variant = variant_array.iter().next().flatten();
+                let at_path = variant.as_ref().and_then(|v| v.get_path(&path));
+                let value = variant_length(at_path);
+
+                Ok(ColumnarValue::Scalar(ScalarValue::UInt64(value)))
+            }
+            _ => datafusion::common::exec_err!(
+                "unsupported argument combination: variant_length requires a scalar path"
+            ),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::array::{Array, ArrayRef, UInt64Array};
+    use arrow_schema::{DataType, Field, Fields};
+    use parquet_variant_compute::VariantType;
+
+    use crate::shared::{build_variant_array_from_json_array, variant_scalar_from_json};
+
+    use super::*;
+
+    fn arg_fields() -> Vec<Arc<Field>> {
+        vec![
+            Arc::new(
+                Field::new("input", DataType::Struct(Fields::empty()), true)
+                    .with_extension_type(VariantType),
+            ),
+            Arc::new(Field::new("path", DataType::Utf8, true)),
+        ]
+    }
+
+    #[test]
+    fn test_scalar_array_length() {
+        let udf = VariantLengthUdf::default();
+        let args = ScalarFunctionArgs {
+            args: vec![
+                ColumnarValue::Scalar(variant_scalar_from_json(serde_json::json!([1, 2, 3]))),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some("".to_string()))),
+            ],
+            return_field: Arc::new(Field::new("result", DataType::UInt64, true)),
+            arg_fields: arg_fields(),
+            number_rows: Default::default(),
+            config_options: Default::default(),
+        };
+
+        let result = udf.invoke_with_args(args).unwrap();
+        let ColumnarValue::Scalar(ScalarValue::UInt64(Some(value))) = result else {
+            panic!("expected u64 scalar")
+        };
+        assert_eq!(value, 3);
+    }
+
+    #[test]
+    fn test_scalar_object_length() {
+        let udf = VariantLengthUdf::default();
+        let args = ScalarFunctionArgs {
+            args: vec![
+                ColumnarValue::Scalar(variant_scalar_from_json(
+                    serde_json::json!({"a": 1, "b": 2}),
+                )),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some("".to_string()))),
+            ],
+            return_field: Arc::new(Field::new("result", DataType::UInt64, true)),
+            arg_fields: arg_fields(),
+            number_rows: Default::default(),
+            config_options: Default::default(),
+        };
+
+        let result = udf.invoke_with_args(args).unwrap();
+        let ColumnarValue::Scalar(ScalarValue::UInt64(Some(value))) = result else {
+            panic!("expected u64 scalar")
+        };
+        assert_eq!(value, 2);
+    }
+
+    #[test]
+    fn test_scalar_nested_path() {
+        let udf = VariantLengthUdf::default();
+        let args = ScalarFunctionArgs {
+            args: vec![
+                ColumnarValue::Scalar(variant_scalar_from_json(
+                    serde_json::json!({"a": [1, 2, 3, 4]}),
+                )),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))),
+            ],
+            return_field: Arc::new(Field::new("result", DataType::UInt64, true)),
+            arg_fields: arg_fields(),
+            number_rows: Default::default(),
+            config_options: Default::default(),
+        };
+
+        let result = udf.invoke_with_args(args).unwrap();
+        let ColumnarValue::Scalar(ScalarValue::UInt64(Some(value))) = result else {
+            panic!("expected u64 scalar")
+        };
+        assert_eq!(value, 4);
+    }
+
+    #[test]
+    fn test_scalar_primitive_returns_null() {
+        let udf = VariantLengthUdf::default();
+        let args = ScalarFunctionArgs {
+            args: vec![
+                ColumnarValue::Scalar(variant_scalar_from_json(serde_json::json!(42))),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some("".to_string()))),
+            ],
+            return_field: Arc::new(Field::new("result", DataType::UInt64, true)),
+            arg_fields: arg_fields(),
+            number_rows: Default::default(),
+            config_options: Default::default(),
+        };
+
+        let result = udf.invoke_with_args(args).unwrap();
+        let ColumnarValue::Scalar(ScalarValue::UInt64(value)) = result else {
+            panic!("expected u64 scalar")
+        };
+        assert_eq!(value, None);
+    }
+
+    #[test]
+    fn test_array_variants() {
+        let udf = VariantLengthUdf::default();
+        let input = build_variant_array_from_json_array(&[
+            Some(serde_json::json!([1, 2])),
+            Some(serde_json::json!({"a": 1, "b": 2, "c": 3})),
+            Some(serde_json::json!(42)),
+            None,
+        ]);
+        let args = ScalarFunctionArgs {
+            args: vec![
+                ColumnarValue::Array(Arc::new(arrow::array::StructArray::from(input)) as ArrayRef),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some("".to_string()))),
+            ],
+            return_field: Arc::new(Field::new("result", DataType::UInt64, true)),
+            arg_fields: arg_fields(),
+            number_rows: Default::default(),
+            config_options: Default::default(),
+        };
+
+        let result = udf.invoke_with_args(args).unwrap();
+        let ColumnarValue::Array(values) = result else {
+            panic!("expected array")
+        };
+
+        let values = values.as_any().downcast_ref::<UInt64Array>().unwrap();
+        assert_eq!(
+            values.into_iter().collect::<Vec<_>>(),
+            vec![Some(2), Some(3), None, None]
+        );
+    }
+}

src/variant_list_construct.rs

@@ -28,7 +28,6 @@ impl Default for VariantListConstruct {
 }
 
 impl ScalarUDFImpl for VariantListConstruct {
-
     fn name(&self) -> &str {
         "variant_list_construct"
     }

src/variant_list_delete.rs

@@ -67,7 +67,6 @@ fn delete_list_element(variant_list: Variant, index: usize) -> Result<(Vec<u8>,
 }
 
 impl ScalarUDFImpl for VariantListDelete {
-
     fn name(&self) -> &str {
         "variant_list_delete"
     }

src/variant_list_insert.rs

@@ -29,7 +29,6 @@ impl Default for VariantListInsert {
 }
 
 impl ScalarUDFImpl for VariantListInsert {
-
     fn name(&self) -> &str {
         "variant_list_insert"
     }

src/variant_normalize.rs

@@ -62,7 +62,6 @@ impl Default for VariantNormalizeUdf {
 }
 
 impl ScalarUDFImpl for VariantNormalizeUdf {
-
     fn name(&self) -> &str {
         "variant_normalize"
     }

src/variant_object_construct.rs

@@ -29,7 +29,6 @@ impl Default for VariantObjectConstruct {
 }
 
 impl ScalarUDFImpl for VariantObjectConstruct {
-
     fn name(&self) -> &str {
         "variant_object_construct"
     }

src/variant_object_delete.rs

@@ -29,7 +29,6 @@ impl Default for VariantObjectDelete {
 }
 
 impl ScalarUDFImpl for VariantObjectDelete {
-
     fn name(&self) -> &str {
         "variant_object_delete"
     }

src/variant_object_insert.rs

@@ -29,7 +29,6 @@ impl Default for VariantObjectInsert {
 }
 
 impl ScalarUDFImpl for VariantObjectInsert {
-
     fn name(&self) -> &str {
         "variant_object_insert"
     }

src/variant_object_keys.rs

@@ -45,7 +45,6 @@ fn append_keys_from_variant(v_opt: Option<Variant>, builder: &mut ListBuilder<St
 }
 
 impl ScalarUDFImpl for VariantObjectKeys {
-
     fn name(&self) -> &str {
         "variant_object_keys"
     }

src/variant_pretty.rs

@@ -28,7 +28,6 @@ impl Default for VariantPretty {
 }
 
 impl ScalarUDFImpl for VariantPretty {
-
     fn name(&self) -> &str {
         "variant_pretty"
     }

src/variant_to_json.rs

@@ -39,7 +39,6 @@ impl Default for VariantToJsonUdf {
 }
 
 impl ScalarUDFImpl for VariantToJsonUdf {
-
     fn name(&self) -> &str {
         "variant_to_json"
     }

tests/sqllogictests.rs

@@ -3,7 +3,7 @@ use datafusion_sqllogictest::{DataFusion, TestContext};
 use datafusion_variant::{
     CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantContainsUdf, VariantGetBoolUdf,
     VariantGetFieldUdf, VariantGetFloatUdf, VariantGetIntUdf, VariantGetJsonUdf, VariantGetStrUdf,
-    VariantGetUdf, VariantListConstruct, VariantListDelete, VariantListInsert,
+    VariantGetUdf, VariantLengthUdf, VariantListConstruct, VariantListDelete, VariantListInsert,
     VariantObjectConstruct, VariantObjectDelete, VariantObjectInsert, VariantObjectKeys,
     VariantPretty, VariantToJsonUdf,
 };
@@ -66,6 +66,7 @@ async fn run_sqllogictests() -> Result<(), Box<dyn std::error::Error>> {
         ctx.register_udf(ScalarUDF::new_from_impl(VariantObjectInsert::default()));
         ctx.register_udf(ScalarUDF::new_from_impl(VariantObjectDelete::default()));
         ctx.register_udf(ScalarUDF::new_from_impl(VariantObjectKeys::default()));
+        ctx.register_udf(ScalarUDF::new_from_impl(VariantLengthUdf::default()));
 
         let pb = ProgressBar::new(24);