16139349253

Committed 08 Jul 2025 09:26AM UTC coverage: 78.057% (-0.2%) from 78.253%

Build # 16139349253

Build Type

push

github

Committed by

web-flow

Commit Message

VortexExpr VTables (#3713)

Adds the same vtable machinery as arrays and layouts already use. It
uses the "Encoding" naming scheme from arrays and layouts. I don't
particularly like it, but it's consistent. Open to renames later.

Further, adds an expression registry to the Vortex session that will be
used for deserialization.

Expressions only decide their "options" serialization. So in theory, can
support many container formats, not just proto, provided each expression
can deserialize their own options format.

---------

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Run Details

800 of 1190 new or added lines in 38 files covered. (67.23%)

40 existing lines in 13 files now uncovered.

44100 of 56497 relevant lines covered (78.06%)

54989.55 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

90.69

/vortex-expr/src/exprs/binary.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::fmt::Display;
use std::hash::Hash;

use vortex_array::compute::{Operator as ArrayOperator, add, and_kleene, compare, or_kleene};
use vortex_array::{ArrayRef, DeserializeMetadata, ProstMetadata};
use vortex_dtype::DType;
use vortex_error::{VortexResult, vortex_bail};
use vortex_proto::expr as pb;

use crate::{
    AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Operator, Scope, ScopeDType,
    StatsCatalog, VTable, vtable,
};

vtable!(Binary);

#[allow(clippy::derived_hash_with_manual_eq)]
#[derive(Debug, Clone, Hash)]
pub struct BinaryExpr {
    lhs: ExprRef,
    operator: Operator,
    rhs: ExprRef,
}

impl PartialEq for BinaryExpr {
    fn eq(&self, other: &Self) -> bool {
        self.lhs.eq(&other.lhs) && self.operator == other.operator && self.rhs.eq(&other.rhs)
    }
}

pub struct BinaryExprEncoding;

impl VTable for BinaryVTable {
    type Expr = BinaryExpr;
    type Encoding = BinaryExprEncoding;
    type Metadata = ProstMetadata<pb::BinaryOpts>;

    fn id(_encoding: &Self::Encoding) -> ExprId {
        ExprId::new_ref("binary")
    }

    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
        ExprEncodingRef::new_ref(BinaryExprEncoding.as_ref())
    }

    fn metadata(expr: &Self::Expr) -> Option<Self::Metadata> {
        Some(ProstMetadata(pb::BinaryOpts {
            op: expr.operator.into(),
        }))
    }

    fn children(expr: &Self::Expr) -> Vec<&ExprRef> {
        vec![expr.lhs(), expr.rhs()]
    }

    fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
        if children.len() != 2 {
            vortex_bail!(
                "Binary expression must have exactly 2 children, got {}",
                children.len()
            );
        }

        Ok(BinaryExpr::new(
            children[0].clone(),
            expr.op(),
            children[1].clone(),
        ))
    }

    fn build(
        _encoding: &Self::Encoding,
        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
        children: Vec<ExprRef>,
    ) -> VortexResult<Self::Expr> {
        Ok(BinaryExpr::new(
            children[0].clone(),
            metadata.op().into(),
            children[1].clone(),
        ))
    }

    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
        let lhs = expr.lhs.unchecked_evaluate(scope)?;
        let rhs = expr.rhs.unchecked_evaluate(scope)?;

        match expr.operator {
            Operator::Eq => compare(&lhs, &rhs, ArrayOperator::Eq),
            Operator::NotEq => compare(&lhs, &rhs, ArrayOperator::NotEq),
            Operator::Lt => compare(&lhs, &rhs, ArrayOperator::Lt),
            Operator::Lte => compare(&lhs, &rhs, ArrayOperator::Lte),
            Operator::Gt => compare(&lhs, &rhs, ArrayOperator::Gt),
            Operator::Gte => compare(&lhs, &rhs, ArrayOperator::Gte),
            Operator::And => and_kleene(&lhs, &rhs),
            Operator::Or => or_kleene(&lhs, &rhs),
            Operator::Add => add(&lhs, &rhs),
        }
    }

    fn return_dtype(expr: &Self::Expr, scope: &ScopeDType) -> VortexResult<DType> {
        let lhs = expr.lhs.return_dtype(scope)?;
        let rhs = expr.rhs.return_dtype(scope)?;

        if expr.operator == Operator::Add {
            if lhs.is_primitive() && lhs.eq_ignore_nullability(&rhs) {
                return Ok(lhs.with_nullability(lhs.nullability() | rhs.nullability()));
            }
            vortex_bail!("incompatible types for checked add: {} {}", lhs, rhs);
        }

        Ok(DType::Bool((lhs.is_nullable() || rhs.is_nullable()).into()))
    }
}

impl BinaryExpr {
    pub fn new(lhs: ExprRef, operator: Operator, rhs: ExprRef) -> Self {
        Self { lhs, operator, rhs }
    }

    pub fn new_expr(lhs: ExprRef, operator: Operator, rhs: ExprRef) -> ExprRef {
        Self::new(lhs, operator, rhs).into_expr()
    }

    pub fn lhs(&self) -> &ExprRef {
        &self.lhs
    }

    pub fn rhs(&self) -> &ExprRef {
        &self.rhs
    }

    pub fn op(&self) -> Operator {
        self.operator
    }
}

impl Display for BinaryExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "({} {} {})", self.lhs, self.operator, self.rhs)
    }
}

impl AnalysisExpr for BinaryExpr {
    fn stat_falsification(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
        match self.operator {
            Operator::Eq => {
                let min_lhs = self.lhs.min(catalog);
                let max_lhs = self.lhs.max(catalog);

                let min_rhs = self.rhs.min(catalog);
                let max_rhs = self.rhs.max(catalog);

                let left = min_lhs.zip(max_rhs).map(|(a, b)| gt(a, b));
                let right = min_rhs.zip(max_lhs).map(|(a, b)| gt(a, b));
                left.into_iter().chain(right).reduce(or)
            }
            Operator::NotEq => {
                let min_lhs = self.lhs.min(catalog)?;
                let max_lhs = self.lhs.max(catalog)?;

                let min_rhs = self.rhs.min(catalog)?;
                let max_rhs = self.rhs.max(catalog)?;

                Some(and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs)))
            }
            Operator::Gt => Some(lt_eq(self.lhs.max(catalog)?, self.rhs.min(catalog)?)),
            Operator::Gte => Some(lt(self.lhs.max(catalog)?, self.rhs.min(catalog)?)),
            Operator::Lt => Some(gt_eq(self.lhs.min(catalog)?, self.rhs.max(catalog)?)),
            Operator::Lte => Some(gt(self.lhs.min(catalog)?, self.rhs.max(catalog)?)),
            Operator::And => self
                .lhs
                .stat_falsification(catalog)
                .into_iter()
                .chain(self.rhs.stat_falsification(catalog))
                .reduce(or),
            Operator::Or => Some(and(
                self.lhs.stat_falsification(catalog)?,
                self.rhs.stat_falsification(catalog)?,
            )),
            Operator::Add => None,
        }
    }
}

/// Create a new `BinaryExpr` using the `Eq` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
/// use vortex_array::{Array, IntoArray, ToCanonical};
/// use vortex_array::validity::Validity;
/// use vortex_buffer::buffer;
/// use vortex_expr::{eq, root, lit, Scope};
///
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
/// let result = eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
///
/// assert_eq!(
///     result.to_bool().unwrap().boolean_buffer(),
///     BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
/// );
/// ```
pub fn eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::Eq, rhs).into_expr()
}

/// Create a new `BinaryExpr` using the `NotEq` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
/// use vortex_array::{IntoArray, ToCanonical};
/// use vortex_array::validity::Validity;
/// use vortex_buffer::buffer;
/// use vortex_expr::{root, lit, not_eq, Scope};
///
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
/// let result = not_eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
///
/// assert_eq!(
///     result.to_bool().unwrap().boolean_buffer(),
///     BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
/// );
/// ```
pub fn not_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::NotEq, rhs).into_expr()
}

/// Create a new `BinaryExpr` using the `Gte` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
/// use vortex_array::{IntoArray, ToCanonical};
/// use vortex_array::validity::Validity;
/// use vortex_buffer::buffer;
/// use vortex_expr::{gt_eq, root, lit, Scope};
///
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
/// let result = gt_eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
///
/// assert_eq!(
///     result.to_bool().unwrap().boolean_buffer(),
///     BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
/// );
/// ```
pub fn gt_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::Gte, rhs).into_expr()
}

/// Create a new `BinaryExpr` using the `Gt` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
/// use vortex_array::{IntoArray, ToCanonical};
/// use vortex_array::validity::Validity;
/// use vortex_buffer::buffer;
/// use vortex_expr::{gt, root, lit, Scope};
///
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
/// let result = gt(root(), lit(2)).evaluate(&Scope::new(xs.to_array())).unwrap();
///
/// assert_eq!(
///     result.to_bool().unwrap().boolean_buffer(),
///     BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
/// );
/// ```
pub fn gt(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::Gt, rhs).into_expr()
}

/// Create a new `BinaryExpr` using the `Lte` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
/// use vortex_array::{IntoArray, ToCanonical};
/// use vortex_array::validity::Validity;
/// use vortex_buffer::buffer;
/// use vortex_expr::{root, lit, lt_eq, Scope};
///
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
/// let result = lt_eq(root(), lit(2)).evaluate(&Scope::new(xs.to_array())).unwrap();
///
/// assert_eq!(
///     result.to_bool().unwrap().boolean_buffer(),
///     BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
/// );
/// ```
pub fn lt_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::Lte, rhs).into_expr()
}

/// Create a new `BinaryExpr` using the `Lt` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
/// use vortex_array::{IntoArray, ToCanonical};
/// use vortex_array::validity::Validity;
/// use vortex_buffer::buffer;
/// use vortex_expr::{root, lit, lt, Scope};
///
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
/// let result = lt(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
///
/// assert_eq!(
///     result.to_bool().unwrap().boolean_buffer(),
///     BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
/// );
/// ```
pub fn lt(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::Lt, rhs).into_expr()
}

/// Create a new `BinaryExpr` using the `Or` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::arrays::BoolArray;
/// use vortex_array::{IntoArray, ToCanonical};
/// use vortex_expr::{root, lit, or, Scope};
///
/// let xs = BoolArray::from_iter(vec![true, false, true]);
/// let result = or(root(), lit(false)).evaluate(&Scope::new(xs.to_array())).unwrap();
///
/// assert_eq!(
///     result.to_bool().unwrap().boolean_buffer(),
///     BoolArray::from_iter(vec![true, false, true]).boolean_buffer(),
/// );
/// ```
pub fn or(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::Or, rhs).into_expr()
}

/// Collects a list of `or`ed values into a single vortex, expr
/// [x, y, z] => x or (y or z)
pub fn or_collect<I>(iter: I) -> Option<ExprRef>
where
    I: IntoIterator<Item = ExprRef>,
    I::IntoIter: DoubleEndedIterator<Item = ExprRef>,
{
    let mut iter = iter.into_iter();
    let first = iter.next_back()?;
    Some(iter.rfold(first, |acc, elem| or(elem, acc)))
}

/// Create a new `BinaryExpr` using the `And` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::arrays::BoolArray;
/// use vortex_array::{IntoArray, ToCanonical};
/// use vortex_expr::{and, root, lit, Scope};
///
/// let xs = BoolArray::from_iter(vec![true, false, true]);
/// let result = and(root(), lit(true)).evaluate(&Scope::new(xs.to_array())).unwrap();
///
/// assert_eq!(
///     result.to_bool().unwrap().boolean_buffer(),
///     BoolArray::from_iter(vec![true, false, true]).boolean_buffer(),
/// );
/// ```
pub fn and(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::And, rhs).into_expr()
}

/// Collects a list of `and`ed values into a single vortex, expr
/// [x, y, z] => x and (y and z)
pub fn and_collect<I>(iter: I) -> Option<ExprRef>
where
    I: IntoIterator<Item = ExprRef>,
    I::IntoIter: DoubleEndedIterator<Item = ExprRef>,
{
    let mut iter = iter.into_iter();
    let first = iter.next_back()?;
    Some(iter.rfold(first, |acc, elem| and(elem, acc)))
}

/// Collects a list of `and`ed values into a single vortex, expr
/// [x, y, z] => x and (y and z)
pub fn and_collect_right<I>(iter: I) -> Option<ExprRef>
where
    I: IntoIterator<Item = ExprRef>,
{
    let iter = iter.into_iter();
    iter.reduce(and)
}

/// Create a new `BinaryExpr` using the `CheckedAdd` operator.
///
/// ## Example usage
///
/// ```
/// use vortex_array::IntoArray;
/// use vortex_array::arrow::IntoArrowArray as _;
/// use vortex_buffer::buffer;
/// use vortex_expr::{Scope, checked_add, lit, root};
///
/// let xs = buffer![1, 2, 3].into_array();
/// let result = checked_add(root(), lit(5))
///     .evaluate(&Scope::new(xs.to_array()))
///     .unwrap();
///
/// assert_eq!(
///     &result.into_arrow_preferred().unwrap(),
///     &buffer![6, 7, 8]
///         .into_array()
///         .into_arrow_preferred()
///         .unwrap()
/// );
/// ```
pub fn checked_add(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
    BinaryExpr::new(lhs, Operator::Add, rhs).into_expr()
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use vortex_dtype::{DType, Nullability};

    use crate::{
        ScopeDType, VortexExpr, and, and_collect, and_collect_right, col, eq, gt, gt_eq, lit, lt,
        lt_eq, not_eq, or, test_harness,
    };

    #[test]
    fn and_collect_left_assoc() {
        let values = vec![lit(1), lit(2), lit(3)];
        assert_eq!(
            Some(and(lit(1), and(lit(2), lit(3)))),
            and_collect(values.into_iter())
        );
    }

    #[test]
    fn and_collect_right_assoc() {
        let values = vec![lit(1), lit(2), lit(3)];
        assert_eq!(
            Some(and(and(lit(1), lit(2)), lit(3))),
            and_collect_right(values.into_iter())
        );
    }

    #[test]
    fn dtype() {
        let dtype = test_harness::struct_dtype();
        let bool1: Arc<dyn VortexExpr> = col("bool1");
        let bool2: Arc<dyn VortexExpr> = col("bool2");
        assert_eq!(
            and(bool1.clone(), bool2.clone())
                .return_dtype(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            DType::Bool(Nullability::NonNullable)
        );
        assert_eq!(
            or(bool1.clone(), bool2.clone())
                .return_dtype(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            DType::Bool(Nullability::NonNullable)
        );

        let col1: Arc<dyn VortexExpr> = col("col1");
        let col2: Arc<dyn VortexExpr> = col("col2");

        assert_eq!(
            eq(col1.clone(), col2.clone())
                .return_dtype(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            not_eq(col1.clone(), col2.clone())
                .return_dtype(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            gt(col1.clone(), col2.clone())
                .return_dtype(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            gt_eq(col1.clone(), col2.clone())
                .return_dtype(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            lt(col1.clone(), col2.clone())
                .return_dtype(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            lt_eq(col1.clone(), col2.clone())
                .return_dtype(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );

        assert_eq!(
            or(
                lt(col1.clone(), col2.clone()),
                not_eq(col1.clone(), col2.clone())
            )
            .return_dtype(&ScopeDType::new(dtype))
            .unwrap(),
            DType::Bool(Nullability::Nullable)
        );
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::fmt::Display;
5	use std::hash::Hash;
6
7	use vortex_array::compute::{Operator as ArrayOperator, add, and_kleene, compare, or_kleene};
8	use vortex_array::{ArrayRef, DeserializeMetadata, ProstMetadata};
9	use vortex_dtype::DType;
10	use vortex_error::{VortexResult, vortex_bail};
11	use vortex_proto::expr as pb;
12
13	use crate::{
14	AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Operator, Scope, ScopeDType,
15	StatsCatalog, VTable, vtable,
16	};
17
18	vtable!(Binary);
19
20	#[allow(clippy::derived_hash_with_manual_eq)]
21	#[derive(Debug, Clone, Hash)]
22	pub struct BinaryExpr {
23	lhs: ExprRef,
24	operator: Operator,
25	rhs: ExprRef,
26	}
27
28	impl PartialEq for BinaryExpr {
29	fn eq(&self, other: &Self) -> bool {	5,751✔
30	self.lhs.eq(&other.lhs) && self.operator == other.operator && self.rhs.eq(&other.rhs)	5,751✔
31	}	5,751✔
32	}
33
34	pub struct BinaryExprEncoding;
35
36	impl VTable for BinaryVTable {
37	type Expr = BinaryExpr;
38	type Encoding = BinaryExprEncoding;
39	type Metadata = ProstMetadata<pb::BinaryOpts>;
40
41	fn id(_encoding: &Self::Encoding) -> ExprId {	109✔
42	ExprId::new_ref("binary")	109✔
43	}	109✔
44
45	fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {	3✔
46	ExprEncodingRef::new_ref(BinaryExprEncoding.as_ref())	3✔
47	}	3✔
48
49	fn metadata(expr: &Self::Expr) -> Option<Self::Metadata> {	3✔
50	Some(ProstMetadata(pb::BinaryOpts {	3✔
51	op: expr.operator.into(),	3✔
52	}))	3✔
53	}	3✔
54
55	fn children(expr: &Self::Expr) -> Vec<&ExprRef> {	17,622✔
56	vec![expr.lhs(), expr.rhs()]	17,622✔
57	}	17,622✔
58
59	fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {	1,189✔
60	if children.len() != 2 {	1,189✔
NEW 61	vortex_bail!(	×
NEW 62	"Binary expression must have exactly 2 children, got {}",	×
NEW 63	children.len()	×
NEW 64	);	×
65	}	1,189✔
66		1,189✔
67	Ok(BinaryExpr::new(	1,189✔
68	children[0].clone(),	1,189✔
69	expr.op(),	1,189✔
70	children[1].clone(),	1,189✔
71	))	1,189✔
72	}	1,189✔
73
74	fn build(	3✔
75	_encoding: &Self::Encoding,	3✔
76	metadata: &<Self::Metadata as DeserializeMetadata>::Output,	3✔
77	children: Vec<ExprRef>,	3✔
78	) -> VortexResult<Self::Expr> {	3✔
79	Ok(BinaryExpr::new(	3✔
80	children[0].clone(),	3✔
81	metadata.op().into(),	3✔
82	children[1].clone(),	3✔
83	))	3✔
84	}	3✔
85
86	fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {	1,387✔
87	let lhs = expr.lhs.unchecked_evaluate(scope)?;	1,387✔
88	let rhs = expr.rhs.unchecked_evaluate(scope)?;	1,387✔
89
90	match expr.operator {	1,387✔
91	Operator::Eq => compare(&lhs, &rhs, ArrayOperator::Eq),	253✔
NEW 92	Operator::NotEq => compare(&lhs, &rhs, ArrayOperator::NotEq),	×
93	Operator::Lt => compare(&lhs, &rhs, ArrayOperator::Lt),	105✔
94	Operator::Lte => compare(&lhs, &rhs, ArrayOperator::Lte),	189✔
95	Operator::Gt => compare(&lhs, &rhs, ArrayOperator::Gt),	462✔
96	Operator::Gte => compare(&lhs, &rhs, ArrayOperator::Gte),	42✔
97	Operator::And => and_kleene(&lhs, &rhs),	42✔
98	Operator::Or => or_kleene(&lhs, &rhs),	294✔
NEW 99	Operator::Add => add(&lhs, &rhs),	×
100	}
101	}	1,387✔
102
103	fn return_dtype(expr: &Self::Expr, scope: &ScopeDType) -> VortexResult<DType> {	1,798✔
104	let lhs = expr.lhs.return_dtype(scope)?;	1,798✔
105	let rhs = expr.rhs.return_dtype(scope)?;	1,798✔
106
107	if expr.operator == Operator::Add {	1,798✔
NEW 108	if lhs.is_primitive() && lhs.eq_ignore_nullability(&rhs) {	×
NEW 109	return Ok(lhs.with_nullability(lhs.nullability() \| rhs.nullability()));	×
NEW 110	}	×
NEW 111	vortex_bail!("incompatible types for checked add: {} {}", lhs, rhs);	×
112	}	1,798✔
113		1,798✔
114	Ok(DType::Bool((lhs.is_nullable() \|\| rhs.is_nullable()).into()))	1,798✔
115	}	1,798✔
116	}
117
118	impl BinaryExpr {
119	pub fn new(lhs: ExprRef, operator: Operator, rhs: ExprRef) -> Self {	4,400✔
120	Self { lhs, operator, rhs }	4,400✔
121	}	4,400✔
122
123	pub fn new_expr(lhs: ExprRef, operator: Operator, rhs: ExprRef) -> ExprRef {	92✔
124	Self::new(lhs, operator, rhs).into_expr()	92✔
125	}	92✔
126
127	pub fn lhs(&self) -> &ExprRef {	18,735✔
128	&self.lhs	18,735✔
129	}	18,735✔
130
131	pub fn rhs(&self) -> &ExprRef {	18,604✔
132	&self.rhs	18,604✔
133	}	18,604✔
134
135	pub fn op(&self) -> Operator {	6,728✔
136	self.operator	6,728✔
137	}	6,728✔
138	}
139
140	impl Display for BinaryExpr {
141	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	11✔
142	write!(f, "({} {} {})", self.lhs, self.operator, self.rhs)	11✔
143	}	11✔
144	}
145
146	impl AnalysisExpr for BinaryExpr {
147	fn stat_falsification(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {	266✔
148	match self.operator {	266✔
149	Operator::Eq => {
150	let min_lhs = self.lhs.min(catalog);	86✔
151	let max_lhs = self.lhs.max(catalog);	86✔
152		86✔
153	let min_rhs = self.rhs.min(catalog);	86✔
154	let max_rhs = self.rhs.max(catalog);	86✔
155		86✔
156	let left = min_lhs.zip(max_rhs).map(\|(a, b)\| gt(a, b));	86✔
157	let right = min_rhs.zip(max_lhs).map(\|(a, b)\| gt(a, b));	86✔
158	left.into_iter().chain(right).reduce(or)	86✔
159	}
160	Operator::NotEq => {
161	let min_lhs = self.lhs.min(catalog)?;	1✔
162	let max_lhs = self.lhs.max(catalog)?;	1✔
163
164	let min_rhs = self.rhs.min(catalog)?;	1✔
165	let max_rhs = self.rhs.max(catalog)?;	1✔
166
167	Some(and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs)))	1✔
168	}
169	Operator::Gt => Some(lt_eq(self.lhs.max(catalog)?, self.rhs.min(catalog)?)),	130✔
170	Operator::Gte => Some(lt(self.lhs.max(catalog)?, self.rhs.min(catalog)?)),	22✔
171	Operator::Lt => Some(gt_eq(self.lhs.min(catalog)?, self.rhs.max(catalog)?)),	4✔
172	Operator::Lte => Some(gt(self.lhs.min(catalog)?, self.rhs.max(catalog)?)),	21✔
173	Operator::And => self	1✔
174	.lhs	1✔
175	.stat_falsification(catalog)	1✔
176	.into_iter()	1✔
177	.chain(self.rhs.stat_falsification(catalog))	1✔
178	.reduce(or),	1✔
179	Operator::Or => Some(and(
180	self.lhs.stat_falsification(catalog)?,	1✔
181	self.rhs.stat_falsification(catalog)?,	1✔
182	)),
183	Operator::Add => None,	×
184	}
185	}	266✔
186	}
187
188	/// Create a new `BinaryExpr` using the `Eq` operator.
189	///
190	/// ## Example usage
191	///
192	/// ```
193	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
194	/// use vortex_array::{Array, IntoArray, ToCanonical};
195	/// use vortex_array::validity::Validity;
196	/// use vortex_buffer::buffer;
197	/// use vortex_expr::{eq, root, lit, Scope};
198	///
199	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
200	/// let result = eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
201	///
202	/// assert_eq!(
203	/// result.to_bool().unwrap().boolean_buffer(),
204	/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
205	/// );
206	/// ```
207	pub fn eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	263✔
208	BinaryExpr::new(lhs, Operator::Eq, rhs).into_expr()	263✔
209	}	263✔
210
211	/// Create a new `BinaryExpr` using the `NotEq` operator.
212	///
213	/// ## Example usage
214	///
215	/// ```
216	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
217	/// use vortex_array::{IntoArray, ToCanonical};
218	/// use vortex_array::validity::Validity;
219	/// use vortex_buffer::buffer;
220	/// use vortex_expr::{root, lit, not_eq, Scope};
221	///
222	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
223	/// let result = not_eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
224	///
225	/// assert_eq!(
226	/// result.to_bool().unwrap().boolean_buffer(),
227	/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
228	/// );
229	/// ```
230	pub fn not_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	110✔
231	BinaryExpr::new(lhs, Operator::NotEq, rhs).into_expr()	110✔
232	}	110✔
233
234	/// Create a new `BinaryExpr` using the `Gte` operator.
235	///
236	/// ## Example usage
237	///
238	/// ```
239	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
240	/// use vortex_array::{IntoArray, ToCanonical};
241	/// use vortex_array::validity::Validity;
242	/// use vortex_buffer::buffer;
243	/// use vortex_expr::{gt_eq, root, lit, Scope};
244	///
245	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
246	/// let result = gt_eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
247	///
248	/// assert_eq!(
249	/// result.to_bool().unwrap().boolean_buffer(),
250	/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
251	/// );
252	/// ```
253	pub fn gt_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	148✔
254	BinaryExpr::new(lhs, Operator::Gte, rhs).into_expr()	148✔
255	}	148✔
256
257	/// Create a new `BinaryExpr` using the `Gt` operator.
258	///
259	/// ## Example usage
260	///
261	/// ```
262	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
263	/// use vortex_array::{IntoArray, ToCanonical};
264	/// use vortex_array::validity::Validity;
265	/// use vortex_buffer::buffer;
266	/// use vortex_expr::{gt, root, lit, Scope};
267	///
268	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
269	/// let result = gt(root(), lit(2)).evaluate(&Scope::new(xs.to_array())).unwrap();
270	///
271	/// assert_eq!(
272	/// result.to_bool().unwrap().boolean_buffer(),
273	/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
274	/// );
275	/// ```
276	pub fn gt(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	504✔
277	BinaryExpr::new(lhs, Operator::Gt, rhs).into_expr()	504✔
278	}	504✔
279
280	/// Create a new `BinaryExpr` using the `Lte` operator.
281	///
282	/// ## Example usage
283	///
284	/// ```
285	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
286	/// use vortex_array::{IntoArray, ToCanonical};
287	/// use vortex_array::validity::Validity;
288	/// use vortex_buffer::buffer;
289	/// use vortex_expr::{root, lit, lt_eq, Scope};
290	///
291	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
292	/// let result = lt_eq(root(), lit(2)).evaluate(&Scope::new(xs.to_array())).unwrap();
293	///
294	/// assert_eq!(
295	/// result.to_bool().unwrap().boolean_buffer(),
296	/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
297	/// );
298	/// ```
299	pub fn lt_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	199✔
300	BinaryExpr::new(lhs, Operator::Lte, rhs).into_expr()	199✔
301	}	199✔
302
303	/// Create a new `BinaryExpr` using the `Lt` operator.
304	///
305	/// ## Example usage
306	///
307	/// ```
308	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
309	/// use vortex_array::{IntoArray, ToCanonical};
310	/// use vortex_array::validity::Validity;
311	/// use vortex_buffer::buffer;
312	/// use vortex_expr::{root, lit, lt, Scope};
313	///
314	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
315	/// let result = lt(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
316	///
317	/// assert_eq!(
318	/// result.to_bool().unwrap().boolean_buffer(),
319	/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
320	/// );
321	/// ```
322	pub fn lt(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	187✔
323	BinaryExpr::new(lhs, Operator::Lt, rhs).into_expr()	187✔
324	}	187✔
325
326	/// Create a new `BinaryExpr` using the `Or` operator.
327	///
328	/// ## Example usage
329	///
330	/// ```
331	/// use vortex_array::arrays::BoolArray;
332	/// use vortex_array::{IntoArray, ToCanonical};
333	/// use vortex_expr::{root, lit, or, Scope};
334	///
335	/// let xs = BoolArray::from_iter(vec![true, false, true]);
336	/// let result = or(root(), lit(false)).evaluate(&Scope::new(xs.to_array())).unwrap();
337	///
338	/// assert_eq!(
339	/// result.to_bool().unwrap().boolean_buffer(),
340	/// BoolArray::from_iter(vec![true, false, true]).boolean_buffer(),
341	/// );
342	/// ```
343	pub fn or(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	1,383✔
344	BinaryExpr::new(lhs, Operator::Or, rhs).into_expr()	1,383✔
345	}	1,383✔
346
347	/// Collects a list of `or`ed values into a single vortex, expr
348	/// [x, y, z] => x or (y or z)
349	pub fn or_collect<I>(iter: I) -> Option<ExprRef>	×
350	where	×
351	I: IntoIterator<Item = ExprRef>,	×
352	I::IntoIter: DoubleEndedIterator<Item = ExprRef>,	×
353	{	×
354	let mut iter = iter.into_iter();	×
355	let first = iter.next_back()?;	×
356	Some(iter.rfold(first, \|acc, elem\| or(elem, acc)))	×
357	}	×
358
359	/// Create a new `BinaryExpr` using the `And` operator.
360	///
361	/// ## Example usage
362	///
363	/// ```
364	/// use vortex_array::arrays::BoolArray;
365	/// use vortex_array::{IntoArray, ToCanonical};
366	/// use vortex_expr::{and, root, lit, Scope};
367	///
368	/// let xs = BoolArray::from_iter(vec![true, false, true]);
369	/// let result = and(root(), lit(true)).evaluate(&Scope::new(xs.to_array())).unwrap();
370	///
371	/// assert_eq!(
372	/// result.to_bool().unwrap().boolean_buffer(),
373	/// BoolArray::from_iter(vec![true, false, true]).boolean_buffer(),
374	/// );
375	/// ```
376	pub fn and(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	320✔
377	BinaryExpr::new(lhs, Operator::And, rhs).into_expr()	320✔
378	}	320✔
379
380	/// Collects a list of `and`ed values into a single vortex, expr
381	/// [x, y, z] => x and (y and z)
382	pub fn and_collect<I>(iter: I) -> Option<ExprRef>	232✔
383	where	232✔
384	I: IntoIterator<Item = ExprRef>,	232✔
385	I::IntoIter: DoubleEndedIterator<Item = ExprRef>,	232✔
386	{	232✔
387	let mut iter = iter.into_iter();	232✔
388	let first = iter.next_back()?;	232✔
389	Some(iter.rfold(first, \|acc, elem\| and(elem, acc)))	127✔
390	}	232✔
391
392	/// Collects a list of `and`ed values into a single vortex, expr
393	/// [x, y, z] => x and (y and z)
394	pub fn and_collect_right<I>(iter: I) -> Option<ExprRef>	1✔
395	where	1✔
396	I: IntoIterator<Item = ExprRef>,	1✔
397	{	1✔
398	let iter = iter.into_iter();	1✔
399	iter.reduce(and)	1✔
400	}	1✔
401
402	/// Create a new `BinaryExpr` using the `CheckedAdd` operator.
403	///
404	/// ## Example usage
405	///
406	/// ```
407	/// use vortex_array::IntoArray;
408	/// use vortex_array::arrow::IntoArrowArray as _;
409	/// use vortex_buffer::buffer;
410	/// use vortex_expr::{Scope, checked_add, lit, root};
411	///
412	/// let xs = buffer![1, 2, 3].into_array();
413	/// let result = checked_add(root(), lit(5))
414	/// .evaluate(&Scope::new(xs.to_array()))
415	/// .unwrap();
416	///
417	/// assert_eq!(
418	/// &result.into_arrow_preferred().unwrap(),
419	/// &buffer![6, 7, 8]
420	/// .into_array()
421	/// .into_arrow_preferred()
422	/// .unwrap()
423	/// );
424	/// ```
425	pub fn checked_add(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	×
NEW 426	BinaryExpr::new(lhs, Operator::Add, rhs).into_expr()	×
427	}	×
428
429	#[cfg(test)]
430	mod tests {
431	use std::sync::Arc;
432
433	use vortex_dtype::{DType, Nullability};
434
435	use crate::{
436	ScopeDType, VortexExpr, and, and_collect, and_collect_right, col, eq, gt, gt_eq, lit, lt,
437	lt_eq, not_eq, or, test_harness,
438	};
439
440	#[test]
441	fn and_collect_left_assoc() {	1✔
442	let values = vec![lit(1), lit(2), lit(3)];	1✔
443	assert_eq!(	1✔
444	Some(and(lit(1), and(lit(2), lit(3)))),	1✔
445	and_collect(values.into_iter())	1✔
446	);	1✔
447	}	1✔
448
449	#[test]
450	fn and_collect_right_assoc() {	1✔
451	let values = vec![lit(1), lit(2), lit(3)];	1✔
452	assert_eq!(	1✔
453	Some(and(and(lit(1), lit(2)), lit(3))),	1✔
454	and_collect_right(values.into_iter())	1✔
455	);	1✔
456	}	1✔
457
458	#[test]
459	fn dtype() {	1✔
460	let dtype = test_harness::struct_dtype();	1✔
461	let bool1: Arc<dyn VortexExpr> = col("bool1");	1✔
462	let bool2: Arc<dyn VortexExpr> = col("bool2");	1✔
463	assert_eq!(	1✔
464	and(bool1.clone(), bool2.clone())	1✔
465	.return_dtype(&ScopeDType::new(dtype.clone()))	1✔
466	.unwrap(),	1✔
467	DType::Bool(Nullability::NonNullable)	1✔
468	);	1✔
469	assert_eq!(	1✔
470	or(bool1.clone(), bool2.clone())	1✔
471	.return_dtype(&ScopeDType::new(dtype.clone()))	1✔
472	.unwrap(),	1✔
473	DType::Bool(Nullability::NonNullable)	1✔
474	);	1✔
475
476	let col1: Arc<dyn VortexExpr> = col("col1");	1✔
477	let col2: Arc<dyn VortexExpr> = col("col2");	1✔
478		1✔
479	assert_eq!(	1✔
480	eq(col1.clone(), col2.clone())	1✔
481	.return_dtype(&ScopeDType::new(dtype.clone()))	1✔
482	.unwrap(),	1✔
483	DType::Bool(Nullability::Nullable)	1✔
484	);	1✔
485	assert_eq!(	1✔
486	not_eq(col1.clone(), col2.clone())	1✔
487	.return_dtype(&ScopeDType::new(dtype.clone()))	1✔
488	.unwrap(),	1✔
489	DType::Bool(Nullability::Nullable)	1✔
490	);	1✔
491	assert_eq!(	1✔
492	gt(col1.clone(), col2.clone())	1✔
493	.return_dtype(&ScopeDType::new(dtype.clone()))	1✔
494	.unwrap(),	1✔
495	DType::Bool(Nullability::Nullable)	1✔
496	);	1✔
497	assert_eq!(	1✔
498	gt_eq(col1.clone(), col2.clone())	1✔
499	.return_dtype(&ScopeDType::new(dtype.clone()))	1✔
500	.unwrap(),	1✔
501	DType::Bool(Nullability::Nullable)	1✔
502	);	1✔
503	assert_eq!(	1✔
504	lt(col1.clone(), col2.clone())	1✔
505	.return_dtype(&ScopeDType::new(dtype.clone()))	1✔
506	.unwrap(),	1✔
507	DType::Bool(Nullability::Nullable)	1✔
508	);	1✔
509	assert_eq!(	1✔
510	lt_eq(col1.clone(), col2.clone())	1✔
511	.return_dtype(&ScopeDType::new(dtype.clone()))	1✔
512	.unwrap(),	1✔
513	DType::Bool(Nullability::Nullable)	1✔
514	);	1✔
515
516	assert_eq!(	1✔
517	or(	1✔
518	lt(col1.clone(), col2.clone()),	1✔
519	not_eq(col1.clone(), col2.clone())	1✔
520	)	1✔
521	.return_dtype(&ScopeDType::new(dtype))	1✔
522	.unwrap(),	1✔
523	DType::Bool(Nullability::Nullable)	1✔
524	);	1✔
525	}	1✔
526	}

vortex-data / vortex / 16139349253

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous