more

Author: AdamGS

vortex-array/src/arrays/filter/mod.rs

@@ -10,6 +10,7 @@ pub use vtable::FilterArray;
 mod execute;
 
 mod kernel;
+mod parent_kernels;
 pub use kernel::FilterExecuteAdaptor;
 pub use kernel::FilterKernel;
 pub use kernel::FilterReduce;

vortex-array/src/arrays/filter/parent_kernels.rs

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Parent kernels for wrapper-side execution over [`Filter`].
+//!
+//! [`VariantGet`](crate::scalar_fn::fns::variant_get::VariantGet) relies on the underlying
+//! variant encoding to perform the real extraction work, so it cannot execute directly once a
+//! `Filter` wrapper sits between the expression and that encoding.
+//!
+//! This module handles that pass-through at the wrapper layer by filtering the wrapped variant
+//! child first and then re-dispatching `VariantGet` on the filtered child.
+
+use vortex_error::VortexResult;
+
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::array::ArrayView;
+use crate::arrays::Filter;
+use crate::arrays::Variant;
+use crate::arrays::filter::FilterArrayExt;
+use crate::arrays::scalar_fn::ExactScalarFn;
+use crate::arrays::scalar_fn::ScalarFnArrayView;
+use crate::arrays::scalar_fn::ScalarFnFactoryExt;
+use crate::arrays::variant::VariantArrayExt;
+use crate::dtype::DType;
+use crate::kernel::ExecuteParentKernel;
+use crate::kernel::ParentKernelSet;
+use crate::scalar_fn::fns::variant_get::VariantGet;
+
+pub(super) static PARENT_KERNELS: ParentKernelSet<Filter> =
+    ParentKernelSet::new(&[ParentKernelSet::lift(&FilterVariantGetExecuteParent)]);
+
+/// Pass `variant_get` through a `Filter` wrapper to the underlying variant child.
+#[derive(Debug)]
+struct FilterVariantGetExecuteParent;
+
+impl ExecuteParentKernel<Filter> for FilterVariantGetExecuteParent {
+    type Parent = ExactScalarFn<VariantGet>;
+
+    fn execute_parent(
+        &self,
+        array: ArrayView<'_, Filter>,
+        parent: ScalarFnArrayView<'_, VariantGet>,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        if child_idx != 0 {
+            return Ok(None);
+        }
+
+        let child = array.child();
+        let inner = if let Some(variant) = child.as_opt::<Variant>() {
+            variant.child().clone()
+        } else if matches!(child.dtype(), DType::Variant(_)) {
+            child.clone()
+        } else {
+            return Ok(None);
+        };
+
+        let filtered = inner
+            .filter(array.filter_mask().clone())?
+            .execute::<ArrayRef>(ctx)?;
+
+        VariantGet
+            .try_new_array(parent.len(), parent.options.clone(), [filtered])?
+            .execute::<ArrayRef>(ctx)
+            .map(Some)
+    }
+}

vortex-array/src/arrays/filter/vtable.rs

@@ -32,6 +32,7 @@ use crate::arrays::filter::array::FilterData;
 use crate::arrays::filter::array::SLOT_NAMES;
 use crate::arrays::filter::execute::execute_filter;
 use crate::arrays::filter::execute::execute_filter_fast_paths;
+use crate::arrays::filter::parent_kernels::PARENT_KERNELS;
 use crate::arrays::filter::rules::PARENT_RULES;
 use crate::arrays::filter::rules::RULES;
 use crate::buffer::BufferHandle;
@@ -170,6 +171,15 @@ impl VTable for Filter {
         PARENT_RULES.evaluate(array, parent, child_idx)
     }
 
+    fn execute_parent(
+        array: ArrayView<'_, Self>,
+        parent: &ArrayRef,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
+    }
+
     fn reduce(array: ArrayView<'_, Self>) -> VortexResult<Option<ArrayRef>> {
         RULES.evaluate(array)
     }

vortex-array/src/arrays/slice/mod.rs

@@ -9,6 +9,7 @@
 //! [`ArrayParentReduceRule`] and [`ExecuteParentKernel`] respectively.
 
 mod array;
+mod parent_kernels;
 mod rules;
 mod slice_;
 mod vtable;

vortex-array/src/arrays/slice/parent_kernels.rs

@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Parent kernels for wrapper-side execution over [`Slice`].
+//!
+//! Most scalar functions can commute with `Slice` via generic scalar-fn rules, but
+//! [`VariantGet`](crate::scalar_fn::fns::variant_get::VariantGet) is unusual: it has to reach the
+//! underlying variant encoding's `execute_parent` kernel rather than executing directly.
+//!
+//! This module keeps that wrapper-specific pass-through on the wrapper side. When a
+//! `VariantGet` parent sits directly above a `Slice`, we slice the wrapped variant child first and
+//! then re-dispatch `VariantGet` on the sliced child so execution can continue at the underlying
+//! variant encoding.
+
+use vortex_error::VortexResult;
+
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::array::ArrayView;
+use crate::arrays::Slice;
+use crate::arrays::Variant;
+use crate::arrays::scalar_fn::ExactScalarFn;
+use crate::arrays::scalar_fn::ScalarFnArrayView;
+use crate::arrays::scalar_fn::ScalarFnFactoryExt;
+use crate::arrays::slice::SliceArrayExt;
+use crate::arrays::variant::VariantArrayExt;
+use crate::dtype::DType;
+use crate::kernel::ExecuteParentKernel;
+use crate::kernel::ParentKernelSet;
+use crate::scalar_fn::fns::variant_get::VariantGet;
+
+pub(super) static PARENT_KERNELS: ParentKernelSet<Slice> =
+    ParentKernelSet::new(&[ParentKernelSet::lift(&SliceVariantGetExecuteParent)]);
+
+/// Pass `variant_get` through a `Slice` wrapper to the underlying variant child.
+#[derive(Debug)]
+struct SliceVariantGetExecuteParent;
+
+impl ExecuteParentKernel<Slice> for SliceVariantGetExecuteParent {
+    type Parent = ExactScalarFn<VariantGet>;
+
+    fn execute_parent(
+        &self,
+        array: ArrayView<'_, Slice>,
+        parent: ScalarFnArrayView<'_, VariantGet>,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        if child_idx != 0 {
+            return Ok(None);
+        }
+
+        let child = array.child();
+        let inner = if let Some(variant) = child.as_opt::<Variant>() {
+            variant.child().clone()
+        } else if matches!(child.dtype(), DType::Variant(_)) {
+            child.clone()
+        } else {
+            return Ok(None);
+        };
+
+        let sliced = inner
+            .slice(array.slice_range().clone())?
+            .execute::<ArrayRef>(ctx)?;
+
+        VariantGet
+            .try_new_array(parent.len(), parent.options.clone(), [sliced])?
+            .execute::<ArrayRef>(ctx)
+            .map(Some)
+    }
+}

vortex-array/src/arrays/slice/vtable.rs

@@ -30,6 +30,7 @@ use crate::arrays::slice::SliceArrayExt;
 use crate::arrays::slice::array::CHILD_SLOT;
 use crate::arrays::slice::array::SLOT_NAMES;
 use crate::arrays::slice::array::SliceData;
+use crate::arrays::slice::parent_kernels::PARENT_KERNELS;
 use crate::arrays::slice::rules::PARENT_RULES;
 use crate::buffer::BufferHandle;
 use crate::dtype::DType;
@@ -158,6 +159,15 @@ impl VTable for Slice {
     ) -> VortexResult<Option<ArrayRef>> {
         PARENT_RULES.evaluate(array, parent, child_idx)
     }
+
+    fn execute_parent(
+        array: ArrayView<'_, Self>,
+        parent: &ArrayRef,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
+    }
 }
 impl OperationsVTable<Slice> for Slice {
     fn scalar_at(

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

@@ -14,12 +14,7 @@ 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;
@@ -46,8 +41,11 @@ pub use path::VariantPathElement;
 /// requested path may not exist in every row.
 ///
 /// Execution is handled by variant encodings (e.g. `ParquetVariantArray`) via `execute_parent`.
-/// The canonical `VariantArray` does not support direct execution; a `reduce` rule unwraps
-/// the `VariantArray` wrapper to expose the underlying encoding.
+/// The canonical `VariantArray` does not support direct execution, so `VariantGet` keeps a small
+/// `reduce` rule that unwraps a direct `VariantArray` child to expose the underlying encoding.
+/// Wrapper arrays such as `Slice` and `Filter` forward `VariantGet` from their own
+/// `execute_parent` hooks so the expression can still reach the underlying variant encoding
+/// without teaching `VariantGet` about wrapper-specific array shapes.
 #[derive(Clone)]
 pub struct VariantGet;
 
@@ -222,12 +220,7 @@ impl ScalarFnVTable for VariantGet {
         args: &dyn ExecutionArgs,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let input = args.get(0)?;
-
-        if let Some(rewritten) = rewrite_wrapped_variant_input(input, options, ctx)? {
-            return Ok(rewritten);
-        }
-
+        let _ = (options, args, ctx);
         vortex_bail!("variant_get cannot be executed directly")
     }
 
@@ -262,48 +255,6 @@ 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;