Is this the right way to execute?

Author: AdamGS

encodings/parquet-variant/src/variant_get/tests.rs

@@ -34,6 +34,7 @@ use vortex_array::expr::variant_get;
 use vortex_array::expr::variant_get_as;
 use vortex_array::scalar_fn::fns::variant_get::VariantPath as VortexVariantPath;
 use vortex_error::VortexResult;
+use vortex_mask::Mask;
 
 use crate::ParquetVariant;
 use crate::ParquetVariantArrayExt;
@@ -458,6 +459,45 @@ fn test_variant_get_different_field() -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn test_variant_get_through_slice_wrapper() -> VortexResult<()> {
+    let arr = make_object_array()?;
+
+    let mut ctx = LEGACY_SESSION.create_execution_ctx();
+
+    let expr = variant_get("a", root());
+    let actual = arr
+        .slice(1..3)?
+        .apply(&expr)?
+        .execute::<ArrayRef>(&mut ctx)?;
+
+    let expected = apply_variant_get(&arr, "a")?;
+
+    assert_eq!(actual.len(), 2);
+    assert_eq!(actual.scalar_at(0)?, expected.scalar_at(1)?);
+    assert_eq!(actual.scalar_at(1)?, expected.scalar_at(2)?);
+    Ok(())
+}
+
+#[test]
+fn test_variant_get_through_filter_wrapper() -> VortexResult<()> {
+    let arr = make_object_array()?;
+    let mask = Mask::from_iter([true, false, true]);
+
+    let expr = variant_get("a", root());
+    let mut ctx = LEGACY_SESSION.create_execution_ctx();
+
+    let array = arr.filter(mask.clone())?.apply(&expr)?;
+    let actual = array.execute::<ArrayRef>(&mut ctx)?;
+    let expected = apply_variant_get(&arr, "a")?;
+
+    assert_eq!(mask.true_count(), 2);
+    assert_eq!(actual.len(), 2);
+    assert_eq!(actual.scalar_at(0)?, expected.scalar_at(0)?);
+    assert_eq!(actual.scalar_at(1)?, expected.scalar_at(2)?);
+    Ok(())
+}
+
 // ---------------------------------------------------------------------------
 // Test data helpers
 // ---------------------------------------------------------------------------

vortex-array/src/scalar_fn/fns/variant_get/mod.rs

@@ -14,7 +14,12 @@ use vortex_session::VortexSession;
 
 use crate::ArrayRef;
 use crate::ExecutionCtx;
+use crate::arrays::Filter;
+use crate::arrays::Slice;
 use crate::arrays::Variant;
+use crate::arrays::filter::FilterArrayExt;
+use crate::arrays::scalar_fn::ScalarFnFactoryExt;
+use crate::arrays::slice::SliceArrayExt;
 use crate::arrays::variant::VariantArrayExt;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
@@ -213,14 +218,17 @@ impl ScalarFnVTable for VariantGet {
 
     fn execute(
         &self,
-        _options: &VariantGetOptions,
-        _args: &dyn ExecutionArgs,
-        _ctx: &mut ExecutionCtx,
+        options: &VariantGetOptions,
+        args: &dyn ExecutionArgs,
+        ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        vortex_bail!(
-            "variant_get cannot be executed directly; \
-             it must be pushed down to a variant encoding via execute_parent"
-        )
+        let input = args.get(0)?;
+
+        if let Some(rewritten) = rewrite_wrapped_variant_input(input, options, ctx)? {
+            return Ok(rewritten);
+        }
+
+        vortex_bail!("variant_get cannot be executed directly")
     }
 
     fn reduce(
@@ -254,6 +262,48 @@ impl ScalarFnVTable for VariantGet {
     }
 }
 
+fn rewrite_wrapped_variant_input(
+    input: ArrayRef,
+    options: &VariantGetOptions,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Option<ArrayRef>> {
+    if let Some(slice) = input.as_opt::<Slice>() {
+        let Some(inner) = unwrap_variant_input(slice.child()) else {
+            return Ok(None);
+        };
+        let sliced = inner
+            .slice(slice.slice_range().clone())?
+            .execute::<ArrayRef>(ctx)?;
+        return VariantGet
+            .try_new_array(input.len(), options.clone(), [sliced])?
+            .execute::<ArrayRef>(ctx)
+            .map(Some);
+    }
+
+    if let Some(filter) = input.as_opt::<Filter>() {
+        let Some(inner) = unwrap_variant_input(filter.child()) else {
+            return Ok(None);
+        };
+        let filtered = inner
+            .filter(filter.filter_mask().clone())?
+            .execute::<ArrayRef>(ctx)?;
+        return VariantGet
+            .try_new_array(input.len(), options.clone(), [filtered])?
+            .execute::<ArrayRef>(ctx)
+            .map(Some);
+    }
+
+    Ok(None)
+}
+
+fn unwrap_variant_input(input: &ArrayRef) -> Option<ArrayRef> {
+    if let Some(variant) = input.as_opt::<Variant>() {
+        return Some(variant.child().clone());
+    }
+
+    matches!(input.dtype(), DType::Variant(_)).then(|| input.clone())
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_session::VortexSession;