From b2eb3a25a6c3bbb27cbdc7444ac40822d62aa338 Mon Sep 17 00:00:00 2001 From: Namgung Chan <9511chn@gmail.com> Date: Sun, 12 Apr 2026 03:05:21 +0900 Subject: [PATCH 1/3] Add variant_contains UDF --- .gitignore | 3 +- examples/cli.rs | 7 +- src/lib.rs | 2 + src/shared.rs | 2 +- src/variant_contains.rs | 280 ++++++++++++++++++++++++++ tests/sqllogictests.rs | 10 +- tests/test_files/variant_contains.slt | 60 ++++++ 7 files changed, 355 insertions(+), 9 deletions(-) create mode 100644 src/variant_contains.rs create mode 100644 tests/test_files/variant_contains.slt diff --git a/.gitignore b/.gitignore index d122bfe..5a32fc5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ /target /data/* -profile.json.gz \ No newline at end of file +profile.json.gz +.idea \ No newline at end of file diff --git a/examples/cli.rs b/examples/cli.rs index 77e0566..1f9672e 100644 --- a/examples/cli.rs +++ b/examples/cli.rs @@ -4,9 +4,9 @@ use arrow::datatypes::{DataType, Field, Schema}; use datafusion::logical_expr::ScalarUDF; use datafusion::prelude::*; use datafusion_variant::{ - CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantGetUdf, VariantListConstruct, - VariantListInsert, VariantObjectConstruct, VariantObjectInsert, VariantObjectKeys, - VariantPretty, VariantToJsonUdf, + CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantContainsUdf, VariantGetUdf, + VariantListConstruct, VariantListInsert, VariantObjectConstruct, VariantObjectInsert, + VariantObjectKeys, VariantPretty, VariantToJsonUdf, }; use flate2::read::GzDecoder; use rustyline::error::ReadlineError; @@ -112,6 +112,7 @@ async fn main() -> Result<()> { ctx.register_udf(ScalarUDF::new_from_impl(JsonToVariantUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(CastToVariantUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(IsVariantNullUdf::default())); + ctx.register_udf(ScalarUDF::new_from_impl(VariantContainsUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(VariantGetUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(VariantPretty::default())); ctx.register_udf(ScalarUDF::new_from_impl(VariantObjectConstruct::default())); diff --git a/src/lib.rs b/src/lib.rs index f944f84..103c5db 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ mod impl_variant_get; mod is_variant_null; mod json_to_variant; mod variant_get; +mod variant_contains; mod variant_list_construct; mod variant_list_delete; mod variant_list_insert; @@ -22,6 +23,7 @@ pub use cast_to_variant::*; pub use is_variant_null::*; pub use json_to_variant::*; pub use variant_get::*; +pub use variant_contains::*; pub use variant_list_construct::*; pub use variant_list_delete::*; pub use variant_list_insert::*; diff --git a/src/shared.rs b/src/shared.rs index b649b04..7de604b 100644 --- a/src/shared.rs +++ b/src/shared.rs @@ -168,7 +168,7 @@ pub fn variant_get_array_values( /// - **List** scalars treat each element as a single field name /// (e.g. `['a.b', 'c']` → path `[a.b, c]`), which is critical for keys that /// contain dots such as OTEL attribute keys like `http.response.status_code`. -fn path_from_scalar(scalar: &ScalarValue) -> Result> { +pub(crate) fn path_from_scalar(scalar: &ScalarValue) -> Result> { match scalar { ScalarValue::Utf8(Some(s)) | ScalarValue::Utf8View(Some(s)) diff --git a/src/variant_contains.rs b/src/variant_contains.rs new file mode 100644 index 0000000..05b153c --- /dev/null +++ b/src/variant_contains.rs @@ -0,0 +1,280 @@ +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray}; +use arrow_schema::DataType; +use datafusion::common::{exec_datafusion_err, exec_err}; +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; +use datafusion::scalar::ScalarValue; +use parquet_variant::{Variant, VariantPath}; +use parquet_variant_compute::VariantArray; + +use crate::shared::{path_from_scalar, try_field_as_variant_array, try_parse_string_columnar}; + +#[derive(Debug, Hash, PartialEq, Eq)] +pub struct VariantContainsUdf { + signature: Signature, +} + +impl Default for VariantContainsUdf { + fn default() -> Self { + Self { + signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable), + } + } +} + +fn variant_contains( + variant: Option<&Variant<'_, '_>>, + path: &VariantPath<'_>, +) -> Option { + variant.map(|value| value.get_path(path).is_some()) +} + +impl ScalarUDFImpl for VariantContainsUdf { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn name(&self) -> &str { + "variant_contains" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let (variant_arg, path_arg) = match args.args.as_slice() { + [variant_arg, path_arg] => (variant_arg, path_arg), + _ => return 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)) => { + if path_scalar.is_null() { + return exec_err!("path argument must be non-null"); + } + + let path = path_from_scalar(path_scalar)?; + let variant_array = VariantArray::try_new(variant_array.as_ref())?; + let values = variant_array + .iter() + .map(|variant| variant_contains(variant.as_ref(), &path)) + .collect::>(); + + Ok(ColumnarValue::Array( + Arc::new(BooleanArray::from(values)) as ArrayRef, + )) + } + (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Scalar(path_scalar)) => { + let ScalarValue::Struct(variant_array) = scalar_variant else { + return exec_err!("expected struct array"); + }; + + if path_scalar.is_null() { + return exec_err!("path argument must be non-null"); + } + + 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 value = variant_contains(variant.as_ref(), &path); + + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(value))) + } + (ColumnarValue::Array(variant_array), ColumnarValue::Array(paths)) => { + if variant_array.len() != paths.len() { + return exec_err!("expected variant array and paths to be of same length"); + } + + let variant_array = VariantArray::try_new(variant_array.as_ref())?; + let paths = try_parse_string_columnar(paths)?; + + let values = variant_array + .iter() + .zip(paths) + .map(|(maybe_variant, path_str)| { + let path_str = path_str.ok_or_else(|| { + exec_datafusion_err!("path argument must be non-null") + })?; + let path = VariantPath::try_from(path_str) + .map_err(Into::::into)?; + + Ok(variant_contains(maybe_variant.as_ref(), &path)) + }) + .collect::>>()?; + + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(values)) as ArrayRef)) + } + (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Array(paths)) => { + let ScalarValue::Struct(variant_array) = scalar_variant else { + return exec_err!("expected struct array"); + }; + + let variant_array = VariantArray::try_new(variant_array.as_ref())?; + let variant = variant_array.iter().next().flatten(); + let paths = try_parse_string_columnar(paths)?; + + let values = paths + .into_iter() + .map(|path_str| { + let path_str = path_str.ok_or_else(|| { + exec_datafusion_err!("path argument must be non-null") + })?; + let path = VariantPath::try_from(path_str) + .map_err(Into::::into)?; + + Ok(variant_contains(variant.as_ref(), &path)) + }) + .collect::>>()?; + + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(values)) as ArrayRef)) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Array, ArrayRef, BooleanArray, StringArray}; + use arrow_schema::{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> { + 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_existing_path_returns_true() { + let udf = VariantContainsUdf::default(); + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(variant_scalar_from_json(serde_json::json!({ + "a": {"b": null} + }))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("a.b".to_string()))), + ], + return_field: Arc::new(Field::new("result", DataType::Boolean, 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::Boolean(Some(value))) = result else { + panic!("expected boolean scalar") + }; + + assert!(value); + } + + #[test] + fn test_scalar_missing_path_returns_false() { + let udf = VariantContainsUdf::default(); + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(variant_scalar_from_json(serde_json::json!({ + "a": 1 + }))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("a.b".to_string()))), + ], + return_field: Arc::new(Field::new("result", DataType::Boolean, 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::Boolean(Some(value))) = result else { + panic!("expected boolean scalar") + }; + + assert!(!value); + } + + #[test] + fn test_array_paths_and_null_variant() { + let udf = VariantContainsUdf::default(); + let input = build_variant_array_from_json_array(&[ + Some(serde_json::json!({"a": 1})), + Some(serde_json::json!({"a": null})), + None, + ]); + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(arrow::array::StructArray::from(input)) as ArrayRef), + ColumnarValue::Array(Arc::new(StringArray::from(vec![ + Some("a"), + Some("a"), + Some("a"), + ])) as ArrayRef), + ], + return_field: Arc::new(Field::new("result", DataType::Boolean, 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 boolean array") + }; + + let values = values.as_any().downcast_ref::().unwrap(); + assert_eq!(values.into_iter().collect::>(), vec![Some(true), Some(true), None]); + } + + #[test] + fn test_array_variant_scalar_path() { + let udf = VariantContainsUdf::default(); + let input = build_variant_array_from_json_array(&[ + Some(serde_json::json!({"a": 1})), + Some(serde_json::json!({"b": 1})), + ]); + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(arrow::array::StructArray::from(input)) as ArrayRef), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))), + ], + return_field: Arc::new(Field::new("result", DataType::Boolean, 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 boolean array") + }; + + let values = values.as_any().downcast_ref::().unwrap(); + assert_eq!(values.into_iter().collect::>(), vec![Some(true), Some(false)]); + } + +} diff --git a/tests/sqllogictests.rs b/tests/sqllogictests.rs index f808d8c..bc0c224 100644 --- a/tests/sqllogictests.rs +++ b/tests/sqllogictests.rs @@ -1,10 +1,11 @@ use datafusion::{logical_expr::ScalarUDF, prelude::*}; use datafusion_sqllogictest::{DataFusion, TestContext}; use datafusion_variant::{ - CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantGetBoolUdf, VariantGetFieldUdf, - VariantGetFloatUdf, VariantGetIntUdf, VariantGetJsonUdf, VariantGetStrUdf, VariantGetUdf, - VariantListConstruct, VariantListDelete, VariantListInsert, VariantObjectConstruct, - VariantObjectDelete, VariantObjectInsert, VariantObjectKeys, VariantPretty, VariantToJsonUdf, + CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantContainsUdf, + VariantGetBoolUdf, VariantGetFieldUdf, VariantGetFloatUdf, VariantGetIntUdf, + VariantGetJsonUdf, VariantGetStrUdf, VariantGetUdf, VariantListConstruct, + VariantListDelete, VariantListInsert, VariantObjectConstruct, VariantObjectDelete, + VariantObjectInsert, VariantObjectKeys, VariantPretty, VariantToJsonUdf, }; use indicatif::ProgressBar; use sqllogictest::strict_column_validator; @@ -49,6 +50,7 @@ async fn run_sqllogictests() -> Result<(), Box> { ctx.register_udf(ScalarUDF::new_from_impl(JsonToVariantUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(CastToVariantUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(IsVariantNullUdf::default())); + ctx.register_udf(ScalarUDF::new_from_impl(VariantContainsUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(VariantGetUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(VariantGetStrUdf::default())); ctx.register_udf(ScalarUDF::new_from_impl(VariantGetFloatUdf::default())); diff --git a/tests/test_files/variant_contains.slt b/tests/test_files/variant_contains.slt new file mode 100644 index 0000000..ec0e247 --- /dev/null +++ b/tests/test_files/variant_contains.slt @@ -0,0 +1,60 @@ +# tests the variant_contains udf + +query BB +SELECT variant_contains(json_to_variant('{"a": 1, "b": null}'), 'a'), + variant_contains(json_to_variant('{"a": 1, "b": null}'), 'b'); +---- +true true + +query B +SELECT variant_contains(json_to_variant('{"a": 1}'), 'b'); +---- +false + +query B +SELECT variant_contains(json_to_variant('{"a": {"b": 1}}'), 'a.b'); +---- +true + +query B +SELECT variant_contains(json_to_variant('{"a": {"b": 1}}'), 'a.c'); +---- +false + +query B +SELECT variant_contains(json_to_variant('[10, 20, 30]'), '[1]'); +---- +true + +query B +SELECT variant_contains(json_to_variant('[10, 20, 30]'), '[3]'); +---- +false + +statement error DataFusion error: Execution error: path argument must be non-null +SELECT variant_contains(json_to_variant('{"a": 1}'), CAST(NULL AS TEXT)); + +statement ok +create table contains_paths(json_data text, path text) as values + ('{"a": 1}', 'a'), + ('{"a": 1}', null); + +statement error DataFusion error: Execution error: path argument must be non-null +select variant_contains(json_to_variant(json_data), path) +from contains_paths; + +statement ok +create table contains_data(json_data text) as values + ('{"a": 1}'), + ('{"b": 2}'), + (null); + +query BBB +select json_data is null, + variant_contains(json_to_variant(json_data), 'a') is null, + variant_contains(json_to_variant(json_data), 'a') +from contains_data; +---- +false false true +false false false +true true NULL From 16d38d2876a5644caa8ee70ef0f87830af4d8a81 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 17 Apr 2026 03:32:12 +0000 Subject: [PATCH 2/3] chore: remove .idea from gitignore Agent-Logs-Url: https://github.com/getChan/datafusion-variant/sessions/4c154f78-d449-4fc1-b0da-6c643e55e364 Co-authored-by: getChan <33323415+getChan@users.noreply.github.com> --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index 5a32fc5..068023a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,3 @@ /target /data/* profile.json.gz -.idea \ No newline at end of file From 4e2fb8f47dff12c97d5cbec675dee4f2d237ad1f Mon Sep 17 00:00:00 2001 From: Namgung Chan <9511chn@gmail.com> Date: Tue, 5 May 2026 23:23:14 +0900 Subject: [PATCH 3/3] cargo fmt --- src/lib.rs | 4 ++-- src/variant_contains.rs | 26 ++++++++++++++++---------- tests/sqllogictests.rs | 10 +++++----- 3 files changed, 23 insertions(+), 17 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 103c5db..6aecc5e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,8 +6,8 @@ mod cast_to_variant; mod impl_variant_get; mod is_variant_null; mod json_to_variant; -mod variant_get; mod variant_contains; +mod variant_get; mod variant_list_construct; mod variant_list_delete; mod variant_list_insert; @@ -22,8 +22,8 @@ mod variant_to_json; pub use cast_to_variant::*; pub use is_variant_null::*; pub use json_to_variant::*; -pub use variant_get::*; pub use variant_contains::*; +pub use variant_get::*; pub use variant_list_construct::*; pub use variant_list_delete::*; pub use variant_list_insert::*; diff --git a/src/variant_contains.rs b/src/variant_contains.rs index 05b153c..44108e5 100644 --- a/src/variant_contains.rs +++ b/src/variant_contains.rs @@ -26,10 +26,7 @@ impl Default for VariantContainsUdf { } } -fn variant_contains( - variant: Option<&Variant<'_, '_>>, - path: &VariantPath<'_>, -) -> Option { +fn variant_contains(variant: Option<&Variant<'_, '_>>, path: &VariantPath<'_>) -> Option { variant.map(|value| value.get_path(path).is_some()) } @@ -77,7 +74,7 @@ impl ScalarUDFImpl for VariantContainsUdf { .collect::>(); Ok(ColumnarValue::Array( - Arc::new(BooleanArray::from(values)) as ArrayRef, + Arc::new(BooleanArray::from(values)) as ArrayRef )) } (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Scalar(path_scalar)) => { @@ -118,7 +115,9 @@ impl ScalarUDFImpl for VariantContainsUdf { }) .collect::>>()?; - Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(values)) as ArrayRef)) + Ok(ColumnarValue::Array( + Arc::new(BooleanArray::from(values)) as ArrayRef + )) } (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Array(paths)) => { let ScalarValue::Struct(variant_array) = scalar_variant else { @@ -142,7 +141,9 @@ impl ScalarUDFImpl for VariantContainsUdf { }) .collect::>>()?; - Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(values)) as ArrayRef)) + Ok(ColumnarValue::Array( + Arc::new(BooleanArray::from(values)) as ArrayRef + )) } } } @@ -247,7 +248,10 @@ mod tests { }; let values = values.as_any().downcast_ref::().unwrap(); - assert_eq!(values.into_iter().collect::>(), vec![Some(true), Some(true), None]); + assert_eq!( + values.into_iter().collect::>(), + vec![Some(true), Some(true), None] + ); } #[test] @@ -274,7 +278,9 @@ mod tests { }; let values = values.as_any().downcast_ref::().unwrap(); - assert_eq!(values.into_iter().collect::>(), vec![Some(true), Some(false)]); + assert_eq!( + values.into_iter().collect::>(), + vec![Some(true), Some(false)] + ); } - } diff --git a/tests/sqllogictests.rs b/tests/sqllogictests.rs index bc0c224..0f37831 100644 --- a/tests/sqllogictests.rs +++ b/tests/sqllogictests.rs @@ -1,11 +1,11 @@ use datafusion::{logical_expr::ScalarUDF, prelude::*}; use datafusion_sqllogictest::{DataFusion, TestContext}; use datafusion_variant::{ - CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantContainsUdf, - VariantGetBoolUdf, VariantGetFieldUdf, VariantGetFloatUdf, VariantGetIntUdf, - VariantGetJsonUdf, VariantGetStrUdf, VariantGetUdf, VariantListConstruct, - VariantListDelete, VariantListInsert, VariantObjectConstruct, VariantObjectDelete, - VariantObjectInsert, VariantObjectKeys, VariantPretty, VariantToJsonUdf, + CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantContainsUdf, VariantGetBoolUdf, + VariantGetFieldUdf, VariantGetFloatUdf, VariantGetIntUdf, VariantGetJsonUdf, VariantGetStrUdf, + VariantGetUdf, VariantListConstruct, VariantListDelete, VariantListInsert, + VariantObjectConstruct, VariantObjectDelete, VariantObjectInsert, VariantObjectKeys, + VariantPretty, VariantToJsonUdf, }; use indicatif::ProgressBar; use sqllogictest::strict_column_validator;