Remove selection mask from batch vtable (#5424)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

encodings/sequence/src/array.rs

@@ -4,8 +4,10 @@
 use std::hash::Hash;
 use std::ops::Range;
 
+use num_traits::One;
 use num_traits::cast::FromPrimitive;
 use vortex_array::arrays::PrimitiveArray;
+use vortex_array::execution::ExecutionCtx;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::stats::{ArrayStats, StatsSetRef};
 use vortex_array::vtable::{
@@ -24,6 +26,8 @@ use vortex_dtype::{
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_mask::Mask;
 use vortex_scalar::{PValue, Scalar, ScalarValue};
+use vortex_vector::Vector;
+use vortex_vector::primitive::PVector;
 
 vtable!(Sequence);
 
@@ -242,6 +246,26 @@ impl VTable for SequenceVTable {
             len,
         ))
     }
+
+    fn execute(array: &Self::Array, _ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        Ok(match_each_native_ptype!(array.ptype(), |P| {
+            let base = array.base().cast::<P>();
+            let multiplier = array.multiplier().cast::<P>();
+
+            let values = if multiplier == <P>::one() {
+                BufferMut::from_iter(
+                    (0..array.len()).map(|i| base + <P>::from_usize(i).vortex_expect("must fit")),
+                )
+            } else {
+                BufferMut::from_iter(
+                    (0..array.len())
+                        .map(|i| base + <P>::from_usize(i).vortex_expect("must fit") * multiplier),
+                )
+            };
+
+            PVector::<P>::new(values.freeze(), Mask::new_true(array.len())).into()
+        }))
+    }
 }
 
 impl ArrayVTable<SequenceVTable> for SequenceVTable {

vortex-array/src/array/operator.rs

@@ -3,9 +3,10 @@
 
 use std::sync::Arc;
 
+use vortex_compute::filter::Filter;
 use vortex_error::{VortexResult, vortex_panic};
 use vortex_mask::Mask;
-use vortex_vector::{Vector, VectorOps, vector_matches_dtype};
+use vortex_vector::{Vector, vector_matches_dtype};
 
 use crate::execution::{BatchKernelRef, BindCtx, DummyExecutionCtx, ExecutionCtx};
 use crate::pipeline::PipelinedNode;
@@ -24,7 +25,7 @@ pub trait ArrayOperator: 'static + Send + Sync {
     ///
     /// If the mask length does not match the array length.
     /// If the array's implementation returns an invalid vector (wrong length, wrong type, etc.).
-    fn execute_batch(&self, selection: &Mask, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector>;
+    fn execute_batch(&self, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector>;
 
     /// Optimize the array by running the optimization rules.
     fn reduce(&self) -> VortexResult<Option<ArrayRef>>;
@@ -44,8 +45,8 @@ pub trait ArrayOperator: 'static + Send + Sync {
 }
 
 impl ArrayOperator for Arc<dyn Array> {
-    fn execute_batch(&self, selection: &Mask, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
-        self.as_ref().execute_batch(selection, ctx)
+    fn execute_batch(&self, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        self.as_ref().execute_batch(ctx)
     }
 
     fn reduce(&self) -> VortexResult<Option<ArrayRef>> {
@@ -70,17 +71,8 @@ impl ArrayOperator for Arc<dyn Array> {
 }
 
 impl<V: VTable> ArrayOperator for ArrayAdapter<V> {
-    fn execute_batch(&self, selection: &Mask, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
-        let vector =
-            <V::OperatorVTable as OperatorVTable<V>>::execute_batch(&self.0, selection, ctx)?;
-
-        // Such a cheap check that we run it always. More expensive DType checks live in
-        // debug_assertions.
-        assert_eq!(
-            vector.len(),
-            selection.true_count(),
-            "Batch execution returned vector of incorrect length"
-        );
+    fn execute_batch(&self, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        let vector = V::execute(&self.0, ctx)?;
 
         if cfg!(debug_assertions) {
             // Checks for correct type and nullability.
@@ -130,17 +122,20 @@ impl BindCtx for () {
 
 impl dyn Array + '_ {
     pub fn execute(&self) -> VortexResult<Vector> {
-        self.execute_with_selection(&Mask::new_true(self.len()))
+        // Check if the array is a pipeline node
+        if self.as_pipelined().is_some() {
+            return PipelineDriver::new(self.to_array()).execute(&Mask::new_true(self.len()));
+        }
+        self.execute_batch(&mut DummyExecutionCtx)
     }
 
     pub fn execute_with_selection(&self, selection: &Mask) -> VortexResult<Vector> {
-        assert_eq!(self.len(), selection.len());
-
         // Check if the array is a pipeline node
         if self.as_pipelined().is_some() {
             return PipelineDriver::new(self.to_array()).execute(selection);
         }
-
-        self.execute_batch(selection, &mut DummyExecutionCtx)
+        Ok(self
+            .execute_batch(&mut DummyExecutionCtx)?
+            .filter(selection))
     }
 }

vortex-array/src/arrays/bool/vtable/mod.rs

@@ -4,8 +4,11 @@
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult, vortex_bail};
+use vortex_vector::Vector;
+use vortex_vector::bool::BoolVector;
 
 use crate::arrays::BoolArray;
+use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
 use crate::validity::Validity;
 use crate::vtable::{NotSupported, VTable, ValidityVTableFromValidityHelper};
@@ -91,6 +94,10 @@ impl VTable for BoolVTable {
 
         BoolArray::try_new(buffers[0].clone(), metadata.offset as usize, len, validity)
     }
+
+    fn execute(array: &Self::Array, _ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        Ok(BoolVector::new(array.bit_buffer().clone(), array.validity_mask()).into())
+    }
 }
 
 #[derive(Clone, Debug)]

vortex-array/src/arrays/chunked/vtable/mod.rs

@@ -5,11 +5,13 @@ use itertools::Itertools;
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::{DType, Nullability, PType};
 use vortex_error::{VortexResult, vortex_bail, vortex_err};
+use vortex_vector::{Vector, VectorMut, VectorMutOps};
 
 use crate::arrays::ChunkedArray;
+use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
 use crate::vtable::{NotSupported, VTable};
-use crate::{EmptyMetadata, EncodingId, EncodingRef, ToCanonical, vtable};
+use crate::{ArrayOperator, EmptyMetadata, EncodingId, EncodingRef, ToCanonical, vtable};
 
 mod array;
 mod canonical;
@@ -95,6 +97,15 @@ impl VTable for ChunkedVTable {
         // Each chunk was validated during deserialization to match the expected dtype.
         unsafe { Ok(ChunkedArray::new_unchecked(chunks, dtype.clone())) }
     }
+
+    fn execute(array: &Self::Array, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        let mut vector = VectorMut::with_capacity(array.dtype(), 0);
+        for chunk in array.chunks() {
+            let chunk_vector = chunk.execute_batch(ctx)?;
+            vector.extend_from_vector(&chunk_vector);
+        }
+        Ok(vector.freeze())
+    }
 }
 
 #[derive(Clone, Debug)]

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

@@ -6,5 +6,7 @@ pub use array::ConstantArray;
 
 mod compute;
 
+mod vector;
 mod vtable;
+
 pub use vtable::{ConstantEncoding, ConstantVTable};

…y/src/arrays/constant/vtable/operator.rs …rtex-array/src/arrays/constant/vector.rsvortex-array/src/arrays/constant/vtable/operator.rs renamed to vortex-array/src/arrays/constant/vector.rs vortex-array/src/arrays/constant/vtable/operator.rs renamed to vortex-array/src/arrays/constant/vector.rs

@@ -4,7 +4,7 @@
 use vortex_dtype::{
     DType, DecimalType, PrecisionScale, match_each_decimal_value_type, match_each_native_ptype,
 };
-use vortex_error::{VortexExpect, VortexResult};
+use vortex_error::VortexExpect;
 use vortex_scalar::{BinaryScalar, BoolScalar, DecimalScalar, PrimitiveScalar, Scalar, Utf8Scalar};
 use vortex_vector::binaryview::{BinaryVectorMut, StringVectorMut};
 use vortex_vector::bool::BoolVectorMut;
@@ -13,29 +13,7 @@ use vortex_vector::null::NullVectorMut;
 use vortex_vector::primitive::{PVectorMut, PrimitiveVectorMut};
 use vortex_vector::{VectorMut, VectorMutOps};
 
-use crate::ArrayRef;
-use crate::arrays::{ConstantArray, ConstantVTable};
-use crate::execution::{BatchKernelRef, BindCtx, kernel};
-use crate::vtable::OperatorVTable;
-
-impl OperatorVTable<ConstantVTable> for ConstantVTable {
-    fn bind(
-        array: &ConstantArray,
-        selection: Option<&ArrayRef>,
-        ctx: &mut dyn BindCtx,
-    ) -> VortexResult<BatchKernelRef> {
-        let mask = ctx.bind_selection(array.len, selection)?;
-        let scalar = array.scalar().clone();
-
-        Ok(kernel(move || {
-            // TODO(ngates): would be good to do a sum aggregation, rather than execution.
-            let mask = mask.execute()?;
-            Ok(to_vector(scalar, mask.true_count()).freeze())
-        }))
-    }
-}
-
-fn to_vector(scalar: Scalar, len: usize) -> VectorMut {
+pub(super) fn to_vector(scalar: Scalar, len: usize) -> VectorMut {
     match scalar.dtype() {
         DType::Null => NullVectorMut::new(len).into(),
         DType::Bool(_) => to_vector_bool(scalar.as_bool(), len).into(),

vortex-array/src/arrays/constant/vtable/mod.rs

@@ -5,8 +5,11 @@ use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail};
 use vortex_scalar::{Scalar, ScalarValue};
+use vortex_vector::{Vector, VectorMutOps};
 
 use crate::arrays::ConstantArray;
+use crate::arrays::constant::vector::to_vector;
+use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
 use crate::vtable::{NotSupported, VTable};
 use crate::{EmptyMetadata, EncodingId, EncodingRef, vtable};
@@ -15,7 +18,6 @@ mod array;
 mod canonical;
 mod encode;
 mod operations;
-mod operator;
 mod validity;
 mod visitor;
 
@@ -37,7 +39,7 @@ impl VTable for ConstantVTable {
     // TODO(ngates): implement a compute kernel for elementwise operations
     type ComputeVTable = NotSupported;
     type EncodeVTable = Self;
-    type OperatorVTable = Self;
+    type OperatorVTable = NotSupported;
 
     fn id(_encoding: &Self::Encoding) -> EncodingId {
         EncodingId::new_ref("vortex.constant")
@@ -74,4 +76,8 @@ impl VTable for ConstantVTable {
         let scalar = Scalar::new(dtype.clone(), sv);
         Ok(ConstantArray::new(scalar, len))
     }
+
+    fn execute(array: &Self::Array, _ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        Ok(to_vector(array.scalar().clone(), array.len()).freeze())
+    }
 }

vortex-array/src/arrays/decimal/vtable/mod.rs

@@ -2,11 +2,14 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_buffer::{Alignment, Buffer, ByteBuffer};
-use vortex_dtype::{DType, NativeDecimalType, match_each_decimal_value_type};
+use vortex_dtype::{DType, NativeDecimalType, PrecisionScale, match_each_decimal_value_type};
 use vortex_error::{VortexResult, vortex_bail, vortex_ensure};
 use vortex_scalar::DecimalType;
+use vortex_vector::Vector;
+use vortex_vector::decimal::DVector;
 
 use crate::arrays::DecimalArray;
+use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
 use crate::validity::Validity;
 use crate::vtable::{NotSupported, VTable, ValidityVTableFromValidityHelper};
@@ -104,6 +107,19 @@ impl VTable for DecimalVTable {
             DecimalArray::try_new::<D>(buffer, *decimal_dtype, validity)
         })
     }
+
+    fn execute(array: &Self::Array, _ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        match_each_decimal_value_type!(array.values_type(), |D| {
+            Ok(unsafe {
+                DVector::<D>::new_unchecked(
+                    PrecisionScale::new_unchecked(array.precision(), array.scale()),
+                    array.buffer::<D>(),
+                    array.validity_mask(),
+                )
+            }
+            .into())
+        })
+    }
 }
 
 #[derive(Clone, Debug)]

vortex-array/src/arrays/expr/vtable/mod.rs

@@ -12,12 +12,14 @@ use std::fmt::Debug;
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail};
+use vortex_vector::Vector;
 
 use crate::arrays::expr::ExprArray;
+use crate::execution::ExecutionCtx;
 use crate::expr::Expression;
 use crate::serde::ArrayChildren;
 use crate::vtable::{NotSupported, VTable};
-use crate::{EncodingId, EncodingRef, vtable};
+use crate::{Array, ArrayOperator, EncodingId, EncodingRef, vtable};
 
 vtable!(Expr);
 
@@ -76,6 +78,11 @@ impl VTable for ExprVTable {
 
         ExprArray::try_new(child, expr.clone(), dtype.clone())
     }
+
+    fn execute(array: &Self::Array, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        let scope = array.child().execute_batch(ctx)?;
+        array.expr().execute(&scope, array.child().dtype())
+    }
 }
 
 pub struct ExprArrayMetadata((Expression, DType));

vortex-array/src/arrays/extension/vtable/mod.rs

@@ -11,11 +11,13 @@ mod visitor;
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail};
+use vortex_vector::Vector;
 
 use crate::arrays::extension::ExtensionArray;
+use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
 use crate::vtable::{NotSupported, VTable, ValidityVTableFromChild};
-use crate::{EmptyMetadata, EncodingId, EncodingRef, vtable};
+use crate::{ArrayOperator, EmptyMetadata, EncodingId, EncodingRef, vtable};
 
 vtable!(Extension);
 
@@ -70,6 +72,10 @@ impl VTable for ExtensionVTable {
         let storage = children.get(0, ext_dtype.storage_dtype(), len)?;
         Ok(ExtensionArray::new(ext_dtype.clone(), storage))
     }
+
+    fn execute(array: &Self::Array, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        array.storage().execute_batch(ctx)
+    }
 }
 
 #[derive(Clone, Debug)]