vortex_array/scalar_fn/fns/binary/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7#[expect(deprecated)]
8pub use boolean::and_kleene;
9#[expect(deprecated)]
10pub use boolean::or_kleene;
11use prost::Message;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_session::VortexSession;
15use vortex_session::registry::CachedId;
16
17use crate::ArrayRef;
18use crate::ExecutionCtx;
19use crate::arrays::ScalarFnArray;
20use crate::dtype::DType;
21use crate::dtype::Nullability;
22use crate::expr::BoundExpression;
23use crate::expr::bound;
24use crate::expr::display::ExprDisplay;
25use crate::proto::expr as pb;
26use crate::scalar_fn::Arity;
27use crate::scalar_fn::ChildName;
28use crate::scalar_fn::ExecutionArgs;
29use crate::scalar_fn::ReduceNode;
30use crate::scalar_fn::ScalarFnId;
31use crate::scalar_fn::ScalarFnVTable;
32use crate::scalar_fn::ScalarFnVTableExt;
33use crate::scalar_fn::fns::literal::Literal;
34use crate::scalar_fn::fns::operators::CompareOperator;
35use crate::scalar_fn::fns::operators::Operator;
36
37pub mod boolean;
38pub use boolean::BooleanExecuteAdaptor;
39pub use boolean::BooleanKernel;
40pub(crate) use boolean::execute_boolean;
41pub use boolean::kleene_boolean_buffer_scalar;
42pub use boolean::kleene_boolean_buffers;
43mod compare;
44pub use compare::*;
45mod numeric;
46pub(crate) use numeric::*;
47mod primitive_operand;
48
49use crate::scalar::NumericOperator;
50use crate::scalar::Scalar;
51
52#[derive(Clone)]
53pub struct Binary;
54
55impl Binary {
56    /// Creates a lazy binary operation over `lhs` and `rhs`.
57    ///
58    /// # Errors
59    ///
60    /// Returns an error if the children have different lengths or incompatible dtypes.
61    pub fn try_new(
62        lhs: ArrayRef,
63        rhs: ArrayRef,
64        operator: Operator,
65    ) -> VortexResult<ScalarFnArray> {
66        ScalarFnArray::try_new(Binary.bind(operator), vec![lhs, rhs])
67    }
68}
69
70/// Kleene and/or lookup table where both arguments are constant and non-NULL.
71/// is_and, left, right
72const KLEENE_LUT: [[[bool; 2]; 2]; 2] = [
73    [[false, true], [true, true]],   // or
74    [[false, false], [false, true]], // and
75];
76
77/// Kleene and/or reduction where one argument is constant non-NULL and other
78/// is non-constant.
79fn kleene_one_const<T: ReduceNode>(node: T, constant: bool, is_and: bool) -> T {
80    match (is_and, constant) {
81        (true, true) | (false, false) => node,
82        (is_and, constant) => node.new_constant((constant && !is_and).into()),
83    }
84}
85
86impl ScalarFnVTable for Binary {
87    type Options = Operator;
88
89    fn id(&self) -> ScalarFnId {
90        static ID: CachedId = CachedId::new("vortex.binary");
91        *ID
92    }
93
94    fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
95        Ok(Some(
96            pb::BinaryOpts {
97                op: (*instance).into(),
98            }
99            .encode_to_vec(),
100        ))
101    }
102
103    fn deserialize(
104        &self,
105        _metadata: &[u8],
106        _session: &VortexSession,
107    ) -> VortexResult<Self::Options> {
108        let opts = pb::BinaryOpts::decode(_metadata)?;
109        Operator::try_from(opts.op)
110    }
111
112    fn arity(&self, _options: &Self::Options) -> Arity {
113        Arity::Exact(2)
114    }
115
116    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
117        match child_idx {
118            0 => ChildName::from("lhs"),
119            1 => ChildName::from("rhs"),
120            _ => unreachable!("Binary has only two children"),
121        }
122    }
123
124    fn fmt_sql(
125        &self,
126        operator: &Operator,
127        expr: &dyn ExprDisplay,
128        f: &mut Formatter<'_>,
129    ) -> std::fmt::Result {
130        write!(f, "(")?;
131        Display::fmt(expr.display_child(0), f)?;
132        write!(f, " {} ", operator)?;
133        Display::fmt(expr.display_child(1), f)?;
134        write!(f, ")")
135    }
136
137    fn return_dtype(&self, operator: &Operator, arg_dtypes: &[DType]) -> VortexResult<DType> {
138        let lhs = &arg_dtypes[0];
139        let rhs = &arg_dtypes[1];
140
141        if operator.is_arithmetic() {
142            if lhs.is_primitive() && lhs.eq_ignore_nullability(rhs) {
143                return Ok(lhs.with_nullability(lhs.nullability() | rhs.nullability()));
144            }
145
146            if let DType::Decimal(decimal_dtype, _) = lhs
147                && lhs.eq_ignore_nullability(rhs)
148            {
149                let numeric_op = NumericOperator::try_from(*operator)?;
150                return Ok(DType::Decimal(
151                    numeric_op_result_decimal_dtype(*decimal_dtype, numeric_op)?,
152                    lhs.nullability() | rhs.nullability(),
153                ));
154            }
155            vortex_bail!(
156                "incompatible types for arithmetic operation: {} {}",
157                lhs,
158                rhs
159            );
160        }
161
162        if operator.is_comparison()
163            && !lhs.eq_ignore_nullability(rhs)
164            && !lhs.is_extension()
165            && !rhs.is_extension()
166        {
167            vortex_bail!("Cannot compare different DTypes {} and {}", lhs, rhs);
168        }
169
170        if matches!(operator, Operator::And | Operator::Or)
171            && !(lhs.is_boolean() && rhs.is_boolean())
172        {
173            vortex_bail!("'{operator}' requires bool operands, got {lhs} and {rhs}",);
174        }
175
176        Ok(DType::Bool((lhs.is_nullable() || rhs.is_nullable()).into()))
177    }
178
179    fn execute(
180        &self,
181        op: &Operator,
182        args: &dyn ExecutionArgs,
183        ctx: &mut ExecutionCtx,
184    ) -> VortexResult<ArrayRef> {
185        let lhs = args.get(0)?;
186        let rhs = args.get(1)?;
187
188        match op {
189            Operator::Eq => execute_compare(&lhs, &rhs, CompareOperator::Eq, ctx),
190            Operator::NotEq => execute_compare(&lhs, &rhs, CompareOperator::NotEq, ctx),
191            Operator::Lt => execute_compare(&lhs, &rhs, CompareOperator::Lt, ctx),
192            Operator::Lte => execute_compare(&lhs, &rhs, CompareOperator::Lte, ctx),
193            Operator::Gt => execute_compare(&lhs, &rhs, CompareOperator::Gt, ctx),
194            Operator::Gte => execute_compare(&lhs, &rhs, CompareOperator::Gte, ctx),
195            Operator::And => execute_boolean(lhs, rhs, Operator::And, ctx),
196            Operator::Or => execute_boolean(lhs, rhs, Operator::Or, ctx),
197            Operator::Add => execute_numeric(&lhs, &rhs, NumericOperator::Add, ctx),
198            Operator::Sub => execute_numeric(&lhs, &rhs, NumericOperator::Sub, ctx),
199            Operator::Mul => execute_numeric(&lhs, &rhs, NumericOperator::Mul, ctx),
200            Operator::Div => execute_numeric(&lhs, &rhs, NumericOperator::Div, ctx),
201        }
202    }
203
204    fn simplify(
205        &self,
206        operator: &Operator,
207        expr: &BoundExpression,
208    ) -> VortexResult<Option<BoundExpression>> {
209        let lhs = expr.child(0);
210        let rhs = expr.child(1);
211
212        let is_literal_null =
213            |expr: &BoundExpression| expr.as_opt::<Literal>().is_some_and(Scalar::is_null);
214
215        // Binding already rejected ill-typed comparisons such as `int_col = null_utf8`, so a
216        // comparison against a null literal always evaluates to null.
217        if operator.is_comparison() && (is_literal_null(lhs) || is_literal_null(rhs)) {
218            return Ok(Some(bound::lit(Scalar::null(DType::Bool(
219                Nullability::Nullable,
220            )))));
221        }
222
223        let bool_literal = |expr: &BoundExpression| {
224            expr.as_opt::<Literal>()?
225                .as_bool_opt()
226                .map(|value| value.value())
227        };
228
229        // AND/OR use Kleene three-valued logic. `None` below is a boolean null.
230        //
231        // AND:
232        // - false AND x => false
233        // - true  AND x => x
234        // - null  AND null => null
235        //
236        // OR:
237        // - true  OR x => true
238        // - false OR x => x
239        // - null  OR null => null
240        //
241        // Other null cases either fall out of the identity/annihilator rules
242        // above (`null AND true`, `null OR false`) or cannot be simplified under
243        // Kleene semantics (`null AND x`, `null OR x` for non-literal `x`).
244        Ok(match operator {
245            Operator::And => match (bool_literal(lhs), bool_literal(rhs)) {
246                (Some(Some(false)), _) | (_, Some(Some(false))) => Some(bound::lit(false)),
247                (Some(Some(true)), _) => Some(rhs.clone()),
248                (_, Some(Some(true))) => Some(lhs.clone()),
249                (Some(None), Some(None)) => Some(lhs.clone()),
250                _ => None,
251            },
252            Operator::Or => match (bool_literal(lhs), bool_literal(rhs)) {
253                (Some(Some(true)), _) | (_, Some(Some(true))) => Some(bound::lit(true)),
254                (Some(Some(false)), _) => Some(rhs.clone()),
255                (_, Some(Some(false))) => Some(lhs.clone()),
256                (Some(None), Some(None)) => Some(lhs.clone()),
257                _ => None,
258            },
259            _ => None,
260        })
261    }
262
263    fn reduce<T: ReduceNode>(&self, operator: &Operator, node: &T) -> VortexResult<Option<T>> {
264        if !matches!(operator, Operator::And | Operator::Or) {
265            return Ok(None);
266        }
267        let left = node.child(0);
268        let right = node.child(1);
269
270        let left_const = left.as_constant();
271        let right_const = right.as_constant();
272
273        // We don't handle Kleene NULL reduction here (.value() returns None
274        // for NULL). This will be reduced in the boolean kernel during
275        // execution. Not handling the case keeps the code much simpler.
276        let left_const = left_const
277            .and_then(|s| s.value().cloned())
278            .map(|v| v.as_bool());
279        let right_const = right_const
280            .and_then(|s| s.value().cloned())
281            .map(|v| v.as_bool());
282        let is_and = *operator == Operator::And;
283
284        Ok(Some(match (left_const, right_const) {
285            (None, None) => return Ok(None),
286            (Some(left_const), Some(right_const)) => left.new_constant(
287                KLEENE_LUT[is_and as usize][left_const as usize][right_const as usize].into(),
288            ),
289            (Some(constant), None) => kleene_one_const(right, constant, is_and),
290            (None, Some(constant)) => kleene_one_const(left, constant, is_and),
291        }))
292    }
293
294    fn is_strict(&self, operator: &Operator) -> bool {
295        !matches!(operator, Operator::And | Operator::Or)
296    }
297
298    fn is_infallible(&self, operator: &Operator) -> bool {
299        // Arithmetic operations could be better modelled here.
300        matches!(
301            operator,
302            Operator::Eq
303                | Operator::NotEq
304                | Operator::Gt
305                | Operator::Gte
306                | Operator::Lt
307                | Operator::Lte
308                | Operator::And
309                | Operator::Or
310        )
311    }
312}
313
314#[cfg(test)]
315mod tests {
316    use vortex_error::VortexExpect;
317    use vortex_error::VortexResult;
318
319    use super::*;
320    use crate::IntoArray;
321    use crate::VortexSessionExecute;
322    use crate::array_session;
323    use crate::arrays::Bool;
324    use crate::arrays::BoolArray;
325    use crate::arrays::ConstantArray;
326    use crate::arrays::PrimitiveArray;
327    use crate::arrays::ScalarFn;
328    use crate::arrays::scalar_fn::ExactScalarFn;
329    use crate::arrays::scalar_fn::ScalarFnArrayExt;
330    use crate::assert_arrays_eq;
331    use crate::builtins::ArrayBuiltins;
332    use crate::dtype::DType;
333    use crate::dtype::Nullability;
334    use crate::dtype::PType;
335    use crate::expr::Expression;
336    use crate::expr::and;
337    use crate::expr::and_collect;
338    use crate::expr::col;
339    use crate::expr::eq;
340    use crate::expr::gt;
341    use crate::expr::gt_eq;
342    use crate::expr::lit;
343    use crate::expr::lt;
344    use crate::expr::lt_eq;
345    use crate::expr::not_eq;
346    use crate::expr::or;
347    use crate::expr::or_collect;
348    use crate::expr::test_harness;
349    use crate::optimizer::ArrayOptimizer;
350    use crate::scalar::Scalar;
351
352    #[test]
353    fn and_collect_balanced() {
354        let values = vec![lit(1), lit(2), lit(3), lit(4), lit(5)];
355
356        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @r"
357        vortex.binary(and)
358        ├── lhs: vortex.binary(and)
359        │   ├── lhs: vortex.literal(1i32)
360        │   └── rhs: vortex.literal(2i32)
361        └── rhs: vortex.binary(and)
362            ├── lhs: vortex.binary(and)
363            │   ├── lhs: vortex.literal(3i32)
364            │   └── rhs: vortex.literal(4i32)
365            └── rhs: vortex.literal(5i32)
366        ");
367
368        // 4 elements: and(and(1, 2), and(3, 4)) - perfectly balanced
369        let values = vec![lit(1), lit(2), lit(3), lit(4)];
370        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @r"
371        vortex.binary(and)
372        ├── lhs: vortex.binary(and)
373        │   ├── lhs: vortex.literal(1i32)
374        │   └── rhs: vortex.literal(2i32)
375        └── rhs: vortex.binary(and)
376            ├── lhs: vortex.literal(3i32)
377            └── rhs: vortex.literal(4i32)
378        ");
379
380        // 1 element: just the element
381        let values = vec![lit(1)];
382        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @"vortex.literal(1i32)");
383
384        // 0 elements: None
385        let values: Vec<Expression> = vec![];
386        assert!(and_collect(values.into_iter()).is_none());
387    }
388
389    #[test]
390    fn or_collect_balanced() {
391        // 4 elements: or(or(1, 2), or(3, 4)) - perfectly balanced
392        let values = vec![lit(1), lit(2), lit(3), lit(4)];
393        insta::assert_snapshot!(or_collect(values.into_iter()).unwrap().display_tree(), @r"
394        vortex.binary(or)
395        ├── lhs: vortex.binary(or)
396        │   ├── lhs: vortex.literal(1i32)
397        │   └── rhs: vortex.literal(2i32)
398        └── rhs: vortex.binary(or)
399            ├── lhs: vortex.literal(3i32)
400            └── rhs: vortex.literal(4i32)
401        ");
402    }
403
404    #[test]
405    fn dtype() {
406        let dtype = test_harness::struct_dtype();
407        let bool1: Expression = col("bool1");
408        let bool2: Expression = col("bool2");
409        assert_eq!(
410            and(bool1.clone(), bool2.clone())
411                .return_dtype(&dtype)
412                .unwrap(),
413            DType::Bool(Nullability::NonNullable)
414        );
415        assert_eq!(
416            or(bool1, bool2).return_dtype(&dtype).unwrap(),
417            DType::Bool(Nullability::NonNullable)
418        );
419
420        let col1: Expression = col("col1");
421        let col2: Expression = col("col2");
422
423        assert_eq!(
424            eq(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
425            DType::Bool(Nullability::Nullable)
426        );
427        assert_eq!(
428            not_eq(col1.clone(), col2.clone())
429                .return_dtype(&dtype)
430                .unwrap(),
431            DType::Bool(Nullability::Nullable)
432        );
433        assert_eq!(
434            gt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
435            DType::Bool(Nullability::Nullable)
436        );
437        assert_eq!(
438            gt_eq(col1.clone(), col2.clone())
439                .return_dtype(&dtype)
440                .unwrap(),
441            DType::Bool(Nullability::Nullable)
442        );
443        assert_eq!(
444            lt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
445            DType::Bool(Nullability::Nullable)
446        );
447        assert_eq!(
448            lt_eq(col1.clone(), col2.clone())
449                .return_dtype(&dtype)
450                .unwrap(),
451            DType::Bool(Nullability::Nullable)
452        );
453
454        assert_eq!(
455            or(lt(col1.clone(), col2.clone()), not_eq(col1, col2))
456                .return_dtype(&dtype)
457                .unwrap(),
458            DType::Bool(Nullability::Nullable)
459        );
460    }
461
462    #[test]
463    fn comparison_with_typed_null_simplifies_after_type_check() -> VortexResult<()> {
464        let dtype = test_harness::struct_dtype();
465
466        let expr = eq(
467            col("col1"),
468            lit(Scalar::null(DType::Primitive(
469                PType::U16,
470                Nullability::Nullable,
471            ))),
472        );
473
474        assert_eq!(
475            expr.bind(&dtype)?.optimize_recursive()?,
476            lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&dtype)?
477        );
478        Ok(())
479    }
480
481    #[test]
482    fn comparison_with_incompatible_null_still_type_checks() {
483        let dtype = test_harness::struct_dtype();
484        let expr = eq(
485            col("col1"),
486            lit(Scalar::null(DType::Utf8(Nullability::Nullable))),
487        );
488
489        assert!(expr.bind(&dtype).is_err());
490    }
491
492    #[test]
493    fn test_display_print() {
494        let expr = gt(lit(1), lit(2));
495        assert_eq!(format!("{expr}"), "(1i32 > 2i32)");
496    }
497
498    /// Regression test for GitHub issue #5947: struct comparison in filter expressions should work
499    /// using `make_comparator` instead of Arrow's `cmp` functions which don't support nested types.
500    #[test]
501    fn test_struct_comparison() {
502        use crate::IntoArray;
503        use crate::arrays::StructArray;
504
505        // Create a struct array with one element for testing.
506        let lhs_struct = StructArray::from_fields(&[
507            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
508            ("b", PrimitiveArray::from_iter([3i32]).into_array()),
509        ])
510        .unwrap()
511        .into_array();
512
513        let rhs_struct_equal = StructArray::from_fields(&[
514            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
515            ("b", PrimitiveArray::from_iter([3i32]).into_array()),
516        ])
517        .unwrap()
518        .into_array();
519
520        let rhs_struct_different = StructArray::from_fields(&[
521            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
522            ("b", PrimitiveArray::from_iter([4i32]).into_array()),
523        ])
524        .unwrap()
525        .into_array();
526
527        // Test using binary method directly
528        let result_equal = lhs_struct.binary(rhs_struct_equal, Operator::Eq).unwrap();
529        assert_eq!(
530            result_equal
531                .execute_scalar(0, &mut array_session().create_execution_ctx())
532                .vortex_expect("value"),
533            Scalar::bool(true, Nullability::NonNullable),
534            "Equal structs should be equal"
535        );
536
537        let result_different = lhs_struct
538            .binary(rhs_struct_different, Operator::Eq)
539            .unwrap();
540        assert_eq!(
541            result_different
542                .execute_scalar(0, &mut array_session().create_execution_ctx())
543                .vortex_expect("value"),
544            Scalar::bool(false, Nullability::NonNullable),
545            "Different structs should not be equal"
546        );
547    }
548
549    #[test]
550    fn test_or_kleene_validity() {
551        let mut ctx = array_session().create_execution_ctx();
552        use crate::arrays::StructArray;
553        use crate::expr::col;
554
555        let struct_arr = StructArray::from_fields(&[
556            ("a", BoolArray::from_iter([Some(true)]).into_array()),
557            (
558                "b",
559                BoolArray::from_iter([Option::<bool>::None]).into_array(),
560            ),
561        ])
562        .unwrap()
563        .into_array();
564
565        let expr = or(col("a"), col("b"));
566        let result = struct_arr.apply(&expr).unwrap();
567
568        assert_arrays_eq!(
569            result,
570            BoolArray::from_iter([Some(true)]).into_array(),
571            &mut ctx
572        )
573    }
574
575    #[test]
576    fn test_scalar_subtract_unsigned() {
577        let mut ctx = array_session().create_execution_ctx();
578        use vortex_buffer::buffer;
579
580        use crate::IntoArray;
581        use crate::arrays::ConstantArray;
582        use crate::arrays::PrimitiveArray;
583
584        let values = buffer![1u16, 2, 3].into_array();
585        let rhs = ConstantArray::new(Scalar::from(1u16), 3).into_array();
586        let result = values.binary(rhs, Operator::Sub).unwrap();
587        assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx);
588    }
589
590    #[test]
591    fn test_scalar_subtract_signed() {
592        let mut ctx = array_session().create_execution_ctx();
593        use vortex_buffer::buffer;
594
595        use crate::IntoArray;
596        use crate::arrays::ConstantArray;
597        use crate::arrays::PrimitiveArray;
598
599        let values = buffer![1i64, 2, 3].into_array();
600        let rhs = ConstantArray::new(Scalar::from(-1i64), 3).into_array();
601        let result = values.binary(rhs, Operator::Sub).unwrap();
602        assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx);
603    }
604
605    #[test]
606    fn test_scalar_subtract_nullable() {
607        let mut ctx = array_session().create_execution_ctx();
608        use crate::IntoArray;
609        use crate::arrays::ConstantArray;
610        use crate::arrays::PrimitiveArray;
611
612        let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]);
613        let rhs = ConstantArray::new(Scalar::from(Some(1u16)), 4).into_array();
614        let result = values.into_array().binary(rhs, Operator::Sub).unwrap();
615        assert_arrays_eq!(
616            result,
617            PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]),
618            &mut ctx
619        );
620    }
621
622    #[test]
623    fn test_scalar_subtract_float() {
624        let mut ctx = array_session().create_execution_ctx();
625        use vortex_buffer::buffer;
626
627        use crate::IntoArray;
628        use crate::arrays::ConstantArray;
629        use crate::arrays::PrimitiveArray;
630
631        let values = buffer![1.0f64, 2.0, 3.0].into_array();
632        let rhs = ConstantArray::new(Scalar::from(-1f64), 3).into_array();
633        let result = values.binary(rhs, Operator::Sub).unwrap();
634        assert_arrays_eq!(
635            result,
636            PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]),
637            &mut ctx
638        );
639    }
640
641    #[test]
642    fn test_scalar_subtract_float_underflow_is_ok() {
643        use vortex_buffer::buffer;
644
645        let values = buffer![f32::MIN, 2.0, 3.0].into_array();
646        let rhs1 = ConstantArray::new(Scalar::from(1.0f32), 3).into_array();
647        let _results = values.binary(rhs1, Operator::Sub).unwrap();
648        let values = buffer![f32::MIN, 2.0, 3.0].into_array();
649        let rhs2 = ConstantArray::new(Scalar::from(f32::MAX), 3).into_array();
650        let _results = values.binary(rhs2, Operator::Sub).unwrap();
651    }
652
653    #[test]
654    fn test_and_reduce() -> VortexResult<()> {
655        let mut ctx = array_session().create_execution_ctx();
656        let left = BoolArray::from_iter([true, true, false]).into_array();
657        let right = ConstantArray::new(true, 3).into_array();
658
659        let array = Binary::try_new(left.clone(), right.clone(), Operator::And)?
660            .into_array()
661            .optimize()?; // calls reduce()
662        assert!(array.is::<Bool>()); // and(left, const) -> left
663        assert_arrays_eq!(array, left, &mut ctx);
664
665        let array = Binary::try_new(right.clone(), left.clone(), Operator::And)?
666            .into_array()
667            .optimize()?;
668        assert!(array.is::<Bool>()); // and(const, left) -> left
669        assert_arrays_eq!(array, left, &mut ctx);
670
671        let array = Binary::try_new(left.clone(), right.clone(), Operator::Or)?
672            .into_array()
673            .optimize()?;
674        assert_eq!(array.as_constant(), Some(true.into())); // or(left, const) -> true
675
676        let array = Binary::try_new(right.clone(), left, Operator::Or)?
677            .into_array()
678            .optimize()?;
679        assert_eq!(array.as_constant(), Some(true.into())); // or(const, left) -> true
680
681        let array = Binary::try_new(right.clone(), right.clone(), Operator::Or)?
682            .into_array()
683            .optimize()?;
684        assert_eq!(array.as_constant(), Some(true.into())); // or(const, const) -> const
685
686        let left_false = ConstantArray::new(false, 3).into_array();
687
688        let array = Binary::try_new(right.clone(), left_false.clone(), Operator::Or)?
689            .into_array()
690            .optimize()?;
691        assert_eq!(array.as_constant(), Some(true.into()));
692        let array = Binary::try_new(left_false, right, Operator::And)?
693            .into_array()
694            .optimize()?;
695        assert_eq!(array.as_constant(), Some(false.into()));
696
697        Ok(())
698    }
699
700    #[test]
701    fn test_and_reduce_nullable() -> VortexResult<()> {
702        let mut ctx = array_session().create_execution_ctx();
703        let nullable_true =
704            ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array();
705        let right = BoolArray::from_iter([true, false, true]).into_array();
706
707        let array = Binary::try_new(nullable_true, right.clone(), Operator::And)?
708            .into_array()
709            .optimize()?;
710        assert_arrays_eq!(array, right, &mut ctx);
711        Ok(())
712    }
713
714    #[test]
715    fn test_and_reject_non_bool() {
716        let lhs = ConstantArray::new(7i32, 3).into_array();
717        let rhs = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
718        assert!(Binary::try_new(lhs, rhs, Operator::And).is_err());
719    }
720
721    #[test]
722    fn test_isnull_and_reduce() -> VortexResult<()> {
723        let mut ctx = array_session().create_execution_ctx();
724        let left = BoolArray::from_iter([Some(true), Some(true), None]).into_array();
725        let right = BoolArray::from_iter([true, false, true]).into_array();
726
727        // IsNull(and(x, y)) -> and(IsNull(x), y) -> and(not(x.validity), y)
728        //            ^ nullable
729        for (lhs, rhs) in [(left.clone(), right.clone()), (right.clone(), left)] {
730            let array = lhs
731                .binary(rhs.clone(), Operator::And)?
732                .is_null()?
733                .optimize()?;
734            assert_eq!(*array.as_::<ExactScalarFn<Binary>>().options, Operator::And);
735            assert_arrays_eq!(
736                array.as_::<ScalarFn>().get_child(0), // not(left.validity())
737                BoolArray::from_iter([false, false, true]),
738                &mut ctx
739            );
740            assert_arrays_eq!(array.as_::<ScalarFn>().get_child(1), right, &mut ctx);
741        }
742        Ok(())
743    }
744
745    #[test]
746    fn test_isnotnull_or_reduce() -> VortexResult<()> {
747        let mut ctx = array_session().create_execution_ctx();
748        let left = BoolArray::from_iter([Some(true), Some(true), None]).into_array();
749        let right = BoolArray::from_iter([true, false, true]).into_array();
750
751        // IsNotNull(or(x, y)) -> or(IsNotNull(x), y) -> or(x.validity, y)
752        //              ^ nullable
753        for (lhs, rhs) in [(left.clone(), right.clone()), (right.clone(), left)] {
754            let array = lhs
755                .binary(rhs.clone(), Operator::Or)?
756                .is_not_null()?
757                .optimize()?;
758            assert_eq!(*array.as_::<ExactScalarFn<Binary>>().options, Operator::Or);
759            assert_arrays_eq!(
760                array.as_::<ScalarFn>().get_child(0), // left.validity()
761                BoolArray::from_iter([true, true, false]),
762                &mut ctx
763            );
764            assert_arrays_eq!(array.as_::<ScalarFn>().get_child(1), right, &mut ctx);
765        }
766        Ok(())
767    }
768}
vortex_array/scalar_fn/fns/binary/mod.rs

vortex_array/scalar_fn/fns/binary/
mod.rs
