16992684502

Committed 15 Aug 2025 02:56PM UTC coverage: 87.875% (+0.2%) from 87.72%

Build # 16992684502

Build Type

Pull #2456

github

Committed by

web-flow

Commit Message

Merge 2d540e578 into 4a23f65b3

Pull Request Pull Request #2456: feat: basic BoolBuffer / BoolBufferMut

Run Details

1275 of 1428 new or added lines in 110 files covered. (89.29%)

334 existing lines in 31 files now uncovered.

57169 of 65057 relevant lines covered (87.88%)

658056.52 hits per line

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

97.32

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

vortex-data / vortex / 16992684502

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