Merge remote-tracking branch 'origin/develop' into ji/assert-array-eq-more-places

Author: joseph-isaacs

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/arrays/primitive/vtable/mod.rs

@@ -8,7 +8,7 @@ use crate::{EncodingId, EncodingRef, vtable};
 mod array;
 mod canonical;
 mod operations;
-mod pipeline;
+mod operator;
 mod serde;
 mod validity;
 mod visitor;
@@ -26,8 +26,8 @@ impl VTable for PrimitiveVTable {
     type VisitorVTable = Self;
     type ComputeVTable = NotSupported;
     type EncodeVTable = NotSupported;
-    type OperatorVTable = Self;
     type SerdeVTable = Self;
+    type OperatorVTable = Self;
 
     fn id(_encoding: &Self::Encoding) -> EncodingId {
         EncodingId::new_ref("vortex.primitive")

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

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_compute::filter::Filter;
+use vortex_dtype::match_each_native_ptype;
+use vortex_error::VortexResult;
+use vortex_vector::PVector;
+
+use crate::ArrayRef;
+use crate::arrays::{PrimitiveArray, PrimitiveVTable};
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::vtable::{OperatorVTable, ValidityHelper};
+
+impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
+    fn bind(
+        array: &PrimitiveArray,
+        selection: Option<&ArrayRef>,
+        ctx: &mut dyn BindCtx,
+    ) -> VortexResult<BatchKernelRef> {
+        let mask = ctx.bind_selection(array.len(), selection)?;
+        let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
+
+        match_each_native_ptype!(array.ptype(), |T| {
+            let elements = array.buffer::<T>();
+            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
+                // the elements.
+                let elements = elements.filter(&mask);
+
+                Ok(PVector::new(elements, validity).into())
+            }))
+        })
+    }
+}

vortex-array/src/arrays/primitive/vtable/pipeline.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]);
     }
 }