Skip to content

Commit

Permalink
New compute functions, no more shortcircuting
Browse files Browse the repository at this point in the history
  • Loading branch information
AdamGS committed Jul 18, 2024
1 parent 952d883 commit acef812
Showing 6 changed files with 46 additions and 74 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
@@ -34,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"
1 change: 1 addition & 0 deletions vortex-array/Cargo.toml
Original file line number Diff line number Diff line change
@@ -25,6 +25,7 @@ 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 }
4 changes: 0 additions & 4 deletions vortex-array/src/array/bool/mod.rs
Original file line number Diff line number Diff line change
@@ -72,10 +72,6 @@ impl BoolArray {
let buffer = BooleanBuffer::from(bools);
Self::try_new(buffer, validity).unwrap()
}

pub fn true_count(&self) -> usize {
self.statistics().compute_true_count().unwrap()
}
}

impl ArrayTrait for BoolArray {}
29 changes: 29 additions & 0 deletions vortex-array/src/compute/mod.rs
Original file line number Diff line number Diff line change
@@ -7,6 +7,7 @@
//! implementations of these operators, else we will decode, and perform the equivalent operator
//! from Arrow.
use arrow_array::cast::AsArray;
pub use compare::{compare, CompareFn};
pub use filter::{filter, FilterFn};
pub use filter_indices::{filter_indices, FilterIndicesFn};
@@ -17,6 +18,10 @@ use unary::cast::CastFn;
use unary::fill_forward::FillForwardFn;
use unary::scalar_at::ScalarAtFn;
use unary::scalar_subtract::SubtractScalarFn;
use vortex_error::VortexResult;

use crate::arrow::FromArrowArray;
use crate::{Array, ArrayData, IntoArray, IntoArrayVariant, IntoCanonical};

mod compare;
mod filter;
@@ -99,3 +104,27 @@ pub trait ArrayCompute {
None
}
}

pub fn and(lhs: &Array, rhs: &Array) -> VortexResult<Array> {
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> {
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)
}
84 changes: 14 additions & 70 deletions vortex-datafusion/src/eval.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use datafusion_expr::{Expr, Operator as DFOperator};
use vortex::{
array::{bool::BoolArray, constant::ConstantArray},
compute::compare,
Array, IntoArray, IntoArrayVariant,
array::constant::ConstantArray,
compute::{and, compare, or},
Array, IntoArray,
};
use vortex_error::{vortex_bail, vortex_err, VortexResult};
use vortex_expr::Operator;
@@ -14,77 +14,21 @@ pub struct ExpressionEvaluator;
impl ExpressionEvaluator {
pub fn eval(array: Array, expr: &Expr) -> VortexResult<Array> {
debug_assert!(can_be_pushed_down(expr));
let original_len = array.len();

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 = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
let lhs = lhs.into_bool()?;

if lhs.true_count() == 0 {
return Ok(ConstantArray::new(false, original_len).into_array());
}

let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
let rhs = rhs.into_bool()?;

if rhs.true_count() == 0 {
return Ok(ConstantArray::new(false, original_len).into_array());
}

let buffer = &lhs.boolean_buffer() & &rhs.boolean_buffer();
Ok(BoolArray::from(buffer).into_array())
}
DFOperator::Or => {
let lhs = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
let lhs = lhs.into_bool()?;

if lhs.true_count() == original_len {
return Ok(ConstantArray::new(true, original_len).into_array());
}

let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
let rhs = rhs.into_bool()?;

if lhs.true_count() == original_len {
return Ok(ConstantArray::new(true, original_len).into_array());
}

let buffer = &lhs.boolean_buffer() | &rhs.boolean_buffer();
Ok(BoolArray::from(buffer).into_array())
}
DFOperator::Eq => {
let lhs = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
compare(&lhs, &rhs, Operator::Eq)
}
DFOperator::Gt => {
let lhs = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
compare(&lhs, &rhs, Operator::Gt)
}
DFOperator::GtEq => {
let lhs = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
compare(&lhs, &rhs, Operator::Gte)
}
DFOperator::Lt => {
let lhs = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
compare(&lhs, &rhs, Operator::Lt)
}
DFOperator::LtEq => {
let lhs = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
compare(&lhs, &rhs, Operator::Lte)
}
DFOperator::NotEq => {
let lhs = ExpressionEvaluator::eval(array.clone(), expr.left.as_ref())?;
let rhs = ExpressionEvaluator::eval(array, expr.right.as_ref())?;
compare(&lhs, &rhs, Operator::NotEq)
}
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),
DFOperator::Lt => compare(&lhs, &rhs, Operator::Lt),
DFOperator::LtEq => compare(&lhs, &rhs, Operator::Lte),
DFOperator::NotEq => compare(&lhs, &rhs, Operator::NotEq),
_ => vortex_bail!("{} is an unsupported operator", expr.op),
}
}

0 comments on commit acef812

Please sign in to comment.