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

vortex_array/scalar_fn/fns/binary/
mod.rs
