variant get

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

encodings/parquet-variant/src/kernel.rs

@@ -19,11 +19,13 @@ use vortex_mask::Mask;
 
 use crate::ParquetVariant;
 use crate::ParquetVariantArrayExt;
+use crate::variant_get::VariantGetExecuteParent;
 
 pub(crate) static PARENT_KERNELS: ParentKernelSet<ParquetVariant> = ParentKernelSet::new(&[
     ParentKernelSet::lift(&FilterExecuteAdaptor(ParquetVariant)),
     ParentKernelSet::lift(&SliceExecuteAdaptor(ParquetVariant)),
     ParentKernelSet::lift(&TakeExecuteAdaptor(ParquetVariant)),
+    ParentKernelSet::lift(&VariantGetExecuteParent),
 ]);
 
 impl SliceKernel for ParquetVariant {

encodings/parquet-variant/src/lib.rs

@@ -28,6 +28,9 @@ mod array;
 mod kernel;
 mod operations;
 mod validity;
+mod variant_get;
+#[cfg(test)]
+mod variant_get_tests;
 mod vtable;
 
 pub use array::ParquetVariantArrayExt;

encodings/parquet-variant/src/variant_get.rs

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Execute-parent kernel for `variant_get` on `ParquetVariantArray`.
+//!
+//! Delegates to `parquet_variant_compute::variant_get` after converting to Arrow.
+
+use std::sync::Arc;
+
+use parquet_variant::VariantPathElement;
+use parquet_variant_compute::GetOptions;
+use parquet_variant_compute::VariantArray as ArrowVariantArray;
+use vortex_array::ArrayRef;
+use vortex_array::ArrayView;
+use vortex_array::ExecutionCtx;
+use vortex_array::arrays::scalar_fn::ExactScalarFn;
+use vortex_array::arrays::scalar_fn::ScalarFnArrayView;
+use vortex_array::dtype::FieldName;
+use vortex_array::kernel::ExecuteParentKernel;
+use vortex_array::scalar_fn::fns::variant_get::VariantGet;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use crate::ParquetVariant;
+use crate::ParquetVariantArrayExt;
+use crate::ParquetVariantData;
+
+#[derive(Debug)]
+pub(crate) struct VariantGetExecuteParent;
+
+impl ExecuteParentKernel<ParquetVariant> for VariantGetExecuteParent {
+    type Parent = ExactScalarFn<VariantGet>;
+
+    fn execute_parent(
+        &self,
+        array: ArrayView<'_, ParquetVariant>,
+        parent: ScalarFnArrayView<'_, VariantGet>,
+        _child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let field_name: &FieldName = parent.options;
+        variant_get_impl(array, field_name, ctx).map(Some)
+    }
+}
+
+fn variant_get_impl(
+    array: ArrayView<'_, ParquetVariant>,
+    field_name: &FieldName,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    // Convert to Arrow VariantArray
+    let arrow_variant = array.to_arrow(ctx)?;
+
+    // Build path for a single field access
+    let path_element = VariantPathElement::Field {
+        name: field_name.as_ref().into(),
+    };
+    let options = GetOptions::new_with_path(vec![path_element].into());
+
+    // Delegate to the parquet-variant-compute kernel.
+    // With as_type = None, the result is itself a VariantArray.
+    let inner: Arc<dyn arrow_array::Array> = Arc::new(arrow_variant.into_inner());
+    let arrow_result = parquet_variant_compute::variant_get(&inner, options)
+        .map_err(|e| vortex_err!("variant_get failed: {e}"))?;
+
+    // Convert back to Vortex
+    let result_variant = ArrowVariantArray::try_new(
+        arrow_result
+            .as_any()
+            .downcast_ref::<arrow_array::StructArray>()
+            .ok_or_else(|| vortex_err!("variant_get did not return a StructArray"))?,
+    )
+    .map_err(|e| vortex_err!("failed to create VariantArray from result: {e}"))?;
+
+    ParquetVariantData::from_arrow_variant(&result_variant)
+}

encodings/parquet-variant/src/variant_get_tests.rs

@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#[cfg(test)]
+mod tests {
+    use arrow_array::StructArray;
+    use arrow_buffer::NullBuffer;
+    use parquet_variant::Variant as PqVariant;
+    use parquet_variant::VariantBuilderExt;
+    use parquet_variant_compute::VariantArrayBuilder;
+    use vortex_array::ArrayRef;
+    use vortex_array::LEGACY_SESSION;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::expr::root;
+    use vortex_array::expr::variant_get;
+    use vortex_error::VortexResult;
+
+    use crate::ParquetVariantData;
+
+    /// Apply variant_get and execute through the full pipeline (including execute_parent).
+    fn apply_variant_get(arr: &ArrayRef, field: &str) -> VortexResult<ArrayRef> {
+        let expr = variant_get(field, root());
+        let lazy = arr.clone().apply(&expr)?;
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        lazy.execute::<ArrayRef>(&mut ctx)
+    }
+
+    /// Build a VariantArray of objects: [{"a": 1, "b": "x"}, {"a": 2, "c": true}, {"b": "y"}]
+    fn make_object_array() -> VortexResult<ArrayRef> {
+        let mut builder = VariantArrayBuilder::new(3);
+
+        builder
+            .new_object()
+            .with_field("a", 1i32)
+            .with_field("b", "x")
+            .finish();
+
+        builder
+            .new_object()
+            .with_field("a", 2i32)
+            .with_field("c", true)
+            .finish();
+
+        builder.new_object().with_field("b", "y").finish();
+
+        ParquetVariantData::from_arrow_variant(&builder.build())
+    }
+
+    /// Build a nullable VariantArray: [{"a": 10}, NULL, {"a": 30}]
+    fn make_nullable_object_array() -> VortexResult<ArrayRef> {
+        let mut builder = VariantArrayBuilder::new(3);
+
+        builder.new_object().with_field("a", 10i32).finish();
+
+        builder.new_object().with_field("a", 20i32).finish();
+
+        builder.new_object().with_field("a", 30i32).finish();
+
+        let inner = builder.build().into_inner();
+        let null_struct = StructArray::try_new(
+            inner.fields().clone(),
+            inner.columns().to_vec(),
+            Some(NullBuffer::from(vec![true, false, true])),
+        )
+        .unwrap();
+        let arrow_variant = parquet_variant_compute::VariantArray::try_new(&null_struct).unwrap();
+        ParquetVariantData::from_arrow_variant(&arrow_variant)
+    }
+
+    #[test]
+    fn test_variant_get_basic() -> VortexResult<()> {
+        let arr = make_object_array()?;
+        let result = apply_variant_get(&arr, "a")?;
+
+        assert_eq!(result.len(), 3);
+
+        // Row 0: {"a": 1, ...} → variant(1)
+        let s0 = result.scalar_at(0)?;
+        assert!(!s0.is_null());
+        let inner0 = s0.as_variant().value().unwrap();
+        assert_eq!(*inner0, 1i32.into());
+
+        // Row 1: {"a": 2, ...} → variant(2)
+        let s1 = result.scalar_at(1)?;
+        assert!(!s1.is_null());
+        let inner1 = s1.as_variant().value().unwrap();
+        assert_eq!(*inner1, 2i32.into());
+
+        // Row 2: {"b": "y"} → null (field "a" missing)
+        let s2 = result.scalar_at(2)?;
+        assert!(s2.is_null());
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_variant_get_missing_field() -> VortexResult<()> {
+        let arr = make_object_array()?;
+        let result = apply_variant_get(&arr, "nonexistent")?;
+
+        assert_eq!(result.len(), 3);
+        for i in 0..3 {
+            assert!(result.scalar_at(i)?.is_null(), "row {i} should be null");
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_variant_get_null_input() -> VortexResult<()> {
+        let arr = make_nullable_object_array()?;
+        let result = apply_variant_get(&arr, "a")?;
+
+        assert_eq!(result.len(), 3);
+
+        // Row 0: {"a": 10} → variant(10)
+        assert!(!result.scalar_at(0)?.is_null());
+
+        // Row 1: NULL → null
+        assert!(result.scalar_at(1)?.is_null());
+
+        // Row 2: {"a": 30} → variant(30)
+        assert!(!result.scalar_at(2)?.is_null());
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_variant_get_non_object() -> VortexResult<()> {
+        // Array of primitive variants (not objects)
+        let mut builder = VariantArrayBuilder::new(2);
+        builder.append_variant(PqVariant::from(42i32));
+        builder.append_variant(PqVariant::from("hello"));
+        let arr = ParquetVariantData::from_arrow_variant(&builder.build())?;
+
+        let result = apply_variant_get(&arr, "a")?;
+
+        assert_eq!(result.len(), 2);
+        assert!(result.scalar_at(0)?.is_null());
+        assert!(result.scalar_at(1)?.is_null());
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_variant_get_different_field() -> VortexResult<()> {
+        let arr = make_object_array()?;
+        let result = apply_variant_get(&arr, "b")?;
+
+        assert_eq!(result.len(), 3);
+
+        // Row 0: {"a": 1, "b": "x"} → variant("x")
+        assert!(!result.scalar_at(0)?.is_null());
+
+        // Row 1: {"a": 2, "c": true} → null (no "b")
+        assert!(result.scalar_at(1)?.is_null());
+
+        // Row 2: {"b": "y"} → variant("y")
+        assert!(!result.scalar_at(2)?.is_null());
+
+        Ok(())
+    }
+}

vortex-array/public-api.lock

@@ -12584,6 +12584,8 @@ pub fn vortex_array::expr::select_exclude(fields: impl core::convert::Into<vorte
 
 pub fn vortex_array::expr::split_conjunction(expr: &vortex_array::expr::Expression) -> alloc::vec::Vec<vortex_array::expr::Expression>
 
+pub fn vortex_array::expr::variant_get(field: impl core::convert::Into<vortex_array::dtype::FieldName>, child: vortex_array::expr::Expression) -> vortex_array::expr::Expression
+
 pub fn vortex_array::expr::zip_expr(mask: vortex_array::expr::Expression, if_true: vortex_array::expr::Expression, if_false: vortex_array::expr::Expression) -> vortex_array::expr::Expression
 
 pub type vortex_array::expr::Annotations<'a, A> = vortex_utils::aliases::hash_map::HashMap<&'a vortex_array::expr::Expression, vortex_utils::aliases::hash_set::HashSet<A>>
@@ -17352,6 +17354,52 @@ pub fn vortex_array::scalar_fn::fns::select::Select::stat_falsification(&self, o
 
 pub fn vortex_array::scalar_fn::fns::select::Select::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
 
+pub mod vortex_array::scalar_fn::fns::variant_get
+
+pub struct vortex_array::scalar_fn::fns::variant_get::VariantGet
+
+impl core::clone::Clone for vortex_array::scalar_fn::fns::variant_get::VariantGet
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::clone(&self) -> vortex_array::scalar_fn::fns::variant_get::VariantGet
+
+impl vortex_array::scalar_fn::ScalarFnVTable for vortex_array::scalar_fn::fns::variant_get::VariantGet
+
+pub type vortex_array::scalar_fn::fns::variant_get::VariantGet::Options = vortex_array::dtype::FieldName
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::arity(&self, _field_name: &vortex_array::dtype::FieldName) -> vortex_array::scalar_fn::Arity
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::child_name(&self, _instance: &Self::Options, child_idx: usize) -> vortex_array::scalar_fn::ChildName
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::coerce_args(&self, options: &Self::Options, args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_array::dtype::DType>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::deserialize(&self, metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::execute(&self, _field_name: &vortex_array::dtype::FieldName, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::fmt_sql(&self, field_name: &vortex_array::dtype::FieldName, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::id(&self) -> vortex_array::scalar_fn::ScalarFnId
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::is_fallible(&self, _field_name: &vortex_array::dtype::FieldName) -> bool
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::is_null_sensitive(&self, _field_name: &vortex_array::dtype::FieldName) -> bool
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::reduce(&self, field_name: &vortex_array::dtype::FieldName, node: &dyn vortex_array::scalar_fn::ReduceNode, ctx: &dyn vortex_array::scalar_fn::ReduceCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::scalar_fn::ReduceNodeRef>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::return_dtype(&self, _field_name: &vortex_array::dtype::FieldName, arg_dtypes: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::serialize(&self, instance: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::simplify(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, ctx: &dyn vortex_array::scalar_fn::SimplifyCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::simplify_untyped(&self, options: &Self::Options, expr: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::stat_expression(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, stat: vortex_array::expr::stats::Stat, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option<vortex_array::expr::Expression>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::stat_falsification(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option<vortex_array::expr::Expression>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
 pub mod vortex_array::scalar_fn::fns::zip
 
 pub struct vortex_array::scalar_fn::fns::zip::Zip
@@ -18560,6 +18608,44 @@ pub fn vortex_array::scalar_fn::fns::select::Select::stat_falsification(&self, o
 
 pub fn vortex_array::scalar_fn::fns::select::Select::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
 
+impl vortex_array::scalar_fn::ScalarFnVTable for vortex_array::scalar_fn::fns::variant_get::VariantGet
+
+pub type vortex_array::scalar_fn::fns::variant_get::VariantGet::Options = vortex_array::dtype::FieldName
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::arity(&self, _field_name: &vortex_array::dtype::FieldName) -> vortex_array::scalar_fn::Arity
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::child_name(&self, _instance: &Self::Options, child_idx: usize) -> vortex_array::scalar_fn::ChildName
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::coerce_args(&self, options: &Self::Options, args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_array::dtype::DType>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::deserialize(&self, metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::execute(&self, _field_name: &vortex_array::dtype::FieldName, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::fmt_sql(&self, field_name: &vortex_array::dtype::FieldName, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::id(&self) -> vortex_array::scalar_fn::ScalarFnId
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::is_fallible(&self, _field_name: &vortex_array::dtype::FieldName) -> bool
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::is_null_sensitive(&self, _field_name: &vortex_array::dtype::FieldName) -> bool
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::reduce(&self, field_name: &vortex_array::dtype::FieldName, node: &dyn vortex_array::scalar_fn::ReduceNode, ctx: &dyn vortex_array::scalar_fn::ReduceCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::scalar_fn::ReduceNodeRef>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::return_dtype(&self, _field_name: &vortex_array::dtype::FieldName, arg_dtypes: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::serialize(&self, instance: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::simplify(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, ctx: &dyn vortex_array::scalar_fn::SimplifyCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::simplify_untyped(&self, options: &Self::Options, expr: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::stat_expression(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, stat: vortex_array::expr::stats::Stat, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option<vortex_array::expr::Expression>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::stat_falsification(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option<vortex_array::expr::Expression>
+
+pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+
 impl vortex_array::scalar_fn::ScalarFnVTable for vortex_array::scalar_fn::fns::zip::Zip
 
 pub type vortex_array::scalar_fn::fns::zip::Zip::Options = vortex_array::scalar_fn::EmptyOptions

vortex-array/src/expr/exprs.rs

@@ -46,6 +46,7 @@ use crate::scalar_fn::fns::pack::PackOptions;
 use crate::scalar_fn::fns::root::Root;
 use crate::scalar_fn::fns::select::FieldSelection;
 use crate::scalar_fn::fns::select::Select;
+use crate::scalar_fn::fns::variant_get::VariantGet;
 use crate::scalar_fn::fns::zip::Zip;
 
 // ---- Root ----
@@ -699,3 +700,17 @@ pub fn dynamic(
 pub fn list_contains(list: Expression, value: Expression) -> Expression {
     ListContains.new_expr(EmptyOptions, [list, value])
 }
+
+// ---- VariantGet ----
+
+/// Creates an expression that extracts a field from a variant object by name.
+///
+/// Returns a new variant containing the field's value, or null if the field does not exist.
+///
+/// ```rust
+/// # use vortex_array::expr::{variant_get, root};
+/// let expr = variant_get("field_name", root());
+/// ```
+pub fn variant_get(field: impl Into<FieldName>, child: Expression) -> Expression {
+    VariantGet.new_expr(field.into(), vec![child])
+}