Move Array BatchKernel to sync API (#5090)

* It's a little aggressive to require an async runtime simply to
decompress in-memory data.
* Async expressions tend to require more powerful APIs anyway. For
example, returning partial results since the output can be much larger
than the input (not common for in-memory scalar compute). As such, I
propose implementing async expressions as layout readers (with an
improved API), see #5089

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

Author: gatesn

vortex-array/src/array/operator.rs

@@ -3,62 +3,58 @@
 
 use std::sync::Arc;
 
-use async_trait::async_trait;
 use vortex_error::VortexResult;
 use vortex_vector::Vector;
 
-use crate::execution::{BatchKernel, BindCtx};
+use crate::execution::{BatchKernelRef, BindCtx};
 use crate::vtable::{OperatorVTable, VTable};
 use crate::{Array, ArrayAdapter, ArrayRef};
 
 /// Array functions as provided by the `OperatorVTable`.
 ///
 /// Note: the public functions such as "execute" should move onto the main `Array` trait when
 /// operators is stabilized. The other functions should remain on a `pub(crate)` trait.
-#[async_trait]
 pub trait ArrayOperator: 'static + Send + Sync {
     /// Execute the array producing a canonical vector.
-    async fn execute(&self) -> VortexResult<Vector> {
-        self.execute_with_selection(None).await
+    fn execute(&self) -> VortexResult<Vector> {
+        self.execute_with_selection(None)
     }
 
     /// Execute the array with a selection mask, producing a canonical vector.
-    async fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector>;
+    fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector>;
 
     /// Bind the array to a batch kernel. This is an internal function
     fn bind(
         &self,
         selection: Option<&ArrayRef>,
         ctx: &mut dyn BindCtx,
-    ) -> VortexResult<BatchKernel>;
+    ) -> VortexResult<BatchKernelRef>;
 }
 
-#[async_trait]
 impl ArrayOperator for Arc<dyn Array> {
-    async fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector> {
-        self.as_ref().execute_with_selection(selection).await
+    fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector> {
+        self.as_ref().execute_with_selection(selection)
     }
 
     fn bind(
         &self,
         selection: Option<&ArrayRef>,
         ctx: &mut dyn BindCtx,
-    ) -> VortexResult<BatchKernel> {
+    ) -> VortexResult<BatchKernelRef> {
         self.as_ref().bind(selection, ctx)
     }
 }
 
-#[async_trait]
 impl<V: VTable> ArrayOperator for ArrayAdapter<V> {
-    async fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector> {
-        self.bind(selection, &mut ())?.await
+    fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector> {
+        self.bind(selection, &mut ())?.execute()
     }
 
     fn bind(
         &self,
         selection: Option<&ArrayRef>,
         ctx: &mut dyn BindCtx,
-    ) -> VortexResult<BatchKernel> {
+    ) -> VortexResult<BatchKernelRef> {
         <V::OperatorVTable as OperatorVTable<V>>::bind(&self.0, selection, ctx)
     }
 }
@@ -69,7 +65,7 @@ impl BindCtx for () {
         &mut self,
         array: &ArrayRef,
         selection: Option<&ArrayRef>,
-    ) -> VortexResult<BatchKernel> {
+    ) -> VortexResult<BatchKernelRef> {
         array.bind(selection, self)
     }
 }

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

@@ -1,34 +1,33 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use futures::{FutureExt, try_join};
 use vortex_compute::filter::Filter;
 use vortex_error::VortexResult;
 use vortex_vector::BoolVector;
 
 use crate::ArrayRef;
 use crate::arrays::{BoolArray, BoolVTable};
-use crate::execution::{BatchKernel, BindCtx};
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
 use crate::vtable::{OperatorVTable, ValidityHelper};
 
 impl OperatorVTable<BoolVTable> for BoolVTable {
     fn bind(
         array: &BoolArray,
         selection: Option<&ArrayRef>,
         ctx: &mut dyn BindCtx,
-    ) -> VortexResult<BatchKernel> {
+    ) -> VortexResult<BatchKernelRef> {
         let bits = array.buffer.clone();
         let mask = ctx.bind_selection(array.len(), selection)?;
         let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
 
-        Ok(async move {
-            let (mask, validity) = try_join!(mask, validity)?;
+        Ok(kernel(move || {
+            let mask = mask.execute()?;
+            let validity = validity.execute()?;
 
             // Note that validity already has the mask applied so we only need to apply it to bits.
             let bits = bits.filter(&mask);
 
             Ok(BoolVector::new(bits, validity).into())
-        }
-        .boxed())
+        }))
     }
 }

vortex-array/src/compute/arrays/logical.rs

@@ -5,7 +5,6 @@ use std::hash::{Hash, Hasher};
 use std::sync::LazyLock;
 
 use enum_map::{Enum, EnumMap, enum_map};
-use futures::{FutureExt, try_join};
 use vortex_buffer::ByteBuffer;
 use vortex_compute::logical::{
     LogicalAnd, LogicalAndKleene, LogicalAndNot, LogicalOr, LogicalOrKleene,
@@ -14,7 +13,7 @@ use vortex_dtype::DType;
 use vortex_error::VortexResult;
 use vortex_vector::BoolVector;
 
-use crate::execution::{BatchKernel, BindCtx};
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
 use crate::serde::ArrayChildren;
 use crate::stats::{ArrayStats, StatsSetRef};
 use crate::vtable::{
@@ -183,37 +182,35 @@ impl OperatorVTable<LogicalVTable> for LogicalVTable {
         array: &LogicalArray,
         selection: Option<&ArrayRef>,
         ctx: &mut dyn BindCtx,
-    ) -> VortexResult<BatchKernel> {
+    ) -> VortexResult<BatchKernelRef> {
         let lhs = ctx.bind(&array.lhs, selection)?;
         let rhs = ctx.bind(&array.rhs, selection)?;
 
         Ok(match array.operator() {
-            LogicalOperator::And => kernel(lhs, rhs, |l, r| l.and(&r)),
-            LogicalOperator::AndKleene => kernel(lhs, rhs, |l, r| l.and_kleene(&r)),
-            LogicalOperator::Or => kernel(lhs, rhs, |l, r| l.or(&r)),
-            LogicalOperator::OrKleene => kernel(lhs, rhs, |l, r| l.or_kleene(&r)),
-            LogicalOperator::AndNot => kernel(lhs, rhs, |l, r| l.and_not(&r)),
+            LogicalOperator::And => logical_kernel(lhs, rhs, |l, r| l.and(&r)),
+            LogicalOperator::AndKleene => logical_kernel(lhs, rhs, |l, r| l.and_kleene(&r)),
+            LogicalOperator::Or => logical_kernel(lhs, rhs, |l, r| l.or(&r)),
+            LogicalOperator::OrKleene => logical_kernel(lhs, rhs, |l, r| l.or_kleene(&r)),
+            LogicalOperator::AndNot => logical_kernel(lhs, rhs, |l, r| l.and_not(&r)),
         })
     }
 }
 
 /// Batch execution kernel for logical operations.
-fn kernel<O>(lhs: BatchKernel, rhs: BatchKernel, op: O) -> BatchKernel
+fn logical_kernel<O>(lhs: BatchKernelRef, rhs: BatchKernelRef, op: O) -> BatchKernelRef
 where
     O: Fn(BoolVector, BoolVector) -> BoolVector + Send + 'static,
 {
-    async move {
-        let (lhs, rhs) = try_join!(lhs, rhs)?;
-        let (lhs, rhs) = (lhs.into_bool(), rhs.into_bool());
+    kernel(move || {
+        let lhs = lhs.execute()?.into_bool();
+        let rhs = rhs.execute()?.into_bool();
         Ok(op(lhs, rhs).into())
-    }
-    .boxed()
+    })
 }
 
 #[cfg(test)]
 mod tests {
     use vortex_buffer::bitbuffer;
-    use vortex_io::runtime::single::block_on;
 
     use crate::compute::arrays::logical::{LogicalArray, LogicalOperator};
     use crate::{ArrayOperator, ArrayRef, IntoArray};
@@ -224,28 +221,23 @@ mod tests {
 
     #[test]
     fn test_and() {
-        block_on(|_| async {
-            let lhs = bitbuffer![0 1 0].into_array();
-            let rhs = bitbuffer![0 1 1].into_array();
-            let result = and_(lhs, rhs).execute().await.unwrap().into_bool();
-            assert_eq!(result.bits(), &bitbuffer![0 1 0]);
-        })
+        let lhs = bitbuffer![0 1 0].into_array();
+        let rhs = bitbuffer![0 1 1].into_array();
+        let result = and_(lhs, rhs).execute().unwrap().into_bool();
+        assert_eq!(result.bits(), &bitbuffer![0 1 0]);
     }
 
     #[test]
     fn test_and_selected() {
-        block_on(|_| async {
-            let lhs = bitbuffer![0 1 0].into_array();
-            let rhs = bitbuffer![0 1 1].into_array();
-
-            let selection = bitbuffer![0 1 1].into_array();
-
-            let result = and_(lhs, rhs)
-                .execute_with_selection(Some(&selection))
-                .await
-                .unwrap()
-                .into_bool();
-            assert_eq!(result.bits(), &bitbuffer![1 0]);
-        })
+        let lhs = bitbuffer![0 1 0].into_array();
+        let rhs = bitbuffer![0 1 1].into_array();
+
+        let selection = bitbuffer![0 1 1].into_array();
+
+        let result = and_(lhs, rhs)
+            .execute_with_selection(Some(&selection))
+            .unwrap()
+            .into_bool();
+        assert_eq!(result.bits(), &bitbuffer![1 0]);
     }
 }

vortex-array/src/execution/batch.rs

@@ -1,14 +1,31 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use futures::future::BoxFuture;
 use vortex_error::VortexResult;
 use vortex_vector::Vector;
 
 use crate::ArrayRef;
 
 /// Type-alias for heap-allocated batch execution kernels.
-pub type BatchKernel = BoxFuture<'static, VortexResult<Vector>>;
+pub type BatchKernelRef = Box<dyn BatchKernel>;
+
+/// Trait for batch execution kernels that produce a vector result.
+pub trait BatchKernel: 'static + Send {
+    fn execute(self: Box<Self>) -> VortexResult<Vector>;
+}
+
+/// Adapter to create a batch kernel from a closure.
+pub struct BatchKernelAdapter<F>(F);
+impl<F: FnOnce() -> VortexResult<Vector> + Send + 'static> BatchKernel for BatchKernelAdapter<F> {
+    fn execute(self: Box<Self>) -> VortexResult<Vector> {
+        self.0()
+    }
+}
+
+/// Create a batch execution kernel from the given closure.
+pub fn kernel<F: FnOnce() -> VortexResult<Vector> + Send + 'static>(f: F) -> BatchKernelRef {
+    Box::new(BatchKernelAdapter(f))
+}
 
 /// Context for binding batch execution kernels.
 ///
@@ -17,6 +34,9 @@ pub type BatchKernel = BoxFuture<'static, VortexResult<Vector>>;
 pub trait BindCtx {
     /// Bind the given array and optional selection to produce a batch kernel, possibly reusing
     /// previously bound results from this context.
-    fn bind(&mut self, array: &ArrayRef, selection: Option<&ArrayRef>)
-    -> VortexResult<BatchKernel>;
+    fn bind(
+        &mut self,
+        array: &ArrayRef,
+        selection: Option<&ArrayRef>,
+    ) -> VortexResult<BatchKernelRef>;
 }

vortex-array/src/execution/mask.rs

@@ -1,11 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::pin::Pin;
-use std::task::{Context, Poll};
-
-use futures::FutureExt;
-use futures::future::BoxFuture;
 use vortex_dtype::DType;
 use vortex_dtype::Nullability::NonNullable;
 use vortex_error::{VortexExpect, VortexResult, vortex_bail};
@@ -17,23 +12,19 @@ use crate::execution::BindCtx;
 pub enum MaskExecution {
     AllTrue(usize),
     AllFalse(usize),
-    Future(BoxFuture<'static, VortexResult<Mask>>),
+    Lazy(Box<dyn FnOnce() -> VortexResult<Mask> + Send + 'static>),
 }
 
-impl Future for MaskExecution {
-    type Output = VortexResult<Mask>;
+impl MaskExecution {
+    pub fn lazy<F: FnOnce() -> VortexResult<Mask> + Send + 'static>(f: F) -> MaskExecution {
+        MaskExecution::Lazy(Box::new(f))
+    }
 
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match self.get_mut() {
-            MaskExecution::AllTrue(len) => {
-                let mask = Mask::new_true(*len);
-                Poll::Ready(Ok(mask))
-            }
-            MaskExecution::AllFalse(len) => {
-                let mask = Mask::new_false(*len);
-                Poll::Ready(Ok(mask))
-            }
-            MaskExecution::Future(fut) => fut.poll_unpin(cx),
+    pub fn execute(self) -> VortexResult<Mask> {
+        match self {
+            MaskExecution::AllTrue(len) => Ok(Mask::new_true(len)),
+            MaskExecution::AllFalse(len) => Ok(Mask::new_false(len)),
+            MaskExecution::Lazy(f) => f(),
         }
     }
 }
@@ -87,12 +78,9 @@ impl dyn BindCtx + '_ {
 
         // If none of the above patterns match, we fall back to canonicalizing.
         let execution = self.bind(mask, None)?;
-        Ok(MaskExecution::Future(
-            async move {
-                let mask = execution.await?.into_bool();
-                Ok(Mask::from(mask.bits().clone()))
-            }
-            .boxed(),
-        ))
+        Ok(MaskExecution::lazy(move || {
+            let mask = execution.execute()?.into_bool();
+            Ok(Mask::from(mask.bits().clone()))
+        }))
     }
 }

vortex-array/src/execution/validity.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use futures::future::FutureExt;
 use vortex_compute::filter::Filter;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
@@ -27,23 +26,22 @@ impl dyn BindCtx + '_ {
             Some(selection) => {
                 let selection = self.bind_mask(selection)?;
                 match validity {
-                    Validity::NonNullable | Validity::AllValid => Ok(MaskExecution::Future(
-                        async move { Ok(Mask::AllTrue(selection.await?.true_count())) }.boxed(),
-                    )),
-                    Validity::AllInvalid => Ok(MaskExecution::Future(
-                        async move { Ok(Mask::AllFalse(selection.await?.true_count())) }.boxed(),
-                    )),
+                    Validity::NonNullable | Validity::AllValid => {
+                        Ok(MaskExecution::lazy(move || {
+                            Ok(Mask::AllTrue(selection.execute()?.true_count()))
+                        }))
+                    }
+                    Validity::AllInvalid => Ok(MaskExecution::lazy(move || {
+                        Ok(Mask::AllFalse(selection.execute()?.true_count()))
+                    })),
                     Validity::Array(validity) => {
                         let validity = self.bind_mask(validity)?;
-                        Ok(MaskExecution::Future(
-                            async move {
-                                let validity = validity.await?;
-                                let selection = selection.await?;
-                                // We perform a take on the validity mask using the selection mask.
-                                Ok(validity.filter(&selection))
-                            }
-                            .boxed(),
-                        ))
+                        Ok(MaskExecution::lazy(move || {
+                            let validity = validity.execute()?;
+                            let selection = selection.execute()?;
+                            // We perform a take on the validity mask using the selection mask.
+                            Ok(validity.filter(&selection))
+                        }))
                     }
                 }
             }