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), [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        // Simplification must keep the result nullability binding inferred: the annihilator
245        // literal and the identity child can both be less nullable than `lhs OP rhs`.
246        let result_nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
247        let nullability = Nullability::from(result_nullable);
248        let bool_lit = |value: bool| bound::lit(Scalar::bool(value, nullability));
249        let widen = |child: &BoundExpression| {
250            if child.dtype().is_nullable() == result_nullable {
251                child.clone()
252            } else {
253                bound::cast(child.clone(), DType::Bool(nullability))
254            }
255        };
256
257        Ok(match operator {
258            Operator::And => match (bool_literal(lhs), bool_literal(rhs)) {
259                (Some(Some(false)), _) | (_, Some(Some(false))) => Some(bool_lit(false)),
260                (Some(Some(true)), _) => Some(widen(rhs)),
261                (_, Some(Some(true))) => Some(widen(lhs)),
262                (Some(None), Some(None)) => Some(widen(lhs)),
263                _ => None,
264            },
265            Operator::Or => match (bool_literal(lhs), bool_literal(rhs)) {
266                (Some(Some(true)), _) | (_, Some(Some(true))) => Some(bool_lit(true)),
267                (Some(Some(false)), _) => Some(widen(rhs)),
268                (_, Some(Some(false))) => Some(widen(lhs)),
269                (Some(None), Some(None)) => Some(widen(lhs)),
270                _ => None,
271            },
272            _ => None,
273        })
274    }
275
276    fn reduce<T: ReduceNode>(&self, operator: &Operator, node: &T) -> VortexResult<Option<T>> {
277        if !matches!(operator, Operator::And | Operator::Or) {
278            return Ok(None);
279        }
280        let left = node.child(0);
281        let right = node.child(1);
282
283        let left_const = left.as_constant();
284        let right_const = right.as_constant();
285
286        // We don't handle Kleene NULL reduction here (.value() returns None
287        // for NULL). This will be reduced in the boolean kernel during
288        // execution. Not handling the case keeps the code much simpler.
289        let left_const = left_const
290            .and_then(|s| s.value().cloned())
291            .map(|v| v.as_bool());
292        let right_const = right_const
293            .and_then(|s| s.value().cloned())
294            .map(|v| v.as_bool());
295        let is_and = *operator == Operator::And;
296
297        Ok(Some(match (left_const, right_const) {
298            (None, None) => return Ok(None),
299            (Some(left_const), Some(right_const)) => left.new_constant(
300                KLEENE_LUT[is_and as usize][left_const as usize][right_const as usize].into(),
301            ),
302            (Some(constant), None) => kleene_one_const(right, constant, is_and),
303            (None, Some(constant)) => kleene_one_const(left, constant, is_and),
304        }))
305    }
306
307    fn is_strict(&self, operator: &Operator) -> bool {
308        !matches!(operator, Operator::And | Operator::Or)
309    }
310
311    fn is_infallible(&self, operator: &Operator) -> bool {
312        // Arithmetic operations could be better modelled here.
313        matches!(
314            operator,
315            Operator::Eq
316                | Operator::NotEq
317                | Operator::Gt
318                | Operator::Gte
319                | Operator::Lt
320                | Operator::Lte
321                | Operator::And
322                | Operator::Or
323        )
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use vortex_error::VortexExpect;
330    use vortex_error::VortexResult;
331
332    use super::*;
333    use crate::IntoArray;
334    use crate::VortexSessionExecute;
335    use crate::array_session;
336    use crate::arrays::Bool;
337    use crate::arrays::BoolArray;
338    use crate::arrays::ConstantArray;
339    use crate::arrays::PrimitiveArray;
340    use crate::arrays::ScalarFn;
341    use crate::arrays::scalar_fn::ExactScalarFn;
342    use crate::arrays::scalar_fn::ScalarFnArrayExt;
343    use crate::assert_arrays_eq;
344    use crate::builtins::ArrayBuiltins;
345    use crate::dtype::DType;
346    use crate::dtype::Nullability;
347    use crate::dtype::PType;
348    use crate::expr::Expression;
349    use crate::expr::and;
350    use crate::expr::and_collect;
351    use crate::expr::col;
352    use crate::expr::eq;
353    use crate::expr::gt;
354    use crate::expr::gt_eq;
355    use crate::expr::lit;
356    use crate::expr::lt;
357    use crate::expr::lt_eq;
358    use crate::expr::not_eq;
359    use crate::expr::or;
360    use crate::expr::or_collect;
361    use crate::expr::test_harness;
362    use crate::optimizer::ArrayOptimizer;
363    use crate::scalar::Scalar;
364
365    #[test]
366    fn and_collect_balanced() {
367        let values = vec![lit(1), lit(2), lit(3), lit(4), lit(5)];
368
369        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @r"
370        vortex.binary(and)
371        ├── lhs: vortex.binary(and)
372        │   ├── lhs: vortex.literal(1i32)
373        │   └── rhs: vortex.literal(2i32)
374        └── rhs: vortex.binary(and)
375            ├── lhs: vortex.binary(and)
376            │   ├── lhs: vortex.literal(3i32)
377            │   └── rhs: vortex.literal(4i32)
378            └── rhs: vortex.literal(5i32)
379        ");
380
381        // 4 elements: and(and(1, 2), and(3, 4)) - perfectly balanced
382        let values = vec![lit(1), lit(2), lit(3), lit(4)];
383        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @r"
384        vortex.binary(and)
385        ├── lhs: vortex.binary(and)
386        │   ├── lhs: vortex.literal(1i32)
387        │   └── rhs: vortex.literal(2i32)
388        └── rhs: vortex.binary(and)
389            ├── lhs: vortex.literal(3i32)
390            └── rhs: vortex.literal(4i32)
391        ");
392
393        // 1 element: just the element
394        let values = vec![lit(1)];
395        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @"vortex.literal(1i32)");
396
397        // 0 elements: None
398        let values: Vec<Expression> = vec![];
399        assert!(and_collect(values.into_iter()).is_none());
400    }
401
402    #[test]
403    fn or_collect_balanced() {
404        // 4 elements: or(or(1, 2), or(3, 4)) - perfectly balanced
405        let values = vec![lit(1), lit(2), lit(3), lit(4)];
406        insta::assert_snapshot!(or_collect(values.into_iter()).unwrap().display_tree(), @r"
407        vortex.binary(or)
408        ├── lhs: vortex.binary(or)
409        │   ├── lhs: vortex.literal(1i32)
410        │   └── rhs: vortex.literal(2i32)
411        └── rhs: vortex.binary(or)
412            ├── lhs: vortex.literal(3i32)
413            └── rhs: vortex.literal(4i32)
414        ");
415    }
416
417    #[test]
418    fn dtype() {
419        let dtype = test_harness::struct_dtype();
420        let bool1: Expression = col("bool1");
421        let bool2: Expression = col("bool2");
422        assert_eq!(
423            and(bool1.clone(), bool2.clone())
424                .return_dtype(&dtype)
425                .unwrap(),
426            DType::Bool(Nullability::NonNullable)
427        );
428        assert_eq!(
429            or(bool1, bool2).return_dtype(&dtype).unwrap(),
430            DType::Bool(Nullability::NonNullable)
431        );
432
433        let col1: Expression = col("col1");
434        let col2: Expression = col("col2");
435
436        assert_eq!(
437            eq(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
438            DType::Bool(Nullability::Nullable)
439        );
440        assert_eq!(
441            not_eq(col1.clone(), col2.clone())
442                .return_dtype(&dtype)
443                .unwrap(),
444            DType::Bool(Nullability::Nullable)
445        );
446        assert_eq!(
447            gt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
448            DType::Bool(Nullability::Nullable)
449        );
450        assert_eq!(
451            gt_eq(col1.clone(), col2.clone())
452                .return_dtype(&dtype)
453                .unwrap(),
454            DType::Bool(Nullability::Nullable)
455        );
456        assert_eq!(
457            lt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
458            DType::Bool(Nullability::Nullable)
459        );
460        assert_eq!(
461            lt_eq(col1.clone(), col2.clone())
462                .return_dtype(&dtype)
463                .unwrap(),
464            DType::Bool(Nullability::Nullable)
465        );
466
467        assert_eq!(
468            or(lt(col1.clone(), col2.clone()), not_eq(col1, col2))
469                .return_dtype(&dtype)
470                .unwrap(),
471            DType::Bool(Nullability::Nullable)
472        );
473    }
474
475    #[test]
476    fn comparison_with_typed_null_simplifies_after_type_check() -> VortexResult<()> {
477        let dtype = test_harness::struct_dtype();
478
479        let expr = eq(
480            col("col1"),
481            lit(Scalar::null(DType::Primitive(
482                PType::U16,
483                Nullability::Nullable,
484            ))),
485        );
486
487        assert_eq!(
488            expr.bind(&dtype)?.optimize_recursive()?,
489            lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&dtype)?
490        );
491        Ok(())
492    }
493
494    #[test]
495    fn comparison_with_incompatible_null_still_type_checks() {
496        let dtype = test_harness::struct_dtype();
497        let expr = eq(
498            col("col1"),
499            lit(Scalar::null(DType::Utf8(Nullability::Nullable))),
500        );
501
502        assert!(expr.bind(&dtype).is_err());
503    }
504
505    #[test]
506    fn test_display_print() {
507        let expr = gt(lit(1), lit(2));
508        assert_eq!(format!("{expr}"), "(1i32 > 2i32)");
509    }
510
511    /// Regression test for GitHub issue #5947: struct comparison in filter expressions should work
512    /// using `make_comparator` instead of Arrow's `cmp` functions which don't support nested types.
513    #[test]
514    fn test_struct_comparison() {
515        use crate::IntoArray;
516        use crate::arrays::StructArray;
517
518        // Create a struct array with one element for testing.
519        let lhs_struct = StructArray::from_fields(&[
520            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
521            ("b", PrimitiveArray::from_iter([3i32]).into_array()),
522        ])
523        .unwrap()
524        .into_array();
525
526        let rhs_struct_equal = StructArray::from_fields(&[
527            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
528            ("b", PrimitiveArray::from_iter([3i32]).into_array()),
529        ])
530        .unwrap()
531        .into_array();
532
533        let rhs_struct_different = StructArray::from_fields(&[
534            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
535            ("b", PrimitiveArray::from_iter([4i32]).into_array()),
536        ])
537        .unwrap()
538        .into_array();
539
540        // Test using binary method directly
541        let result_equal = lhs_struct.binary(rhs_struct_equal, Operator::Eq).unwrap();
542        assert_eq!(
543            result_equal
544                .execute_scalar(0, &mut array_session().create_execution_ctx())
545                .vortex_expect("value"),
546            Scalar::bool(true, Nullability::NonNullable),
547            "Equal structs should be equal"
548        );
549
550        let result_different = lhs_struct
551            .binary(rhs_struct_different, Operator::Eq)
552            .unwrap();
553        assert_eq!(
554            result_different
555                .execute_scalar(0, &mut array_session().create_execution_ctx())
556                .vortex_expect("value"),
557            Scalar::bool(false, Nullability::NonNullable),
558            "Different structs should not be equal"
559        );
560    }
561
562    #[test]
563    fn test_or_kleene_validity() {
564        let mut ctx = array_session().create_execution_ctx();
565        use crate::arrays::StructArray;
566        use crate::expr::col;
567
568        let struct_arr = StructArray::from_fields(&[
569            ("a", BoolArray::from_iter([Some(true)]).into_array()),
570            (
571                "b",
572                BoolArray::from_iter([Option::<bool>::None]).into_array(),
573            ),
574        ])
575        .unwrap()
576        .into_array();
577
578        let expr = or(col("a"), col("b"));
579        let result = struct_arr.apply(&expr).unwrap();
580
581        assert_arrays_eq!(
582            result,
583            BoolArray::from_iter([Some(true)]).into_array(),
584            &mut ctx
585        )
586    }
587
588    #[test]
589    fn test_scalar_subtract_unsigned() {
590        let mut ctx = array_session().create_execution_ctx();
591        use vortex_buffer::buffer;
592
593        use crate::IntoArray;
594        use crate::arrays::ConstantArray;
595        use crate::arrays::PrimitiveArray;
596
597        let values = buffer![1u16, 2, 3].into_array();
598        let rhs = ConstantArray::new(Scalar::from(1u16), 3).into_array();
599        let result = values.binary(rhs, Operator::Sub).unwrap();
600        assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx);
601    }
602
603    #[test]
604    fn test_scalar_subtract_signed() {
605        let mut ctx = array_session().create_execution_ctx();
606        use vortex_buffer::buffer;
607
608        use crate::IntoArray;
609        use crate::arrays::ConstantArray;
610        use crate::arrays::PrimitiveArray;
611
612        let values = buffer![1i64, 2, 3].into_array();
613        let rhs = ConstantArray::new(Scalar::from(-1i64), 3).into_array();
614        let result = values.binary(rhs, Operator::Sub).unwrap();
615        assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx);
616    }
617
618    #[test]
619    fn test_scalar_subtract_nullable() {
620        let mut ctx = array_session().create_execution_ctx();
621        use crate::IntoArray;
622        use crate::arrays::ConstantArray;
623        use crate::arrays::PrimitiveArray;
624
625        let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]);
626        let rhs = ConstantArray::new(Scalar::from(Some(1u16)), 4).into_array();
627        let result = values.into_array().binary(rhs, Operator::Sub).unwrap();
628        assert_arrays_eq!(
629            result,
630            PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]),
631            &mut ctx
632        );
633    }
634
635    #[test]
636    fn test_scalar_subtract_float() {
637        let mut ctx = array_session().create_execution_ctx();
638        use vortex_buffer::buffer;
639
640        use crate::IntoArray;
641        use crate::arrays::ConstantArray;
642        use crate::arrays::PrimitiveArray;
643
644        let values = buffer![1.0f64, 2.0, 3.0].into_array();
645        let rhs = ConstantArray::new(Scalar::from(-1f64), 3).into_array();
646        let result = values.binary(rhs, Operator::Sub).unwrap();
647        assert_arrays_eq!(
648            result,
649            PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]),
650            &mut ctx
651        );
652    }
653
654    #[test]
655    fn test_scalar_subtract_float_underflow_is_ok() {
656        use vortex_buffer::buffer;
657
658        let values = buffer![f32::MIN, 2.0, 3.0].into_array();
659        let rhs1 = ConstantArray::new(Scalar::from(1.0f32), 3).into_array();
660        let _results = values.binary(rhs1, Operator::Sub).unwrap();
661        let values = buffer![f32::MIN, 2.0, 3.0].into_array();
662        let rhs2 = ConstantArray::new(Scalar::from(f32::MAX), 3).into_array();
663        let _results = values.binary(rhs2, Operator::Sub).unwrap();
664    }
665
666    #[test]
667    fn test_and_reduce() -> VortexResult<()> {
668        let mut ctx = array_session().create_execution_ctx();
669        let left = BoolArray::from_iter([true, true, false]).into_array();
670        let right = ConstantArray::new(true, 3).into_array();
671
672        let array = Binary::try_new(left.clone(), right.clone(), Operator::And)?
673            .into_array()
674            .optimize()?; // calls reduce()
675        assert!(array.is::<Bool>()); // and(left, const) -> left
676        assert_arrays_eq!(array, left, &mut ctx);
677
678        let array = Binary::try_new(right.clone(), left.clone(), Operator::And)?
679            .into_array()
680            .optimize()?;
681        assert!(array.is::<Bool>()); // and(const, left) -> left
682        assert_arrays_eq!(array, left, &mut ctx);
683
684        let array = Binary::try_new(left.clone(), right.clone(), Operator::Or)?
685            .into_array()
686            .optimize()?;
687        assert_eq!(array.as_constant(), Some(true.into())); // or(left, const) -> true
688
689        let array = Binary::try_new(right.clone(), left, Operator::Or)?
690            .into_array()
691            .optimize()?;
692        assert_eq!(array.as_constant(), Some(true.into())); // or(const, left) -> true
693
694        let array = Binary::try_new(right.clone(), right.clone(), Operator::Or)?
695            .into_array()
696            .optimize()?;
697        assert_eq!(array.as_constant(), Some(true.into())); // or(const, const) -> const
698
699        let left_false = ConstantArray::new(false, 3).into_array();
700
701        let array = Binary::try_new(right.clone(), left_false.clone(), Operator::Or)?
702            .into_array()
703            .optimize()?;
704        assert_eq!(array.as_constant(), Some(true.into()));
705        let array = Binary::try_new(left_false, right, Operator::And)?
706            .into_array()
707            .optimize()?;
708        assert_eq!(array.as_constant(), Some(false.into()));
709
710        Ok(())
711    }
712
713    #[test]
714    fn test_and_reduce_nullable() -> VortexResult<()> {
715        let mut ctx = array_session().create_execution_ctx();
716        let nullable_true =
717            ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array();
718        let right = BoolArray::from_iter([true, false, true]).into_array();
719
720        let array = Binary::try_new(nullable_true, right.clone(), Operator::And)?
721            .into_array()
722            .optimize()?;
723        assert_arrays_eq!(array, right, &mut ctx);
724        Ok(())
725    }
726
727    #[test]
728    fn test_and_reject_non_bool() {
729        let lhs = ConstantArray::new(7i32, 3).into_array();
730        let rhs = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
731        assert!(Binary::try_new(lhs, rhs, Operator::And).is_err());
732    }
733
734    #[test]
735    fn test_isnull_and_reduce() -> VortexResult<()> {
736        let mut ctx = array_session().create_execution_ctx();
737        let left = BoolArray::from_iter([Some(true), Some(true), None]).into_array();
738        let right = BoolArray::from_iter([true, false, true]).into_array();
739
740        // IsNull(and(x, y)) -> and(IsNull(x), y) -> and(not(x.validity), y)
741        //            ^ nullable
742        for (lhs, rhs) in [(left.clone(), right.clone()), (right.clone(), left)] {
743            let array = lhs
744                .binary(rhs.clone(), Operator::And)?
745                .is_null()?
746                .optimize()?;
747            assert_eq!(*array.as_::<ExactScalarFn<Binary>>().options, Operator::And);
748            assert_arrays_eq!(
749                array.as_::<ScalarFn>().get_child(0), // not(left.validity())
750                BoolArray::from_iter([false, false, true]),
751                &mut ctx
752            );
753            assert_arrays_eq!(array.as_::<ScalarFn>().get_child(1), right, &mut ctx);
754        }
755        Ok(())
756    }
757
758    #[test]
759    fn test_isnotnull_or_reduce() -> VortexResult<()> {
760        let mut ctx = array_session().create_execution_ctx();
761        let left = BoolArray::from_iter([Some(true), Some(true), None]).into_array();
762        let right = BoolArray::from_iter([true, false, true]).into_array();
763
764        // IsNotNull(or(x, y)) -> or(IsNotNull(x), y) -> or(x.validity, y)
765        //              ^ nullable
766        for (lhs, rhs) in [(left.clone(), right.clone()), (right.clone(), left)] {
767            let array = lhs
768                .binary(rhs.clone(), Operator::Or)?
769                .is_not_null()?
770                .optimize()?;
771            assert_eq!(*array.as_::<ExactScalarFn<Binary>>().options, Operator::Or);
772            assert_arrays_eq!(
773                array.as_::<ScalarFn>().get_child(0), // left.validity()
774                BoolArray::from_iter([true, true, false]),
775                &mut ctx
776            );
777            assert_arrays_eq!(array.as_::<ScalarFn>().get_child(1), right, &mut ctx);
778        }
779        Ok(())
780    }
781}
vortex_array/scalar_fn/fns/binary/mod.rs

vortex_array/scalar_fn/fns/binary/
mod.rs
