Lazy array expression evaluation (#5625)

This long-lived feature branch contain changes required for the
"operators RFC", "Lazy Arrays", whatever we're calling it.

The change is hidden behind the `VORTEX_OPERATORS=true` environment
variable allowing us to run parallel benchmarks until we reach parity
between legacy compute function dispatch and the new optimizer rules
framework.

Summary of changes:
* Rename `execute` -> `execute_vector`, add similar `execute_mask`
function. This pattern will be used when turning a Vortex Array into
concrete physical data via whatever means, e.g. `execute_arrow`,
`execute_duckdb`, and so on.
* Introduces `trait Kernel` to represent a physical tree of CPU-based
vector execution
* `Array::batch_execute -> Array::bind_kernel`
* Implements ScalarFnArray execution using a ScalarFnKernel, binding
constant children to scalar datum inputs for the ScalarFn.
* Reverts the change to pull out a `ScalarFnVTable`, in favor of
allowing all `expr::VTable` to be used as scalar functions. This
provides a smoother migration for end-users, at the expense of having to
special-case the `root` expression during execution.
* Expression now holds ScalarFn instance (vtable + options), as well as
children.
* Removes `ExpressionView`, easier to pass separate Expression + options
into vtable.
* `expr::VTable::execute` takes arguments by ownership.
* Implements all `expr::VTable::execute` functions.
* Adds `Array::filter` to vtable, providing an O(1) function for masking
arrays. Fallback is to wrap up the child in a FilterArray.
* Adds `Datum::ensure_vector` and fleshes out `Scalar::repeat`
implementations.
* Fixes expression display to not use debug formatting of options.
* Removes untyped simplification, it was never called.
* Removes ExprOptimizer in favor of a simple reduce_child
`Expr::simplify` function. Complex reduction rules can be implemented
over the array tree only to avoid duplication.
* Introduces `Array::apply(&Expression)` to turn expressions into
arrays. This function is the new API that will live in vortex-array and
allow us to invert the expr -> array dependency to become array -> expr.

---------

Signed-off-by: Nicholas Gates <[email protected]>
Signed-off-by: Siddharth Kumar <[email protected]>
Co-authored-by: Alfonso Subiotto Marqués <[email protected]>
Co-authored-by: Baris Palaska <[email protected]>
Co-authored-by: sherlockbeard <[email protected]>
Co-authored-by: Connor Tsui <[email protected]>

Author: 5 people

Cargo.lock

@@ -15,7 +15,9 @@ use vortex_array::DeserializeMetadata;
 use vortex_array::Precision;
 use vortex_array::ProstMetadata;
 use vortex_array::SerializeMetadata;
-use vortex_array::execution::ExecutionCtx;
+use vortex_array::kernel::BindCtx;
+use vortex_array::kernel::KernelRef;
+use vortex_array::kernel::kernel;
 use vortex_array::patches::Patches;
 use vortex_array::patches::PatchesMetadata;
 use vortex_array::serde::ArrayChildren;
@@ -41,7 +43,6 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
-use vortex_vector::Vector;
 
 use crate::ALPFloat;
 use crate::alp::Exponents;
@@ -140,17 +141,16 @@ impl VTable for ALPVTable {
         )
     }
 
-    fn batch_execute(array: &ALPArray, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
-        let encoded_vector = array.encoded().batch_execute(ctx)?;
-
-        let patches_vectors = if let Some(patches) = array.patches() {
+    fn bind_kernel(array: &ALPArray, ctx: &mut BindCtx) -> VortexResult<KernelRef> {
+        let encoded = array.encoded().bind_kernel(ctx)?;
+        let patches_kernels = if let Some(patches) = array.patches() {
             Some((
-                patches.indices().batch_execute(ctx)?,
-                patches.values().batch_execute(ctx)?,
+                patches.indices().bind_kernel(ctx)?,
+                patches.values().bind_kernel(ctx)?,
                 patches
                     .chunk_offsets()
                     .as_ref()
-                    .map(|co| co.batch_execute(ctx))
+                    .map(|co| co.bind_kernel(ctx))
                     .transpose()?,
             ))
         } else {
@@ -161,7 +161,24 @@ impl VTable for ALPVTable {
         let exponents = array.exponents();
 
         match_each_alp_float_ptype!(array.dtype().as_ptype(), |T| {
-            decompress_into_vector::<T>(encoded_vector, exponents, patches_vectors, patches_offset)
+            Ok(kernel(move || {
+                let encoded_vector = encoded.execute()?;
+                let patches_vectors = match patches_kernels {
+                    Some((idx_kernel, val_kernel, co_kernel)) => Some((
+                        idx_kernel.execute()?,
+                        val_kernel.execute()?,
+                        co_kernel.map(|k| k.execute()).transpose()?,
+                    )),
+                    None => None,
+                };
+
+                decompress_into_vector::<T>(
+                    encoded_vector,
+                    exponents,
+                    patches_vectors,
+                    patches_offset,
+                )
+            }))
         })
     }
 }
@@ -456,15 +473,18 @@ mod tests {
     use std::sync::LazyLock;
 
     use rstest::rstest;
+    use vortex_array::VectorExecutor;
     use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::session::ArraySession;
     use vortex_array::vtable::ValidityHelper;
     use vortex_dtype::PTypeDowncast;
     use vortex_session::VortexSession;
     use vortex_vector::VectorOps;
 
     use super::*;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSession> =
+        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
 
     #[rstest]
     #[case(0)]
@@ -480,7 +500,7 @@ mod tests {
         let values = PrimitiveArray::from_iter((0..size).map(|i| i as f32));
         let encoded = alp_encode(&values, None).unwrap();
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -504,7 +524,7 @@ mod tests {
         let values = PrimitiveArray::from_iter((0..size).map(|i| i as f64));
         let encoded = alp_encode(&values, None).unwrap();
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -534,7 +554,7 @@ mod tests {
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().unwrap().array_len() > 0);
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -562,7 +582,7 @@ mod tests {
         let array = PrimitiveArray::from_option_iter(values);
         let encoded = alp_encode(&array, None).unwrap();
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -601,7 +621,7 @@ mod tests {
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().unwrap().array_len() > 0);
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -643,7 +663,7 @@ mod tests {
         let slice_len = slice_end - slice_start;
         let sliced_encoded = encoded.slice(slice_start..slice_end);
 
-        let result_vector = sliced_encoded.execute(&SESSION).unwrap();
+        let result_vector = sliced_encoded.execute_vector_optimized(&SESSION).unwrap();
         let result_primitive = result_vector.into_primitive().into_f64();
 
         for idx in 0..slice_len {

encodings/alp/src/alp/array.rs

@@ -204,6 +204,7 @@ mod tests {
     use std::sync::LazyLock;
 
     use vortex_array::IntoArray;
+    use vortex_array::VectorExecutor;
     use vortex_array::assert_arrays_eq;
     use vortex_array::validity::Validity;
     use vortex_buffer::Buffer;
@@ -536,7 +537,7 @@ mod tests {
             let unpacked_array = unpack_array(&bitpacked);
 
             // Method 3: Using the execute() method (this is what would be used in production).
-            let executed = bitpacked.into_array().execute(&SESSION).unwrap();
+            let executed = bitpacked.into_array().execute_vector(&SESSION).unwrap();
 
             // All three should produce the same length.
             assert_eq!(vector_result.len(), array.len(), "vector length mismatch");
@@ -556,7 +557,10 @@ mod tests {
 
             // Verify that the execute() method works correctly by comparing with unpack_array.
             // We convert unpack_array result to a vector using execute() to compare.
-            let unpacked_executed = unpacked_array.into_array().execute(&SESSION).unwrap();
+            let unpacked_executed = unpacked_array
+                .into_array()
+                .execute_vector(&SESSION)
+                .unwrap();
             match (&executed, &unpacked_executed) {
                 (Vector::Primitive(exec_pv), Vector::Primitive(unpack_pv)) => {
                     assert_eq!(
@@ -593,7 +597,7 @@ mod tests {
         let sliced_bp = sliced.as_::<BitPackedVTable>();
         let vector_result = unpack_to_primitive_vector(sliced_bp);
         let unpacked_array = unpack_array(sliced_bp);
-        let executed = sliced.execute(&SESSION).unwrap();
+        let executed = sliced.execute_vector(&SESSION).unwrap();
 
         assert_eq!(
             vector_result.len(),

encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs

@@ -4,7 +4,9 @@
 use vortex_array::DeserializeMetadata;
 use vortex_array::ProstMetadata;
 use vortex_array::SerializeMetadata;
-use vortex_array::execution::ExecutionCtx;
+use vortex_array::kernel::BindCtx;
+use vortex_array::kernel::KernelRef;
+use vortex_array::kernel::kernel;
 use vortex_array::patches::Patches;
 use vortex_array::patches::PatchesMetadata;
 use vortex_array::serde::ArrayChildren;
@@ -23,7 +25,6 @@ use vortex_error::VortexError;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
-use vortex_vector::Vector;
 use vortex_vector::VectorMutOps;
 
 use crate::BitPackedArray;
@@ -172,8 +173,11 @@ impl VTable for BitPackedVTable {
         )
     }
 
-    fn batch_execute(array: &BitPackedArray, _ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
-        Ok(unpack_to_primitive_vector(array).freeze().into())
+    fn bind_kernel(array: &BitPackedArray, _ctx: &mut BindCtx) -> VortexResult<KernelRef> {
+        let array = array.clone();
+        Ok(kernel(move || {
+            Ok(unpack_to_primitive_vector(&array).freeze().into())
+        }))
     }
 }
 

encodings/fastlanes/src/bitpacking/vtable/mod.rs

@@ -24,6 +24,7 @@ vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }
 vortex-mask = { workspace = true }
 vortex-scalar = { workspace = true }
+vortex-session = { workspace = true }
 
 [lints]
 workspace = true

encodings/runend/Cargo.toml

@@ -11,6 +11,7 @@ pub mod compress;
 mod compute;
 mod iter;
 mod ops;
+mod rules;
 
 #[doc(hidden)]
 pub mod _benchmarking {
@@ -23,11 +24,17 @@ pub mod _benchmarking {
 use vortex_array::ArrayBufferVisitor;
 use vortex_array::ArrayChildVisitor;
 use vortex_array::Canonical;
+use vortex_array::session::ArraySession;
+use vortex_array::session::ArraySessionExt;
+use vortex_array::vtable::ArrayVTableExt;
 use vortex_array::vtable::EncodeVTable;
 use vortex_array::vtable::VisitorVTable;
 use vortex_error::VortexResult;
+use vortex_session::SessionExt;
+use vortex_session::VortexSession;
 
 use crate::compress::runend_encode;
+use crate::rules::RunEndScalarFnRule;
 
 impl EncodeVTable<RunEndVTable> for RunEndVTable {
     fn encode(
@@ -59,6 +66,15 @@ impl VisitorVTable<RunEndVTable> for RunEndVTable {
     }
 }
 
+/// Initialize run-end encoding in the given session.
+pub fn initialize(session: &mut VortexSession) {
+    session.arrays().register(RunEndVTable.as_vtable());
+    session
+        .get_mut::<ArraySession>()
+        .optimizer_mut()
+        .register_parent_rule(RunEndScalarFnRule);
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_array::ProstMetadata;

encodings/runend/src/lib.rs

@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::arrays::AnyScalarFn;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::ConstantVTable;
+use vortex_array::arrays::ScalarFnArray;
+use vortex_array::optimizer::rules::ArrayParentReduceRule;
+use vortex_array::optimizer::rules::Exact;
+use vortex_dtype::DType;
+use vortex_error::VortexResult;
+
+use crate::RunEndArray;
+use crate::RunEndVTable;
+
+/// A rule to push down scalar functions through run-end encoding into the values array.
+///
+/// This only works if all other children of the scalar function array are constants.
+#[derive(Debug)]
+pub(crate) struct RunEndScalarFnRule;
+
+impl ArrayParentReduceRule<Exact<RunEndVTable>, AnyScalarFn> for RunEndScalarFnRule {
+    fn child(&self) -> Exact<RunEndVTable> {
+        Exact::from(&RunEndVTable)
+    }
+
+    fn parent(&self) -> AnyScalarFn {
+        AnyScalarFn
+    }
+
+    fn reduce_parent(
+        &self,
+        run_end: &RunEndArray,
+        parent: &ScalarFnArray,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        for (idx, child) in parent.children().iter().enumerate() {
+            if idx == child_idx {
+                // Skip ourselves
+                continue;
+            }
+
+            if !child.is::<ConstantVTable>() {
+                // We can only push down if all other children are constants
+                return Ok(None);
+            }
+        }
+
+        // TODO(ngates): relax this constraint and implement run-end decoding for all vector types.
+        if !matches!(parent.dtype(), DType::Bool(_) | DType::Primitive(..)) {
+            return Ok(None);
+        }
+
+        let values_len = run_end.values().len();
+        let mut new_children = parent.children();
+        for (idx, child) in new_children.iter_mut().enumerate() {
+            if idx == child_idx {
+                // Replace ourselves with run end values
+                *child = run_end.values().clone();
+                continue;
+            }
+
+            // Replace other children with their constant scalar value with length adjusted
+            // to the length of the run end values.
+            let constant = child.as_::<ConstantVTable>();
+            *child = ConstantArray::new(constant.scalar().clone(), values_len).into_array();
+        }
+
+        let new_values =
+            ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children, values_len)?
+                .into_array();
+
+        Ok(Some(
+            RunEndArray::try_new_offset_length(
+                run_end.ends().clone(),
+                new_values,
+                run_end.offset(),
+                run_end.len(),
+            )?
+            .into_array(),
+        ))
+    }
+}