Merge branch 'main' of https://github.com/datafusion-contrib/datafusion-variant into errors

Author: sdf-jkl

src/impl_variant_get.rs

@@ -12,12 +12,24 @@ macro_rules! impl_variant_get_typed {
         #[derive(Debug, Hash, PartialEq, Eq)]
         pub struct $struct_name {
             signature: Signature,
+            path_mode: crate::variant_get::PathMode,
         }
 
         impl Default for $struct_name {
             fn default() -> Self {
                 Self {
                     signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
+                    path_mode: crate::variant_get::PathMode::DotNotation,
+                }
+            }
+        }
+
+        impl $struct_name {
+            /// Create a new instance with the specified path mode.
+            pub fn with_path_mode(path_mode: crate::variant_get::PathMode) -> Self {
+                Self {
+                    signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
+                    path_mode,
                 }
             }
         }
@@ -40,7 +52,7 @@ macro_rules! impl_variant_get_typed {
             }
 
             fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
-                invoke_variant_get_typed(args, $scalar_from, $array_from, $extract)
+                invoke_variant_get_typed(args, $scalar_from, $array_from, $extract, self.path_mode)
             }
         }
     };

src/shared.rs

@@ -11,9 +11,11 @@ use datafusion::common::exec_datafusion_err;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs};
 use datafusion::{common::exec_err, scalar::ScalarValue};
-use parquet_variant::{Variant, VariantPath};
+use parquet_variant::Variant;
 use parquet_variant_compute::{VariantArray, VariantType};
 
+use crate::variant_get::PathMode;
+
 #[cfg(test)]
 use parquet_variant_compute::VariantArrayBuilder;
 
@@ -128,13 +130,14 @@ pub fn variant_get_single_value<T>(
     variant_array: &VariantArray,
     index: usize,
     path: &str,
+    path_mode: PathMode,
     extract: for<'m, 'v> fn(Variant<'m, 'v>) -> Result<Option<T>>,
 ) -> Result<Option<T>> {
     let Some(variant) = variant_array.iter().nth(index).flatten() else {
         return Ok(None);
     };
 
-    let variant_path = VariantPath::try_from(path)?;
+    let variant_path = path_mode.try_build_path(path)?;
     let Some(value) = variant.get_path(&variant_path) else {
         return Ok(None);
     };
@@ -145,9 +148,10 @@ pub fn variant_get_single_value<T>(
 pub fn variant_get_array_values<T>(
     variant_array: &VariantArray,
     path: &str,
+    path_mode: PathMode,
     extract: for<'m, 'v> fn(Variant<'m, 'v>) -> Result<Option<T>>,
 ) -> Result<Vec<Option<T>>> {
-    let variant_path = VariantPath::try_from(path)?;
+    let variant_path = path_mode.try_build_path(path)?;
 
     variant_array
         .iter()
@@ -170,6 +174,7 @@ pub fn invoke_variant_get_typed<T>(
     scalar_from_option: fn(Option<T>) -> ScalarValue,
     array_from_values: fn(Vec<Option<T>>) -> ArrayRef,
     extract: for<'m, 'v> fn(Variant<'m, 'v>) -> Result<Option<T>>,
+    path_mode: PathMode,
 ) -> Result<ColumnarValue> {
     let (variant_arg, path_arg) = match args.args.as_slice() {
         [variant_arg, path_arg] => (variant_arg, path_arg),
@@ -190,7 +195,7 @@ pub fn invoke_variant_get_typed<T>(
                 .unwrap_or_default();
 
             let variant_array = VariantArray::try_new(variant_array.as_ref())?;
-            let values = variant_get_array_values(&variant_array, path, extract)?;
+            let values = variant_get_array_values(&variant_array, path, path_mode, extract)?;
             ColumnarValue::Array(array_from_values(values))
         }
         (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Scalar(path_scalar)) => {
@@ -203,7 +208,7 @@ pub fn invoke_variant_get_typed<T>(
                 .unwrap_or_default();
 
             let variant_array = VariantArray::try_new(variant_array.as_ref())?;
-            let value = variant_get_single_value(&variant_array, 0, path, extract)?;
+            let value = variant_get_single_value(&variant_array, 0, path, path_mode, extract)?;
 
             ColumnarValue::Scalar(scalar_from_option(value))
         }
@@ -218,7 +223,7 @@ pub fn invoke_variant_get_typed<T>(
             let values: Vec<Option<T>> = (0..variant_array.len())
                 .map(|i| {
                     let path = paths[i].unwrap_or_default();
-                    variant_get_single_value(&variant_array, i, path, extract)
+                    variant_get_single_value(&variant_array, i, path, path_mode, extract)
                 })
                 .collect::<Result<_>>()?;
 
@@ -236,7 +241,7 @@ pub fn invoke_variant_get_typed<T>(
                 .iter()
                 .map(|path| {
                     let path = path.unwrap_or_default();
-                    variant_get_single_value(&variant_array, 0, path, extract)
+                    variant_get_single_value(&variant_array, 0, path, path_mode, extract)
                 })
                 .collect::<Result<_>>()?;
 

src/variant_get.rs

@@ -74,7 +74,7 @@ fn build_get_options<'a>(path: VariantPath<'a>, as_type: &Option<FieldRef>) -> G
 
 /// Determines how a string path is converted to a [`VariantPath`].
 #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
-enum PathMode {
+pub enum PathMode {
     /// Splits the path on `.` for dot-notation traversal (e.g., `"a.b.c"` → `["a", "b", "c"]`).
     DotNotation,
     /// Treats the entire path string as a single field name (e.g., `"a.b.c"` → `["a.b.c"]`).
@@ -84,7 +84,7 @@ enum PathMode {
 }
 
 impl PathMode {
-    fn try_build_path<'a>(&self, path: &'a str) -> Result<VariantPath<'a>> {
+    pub fn try_build_path<'a>(&self, path: &'a str) -> Result<VariantPath<'a>> {
         match self {
             PathMode::DotNotation => VariantPath::try_from(path).map_err(Into::into),
             PathMode::SingleField => Ok(VariantPath::new(vec![VariantPathElement::field(path)])),
@@ -1556,4 +1556,134 @@ mod tests {
         assert!(bool_arr.is_null(2));
         assert!(bool_arr.is_null(3));
     }
+
+    #[test]
+    fn test_get_str_with_single_field_mode_dotted_key() {
+        // VariantGetStrUdf with SingleField mode should treat dotted keys as a single field
+        let variant_input = variant_scalar_from_json(serde_json::json!({
+            "http.response.status_code": 200,
+            "service.name": "my-service"
+        }));
+
+        let udf = VariantGetStrUdf::with_path_mode(PathMode::SingleField);
+        let args = build_variant_get_args(
+            ColumnarValue::Scalar(variant_input),
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+                "http.response.status_code".to_string(),
+            ))),
+            DataType::Utf8View,
+            standard_variant_get_arg_fields(),
+        );
+
+        let result = udf.invoke_with_args(args).unwrap();
+
+        let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(s))) = result else {
+            panic!("expected Utf8View scalar, got {result:?}");
+        };
+        // Integer 200 should be JSON-serialized to "200"
+        assert_eq!(s, "200");
+    }
+
+    #[test]
+    fn test_get_str_with_single_field_mode_string_value() {
+        let variant_input = variant_scalar_from_json(serde_json::json!({
+            "service.name": "my-service"
+        }));
+
+        let udf = VariantGetStrUdf::with_path_mode(PathMode::SingleField);
+        let args = build_variant_get_args(
+            ColumnarValue::Scalar(variant_input),
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some("service.name".to_string()))),
+            DataType::Utf8View,
+            standard_variant_get_arg_fields(),
+        );
+
+        let result = udf.invoke_with_args(args).unwrap();
+
+        let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(s))) = result else {
+            panic!("expected Utf8View scalar, got {result:?}");
+        };
+        // String values returned as-is (no JSON quotes)
+        assert_eq!(s, "my-service");
+    }
+
+    #[test]
+    fn test_get_str_with_single_field_mode_array() {
+        // Test array input with SingleField mode
+        let json_rows = vec![
+            serde_json::json!({"http.status": 200, "http.method": "GET"}),
+            serde_json::json!({"http.status": 404, "http.method": "POST"}),
+            serde_json::json!({"http.status": 500}),
+        ];
+
+        let variant_array = variant_array_from_json_rows(&json_rows);
+
+        let udf = VariantGetStrUdf::with_path_mode(PathMode::SingleField);
+        let args = build_variant_get_args(
+            ColumnarValue::Array(variant_array),
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some("http.status".to_string()))),
+            DataType::Utf8View,
+            standard_variant_get_arg_fields(),
+        );
+
+        let result = udf.invoke_with_args(args).unwrap();
+
+        let ColumnarValue::Array(arr) = result else {
+            panic!("expected array output");
+        };
+        let str_arr = arr.as_any().downcast_ref::<StringViewArray>().unwrap();
+        assert_eq!(str_arr.len(), 3);
+        assert_eq!(str_arr.value(0), "200");
+        assert_eq!(str_arr.value(1), "404");
+        assert_eq!(str_arr.value(2), "500");
+    }
+
+    #[test]
+    fn test_get_int_with_single_field_mode() {
+        let variant_input = variant_scalar_from_json(serde_json::json!({
+            "http.status": 200
+        }));
+
+        let udf = VariantGetIntUdf::with_path_mode(PathMode::SingleField);
+        let args = build_variant_get_args(
+            ColumnarValue::Scalar(variant_input),
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some("http.status".to_string()))),
+            DataType::Int64,
+            standard_variant_get_arg_fields(),
+        );
+
+        let result = udf.invoke_with_args(args).unwrap();
+
+        let ColumnarValue::Scalar(ScalarValue::Int64(Some(v))) = result else {
+            panic!("expected Int64 scalar, got {result:?}");
+        };
+        assert_eq!(v, 200);
+    }
+
+    #[test]
+    fn test_get_str_dot_notation_splits_dotted_key() {
+        // With default DotNotation mode, dotted keys are split — should return NULL
+        // for a key like "http.response.status_code" that's stored as a single field
+        let variant_input = variant_scalar_from_json(serde_json::json!({
+            "http.response.status_code": 200
+        }));
+
+        let udf = VariantGetStrUdf::default(); // DotNotation
+        let args = build_variant_get_args(
+            ColumnarValue::Scalar(variant_input),
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+                "http.response.status_code".to_string(),
+            ))),
+            DataType::Utf8View,
+            standard_variant_get_arg_fields(),
+        );
+
+        let result = udf.invoke_with_args(args).unwrap();
+
+        // DotNotation splits on dots, tries to traverse http -> response -> status_code
+        // which doesn't exist, so returns NULL
+        let ColumnarValue::Scalar(ScalarValue::Utf8View(None)) = result else {
+            panic!("expected NULL (dot notation splits the key), got {result:?}");
+        };
+    }
 }