Fix typed UDFs (variant_get_str, etc.) to handle shredded variants

Use variant_get kernel for path navigation in invoke_variant_get_typed
instead of VariantArray::iter() + Variant::get_path(). The iter() API
panics on shredded variant arrays with nested typed_value structs.
The variant_get kernel correctly navigates shredded structures and
returns simple sub-variant arrays that iter() can handle.

Author: friendlymatthew

src/shared.rs

@@ -12,7 +12,7 @@ use datafusion::error::Result;
 use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs};
 use datafusion::{common::exec_err, scalar::ScalarValue};
 use parquet_variant::Variant;
-use parquet_variant_compute::{VariantArray, VariantType};
+use parquet_variant_compute::{GetOptions, VariantArray, VariantType, variant_get};
 
 use crate::variant_get::PathMode;
 
@@ -126,49 +126,43 @@ pub fn try_parse_string_columnar(array: &Arc<dyn Array>) -> Result<Vec<Option<&s
     Err(exec_datafusion_err!("expected string array"))
 }
 
-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 = path_mode.try_build_path(path)?;
-    let Some(value) = variant.get_path(&variant_path) else {
-        return Ok(None);
-    };
-
-    extract(value)
-}
-
-pub fn variant_get_array_values<T>(
-    variant_array: &VariantArray,
+/// Navigate to a path within a variant array and extract values.
+///
+/// Uses `variant_get` for path navigation (which correctly handles shredded variants),
+/// then iterates the resulting simple sub-variant array to extract typed values.
+pub fn variant_get_and_extract<T>(
+    input: &ArrayRef,
     path: &str,
     path_mode: PathMode,
     extract: for<'m, 'v> fn(Variant<'m, 'v>) -> Result<Option<T>>,
 ) -> Result<Vec<Option<T>>> {
     let variant_path = path_mode.try_build_path(path)?;
+    let options = GetOptions::new_with_path(variant_path);
+    let sub_variant_arr = variant_get(input, options)?;
 
-    variant_array
+    let sub_variant = VariantArray::try_new(sub_variant_arr.as_ref())?;
+    sub_variant
         .iter()
         .map(|maybe_variant| {
             let Some(variant) = maybe_variant else {
                 return Ok(None);
             };
-
-            let Some(value) = variant.get_path(&variant_path) else {
-                return Ok(None);
-            };
-
-            extract(value)
+            extract(variant)
         })
         .collect()
 }
 
+/// Navigate to a path within a single-row variant and extract a value.
+pub fn variant_get_and_extract_single<T>(
+    input: &ArrayRef,
+    path: &str,
+    path_mode: PathMode,
+    extract: for<'m, 'v> fn(Variant<'m, 'v>) -> Result<Option<T>>,
+) -> Result<Option<T>> {
+    let values = variant_get_and_extract(input, path, path_mode, extract)?;
+    Ok(values.into_iter().next().flatten())
+}
+
 pub fn invoke_variant_get_typed<T>(
     args: ScalarFunctionArgs,
     scalar_from_option: fn(Option<T>) -> ScalarValue,
@@ -194,8 +188,7 @@ pub fn invoke_variant_get_typed<T>(
                 .map(|s| s.as_str())
                 .unwrap_or_default();
 
-            let variant_array = VariantArray::try_new(variant_array.as_ref())?;
-            let values = variant_get_array_values(&variant_array, path, path_mode, extract)?;
+            let values = variant_get_and_extract(variant_array, path, path_mode, extract)?;
             ColumnarValue::Array(array_from_values(values))
         }
         (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Scalar(path_scalar)) => {
@@ -207,8 +200,8 @@ pub fn invoke_variant_get_typed<T>(
                 .map(|s| s.as_str())
                 .unwrap_or_default();
 
-            let variant_array = VariantArray::try_new(variant_array.as_ref())?;
-            let value = variant_get_single_value(&variant_array, 0, path, path_mode, extract)?;
+            let arr = Arc::clone(variant_array) as ArrayRef;
+            let value = variant_get_and_extract_single(&arr, path, path_mode, extract)?;
 
             ColumnarValue::Scalar(scalar_from_option(value))
         }
@@ -217,13 +210,23 @@ pub fn invoke_variant_get_typed<T>(
                 return exec_err!("expected variant array and paths to be of same length");
             }
 
+            // Per-row paths: must navigate each row individually
             let paths = try_parse_string_columnar(paths)?;
             let variant_array = VariantArray::try_new(variant_array.as_ref())?;
 
-            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, path_mode, extract)
+            let values: Vec<Option<T>> = variant_array
+                .iter()
+                .zip(paths.iter())
+                .map(|(maybe_variant, path)| {
+                    let Some(variant) = maybe_variant else {
+                        return Ok(None);
+                    };
+                    let path = path.unwrap_or_default();
+                    let variant_path = path_mode.try_build_path(path)?;
+                    let Some(value) = variant.get_path(&variant_path) else {
+                        return Ok(None);
+                    };
+                    extract(value)
                 })
                 .collect::<Result<_>>()?;
 
@@ -241,7 +244,14 @@ pub fn invoke_variant_get_typed<T>(
                 .iter()
                 .map(|path| {
                     let path = path.unwrap_or_default();
-                    variant_get_single_value(&variant_array, 0, path, path_mode, extract)
+                    let Some(variant) = variant_array.iter().next().flatten() else {
+                        return Ok(None);
+                    };
+                    let variant_path = path_mode.try_build_path(path)?;
+                    let Some(value) = variant.get_path(&variant_path) else {
+                        return Ok(None);
+                    };
+                    extract(value)
                 })
                 .collect::<Result<_>>()?;