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

vortex_array/scalar_fn/fns/binary/
mod.rs
