16991469020

Committed 15 Aug 2025 01:54PM UTC coverage: 86.018% (-1.8%) from 87.855%

Build # 16991469020

Build Type

Pull #4215

github

Committed by

web-flow

Commit Message

Merge 58201b202 into 62e231a41

Pull Request Pull Request #4215: Ji/vectors

Run Details

90 of 1746 new or added lines in 40 files covered. (5.15%)

117 existing lines in 25 files now uncovered.

56661 of 65871 relevant lines covered (86.02%)

614266.74 hits per line

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

98.07

/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, StatsCatalog,
    VTable, lit, vtable,
};

vtable!(Binary);

#[allow(clippy::derived_hash_with_manual_eq)]
#[derive(Debug, Clone, Hash, Eq)]
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> {
        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: &DType) -> 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> {
        // Wrap another predicate with an optional NaNCount check, if the stat is available.
        //
        // For example, regular pruning conversion for `A >= B` would be
        //
        //      A.max < B.min
        //
        // With NaN predicate introduction, we'd conjunct it with a check for NaNCount, resulting
        // in:
        //
        //      (A.nan_count = 0) AND (B.nan_count = 0) AND A.max < B.min
        //
        // Non-floating point column and literal expressions should be unaffected as they do not
        // have a nan_count statistic defined.
        #[inline]
        fn with_nan_predicate(
            lhs: &ExprRef,
            rhs: &ExprRef,
            value_predicate: ExprRef,
            catalog: &mut dyn StatsCatalog,
        ) -> ExprRef {
            let nan_predicate = lhs
                .nan_count(catalog)
                .into_iter()
                .chain(rhs.nan_count(catalog))
                .map(|nans| eq(nans, lit(0u64)))
                .reduce(and);

            if let Some(nan_check) = nan_predicate {
                and(nan_check, value_predicate)
            } else {
                value_predicate
            }
        }

        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));

                let min_max_check = left.into_iter().chain(right).reduce(or)?;

                // NaN is not captured by the min/max stat, so we must check NaNCount before pruning
                Some(with_nan_predicate(
                    self.lhs(),
                    self.rhs(),
                    min_max_check,
                    catalog,
                ))
            }
            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)?;

                let min_max_check = and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs));

                Some(with_nan_predicate(
                    self.lhs(),
                    self.rhs(),
                    min_max_check,
                    catalog,
                ))
            }
            Operator::Gt => {
                let min_max_check = lt_eq(self.lhs.max(catalog)?, self.rhs.min(catalog)?);

                Some(with_nan_predicate(
                    self.lhs(),
                    self.rhs(),
                    min_max_check,
                    catalog,
                ))
            }
            Operator::Gte => {
                // NaN is not captured by the min/max stat, so we must check NaNCount before pruning
                let min_max_check = lt(self.lhs.max(catalog)?, self.rhs.min(catalog)?);

                Some(with_nan_predicate(
                    self.lhs(),
                    self.rhs(),
                    min_max_check,
                    catalog,
                ))
            }
            Operator::Lt => {
                // NaN is not captured by the min/max stat, so we must check NaNCount before pruning
                let min_max_check = gt_eq(self.lhs.min(catalog)?, self.rhs.max(catalog)?);

                Some(with_nan_predicate(
                    self.lhs(),
                    self.rhs(),
                    min_max_check,
                    catalog,
                ))
            }
            Operator::Lte => {
                // NaN is not captured by the min/max stat, so we must check NaNCount before pruning
                let min_max_check = gt(self.lhs.min(catalog)?, self.rhs.max(catalog)?);

                Some(with_nan_predicate(
                    self.lhs(),
                    self.rhs(),
                    min_max_check,
                    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::{
        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(&dtype)
                .unwrap(),
            DType::Bool(Nullability::NonNullable)
        );
        assert_eq!(
            or(bool1.clone(), bool2.clone())
                .return_dtype(&dtype)
                .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(&dtype).unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            not_eq(col1.clone(), col2.clone())
                .return_dtype(&dtype)
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            gt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            gt_eq(col1.clone(), col2.clone())
                .return_dtype(&dtype)
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            lt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
            DType::Bool(Nullability::Nullable)
        );
        assert_eq!(
            lt_eq(col1.clone(), col2.clone())
                .return_dtype(&dtype)
                .unwrap(),
            DType::Bool(Nullability::Nullable)
        );

        assert_eq!(
            or(
                lt(col1.clone(), col2.clone()),
                not_eq(col1.clone(), col2.clone())
            )
            .return_dtype(&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, StatsCatalog,
15	VTable, lit, vtable,
16	};
17
18	vtable!(Binary);
19
20	#[allow(clippy::derived_hash_with_manual_eq)]
21	#[derive(Debug, Clone, Hash, Eq)]
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 {	6,359✔
30	self.lhs.eq(&other.lhs) && self.operator == other.operator && self.rhs.eq(&other.rhs)	6,359✔
31	}	6,359✔
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 {	187✔
42	ExprId::new_ref("binary")	187✔
43	}	187✔
44
45	fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {	20✔
46	ExprEncodingRef::new_ref(BinaryExprEncoding.as_ref())	20✔
47	}	20✔
48
49	fn metadata(expr: &Self::Expr) -> Option<Self::Metadata> {	20✔
50	Some(ProstMetadata(pb::BinaryOpts {	20✔
51	op: expr.operator.into(),	20✔
52	}))	20✔
53	}	20✔
54
55	fn children(expr: &Self::Expr) -> Vec<&ExprRef> {	53,722✔
56	vec![expr.lhs(), expr.rhs()]	53,722✔
57	}	53,722✔
58
59	fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {	2,303✔
60	Ok(BinaryExpr::new(	2,303✔
61	children[0].clone(),	2,303✔
62	expr.op(),	2,303✔
63	children[1].clone(),	2,303✔
64	))	2,303✔
65	}	2,303✔
66
67	fn build(	20✔
68	_encoding: &Self::Encoding,	20✔
69	metadata: &<Self::Metadata as DeserializeMetadata>::Output,	20✔
70	children: Vec<ExprRef>,	20✔
71	) -> VortexResult<Self::Expr> {	20✔
72	Ok(BinaryExpr::new(	20✔
73	children[0].clone(),	20✔
74	metadata.op().into(),	20✔
75	children[1].clone(),	20✔
76	))	20✔
77	}	20✔
78
79	fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {	8,877✔
80	let lhs = expr.lhs.unchecked_evaluate(scope)?;	8,877✔
81	let rhs = expr.rhs.unchecked_evaluate(scope)?;	8,877✔
82
83	match expr.operator {	8,877✔
84	Operator::Eq => compare(&lhs, &rhs, ArrayOperator::Eq),	534✔
85	Operator::NotEq => compare(&lhs, &rhs, ArrayOperator::NotEq),	6✔
86	Operator::Lt => compare(&lhs, &rhs, ArrayOperator::Lt),	1,740✔
87	Operator::Lte => compare(&lhs, &rhs, ArrayOperator::Lte),	294✔
88	Operator::Gt => compare(&lhs, &rhs, ArrayOperator::Gt),	2,691✔
89	Operator::Gte => compare(&lhs, &rhs, ArrayOperator::Gte),	105✔
90	Operator::And => and_kleene(&lhs, &rhs),	1,368✔
91	Operator::Or => or_kleene(&lhs, &rhs),	2,139✔
UNCOV 92	Operator::Add => add(&lhs, &rhs),	×
93	}
94	}	8,877✔
95
96	fn return_dtype(expr: &Self::Expr, scope: &DType) -> VortexResult<DType> {	10,485✔
97	let lhs = expr.lhs.return_dtype(scope)?;	10,485✔
98	let rhs = expr.rhs.return_dtype(scope)?;	10,485✔
99
100	if expr.operator == Operator::Add {	10,485✔
UNCOV 101	if lhs.is_primitive() && lhs.eq_ignore_nullability(&rhs) {	×
102	return Ok(lhs.with_nullability(lhs.nullability() \| rhs.nullability()));
UNCOV 103	}	×
104	vortex_bail!("incompatible types for checked add: {} {}", lhs, rhs);	×
105	}	10,485✔
106
107	Ok(DType::Bool((lhs.is_nullable() \|\| rhs.is_nullable()).into()))	10,485✔
108	}	10,485✔
109	}
110
111	impl BinaryExpr {
112	pub fn new(lhs: ExprRef, operator: Operator, rhs: ExprRef) -> Self {	56,420✔
113	Self { lhs, operator, rhs }	56,420✔
114	}	56,420✔
115
116	pub fn new_expr(lhs: ExprRef, operator: Operator, rhs: ExprRef) -> ExprRef {	695✔
117	Self::new(lhs, operator, rhs).into_expr()	695✔
118	}	695✔
119
120	pub fn lhs(&self) -> &ExprRef {	78,656✔
121	&self.lhs	78,656✔
122	}	78,656✔
123
124	pub fn rhs(&self) -> &ExprRef {	77,535✔
125	&self.rhs	77,535✔
126	}	77,535✔
127
128	pub fn op(&self) -> Operator {	52,546✔
129	self.operator	52,546✔
130	}	52,546✔
131	}
132
133	impl Display for BinaryExpr {
134	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	17✔
135	write!(f, "({} {} {})", self.lhs, self.operator, self.rhs)	17✔
136	}	17✔
137	}
138
139	impl AnalysisExpr for BinaryExpr {
140	fn stat_falsification(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {	1,326✔
141	// Wrap another predicate with an optional NaNCount check, if the stat is available.
142	//
143	// For example, regular pruning conversion for `A >= B` would be
144	//
145	// A.max < B.min
146	//
147	// With NaN predicate introduction, we'd conjunct it with a check for NaNCount, resulting
148	// in:
149	//
150	// (A.nan_count = 0) AND (B.nan_count = 0) AND A.max < B.min
151	//
152	// Non-floating point column and literal expressions should be unaffected as they do not
153	// have a nan_count statistic defined.
154	#[inline]
155	fn with_nan_predicate(	995✔
156	lhs: &ExprRef,	995✔
157	rhs: &ExprRef,	995✔
158	value_predicate: ExprRef,	995✔
159	catalog: &mut dyn StatsCatalog,	995✔
160	) -> ExprRef {	995✔
161	let nan_predicate = lhs	995✔
162	.nan_count(catalog)	995✔
163	.into_iter()	995✔
164	.chain(rhs.nan_count(catalog))	995✔
165	.map(\|nans\| eq(nans, lit(0u64)))	995✔
166	.reduce(and);	995✔
167
168	if let Some(nan_check) = nan_predicate {	995✔
169	and(nan_check, value_predicate)	44✔
170	} else {
171	value_predicate	951✔
172	}
173	}	995✔
174
175	match self.operator {	1,326✔
176	Operator::Eq => {
177	let min_lhs = self.lhs.min(catalog);	371✔
178	let max_lhs = self.lhs.max(catalog);	371✔
179
180	let min_rhs = self.rhs.min(catalog);	371✔
181	let max_rhs = self.rhs.max(catalog);	371✔
182
183	let left = min_lhs.zip(max_rhs).map(\|(a, b)\| gt(a, b));	371✔
184	let right = min_rhs.zip(max_lhs).map(\|(a, b)\| gt(a, b));	371✔
185
186	let min_max_check = left.into_iter().chain(right).reduce(or)?;	371✔
187
188	// NaN is not captured by the min/max stat, so we must check NaNCount before pruning
189	Some(with_nan_predicate(	371✔
190	self.lhs(),	371✔
191	self.rhs(),	371✔
192	min_max_check,	371✔
193	catalog,	371✔
194	))	371✔
195	}
196	Operator::NotEq => {
197	let min_lhs = self.lhs.min(catalog)?;	10✔
198	let max_lhs = self.lhs.max(catalog)?;	10✔
199
200	let min_rhs = self.rhs.min(catalog)?;	10✔
201	let max_rhs = self.rhs.max(catalog)?;	10✔
202
203	let min_max_check = and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs));	10✔
204
205	Some(with_nan_predicate(	10✔
206	self.lhs(),	10✔
207	self.rhs(),	10✔
208	min_max_check,	10✔
209	catalog,	10✔
210	))	10✔
211	}
212	Operator::Gt => {
213	let min_max_check = lt_eq(self.lhs.max(catalog)?, self.rhs.min(catalog)?);	176✔
214
215	Some(with_nan_predicate(	176✔
216	self.lhs(),	176✔
217	self.rhs(),	176✔
218	min_max_check,	176✔
219	catalog,	176✔
220	))	176✔
221	}
222	Operator::Gte => {
223	// NaN is not captured by the min/max stat, so we must check NaNCount before pruning
224	let min_max_check = lt(self.lhs.max(catalog)?, self.rhs.min(catalog)?);	218✔
225
226	Some(with_nan_predicate(	217✔
227	self.lhs(),	217✔
228	self.rhs(),	217✔
229	min_max_check,	217✔
230	catalog,	217✔
231	))	217✔
232	}
233	Operator::Lt => {
234	// NaN is not captured by the min/max stat, so we must check NaNCount before pruning
235	let min_max_check = gt_eq(self.lhs.min(catalog)?, self.rhs.max(catalog)?);	65✔
236
237	Some(with_nan_predicate(	65✔
238	self.lhs(),	65✔
239	self.rhs(),	65✔
240	min_max_check,	65✔
241	catalog,	65✔
242	))	65✔
243	}
244	Operator::Lte => {
245	// NaN is not captured by the min/max stat, so we must check NaNCount before pruning
246	let min_max_check = gt(self.lhs.min(catalog)?, self.rhs.max(catalog)?);	156✔
247
248	Some(with_nan_predicate(	156✔
249	self.lhs(),	156✔
250	self.rhs(),	156✔
251	min_max_check,	156✔
252	catalog,	156✔
253	))	156✔
254	}
255	Operator::And => self	305✔
256	.lhs	305✔
257	.stat_falsification(catalog)	305✔
258	.into_iter()	305✔
259	.chain(self.rhs.stat_falsification(catalog))	305✔
260	.reduce(or),	305✔
261	Operator::Or => Some(and(	25✔
262	self.lhs.stat_falsification(catalog)?,	25✔
263	self.rhs.stat_falsification(catalog)?,	25✔
264	)),
UNCOV 265	Operator::Add => None,	×
266	}
267	}	1,326✔
268	}
269
270	/// Create a new `BinaryExpr` using the `Eq` operator.
271	///
272	/// ## Example usage
273	///
274	/// ```
275	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
276	/// use vortex_array::{Array, IntoArray, ToCanonical};
277	/// use vortex_array::validity::Validity;
278	/// use vortex_buffer::buffer;
279	/// use vortex_expr::{eq, root, lit, Scope};
280	///
281	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
282	/// let result = eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
283	///
284	/// assert_eq!(
285	/// result.to_bool().unwrap().boolean_buffer(),
286	/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
287	/// );
288	/// ```
289	pub fn eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	354✔
290	BinaryExpr::new(lhs, Operator::Eq, rhs).into_expr()	354✔
291	}	354✔
292
293	/// Create a new `BinaryExpr` using the `NotEq` operator.
294	///
295	/// ## Example usage
296	///
297	/// ```
298	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
299	/// use vortex_array::{IntoArray, ToCanonical};
300	/// use vortex_array::validity::Validity;
301	/// use vortex_buffer::buffer;
302	/// use vortex_expr::{root, lit, not_eq, Scope};
303	///
304	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
305	/// let result = not_eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
306	///
307	/// assert_eq!(
308	/// result.to_bool().unwrap().boolean_buffer(),
309	/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
310	/// );
311	/// ```
312	pub fn not_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	6✔
313	BinaryExpr::new(lhs, Operator::NotEq, rhs).into_expr()	6✔
314	}	6✔
315
316	/// Create a new `BinaryExpr` using the `Gte` operator.
317	///
318	/// ## Example usage
319	///
320	/// ```
321	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
322	/// use vortex_array::{IntoArray, ToCanonical};
323	/// use vortex_array::validity::Validity;
324	/// use vortex_buffer::buffer;
325	/// use vortex_expr::{gt_eq, root, lit, Scope};
326	///
327	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
328	/// let result = gt_eq(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
329	///
330	/// assert_eq!(
331	/// result.to_bool().unwrap().boolean_buffer(),
332	/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
333	/// );
334	/// ```
335	pub fn gt_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	128✔
336	BinaryExpr::new(lhs, Operator::Gte, rhs).into_expr()	128✔
337	}	128✔
338
339	/// Create a new `BinaryExpr` using the `Gt` operator.
340	///
341	/// ## Example usage
342	///
343	/// ```
344	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
345	/// use vortex_array::{IntoArray, ToCanonical};
346	/// use vortex_array::validity::Validity;
347	/// use vortex_buffer::buffer;
348	/// use vortex_expr::{gt, root, lit, Scope};
349	///
350	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
351	/// let result = gt(root(), lit(2)).evaluate(&Scope::new(xs.to_array())).unwrap();
352	///
353	/// assert_eq!(
354	/// result.to_bool().unwrap().boolean_buffer(),
355	/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
356	/// );
357	/// ```
358	pub fn gt(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	20,100✔
359	BinaryExpr::new(lhs, Operator::Gt, rhs).into_expr()	20,100✔
360	}	20,100✔
361
362	/// Create a new `BinaryExpr` using the `Lte` operator.
363	///
364	/// ## Example usage
365	///
366	/// ```
367	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
368	/// use vortex_array::{IntoArray, ToCanonical};
369	/// use vortex_array::validity::Validity;
370	/// use vortex_buffer::buffer;
371	/// use vortex_expr::{root, lit, lt_eq, Scope};
372	///
373	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
374	/// let result = lt_eq(root(), lit(2)).evaluate(&Scope::new(xs.to_array())).unwrap();
375	///
376	/// assert_eq!(
377	/// result.to_bool().unwrap().boolean_buffer(),
378	/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
379	/// );
380	/// ```
381	pub fn lt_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	247✔
382	BinaryExpr::new(lhs, Operator::Lte, rhs).into_expr()	247✔
383	}	247✔
384
385	/// Create a new `BinaryExpr` using the `Lt` operator.
386	///
387	/// ## Example usage
388	///
389	/// ```
390	/// use vortex_array::arrays::{BoolArray, PrimitiveArray };
391	/// use vortex_array::{IntoArray, ToCanonical};
392	/// use vortex_array::validity::Validity;
393	/// use vortex_buffer::buffer;
394	/// use vortex_expr::{root, lit, lt, Scope};
395	///
396	/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
397	/// let result = lt(root(), lit(3)).evaluate(&Scope::new(xs.to_array())).unwrap();
398	///
399	/// assert_eq!(
400	/// result.to_bool().unwrap().boolean_buffer(),
401	/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
402	/// );
403	/// ```
404	pub fn lt(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	10,515✔
405	BinaryExpr::new(lhs, Operator::Lt, rhs).into_expr()	10,515✔
406	}	10,515✔
407
408	/// Create a new `BinaryExpr` using the `Or` operator.
409	///
410	/// ## Example usage
411	///
412	/// ```
413	/// use vortex_array::arrays::BoolArray;
414	/// use vortex_array::{IntoArray, ToCanonical};
415	/// use vortex_expr::{root, lit, or, Scope};
416	///
417	/// let xs = BoolArray::from_iter(vec![true, false, true]);
418	/// let result = or(root(), lit(false)).evaluate(&Scope::new(xs.to_array())).unwrap();
419	///
420	/// assert_eq!(
421	/// result.to_bool().unwrap().boolean_buffer(),
422	/// BoolArray::from_iter(vec![true, false, true]).boolean_buffer(),
423	/// );
424	/// ```
425	pub fn or(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	2,149✔
426	BinaryExpr::new(lhs, Operator::Or, rhs).into_expr()	2,149✔
427	}	2,149✔
428
429	/// Collects a list of `or`ed values into a single vortex, expr
430	/// [x, y, z] => x or (y or z)
431	pub fn or_collect<I>(iter: I) -> Option<ExprRef>	6✔
432	where	6✔
433	I: IntoIterator<Item = ExprRef>,	6✔
434	I::IntoIter: DoubleEndedIterator<Item = ExprRef>,	6✔
435	{
436	let mut iter = iter.into_iter();	6✔
437	let first = iter.next_back()?;	6✔
438	Some(iter.rfold(first, \|acc, elem\| or(elem, acc)))	6✔
439	}	6✔
440
441	/// Create a new `BinaryExpr` using the `And` operator.
442	///
443	/// ## Example usage
444	///
445	/// ```
446	/// use vortex_array::arrays::BoolArray;
447	/// use vortex_array::{IntoArray, ToCanonical};
448	/// use vortex_expr::{and, root, lit, Scope};
449	///
450	/// let xs = BoolArray::from_iter(vec![true, false, true]);
451	/// let result = and(root(), lit(true)).evaluate(&Scope::new(xs.to_array())).unwrap();
452	///
453	/// assert_eq!(
454	/// result.to_bool().unwrap().boolean_buffer(),
455	/// BoolArray::from_iter(vec![true, false, true]).boolean_buffer(),
456	/// );
457	/// ```
458	pub fn and(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	19,894✔
459	BinaryExpr::new(lhs, Operator::And, rhs).into_expr()	19,894✔
460	}	19,894✔
461
462	/// Collects a list of `and`ed values into a single vortex, expr
463	/// [x, y, z] => x and (y and z)
464	pub fn and_collect<I>(iter: I) -> Option<ExprRef>	661✔
465	where	661✔
466	I: IntoIterator<Item = ExprRef>,	661✔
467	I::IntoIter: DoubleEndedIterator<Item = ExprRef>,	661✔
468	{
469	let mut iter = iter.into_iter();	661✔
470	let first = iter.next_back()?;	661✔
471	Some(iter.rfold(first, \|acc, elem\| and(elem, acc)))	386✔
472	}	661✔
473
474	/// Collects a list of `and`ed values into a single vortex, expr
475	/// [x, y, z] => x and (y and z)
476	pub fn and_collect_right<I>(iter: I) -> Option<ExprRef>	1✔
477	where	1✔
478	I: IntoIterator<Item = ExprRef>,	1✔
479	{
480	let iter = iter.into_iter();	1✔
481	iter.reduce(and)	1✔
482	}	1✔
483
484	/// Create a new `BinaryExpr` using the `CheckedAdd` operator.
485	///
486	/// ## Example usage
487	///
488	/// ```
489	/// use vortex_array::IntoArray;
490	/// use vortex_array::arrow::IntoArrowArray as _;
491	/// use vortex_buffer::buffer;
492	/// use vortex_expr::{Scope, checked_add, lit, root};
493	///
494	/// let xs = buffer![1, 2, 3].into_array();
495	/// let result = checked_add(root(), lit(5))
496	/// .evaluate(&Scope::new(xs.to_array()))
497	/// .unwrap();
498	///
499	/// assert_eq!(
500	/// &result.into_arrow_preferred().unwrap(),
501	/// &buffer![6, 7, 8]
502	/// .into_array()
503	/// .into_arrow_preferred()
504	/// .unwrap()
505	/// );
506	/// ```
507	pub fn checked_add(lhs: ExprRef, rhs: ExprRef) -> ExprRef {	7✔
508	BinaryExpr::new(lhs, Operator::Add, rhs).into_expr()	7✔
509	}	7✔
510
511	#[cfg(test)]
512	mod tests {
513	use std::sync::Arc;
514
515	use vortex_dtype::{DType, Nullability};
516
517	use crate::{
518	VortexExpr, and, and_collect, and_collect_right, col, eq, gt, gt_eq, lit, lt, lt_eq,
519	not_eq, or, test_harness,
520	};
521
522	#[test]
523	fn and_collect_left_assoc() {	1✔
524	let values = vec![lit(1), lit(2), lit(3)];	1✔
525	assert_eq!(	1✔
526	Some(and(lit(1), and(lit(2), lit(3)))),	1✔
527	and_collect(values.into_iter())	1✔
528	);
529	}	1✔
530
531	#[test]
532	fn and_collect_right_assoc() {	1✔
533	let values = vec![lit(1), lit(2), lit(3)];	1✔
534	assert_eq!(	1✔
535	Some(and(and(lit(1), lit(2)), lit(3))),	1✔
536	and_collect_right(values.into_iter())	1✔
537	);
538	}	1✔
539
540	#[test]
541	fn dtype() {	1✔
542	let dtype = test_harness::struct_dtype();	1✔
543	let bool1: Arc<dyn VortexExpr> = col("bool1");	1✔
544	let bool2: Arc<dyn VortexExpr> = col("bool2");	1✔
545	assert_eq!(	1✔
546	and(bool1.clone(), bool2.clone())	1✔
547	.return_dtype(&dtype)	1✔
548	.unwrap(),	1✔
549	DType::Bool(Nullability::NonNullable)
550	);
551	assert_eq!(	1✔
552	or(bool1.clone(), bool2.clone())	1✔
553	.return_dtype(&dtype)	1✔
554	.unwrap(),	1✔
555	DType::Bool(Nullability::NonNullable)
556	);
557
558	let col1: Arc<dyn VortexExpr> = col("col1");	1✔
559	let col2: Arc<dyn VortexExpr> = col("col2");	1✔
560
561	assert_eq!(	1✔
562	eq(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),	1✔
563	DType::Bool(Nullability::Nullable)
564	);
565	assert_eq!(	1✔
566	not_eq(col1.clone(), col2.clone())	1✔
567	.return_dtype(&dtype)	1✔
568	.unwrap(),	1✔
569	DType::Bool(Nullability::Nullable)
570	);
571	assert_eq!(	1✔
572	gt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),	1✔
573	DType::Bool(Nullability::Nullable)
574	);
575	assert_eq!(	1✔
576	gt_eq(col1.clone(), col2.clone())	1✔
577	.return_dtype(&dtype)	1✔
578	.unwrap(),	1✔
579	DType::Bool(Nullability::Nullable)
580	);
581	assert_eq!(	1✔
582	lt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),	1✔
583	DType::Bool(Nullability::Nullable)
584	);
585	assert_eq!(	1✔
586	lt_eq(col1.clone(), col2.clone())	1✔
587	.return_dtype(&dtype)	1✔
588	.unwrap(),	1✔
589	DType::Bool(Nullability::Nullable)
590	);
591
592	assert_eq!(	1✔
593	or(	1✔
594	lt(col1.clone(), col2.clone()),	1✔
595	not_eq(col1.clone(), col2.clone())	1✔
596	)	1✔
597	.return_dtype(&dtype)	1✔
598	.unwrap(),	1✔
599	DType::Bool(Nullability::Nullable)
600	);
601	}	1✔
602	}

vortex-data / vortex / 16991469020

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