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_proto::expr as pb;
15use vortex_session::VortexSession;
16use vortex_session::registry::CachedId;
17
18use crate::ArrayRef;
19use crate::ExecutionCtx;
20use crate::dtype::DType;
21use crate::dtype::Nullability;
22use crate::expr::and;
23use crate::expr::display::ExprDisplay;
24use crate::expr::expression::Expression;
25use crate::expr::lit;
26use crate::scalar_fn::Arity;
27use crate::scalar_fn::ChildName;
28use crate::scalar_fn::ExecutionArgs;
29use crate::scalar_fn::ScalarFnId;
30use crate::scalar_fn::ScalarFnVTable;
31use crate::scalar_fn::SimplifyCtx;
32use crate::scalar_fn::fns::literal::Literal;
33use crate::scalar_fn::fns::operators::CompareOperator;
34use crate::scalar_fn::fns::operators::Operator;
35
36pub mod boolean;
37pub use boolean::BooleanExecuteAdaptor;
38pub use boolean::BooleanKernel;
39pub(crate) use boolean::execute_boolean;
40pub use boolean::kleene_boolean_buffer_scalar;
41pub use boolean::kleene_boolean_buffers;
42mod compare;
43pub use compare::*;
44mod numeric;
45pub(crate) use numeric::*;
46
47use crate::scalar::NumericOperator;
48use crate::scalar::Scalar;
49
50#[derive(Clone)]
51pub struct Binary;
52
53impl ScalarFnVTable for Binary {
54    type Options = Operator;
55
56    fn id(&self) -> ScalarFnId {
57        static ID: CachedId = CachedId::new("vortex.binary");
58        *ID
59    }
60
61    fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
62        Ok(Some(
63            pb::BinaryOpts {
64                op: (*instance).into(),
65            }
66            .encode_to_vec(),
67        ))
68    }
69
70    fn deserialize(
71        &self,
72        _metadata: &[u8],
73        _session: &VortexSession,
74    ) -> VortexResult<Self::Options> {
75        let opts = pb::BinaryOpts::decode(_metadata)?;
76        Operator::try_from(opts.op)
77    }
78
79    fn arity(&self, _options: &Self::Options) -> Arity {
80        Arity::Exact(2)
81    }
82
83    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
84        match child_idx {
85            0 => ChildName::from("lhs"),
86            1 => ChildName::from("rhs"),
87            _ => unreachable!("Binary has only two children"),
88        }
89    }
90
91    fn fmt_sql(
92        &self,
93        operator: &Operator,
94        expr: &dyn ExprDisplay,
95        f: &mut Formatter<'_>,
96    ) -> std::fmt::Result {
97        write!(f, "(")?;
98        Display::fmt(expr.display_child(0), f)?;
99        write!(f, " {} ", operator)?;
100        Display::fmt(expr.display_child(1), f)?;
101        write!(f, ")")
102    }
103
104    fn coerce_args(&self, operator: &Self::Options, args: &[DType]) -> VortexResult<Vec<DType>> {
105        let lhs = &args[0];
106        let rhs = &args[1];
107        if operator.is_arithmetic() || operator.is_comparison() {
108            let supertype = lhs.least_supertype(rhs).ok_or_else(|| {
109                vortex_error::vortex_err!("No common supertype for {} and {}", lhs, rhs)
110            })?;
111            Ok(vec![supertype.clone(), supertype])
112        } else {
113            // Boolean And/Or: no coercion
114            Ok(args.to_vec())
115        }
116    }
117
118    fn return_dtype(&self, operator: &Operator, arg_dtypes: &[DType]) -> VortexResult<DType> {
119        let lhs = &arg_dtypes[0];
120        let rhs = &arg_dtypes[1];
121
122        if operator.is_arithmetic() {
123            if lhs.is_primitive() && lhs.eq_ignore_nullability(rhs) {
124                return Ok(lhs.with_nullability(lhs.nullability() | rhs.nullability()));
125            }
126
127            if let DType::Decimal(decimal_dtype, _) = lhs
128                && lhs.eq_ignore_nullability(rhs)
129            {
130                let numeric_op = NumericOperator::try_from(*operator)?;
131                return Ok(DType::Decimal(
132                    numeric_op_result_decimal_dtype(*decimal_dtype, numeric_op)?,
133                    lhs.nullability() | rhs.nullability(),
134                ));
135            }
136            vortex_bail!(
137                "incompatible types for arithmetic operation: {} {}",
138                lhs,
139                rhs
140            );
141        }
142
143        if operator.is_comparison()
144            && !lhs.eq_ignore_nullability(rhs)
145            && !lhs.is_extension()
146            && !rhs.is_extension()
147        {
148            vortex_bail!("Cannot compare different DTypes {} and {}", lhs, rhs);
149        }
150
151        Ok(DType::Bool((lhs.is_nullable() || rhs.is_nullable()).into()))
152    }
153
154    fn execute(
155        &self,
156        op: &Operator,
157        args: &dyn ExecutionArgs,
158        ctx: &mut ExecutionCtx,
159    ) -> VortexResult<ArrayRef> {
160        let lhs = args.get(0)?;
161        let rhs = args.get(1)?;
162
163        match op {
164            Operator::Eq => execute_compare(&lhs, &rhs, CompareOperator::Eq, ctx),
165            Operator::NotEq => execute_compare(&lhs, &rhs, CompareOperator::NotEq, ctx),
166            Operator::Lt => execute_compare(&lhs, &rhs, CompareOperator::Lt, ctx),
167            Operator::Lte => execute_compare(&lhs, &rhs, CompareOperator::Lte, ctx),
168            Operator::Gt => execute_compare(&lhs, &rhs, CompareOperator::Gt, ctx),
169            Operator::Gte => execute_compare(&lhs, &rhs, CompareOperator::Gte, ctx),
170            Operator::And => execute_boolean(lhs, rhs, Operator::And, ctx),
171            Operator::Or => execute_boolean(lhs, rhs, Operator::Or, ctx),
172            Operator::Add => execute_numeric(&lhs, &rhs, NumericOperator::Add, ctx),
173            Operator::Sub => execute_numeric(&lhs, &rhs, NumericOperator::Sub, ctx),
174            Operator::Mul => execute_numeric(&lhs, &rhs, NumericOperator::Mul, ctx),
175            Operator::Div => execute_numeric(&lhs, &rhs, NumericOperator::Div, ctx),
176        }
177    }
178
179    fn simplify_untyped(
180        &self,
181        operator: &Operator,
182        expr: &Expression,
183    ) -> VortexResult<Option<Expression>> {
184        let lhs = expr.child(0);
185        let rhs = expr.child(1);
186
187        let bool_literal = |expr: &Expression| {
188            expr.as_opt::<Literal>()?
189                .as_bool_opt()
190                .map(|value| value.value())
191        };
192
193        // AND/OR use Kleene three-valued logic. `None` below is a boolean null.
194        //
195        // AND:
196        // - false AND x => false
197        // - true  AND x => x
198        // - null  AND null => null
199        //
200        // OR:
201        // - true  OR x => true
202        // - false OR x => x
203        // - null  OR null => null
204        //
205        // Other null cases either fall out of the identity/annihilator rules
206        // above (`null AND true`, `null OR false`) or cannot be simplified under
207        // Kleene semantics (`null AND x`, `null OR x` for non-literal `x`).
208        Ok(match operator {
209            Operator::And => match (bool_literal(lhs), bool_literal(rhs)) {
210                (Some(Some(false)), _) | (_, Some(Some(false))) => Some(lit(false)),
211                (Some(Some(true)), _) => Some(rhs.clone()),
212                (_, Some(Some(true))) => Some(lhs.clone()),
213                (Some(None), Some(None)) => Some(lhs.clone()),
214                _ => None,
215            },
216            Operator::Or => match (bool_literal(lhs), bool_literal(rhs)) {
217                (Some(Some(true)), _) | (_, Some(Some(true))) => Some(lit(true)),
218                (Some(Some(false)), _) => Some(rhs.clone()),
219                (_, Some(Some(false))) => Some(lhs.clone()),
220                (Some(None), Some(None)) => Some(lhs.clone()),
221                _ => None,
222            },
223            _ => None,
224        })
225    }
226
227    fn simplify(
228        &self,
229        operator: &Operator,
230        expr: &Expression,
231        ctx: &dyn SimplifyCtx,
232    ) -> VortexResult<Option<Expression>> {
233        let is_literal_null =
234            |expr: &Expression| expr.as_opt::<Literal>().is_some_and(Scalar::is_null);
235
236        if operator.is_comparison()
237            && (is_literal_null(expr.child(0)) || is_literal_null(expr.child(1)))
238        {
239            // Validate the comparison before reducing it. This preserves type
240            // errors for expressions like `int_col = null_utf8`.
241            ctx.return_dtype(expr)?;
242            return Ok(Some(lit(Scalar::null(DType::Bool(Nullability::Nullable)))));
243        }
244
245        Ok(None)
246    }
247
248    fn validity(
249        &self,
250        operator: &Operator,
251        expression: &Expression,
252    ) -> VortexResult<Option<Expression>> {
253        let lhs = expression.child(0).validity()?;
254        let rhs = expression.child(1).validity()?;
255
256        Ok(match operator {
257            // AND and OR are kleene logic.
258            Operator::And => None,
259            Operator::Or => None,
260            _ => {
261                // All other binary operators are null if either side is null.
262                Some(and(lhs, rhs))
263            }
264        })
265    }
266
267    fn is_strict(&self, operator: &Operator) -> bool {
268        // Kleene AND/OR is not strict (`false AND null = false`, `true OR null = true`), which is
269        // consistent with `validity` returning `None` for these operators above.
270        !matches!(operator, Operator::And | Operator::Or)
271    }
272
273    fn is_fallible(&self, operator: &Operator) -> bool {
274        // Opt-in not out for fallibility.
275        // Arithmetic operations could be better modelled here.
276        let infallible = matches!(
277            operator,
278            Operator::Eq
279                | Operator::NotEq
280                | Operator::Gt
281                | Operator::Gte
282                | Operator::Lt
283                | Operator::Lte
284                | Operator::And
285                | Operator::Or
286        );
287
288        !infallible
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use vortex_error::VortexExpect;
295    use vortex_error::VortexResult;
296
297    use super::*;
298    use crate::VortexSessionExecute;
299    use crate::array_session;
300    use crate::assert_arrays_eq;
301    use crate::builtins::ArrayBuiltins;
302    use crate::dtype::DType;
303    use crate::dtype::Nullability;
304    use crate::dtype::PType;
305    use crate::expr::Expression;
306    use crate::expr::and_collect;
307    use crate::expr::col;
308    use crate::expr::eq;
309    use crate::expr::gt;
310    use crate::expr::gt_eq;
311    use crate::expr::lit;
312    use crate::expr::lt;
313    use crate::expr::lt_eq;
314    use crate::expr::not_eq;
315    use crate::expr::or;
316    use crate::expr::or_collect;
317    use crate::expr::test_harness;
318    use crate::scalar::Scalar;
319    #[test]
320    fn and_collect_balanced() {
321        let values = vec![lit(1), lit(2), lit(3), lit(4), lit(5)];
322
323        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @r"
324        vortex.binary(and)
325        ├── lhs: vortex.binary(and)
326        │   ├── lhs: vortex.literal(1i32)
327        │   └── rhs: vortex.literal(2i32)
328        └── rhs: vortex.binary(and)
329            ├── lhs: vortex.binary(and)
330            │   ├── lhs: vortex.literal(3i32)
331            │   └── rhs: vortex.literal(4i32)
332            └── rhs: vortex.literal(5i32)
333        ");
334
335        // 4 elements: and(and(1, 2), and(3, 4)) - perfectly balanced
336        let values = vec![lit(1), lit(2), lit(3), lit(4)];
337        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @r"
338        vortex.binary(and)
339        ├── lhs: vortex.binary(and)
340        │   ├── lhs: vortex.literal(1i32)
341        │   └── rhs: vortex.literal(2i32)
342        └── rhs: vortex.binary(and)
343            ├── lhs: vortex.literal(3i32)
344            └── rhs: vortex.literal(4i32)
345        ");
346
347        // 1 element: just the element
348        let values = vec![lit(1)];
349        insta::assert_snapshot!(and_collect(values.into_iter()).unwrap().display_tree(), @"vortex.literal(1i32)");
350
351        // 0 elements: None
352        let values: Vec<Expression> = vec![];
353        assert!(and_collect(values.into_iter()).is_none());
354    }
355
356    #[test]
357    fn or_collect_balanced() {
358        // 4 elements: or(or(1, 2), or(3, 4)) - perfectly balanced
359        let values = vec![lit(1), lit(2), lit(3), lit(4)];
360        insta::assert_snapshot!(or_collect(values.into_iter()).unwrap().display_tree(), @r"
361        vortex.binary(or)
362        ├── lhs: vortex.binary(or)
363        │   ├── lhs: vortex.literal(1i32)
364        │   └── rhs: vortex.literal(2i32)
365        └── rhs: vortex.binary(or)
366            ├── lhs: vortex.literal(3i32)
367            └── rhs: vortex.literal(4i32)
368        ");
369    }
370
371    #[test]
372    fn dtype() {
373        let dtype = test_harness::struct_dtype();
374        let bool1: Expression = col("bool1");
375        let bool2: Expression = col("bool2");
376        assert_eq!(
377            and(bool1.clone(), bool2.clone())
378                .return_dtype(&dtype)
379                .unwrap(),
380            DType::Bool(Nullability::NonNullable)
381        );
382        assert_eq!(
383            or(bool1, bool2).return_dtype(&dtype).unwrap(),
384            DType::Bool(Nullability::NonNullable)
385        );
386
387        let col1: Expression = col("col1");
388        let col2: Expression = col("col2");
389
390        assert_eq!(
391            eq(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
392            DType::Bool(Nullability::Nullable)
393        );
394        assert_eq!(
395            not_eq(col1.clone(), col2.clone())
396                .return_dtype(&dtype)
397                .unwrap(),
398            DType::Bool(Nullability::Nullable)
399        );
400        assert_eq!(
401            gt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
402            DType::Bool(Nullability::Nullable)
403        );
404        assert_eq!(
405            gt_eq(col1.clone(), col2.clone())
406                .return_dtype(&dtype)
407                .unwrap(),
408            DType::Bool(Nullability::Nullable)
409        );
410        assert_eq!(
411            lt(col1.clone(), col2.clone()).return_dtype(&dtype).unwrap(),
412            DType::Bool(Nullability::Nullable)
413        );
414        assert_eq!(
415            lt_eq(col1.clone(), col2.clone())
416                .return_dtype(&dtype)
417                .unwrap(),
418            DType::Bool(Nullability::Nullable)
419        );
420
421        assert_eq!(
422            or(lt(col1.clone(), col2.clone()), not_eq(col1, col2))
423                .return_dtype(&dtype)
424                .unwrap(),
425            DType::Bool(Nullability::Nullable)
426        );
427    }
428
429    #[test]
430    fn comparison_with_typed_null_simplifies_after_type_check() -> VortexResult<()> {
431        let dtype = test_harness::struct_dtype();
432
433        let expr = eq(
434            col("col1"),
435            lit(Scalar::null(DType::Primitive(
436                PType::U16,
437                Nullability::Nullable,
438            ))),
439        );
440
441        assert_eq!(
442            expr.optimize_recursive(&dtype)?,
443            lit(Scalar::null(DType::Bool(Nullability::Nullable)))
444        );
445        Ok(())
446    }
447
448    #[test]
449    fn comparison_with_incompatible_null_still_type_checks() {
450        let dtype = test_harness::struct_dtype();
451        let expr = eq(
452            col("col1"),
453            lit(Scalar::null(DType::Utf8(Nullability::Nullable))),
454        );
455
456        assert!(expr.optimize_recursive(&dtype).is_err());
457    }
458
459    #[test]
460    fn test_display_print() {
461        let expr = gt(lit(1), lit(2));
462        assert_eq!(format!("{expr}"), "(1i32 > 2i32)");
463    }
464
465    /// Regression test for GitHub issue #5947: struct comparison in filter expressions should work
466    /// using `make_comparator` instead of Arrow's `cmp` functions which don't support nested types.
467    #[test]
468    fn test_struct_comparison() {
469        use crate::IntoArray;
470        use crate::arrays::StructArray;
471
472        // Create a struct array with one element for testing.
473        let lhs_struct = StructArray::from_fields(&[
474            (
475                "a",
476                crate::arrays::PrimitiveArray::from_iter([1i32]).into_array(),
477            ),
478            (
479                "b",
480                crate::arrays::PrimitiveArray::from_iter([3i32]).into_array(),
481            ),
482        ])
483        .unwrap()
484        .into_array();
485
486        let rhs_struct_equal = StructArray::from_fields(&[
487            (
488                "a",
489                crate::arrays::PrimitiveArray::from_iter([1i32]).into_array(),
490            ),
491            (
492                "b",
493                crate::arrays::PrimitiveArray::from_iter([3i32]).into_array(),
494            ),
495        ])
496        .unwrap()
497        .into_array();
498
499        let rhs_struct_different = StructArray::from_fields(&[
500            (
501                "a",
502                crate::arrays::PrimitiveArray::from_iter([1i32]).into_array(),
503            ),
504            (
505                "b",
506                crate::arrays::PrimitiveArray::from_iter([4i32]).into_array(),
507            ),
508        ])
509        .unwrap()
510        .into_array();
511
512        // Test using binary method directly
513        let result_equal = lhs_struct.binary(rhs_struct_equal, Operator::Eq).unwrap();
514        assert_eq!(
515            result_equal
516                .execute_scalar(0, &mut array_session().create_execution_ctx())
517                .vortex_expect("value"),
518            Scalar::bool(true, Nullability::NonNullable),
519            "Equal structs should be equal"
520        );
521
522        let result_different = lhs_struct
523            .binary(rhs_struct_different, Operator::Eq)
524            .unwrap();
525        assert_eq!(
526            result_different
527                .execute_scalar(0, &mut array_session().create_execution_ctx())
528                .vortex_expect("value"),
529            Scalar::bool(false, Nullability::NonNullable),
530            "Different structs should not be equal"
531        );
532    }
533
534    #[test]
535    fn test_or_kleene_validity() {
536        let mut ctx = array_session().create_execution_ctx();
537        use crate::IntoArray;
538        use crate::arrays::BoolArray;
539        use crate::arrays::StructArray;
540        use crate::expr::col;
541
542        let struct_arr = StructArray::from_fields(&[
543            ("a", BoolArray::from_iter([Some(true)]).into_array()),
544            (
545                "b",
546                BoolArray::from_iter([Option::<bool>::None]).into_array(),
547            ),
548        ])
549        .unwrap()
550        .into_array();
551
552        let expr = or(col("a"), col("b"));
553        let result = struct_arr.apply(&expr).unwrap();
554
555        assert_arrays_eq!(
556            result,
557            BoolArray::from_iter([Some(true)]).into_array(),
558            &mut ctx
559        )
560    }
561
562    #[test]
563    fn test_scalar_subtract_unsigned() {
564        let mut ctx = array_session().create_execution_ctx();
565        use vortex_buffer::buffer;
566
567        use crate::IntoArray;
568        use crate::arrays::ConstantArray;
569        use crate::arrays::PrimitiveArray;
570
571        let values = buffer![1u16, 2, 3].into_array();
572        let rhs = ConstantArray::new(Scalar::from(1u16), 3).into_array();
573        let result = values.binary(rhs, Operator::Sub).unwrap();
574        assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx);
575    }
576
577    #[test]
578    fn test_scalar_subtract_signed() {
579        let mut ctx = array_session().create_execution_ctx();
580        use vortex_buffer::buffer;
581
582        use crate::IntoArray;
583        use crate::arrays::ConstantArray;
584        use crate::arrays::PrimitiveArray;
585
586        let values = buffer![1i64, 2, 3].into_array();
587        let rhs = ConstantArray::new(Scalar::from(-1i64), 3).into_array();
588        let result = values.binary(rhs, Operator::Sub).unwrap();
589        assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx);
590    }
591
592    #[test]
593    fn test_scalar_subtract_nullable() {
594        let mut ctx = array_session().create_execution_ctx();
595        use crate::IntoArray;
596        use crate::arrays::ConstantArray;
597        use crate::arrays::PrimitiveArray;
598
599        let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]);
600        let rhs = ConstantArray::new(Scalar::from(Some(1u16)), 4).into_array();
601        let result = values.into_array().binary(rhs, Operator::Sub).unwrap();
602        assert_arrays_eq!(
603            result,
604            PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]),
605            &mut ctx
606        );
607    }
608
609    #[test]
610    fn test_scalar_subtract_float() {
611        let mut ctx = array_session().create_execution_ctx();
612        use vortex_buffer::buffer;
613
614        use crate::IntoArray;
615        use crate::arrays::ConstantArray;
616        use crate::arrays::PrimitiveArray;
617
618        let values = buffer![1.0f64, 2.0, 3.0].into_array();
619        let rhs = ConstantArray::new(Scalar::from(-1f64), 3).into_array();
620        let result = values.binary(rhs, Operator::Sub).unwrap();
621        assert_arrays_eq!(
622            result,
623            PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]),
624            &mut ctx
625        );
626    }
627
628    #[test]
629    fn test_scalar_subtract_float_underflow_is_ok() {
630        use vortex_buffer::buffer;
631
632        use crate::IntoArray;
633        use crate::arrays::ConstantArray;
634
635        let values = buffer![f32::MIN, 2.0, 3.0].into_array();
636        let rhs1 = ConstantArray::new(Scalar::from(1.0f32), 3).into_array();
637        let _results = values.binary(rhs1, Operator::Sub).unwrap();
638        let values = buffer![f32::MIN, 2.0, 3.0].into_array();
639        let rhs2 = ConstantArray::new(Scalar::from(f32::MAX), 3).into_array();
640        let _results = values.binary(rhs2, Operator::Sub).unwrap();
641    }
642}
vortex_array/scalar_fn/fns/binary/mod.rs

vortex_array/scalar_fn/fns/binary/
mod.rs
