vortex_array/compute/
boolean.rs1use std::any::Any;
5
6use arrow_array::cast::AsArray;
7use arrow_schema::DataType;
8use vortex_error::VortexResult;
9
10use crate::Array;
11use crate::ArrayRef;
12use crate::IntoArray;
13use crate::arrays::ScalarFnArray;
14use crate::arrow::FromArrowArray;
15use crate::arrow::IntoArrowArray;
16use crate::compute::Options;
17use crate::expr::Binary;
18use crate::expr::ScalarFn;
19use crate::expr::operators::Operator;
20
21#[deprecated(note = "Use and_kleene instead. Non-Kleene boolean ops cannot be lazily evaluated.")]
28pub fn and(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
29 boolean(lhs, rhs, BooleanOperator::And)
30}
31
32pub fn and_kleene(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
36 boolean(lhs, rhs, BooleanOperator::AndKleene)
37}
38
39#[deprecated(note = "Use or_kleene instead. Non-Kleene boolean ops cannot be lazily evaluated.")]
46pub fn or(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
47 boolean(lhs, rhs, BooleanOperator::Or)
48}
49
50pub fn or_kleene(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
54 boolean(lhs, rhs, BooleanOperator::OrKleene)
55}
56
57pub fn boolean(lhs: &dyn Array, rhs: &dyn Array, op: BooleanOperator) -> VortexResult<ArrayRef> {
59 match Operator::try_from(op) {
60 Ok(expr_op) => Ok(ScalarFnArray::try_new(
61 ScalarFn::new(Binary, expr_op),
62 vec![lhs.to_array(), rhs.to_array()],
63 lhs.len(),
64 )?
65 .into_array()),
66 Err(_) => {
67 tracing::trace!(
68 "non-Kleene boolean op {op:?} cannot be lazily evaluated, falling back to eager Arrow evaluation"
69 );
70 arrow_boolean(lhs.to_array(), rhs.to_array(), op)
71 }
72 }
73}
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum BooleanOperator {
80 And,
89 AndKleene,
98 Or,
107 OrKleene,
116 }
120
121impl Options for BooleanOperator {
122 fn as_any(&self) -> &dyn Any {
123 self
124 }
125}
126
127pub(crate) fn arrow_boolean(
132 lhs: ArrayRef,
133 rhs: ArrayRef,
134 operator: BooleanOperator,
135) -> VortexResult<ArrayRef> {
136 let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
137
138 let lhs = lhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
139 let rhs = rhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
140
141 let array = match operator {
142 BooleanOperator::And => arrow_arith::boolean::and(&lhs, &rhs)?,
143 BooleanOperator::AndKleene => arrow_arith::boolean::and_kleene(&lhs, &rhs)?,
144 BooleanOperator::Or => arrow_arith::boolean::or(&lhs, &rhs)?,
145 BooleanOperator::OrKleene => arrow_arith::boolean::or_kleene(&lhs, &rhs)?,
146 };
147
148 ArrayRef::from_arrow(&array, nullable)
149}
150
151#[cfg(test)]
152mod tests {
153 use rstest::rstest;
154
155 use super::*;
156 use crate::IntoArray;
157 use crate::arrays::BoolArray;
158 use crate::canonical::ToCanonical;
159 #[rstest]
160 #[case(BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter())
161 .into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter())
162 .into_array())]
163 #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(),
164 BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter()).into_array())]
165 fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
166 let r = or_kleene(&lhs, &rhs).unwrap();
167
168 let r = r.to_bool().into_array();
169
170 let v0 = r.scalar_at(0).unwrap().as_bool().value();
171 let v1 = r.scalar_at(1).unwrap().as_bool().value();
172 let v2 = r.scalar_at(2).unwrap().as_bool().value();
173 let v3 = r.scalar_at(3).unwrap().as_bool().value();
174
175 assert!(v0.unwrap());
176 assert!(v1.unwrap());
177 assert!(v2.unwrap());
178 assert!(!v3.unwrap());
179 }
180
181 #[rstest]
182 #[case(BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter())
183 .into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter())
184 .into_array())]
185 #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(),
186 BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter()).into_array())]
187 fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
188 let r = and_kleene(&lhs, &rhs).unwrap().to_bool().into_array();
189
190 let v0 = r.scalar_at(0).unwrap().as_bool().value();
191 let v1 = r.scalar_at(1).unwrap().as_bool().value();
192 let v2 = r.scalar_at(2).unwrap().as_bool().value();
193 let v3 = r.scalar_at(3).unwrap().as_bool().value();
194
195 assert!(v0.unwrap());
196 assert!(!v1.unwrap());
197 assert!(!v2.unwrap());
198 assert!(!v3.unwrap());
199 }
200}