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 = "1"
 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/constant/compute.rs

@@ -1,13 +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::boolean::{and, AndFn, OrFn};
 use crate::compute::unary::scalar_at::ScalarAtFn;
 use crate::compute::{ArrayCompute, SliceFn, TakeFn};
 use crate::compute::{SearchResult, SearchSortedFn, SearchSortedSide};
-use crate::{Array, IntoArray};
+use crate::{Array, ArrayDType, AsArray, IntoArray, IntoArrayVariant};
 
 impl ArrayCompute for ConstantArray {
     fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
@@ -25,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 {
@@ -58,6 +68,48 @@ 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 {
+            let lhs = constant_array.clone().into_bool()?;
+            let rhs = other.clone().into_bool()?;
+
+            and(lhs.as_array_ref(), rhs.as_array_ref())
+        }
+    } 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;
 use vortex_dtype::NativePType;
 
@@ -11,7 +10,7 @@ use crate::IntoArray;
 
 pub struct VarBinBuilder<O: NativePType> {
     offsets: Vec<O>,
-    data: Vec<u8>,
+    data: BytesMut,
     validity: NullBufferBuilder,
 }
 
@@ -21,7 +20,7 @@ impl<O: NativePType> VarBinBuilder<O> {
         offsets.push(O::zero());
         Self {
             offsets,
-            data: Vec::new(),
+            data: BytesMut::new(),
             validity: NullBufferBuilder::new(len),
         }
     }
@@ -48,9 +47,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

@@ -215,10 +215,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,44 @@
+use arrow_array::cast::AsArray;
+use vortex_error::VortexResult;
+
+use crate::{arrow::FromArrowArray, Array, ArrayData, IntoArray, IntoArrayVariant, IntoCanonical};
+
+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;
+    }
+
+    let lhs = lhs.clone().into_bool()?.into_canonical()?.into_arrow();
+    let lhs_bool = lhs.as_boolean();
+    let rhs = rhs.clone().into_bool()?.into_canonical()?.into_arrow();
+    let rhs_bool = rhs.as_boolean();
+
+    let data =
+        ArrayData::from_arrow(&arrow_arith::boolean::and(lhs_bool, rhs_bool)?, true).into_array();
+
+    Ok(data)
+}
+
+pub fn or(lhs: &Array, rhs: &Array) -> VortexResult<Array> {
+    if let Some(selection) = lhs.with_dyn(|lhs| lhs.and().map(|lhs| lhs.and(rhs))) {
+        return selection;
+    }
+
+    let lhs = lhs.clone().into_bool()?.into_canonical()?.into_arrow();
+    let lhs_bool = lhs.as_boolean();
+    let rhs = rhs.clone().into_bool()?.into_canonical()?.into_arrow();
+    let rhs_bool = rhs.as_boolean();
+
+    let data =
+        ArrayData::from_arrow(&arrow_arith::boolean::or(lhs_bool, rhs_bool)?, true).into_array();
+
+    Ok(data)
+}

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.
 
+use boolean::{AndFn, OrFn};
 pub use compare::{compare, CompareFn};
 pub use filter::{filter, FilterFn};
 pub use filter_indices::{filter_indices, FilterIndicesFn};
@@ -21,10 +22,11 @@ use unary::scalar_subtract::SubtractScalarFn;
 mod compare;
 mod filter;
 mod filter_indices;
+mod search_sorted;
 mod slice;
 mod take;
 
-mod search_sorted;
+pub mod boolean;
 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,35 +1,29 @@
 use datafusion_expr::{Expr, Operator as DFOperator};
 use vortex::{
-    array::{bool::BoolArray, constant::ConstantArray},
+    array::constant::ConstantArray,
+    compute::boolean::{and, or},
     compute::compare,
-    Array, IntoArray, IntoArrayVariant,
+    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),

vortex-datafusion/src/lib.rs

@@ -299,9 +299,10 @@ fn make_filter_then_take_plan(
 
 fn supported_data_types(dt: DataType) -> bool {
     dt.is_integer()
-        || dt.is_floating()
         || dt.is_signed_integer()
+        || dt.is_floating()
         || dt.is_null()
+        || dt == DataType::Boolean
         || dt == DataType::Binary
         || dt == DataType::Utf8
         || dt == DataType::Binary
@@ -474,7 +475,7 @@ mod test {
     use datafusion::functions_aggregate::count::count_distinct;
     use datafusion::prelude::SessionContext;
     use datafusion_common::{Column, TableReference};
-    use datafusion_expr::{col, lit, BinaryExpr, Expr, Operator};
+    use datafusion_expr::{and, col, lit, BinaryExpr, Expr, Operator};
     use vortex::array::primitive::PrimitiveArray;
     use vortex::array::struct_::StructArray;
     use vortex::array::varbin::VarBinArray;
@@ -618,4 +619,10 @@ mod test {
 
         assert!(!can_be_pushed_down(&e));
     }
+
+    #[test]
+    fn test_can_be_pushed_down4() {
+        let e = and((col("a")).eq(lit(2u64)), col("b").eq(lit(true)));
+        assert!(can_be_pushed_down(&e));
+    }
 }

vortex-dtype/src/dtype.rs

@@ -83,6 +83,10 @@ impl DType {
     pub fn is_struct(&self) -> bool {
         matches!(self, Struct(_, _))
     }
+
+    pub fn is_boolean(&self) -> bool {
+        matches!(self, Bool(_))
+    }
 }
 
 impl Display for DType {