Add and/or compute functions

Cargo.toml

@@ -20,14 +20,12 @@ resolver = "2"
 version = "0.1.0"
 homepage = "https://github.com/fulcrum-so/vortex"
 repository = "https://github.com/fulcrum-so/vortex"
-authors = ["Robert Kruszewski <[email protected]>, Nicholas Gates <[email protected]>, Will Manning <[email protected]>"]
+authors = [
+    "Robert Kruszewski <[email protected]>, Nicholas Gates <[email protected]>, Will Manning <[email protected]>",
+]
 license = "Apache-2.0"
 keywords = ["vortex"]
-include = [
-    "benches/*.rs",
-    "src/**/*.rs",
-    "Cargo.toml",
-]
+include = ["benches/*.rs", "src/**/*.rs", "Cargo.toml"]
 edition = "2021"
 rust-version = "1.76"
 
@@ -36,6 +34,7 @@ ahash = "0.8.11"
 allocator-api2 = "0.2.16"
 arrayref = "0.3.7"
 arrow = { version = "52.0.0", features = ["pyarrow"] }
+arrow-arith = "52.0.0"
 arrow-array = "52.0.0"
 arrow-buffer = "52.0.0"
 arrow-cast = "52.0.0"
@@ -116,3 +115,6 @@ warnings = "deny"
 [workspace.lints.clippy]
 all = { level = "deny", priority = -1 }
 or_fun_call = "deny"
+
+[profile.release]
+debug = true

vortex-array/Cargo.toml

@@ -25,6 +25,8 @@ arrow-cast = { workspace = true }
 arrow-select = { workspace = true }
 arrow-schema = { workspace = true }
 arrow-ord = { workspace = true }
+arrow-arith = { workspace = true }
+bytes = { workspace = true }
 enum-iterator = { workspace = true }
 flatbuffers = { workspace = true }
 flexbuffers = { workspace = true }
@@ -43,7 +45,10 @@ vortex-dtype = { path = "../vortex-dtype", features = ["flatbuffers", "serde"] }
 vortex-error = { path = "../vortex-error", features = ["flexbuffers"] }
 vortex-expr = { path = "../vortex-expr" }
 vortex-flatbuffers = { path = "../vortex-flatbuffers" }
-vortex-scalar = { path = "../vortex-scalar", features = ["flatbuffers", "serde"] }
+vortex-scalar = { path = "../vortex-scalar", features = [
+    "flatbuffers",
+    "serde",
+] }
 serde = { workspace = true, features = ["derive"] }
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]

vortex-array/src/array/bool/compute/boolean.rs

@@ -0,0 +1,36 @@
+use arrow_arith::boolean;
+use arrow_array::cast::AsArray;
+use vortex_error::VortexResult;
+
+use crate::array::bool::BoolArray;
+use crate::arrow::FromArrowArray;
+use crate::compute::{AndFn, OrFn};
+use crate::{Array, ArrayData, IntoArray, IntoCanonical};
+
+impl OrFn for BoolArray {
+    fn or(&self, array: &Array) -> VortexResult<Array> {
+        let lhs = self.clone().into_canonical()?.into_arrow();
+        let lhs = lhs.as_boolean();
+
+        let rhs = array.clone().into_canonical()?.into_arrow();
+        let rhs = rhs.as_boolean();
+
+        let array = boolean::or(lhs, rhs)?;
+
+        Ok(ArrayData::from_arrow(&array, true).into_array())
+    }
+}
+
+impl AndFn for BoolArray {
+    fn and(&self, array: &crate::Array) -> vortex_error::VortexResult<crate::Array> {
+        let lhs = self.clone().into_canonical()?.into_arrow();
+        let lhs = lhs.as_boolean();
+
+        let rhs = array.clone().into_canonical()?.into_arrow();
+        let rhs = rhs.as_boolean();
+
+        let array = boolean::and(lhs, rhs)?;
+
+        Ok(ArrayData::from_arrow(&array, true).into_array())
+    }
+}

vortex-array/src/array/bool/compute/mod.rs

@@ -3,6 +3,7 @@ use crate::compute::unary::fill_forward::FillForwardFn;
 use crate::compute::unary::scalar_at::ScalarAtFn;
 use crate::compute::{ArrayCompute, CompareFn, SliceFn, TakeFn};
 
+mod boolean;
 mod compare;
 mod fill;
 mod flatten;
@@ -30,4 +31,12 @@ impl ArrayCompute for BoolArray {
     fn take(&self) -> Option<&dyn TakeFn> {
         Some(self)
     }
+
+    fn and(&self) -> Option<&dyn crate::compute::AndFn> {
+        Some(self)
+    }
+
+    fn or(&self) -> Option<&dyn crate::compute::OrFn> {
+        Some(self)
+    }
 }

vortex-array/src/array/constant/compute.rs

@@ -1,14 +1,15 @@
 use std::cmp::Ordering;
 
-use vortex_error::VortexResult;
+use vortex_dtype::Nullability;
+use vortex_error::{vortex_bail, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::array::constant::ConstantArray;
 use crate::compute::unary::scalar_at::ScalarAtFn;
 use crate::compute::{
-    ArrayCompute, SearchResult, SearchSortedFn, SearchSortedSide, SliceFn, TakeFn,
+    AndFn, ArrayCompute, OrFn, SearchResult, SearchSortedFn, SearchSortedSide, SliceFn, TakeFn,
 };
-use crate::{Array, IntoArray};
+use crate::{Array, ArrayDType, IntoArray, IntoArrayVariant};
 
 impl ArrayCompute for ConstantArray {
     fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
@@ -26,6 +27,14 @@ impl ArrayCompute for ConstantArray {
     fn take(&self) -> Option<&dyn TakeFn> {
         Some(self)
     }
+
+    fn and(&self) -> Option<&dyn AndFn> {
+        Some(self)
+    }
+
+    fn or(&self) -> Option<&dyn OrFn> {
+        Some(self)
+    }
 }
 
 impl ScalarAtFn for ConstantArray {
@@ -59,6 +68,45 @@ impl SearchSortedFn for ConstantArray {
     }
 }
 
+impl AndFn for ConstantArray {
+    fn and(&self, array: &Array) -> VortexResult<Array> {
+        constant_array_bool_impl(self, array, |(l, r)| l & r)
+    }
+}
+
+impl OrFn for ConstantArray {
+    fn or(&self, array: &Array) -> VortexResult<Array> {
+        constant_array_bool_impl(self, array, |(l, r)| l | r)
+    }
+}
+
+fn constant_array_bool_impl(
+    constant_array: &ConstantArray,
+    other: &Array,
+    bool_op: impl Fn((bool, bool)) -> bool,
+) -> VortexResult<Array> {
+    if constant_array.dtype().is_boolean()
+        && other.dtype().is_boolean()
+        && constant_array.len() == other.len()
+    {
+        if let Ok(array) = ConstantArray::try_from(other.clone()) {
+            let lhs = constant_array.scalar().value().as_bool()?;
+            let rhs = array.scalar().value().as_bool()?;
+
+            let scalar = match lhs.zip(rhs).map(bool_op) {
+                Some(b) => Scalar::bool(b, Nullability::Nullable),
+                None => Scalar::null(constant_array.dtype().as_nullable()),
+            };
+
+            Ok(ConstantArray::new(scalar, constant_array.len()).into_array())
+        } else {
+            AndFn::and(&constant_array.clone().into_bool()?, other)
+        }
+    } else {
+        vortex_bail!("Boolean operations aren't supported on arrays of different lengths")
+    }
+}
+
 #[cfg(test)]
 mod test {
     use crate::array::constant::ConstantArray;

vortex-array/src/array/varbin/builder.rs

@@ -1,6 +1,5 @@
-use std::mem;
-
 use arrow_buffer::NullBufferBuilder;
+use bytes::BytesMut;
 use vortex_dtype::{DType, NativePType};
 
 use crate::array::primitive::PrimitiveArray;
@@ -10,7 +9,7 @@ use crate::IntoArray;
 
 pub struct VarBinBuilder<O: NativePType> {
     offsets: Vec<O>,
-    data: Vec<u8>,
+    data: BytesMut,
     validity: NullBufferBuilder,
 }
 
@@ -20,7 +19,7 @@ impl<O: NativePType> VarBinBuilder<O> {
         offsets.push(O::zero());
         Self {
             offsets,
-            data: Vec::new(),
+            data: BytesMut::new(),
             validity: NullBufferBuilder::new(len),
         }
     }
@@ -47,9 +46,9 @@ impl<O: NativePType> VarBinBuilder<O> {
         self.validity.append_null();
     }
 
-    pub fn finish(&mut self, dtype: DType) -> VarBinArray {
-        let offsets = PrimitiveArray::from(mem::take(&mut self.offsets));
-        let data = PrimitiveArray::from(mem::take(&mut self.data));
+    pub fn finish(mut self, dtype: DType) -> VarBinArray {
+        let offsets = PrimitiveArray::from(self.offsets);
+        let data = PrimitiveArray::from(Vec::from(self.data.freeze()));
 
         let nulls = self.validity.finish();
 

vortex-array/src/canonical.rs

@@ -222,10 +222,8 @@ fn varbin_to_arrow(varbin_array: VarBinArray) -> ArrayRef {
         PType::I32 | PType::I64 => offsets,
         // Unless it's u64, everything else can be converted into an i32.
         // FIXME(ngates): do not copy offsets again
-        PType::U64 => try_cast(&offsets.to_array(), PType::I64.into())
-            .expect("cast to i64")
-            .into_primitive()
-            .expect("flatten_primitive"),
+        PType::U64 => offsets.reinterpret_cast(PType::I64),
+        PType::U32 => offsets.reinterpret_cast(PType::I32),
         _ => try_cast(&offsets.to_array(), PType::I32.into())
             .expect("cast to i32")
             .into_primitive()

vortex-array/src/compute/boolean.rs

@@ -0,0 +1,39 @@
+use vortex_error::VortexResult;
+
+use crate::{Array, IntoArrayVariant};
+
+pub trait AndFn {
+    fn and(&self, array: &Array) -> VortexResult<Array>;
+}
+
+pub trait OrFn {
+    fn or(&self, array: &Array) -> VortexResult<Array>;
+}
+
+pub fn and(lhs: &Array, rhs: &Array) -> VortexResult<Array> {
+    if let Some(selection) = lhs.with_dyn(|lhs| lhs.and().map(|lhs| lhs.and(rhs))) {
+        return selection;
+    }
+
+    if let Some(selection) = rhs.with_dyn(|rhs| rhs.and().map(|rhs| rhs.and(lhs))) {
+        return selection;
+    }
+
+    let lhs = lhs.clone().into_bool()?;
+
+    lhs.and(rhs)
+}
+
+pub fn or(lhs: &Array, rhs: &Array) -> VortexResult<Array> {
+    if let Some(selection) = lhs.with_dyn(|lhs| lhs.or().map(|lhs| lhs.or(rhs))) {
+        return selection;
+    }
+
+    if let Some(selection) = rhs.with_dyn(|rhs| rhs.or().map(|rhs| rhs.or(lhs))) {
+        return selection;
+    }
+
+    let lhs = lhs.clone().into_bool()?;
+
+    lhs.or(rhs)
+}

vortex-array/src/compute/mod.rs

@@ -7,6 +7,7 @@
 //! implementations of these operators, else we will decode, and perform the equivalent operator
 //! from Arrow.
 
+pub use boolean::{and, or, AndFn, OrFn};
 pub use compare::{compare, CompareFn};
 pub use filter::{filter, FilterFn};
 pub use filter_indices::{filter_indices, FilterIndicesFn};
@@ -18,13 +19,14 @@ use unary::fill_forward::FillForwardFn;
 use unary::scalar_at::ScalarAtFn;
 use unary::scalar_subtract::SubtractScalarFn;
 
+mod boolean;
 mod compare;
 mod filter;
 mod filter_indices;
+mod search_sorted;
 mod slice;
 mod take;
 
-mod search_sorted;
 pub mod unary;
 
 /// Trait providing compute functions on top of Vortex arrays.
@@ -98,4 +100,18 @@ pub trait ArrayCompute {
     fn take(&self) -> Option<&dyn TakeFn> {
         None
     }
+
+    /// Perform a boolean AND operation over two arrays
+    ///
+    /// See: [AndFn].
+    fn and(&self) -> Option<&dyn AndFn> {
+        None
+    }
+
+    /// Perform a boolean OR operation over two arrays
+    ///
+    /// See: [OrFn].
+    fn or(&self) -> Option<&dyn OrFn> {
+        None
+    }
 }

vortex-datafusion/src/eval.rs

@@ -1,34 +1,26 @@
 use datafusion_expr::{Expr, Operator as DFOperator};
-use vortex::array::bool::BoolArray;
 use vortex::array::constant::ConstantArray;
-use vortex::compute::compare;
-use vortex::{Array, IntoArray, IntoArrayVariant};
+use vortex::compute::{and, compare, or};
+use vortex::{Array, IntoArray};
 use vortex_error::{vortex_bail, vortex_err, VortexResult};
 use vortex_expr::Operator;
 
+use crate::can_be_pushed_down;
+
 pub struct ExpressionEvaluator;
 
 impl ExpressionEvaluator {
     pub fn eval(array: Array, expr: &Expr) -> VortexResult<Array> {
+        debug_assert!(can_be_pushed_down(expr));
+
         match expr {
             Expr::BinaryExpr(expr) => {
                 let lhs = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
                 let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
-
                 // TODO(adamg): turn and/or into more general compute functions
                 match expr.op {
-                    DFOperator::And => {
-                        let lhs = lhs.into_bool()?;
-                        let rhs = rhs.into_bool()?;
-                        let buffer = &lhs.boolean_buffer() & &rhs.boolean_buffer();
-                        Ok(BoolArray::from(buffer).into_array())
-                    }
-                    DFOperator::Or => {
-                        let lhs = lhs.into_bool()?;
-                        let rhs = rhs.into_bool()?;
-                        let buffer = &lhs.boolean_buffer() | &rhs.boolean_buffer();
-                        Ok(BoolArray::from(buffer).into_array())
-                    }
+                    DFOperator::And => and(&lhs, &rhs),
+                    DFOperator::Or => or(&lhs, &rhs),
                     DFOperator::Eq => compare(&lhs, &rhs, Operator::Eq),
                     DFOperator::Gt => compare(&lhs, &rhs, Operator::Gt),
                     DFOperator::GtEq => compare(&lhs, &rhs, Operator::Gte),