1mod kernels;
19
20use crate::PhysicalExpr;
21use crate::intervals::cp_solver::{propagate_arithmetic, propagate_comparison};
22use std::cmp::Ordering;
23use std::hash::Hash;
24use std::sync::Arc;
25
26use arrow::array::*;
27use arrow::compute::kernels::boolean::{and_kleene, or_kleene};
28use arrow::compute::kernels::concat_elements::concat_elements_dyn;
29use arrow::compute::{SlicesIterator, cast, filter_record_batch};
30use arrow::datatypes::*;
31use arrow::error::ArrowError;
32use datafusion_common::cast::as_boolean_array;
33use datafusion_common::{Result, ScalarValue, internal_err, not_impl_err};
34
35use datafusion_expr::binary::BinaryTypeCoercer;
36use datafusion_expr::interval_arithmetic::{Interval, apply_operator};
37use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
38#[expect(deprecated)]
39use datafusion_expr::statistics::Distribution::{Bernoulli, Gaussian};
40#[expect(deprecated)]
41use datafusion_expr::statistics::{
42 Distribution, combine_bernoullis, combine_gaussians,
43 create_bernoulli_from_comparison, new_generic_from_binary_op,
44};
45use datafusion_expr::{ColumnarValue, Operator};
46use datafusion_physical_expr_common::datum::{apply, apply_cmp};
47
48use kernels::{
49 bitwise_and_dyn, bitwise_and_dyn_scalar, bitwise_or_dyn, bitwise_or_dyn_scalar,
50 bitwise_shift_left_dyn, bitwise_shift_left_dyn_scalar, bitwise_shift_right_dyn,
51 bitwise_shift_right_dyn_scalar, bitwise_xor_dyn, bitwise_xor_dyn_scalar,
52 regex_match_dyn, regex_match_dyn_scalar,
53};
54
55#[derive(Debug, Clone, Eq)]
57pub struct BinaryExpr {
58 left: Arc<dyn PhysicalExpr>,
59 op: Operator,
60 right: Arc<dyn PhysicalExpr>,
61 fail_on_overflow: bool,
63}
64
65impl PartialEq for BinaryExpr {
67 fn eq(&self, other: &Self) -> bool {
68 self.left.eq(&other.left)
69 && self.op.eq(&other.op)
70 && self.right.eq(&other.right)
71 && self.fail_on_overflow.eq(&other.fail_on_overflow)
72 }
73}
74impl Hash for BinaryExpr {
75 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
76 self.left.hash(state);
77 self.op.hash(state);
78 self.right.hash(state);
79 self.fail_on_overflow.hash(state);
80 }
81}
82
83impl BinaryExpr {
84 pub fn new(
86 left: Arc<dyn PhysicalExpr>,
87 op: Operator,
88 right: Arc<dyn PhysicalExpr>,
89 ) -> Self {
90 Self {
91 left,
92 op,
93 right,
94 fail_on_overflow: false,
95 }
96 }
97
98 pub fn with_fail_on_overflow(self, fail_on_overflow: bool) -> Self {
100 Self {
101 left: self.left,
102 op: self.op,
103 right: self.right,
104 fail_on_overflow,
105 }
106 }
107
108 pub fn left(&self) -> &Arc<dyn PhysicalExpr> {
110 &self.left
111 }
112
113 pub fn right(&self) -> &Arc<dyn PhysicalExpr> {
115 &self.right
116 }
117
118 pub fn op(&self) -> &Operator {
120 &self.op
121 }
122
123 fn arithmetic_sort_properties(
129 &self,
130 sort_properties: SortProperties,
131 l_range: &Interval,
132 r_range: &Interval,
133 range: &Interval,
134 ) -> SortProperties {
135 if sort_properties == SortProperties::Singleton {
136 return sort_properties;
137 }
138 let wraps_in_domain = match self.op {
139 Operator::Plus => {
140 is_time_plus_interval(&l_range.data_type(), &r_range.data_type())
141 }
142 Operator::Minus => {
143 is_time_minus_interval(&l_range.data_type(), &r_range.data_type())
144 }
145 _ => false,
146 };
147 let cannot_overflow = !range.is_unbounded()
148 && !unsigned_subtraction_may_underflow(self.op, l_range, r_range, range);
149 if !wraps_in_domain && (self.fail_on_overflow || cannot_overflow) {
150 sort_properties
151 } else {
152 SortProperties::Unordered
153 }
154 }
155}
156
157fn unsigned_subtraction_may_underflow(
166 op: Operator,
167 l_range: &Interval,
168 r_range: &Interval,
169 range: &Interval,
170) -> bool {
171 if op != Operator::Minus || !range.data_type().is_unsigned_integer() {
172 return false;
173 }
174 let (smallest_lhs, largest_rhs) = (l_range.lower(), r_range.upper());
175 if smallest_lhs.is_null() || largest_rhs.is_null() {
176 return true;
177 }
178 !matches!(
181 smallest_lhs.partial_cmp(largest_rhs),
182 Some(Ordering::Greater | Ordering::Equal)
183 )
184}
185
186impl std::fmt::Display for BinaryExpr {
187 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
188 fn write_child(
194 f: &mut std::fmt::Formatter,
195 expr: &dyn PhysicalExpr,
196 precedence: u8,
197 ) -> std::fmt::Result {
198 if let Some(child) = expr.downcast_ref::<BinaryExpr>() {
199 let p = child.op.precedence();
200 if p == 0 || p < precedence {
201 write!(f, "({child})")?;
202 } else {
203 write!(f, "{child}")?;
204 }
205 } else {
206 write!(f, "{expr}")?;
207 }
208
209 Ok(())
210 }
211
212 let precedence = self.op.precedence();
213 write_child(f, self.left.as_ref(), precedence)?;
214 write!(f, " {} ", self.op)?;
215 write_child(f, self.right.as_ref(), precedence)
216 }
217}
218
219#[inline]
221fn boolean_op(
222 left: &dyn Array,
223 right: &dyn Array,
224 op: impl FnOnce(&BooleanArray, &BooleanArray) -> Result<BooleanArray, ArrowError>,
225) -> Result<Arc<dyn Array + 'static>, ArrowError> {
226 let ll = as_boolean_array(left).expect("boolean_op failed to downcast left array");
227 let rr = as_boolean_array(right).expect("boolean_op failed to downcast right array");
228 op(ll, rr).map(|t| Arc::new(t) as _)
229}
230
231fn is_date_minus_date(lhs: &DataType, rhs: &DataType) -> bool {
234 matches!(
235 (lhs, rhs),
236 (DataType::Date32, DataType::Date32) | (DataType::Date64, DataType::Date64)
237 )
238}
239
240const MILLIS_PER_DAY: i64 = 86_400_000;
242
243fn apply_date_subtraction(
249 lhs: &ColumnarValue,
250 rhs: &ColumnarValue,
251) -> Result<ColumnarValue> {
252 match (lhs.data_type(), rhs.data_type()) {
253 (DataType::Date32, DataType::Date32) => {
254 subtract_date_to_days::<Date32Type>(lhs, rhs, |l, r| l - r)
255 }
256 (DataType::Date64, DataType::Date64) => {
257 subtract_date_to_days::<Date64Type>(lhs, rhs, |l, r| {
258 l.wrapping_sub(r) / MILLIS_PER_DAY
259 })
260 }
261 (_, _) => unreachable!("apply_date_subtraction called with non-date types"),
262 }
263}
264
265fn subtract_date_to_days<T: ArrowPrimitiveType>(
269 lhs: &ColumnarValue,
270 rhs: &ColumnarValue,
271 day_diff_fn: impl Fn(i64, i64) -> i64,
272) -> Result<ColumnarValue>
273where
274 T::Native: Copy + Into<i64>,
275{
276 fn date_scalar_to_i64<P: ArrowPrimitiveType>(
278 scalar: &ScalarValue,
279 ) -> Result<Option<i64>> {
280 match scalar {
281 ScalarValue::Date32(value) if P::DATA_TYPE == DataType::Date32 => {
282 Ok(value.map(i64::from))
283 }
284 ScalarValue::Date64(value) if P::DATA_TYPE == DataType::Date64 => Ok(*value),
285 other => {
286 internal_err!(
287 "{} date scalar expected, got: {}",
288 P::DATA_TYPE,
289 other.data_type()
290 )
291 }
292 }
293 }
294
295 match (lhs, rhs) {
296 (ColumnarValue::Array(left), ColumnarValue::Array(right)) => {
297 let left = left.as_primitive::<T>();
298 let right = right.as_primitive::<T>();
299 let result: Int64Array =
300 arrow::compute::binary::<_, _, _, Int64Type>(left, right, |l, r| {
301 day_diff_fn(l.into(), r.into())
302 })?;
303 Ok(ColumnarValue::Array(Arc::new(result)))
304 }
305 (ColumnarValue::Array(left), ColumnarValue::Scalar(right)) => {
306 let left = left.as_primitive::<T>();
307 match date_scalar_to_i64::<T>(right)? {
308 Some(right_val) => {
309 let result: Int64Array =
310 left.unary(|l| day_diff_fn(l.into(), right_val));
311 Ok(ColumnarValue::Array(Arc::new(result)))
312 }
313 None => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))),
314 }
315 }
316 (ColumnarValue::Scalar(left), ColumnarValue::Array(right)) => {
317 let right = right.as_primitive::<T>();
318 match date_scalar_to_i64::<T>(left)? {
319 Some(left_val) => {
320 let result: Int64Array =
321 right.unary(|r| day_diff_fn(left_val, r.into()));
322 Ok(ColumnarValue::Array(Arc::new(result)))
323 }
324 None => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))),
325 }
326 }
327 (ColumnarValue::Scalar(left), ColumnarValue::Scalar(right)) => {
328 let left_val = date_scalar_to_i64::<T>(left)?;
329 let right_val = date_scalar_to_i64::<T>(right)?;
330 Ok(ColumnarValue::Scalar(ScalarValue::Int64(
331 left_val.zip(right_val).map(|(l, r)| day_diff_fn(l, r)),
332 )))
333 }
334 }
335}
336
337fn is_time_plus_interval(lhs: &DataType, rhs: &DataType) -> bool {
339 matches!(
340 (lhs, rhs),
341 (
342 DataType::Time32(_) | DataType::Time64(_),
343 DataType::Interval(_)
344 ) | (
345 DataType::Interval(_),
346 DataType::Time32(_) | DataType::Time64(_)
347 )
348 )
349}
350
351fn is_time_minus_interval(lhs: &DataType, rhs: &DataType) -> bool {
353 matches!(
354 (lhs, rhs),
355 (
356 DataType::Time32(_) | DataType::Time64(_),
357 DataType::Interval(_)
358 )
359 )
360}
361
362fn apply_time_interval(
375 lhs: &ColumnarValue,
376 rhs: &ColumnarValue,
377 subtract: bool,
378) -> Result<ColumnarValue> {
379 let (time, interval) = if matches!(lhs.data_type(), DataType::Interval(_)) {
381 (rhs, lhs)
382 } else {
383 (lhs, rhs)
384 };
385
386 match time.data_type() {
389 DataType::Time32(TimeUnit::Second) => wrap_time_interval::<Time32SecondType>(
390 time,
391 interval,
392 subtract,
393 1_000_000_000,
394 ),
395 DataType::Time32(TimeUnit::Millisecond) => {
396 wrap_time_interval::<Time32MillisecondType>(
397 time, interval, subtract, 1_000_000,
398 )
399 }
400 DataType::Time64(TimeUnit::Microsecond) => {
401 wrap_time_interval::<Time64MicrosecondType>(time, interval, subtract, 1_000)
402 }
403 DataType::Time64(TimeUnit::Nanosecond) => {
404 wrap_time_interval::<Time64NanosecondType>(time, interval, subtract, 1)
405 }
406 other => internal_err!("time operand expected, got: {other}"),
407 }
408}
409
410fn wrap_time_interval<T: ArrowPrimitiveType>(
414 time: &ColumnarValue,
415 interval: &ColumnarValue,
416 subtract: bool,
417 ns_per_unit: i64,
418) -> Result<ColumnarValue>
419where
420 T::Native: Copy + Into<i64> + TryFrom<i64>,
421{
422 const DAY_NANOS: i64 = 86_400_000_000_000;
424 let day_units = DAY_NANOS / ns_per_unit;
426
427 let wrap = |time_unit: i64, iv: IntervalMonthDayNano| -> T::Native {
435 let iv_ns = iv.nanoseconds % DAY_NANOS;
436 let signed_ns = if subtract { -iv_ns } else { iv_ns };
437 let delta = signed_ns.div_euclid(ns_per_unit);
438 let wrapped = (time_unit + delta).rem_euclid(day_units);
439 T::Native::try_from(wrapped).unwrap_or_default()
440 };
441
442 fn interval_scalar(scalar: &ScalarValue) -> Result<Option<IntervalMonthDayNano>> {
444 match scalar {
445 ScalarValue::IntervalMonthDayNano(value) => Ok(*value),
446 other => internal_err!(
447 "Interval(MonthDayNano) scalar expected, got: {}",
448 other.data_type()
449 ),
450 }
451 }
452
453 fn time_scalar_units(scalar: &ScalarValue) -> Result<Option<i64>> {
455 match scalar {
456 ScalarValue::Time32Second(value) | ScalarValue::Time32Millisecond(value) => {
457 Ok(value.map(i64::from))
458 }
459 ScalarValue::Time64Microsecond(value)
460 | ScalarValue::Time64Nanosecond(value) => Ok(*value),
461 other => {
462 internal_err!("time scalar expected, got: {}", other.data_type())
463 }
464 }
465 }
466
467 fn time_scalar<P: ArrowPrimitiveType>(value: Option<i64>) -> ScalarValue {
469 match P::DATA_TYPE {
470 DataType::Time32(TimeUnit::Second) => {
471 ScalarValue::Time32Second(value.map(|v| v as i32))
472 }
473 DataType::Time32(TimeUnit::Millisecond) => {
474 ScalarValue::Time32Millisecond(value.map(|v| v as i32))
475 }
476 DataType::Time64(TimeUnit::Microsecond) => {
477 ScalarValue::Time64Microsecond(value)
478 }
479 _ => ScalarValue::Time64Nanosecond(value),
480 }
481 }
482
483 match (time, interval) {
484 (ColumnarValue::Array(time), ColumnarValue::Array(interval)) => {
485 let time = time.as_primitive::<T>();
486 let interval = interval.as_primitive::<IntervalMonthDayNanoType>();
487 let result: PrimitiveArray<T> =
488 arrow::compute::binary(time, interval, |t, iv| wrap(t.into(), iv))?;
489 Ok(ColumnarValue::Array(Arc::new(result)))
490 }
491 (ColumnarValue::Array(time), ColumnarValue::Scalar(interval)) => {
492 let time = time.as_primitive::<T>();
493 match interval_scalar(interval)? {
494 Some(iv) => {
495 let result: PrimitiveArray<T> = time.unary(|t| wrap(t.into(), iv));
496 Ok(ColumnarValue::Array(Arc::new(result)))
497 }
498 None => Ok(ColumnarValue::Scalar(time_scalar::<T>(None))),
499 }
500 }
501 (ColumnarValue::Scalar(time), ColumnarValue::Array(interval)) => {
502 let interval = interval.as_primitive::<IntervalMonthDayNanoType>();
503 match time_scalar_units(time)? {
504 Some(t) => {
505 let result: PrimitiveArray<T> = interval.unary(|iv| wrap(t, iv));
506 Ok(ColumnarValue::Array(Arc::new(result)))
507 }
508 None => Ok(ColumnarValue::Scalar(time_scalar::<T>(None))),
509 }
510 }
511 (ColumnarValue::Scalar(time), ColumnarValue::Scalar(interval)) => {
512 let result = time_scalar_units(time)?
513 .zip(interval_scalar(interval)?)
514 .map(|(t, iv)| wrap(t, iv).into());
515 Ok(ColumnarValue::Scalar(time_scalar::<T>(result)))
516 }
517 }
518}
519
520impl PhysicalExpr for BinaryExpr {
521 fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
522 BinaryTypeCoercer::new(
523 &self.left.data_type(input_schema)?,
524 &self.op,
525 &self.right.data_type(input_schema)?,
526 )
527 .get_result_type()
528 }
529
530 fn nullable(&self, input_schema: &Schema) -> Result<bool> {
531 Ok(self.left.nullable(input_schema)? || self.right.nullable(input_schema)?)
532 }
533
534 fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
535 use arrow::compute::kernels::numeric::*;
536
537 let lhs = self.left.evaluate(batch)?;
539
540 match check_short_circuit(&lhs, &self.op) {
542 ShortCircuitStrategy::None => {}
543 ShortCircuitStrategy::ReturnLeft => return Ok(lhs),
544 ShortCircuitStrategy::ReturnRight => {
545 let rhs = self.right.evaluate(batch)?;
546 return Ok(rhs);
547 }
548 ShortCircuitStrategy::PreSelection { mask, fill_value } => {
549 let selection_batch = filter_record_batch(batch, &mask)?;
555 let right_ret = self.right.evaluate(&selection_batch)?;
556
557 match &right_ret {
558 ColumnarValue::Array(array) => {
559 let boolean_array = array.as_boolean();
560 if boolean_array.null_count() == 0 {
563 let rhs_value = if !boolean_array.has_false() {
564 Some(true)
565 } else if !boolean_array.has_true() {
566 Some(false)
567 } else {
568 None
569 };
570 if let Some(rhs_value) = rhs_value {
571 return Ok(uniform_pre_selection_result(
572 rhs_value, fill_value, lhs,
573 ));
574 }
575 }
576
577 return pre_selection_scatter(
578 &mask,
579 Some(boolean_array),
580 fill_value,
581 );
582 }
583 ColumnarValue::Scalar(scalar) => {
584 if let ScalarValue::Boolean(v) = scalar {
585 if let Some(v) = v {
587 return Ok(uniform_pre_selection_result(
588 *v, fill_value, lhs,
589 ));
590 } else {
591 return pre_selection_scatter(&mask, None, fill_value);
592 }
593 } else {
594 return internal_err!(
595 "Expected boolean scalar value, found: {right_ret:?}"
596 );
597 }
598 }
599 }
600 }
601 }
602
603 let rhs = self.right.evaluate(batch)?;
604 let left_data_type = lhs.data_type();
605 let right_data_type = rhs.data_type();
606
607 let schema = batch.schema();
608 let input_schema = schema.as_ref();
609
610 match self.op {
611 Operator::Plus
614 if is_time_plus_interval(&left_data_type, &right_data_type) =>
615 {
616 return apply_time_interval(&lhs, &rhs, false);
617 }
618 Operator::Minus
619 if is_time_minus_interval(&left_data_type, &right_data_type) =>
620 {
621 return apply_time_interval(&lhs, &rhs, true);
622 }
623 Operator::Plus if self.fail_on_overflow => return apply(&lhs, &rhs, add),
624 Operator::Plus => return apply(&lhs, &rhs, add_wrapping),
625 Operator::Minus if is_date_minus_date(&left_data_type, &right_data_type) => {
628 return apply_date_subtraction(&lhs, &rhs);
629 }
630 Operator::Minus if self.fail_on_overflow => return apply(&lhs, &rhs, sub),
631 Operator::Minus => return apply(&lhs, &rhs, sub_wrapping),
632 Operator::Multiply if self.fail_on_overflow => return apply(&lhs, &rhs, mul),
633 Operator::Multiply => return apply(&lhs, &rhs, mul_wrapping),
634 Operator::Divide => return apply(&lhs, &rhs, div),
635 Operator::Modulo => return apply(&lhs, &rhs, rem),
636
637 Operator::Eq
638 | Operator::NotEq
639 | Operator::Lt
640 | Operator::Gt
641 | Operator::LtEq
642 | Operator::GtEq
643 | Operator::IsDistinctFrom
644 | Operator::IsNotDistinctFrom
645 | Operator::LikeMatch
646 | Operator::ILikeMatch
647 | Operator::NotLikeMatch
648 | Operator::NotILikeMatch => {
649 return apply_cmp(self.op, &lhs, &rhs);
650 }
651 _ => {}
652 }
653
654 let result_type = self.data_type(input_schema)?;
655
656 if let (ColumnarValue::Array(array), ColumnarValue::Scalar(scalar)) = (&lhs, &rhs)
658 && !scalar.is_null()
659 && let Some(result_array) =
660 self.evaluate_array_scalar(array, scalar.clone())?
661 {
662 let final_array = result_array
663 .and_then(|a| to_result_type_array(&self.op, a, &result_type));
664 return final_array.map(ColumnarValue::Array);
665 }
666
667 let (left, right) = (
669 lhs.into_array(batch.num_rows())?,
670 rhs.into_array(batch.num_rows())?,
671 );
672 self.evaluate_with_resolved_args(left, &left_data_type, right, &right_data_type)
673 .map(ColumnarValue::Array)
674 }
675
676 fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
677 vec![&self.left, &self.right]
678 }
679
680 fn with_new_children(
681 self: Arc<Self>,
682 children: Vec<Arc<dyn PhysicalExpr>>,
683 ) -> Result<Arc<dyn PhysicalExpr>> {
684 Ok(Arc::new(
685 BinaryExpr::new(Arc::clone(&children[0]), self.op, Arc::clone(&children[1]))
686 .with_fail_on_overflow(self.fail_on_overflow),
687 ))
688 }
689
690 fn evaluate_bounds(&self, children: &[&Interval]) -> Result<Interval> {
691 let left_interval = children[0];
693 let right_interval = children[1];
694 apply_operator(&self.op, left_interval, right_interval)
696 }
697
698 fn propagate_constraints(
699 &self,
700 interval: &Interval,
701 children: &[&Interval],
702 ) -> Result<Option<Vec<Interval>>> {
703 let left_interval = children[0];
705 let right_interval = children[1];
706
707 if self.op.eq(&Operator::And) {
708 if interval.eq(&Interval::TRUE) {
709 Ok((!left_interval.eq(&Interval::FALSE)
712 && !right_interval.eq(&Interval::FALSE))
713 .then(|| vec![Interval::TRUE, Interval::TRUE]))
714 } else if interval.eq(&Interval::FALSE) {
715 if left_interval.eq(&Interval::TRUE)
722 && right_interval.eq(&Interval::TRUE_OR_FALSE)
723 {
724 Ok(Some(vec![Interval::TRUE, Interval::FALSE]))
725 } else if right_interval.eq(&Interval::TRUE)
726 && left_interval.eq(&Interval::TRUE_OR_FALSE)
727 {
728 Ok(Some(vec![Interval::FALSE, Interval::TRUE]))
729 }
730 else {
734 Ok(Some(vec![]))
735 }
736 } else {
737 Ok(Some(vec![]))
740 }
741 } else if self.op.eq(&Operator::Or) {
742 if interval.eq(&Interval::FALSE) {
743 Ok((!left_interval.eq(&Interval::TRUE)
746 && !right_interval.eq(&Interval::TRUE))
747 .then(|| vec![Interval::FALSE, Interval::FALSE]))
748 } else if interval.eq(&Interval::TRUE) {
749 if left_interval.eq(&Interval::FALSE)
756 && right_interval.eq(&Interval::TRUE_OR_FALSE)
757 {
758 Ok(Some(vec![Interval::FALSE, Interval::TRUE]))
759 } else if right_interval.eq(&Interval::FALSE)
760 && left_interval.eq(&Interval::TRUE_OR_FALSE)
761 {
762 Ok(Some(vec![Interval::TRUE, Interval::FALSE]))
763 }
764 else {
768 Ok(Some(vec![]))
769 }
770 } else {
771 Ok(Some(vec![]))
774 }
775 } else if self.op.supports_propagation() {
776 Ok(
777 propagate_comparison(&self.op, interval, left_interval, right_interval)?
778 .map(|(left, right)| vec![left, right]),
779 )
780 } else {
781 Ok(
782 propagate_arithmetic(&self.op, interval, left_interval, right_interval)?
783 .map(|(left, right)| vec![left, right]),
784 )
785 }
786 }
787
788 #[expect(deprecated)]
789 fn evaluate_statistics(&self, children: &[&Distribution]) -> Result<Distribution> {
790 let (left, right) = (children[0], children[1]);
791
792 if self.op.is_numerical_operators() {
793 if let (Gaussian(left), Gaussian(right)) = (left, right)
797 && let Some(result) = combine_gaussians(&self.op, left, right)?
798 {
799 return Ok(Gaussian(result));
800 }
801 } else if self.op.is_logic_operator() {
802 return if let (Bernoulli(left), Bernoulli(right)) = (left, right) {
805 combine_bernoullis(&self.op, left, right).map(Bernoulli)
806 } else {
807 internal_err!(
808 "Logical operators are only compatible with `Bernoulli` distributions"
809 )
810 };
811 } else if self.op.supports_propagation() {
812 return create_bernoulli_from_comparison(&self.op, left, right);
815 }
816 new_generic_from_binary_op(&self.op, left, right)
818 }
819
820 fn get_properties(&self, children: &[ExprProperties]) -> Result<ExprProperties> {
823 let (l_order, l_range) = (children[0].sort_properties, &children[0].range);
824 let (r_order, r_range) = (children[1].sort_properties, &children[1].range);
825 match self.op() {
826 Operator::Plus => {
827 let range = l_range.add(r_range)?;
828 Ok(ExprProperties {
829 sort_properties: self.arithmetic_sort_properties(
830 l_order.add(&r_order),
831 l_range,
832 r_range,
833 &range,
834 ),
835 range,
836 preserves_lex_ordering: false,
837 strictly_order_preserving: false,
838 })
839 }
840 Operator::Minus => {
841 let range = l_range.sub(r_range)?;
842 Ok(ExprProperties {
843 sort_properties: self.arithmetic_sort_properties(
844 l_order.sub(&r_order),
845 l_range,
846 r_range,
847 &range,
848 ),
849 range,
850 preserves_lex_ordering: false,
851 strictly_order_preserving: false,
852 })
853 }
854 Operator::Gt => Ok(ExprProperties {
855 sort_properties: l_order.gt_or_gteq(&r_order),
856 range: l_range.gt(r_range)?,
857 preserves_lex_ordering: false,
858 strictly_order_preserving: false,
859 }),
860 Operator::GtEq => Ok(ExprProperties {
861 sort_properties: l_order.gt_or_gteq(&r_order),
862 range: l_range.gt_eq(r_range)?,
863 preserves_lex_ordering: false,
864 strictly_order_preserving: false,
865 }),
866 Operator::Lt => Ok(ExprProperties {
867 sort_properties: r_order.gt_or_gteq(&l_order),
868 range: l_range.lt(r_range)?,
869 preserves_lex_ordering: false,
870 strictly_order_preserving: false,
871 }),
872 Operator::LtEq => Ok(ExprProperties {
873 sort_properties: r_order.gt_or_gteq(&l_order),
874 range: l_range.lt_eq(r_range)?,
875 preserves_lex_ordering: false,
876 strictly_order_preserving: false,
877 }),
878 Operator::And => Ok(ExprProperties {
879 sort_properties: r_order.and_or(&l_order),
880 range: l_range.and(r_range)?,
881 preserves_lex_ordering: false,
882 strictly_order_preserving: false,
883 }),
884 Operator::Or => Ok(ExprProperties {
885 sort_properties: r_order.and_or(&l_order),
886 range: l_range.or(r_range)?,
887 preserves_lex_ordering: false,
888 strictly_order_preserving: false,
889 }),
890 _ => Ok(ExprProperties::new_unknown()),
891 }
892 }
893
894 fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
895 fn write_child(
896 f: &mut std::fmt::Formatter,
897 expr: &dyn PhysicalExpr,
898 precedence: u8,
899 ) -> std::fmt::Result {
900 if let Some(child) = expr.downcast_ref::<BinaryExpr>() {
901 let p = child.op.precedence();
902 if p == 0 || p < precedence {
903 write!(f, "(")?;
904 child.fmt_sql(f)?;
905 write!(f, ")")
906 } else {
907 child.fmt_sql(f)
908 }
909 } else {
910 expr.fmt_sql(f)
911 }
912 }
913
914 let precedence = self.op.precedence();
915 write_child(f, self.left.as_ref(), precedence)?;
916 write!(f, " {} ", self.op)?;
917 write_child(f, self.right.as_ref(), precedence)
918 }
919
920 #[cfg(feature = "proto")]
921 fn try_to_proto(
922 &self,
923 ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
924 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
925 use datafusion_proto_models::protobuf;
926
927 let op = self.op;
930 let mut operand_refs: Vec<&Arc<dyn PhysicalExpr>> = vec![&self.right];
931 let mut current_expr: &BinaryExpr = self;
932 loop {
933 match current_expr.left.downcast_ref::<BinaryExpr>() {
934 Some(bin) if bin.op == op => {
935 operand_refs.push(&bin.right);
936 current_expr = bin;
937 }
938 _ => {
939 operand_refs.push(¤t_expr.left);
940 break;
941 }
942 }
943 }
944 operand_refs.reverse();
946
947 let operands = ctx.encode_children_expressions(operand_refs)?;
948
949 Ok(Some(protobuf::PhysicalExprNode {
950 expr_id: None,
951 expr_type: Some(protobuf::physical_expr_node::ExprType::BinaryExpr(
952 Box::new(protobuf::PhysicalBinaryExprNode {
953 l: None,
954 r: None,
955 op: format!("{op:?}"),
956 operands,
957 }),
958 )),
959 }))
960 }
961}
962
963#[cfg(feature = "proto")]
964impl BinaryExpr {
965 pub fn try_from_proto(
978 node: &datafusion_proto_models::protobuf::PhysicalExprNode,
979 ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
980 ) -> Result<Arc<dyn PhysicalExpr>> {
981 use datafusion_physical_expr_common::expect_expr_variant;
982 use datafusion_proto_models::protobuf;
983 let node = expect_expr_variant!(
984 node,
985 protobuf::physical_expr_node::ExprType::BinaryExpr,
986 "BinaryExpr",
987 );
988 let op = Operator::from_proto_name(&node.op).ok_or_else(|| {
989 datafusion_common::DataFusionError::Internal(format!(
990 "Unsupported binary operator '{}'",
991 node.op
992 ))
993 })?;
994
995 if !node.operands.is_empty() {
996 let operands = ctx.decode_children_expressions(&node.operands)?;
999
1000 if operands.len() < 2 {
1001 return internal_err!(
1002 "A binary expression must always have at least 2 operands"
1003 );
1004 }
1005
1006 Ok(operands
1007 .into_iter()
1008 .reduce(|left, right| {
1009 Arc::new(BinaryExpr::new(left, op, right)) as Arc<dyn PhysicalExpr>
1010 })
1011 .expect("Binary expression could not be reduced to a single expression."))
1012 } else {
1013 let left =
1015 ctx.decode_required_expression(node.l.as_deref(), "BinaryExpr", "left")?;
1016 let right =
1017 ctx.decode_required_expression(node.r.as_deref(), "BinaryExpr", "right")?;
1018 Ok(Arc::new(BinaryExpr::new(left, op, right)))
1019 }
1020 }
1021}
1022
1023fn to_result_type_array(
1029 op: &Operator,
1030 array: ArrayRef,
1031 result_type: &DataType,
1032) -> Result<ArrayRef> {
1033 if array.data_type() == result_type {
1034 Ok(array)
1035 } else if op.is_numerical_operators() {
1036 match array.data_type() {
1037 DataType::Dictionary(_, value_type) => {
1038 if value_type.as_ref() == result_type {
1039 Ok(cast(&array, result_type)?)
1040 } else {
1041 internal_err!(
1042 "Incompatible Dictionary value type {value_type} with result type {result_type} of Binary operator {op:?}"
1043 )
1044 }
1045 }
1046 _ => Ok(array),
1047 }
1048 } else {
1049 Ok(array)
1050 }
1051}
1052
1053impl BinaryExpr {
1054 fn evaluate_array_scalar(
1057 &self,
1058 array: &dyn Array,
1059 scalar: ScalarValue,
1060 ) -> Result<Option<Result<ArrayRef>>> {
1061 use Operator::*;
1062 let scalar_result = match &self.op {
1063 RegexMatch => regex_match_dyn_scalar(array, &scalar, false, false),
1064 RegexIMatch => regex_match_dyn_scalar(array, &scalar, false, true),
1065 RegexNotMatch => regex_match_dyn_scalar(array, &scalar, true, false),
1066 RegexNotIMatch => regex_match_dyn_scalar(array, &scalar, true, true),
1067 BitwiseAnd => bitwise_and_dyn_scalar(array, scalar),
1068 BitwiseOr => bitwise_or_dyn_scalar(array, scalar),
1069 BitwiseXor => bitwise_xor_dyn_scalar(array, scalar),
1070 BitwiseShiftRight => bitwise_shift_right_dyn_scalar(array, scalar),
1071 BitwiseShiftLeft => bitwise_shift_left_dyn_scalar(array, scalar),
1072 _ => None,
1074 };
1075
1076 Ok(scalar_result)
1077 }
1078
1079 fn evaluate_with_resolved_args(
1080 &self,
1081 left: Arc<dyn Array>,
1082 left_data_type: &DataType,
1083 right: Arc<dyn Array>,
1084 right_data_type: &DataType,
1085 ) -> Result<ArrayRef> {
1086 use Operator::*;
1087 match &self.op {
1088 IsDistinctFrom | IsNotDistinctFrom | Lt | LtEq | Gt | GtEq | Eq | NotEq
1089 | Plus | Minus | Multiply | Divide | Modulo | LikeMatch | ILikeMatch
1090 | NotLikeMatch | NotILikeMatch => unreachable!(),
1091 And => {
1092 if left_data_type == &DataType::Boolean {
1093 Ok(boolean_op(&left, &right, and_kleene)?)
1094 } else {
1095 internal_err!(
1096 "Cannot evaluate binary expression {:?} with types {:?} and {:?}",
1097 self.op,
1098 left.data_type(),
1099 right.data_type()
1100 )
1101 }
1102 }
1103 Or => {
1104 if left_data_type == &DataType::Boolean {
1105 Ok(boolean_op(&left, &right, or_kleene)?)
1106 } else {
1107 internal_err!(
1108 "Cannot evaluate binary expression {:?} with types {:?} and {:?}",
1109 self.op,
1110 left_data_type,
1111 right_data_type
1112 )
1113 }
1114 }
1115 RegexMatch => regex_match_dyn(&left, &right, false, false),
1116 RegexIMatch => regex_match_dyn(&left, &right, false, true),
1117 RegexNotMatch => regex_match_dyn(&left, &right, true, false),
1118 RegexNotIMatch => regex_match_dyn(&left, &right, true, true),
1119 BitwiseAnd => bitwise_and_dyn(left, right),
1120 BitwiseOr => bitwise_or_dyn(left, right),
1121 BitwiseXor => bitwise_xor_dyn(left, right),
1122 BitwiseShiftRight => bitwise_shift_right_dyn(left, right),
1123 BitwiseShiftLeft => bitwise_shift_left_dyn(left, right),
1124 StringConcat => concat_elements_dyn(&left, &right).map_err(|e| e.into()),
1125 AtArrow | ArrowAt | Arrow | LongArrow | HashArrow | HashLongArrow | AtAt
1126 | HashMinus | AtQuestion | Question | QuestionAnd | QuestionPipe
1127 | IntegerDivide | Colon => {
1128 not_impl_err!(
1129 "Binary operator '{:?}' is not supported in the physical expr",
1130 self.op
1131 )
1132 }
1133 }
1134 }
1135}
1136
1137enum ShortCircuitStrategy {
1138 None,
1139 ReturnLeft,
1140 ReturnRight,
1141 PreSelection {
1149 mask: BooleanArray,
1150 fill_value: bool,
1151 },
1152}
1153
1154const PRE_SELECTION_THRESHOLD: f32 = 0.2;
1160
1161fn check_short_circuit(lhs: &ColumnarValue, op: &Operator) -> ShortCircuitStrategy {
1181 let is_and = match op {
1183 Operator::And => true,
1184 Operator::Or => false,
1185 _ => return ShortCircuitStrategy::None,
1186 };
1187
1188 if lhs.data_type() != DataType::Boolean {
1190 return ShortCircuitStrategy::None;
1191 }
1192
1193 match lhs {
1194 ColumnarValue::Array(array) => {
1195 if let Ok(bool_array) = as_boolean_array(array) {
1197 if bool_array.null_count() > 0 {
1199 return ShortCircuitStrategy::None;
1200 }
1201
1202 let len = bool_array.len();
1203 if len == 0 {
1204 return ShortCircuitStrategy::None;
1205 }
1206
1207 let true_count = bool_array.values().count_set_bits();
1208 if is_and {
1209 if true_count == 0 {
1210 return ShortCircuitStrategy::ReturnLeft;
1211 }
1212
1213 if true_count == len {
1214 return ShortCircuitStrategy::ReturnRight;
1215 }
1216
1217 if true_count as f32 / len as f32 <= PRE_SELECTION_THRESHOLD {
1218 return ShortCircuitStrategy::PreSelection {
1221 mask: bool_array.clone(),
1222 fill_value: false,
1223 };
1224 }
1225 } else {
1226 if true_count == len {
1227 return ShortCircuitStrategy::ReturnLeft;
1228 }
1229
1230 if true_count == 0 {
1231 return ShortCircuitStrategy::ReturnRight;
1232 }
1233
1234 let false_count = len - true_count;
1235 if false_count as f32 / len as f32 <= PRE_SELECTION_THRESHOLD {
1236 let mask = BooleanArray::new(!bool_array.values(), None);
1240 return ShortCircuitStrategy::PreSelection {
1241 mask,
1242 fill_value: true,
1243 };
1244 }
1245 }
1246 }
1247 }
1248 ColumnarValue::Scalar(scalar) => {
1249 if let ScalarValue::Boolean(Some(is_true)) = scalar {
1251 if (is_and && !is_true) || (!is_and && *is_true) {
1255 return ShortCircuitStrategy::ReturnLeft;
1256 } else {
1257 return ShortCircuitStrategy::ReturnRight;
1258 }
1259 }
1260 }
1261 }
1262
1263 ShortCircuitStrategy::None
1265}
1266
1267fn uniform_pre_selection_result(
1273 rhs_value: bool,
1274 fill_value: bool,
1275 lhs: ColumnarValue,
1276) -> ColumnarValue {
1277 if rhs_value == fill_value {
1278 ColumnarValue::Scalar(ScalarValue::Boolean(Some(fill_value)))
1279 } else {
1280 lhs
1281 }
1282}
1283
1284fn pre_selection_scatter(
1300 mask: &BooleanArray,
1301 right_result: Option<&BooleanArray>,
1302 fill_value: bool,
1303) -> Result<ColumnarValue> {
1304 let result_len = mask.len();
1305
1306 let mut result_array_builder = BooleanArray::builder(result_len);
1307
1308 let mut right_array_pos = 0;
1309 let mut last_end = 0;
1310 match right_result {
1311 Some(right_result) => {
1312 SlicesIterator::new(mask).for_each(|(start, end)| {
1313 if start > last_end {
1314 result_array_builder.append_n(start - last_end, fill_value);
1315 }
1316
1317 let len = end - start;
1319 right_result
1320 .slice(right_array_pos, len)
1321 .iter()
1322 .for_each(|v| result_array_builder.append_option(v));
1323
1324 right_array_pos += len;
1325 last_end = end;
1326 });
1327 }
1328 None => SlicesIterator::new(mask).for_each(|(start, end)| {
1329 if start > last_end {
1330 result_array_builder.append_n(start - last_end, fill_value);
1331 }
1332
1333 let len = end - start;
1334 result_array_builder.append_nulls(len);
1335
1336 last_end = end;
1337 }),
1338 }
1339
1340 if last_end < result_len {
1342 result_array_builder.append_n(result_len - last_end, fill_value);
1343 }
1344 let boolean_result = result_array_builder.finish();
1345
1346 Ok(ColumnarValue::Array(Arc::new(boolean_result)))
1347}
1348
1349pub fn binary(
1353 lhs: Arc<dyn PhysicalExpr>,
1354 op: Operator,
1355 rhs: Arc<dyn PhysicalExpr>,
1356 _input_schema: &Schema,
1357) -> Result<Arc<dyn PhysicalExpr>> {
1358 Ok(Arc::new(BinaryExpr::new(lhs, op, rhs)))
1359}
1360
1361pub fn similar_to(
1363 negated: bool,
1364 case_insensitive: bool,
1365 expr: Arc<dyn PhysicalExpr>,
1366 pattern: Arc<dyn PhysicalExpr>,
1367) -> Result<Arc<dyn PhysicalExpr>> {
1368 let binary_op = match (negated, case_insensitive) {
1369 (false, false) => Operator::RegexMatch,
1370 (false, true) => Operator::RegexIMatch,
1371 (true, false) => Operator::RegexNotMatch,
1372 (true, true) => Operator::RegexNotIMatch,
1373 };
1374 Ok(Arc::new(BinaryExpr::new(expr, binary_op, pattern)))
1375}
1376
1377#[cfg(test)]
1378mod tests {
1379 use super::*;
1380 use crate::expressions::{Column, Literal, col, lit, try_cast};
1381 use datafusion_expr::lit as expr_lit;
1382
1383 use datafusion_common::{assert_contains, plan_datafusion_err};
1384 use datafusion_physical_expr_common::physical_expr::fmt_sql;
1385
1386 use crate::planner::logical2physical;
1387 use arrow::array::BooleanArray;
1388 use arrow::compute::SortOptions;
1389 use datafusion_expr::col as logical_col;
1390
1391 #[test]
1392 fn test_arithmetic_ordering_overflow() -> Result<()> {
1393 let asc = SortProperties::Ordered(Default::default());
1394 let ordered = |range: Interval| ExprProperties {
1395 sort_properties: asc,
1396 range,
1397 preserves_lex_ordering: false,
1398 strictly_order_preserving: false,
1399 };
1400
1401 let schema = Schema::new(vec![
1402 Field::new("a", DataType::Int32, false),
1403 Field::new("b", DataType::Int32, false),
1404 ]);
1405 let a_plus_b =
1406 BinaryExpr::new(col("a", &schema)?, Operator::Plus, col("b", &schema)?);
1407 let unbounded = [
1408 ordered(Interval::make_unbounded(&DataType::Int32)?),
1409 ordered(Interval::make_unbounded(&DataType::Int32)?),
1410 ];
1411 let bounded = [
1412 ordered(Interval::make(Some(0), Some(10))?),
1413 ordered(Interval::make(Some(0), Some(10))?),
1414 ];
1415
1416 assert_eq!(
1418 a_plus_b.get_properties(&unbounded)?.sort_properties,
1419 SortProperties::Unordered
1420 );
1421 assert_eq!(a_plus_b.get_properties(&bounded)?.sort_properties, asc);
1424 let checked = a_plus_b.with_fail_on_overflow(true);
1425 assert_eq!(checked.get_properties(&unbounded)?.sort_properties, asc);
1426
1427 let time = DataType::Time64(TimeUnit::Nanosecond);
1430 let interval = DataType::Interval(IntervalUnit::MonthDayNano);
1431 let schema = Schema::new(vec![
1432 Field::new("t", time.clone(), false),
1433 Field::new("i", interval.clone(), false),
1434 ]);
1435 let time_plus_interval =
1436 BinaryExpr::new(col("t", &schema)?, Operator::Plus, col("i", &schema)?)
1437 .with_fail_on_overflow(true);
1438 let time_props = [
1439 ordered(Interval::make_unbounded(&time)?),
1440 ordered(Interval::make_unbounded(&interval)?),
1441 ];
1442 assert_eq!(
1443 time_plus_interval
1444 .get_properties(&time_props)?
1445 .sort_properties,
1446 SortProperties::Unordered
1447 );
1448
1449 Ok(())
1450 }
1451
1452 #[test]
1456 fn test_subtraction_ordering_overflow() -> Result<()> {
1457 let asc = SortProperties::Ordered(SortOptions {
1458 descending: false,
1459 nulls_first: true,
1460 });
1461 let desc = SortProperties::Ordered(SortOptions {
1462 descending: true,
1463 nulls_first: true,
1464 });
1465 let props = |sort_properties, range| ExprProperties {
1466 sort_properties,
1467 range,
1468 preserves_lex_ordering: false,
1469 strictly_order_preserving: false,
1470 };
1471
1472 let schema = Schema::new(vec![
1473 Field::new("a", DataType::Int32, false),
1474 Field::new("b", DataType::Int32, false),
1475 ]);
1476 let a_minus_b =
1477 BinaryExpr::new(col("a", &schema)?, Operator::Minus, col("b", &schema)?);
1478
1479 let signed_underflow = [
1482 props(asc, Interval::make(Some(i32::MIN), Some(0))?),
1483 props(desc, Interval::make(Some(0), Some(i32::MAX))?),
1484 ];
1485 assert_eq!(
1486 a_minus_b.get_properties(&signed_underflow)?.sort_properties,
1487 SortProperties::Unordered
1488 );
1489 let checked = a_minus_b.clone().with_fail_on_overflow(true);
1492 assert_eq!(
1493 checked.get_properties(&signed_underflow)?.sort_properties,
1494 asc
1495 );
1496 let signed_safe = [
1498 props(asc, Interval::make(Some(0), Some(10))?),
1499 props(desc, Interval::make(Some(0), Some(10))?),
1500 ];
1501 assert_eq!(a_minus_b.get_properties(&signed_safe)?.sort_properties, asc);
1502
1503 let schema = Schema::new(vec![
1504 Field::new("a", DataType::UInt32, false),
1505 Field::new("b", DataType::UInt32, false),
1506 ]);
1507 let a_minus_b =
1508 BinaryExpr::new(col("a", &schema)?, Operator::Minus, col("b", &schema)?);
1509
1510 let unsigned_underflow = [
1513 props(asc, Interval::make(Some(0_u32), Some(10_u32))?),
1514 props(desc, Interval::make(Some(0_u32), Some(10_u32))?),
1515 ];
1516 assert_eq!(
1517 a_minus_b
1518 .get_properties(&unsigned_underflow)?
1519 .sort_properties,
1520 SortProperties::Unordered
1521 );
1522 let unsigned_safe = [
1525 props(asc, Interval::make(Some(10_u32), Some(20_u32))?),
1526 props(desc, Interval::make(Some(0_u32), Some(5_u32))?),
1527 ];
1528 assert_eq!(
1529 a_minus_b.get_properties(&unsigned_safe)?.sort_properties,
1530 asc
1531 );
1532
1533 let time = DataType::Time64(TimeUnit::Nanosecond);
1536 let interval = DataType::Interval(IntervalUnit::MonthDayNano);
1537 let schema = Schema::new(vec![
1538 Field::new("t", time.clone(), false),
1539 Field::new("i", interval.clone(), false),
1540 ]);
1541 let time_minus_interval =
1542 BinaryExpr::new(col("t", &schema)?, Operator::Minus, col("i", &schema)?)
1543 .with_fail_on_overflow(true);
1544 let time_props = [
1545 props(asc, Interval::make_unbounded(&time)?),
1546 props(desc, Interval::make_unbounded(&interval)?),
1547 ];
1548 assert_eq!(
1549 time_minus_interval
1550 .get_properties(&time_props)?
1551 .sort_properties,
1552 SortProperties::Unordered
1553 );
1554
1555 Ok(())
1556 }
1557
1558 fn binary_op(
1560 left: Arc<dyn PhysicalExpr>,
1561 op: Operator,
1562 right: Arc<dyn PhysicalExpr>,
1563 schema: &Schema,
1564 ) -> Result<Arc<dyn PhysicalExpr>> {
1565 let left_type = left.data_type(schema)?;
1566 let right_type = right.data_type(schema)?;
1567 let (lhs, rhs) =
1568 BinaryTypeCoercer::new(&left_type, &op, &right_type).get_input_types()?;
1569
1570 let left_expr = try_cast(left, schema, lhs)?;
1571 let right_expr = try_cast(right, schema, rhs)?;
1572 binary(left_expr, op, right_expr, schema)
1573 }
1574
1575 #[test]
1576 fn binary_comparison() -> Result<()> {
1577 let schema = Schema::new(vec![
1578 Field::new("a", DataType::Int32, false),
1579 Field::new("b", DataType::Int32, false),
1580 ]);
1581 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
1582 let b = Int32Array::from(vec![1, 2, 4, 8, 16]);
1583
1584 let lt = binary(
1586 col("a", &schema)?,
1587 Operator::Lt,
1588 col("b", &schema)?,
1589 &schema,
1590 )?;
1591 let batch =
1592 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)])?;
1593
1594 let result = lt
1595 .evaluate(&batch)?
1596 .into_array(batch.num_rows())
1597 .expect("Failed to convert to array");
1598 assert_eq!(result.len(), 5);
1599
1600 let expected = [false, false, true, true, true];
1601 let result =
1602 as_boolean_array(&result).expect("failed to downcast to BooleanArray");
1603 for (i, &expected_item) in expected.iter().enumerate().take(5) {
1604 assert_eq!(result.value(i), expected_item);
1605 }
1606
1607 Ok(())
1608 }
1609
1610 #[test]
1611 fn binary_nested() -> Result<()> {
1612 let schema = Schema::new(vec![
1613 Field::new("a", DataType::Int32, false),
1614 Field::new("b", DataType::Int32, false),
1615 ]);
1616 let a = Int32Array::from(vec![2, 4, 6, 8, 10]);
1617 let b = Int32Array::from(vec![2, 5, 4, 8, 8]);
1618
1619 let expr = binary(
1621 binary(
1622 col("a", &schema)?,
1623 Operator::Lt,
1624 col("b", &schema)?,
1625 &schema,
1626 )?,
1627 Operator::Or,
1628 binary(
1629 col("a", &schema)?,
1630 Operator::Eq,
1631 col("b", &schema)?,
1632 &schema,
1633 )?,
1634 &schema,
1635 )?;
1636 let batch =
1637 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)])?;
1638
1639 assert_eq!("a@0 < b@1 OR a@0 = b@1", format!("{expr}"));
1640
1641 let result = expr
1642 .evaluate(&batch)?
1643 .into_array(batch.num_rows())
1644 .expect("Failed to convert to array");
1645 assert_eq!(result.len(), 5);
1646
1647 let expected = [true, true, false, true, false];
1648 let result =
1649 as_boolean_array(&result).expect("failed to downcast to BooleanArray");
1650 for (i, &expected_item) in expected.iter().enumerate().take(5) {
1651 assert_eq!(result.value(i), expected_item);
1652 }
1653
1654 Ok(())
1655 }
1656
1657 macro_rules! test_coercion {
1665 ($A_ARRAY:ident, $A_TYPE:expr, $A_VEC:expr, $B_ARRAY:ident, $B_TYPE:expr, $B_VEC:expr, $OP:expr, $C_ARRAY:ident, $C_TYPE:expr, $VEC:expr,) => {{
1666 let schema = Schema::new(vec![
1667 Field::new("a", $A_TYPE, false),
1668 Field::new("b", $B_TYPE, false),
1669 ]);
1670 let a = $A_ARRAY::from($A_VEC);
1671 let b = $B_ARRAY::from($B_VEC);
1672 let (lhs, rhs) =
1673 BinaryTypeCoercer::new(&$A_TYPE, &$OP, &$B_TYPE).get_input_types()?;
1674
1675 let left = try_cast(col("a", &schema)?, &schema, lhs)?;
1676 let right = try_cast(col("b", &schema)?, &schema, rhs)?;
1677
1678 let expression = binary(left, $OP, right, &schema)?;
1680 let batch = RecordBatch::try_new(
1681 Arc::new(schema.clone()),
1682 vec![Arc::new(a), Arc::new(b)],
1683 )?;
1684
1685 assert_eq!(expression.data_type(&schema)?, $C_TYPE);
1687
1688 let result = expression
1690 .evaluate(&batch)?
1691 .into_array(batch.num_rows())
1692 .expect("Failed to convert to array");
1693
1694 assert_eq!(*result.data_type(), $C_TYPE);
1696
1697 let result = result
1699 .as_any()
1700 .downcast_ref::<$C_ARRAY>()
1701 .expect("failed to downcast");
1702 for (i, x) in $VEC.iter().enumerate() {
1704 let v = result.value(i);
1705 assert_eq!(
1706 v, *x,
1707 "Unexpected output at position {i}:\n\nActual:\n{v}\n\nExpected:\n{x}"
1708 );
1709 }
1710 }};
1711 }
1712
1713 #[test]
1714 fn test_type_coercion() -> Result<()> {
1715 test_coercion!(
1716 Int32Array,
1717 DataType::Int32,
1718 vec![1i32, 2i32],
1719 UInt32Array,
1720 DataType::UInt32,
1721 vec![1u32, 2u32],
1722 Operator::Plus,
1723 Int64Array,
1724 DataType::Int64,
1725 [2i64, 4i64],
1726 );
1727 test_coercion!(
1728 Int32Array,
1729 DataType::Int32,
1730 vec![1i32],
1731 UInt16Array,
1732 DataType::UInt16,
1733 vec![1u16],
1734 Operator::Plus,
1735 Int32Array,
1736 DataType::Int32,
1737 [2i32],
1738 );
1739 test_coercion!(
1740 Float32Array,
1741 DataType::Float32,
1742 vec![1f32],
1743 UInt16Array,
1744 DataType::UInt16,
1745 vec![1u16],
1746 Operator::Plus,
1747 Float32Array,
1748 DataType::Float32,
1749 [2f32],
1750 );
1751 test_coercion!(
1752 Float32Array,
1753 DataType::Float32,
1754 vec![2f32],
1755 UInt16Array,
1756 DataType::UInt16,
1757 vec![1u16],
1758 Operator::Multiply,
1759 Float32Array,
1760 DataType::Float32,
1761 [2f32],
1762 );
1763 test_coercion!(
1764 StringArray,
1765 DataType::Utf8,
1766 vec!["1994-12-13", "1995-01-26"],
1767 Date32Array,
1768 DataType::Date32,
1769 vec![9112, 9156],
1770 Operator::Eq,
1771 BooleanArray,
1772 DataType::Boolean,
1773 [true, true],
1774 );
1775 test_coercion!(
1776 StringArray,
1777 DataType::Utf8,
1778 vec!["1994-12-13", "1995-01-26"],
1779 Date32Array,
1780 DataType::Date32,
1781 vec![9113, 9154],
1782 Operator::Lt,
1783 BooleanArray,
1784 DataType::Boolean,
1785 [true, false],
1786 );
1787 test_coercion!(
1788 StringArray,
1789 DataType::Utf8,
1790 vec!["1994-12-13T12:34:56", "1995-01-26T01:23:45"],
1791 Date64Array,
1792 DataType::Date64,
1793 vec![787322096000, 791083425000],
1794 Operator::Eq,
1795 BooleanArray,
1796 DataType::Boolean,
1797 [true, true],
1798 );
1799 test_coercion!(
1800 StringArray,
1801 DataType::Utf8,
1802 vec!["1994-12-13T12:34:56", "1995-01-26T01:23:45"],
1803 Date64Array,
1804 DataType::Date64,
1805 vec![787322096001, 791083424999],
1806 Operator::Lt,
1807 BooleanArray,
1808 DataType::Boolean,
1809 [true, false],
1810 );
1811 test_coercion!(
1812 StringViewArray,
1813 DataType::Utf8View,
1814 vec!["abc"; 5],
1815 StringArray,
1816 DataType::Utf8,
1817 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1818 Operator::RegexMatch,
1819 BooleanArray,
1820 DataType::Boolean,
1821 [true, false, true, false, false],
1822 );
1823 test_coercion!(
1824 StringViewArray,
1825 DataType::Utf8View,
1826 vec!["abc"; 5],
1827 StringArray,
1828 DataType::Utf8,
1829 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1830 Operator::RegexIMatch,
1831 BooleanArray,
1832 DataType::Boolean,
1833 [true, true, true, true, false],
1834 );
1835 test_coercion!(
1836 StringArray,
1837 DataType::Utf8,
1838 vec!["abc"; 5],
1839 StringViewArray,
1840 DataType::Utf8View,
1841 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1842 Operator::RegexNotMatch,
1843 BooleanArray,
1844 DataType::Boolean,
1845 [false, true, false, true, true],
1846 );
1847 test_coercion!(
1848 StringArray,
1849 DataType::Utf8,
1850 vec!["abc"; 5],
1851 StringViewArray,
1852 DataType::Utf8View,
1853 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1854 Operator::RegexNotIMatch,
1855 BooleanArray,
1856 DataType::Boolean,
1857 [false, false, false, false, true],
1858 );
1859 test_coercion!(
1860 StringArray,
1861 DataType::Utf8,
1862 vec!["abc"; 5],
1863 StringArray,
1864 DataType::Utf8,
1865 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1866 Operator::RegexMatch,
1867 BooleanArray,
1868 DataType::Boolean,
1869 [true, false, true, false, false],
1870 );
1871 test_coercion!(
1872 StringArray,
1873 DataType::Utf8,
1874 vec!["abc"; 5],
1875 StringArray,
1876 DataType::Utf8,
1877 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1878 Operator::RegexIMatch,
1879 BooleanArray,
1880 DataType::Boolean,
1881 [true, true, true, true, false],
1882 );
1883 test_coercion!(
1884 StringArray,
1885 DataType::Utf8,
1886 vec!["abc"; 5],
1887 StringArray,
1888 DataType::Utf8,
1889 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1890 Operator::RegexNotMatch,
1891 BooleanArray,
1892 DataType::Boolean,
1893 [false, true, false, true, true],
1894 );
1895 test_coercion!(
1896 StringArray,
1897 DataType::Utf8,
1898 vec!["abc"; 5],
1899 StringArray,
1900 DataType::Utf8,
1901 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1902 Operator::RegexNotIMatch,
1903 BooleanArray,
1904 DataType::Boolean,
1905 [false, false, false, false, true],
1906 );
1907 test_coercion!(
1908 LargeStringArray,
1909 DataType::LargeUtf8,
1910 vec!["abc"; 5],
1911 LargeStringArray,
1912 DataType::LargeUtf8,
1913 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1914 Operator::RegexMatch,
1915 BooleanArray,
1916 DataType::Boolean,
1917 [true, false, true, false, false],
1918 );
1919 test_coercion!(
1920 LargeStringArray,
1921 DataType::LargeUtf8,
1922 vec!["abc"; 5],
1923 LargeStringArray,
1924 DataType::LargeUtf8,
1925 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1926 Operator::RegexIMatch,
1927 BooleanArray,
1928 DataType::Boolean,
1929 [true, true, true, true, false],
1930 );
1931 test_coercion!(
1932 LargeStringArray,
1933 DataType::LargeUtf8,
1934 vec!["abc"; 5],
1935 LargeStringArray,
1936 DataType::LargeUtf8,
1937 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1938 Operator::RegexNotMatch,
1939 BooleanArray,
1940 DataType::Boolean,
1941 [false, true, false, true, true],
1942 );
1943 test_coercion!(
1944 LargeStringArray,
1945 DataType::LargeUtf8,
1946 vec!["abc"; 5],
1947 LargeStringArray,
1948 DataType::LargeUtf8,
1949 vec!["^a", "^A", "(b|d)", "(B|D)", "^(b|c)"],
1950 Operator::RegexNotIMatch,
1951 BooleanArray,
1952 DataType::Boolean,
1953 [false, false, false, false, true],
1954 );
1955 test_coercion!(
1956 StringArray,
1957 DataType::Utf8,
1958 vec!["abc"; 5],
1959 StringArray,
1960 DataType::Utf8,
1961 vec!["a__", "A%BC", "A_BC", "abc", "a%C"],
1962 Operator::LikeMatch,
1963 BooleanArray,
1964 DataType::Boolean,
1965 [true, false, false, true, false],
1966 );
1967 test_coercion!(
1968 StringArray,
1969 DataType::Utf8,
1970 vec!["abc"; 5],
1971 StringArray,
1972 DataType::Utf8,
1973 vec!["a__", "A%BC", "A_BC", "abc", "a%C"],
1974 Operator::ILikeMatch,
1975 BooleanArray,
1976 DataType::Boolean,
1977 [true, true, false, true, true],
1978 );
1979 test_coercion!(
1980 StringArray,
1981 DataType::Utf8,
1982 vec!["abc"; 5],
1983 StringArray,
1984 DataType::Utf8,
1985 vec!["a__", "A%BC", "A_BC", "abc", "a%C"],
1986 Operator::NotLikeMatch,
1987 BooleanArray,
1988 DataType::Boolean,
1989 [false, true, true, false, true],
1990 );
1991 test_coercion!(
1992 StringArray,
1993 DataType::Utf8,
1994 vec!["abc"; 5],
1995 StringArray,
1996 DataType::Utf8,
1997 vec!["a__", "A%BC", "A_BC", "abc", "a%C"],
1998 Operator::NotILikeMatch,
1999 BooleanArray,
2000 DataType::Boolean,
2001 [false, false, true, false, false],
2002 );
2003 test_coercion!(
2004 LargeStringArray,
2005 DataType::LargeUtf8,
2006 vec!["abc"; 5],
2007 LargeStringArray,
2008 DataType::LargeUtf8,
2009 vec!["a__", "A%BC", "A_BC", "abc", "a%C"],
2010 Operator::LikeMatch,
2011 BooleanArray,
2012 DataType::Boolean,
2013 [true, false, false, true, false],
2014 );
2015 test_coercion!(
2016 LargeStringArray,
2017 DataType::LargeUtf8,
2018 vec!["abc"; 5],
2019 LargeStringArray,
2020 DataType::LargeUtf8,
2021 vec!["a__", "A%BC", "A_BC", "abc", "a%C"],
2022 Operator::ILikeMatch,
2023 BooleanArray,
2024 DataType::Boolean,
2025 [true, true, false, true, true],
2026 );
2027 test_coercion!(
2028 LargeStringArray,
2029 DataType::LargeUtf8,
2030 vec!["abc"; 5],
2031 LargeStringArray,
2032 DataType::LargeUtf8,
2033 vec!["a__", "A%BC", "A_BC", "abc", "a%C"],
2034 Operator::NotLikeMatch,
2035 BooleanArray,
2036 DataType::Boolean,
2037 [false, true, true, false, true],
2038 );
2039 test_coercion!(
2040 LargeStringArray,
2041 DataType::LargeUtf8,
2042 vec!["abc"; 5],
2043 LargeStringArray,
2044 DataType::LargeUtf8,
2045 vec!["a__", "A%BC", "A_BC", "abc", "a%C"],
2046 Operator::NotILikeMatch,
2047 BooleanArray,
2048 DataType::Boolean,
2049 [false, false, true, false, false],
2050 );
2051 test_coercion!(
2052 Int16Array,
2053 DataType::Int16,
2054 vec![1i16, 2i16, 3i16],
2055 Int64Array,
2056 DataType::Int64,
2057 vec![10i64, 4i64, 5i64],
2058 Operator::BitwiseAnd,
2059 Int64Array,
2060 DataType::Int64,
2061 [0i64, 0i64, 1i64],
2062 );
2063 test_coercion!(
2064 UInt16Array,
2065 DataType::UInt16,
2066 vec![1u16, 2u16, 3u16],
2067 UInt64Array,
2068 DataType::UInt64,
2069 vec![10u64, 4u64, 5u64],
2070 Operator::BitwiseAnd,
2071 UInt64Array,
2072 DataType::UInt64,
2073 [0u64, 0u64, 1u64],
2074 );
2075 test_coercion!(
2076 Int16Array,
2077 DataType::Int16,
2078 vec![3i16, 2i16, 3i16],
2079 Int64Array,
2080 DataType::Int64,
2081 vec![10i64, 6i64, 5i64],
2082 Operator::BitwiseOr,
2083 Int64Array,
2084 DataType::Int64,
2085 [11i64, 6i64, 7i64],
2086 );
2087 test_coercion!(
2088 UInt16Array,
2089 DataType::UInt16,
2090 vec![1u16, 2u16, 3u16],
2091 UInt64Array,
2092 DataType::UInt64,
2093 vec![10u64, 4u64, 5u64],
2094 Operator::BitwiseOr,
2095 UInt64Array,
2096 DataType::UInt64,
2097 [11u64, 6u64, 7u64],
2098 );
2099 test_coercion!(
2100 Int16Array,
2101 DataType::Int16,
2102 vec![3i16, 2i16, 3i16],
2103 Int64Array,
2104 DataType::Int64,
2105 vec![10i64, 6i64, 5i64],
2106 Operator::BitwiseXor,
2107 Int64Array,
2108 DataType::Int64,
2109 [9i64, 4i64, 6i64],
2110 );
2111 test_coercion!(
2112 UInt16Array,
2113 DataType::UInt16,
2114 vec![3u16, 2u16, 3u16],
2115 UInt64Array,
2116 DataType::UInt64,
2117 vec![10u64, 6u64, 5u64],
2118 Operator::BitwiseXor,
2119 UInt64Array,
2120 DataType::UInt64,
2121 [9u64, 4u64, 6u64],
2122 );
2123 test_coercion!(
2124 Int16Array,
2125 DataType::Int16,
2126 vec![4i16, 27i16, 35i16],
2127 Int64Array,
2128 DataType::Int64,
2129 vec![2i64, 3i64, 4i64],
2130 Operator::BitwiseShiftRight,
2131 Int64Array,
2132 DataType::Int64,
2133 [1i64, 3i64, 2i64],
2134 );
2135 test_coercion!(
2136 UInt16Array,
2137 DataType::UInt16,
2138 vec![4u16, 27u16, 35u16],
2139 UInt64Array,
2140 DataType::UInt64,
2141 vec![2u64, 3u64, 4u64],
2142 Operator::BitwiseShiftRight,
2143 UInt64Array,
2144 DataType::UInt64,
2145 [1u64, 3u64, 2u64],
2146 );
2147 test_coercion!(
2148 Int16Array,
2149 DataType::Int16,
2150 vec![2i16, 3i16, 4i16],
2151 Int64Array,
2152 DataType::Int64,
2153 vec![4i64, 12i64, 7i64],
2154 Operator::BitwiseShiftLeft,
2155 Int64Array,
2156 DataType::Int64,
2157 [32i64, 12288i64, 512i64],
2158 );
2159 test_coercion!(
2160 UInt16Array,
2161 DataType::UInt16,
2162 vec![2u16, 3u16, 4u16],
2163 UInt64Array,
2164 DataType::UInt64,
2165 vec![4u64, 12u64, 7u64],
2166 Operator::BitwiseShiftLeft,
2167 UInt64Array,
2168 DataType::UInt64,
2169 [32u64, 12288u64, 512u64],
2170 );
2171 Ok(())
2172 }
2173
2174 #[test]
2180 fn test_dictionary_type_to_array_coercion() -> Result<()> {
2181 let dict_type =
2183 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
2184 let string_type = DataType::Utf8;
2185
2186 let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
2188
2189 dict_builder.append("one")?;
2190 dict_builder.append_null();
2191 dict_builder.append("three")?;
2192 dict_builder.append("four")?;
2193 let dict_array = Arc::new(dict_builder.finish()) as ArrayRef;
2194
2195 let str_array = Arc::new(StringArray::from(vec![
2196 Some("not one"),
2197 Some("two"),
2198 None,
2199 Some("four"),
2200 ])) as ArrayRef;
2201
2202 let schema = Arc::new(Schema::new(vec![
2203 Field::new("a", dict_type.clone(), true),
2204 Field::new("b", string_type.clone(), true),
2205 ]));
2206
2207 let result = BooleanArray::from(vec![Some(false), None, None, Some(true)]);
2209 apply_logic_op(&schema, &dict_array, &str_array, Operator::Eq, result)?;
2210
2211 let schema = Arc::new(Schema::new(vec![
2214 Field::new("a", string_type, true),
2215 Field::new("b", dict_type, true),
2216 ]));
2217 let result = BooleanArray::from(vec![Some(false), None, None, Some(true)]);
2218 apply_logic_op(&schema, &str_array, &dict_array, Operator::Eq, result)?;
2219
2220 Ok(())
2221 }
2222
2223 #[test]
2224 fn plus_op() -> Result<()> {
2225 let schema = Schema::new(vec![
2226 Field::new("a", DataType::Int32, false),
2227 Field::new("b", DataType::Int32, false),
2228 ]);
2229 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2230 let b = Int32Array::from(vec![1, 2, 4, 8, 16]);
2231
2232 apply_arithmetic::<Int32Type>(
2233 Arc::new(schema),
2234 vec![Arc::new(a), Arc::new(b)],
2235 Operator::Plus,
2236 Int32Array::from(vec![2, 4, 7, 12, 21]),
2237 )?;
2238
2239 Ok(())
2240 }
2241
2242 #[test]
2243 fn plus_op_dict() -> Result<()> {
2244 let schema = Schema::new(vec![
2245 Field::new(
2246 "a",
2247 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2248 true,
2249 ),
2250 Field::new(
2251 "b",
2252 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2253 true,
2254 ),
2255 ]);
2256
2257 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2258 let keys = Int8Array::from(vec![Some(0), None, Some(1), Some(3), None]);
2259 let a = DictionaryArray::try_new(keys, Arc::new(a))?;
2260
2261 let b = Int32Array::from(vec![1, 2, 4, 8, 16]);
2262 let keys = Int8Array::from(vec![0, 1, 1, 2, 1]);
2263 let b = DictionaryArray::try_new(keys, Arc::new(b))?;
2264
2265 apply_arithmetic::<Int32Type>(
2266 Arc::new(schema),
2267 vec![Arc::new(a), Arc::new(b)],
2268 Operator::Plus,
2269 Int32Array::from(vec![Some(2), None, Some(4), Some(8), None]),
2270 )?;
2271
2272 Ok(())
2273 }
2274
2275 #[test]
2276 fn plus_op_dict_decimal() -> Result<()> {
2277 let schema = Schema::new(vec![
2278 Field::new(
2279 "a",
2280 DataType::Dictionary(
2281 Box::new(DataType::Int8),
2282 Box::new(DataType::Decimal128(10, 0)),
2283 ),
2284 true,
2285 ),
2286 Field::new(
2287 "b",
2288 DataType::Dictionary(
2289 Box::new(DataType::Int8),
2290 Box::new(DataType::Decimal128(10, 0)),
2291 ),
2292 true,
2293 ),
2294 ]);
2295
2296 let value = 123;
2297 let decimal_array = Arc::new(create_decimal_array(
2298 &[
2299 Some(value),
2300 Some(value + 2),
2301 Some(value - 1),
2302 Some(value + 1),
2303 ],
2304 10,
2305 0,
2306 ));
2307
2308 let keys = Int8Array::from(vec![Some(0), Some(2), None, Some(3), Some(0)]);
2309 let a = DictionaryArray::try_new(keys, decimal_array)?;
2310
2311 let keys = Int8Array::from(vec![Some(0), None, Some(3), Some(2), Some(2)]);
2312 let decimal_array = Arc::new(create_decimal_array(
2313 &[
2314 Some(value + 1),
2315 Some(value + 3),
2316 Some(value),
2317 Some(value + 2),
2318 ],
2319 10,
2320 0,
2321 ));
2322 let b = DictionaryArray::try_new(keys, decimal_array)?;
2323
2324 apply_arithmetic(
2325 Arc::new(schema),
2326 vec![Arc::new(a), Arc::new(b)],
2327 Operator::Plus,
2328 create_decimal_array(&[Some(247), None, None, Some(247), Some(246)], 11, 0),
2329 )?;
2330
2331 Ok(())
2332 }
2333
2334 #[test]
2335 fn plus_op_scalar() -> Result<()> {
2336 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2337 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2338
2339 apply_arithmetic_scalar(
2340 Arc::new(schema),
2341 vec![Arc::new(a)],
2342 Operator::Plus,
2343 ScalarValue::Int32(Some(1)),
2344 Arc::new(Int32Array::from(vec![2, 3, 4, 5, 6])),
2345 )?;
2346
2347 Ok(())
2348 }
2349
2350 #[test]
2351 fn plus_op_dict_scalar() -> Result<()> {
2352 let schema = Schema::new(vec![Field::new(
2353 "a",
2354 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2355 true,
2356 )]);
2357
2358 let mut dict_builder = PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::new();
2359
2360 dict_builder.append(1)?;
2361 dict_builder.append_null();
2362 dict_builder.append(2)?;
2363 dict_builder.append(5)?;
2364
2365 let a = dict_builder.finish();
2366
2367 let expected: PrimitiveArray<Int32Type> =
2368 PrimitiveArray::from(vec![Some(2), None, Some(3), Some(6)]);
2369
2370 apply_arithmetic_scalar(
2371 Arc::new(schema),
2372 vec![Arc::new(a)],
2373 Operator::Plus,
2374 ScalarValue::Dictionary(
2375 Box::new(DataType::Int8),
2376 Box::new(ScalarValue::Int32(Some(1))),
2377 ),
2378 Arc::new(expected),
2379 )?;
2380
2381 Ok(())
2382 }
2383
2384 #[test]
2385 fn plus_op_dict_scalar_decimal() -> Result<()> {
2386 let schema = Schema::new(vec![Field::new(
2387 "a",
2388 DataType::Dictionary(
2389 Box::new(DataType::Int8),
2390 Box::new(DataType::Decimal128(10, 0)),
2391 ),
2392 true,
2393 )]);
2394
2395 let value = 123;
2396 let decimal_array = Arc::new(create_decimal_array(
2397 &[Some(value), None, Some(value - 1), Some(value + 1)],
2398 10,
2399 0,
2400 ));
2401
2402 let keys = Int8Array::from(vec![0, 2, 1, 3, 0]);
2403 let a = DictionaryArray::try_new(keys, decimal_array)?;
2404
2405 let decimal_array = Arc::new(create_decimal_array(
2406 &[
2407 Some(value + 1),
2408 Some(value),
2409 None,
2410 Some(value + 2),
2411 Some(value + 1),
2412 ],
2413 11,
2414 0,
2415 ));
2416
2417 apply_arithmetic_scalar(
2418 Arc::new(schema),
2419 vec![Arc::new(a)],
2420 Operator::Plus,
2421 ScalarValue::Dictionary(
2422 Box::new(DataType::Int8),
2423 Box::new(ScalarValue::Decimal128(Some(1), 10, 0)),
2424 ),
2425 decimal_array,
2426 )?;
2427
2428 Ok(())
2429 }
2430
2431 #[test]
2432 fn minus_op() -> Result<()> {
2433 let schema = Arc::new(Schema::new(vec![
2434 Field::new("a", DataType::Int32, false),
2435 Field::new("b", DataType::Int32, false),
2436 ]));
2437 let a = Arc::new(Int32Array::from(vec![1, 2, 4, 8, 16]));
2438 let b = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
2439
2440 apply_arithmetic::<Int32Type>(
2441 Arc::clone(&schema),
2442 vec![
2443 Arc::clone(&a) as Arc<dyn Array>,
2444 Arc::clone(&b) as Arc<dyn Array>,
2445 ],
2446 Operator::Minus,
2447 Int32Array::from(vec![0, 0, 1, 4, 11]),
2448 )?;
2449
2450 apply_arithmetic::<Int32Type>(
2452 schema,
2453 vec![b, a],
2454 Operator::Minus,
2455 Int32Array::from(vec![0, 0, -1, -4, -11]),
2456 )?;
2457
2458 Ok(())
2459 }
2460
2461 #[test]
2462 fn date32_minus_date32_returns_int64_days() -> Result<()> {
2463 let schema = Arc::new(Schema::new(vec![
2464 Field::new("a", DataType::Date32, true),
2465 Field::new("b", DataType::Date32, true),
2466 ]));
2467 let a = Arc::new(Date32Array::from(vec![
2468 Some(18_901),
2469 Some(18_901),
2470 None,
2471 Some(18_900),
2472 ]));
2473 let b = Arc::new(Date32Array::from(vec![
2474 Some(18_898),
2475 Some(18_904),
2476 Some(18_900),
2477 None,
2478 ]));
2479
2480 apply_arithmetic::<Int64Type>(
2481 schema,
2482 vec![a, b],
2483 Operator::Minus,
2484 Int64Array::from(vec![Some(3), Some(-3), None, None]),
2485 )?;
2486
2487 Ok(())
2488 }
2489
2490 #[test]
2491 fn date64_minus_date64_returns_int64_days() -> Result<()> {
2492 let schema = Arc::new(Schema::new(vec![
2493 Field::new("a", DataType::Date64, true),
2494 Field::new("b", DataType::Date64, true),
2495 ]));
2496 let a = Arc::new(Date64Array::from(vec![
2497 Some(18_901 * MILLIS_PER_DAY),
2498 Some(18_901 * MILLIS_PER_DAY),
2499 None,
2500 Some(18_900 * MILLIS_PER_DAY),
2501 ]));
2502 let b = Arc::new(Date64Array::from(vec![
2503 Some(18_898 * MILLIS_PER_DAY),
2504 Some(18_904 * MILLIS_PER_DAY),
2505 Some(18_900 * MILLIS_PER_DAY),
2506 None,
2507 ]));
2508
2509 apply_arithmetic::<Int64Type>(
2510 schema,
2511 vec![a, b],
2512 Operator::Minus,
2513 Int64Array::from(vec![Some(3), Some(-3), None, None]),
2514 )?;
2515
2516 Ok(())
2517 }
2518
2519 #[test]
2520 fn date32_minus_null_scalar_returns_int64_null_scalar() -> Result<()> {
2521 let result = apply_date_subtraction(
2522 &ColumnarValue::Array(Arc::new(Date32Array::from(vec![
2523 Some(18_901),
2524 Some(18_900),
2525 ]))),
2526 &ColumnarValue::Scalar(ScalarValue::Date32(None)),
2527 )?;
2528
2529 assert!(matches!(
2530 result,
2531 ColumnarValue::Scalar(ScalarValue::Int64(None))
2532 ));
2533
2534 Ok(())
2535 }
2536
2537 #[test]
2538 fn minus_op_dict() -> Result<()> {
2539 let schema = Schema::new(vec![
2540 Field::new(
2541 "a",
2542 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2543 true,
2544 ),
2545 Field::new(
2546 "b",
2547 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2548 true,
2549 ),
2550 ]);
2551
2552 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2553 let keys = Int8Array::from(vec![Some(0), None, Some(1), Some(3), None]);
2554 let a = DictionaryArray::try_new(keys, Arc::new(a))?;
2555
2556 let b = Int32Array::from(vec![1, 2, 4, 8, 16]);
2557 let keys = Int8Array::from(vec![0, 1, 1, 2, 1]);
2558 let b = DictionaryArray::try_new(keys, Arc::new(b))?;
2559
2560 apply_arithmetic::<Int32Type>(
2561 Arc::new(schema),
2562 vec![Arc::new(a), Arc::new(b)],
2563 Operator::Minus,
2564 Int32Array::from(vec![Some(0), None, Some(0), Some(0), None]),
2565 )?;
2566
2567 Ok(())
2568 }
2569
2570 #[test]
2571 fn minus_op_dict_decimal() -> Result<()> {
2572 let schema = Schema::new(vec![
2573 Field::new(
2574 "a",
2575 DataType::Dictionary(
2576 Box::new(DataType::Int8),
2577 Box::new(DataType::Decimal128(10, 0)),
2578 ),
2579 true,
2580 ),
2581 Field::new(
2582 "b",
2583 DataType::Dictionary(
2584 Box::new(DataType::Int8),
2585 Box::new(DataType::Decimal128(10, 0)),
2586 ),
2587 true,
2588 ),
2589 ]);
2590
2591 let value = 123;
2592 let decimal_array = Arc::new(create_decimal_array(
2593 &[
2594 Some(value),
2595 Some(value + 2),
2596 Some(value - 1),
2597 Some(value + 1),
2598 ],
2599 10,
2600 0,
2601 ));
2602
2603 let keys = Int8Array::from(vec![Some(0), Some(2), None, Some(3), Some(0)]);
2604 let a = DictionaryArray::try_new(keys, decimal_array)?;
2605
2606 let keys = Int8Array::from(vec![Some(0), None, Some(3), Some(2), Some(2)]);
2607 let decimal_array = Arc::new(create_decimal_array(
2608 &[
2609 Some(value + 1),
2610 Some(value + 3),
2611 Some(value),
2612 Some(value + 2),
2613 ],
2614 10,
2615 0,
2616 ));
2617 let b = DictionaryArray::try_new(keys, decimal_array)?;
2618
2619 apply_arithmetic(
2620 Arc::new(schema),
2621 vec![Arc::new(a), Arc::new(b)],
2622 Operator::Minus,
2623 create_decimal_array(&[Some(-1), None, None, Some(1), Some(0)], 11, 0),
2624 )?;
2625
2626 Ok(())
2627 }
2628
2629 #[test]
2630 fn minus_op_scalar() -> Result<()> {
2631 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2632 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2633
2634 apply_arithmetic_scalar(
2635 Arc::new(schema),
2636 vec![Arc::new(a)],
2637 Operator::Minus,
2638 ScalarValue::Int32(Some(1)),
2639 Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4])),
2640 )?;
2641
2642 Ok(())
2643 }
2644
2645 #[test]
2646 fn minus_op_dict_scalar() -> Result<()> {
2647 let schema = Schema::new(vec![Field::new(
2648 "a",
2649 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2650 true,
2651 )]);
2652
2653 let mut dict_builder = PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::new();
2654
2655 dict_builder.append(1)?;
2656 dict_builder.append_null();
2657 dict_builder.append(2)?;
2658 dict_builder.append(5)?;
2659
2660 let a = dict_builder.finish();
2661
2662 let expected: PrimitiveArray<Int32Type> =
2663 PrimitiveArray::from(vec![Some(0), None, Some(1), Some(4)]);
2664
2665 apply_arithmetic_scalar(
2666 Arc::new(schema),
2667 vec![Arc::new(a)],
2668 Operator::Minus,
2669 ScalarValue::Dictionary(
2670 Box::new(DataType::Int8),
2671 Box::new(ScalarValue::Int32(Some(1))),
2672 ),
2673 Arc::new(expected),
2674 )?;
2675
2676 Ok(())
2677 }
2678
2679 #[test]
2680 fn minus_op_dict_scalar_decimal() -> Result<()> {
2681 let schema = Schema::new(vec![Field::new(
2682 "a",
2683 DataType::Dictionary(
2684 Box::new(DataType::Int8),
2685 Box::new(DataType::Decimal128(10, 0)),
2686 ),
2687 true,
2688 )]);
2689
2690 let value = 123;
2691 let decimal_array = Arc::new(create_decimal_array(
2692 &[Some(value), None, Some(value - 1), Some(value + 1)],
2693 10,
2694 0,
2695 ));
2696
2697 let keys = Int8Array::from(vec![0, 2, 1, 3, 0]);
2698 let a = DictionaryArray::try_new(keys, decimal_array)?;
2699
2700 let decimal_array = Arc::new(create_decimal_array(
2701 &[
2702 Some(value - 1),
2703 Some(value - 2),
2704 None,
2705 Some(value),
2706 Some(value - 1),
2707 ],
2708 11,
2709 0,
2710 ));
2711
2712 apply_arithmetic_scalar(
2713 Arc::new(schema),
2714 vec![Arc::new(a)],
2715 Operator::Minus,
2716 ScalarValue::Dictionary(
2717 Box::new(DataType::Int8),
2718 Box::new(ScalarValue::Decimal128(Some(1), 10, 0)),
2719 ),
2720 decimal_array,
2721 )?;
2722
2723 Ok(())
2724 }
2725
2726 #[test]
2727 fn multiply_op() -> Result<()> {
2728 let schema = Arc::new(Schema::new(vec![
2729 Field::new("a", DataType::Int32, false),
2730 Field::new("b", DataType::Int32, false),
2731 ]));
2732 let a = Arc::new(Int32Array::from(vec![4, 8, 16, 32, 64]));
2733 let b = Arc::new(Int32Array::from(vec![2, 4, 8, 16, 32]));
2734
2735 apply_arithmetic::<Int32Type>(
2736 schema,
2737 vec![a, b],
2738 Operator::Multiply,
2739 Int32Array::from(vec![8, 32, 128, 512, 2048]),
2740 )?;
2741
2742 Ok(())
2743 }
2744
2745 #[test]
2746 fn multiply_op_dict() -> Result<()> {
2747 let schema = Schema::new(vec![
2748 Field::new(
2749 "a",
2750 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2751 true,
2752 ),
2753 Field::new(
2754 "b",
2755 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2756 true,
2757 ),
2758 ]);
2759
2760 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2761 let keys = Int8Array::from(vec![Some(0), None, Some(1), Some(3), None]);
2762 let a = DictionaryArray::try_new(keys, Arc::new(a))?;
2763
2764 let b = Int32Array::from(vec![1, 2, 4, 8, 16]);
2765 let keys = Int8Array::from(vec![0, 1, 1, 2, 1]);
2766 let b = DictionaryArray::try_new(keys, Arc::new(b))?;
2767
2768 apply_arithmetic::<Int32Type>(
2769 Arc::new(schema),
2770 vec![Arc::new(a), Arc::new(b)],
2771 Operator::Multiply,
2772 Int32Array::from(vec![Some(1), None, Some(4), Some(16), None]),
2773 )?;
2774
2775 Ok(())
2776 }
2777
2778 #[test]
2779 fn multiply_op_dict_decimal() -> Result<()> {
2780 let schema = Schema::new(vec![
2781 Field::new(
2782 "a",
2783 DataType::Dictionary(
2784 Box::new(DataType::Int8),
2785 Box::new(DataType::Decimal128(10, 0)),
2786 ),
2787 true,
2788 ),
2789 Field::new(
2790 "b",
2791 DataType::Dictionary(
2792 Box::new(DataType::Int8),
2793 Box::new(DataType::Decimal128(10, 0)),
2794 ),
2795 true,
2796 ),
2797 ]);
2798
2799 let value = 123;
2800 let decimal_array = Arc::new(create_decimal_array(
2801 &[
2802 Some(value),
2803 Some(value + 2),
2804 Some(value - 1),
2805 Some(value + 1),
2806 ],
2807 10,
2808 0,
2809 )) as ArrayRef;
2810
2811 let keys = Int8Array::from(vec![Some(0), Some(2), None, Some(3), Some(0)]);
2812 let a = DictionaryArray::try_new(keys, decimal_array)?;
2813
2814 let keys = Int8Array::from(vec![Some(0), None, Some(3), Some(2), Some(2)]);
2815 let decimal_array = Arc::new(create_decimal_array(
2816 &[
2817 Some(value + 1),
2818 Some(value + 3),
2819 Some(value),
2820 Some(value + 2),
2821 ],
2822 10,
2823 0,
2824 ));
2825 let b = DictionaryArray::try_new(keys, decimal_array)?;
2826
2827 apply_arithmetic(
2828 Arc::new(schema),
2829 vec![Arc::new(a), Arc::new(b)],
2830 Operator::Multiply,
2831 create_decimal_array(
2832 &[Some(15252), None, None, Some(15252), Some(15129)],
2833 21,
2834 0,
2835 ),
2836 )?;
2837
2838 Ok(())
2839 }
2840
2841 #[test]
2842 fn multiply_op_scalar() -> Result<()> {
2843 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2844 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2845
2846 apply_arithmetic_scalar(
2847 Arc::new(schema),
2848 vec![Arc::new(a)],
2849 Operator::Multiply,
2850 ScalarValue::Int32(Some(2)),
2851 Arc::new(Int32Array::from(vec![2, 4, 6, 8, 10])),
2852 )?;
2853
2854 Ok(())
2855 }
2856
2857 #[test]
2858 fn multiply_op_dict_scalar() -> Result<()> {
2859 let schema = Schema::new(vec![Field::new(
2860 "a",
2861 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2862 true,
2863 )]);
2864
2865 let mut dict_builder = PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::new();
2866
2867 dict_builder.append(1)?;
2868 dict_builder.append_null();
2869 dict_builder.append(2)?;
2870 dict_builder.append(5)?;
2871
2872 let a = dict_builder.finish();
2873
2874 let expected: PrimitiveArray<Int32Type> =
2875 PrimitiveArray::from(vec![Some(2), None, Some(4), Some(10)]);
2876
2877 apply_arithmetic_scalar(
2878 Arc::new(schema),
2879 vec![Arc::new(a)],
2880 Operator::Multiply,
2881 ScalarValue::Dictionary(
2882 Box::new(DataType::Int8),
2883 Box::new(ScalarValue::Int32(Some(2))),
2884 ),
2885 Arc::new(expected),
2886 )?;
2887
2888 Ok(())
2889 }
2890
2891 #[test]
2892 fn multiply_op_dict_scalar_decimal() -> Result<()> {
2893 let schema = Schema::new(vec![Field::new(
2894 "a",
2895 DataType::Dictionary(
2896 Box::new(DataType::Int8),
2897 Box::new(DataType::Decimal128(10, 0)),
2898 ),
2899 true,
2900 )]);
2901
2902 let value = 123;
2903 let decimal_array = Arc::new(create_decimal_array(
2904 &[Some(value), None, Some(value - 1), Some(value + 1)],
2905 10,
2906 0,
2907 ));
2908
2909 let keys = Int8Array::from(vec![0, 2, 1, 3, 0]);
2910 let a = DictionaryArray::try_new(keys, decimal_array)?;
2911
2912 let decimal_array = Arc::new(create_decimal_array(
2913 &[Some(246), Some(244), None, Some(248), Some(246)],
2914 21,
2915 0,
2916 ));
2917
2918 apply_arithmetic_scalar(
2919 Arc::new(schema),
2920 vec![Arc::new(a)],
2921 Operator::Multiply,
2922 ScalarValue::Dictionary(
2923 Box::new(DataType::Int8),
2924 Box::new(ScalarValue::Decimal128(Some(2), 10, 0)),
2925 ),
2926 decimal_array,
2927 )?;
2928
2929 Ok(())
2930 }
2931
2932 #[test]
2933 fn divide_op() -> Result<()> {
2934 let schema = Arc::new(Schema::new(vec![
2935 Field::new("a", DataType::Int32, false),
2936 Field::new("b", DataType::Int32, false),
2937 ]));
2938 let a = Arc::new(Int32Array::from(vec![8, 32, 128, 512, 2048]));
2939 let b = Arc::new(Int32Array::from(vec![2, 4, 8, 16, 32]));
2940
2941 apply_arithmetic::<Int32Type>(
2942 schema,
2943 vec![a, b],
2944 Operator::Divide,
2945 Int32Array::from(vec![4, 8, 16, 32, 64]),
2946 )?;
2947
2948 Ok(())
2949 }
2950
2951 #[test]
2952 fn divide_op_dict() -> Result<()> {
2953 let schema = Schema::new(vec![
2954 Field::new(
2955 "a",
2956 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2957 true,
2958 ),
2959 Field::new(
2960 "b",
2961 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
2962 true,
2963 ),
2964 ]);
2965
2966 let mut dict_builder = PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::new();
2967
2968 dict_builder.append(1)?;
2969 dict_builder.append_null();
2970 dict_builder.append(2)?;
2971 dict_builder.append(5)?;
2972 dict_builder.append(0)?;
2973
2974 let a = dict_builder.finish();
2975
2976 let b = Int32Array::from(vec![1, 2, 4, 8, 16]);
2977 let keys = Int8Array::from(vec![0, 1, 1, 2, 1]);
2978 let b = DictionaryArray::try_new(keys, Arc::new(b))?;
2979
2980 apply_arithmetic::<Int32Type>(
2981 Arc::new(schema),
2982 vec![Arc::new(a), Arc::new(b)],
2983 Operator::Divide,
2984 Int32Array::from(vec![Some(1), None, Some(1), Some(1), Some(0)]),
2985 )?;
2986
2987 Ok(())
2988 }
2989
2990 #[test]
2991 fn divide_op_dict_decimal() -> Result<()> {
2992 let schema = Schema::new(vec![
2993 Field::new(
2994 "a",
2995 DataType::Dictionary(
2996 Box::new(DataType::Int8),
2997 Box::new(DataType::Decimal128(10, 0)),
2998 ),
2999 true,
3000 ),
3001 Field::new(
3002 "b",
3003 DataType::Dictionary(
3004 Box::new(DataType::Int8),
3005 Box::new(DataType::Decimal128(10, 0)),
3006 ),
3007 true,
3008 ),
3009 ]);
3010
3011 let value = 123;
3012 let decimal_array = Arc::new(create_decimal_array(
3013 &[
3014 Some(value),
3015 Some(value + 2),
3016 Some(value - 1),
3017 Some(value + 1),
3018 ],
3019 10,
3020 0,
3021 ));
3022
3023 let keys = Int8Array::from(vec![Some(0), Some(2), None, Some(3), Some(0)]);
3024 let a = DictionaryArray::try_new(keys, decimal_array)?;
3025
3026 let keys = Int8Array::from(vec![Some(0), None, Some(3), Some(2), Some(2)]);
3027 let decimal_array = Arc::new(create_decimal_array(
3028 &[
3029 Some(value + 1),
3030 Some(value + 3),
3031 Some(value),
3032 Some(value + 2),
3033 ],
3034 10,
3035 0,
3036 ));
3037 let b = DictionaryArray::try_new(keys, decimal_array)?;
3038
3039 apply_arithmetic(
3040 Arc::new(schema),
3041 vec![Arc::new(a), Arc::new(b)],
3042 Operator::Divide,
3043 create_decimal_array(
3044 &[
3045 Some(9919), None,
3047 None,
3048 Some(10081), Some(10000), ],
3051 14,
3052 4,
3053 ),
3054 )?;
3055
3056 Ok(())
3057 }
3058
3059 #[test]
3060 fn divide_op_scalar() -> Result<()> {
3061 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
3062 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
3063
3064 apply_arithmetic_scalar(
3065 Arc::new(schema),
3066 vec![Arc::new(a)],
3067 Operator::Divide,
3068 ScalarValue::Int32(Some(2)),
3069 Arc::new(Int32Array::from(vec![0, 1, 1, 2, 2])),
3070 )?;
3071
3072 Ok(())
3073 }
3074
3075 #[test]
3076 fn divide_op_dict_scalar() -> Result<()> {
3077 let schema = Schema::new(vec![Field::new(
3078 "a",
3079 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3080 true,
3081 )]);
3082
3083 let mut dict_builder = PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::new();
3084
3085 dict_builder.append(1)?;
3086 dict_builder.append_null();
3087 dict_builder.append(2)?;
3088 dict_builder.append(5)?;
3089
3090 let a = dict_builder.finish();
3091
3092 let expected: PrimitiveArray<Int32Type> =
3093 PrimitiveArray::from(vec![Some(0), None, Some(1), Some(2)]);
3094
3095 apply_arithmetic_scalar(
3096 Arc::new(schema),
3097 vec![Arc::new(a)],
3098 Operator::Divide,
3099 ScalarValue::Dictionary(
3100 Box::new(DataType::Int8),
3101 Box::new(ScalarValue::Int32(Some(2))),
3102 ),
3103 Arc::new(expected),
3104 )?;
3105
3106 Ok(())
3107 }
3108
3109 #[test]
3110 fn divide_op_dict_scalar_decimal() -> Result<()> {
3111 let schema = Schema::new(vec![Field::new(
3112 "a",
3113 DataType::Dictionary(
3114 Box::new(DataType::Int8),
3115 Box::new(DataType::Decimal128(10, 0)),
3116 ),
3117 true,
3118 )]);
3119
3120 let value = 123;
3121 let decimal_array = Arc::new(create_decimal_array(
3122 &[Some(value), None, Some(value - 1), Some(value + 1)],
3123 10,
3124 0,
3125 ));
3126
3127 let keys = Int8Array::from(vec![0, 2, 1, 3, 0]);
3128 let a = DictionaryArray::try_new(keys, decimal_array)?;
3129
3130 let decimal_array = Arc::new(create_decimal_array(
3131 &[Some(615000), Some(610000), None, Some(620000), Some(615000)],
3132 14,
3133 4,
3134 ));
3135
3136 apply_arithmetic_scalar(
3137 Arc::new(schema),
3138 vec![Arc::new(a)],
3139 Operator::Divide,
3140 ScalarValue::Dictionary(
3141 Box::new(DataType::Int8),
3142 Box::new(ScalarValue::Decimal128(Some(2), 10, 0)),
3143 ),
3144 decimal_array,
3145 )?;
3146
3147 Ok(())
3148 }
3149
3150 #[test]
3151 fn modulus_op() -> Result<()> {
3152 let schema = Arc::new(Schema::new(vec![
3153 Field::new("a", DataType::Int32, false),
3154 Field::new("b", DataType::Int32, false),
3155 ]));
3156 let a = Arc::new(Int32Array::from(vec![8, 32, 128, 512, 2048]));
3157 let b = Arc::new(Int32Array::from(vec![2, 4, 7, 14, 32]));
3158
3159 apply_arithmetic::<Int32Type>(
3160 schema,
3161 vec![a, b],
3162 Operator::Modulo,
3163 Int32Array::from(vec![0, 0, 2, 8, 0]),
3164 )?;
3165
3166 Ok(())
3167 }
3168
3169 #[test]
3170 fn modulus_op_dict() -> Result<()> {
3171 let schema = Schema::new(vec![
3172 Field::new(
3173 "a",
3174 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3175 true,
3176 ),
3177 Field::new(
3178 "b",
3179 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3180 true,
3181 ),
3182 ]);
3183
3184 let mut dict_builder = PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::new();
3185
3186 dict_builder.append(1)?;
3187 dict_builder.append_null();
3188 dict_builder.append(2)?;
3189 dict_builder.append(5)?;
3190 dict_builder.append(0)?;
3191
3192 let a = dict_builder.finish();
3193
3194 let b = Int32Array::from(vec![1, 2, 4, 8, 16]);
3195 let keys = Int8Array::from(vec![0, 1, 1, 2, 1]);
3196 let b = DictionaryArray::try_new(keys, Arc::new(b))?;
3197
3198 apply_arithmetic::<Int32Type>(
3199 Arc::new(schema),
3200 vec![Arc::new(a), Arc::new(b)],
3201 Operator::Modulo,
3202 Int32Array::from(vec![Some(0), None, Some(0), Some(1), Some(0)]),
3203 )?;
3204
3205 Ok(())
3206 }
3207
3208 #[test]
3209 fn modulus_op_dict_decimal() -> Result<()> {
3210 let schema = Schema::new(vec![
3211 Field::new(
3212 "a",
3213 DataType::Dictionary(
3214 Box::new(DataType::Int8),
3215 Box::new(DataType::Decimal128(10, 0)),
3216 ),
3217 true,
3218 ),
3219 Field::new(
3220 "b",
3221 DataType::Dictionary(
3222 Box::new(DataType::Int8),
3223 Box::new(DataType::Decimal128(10, 0)),
3224 ),
3225 true,
3226 ),
3227 ]);
3228
3229 let value = 123;
3230 let decimal_array = Arc::new(create_decimal_array(
3231 &[
3232 Some(value),
3233 Some(value + 2),
3234 Some(value - 1),
3235 Some(value + 1),
3236 ],
3237 10,
3238 0,
3239 ));
3240
3241 let keys = Int8Array::from(vec![Some(0), Some(2), None, Some(3), Some(0)]);
3242 let a = DictionaryArray::try_new(keys, decimal_array)?;
3243
3244 let keys = Int8Array::from(vec![Some(0), None, Some(3), Some(2), Some(2)]);
3245 let decimal_array = Arc::new(create_decimal_array(
3246 &[
3247 Some(value + 1),
3248 Some(value + 3),
3249 Some(value),
3250 Some(value + 2),
3251 ],
3252 10,
3253 0,
3254 ));
3255 let b = DictionaryArray::try_new(keys, decimal_array)?;
3256
3257 apply_arithmetic(
3258 Arc::new(schema),
3259 vec![Arc::new(a), Arc::new(b)],
3260 Operator::Modulo,
3261 create_decimal_array(&[Some(123), None, None, Some(1), Some(0)], 10, 0),
3262 )?;
3263
3264 Ok(())
3265 }
3266
3267 #[test]
3268 fn modulus_op_scalar() -> Result<()> {
3269 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
3270 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
3271
3272 apply_arithmetic_scalar(
3273 Arc::new(schema),
3274 vec![Arc::new(a)],
3275 Operator::Modulo,
3276 ScalarValue::Int32(Some(2)),
3277 Arc::new(Int32Array::from(vec![1, 0, 1, 0, 1])),
3278 )?;
3279
3280 Ok(())
3281 }
3282
3283 #[test]
3284 fn modules_op_dict_scalar() -> Result<()> {
3285 let schema = Schema::new(vec![Field::new(
3286 "a",
3287 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3288 true,
3289 )]);
3290
3291 let mut dict_builder = PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::new();
3292
3293 dict_builder.append(1)?;
3294 dict_builder.append_null();
3295 dict_builder.append(2)?;
3296 dict_builder.append(5)?;
3297
3298 let a = dict_builder.finish();
3299
3300 let expected: PrimitiveArray<Int32Type> =
3301 PrimitiveArray::from(vec![Some(1), None, Some(0), Some(1)]);
3302
3303 apply_arithmetic_scalar(
3304 Arc::new(schema),
3305 vec![Arc::new(a)],
3306 Operator::Modulo,
3307 ScalarValue::Dictionary(
3308 Box::new(DataType::Int8),
3309 Box::new(ScalarValue::Int32(Some(2))),
3310 ),
3311 Arc::new(expected),
3312 )?;
3313
3314 Ok(())
3315 }
3316
3317 #[test]
3318 fn modulus_op_dict_scalar_decimal() -> Result<()> {
3319 let schema = Schema::new(vec![Field::new(
3320 "a",
3321 DataType::Dictionary(
3322 Box::new(DataType::Int8),
3323 Box::new(DataType::Decimal128(10, 0)),
3324 ),
3325 true,
3326 )]);
3327
3328 let value = 123;
3329 let decimal_array = Arc::new(create_decimal_array(
3330 &[Some(value), None, Some(value - 1), Some(value + 1)],
3331 10,
3332 0,
3333 ));
3334
3335 let keys = Int8Array::from(vec![0, 2, 1, 3, 0]);
3336 let a = DictionaryArray::try_new(keys, decimal_array)?;
3337
3338 let decimal_array = Arc::new(create_decimal_array(
3339 &[Some(1), Some(0), None, Some(0), Some(1)],
3340 10,
3341 0,
3342 ));
3343
3344 apply_arithmetic_scalar(
3345 Arc::new(schema),
3346 vec![Arc::new(a)],
3347 Operator::Modulo,
3348 ScalarValue::Dictionary(
3349 Box::new(DataType::Int8),
3350 Box::new(ScalarValue::Decimal128(Some(2), 10, 0)),
3351 ),
3352 decimal_array,
3353 )?;
3354
3355 Ok(())
3356 }
3357
3358 fn apply_arithmetic<T: ArrowNumericType>(
3359 schema: SchemaRef,
3360 data: Vec<ArrayRef>,
3361 op: Operator,
3362 expected: PrimitiveArray<T>,
3363 ) -> Result<()> {
3364 let arithmetic_op =
3365 binary_op(col("a", &schema)?, op, col("b", &schema)?, &schema)?;
3366 let batch = RecordBatch::try_new(schema, data)?;
3367 let result = arithmetic_op
3368 .evaluate(&batch)?
3369 .into_array(batch.num_rows())
3370 .expect("Failed to convert to array");
3371
3372 assert_eq!(result.as_ref(), &expected);
3373 Ok(())
3374 }
3375
3376 fn apply_arithmetic_scalar(
3377 schema: SchemaRef,
3378 data: Vec<ArrayRef>,
3379 op: Operator,
3380 literal: ScalarValue,
3381 expected: ArrayRef,
3382 ) -> Result<()> {
3383 let lit = Arc::new(Literal::new(literal));
3384 let arithmetic_op = binary_op(col("a", &schema)?, op, lit, &schema)?;
3385 let batch = RecordBatch::try_new(schema, data)?;
3386 let result = arithmetic_op
3387 .evaluate(&batch)?
3388 .into_array(batch.num_rows())
3389 .expect("Failed to convert to array");
3390
3391 assert_eq!(&result, &expected);
3392 Ok(())
3393 }
3394
3395 fn apply_logic_op(
3396 schema: &SchemaRef,
3397 left: &ArrayRef,
3398 right: &ArrayRef,
3399 op: Operator,
3400 expected: BooleanArray,
3401 ) -> Result<()> {
3402 let op = binary_op(col("a", schema)?, op, col("b", schema)?, schema)?;
3403 let data: Vec<ArrayRef> = vec![Arc::clone(left), Arc::clone(right)];
3404 let batch = RecordBatch::try_new(Arc::clone(schema), data)?;
3405 let result = op
3406 .evaluate(&batch)?
3407 .into_array(batch.num_rows())
3408 .expect("Failed to convert to array");
3409
3410 assert_eq!(result.as_ref(), &expected);
3411 Ok(())
3412 }
3413
3414 fn apply_logic_op_scalar_arr(
3416 schema: &SchemaRef,
3417 scalar: &ScalarValue,
3418 arr: &ArrayRef,
3419 op: Operator,
3420 expected: &BooleanArray,
3421 ) -> Result<()> {
3422 let scalar = lit(scalar.clone());
3423 let op = binary_op(scalar, op, col("a", schema)?, schema)?;
3424 let batch = RecordBatch::try_new(Arc::clone(schema), vec![Arc::clone(arr)])?;
3425 let result = op
3426 .evaluate(&batch)?
3427 .into_array(batch.num_rows())
3428 .expect("Failed to convert to array");
3429 assert_eq!(result.as_ref(), expected);
3430
3431 Ok(())
3432 }
3433
3434 fn apply_logic_op_arr_scalar(
3436 schema: &SchemaRef,
3437 arr: &ArrayRef,
3438 scalar: &ScalarValue,
3439 op: Operator,
3440 expected: &BooleanArray,
3441 ) -> Result<()> {
3442 let scalar = lit(scalar.clone());
3443 let op = binary_op(col("a", schema)?, op, scalar, schema)?;
3444 let batch = RecordBatch::try_new(Arc::clone(schema), vec![Arc::clone(arr)])?;
3445 let result = op
3446 .evaluate(&batch)?
3447 .into_array(batch.num_rows())
3448 .expect("Failed to convert to array");
3449 assert_eq!(result.as_ref(), expected);
3450
3451 Ok(())
3452 }
3453
3454 #[test]
3455 fn and_with_nulls_op() -> Result<()> {
3456 let schema = Schema::new(vec![
3457 Field::new("a", DataType::Boolean, true),
3458 Field::new("b", DataType::Boolean, true),
3459 ]);
3460 let a = Arc::new(BooleanArray::from(vec![
3461 Some(true),
3462 Some(false),
3463 None,
3464 Some(true),
3465 Some(false),
3466 None,
3467 Some(true),
3468 Some(false),
3469 None,
3470 ])) as ArrayRef;
3471 let b = Arc::new(BooleanArray::from(vec![
3472 Some(true),
3473 Some(true),
3474 Some(true),
3475 Some(false),
3476 Some(false),
3477 Some(false),
3478 None,
3479 None,
3480 None,
3481 ])) as ArrayRef;
3482
3483 let expected = BooleanArray::from(vec![
3484 Some(true),
3485 Some(false),
3486 None,
3487 Some(false),
3488 Some(false),
3489 Some(false),
3490 None,
3491 Some(false),
3492 None,
3493 ]);
3494 apply_logic_op(&Arc::new(schema), &a, &b, Operator::And, expected)?;
3495
3496 Ok(())
3497 }
3498
3499 #[test]
3500 fn regex_with_nulls() -> Result<()> {
3501 let schema = Schema::new(vec![
3502 Field::new("a", DataType::Utf8, true),
3503 Field::new("b", DataType::Utf8, true),
3504 ]);
3505 let a = Arc::new(StringArray::from(vec![
3506 Some("abc"),
3507 None,
3508 Some("abc"),
3509 None,
3510 Some("abc"),
3511 ])) as ArrayRef;
3512 let b = Arc::new(StringArray::from(vec![
3513 Some("^a"),
3514 Some("^A"),
3515 None,
3516 None,
3517 Some("^(b|c)"),
3518 ])) as ArrayRef;
3519
3520 let regex_expected =
3521 BooleanArray::from(vec![Some(true), None, None, None, Some(false)]);
3522 let regex_not_expected =
3523 BooleanArray::from(vec![Some(false), None, None, None, Some(true)]);
3524 apply_logic_op(
3525 &Arc::new(schema.clone()),
3526 &a,
3527 &b,
3528 Operator::RegexMatch,
3529 regex_expected.clone(),
3530 )?;
3531 apply_logic_op(
3532 &Arc::new(schema.clone()),
3533 &a,
3534 &b,
3535 Operator::RegexIMatch,
3536 regex_expected.clone(),
3537 )?;
3538 apply_logic_op(
3539 &Arc::new(schema.clone()),
3540 &a,
3541 &b,
3542 Operator::RegexNotMatch,
3543 regex_not_expected.clone(),
3544 )?;
3545 apply_logic_op(
3546 &Arc::new(schema),
3547 &a,
3548 &b,
3549 Operator::RegexNotIMatch,
3550 regex_not_expected.clone(),
3551 )?;
3552
3553 let schema = Schema::new(vec![
3554 Field::new("a", DataType::LargeUtf8, true),
3555 Field::new("b", DataType::LargeUtf8, true),
3556 ]);
3557 let a = Arc::new(LargeStringArray::from(vec![
3558 Some("abc"),
3559 None,
3560 Some("abc"),
3561 None,
3562 Some("abc"),
3563 ])) as ArrayRef;
3564 let b = Arc::new(LargeStringArray::from(vec![
3565 Some("^a"),
3566 Some("^A"),
3567 None,
3568 None,
3569 Some("^(b|c)"),
3570 ])) as ArrayRef;
3571
3572 apply_logic_op(
3573 &Arc::new(schema.clone()),
3574 &a,
3575 &b,
3576 Operator::RegexMatch,
3577 regex_expected.clone(),
3578 )?;
3579 apply_logic_op(
3580 &Arc::new(schema.clone()),
3581 &a,
3582 &b,
3583 Operator::RegexIMatch,
3584 regex_expected,
3585 )?;
3586 apply_logic_op(
3587 &Arc::new(schema.clone()),
3588 &a,
3589 &b,
3590 Operator::RegexNotMatch,
3591 regex_not_expected.clone(),
3592 )?;
3593 apply_logic_op(
3594 &Arc::new(schema),
3595 &a,
3596 &b,
3597 Operator::RegexNotIMatch,
3598 regex_not_expected,
3599 )?;
3600
3601 Ok(())
3602 }
3603
3604 #[test]
3605 fn regex_scalar_with_dictionary_nulls() -> Result<()> {
3606 let dictionary_values = Arc::new(StringArray::from(vec![
3607 Some("abc"),
3608 None,
3609 Some("ABC"),
3610 Some("def"),
3611 ]));
3612 let keys = UInt32Array::from(vec![Some(0), None, Some(1), Some(2), Some(3)]);
3613 let dictionary =
3614 Arc::new(DictionaryArray::try_new(keys, dictionary_values)?) as ArrayRef;
3615 let utf8 = cast(&dictionary, &DataType::Utf8)?;
3616 let pattern = ScalarValue::Utf8(Some("^abc$".to_string()));
3617 let dictionary_schema = Arc::new(Schema::new(vec![Field::new(
3618 "a",
3619 dictionary.data_type().clone(),
3620 true,
3621 )]));
3622 let utf8_schema =
3623 Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)]));
3624
3625 let evaluate =
3626 |schema: &SchemaRef, array: &ArrayRef, op: Operator| -> Result<ArrayRef> {
3627 let expr = binary(col("a", schema)?, op, lit(pattern.clone()), schema)?;
3628 let batch =
3629 RecordBatch::try_new(Arc::clone(schema), vec![Arc::clone(array)])?;
3630 Ok(expr
3631 .evaluate(&batch)?
3632 .into_array(batch.num_rows())
3633 .expect("Failed to convert to array"))
3634 };
3635
3636 for (op, expected) in [
3637 (
3638 Operator::RegexMatch,
3639 BooleanArray::from(vec![
3640 Some(true),
3641 None,
3642 None,
3643 Some(false),
3644 Some(false),
3645 ]),
3646 ),
3647 (
3648 Operator::RegexIMatch,
3649 BooleanArray::from(vec![Some(true), None, None, Some(true), Some(false)]),
3650 ),
3651 (
3652 Operator::RegexNotMatch,
3653 BooleanArray::from(vec![Some(false), None, None, Some(true), Some(true)]),
3654 ),
3655 (
3656 Operator::RegexNotIMatch,
3657 BooleanArray::from(vec![
3658 Some(false),
3659 None,
3660 None,
3661 Some(false),
3662 Some(true),
3663 ]),
3664 ),
3665 ] {
3666 let dictionary_result = evaluate(&dictionary_schema, &dictionary, op)?;
3667 let utf8_result = evaluate(&utf8_schema, &utf8, op)?;
3668
3669 assert_eq!(dictionary_result.as_ref(), &expected);
3670 assert_eq!(&dictionary_result, &utf8_result);
3671 }
3672
3673 Ok(())
3674 }
3675
3676 #[test]
3677 fn regex_mismatched_array_types_error() -> Result<()> {
3678 let schema = Schema::new(vec![
3683 Field::new("a", DataType::Utf8View, true),
3684 Field::new("b", DataType::Utf8, true),
3685 ]);
3686 let a = Arc::new(StringViewArray::from(vec!["user auth failed"])) as ArrayRef;
3687 let b = Arc::new(StringArray::from(vec!["(auth|login)"])) as ArrayRef;
3688
3689 let expr = binary(
3691 col("a", &schema)?,
3692 Operator::RegexMatch,
3693 col("b", &schema)?,
3694 &schema,
3695 )?;
3696 let batch = RecordBatch::try_new(Arc::new(schema), vec![a, b])?;
3697 let err = expr.evaluate(&batch).unwrap_err();
3698 assert_contains!(err.to_string(), "failed to downcast array");
3699
3700 Ok(())
3701 }
3702
3703 #[test]
3704 fn or_with_nulls_op() -> Result<()> {
3705 let schema = Schema::new(vec![
3706 Field::new("a", DataType::Boolean, true),
3707 Field::new("b", DataType::Boolean, true),
3708 ]);
3709 let a = Arc::new(BooleanArray::from(vec![
3710 Some(true),
3711 Some(false),
3712 None,
3713 Some(true),
3714 Some(false),
3715 None,
3716 Some(true),
3717 Some(false),
3718 None,
3719 ])) as ArrayRef;
3720 let b = Arc::new(BooleanArray::from(vec![
3721 Some(true),
3722 Some(true),
3723 Some(true),
3724 Some(false),
3725 Some(false),
3726 Some(false),
3727 None,
3728 None,
3729 None,
3730 ])) as ArrayRef;
3731
3732 let expected = BooleanArray::from(vec![
3733 Some(true),
3734 Some(true),
3735 Some(true),
3736 Some(true),
3737 Some(false),
3738 None,
3739 Some(true),
3740 None,
3741 None,
3742 ]);
3743 apply_logic_op(&Arc::new(schema), &a, &b, Operator::Or, expected)?;
3744
3745 Ok(())
3746 }
3747
3748 fn bool_test_arrays() -> (SchemaRef, ArrayRef, ArrayRef) {
3753 let schema = Schema::new(vec![
3754 Field::new("a", DataType::Boolean, true),
3755 Field::new("b", DataType::Boolean, true),
3756 ]);
3757 let a: BooleanArray = [
3758 Some(true),
3759 Some(true),
3760 Some(true),
3761 None,
3762 None,
3763 None,
3764 Some(false),
3765 Some(false),
3766 Some(false),
3767 ]
3768 .iter()
3769 .collect();
3770 let b: BooleanArray = [
3771 Some(true),
3772 None,
3773 Some(false),
3774 Some(true),
3775 None,
3776 Some(false),
3777 Some(true),
3778 None,
3779 Some(false),
3780 ]
3781 .iter()
3782 .collect();
3783 (Arc::new(schema), Arc::new(a), Arc::new(b))
3784 }
3785
3786 fn scalar_bool_test_array() -> (SchemaRef, ArrayRef) {
3788 let schema = Schema::new(vec![Field::new("a", DataType::Boolean, true)]);
3789 let a: BooleanArray = [Some(true), None, Some(false)].iter().collect();
3790 (Arc::new(schema), Arc::new(a))
3791 }
3792
3793 #[test]
3794 fn eq_op_bool() {
3795 let (schema, a, b) = bool_test_arrays();
3796 let expected = [
3797 Some(true),
3798 None,
3799 Some(false),
3800 None,
3801 None,
3802 None,
3803 Some(false),
3804 None,
3805 Some(true),
3806 ]
3807 .iter()
3808 .collect();
3809 apply_logic_op(&schema, &a, &b, Operator::Eq, expected).unwrap();
3810 }
3811
3812 #[test]
3813 fn eq_op_bool_scalar() {
3814 let (schema, a) = scalar_bool_test_array();
3815 let expected = [Some(true), None, Some(false)].iter().collect();
3816 apply_logic_op_scalar_arr(
3817 &schema,
3818 &ScalarValue::from(true),
3819 &a,
3820 Operator::Eq,
3821 &expected,
3822 )
3823 .unwrap();
3824 apply_logic_op_arr_scalar(
3825 &schema,
3826 &a,
3827 &ScalarValue::from(true),
3828 Operator::Eq,
3829 &expected,
3830 )
3831 .unwrap();
3832
3833 let expected = [Some(false), None, Some(true)].iter().collect();
3834 apply_logic_op_scalar_arr(
3835 &schema,
3836 &ScalarValue::from(false),
3837 &a,
3838 Operator::Eq,
3839 &expected,
3840 )
3841 .unwrap();
3842 apply_logic_op_arr_scalar(
3843 &schema,
3844 &a,
3845 &ScalarValue::from(false),
3846 Operator::Eq,
3847 &expected,
3848 )
3849 .unwrap();
3850 }
3851
3852 #[test]
3853 fn neq_op_bool() {
3854 let (schema, a, b) = bool_test_arrays();
3855 let expected = [
3856 Some(false),
3857 None,
3858 Some(true),
3859 None,
3860 None,
3861 None,
3862 Some(true),
3863 None,
3864 Some(false),
3865 ]
3866 .iter()
3867 .collect();
3868 apply_logic_op(&schema, &a, &b, Operator::NotEq, expected).unwrap();
3869 }
3870
3871 #[test]
3872 fn neq_op_bool_scalar() {
3873 let (schema, a) = scalar_bool_test_array();
3874 let expected = [Some(false), None, Some(true)].iter().collect();
3875 apply_logic_op_scalar_arr(
3876 &schema,
3877 &ScalarValue::from(true),
3878 &a,
3879 Operator::NotEq,
3880 &expected,
3881 )
3882 .unwrap();
3883 apply_logic_op_arr_scalar(
3884 &schema,
3885 &a,
3886 &ScalarValue::from(true),
3887 Operator::NotEq,
3888 &expected,
3889 )
3890 .unwrap();
3891
3892 let expected = [Some(true), None, Some(false)].iter().collect();
3893 apply_logic_op_scalar_arr(
3894 &schema,
3895 &ScalarValue::from(false),
3896 &a,
3897 Operator::NotEq,
3898 &expected,
3899 )
3900 .unwrap();
3901 apply_logic_op_arr_scalar(
3902 &schema,
3903 &a,
3904 &ScalarValue::from(false),
3905 Operator::NotEq,
3906 &expected,
3907 )
3908 .unwrap();
3909 }
3910
3911 #[test]
3912 fn lt_op_bool() {
3913 let (schema, a, b) = bool_test_arrays();
3914 let expected = [
3915 Some(false),
3916 None,
3917 Some(false),
3918 None,
3919 None,
3920 None,
3921 Some(true),
3922 None,
3923 Some(false),
3924 ]
3925 .iter()
3926 .collect();
3927 apply_logic_op(&schema, &a, &b, Operator::Lt, expected).unwrap();
3928 }
3929
3930 #[test]
3931 fn lt_op_bool_scalar() {
3932 let (schema, a) = scalar_bool_test_array();
3933 let expected = [Some(false), None, Some(false)].iter().collect();
3934 apply_logic_op_scalar_arr(
3935 &schema,
3936 &ScalarValue::from(true),
3937 &a,
3938 Operator::Lt,
3939 &expected,
3940 )
3941 .unwrap();
3942
3943 let expected = [Some(false), None, Some(true)].iter().collect();
3944 apply_logic_op_arr_scalar(
3945 &schema,
3946 &a,
3947 &ScalarValue::from(true),
3948 Operator::Lt,
3949 &expected,
3950 )
3951 .unwrap();
3952
3953 let expected = [Some(true), None, Some(false)].iter().collect();
3954 apply_logic_op_scalar_arr(
3955 &schema,
3956 &ScalarValue::from(false),
3957 &a,
3958 Operator::Lt,
3959 &expected,
3960 )
3961 .unwrap();
3962
3963 let expected = [Some(false), None, Some(false)].iter().collect();
3964 apply_logic_op_arr_scalar(
3965 &schema,
3966 &a,
3967 &ScalarValue::from(false),
3968 Operator::Lt,
3969 &expected,
3970 )
3971 .unwrap();
3972 }
3973
3974 #[test]
3975 fn lt_eq_op_bool() {
3976 let (schema, a, b) = bool_test_arrays();
3977 let expected = [
3978 Some(true),
3979 None,
3980 Some(false),
3981 None,
3982 None,
3983 None,
3984 Some(true),
3985 None,
3986 Some(true),
3987 ]
3988 .iter()
3989 .collect();
3990 apply_logic_op(&schema, &a, &b, Operator::LtEq, expected).unwrap();
3991 }
3992
3993 #[test]
3994 fn lt_eq_op_bool_scalar() {
3995 let (schema, a) = scalar_bool_test_array();
3996 let expected = [Some(true), None, Some(false)].iter().collect();
3997 apply_logic_op_scalar_arr(
3998 &schema,
3999 &ScalarValue::from(true),
4000 &a,
4001 Operator::LtEq,
4002 &expected,
4003 )
4004 .unwrap();
4005
4006 let expected = [Some(true), None, Some(true)].iter().collect();
4007 apply_logic_op_arr_scalar(
4008 &schema,
4009 &a,
4010 &ScalarValue::from(true),
4011 Operator::LtEq,
4012 &expected,
4013 )
4014 .unwrap();
4015
4016 let expected = [Some(true), None, Some(true)].iter().collect();
4017 apply_logic_op_scalar_arr(
4018 &schema,
4019 &ScalarValue::from(false),
4020 &a,
4021 Operator::LtEq,
4022 &expected,
4023 )
4024 .unwrap();
4025
4026 let expected = [Some(false), None, Some(true)].iter().collect();
4027 apply_logic_op_arr_scalar(
4028 &schema,
4029 &a,
4030 &ScalarValue::from(false),
4031 Operator::LtEq,
4032 &expected,
4033 )
4034 .unwrap();
4035 }
4036
4037 #[test]
4038 fn gt_op_bool() {
4039 let (schema, a, b) = bool_test_arrays();
4040 let expected = [
4041 Some(false),
4042 None,
4043 Some(true),
4044 None,
4045 None,
4046 None,
4047 Some(false),
4048 None,
4049 Some(false),
4050 ]
4051 .iter()
4052 .collect();
4053 apply_logic_op(&schema, &a, &b, Operator::Gt, expected).unwrap();
4054 }
4055
4056 #[test]
4057 fn gt_op_bool_scalar() {
4058 let (schema, a) = scalar_bool_test_array();
4059 let expected = [Some(false), None, Some(true)].iter().collect();
4060 apply_logic_op_scalar_arr(
4061 &schema,
4062 &ScalarValue::from(true),
4063 &a,
4064 Operator::Gt,
4065 &expected,
4066 )
4067 .unwrap();
4068
4069 let expected = [Some(false), None, Some(false)].iter().collect();
4070 apply_logic_op_arr_scalar(
4071 &schema,
4072 &a,
4073 &ScalarValue::from(true),
4074 Operator::Gt,
4075 &expected,
4076 )
4077 .unwrap();
4078
4079 let expected = [Some(false), None, Some(false)].iter().collect();
4080 apply_logic_op_scalar_arr(
4081 &schema,
4082 &ScalarValue::from(false),
4083 &a,
4084 Operator::Gt,
4085 &expected,
4086 )
4087 .unwrap();
4088
4089 let expected = [Some(true), None, Some(false)].iter().collect();
4090 apply_logic_op_arr_scalar(
4091 &schema,
4092 &a,
4093 &ScalarValue::from(false),
4094 Operator::Gt,
4095 &expected,
4096 )
4097 .unwrap();
4098 }
4099
4100 #[test]
4101 fn gt_eq_op_bool() {
4102 let (schema, a, b) = bool_test_arrays();
4103 let expected = [
4104 Some(true),
4105 None,
4106 Some(true),
4107 None,
4108 None,
4109 None,
4110 Some(false),
4111 None,
4112 Some(true),
4113 ]
4114 .iter()
4115 .collect();
4116 apply_logic_op(&schema, &a, &b, Operator::GtEq, expected).unwrap();
4117 }
4118
4119 #[test]
4120 fn gt_eq_op_bool_scalar() {
4121 let (schema, a) = scalar_bool_test_array();
4122 let expected = [Some(true), None, Some(true)].iter().collect();
4123 apply_logic_op_scalar_arr(
4124 &schema,
4125 &ScalarValue::from(true),
4126 &a,
4127 Operator::GtEq,
4128 &expected,
4129 )
4130 .unwrap();
4131
4132 let expected = [Some(true), None, Some(false)].iter().collect();
4133 apply_logic_op_arr_scalar(
4134 &schema,
4135 &a,
4136 &ScalarValue::from(true),
4137 Operator::GtEq,
4138 &expected,
4139 )
4140 .unwrap();
4141
4142 let expected = [Some(false), None, Some(true)].iter().collect();
4143 apply_logic_op_scalar_arr(
4144 &schema,
4145 &ScalarValue::from(false),
4146 &a,
4147 Operator::GtEq,
4148 &expected,
4149 )
4150 .unwrap();
4151
4152 let expected = [Some(true), None, Some(true)].iter().collect();
4153 apply_logic_op_arr_scalar(
4154 &schema,
4155 &a,
4156 &ScalarValue::from(false),
4157 Operator::GtEq,
4158 &expected,
4159 )
4160 .unwrap();
4161 }
4162
4163 #[test]
4164 fn is_distinct_from_op_bool() {
4165 let (schema, a, b) = bool_test_arrays();
4166 let expected = [
4167 Some(false),
4168 Some(true),
4169 Some(true),
4170 Some(true),
4171 Some(false),
4172 Some(true),
4173 Some(true),
4174 Some(true),
4175 Some(false),
4176 ]
4177 .iter()
4178 .collect();
4179 apply_logic_op(&schema, &a, &b, Operator::IsDistinctFrom, expected).unwrap();
4180 }
4181
4182 #[test]
4183 fn is_not_distinct_from_op_bool() {
4184 let (schema, a, b) = bool_test_arrays();
4185 let expected = [
4186 Some(true),
4187 Some(false),
4188 Some(false),
4189 Some(false),
4190 Some(true),
4191 Some(false),
4192 Some(false),
4193 Some(false),
4194 Some(true),
4195 ]
4196 .iter()
4197 .collect();
4198 apply_logic_op(&schema, &a, &b, Operator::IsNotDistinctFrom, expected).unwrap();
4199 }
4200
4201 #[test]
4202 fn relatively_deeply_nested() {
4203 let input: Vec<_> = vec![1, 2, 3, 4, 5].into_iter().map(Some).collect();
4209 let a: Int32Array = input.iter().collect();
4210
4211 let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(a) as _)]).unwrap();
4212 let schema = batch.schema();
4213
4214 let tree_depth: i32 = 100;
4216 let expr = (0..tree_depth)
4217 .map(|_| col("a", schema.as_ref()).unwrap())
4218 .reduce(|l, r| binary(l, Operator::Plus, r, &schema).unwrap())
4219 .unwrap();
4220
4221 let result = expr
4222 .evaluate(&batch)
4223 .expect("evaluation")
4224 .into_array(batch.num_rows())
4225 .expect("Failed to convert to array");
4226
4227 let expected: Int32Array = input
4228 .into_iter()
4229 .map(|i| i.map(|i| i * tree_depth))
4230 .collect();
4231 assert_eq!(result.as_ref(), &expected);
4232 }
4233
4234 fn create_decimal_array(
4235 array: &[Option<i128>],
4236 precision: u8,
4237 scale: i8,
4238 ) -> Decimal128Array {
4239 let mut decimal_builder = Decimal128Builder::with_capacity(array.len());
4240 for value in array.iter().copied() {
4241 decimal_builder.append_option(value)
4242 }
4243 decimal_builder
4244 .finish()
4245 .with_precision_and_scale(precision, scale)
4246 .unwrap()
4247 }
4248
4249 #[test]
4250 fn comparison_dict_decimal_scalar_expr_test() -> Result<()> {
4251 let value_i128 = 123;
4253 let decimal_scalar = ScalarValue::Dictionary(
4254 Box::new(DataType::Int8),
4255 Box::new(ScalarValue::Decimal128(Some(value_i128), 25, 3)),
4256 );
4257 let schema = Arc::new(Schema::new(vec![Field::new(
4258 "a",
4259 DataType::Dictionary(
4260 Box::new(DataType::Int8),
4261 Box::new(DataType::Decimal128(25, 3)),
4262 ),
4263 true,
4264 )]));
4265 let decimal_array = Arc::new(create_decimal_array(
4266 &[
4267 Some(value_i128),
4268 None,
4269 Some(value_i128 - 1),
4270 Some(value_i128 + 1),
4271 ],
4272 25,
4273 3,
4274 ));
4275
4276 let keys = Int8Array::from(vec![Some(0), None, Some(2), Some(3)]);
4277 let dictionary =
4278 Arc::new(DictionaryArray::try_new(keys, decimal_array)?) as ArrayRef;
4279
4280 apply_logic_op_arr_scalar(
4282 &schema,
4283 &dictionary,
4284 &decimal_scalar,
4285 Operator::Eq,
4286 &BooleanArray::from(vec![Some(true), None, Some(false), Some(false)]),
4287 )
4288 .unwrap();
4289 apply_logic_op_arr_scalar(
4291 &schema,
4292 &dictionary,
4293 &decimal_scalar,
4294 Operator::NotEq,
4295 &BooleanArray::from(vec![Some(false), None, Some(true), Some(true)]),
4296 )
4297 .unwrap();
4298 apply_logic_op_arr_scalar(
4300 &schema,
4301 &dictionary,
4302 &decimal_scalar,
4303 Operator::Lt,
4304 &BooleanArray::from(vec![Some(false), None, Some(true), Some(false)]),
4305 )
4306 .unwrap();
4307
4308 apply_logic_op_arr_scalar(
4310 &schema,
4311 &dictionary,
4312 &decimal_scalar,
4313 Operator::LtEq,
4314 &BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]),
4315 )
4316 .unwrap();
4317 apply_logic_op_arr_scalar(
4319 &schema,
4320 &dictionary,
4321 &decimal_scalar,
4322 Operator::Gt,
4323 &BooleanArray::from(vec![Some(false), None, Some(false), Some(true)]),
4324 )
4325 .unwrap();
4326
4327 apply_logic_op_arr_scalar(
4329 &schema,
4330 &dictionary,
4331 &decimal_scalar,
4332 Operator::GtEq,
4333 &BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]),
4334 )
4335 .unwrap();
4336
4337 Ok(())
4338 }
4339
4340 #[test]
4341 fn comparison_decimal_expr_test() -> Result<()> {
4342 let value_i128 = 123;
4344 let decimal_scalar = ScalarValue::Decimal128(Some(value_i128), 25, 3);
4345 let schema = Arc::new(Schema::new(vec![Field::new(
4346 "a",
4347 DataType::Decimal128(25, 3),
4348 true,
4349 )]));
4350 let decimal_array = Arc::new(create_decimal_array(
4351 &[
4352 Some(value_i128),
4353 None,
4354 Some(value_i128 - 1),
4355 Some(value_i128 + 1),
4356 ],
4357 25,
4358 3,
4359 )) as ArrayRef;
4360 apply_logic_op_arr_scalar(
4362 &schema,
4363 &decimal_array,
4364 &decimal_scalar,
4365 Operator::Eq,
4366 &BooleanArray::from(vec![Some(true), None, Some(false), Some(false)]),
4367 )
4368 .unwrap();
4369 apply_logic_op_arr_scalar(
4371 &schema,
4372 &decimal_array,
4373 &decimal_scalar,
4374 Operator::NotEq,
4375 &BooleanArray::from(vec![Some(false), None, Some(true), Some(true)]),
4376 )
4377 .unwrap();
4378 apply_logic_op_arr_scalar(
4380 &schema,
4381 &decimal_array,
4382 &decimal_scalar,
4383 Operator::Lt,
4384 &BooleanArray::from(vec![Some(false), None, Some(true), Some(false)]),
4385 )
4386 .unwrap();
4387
4388 apply_logic_op_arr_scalar(
4390 &schema,
4391 &decimal_array,
4392 &decimal_scalar,
4393 Operator::LtEq,
4394 &BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]),
4395 )
4396 .unwrap();
4397 apply_logic_op_arr_scalar(
4399 &schema,
4400 &decimal_array,
4401 &decimal_scalar,
4402 Operator::Gt,
4403 &BooleanArray::from(vec![Some(false), None, Some(false), Some(true)]),
4404 )
4405 .unwrap();
4406
4407 apply_logic_op_arr_scalar(
4409 &schema,
4410 &decimal_array,
4411 &decimal_scalar,
4412 Operator::GtEq,
4413 &BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]),
4414 )
4415 .unwrap();
4416
4417 let decimal_scalar = ScalarValue::Decimal128(Some(123_456), 10, 3);
4419 let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)]));
4420 apply_logic_op_scalar_arr(
4422 &schema,
4423 &decimal_scalar,
4424 &(Arc::new(Int64Array::from(vec![Some(124), None])) as ArrayRef),
4425 Operator::Eq,
4426 &BooleanArray::from(vec![Some(false), None]),
4427 )
4428 .unwrap();
4429
4430 apply_logic_op_arr_scalar(
4432 &schema,
4433 &(Arc::new(Int64Array::from(vec![Some(123), None, Some(1)])) as ArrayRef),
4434 &decimal_scalar,
4435 Operator::NotEq,
4436 &BooleanArray::from(vec![Some(true), None, Some(true)]),
4437 )
4438 .unwrap();
4439
4440 apply_logic_op_arr_scalar(
4442 &schema,
4443 &(Arc::new(Int64Array::from(vec![Some(123), None, Some(124)])) as ArrayRef),
4444 &decimal_scalar,
4445 Operator::Lt,
4446 &BooleanArray::from(vec![Some(true), None, Some(false)]),
4447 )
4448 .unwrap();
4449
4450 apply_logic_op_arr_scalar(
4452 &schema,
4453 &(Arc::new(Int64Array::from(vec![Some(123), None, Some(124)])) as ArrayRef),
4454 &decimal_scalar,
4455 Operator::Gt,
4456 &BooleanArray::from(vec![Some(false), None, Some(true)]),
4457 )
4458 .unwrap();
4459
4460 let schema =
4461 Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)]));
4462 apply_logic_op_arr_scalar(
4464 &schema,
4465 &(Arc::new(Float64Array::from(vec![Some(123.456), None, Some(123.457)]))
4466 as ArrayRef),
4467 &decimal_scalar,
4468 Operator::Eq,
4469 &BooleanArray::from(vec![Some(true), None, Some(false)]),
4470 )
4471 .unwrap();
4472
4473 apply_logic_op_arr_scalar(
4475 &schema,
4476 &(Arc::new(Float64Array::from(vec![
4477 Some(123.456),
4478 None,
4479 Some(123.457),
4480 Some(123.45),
4481 ])) as ArrayRef),
4482 &decimal_scalar,
4483 Operator::LtEq,
4484 &BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]),
4485 )
4486 .unwrap();
4487 apply_logic_op_arr_scalar(
4489 &schema,
4490 &(Arc::new(Float64Array::from(vec![
4491 Some(123.456),
4492 None,
4493 Some(123.457),
4494 Some(123.45),
4495 ])) as ArrayRef),
4496 &decimal_scalar,
4497 Operator::GtEq,
4498 &BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]),
4499 )
4500 .unwrap();
4501
4502 let value: i128 = 123;
4503 let decimal_array = Arc::new(create_decimal_array(
4504 &[Some(value), None, Some(value - 1), Some(value + 1)],
4505 10,
4506 0,
4507 )) as ArrayRef;
4508
4509 let schema = Arc::new(Schema::new(vec![
4511 Field::new("a", DataType::Decimal128(10, 0), true),
4512 Field::new("b", DataType::Decimal128(10, 0), true),
4513 ]));
4514 let right_decimal_array = Arc::new(create_decimal_array(
4515 &[
4516 Some(value - 1),
4517 Some(value),
4518 Some(value + 1),
4519 Some(value + 1),
4520 ],
4521 10,
4522 0,
4523 )) as ArrayRef;
4524
4525 apply_logic_op(
4526 &schema,
4527 &decimal_array,
4528 &right_decimal_array,
4529 Operator::Eq,
4530 BooleanArray::from(vec![Some(false), None, Some(false), Some(true)]),
4531 )
4532 .unwrap();
4533
4534 apply_logic_op(
4535 &schema,
4536 &decimal_array,
4537 &right_decimal_array,
4538 Operator::NotEq,
4539 BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]),
4540 )
4541 .unwrap();
4542
4543 apply_logic_op(
4544 &schema,
4545 &decimal_array,
4546 &right_decimal_array,
4547 Operator::Lt,
4548 BooleanArray::from(vec![Some(false), None, Some(true), Some(false)]),
4549 )
4550 .unwrap();
4551
4552 apply_logic_op(
4553 &schema,
4554 &decimal_array,
4555 &right_decimal_array,
4556 Operator::LtEq,
4557 BooleanArray::from(vec![Some(false), None, Some(true), Some(true)]),
4558 )
4559 .unwrap();
4560
4561 apply_logic_op(
4562 &schema,
4563 &decimal_array,
4564 &right_decimal_array,
4565 Operator::Gt,
4566 BooleanArray::from(vec![Some(true), None, Some(false), Some(false)]),
4567 )
4568 .unwrap();
4569
4570 apply_logic_op(
4571 &schema,
4572 &decimal_array,
4573 &right_decimal_array,
4574 Operator::GtEq,
4575 BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]),
4576 )
4577 .unwrap();
4578
4579 let value: i64 = 123;
4581 let schema = Arc::new(Schema::new(vec![
4582 Field::new("a", DataType::Int64, true),
4583 Field::new("b", DataType::Decimal128(10, 0), true),
4584 ]));
4585
4586 let int64_array = Arc::new(Int64Array::from(vec![
4587 Some(value),
4588 Some(value - 1),
4589 Some(value),
4590 Some(value + 1),
4591 ])) as ArrayRef;
4592
4593 apply_logic_op(
4595 &schema,
4596 &int64_array,
4597 &decimal_array,
4598 Operator::Eq,
4599 BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]),
4600 )
4601 .unwrap();
4602 apply_logic_op(
4604 &schema,
4605 &int64_array,
4606 &decimal_array,
4607 Operator::NotEq,
4608 BooleanArray::from(vec![Some(false), None, Some(true), Some(false)]),
4609 )
4610 .unwrap();
4611
4612 let schema = Arc::new(Schema::new(vec![
4613 Field::new("a", DataType::Float64, true),
4614 Field::new("b", DataType::Decimal128(10, 2), true),
4615 ]));
4616
4617 let value: i128 = 123;
4618 let decimal_array = Arc::new(create_decimal_array(
4619 &[
4620 Some(value), None,
4622 Some(value - 1), Some(value + 1), ],
4625 10,
4626 2,
4627 )) as ArrayRef;
4628 let float64_array = Arc::new(Float64Array::from(vec![
4629 Some(1.23),
4630 Some(1.22),
4631 Some(1.23),
4632 Some(1.24),
4633 ])) as ArrayRef;
4634 apply_logic_op(
4636 &schema,
4637 &float64_array,
4638 &decimal_array,
4639 Operator::Lt,
4640 BooleanArray::from(vec![Some(false), None, Some(false), Some(false)]),
4641 )
4642 .unwrap();
4643 apply_logic_op(
4645 &schema,
4646 &float64_array,
4647 &decimal_array,
4648 Operator::LtEq,
4649 BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]),
4650 )
4651 .unwrap();
4652 apply_logic_op(
4654 &schema,
4655 &float64_array,
4656 &decimal_array,
4657 Operator::Gt,
4658 BooleanArray::from(vec![Some(false), None, Some(true), Some(false)]),
4659 )
4660 .unwrap();
4661 apply_logic_op(
4662 &schema,
4663 &float64_array,
4664 &decimal_array,
4665 Operator::GtEq,
4666 BooleanArray::from(vec![Some(true), None, Some(true), Some(true)]),
4667 )
4668 .unwrap();
4669 apply_logic_op(
4674 &schema,
4675 &float64_array,
4676 &decimal_array,
4677 Operator::IsDistinctFrom,
4678 BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]),
4679 )
4680 .unwrap();
4681 apply_logic_op(
4683 &schema,
4684 &float64_array,
4685 &decimal_array,
4686 Operator::IsNotDistinctFrom,
4687 BooleanArray::from(vec![Some(true), Some(false), Some(false), Some(true)]),
4688 )
4689 .unwrap();
4690
4691 Ok(())
4692 }
4693
4694 fn apply_decimal_arithmetic_op(
4695 schema: &SchemaRef,
4696 left: &ArrayRef,
4697 right: &ArrayRef,
4698 op: Operator,
4699 expected: ArrayRef,
4700 ) -> Result<()> {
4701 let arithmetic_op = binary_op(col("a", schema)?, op, col("b", schema)?, schema)?;
4702 let data: Vec<ArrayRef> = vec![Arc::clone(left), Arc::clone(right)];
4703 let batch = RecordBatch::try_new(Arc::clone(schema), data)?;
4704 let result = arithmetic_op
4705 .evaluate(&batch)?
4706 .into_array(batch.num_rows())
4707 .expect("Failed to convert to array");
4708
4709 assert_eq!(result.as_ref(), expected.as_ref());
4710 Ok(())
4711 }
4712
4713 #[test]
4714 fn arithmetic_decimal_expr_test() -> Result<()> {
4715 let schema = Arc::new(Schema::new(vec![
4716 Field::new("a", DataType::Int32, true),
4717 Field::new("b", DataType::Decimal128(10, 2), true),
4718 ]));
4719 let value: i128 = 123;
4720 let decimal_array = Arc::new(create_decimal_array(
4721 &[
4722 Some(value), None,
4724 Some(value - 1), Some(value + 1), ],
4727 10,
4728 2,
4729 )) as ArrayRef;
4730 let int32_array = Arc::new(Int32Array::from(vec![
4731 Some(123),
4732 Some(122),
4733 Some(123),
4734 Some(124),
4735 ])) as ArrayRef;
4736
4737 let expect = Arc::new(create_decimal_array(
4739 &[Some(12423), None, Some(12422), Some(12524)],
4740 13,
4741 2,
4742 )) as ArrayRef;
4743 apply_decimal_arithmetic_op(
4744 &schema,
4745 &int32_array,
4746 &decimal_array,
4747 Operator::Plus,
4748 expect,
4749 )
4750 .unwrap();
4751
4752 let schema = Arc::new(Schema::new(vec![
4754 Field::new("a", DataType::Decimal128(10, 2), true),
4755 Field::new("b", DataType::Int32, true),
4756 ]));
4757 let expect = Arc::new(create_decimal_array(
4758 &[Some(-12177), None, Some(-12178), Some(-12276)],
4759 13,
4760 2,
4761 )) as ArrayRef;
4762 apply_decimal_arithmetic_op(
4763 &schema,
4764 &decimal_array,
4765 &int32_array,
4766 Operator::Minus,
4767 expect,
4768 )
4769 .unwrap();
4770
4771 let expect = Arc::new(create_decimal_array(
4773 &[Some(15129), None, Some(15006), Some(15376)],
4774 21,
4775 2,
4776 )) as ArrayRef;
4777 apply_decimal_arithmetic_op(
4778 &schema,
4779 &decimal_array,
4780 &int32_array,
4781 Operator::Multiply,
4782 expect,
4783 )
4784 .unwrap();
4785
4786 let schema = Arc::new(Schema::new(vec![
4788 Field::new("a", DataType::Int32, true),
4789 Field::new("b", DataType::Decimal128(10, 2), true),
4790 ]));
4791 let expect = Arc::new(create_decimal_array(
4792 &[Some(1000000), None, Some(1008196), Some(1000000)],
4793 16,
4794 4,
4795 )) as ArrayRef;
4796 apply_decimal_arithmetic_op(
4797 &schema,
4798 &int32_array,
4799 &decimal_array,
4800 Operator::Divide,
4801 expect,
4802 )
4803 .unwrap();
4804
4805 let schema = Arc::new(Schema::new(vec![
4807 Field::new("a", DataType::Int32, true),
4808 Field::new("b", DataType::Decimal128(10, 2), true),
4809 ]));
4810 let expect = Arc::new(create_decimal_array(
4811 &[Some(000), None, Some(100), Some(000)],
4812 10,
4813 2,
4814 )) as ArrayRef;
4815 apply_decimal_arithmetic_op(
4816 &schema,
4817 &int32_array,
4818 &decimal_array,
4819 Operator::Modulo,
4820 expect,
4821 )
4822 .unwrap();
4823
4824 Ok(())
4825 }
4826
4827 #[test]
4828 fn arithmetic_decimal_float_expr_test() -> Result<()> {
4829 let schema = Arc::new(Schema::new(vec![
4830 Field::new("a", DataType::Float64, true),
4831 Field::new("b", DataType::Decimal128(10, 2), true),
4832 ]));
4833 let value: i128 = 123;
4834 let decimal_array = Arc::new(create_decimal_array(
4835 &[
4836 Some(value), None,
4838 Some(value - 1), Some(value + 1), ],
4841 10,
4842 2,
4843 )) as ArrayRef;
4844 let float64_array = Arc::new(Float64Array::from(vec![
4845 Some(123.0),
4846 Some(122.0),
4847 Some(123.0),
4848 Some(124.0),
4849 ])) as ArrayRef;
4850
4851 let expect = Arc::new(Float64Array::from(vec![
4853 Some(124.23),
4854 None,
4855 Some(124.22),
4856 Some(125.24),
4857 ])) as ArrayRef;
4858 apply_decimal_arithmetic_op(
4859 &schema,
4860 &float64_array,
4861 &decimal_array,
4862 Operator::Plus,
4863 expect,
4864 )
4865 .unwrap();
4866
4867 let schema = Arc::new(Schema::new(vec![
4869 Field::new("a", DataType::Float64, true),
4870 Field::new("b", DataType::Decimal128(10, 2), true),
4871 ]));
4872 let expect = Arc::new(Float64Array::from(vec![
4873 Some(121.77),
4874 None,
4875 Some(121.78),
4876 Some(122.76),
4877 ])) as ArrayRef;
4878 apply_decimal_arithmetic_op(
4879 &schema,
4880 &float64_array,
4881 &decimal_array,
4882 Operator::Minus,
4883 expect,
4884 )
4885 .unwrap();
4886
4887 let expect = Arc::new(Float64Array::from(vec![
4889 Some(151.29),
4890 None,
4891 Some(150.06),
4892 Some(153.76),
4893 ])) as ArrayRef;
4894 apply_decimal_arithmetic_op(
4895 &schema,
4896 &float64_array,
4897 &decimal_array,
4898 Operator::Multiply,
4899 expect,
4900 )
4901 .unwrap();
4902
4903 let schema = Arc::new(Schema::new(vec![
4905 Field::new("a", DataType::Float64, true),
4906 Field::new("b", DataType::Decimal128(10, 2), true),
4907 ]));
4908 let expect = Arc::new(Float64Array::from(vec![
4909 Some(100.0),
4910 None,
4911 Some(100.81967213114754),
4912 Some(100.0),
4913 ])) as ArrayRef;
4914 apply_decimal_arithmetic_op(
4915 &schema,
4916 &float64_array,
4917 &decimal_array,
4918 Operator::Divide,
4919 expect,
4920 )
4921 .unwrap();
4922
4923 let schema = Arc::new(Schema::new(vec![
4925 Field::new("a", DataType::Float64, true),
4926 Field::new("b", DataType::Decimal128(10, 2), true),
4927 ]));
4928 let expect = Arc::new(Float64Array::from(vec![
4929 Some(1.7763568394002505e-15),
4930 None,
4931 Some(1.0000000000000027),
4932 Some(8.881784197001252e-16),
4933 ])) as ArrayRef;
4934 apply_decimal_arithmetic_op(
4935 &schema,
4936 &float64_array,
4937 &decimal_array,
4938 Operator::Modulo,
4939 expect,
4940 )
4941 .unwrap();
4942
4943 Ok(())
4944 }
4945
4946 #[test]
4947 fn arithmetic_divide_zero() -> Result<()> {
4948 let schema = Arc::new(Schema::new(vec![
4950 Field::new("a", DataType::Int32, true),
4951 Field::new("b", DataType::Int32, true),
4952 ]));
4953 let a = Arc::new(Int32Array::from(vec![100]));
4954 let b = Arc::new(Int32Array::from(vec![0]));
4955
4956 let err = apply_arithmetic::<Int32Type>(
4957 schema,
4958 vec![a, b],
4959 Operator::Divide,
4960 Int32Array::from(vec![Some(4), Some(8), Some(16), Some(32), Some(64)]),
4961 )
4962 .unwrap_err();
4963
4964 let _expected = plan_datafusion_err!("Divide by zero");
4965
4966 assert!(matches!(err, ref _expected), "{err}");
4967
4968 let schema = Arc::new(Schema::new(vec![
4970 Field::new("a", DataType::Decimal128(25, 3), true),
4971 Field::new("b", DataType::Decimal128(25, 3), true),
4972 ]));
4973 let left_decimal_array = Arc::new(create_decimal_array(&[Some(1234567)], 25, 3));
4974 let right_decimal_array = Arc::new(create_decimal_array(&[Some(0)], 25, 3));
4975
4976 let err = apply_arithmetic::<Decimal128Type>(
4977 schema,
4978 vec![left_decimal_array, right_decimal_array],
4979 Operator::Divide,
4980 create_decimal_array(
4981 &[Some(12345670000000000000000000000000000), None],
4982 38,
4983 29,
4984 ),
4985 )
4986 .unwrap_err();
4987
4988 assert!(matches!(err, ref _expected), "{err}");
4989
4990 Ok(())
4991 }
4992
4993 #[test]
4994 fn bitwise_array_test() -> Result<()> {
4995 let left = Arc::new(Int32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
4996 let right =
4997 Arc::new(Int32Array::from(vec![Some(1), Some(3), Some(7)])) as ArrayRef;
4998 let mut result = bitwise_and_dyn(Arc::clone(&left), Arc::clone(&right))?;
4999 let expected = Int32Array::from(vec![Some(0), None, Some(3)]);
5000 assert_eq!(result.as_ref(), &expected);
5001
5002 result = bitwise_or_dyn(Arc::clone(&left), Arc::clone(&right))?;
5003 let expected = Int32Array::from(vec![Some(13), None, Some(15)]);
5004 assert_eq!(result.as_ref(), &expected);
5005
5006 result = bitwise_xor_dyn(Arc::clone(&left), Arc::clone(&right))?;
5007 let expected = Int32Array::from(vec![Some(13), None, Some(12)]);
5008 assert_eq!(result.as_ref(), &expected);
5009
5010 let left =
5011 Arc::new(UInt32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
5012 let right =
5013 Arc::new(UInt32Array::from(vec![Some(1), Some(3), Some(7)])) as ArrayRef;
5014 let mut result = bitwise_and_dyn(Arc::clone(&left), Arc::clone(&right))?;
5015 let expected = UInt32Array::from(vec![Some(0), None, Some(3)]);
5016 assert_eq!(result.as_ref(), &expected);
5017
5018 result = bitwise_or_dyn(Arc::clone(&left), Arc::clone(&right))?;
5019 let expected = UInt32Array::from(vec![Some(13), None, Some(15)]);
5020 assert_eq!(result.as_ref(), &expected);
5021
5022 result = bitwise_xor_dyn(Arc::clone(&left), Arc::clone(&right))?;
5023 let expected = UInt32Array::from(vec![Some(13), None, Some(12)]);
5024 assert_eq!(result.as_ref(), &expected);
5025
5026 Ok(())
5027 }
5028
5029 #[test]
5030 fn bitwise_shift_array_test() -> Result<()> {
5031 let input = Arc::new(Int32Array::from(vec![Some(2), None, Some(10)])) as ArrayRef;
5032 let modules =
5033 Arc::new(Int32Array::from(vec![Some(2), Some(4), Some(8)])) as ArrayRef;
5034 let mut result =
5035 bitwise_shift_left_dyn(Arc::clone(&input), Arc::clone(&modules))?;
5036
5037 let expected = Int32Array::from(vec![Some(8), None, Some(2560)]);
5038 assert_eq!(result.as_ref(), &expected);
5039
5040 result = bitwise_shift_right_dyn(Arc::clone(&result), Arc::clone(&modules))?;
5041 assert_eq!(result.as_ref(), &input);
5042
5043 let input =
5044 Arc::new(UInt32Array::from(vec![Some(2), None, Some(10)])) as ArrayRef;
5045 let modules =
5046 Arc::new(UInt32Array::from(vec![Some(2), Some(4), Some(8)])) as ArrayRef;
5047 let mut result =
5048 bitwise_shift_left_dyn(Arc::clone(&input), Arc::clone(&modules))?;
5049
5050 let expected = UInt32Array::from(vec![Some(8), None, Some(2560)]);
5051 assert_eq!(result.as_ref(), &expected);
5052
5053 result = bitwise_shift_right_dyn(Arc::clone(&result), Arc::clone(&modules))?;
5054 assert_eq!(result.as_ref(), &input);
5055 Ok(())
5056 }
5057
5058 #[test]
5059 fn bitwise_shift_array_overflow_test() -> Result<()> {
5060 let input = Arc::new(Int32Array::from(vec![Some(2)])) as ArrayRef;
5061 let modules = Arc::new(Int32Array::from(vec![Some(100)])) as ArrayRef;
5062 let result = bitwise_shift_left_dyn(Arc::clone(&input), Arc::clone(&modules))?;
5063
5064 let expected = Int32Array::from(vec![Some(32)]);
5065 assert_eq!(result.as_ref(), &expected);
5066
5067 let input = Arc::new(UInt32Array::from(vec![Some(2)])) as ArrayRef;
5068 let modules = Arc::new(UInt32Array::from(vec![Some(100)])) as ArrayRef;
5069 let result = bitwise_shift_left_dyn(Arc::clone(&input), Arc::clone(&modules))?;
5070
5071 let expected = UInt32Array::from(vec![Some(32)]);
5072 assert_eq!(result.as_ref(), &expected);
5073 Ok(())
5074 }
5075
5076 #[test]
5077 fn bitwise_scalar_test() -> Result<()> {
5078 let left = Arc::new(Int32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
5079 let right = ScalarValue::from(3i32);
5080 let mut result = bitwise_and_dyn_scalar(&left, right.clone()).unwrap()?;
5081 let expected = Int32Array::from(vec![Some(0), None, Some(3)]);
5082 assert_eq!(result.as_ref(), &expected);
5083
5084 result = bitwise_or_dyn_scalar(&left, right.clone()).unwrap()?;
5085 let expected = Int32Array::from(vec![Some(15), None, Some(11)]);
5086 assert_eq!(result.as_ref(), &expected);
5087
5088 result = bitwise_xor_dyn_scalar(&left, right).unwrap()?;
5089 let expected = Int32Array::from(vec![Some(15), None, Some(8)]);
5090 assert_eq!(result.as_ref(), &expected);
5091
5092 let left =
5093 Arc::new(UInt32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
5094 let right = ScalarValue::from(3u32);
5095 let mut result = bitwise_and_dyn_scalar(&left, right.clone()).unwrap()?;
5096 let expected = UInt32Array::from(vec![Some(0), None, Some(3)]);
5097 assert_eq!(result.as_ref(), &expected);
5098
5099 result = bitwise_or_dyn_scalar(&left, right.clone()).unwrap()?;
5100 let expected = UInt32Array::from(vec![Some(15), None, Some(11)]);
5101 assert_eq!(result.as_ref(), &expected);
5102
5103 result = bitwise_xor_dyn_scalar(&left, right).unwrap()?;
5104 let expected = UInt32Array::from(vec![Some(15), None, Some(8)]);
5105 assert_eq!(result.as_ref(), &expected);
5106 Ok(())
5107 }
5108
5109 #[test]
5110 fn bitwise_shift_scalar_test() -> Result<()> {
5111 let input = Arc::new(Int32Array::from(vec![Some(2), None, Some(4)])) as ArrayRef;
5112 let module = ScalarValue::from(10i32);
5113 let mut result =
5114 bitwise_shift_left_dyn_scalar(&input, module.clone()).unwrap()?;
5115
5116 let expected = Int32Array::from(vec![Some(2048), None, Some(4096)]);
5117 assert_eq!(result.as_ref(), &expected);
5118
5119 result = bitwise_shift_right_dyn_scalar(&result, module).unwrap()?;
5120 assert_eq!(result.as_ref(), &input);
5121
5122 let input = Arc::new(UInt32Array::from(vec![Some(2), None, Some(4)])) as ArrayRef;
5123 let module = ScalarValue::from(10u32);
5124 let mut result =
5125 bitwise_shift_left_dyn_scalar(&input, module.clone()).unwrap()?;
5126
5127 let expected = UInt32Array::from(vec![Some(2048), None, Some(4096)]);
5128 assert_eq!(result.as_ref(), &expected);
5129
5130 result = bitwise_shift_right_dyn_scalar(&result, module).unwrap()?;
5131 assert_eq!(result.as_ref(), &input);
5132 Ok(())
5133 }
5134
5135 #[test]
5136 fn test_display_and_or_combo() {
5137 let expr = BinaryExpr::new(
5138 Arc::new(BinaryExpr::new(
5139 lit(ScalarValue::from(1)),
5140 Operator::And,
5141 lit(ScalarValue::from(2)),
5142 )),
5143 Operator::And,
5144 Arc::new(BinaryExpr::new(
5145 lit(ScalarValue::from(3)),
5146 Operator::And,
5147 lit(ScalarValue::from(4)),
5148 )),
5149 );
5150 assert_eq!(expr.to_string(), "1 AND 2 AND 3 AND 4");
5151
5152 let expr = BinaryExpr::new(
5153 Arc::new(BinaryExpr::new(
5154 lit(ScalarValue::from(1)),
5155 Operator::Or,
5156 lit(ScalarValue::from(2)),
5157 )),
5158 Operator::Or,
5159 Arc::new(BinaryExpr::new(
5160 lit(ScalarValue::from(3)),
5161 Operator::Or,
5162 lit(ScalarValue::from(4)),
5163 )),
5164 );
5165 assert_eq!(expr.to_string(), "1 OR 2 OR 3 OR 4");
5166
5167 let expr = BinaryExpr::new(
5168 Arc::new(BinaryExpr::new(
5169 lit(ScalarValue::from(1)),
5170 Operator::And,
5171 lit(ScalarValue::from(2)),
5172 )),
5173 Operator::Or,
5174 Arc::new(BinaryExpr::new(
5175 lit(ScalarValue::from(3)),
5176 Operator::And,
5177 lit(ScalarValue::from(4)),
5178 )),
5179 );
5180 assert_eq!(expr.to_string(), "1 AND 2 OR 3 AND 4");
5181
5182 let expr = BinaryExpr::new(
5183 Arc::new(BinaryExpr::new(
5184 lit(ScalarValue::from(1)),
5185 Operator::Or,
5186 lit(ScalarValue::from(2)),
5187 )),
5188 Operator::And,
5189 Arc::new(BinaryExpr::new(
5190 lit(ScalarValue::from(3)),
5191 Operator::Or,
5192 lit(ScalarValue::from(4)),
5193 )),
5194 );
5195 assert_eq!(expr.to_string(), "(1 OR 2) AND (3 OR 4)");
5196 }
5197
5198 #[test]
5199 fn test_to_result_type_array() {
5200 let values = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
5201 let keys = Int8Array::from(vec![Some(0), None, Some(2), Some(3)]);
5202 let dictionary =
5203 Arc::new(DictionaryArray::try_new(keys, values).unwrap()) as ArrayRef;
5204
5205 let casted = to_result_type_array(
5207 &Operator::Plus,
5208 Arc::clone(&dictionary),
5209 &DataType::Int32,
5210 )
5211 .unwrap();
5212 assert_eq!(
5213 &casted,
5214 &(Arc::new(Int32Array::from(vec![Some(1), None, Some(3), Some(4)]))
5215 as ArrayRef)
5216 );
5217
5218 let casted = to_result_type_array(
5220 &Operator::Plus,
5221 Arc::clone(&dictionary),
5222 dictionary.data_type(),
5223 )
5224 .unwrap();
5225 assert_eq!(&casted, &dictionary);
5226
5227 let casted = to_result_type_array(
5229 &Operator::Eq,
5230 Arc::clone(&dictionary),
5231 &DataType::Int32,
5232 )
5233 .unwrap();
5234 assert_eq!(&casted, &dictionary);
5235 }
5236
5237 #[test]
5238 fn test_add_with_overflow() -> Result<()> {
5239 let l = Arc::new(Int32Array::from(vec![1, i32::MAX]));
5241 let r = Arc::new(Int32Array::from(vec![2, 1]));
5242 let schema = Arc::new(Schema::new(vec![
5243 Field::new("l", DataType::Int32, false),
5244 Field::new("r", DataType::Int32, false),
5245 ]));
5246 let batch = RecordBatch::try_new(schema, vec![l, r])?;
5247
5248 let expr = BinaryExpr::new(
5250 Arc::new(Column::new("l", 0)),
5251 Operator::Plus,
5252 Arc::new(Column::new("r", 1)),
5253 )
5254 .with_fail_on_overflow(true);
5255
5256 let result = expr.evaluate(&batch);
5258 assert!(
5259 result
5260 .err()
5261 .unwrap()
5262 .to_string()
5263 .contains("Overflow happened on: 2147483647 + 1")
5264 );
5265 Ok(())
5266 }
5267
5268 #[test]
5269 fn test_subtract_with_overflow() -> Result<()> {
5270 let l = Arc::new(Int32Array::from(vec![1, i32::MIN]));
5272 let r = Arc::new(Int32Array::from(vec![2, 1]));
5273 let schema = Arc::new(Schema::new(vec![
5274 Field::new("l", DataType::Int32, false),
5275 Field::new("r", DataType::Int32, false),
5276 ]));
5277 let batch = RecordBatch::try_new(schema, vec![l, r])?;
5278
5279 let expr = BinaryExpr::new(
5281 Arc::new(Column::new("l", 0)),
5282 Operator::Minus,
5283 Arc::new(Column::new("r", 1)),
5284 )
5285 .with_fail_on_overflow(true);
5286
5287 let result = expr.evaluate(&batch);
5289 assert!(
5290 result
5291 .err()
5292 .unwrap()
5293 .to_string()
5294 .contains("Overflow happened on: -2147483648 - 1")
5295 );
5296 Ok(())
5297 }
5298
5299 #[test]
5300 fn test_mul_with_overflow() -> Result<()> {
5301 let l = Arc::new(Int32Array::from(vec![1, i32::MAX]));
5303 let r = Arc::new(Int32Array::from(vec![2, 2]));
5304 let schema = Arc::new(Schema::new(vec![
5305 Field::new("l", DataType::Int32, false),
5306 Field::new("r", DataType::Int32, false),
5307 ]));
5308 let batch = RecordBatch::try_new(schema, vec![l, r])?;
5309
5310 let expr = BinaryExpr::new(
5312 Arc::new(Column::new("l", 0)),
5313 Operator::Multiply,
5314 Arc::new(Column::new("r", 1)),
5315 )
5316 .with_fail_on_overflow(true);
5317
5318 let result = expr.evaluate(&batch);
5320 assert!(
5321 result
5322 .err()
5323 .unwrap()
5324 .to_string()
5325 .contains("Overflow happened on: 2147483647 * 2")
5326 );
5327 Ok(())
5328 }
5329
5330 fn apply_similar_to(
5332 schema: &SchemaRef,
5333 va: Vec<&str>,
5334 vb: Vec<&str>,
5335 negated: bool,
5336 case_insensitive: bool,
5337 expected: &BooleanArray,
5338 ) -> Result<()> {
5339 let a = StringArray::from(va);
5340 let b = StringArray::from(vb);
5341 let op = similar_to(
5342 negated,
5343 case_insensitive,
5344 col("a", schema)?,
5345 col("b", schema)?,
5346 )?;
5347 let batch =
5348 RecordBatch::try_new(Arc::clone(schema), vec![Arc::new(a), Arc::new(b)])?;
5349 let result = op
5350 .evaluate(&batch)?
5351 .into_array(batch.num_rows())
5352 .expect("Failed to convert to array");
5353 assert_eq!(result.as_ref(), expected);
5354
5355 Ok(())
5356 }
5357
5358 #[test]
5359 fn test_similar_to() {
5360 let schema = Arc::new(Schema::new(vec![
5361 Field::new("a", DataType::Utf8, false),
5362 Field::new("b", DataType::Utf8, false),
5363 ]));
5364
5365 let expected = [Some(true), Some(false)].iter().collect();
5366 apply_similar_to(
5368 &schema,
5369 vec!["hello world", "Hello World"],
5370 vec!["hello.*", "hello.*"],
5371 false,
5372 false,
5373 &expected,
5374 )
5375 .unwrap();
5376 apply_similar_to(
5378 &schema,
5379 vec!["hello world", "bye"],
5380 vec!["hello.*", "hello.*"],
5381 false,
5382 true,
5383 &expected,
5384 )
5385 .unwrap();
5386 }
5387
5388 pub fn binary_expr(
5389 left: Arc<dyn PhysicalExpr>,
5390 op: Operator,
5391 right: Arc<dyn PhysicalExpr>,
5392 schema: &Schema,
5393 ) -> Result<BinaryExpr> {
5394 Ok(binary_op(left, op, right, schema)?
5395 .downcast_ref::<BinaryExpr>()
5396 .unwrap()
5397 .clone())
5398 }
5399
5400 #[test]
5402 #[expect(deprecated)]
5403 fn test_evaluate_statistics_combination_of_range_holders() -> Result<()> {
5404 let schema = &Schema::new(vec![Field::new("a", DataType::Float64, false)]);
5405 let a = Arc::new(Column::new("a", 0)) as _;
5406 let b = lit(ScalarValue::from(12.0));
5407
5408 let left_interval = Interval::make(Some(0.0), Some(12.0))?;
5409 let right_interval = Interval::make(Some(12.0), Some(36.0))?;
5410 let (left_mean, right_mean) = (ScalarValue::from(6.0), ScalarValue::from(24.0));
5411 let (left_med, right_med) = (ScalarValue::from(6.0), ScalarValue::from(24.0));
5412
5413 for children in [
5414 vec![
5415 &Distribution::new_uniform(left_interval.clone())?,
5416 &Distribution::new_uniform(right_interval.clone())?,
5417 ],
5418 vec![
5419 &Distribution::new_generic(
5420 left_mean.clone(),
5421 left_med.clone(),
5422 ScalarValue::Float64(None),
5423 left_interval.clone(),
5424 )?,
5425 &Distribution::new_uniform(right_interval.clone())?,
5426 ],
5427 vec![
5428 &Distribution::new_uniform(right_interval.clone())?,
5429 &Distribution::new_generic(
5430 right_mean.clone(),
5431 right_med.clone(),
5432 ScalarValue::Float64(None),
5433 right_interval.clone(),
5434 )?,
5435 ],
5436 vec![
5437 &Distribution::new_generic(
5438 left_mean.clone(),
5439 left_med.clone(),
5440 ScalarValue::Float64(None),
5441 left_interval.clone(),
5442 )?,
5443 &Distribution::new_generic(
5444 right_mean.clone(),
5445 right_med.clone(),
5446 ScalarValue::Float64(None),
5447 right_interval.clone(),
5448 )?,
5449 ],
5450 ] {
5451 let ops = vec![
5452 Operator::Plus,
5453 Operator::Minus,
5454 Operator::Multiply,
5455 Operator::Divide,
5456 ];
5457
5458 for op in ops {
5459 let expr = binary_expr(Arc::clone(&a), op, Arc::clone(&b), schema)?;
5460 assert_eq!(
5461 expr.evaluate_statistics(&children)?,
5462 new_generic_from_binary_op(&op, children[0], children[1])?
5463 );
5464 }
5465 }
5466 Ok(())
5467 }
5468
5469 #[test]
5470 #[expect(deprecated)]
5471 fn test_evaluate_statistics_bernoulli() -> Result<()> {
5472 let schema = &Schema::new(vec![
5473 Field::new("a", DataType::Int64, false),
5474 Field::new("b", DataType::Int64, false),
5475 ]);
5476 let a = Arc::new(Column::new("a", 0)) as _;
5477 let b = Arc::new(Column::new("b", 1)) as _;
5478 let eq = Arc::new(binary_expr(
5479 Arc::clone(&a),
5480 Operator::Eq,
5481 Arc::clone(&b),
5482 schema,
5483 )?);
5484 let neq = Arc::new(binary_expr(a, Operator::NotEq, b, schema)?);
5485
5486 let left_stat = &Distribution::new_uniform(Interval::make(Some(0), Some(7))?)?;
5487 let right_stat = &Distribution::new_uniform(Interval::make(Some(4), Some(11))?)?;
5488
5489 assert_eq!(
5492 eq.evaluate_statistics(&[left_stat, right_stat])?,
5493 Distribution::new_bernoulli(ScalarValue::from(1.0 / 16.0))?
5494 );
5495
5496 assert_eq!(
5498 neq.evaluate_statistics(&[left_stat, right_stat])?,
5499 Distribution::new_bernoulli(ScalarValue::from(15.0 / 16.0))?
5500 );
5501
5502 Ok(())
5503 }
5504
5505 #[test]
5506 #[expect(deprecated)]
5507 fn test_propagate_statistics_combination_of_range_holders_arithmetic() -> Result<()> {
5508 let schema = &Schema::new(vec![Field::new("a", DataType::Float64, false)]);
5509 let a = Arc::new(Column::new("a", 0)) as _;
5510 let b = lit(ScalarValue::from(12.0));
5511
5512 let left_interval = Interval::make(Some(0.0), Some(12.0))?;
5513 let right_interval = Interval::make(Some(12.0), Some(36.0))?;
5514
5515 let parent = Distribution::new_uniform(Interval::make(Some(-432.), Some(432.))?)?;
5516 let children = vec![
5517 vec![
5518 Distribution::new_uniform(left_interval.clone())?,
5519 Distribution::new_uniform(right_interval.clone())?,
5520 ],
5521 vec![
5522 Distribution::new_generic(
5523 ScalarValue::from(6.),
5524 ScalarValue::from(6.),
5525 ScalarValue::Float64(None),
5526 left_interval.clone(),
5527 )?,
5528 Distribution::new_uniform(right_interval.clone())?,
5529 ],
5530 vec![
5531 Distribution::new_uniform(left_interval.clone())?,
5532 Distribution::new_generic(
5533 ScalarValue::from(12.),
5534 ScalarValue::from(12.),
5535 ScalarValue::Float64(None),
5536 right_interval.clone(),
5537 )?,
5538 ],
5539 vec![
5540 Distribution::new_generic(
5541 ScalarValue::from(6.),
5542 ScalarValue::from(6.),
5543 ScalarValue::Float64(None),
5544 left_interval.clone(),
5545 )?,
5546 Distribution::new_generic(
5547 ScalarValue::from(12.),
5548 ScalarValue::from(12.),
5549 ScalarValue::Float64(None),
5550 right_interval.clone(),
5551 )?,
5552 ],
5553 ];
5554
5555 let ops = vec![
5556 Operator::Plus,
5557 Operator::Minus,
5558 Operator::Multiply,
5559 Operator::Divide,
5560 ];
5561
5562 for child_view in children {
5563 let child_refs = child_view.iter().collect::<Vec<_>>();
5564 for op in &ops {
5565 let expr = binary_expr(Arc::clone(&a), *op, Arc::clone(&b), schema)?;
5566 assert_eq!(
5567 expr.propagate_statistics(&parent, child_refs.as_slice())?,
5568 Some(child_view.clone())
5569 );
5570 }
5571 }
5572 Ok(())
5573 }
5574
5575 #[test]
5576 #[expect(deprecated)]
5577 fn test_propagate_statistics_combination_of_range_holders_comparison() -> Result<()> {
5578 let schema = &Schema::new(vec![Field::new("a", DataType::Float64, false)]);
5579 let a = Arc::new(Column::new("a", 0)) as _;
5580 let b = lit(ScalarValue::from(12.0));
5581
5582 let left_interval = Interval::make(Some(0.0), Some(12.0))?;
5583 let right_interval = Interval::make(Some(6.0), Some(18.0))?;
5584
5585 let one = ScalarValue::from(1.0);
5586 let parent = Distribution::new_bernoulli(one)?;
5587 let children = vec![
5588 vec![
5589 Distribution::new_uniform(left_interval.clone())?,
5590 Distribution::new_uniform(right_interval.clone())?,
5591 ],
5592 vec![
5593 Distribution::new_generic(
5594 ScalarValue::from(6.),
5595 ScalarValue::from(6.),
5596 ScalarValue::Float64(None),
5597 left_interval.clone(),
5598 )?,
5599 Distribution::new_uniform(right_interval.clone())?,
5600 ],
5601 vec![
5602 Distribution::new_uniform(left_interval.clone())?,
5603 Distribution::new_generic(
5604 ScalarValue::from(12.),
5605 ScalarValue::from(12.),
5606 ScalarValue::Float64(None),
5607 right_interval.clone(),
5608 )?,
5609 ],
5610 vec![
5611 Distribution::new_generic(
5612 ScalarValue::from(6.),
5613 ScalarValue::from(6.),
5614 ScalarValue::Float64(None),
5615 left_interval.clone(),
5616 )?,
5617 Distribution::new_generic(
5618 ScalarValue::from(12.),
5619 ScalarValue::from(12.),
5620 ScalarValue::Float64(None),
5621 right_interval.clone(),
5622 )?,
5623 ],
5624 ];
5625
5626 let ops = vec![
5627 Operator::Eq,
5628 Operator::Gt,
5629 Operator::GtEq,
5630 Operator::Lt,
5631 Operator::LtEq,
5632 ];
5633
5634 for child_view in children {
5635 let child_refs = child_view.iter().collect::<Vec<_>>();
5636 for op in &ops {
5637 let expr = binary_expr(Arc::clone(&a), *op, Arc::clone(&b), schema)?;
5638 assert!(
5639 expr.propagate_statistics(&parent, child_refs.as_slice())?
5640 .is_some()
5641 );
5642 }
5643 }
5644
5645 Ok(())
5646 }
5647
5648 #[test]
5649 fn test_fmt_sql() -> Result<()> {
5650 let schema = Schema::new(vec![
5651 Field::new("a", DataType::Int32, false),
5652 Field::new("b", DataType::Int32, false),
5653 ]);
5654
5655 let simple_expr = binary_expr(
5657 col("a", &schema)?,
5658 Operator::Plus,
5659 col("b", &schema)?,
5660 &schema,
5661 )?;
5662 let display_string = simple_expr.to_string();
5663 assert_eq!(display_string, "a@0 + b@1");
5664 let sql_string = fmt_sql(&simple_expr).to_string();
5665 assert_eq!(sql_string, "a + b");
5666
5667 let nested_expr = binary_expr(
5669 Arc::new(binary_expr(
5670 col("a", &schema)?,
5671 Operator::Plus,
5672 col("b", &schema)?,
5673 &schema,
5674 )?),
5675 Operator::Multiply,
5676 col("b", &schema)?,
5677 &schema,
5678 )?;
5679 let display_string = nested_expr.to_string();
5680 assert_eq!(display_string, "(a@0 + b@1) * b@1");
5681 let sql_string = fmt_sql(&nested_expr).to_string();
5682 assert_eq!(sql_string, "(a + b) * b");
5683
5684 let nested_same_prec = binary_expr(
5686 Arc::new(binary_expr(
5687 col("a", &schema)?,
5688 Operator::Plus,
5689 col("b", &schema)?,
5690 &schema,
5691 )?),
5692 Operator::Plus,
5693 col("b", &schema)?,
5694 &schema,
5695 )?;
5696 let display_string = nested_same_prec.to_string();
5697 assert_eq!(display_string, "a@0 + b@1 + b@1");
5698 let sql_string = fmt_sql(&nested_same_prec).to_string();
5699 assert_eq!(sql_string, "a + b + b");
5700
5701 let lit_expr = binary_expr(
5703 col("a", &schema)?,
5704 Operator::Eq,
5705 lit(ScalarValue::Int32(Some(42))),
5706 &schema,
5707 )?;
5708 let display_string = lit_expr.to_string();
5709 assert_eq!(display_string, "a@0 = 42");
5710 let sql_string = fmt_sql(&lit_expr).to_string();
5711 assert_eq!(sql_string, "a = 42");
5712
5713 Ok(())
5714 }
5715
5716 #[test]
5717 fn test_check_short_circuit() {
5718 let schema = Arc::new(Schema::new(vec![
5720 Field::new("a", DataType::Int32, false),
5721 Field::new("b", DataType::Int32, false),
5722 ]));
5723 let a_array = Int32Array::from(vec![1, 3, 4, 5, 6]);
5724 let b_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
5725 let batch = RecordBatch::try_new(
5726 Arc::clone(&schema),
5727 vec![Arc::new(a_array), Arc::new(b_array)],
5728 )
5729 .unwrap();
5730
5731 let left_expr = logical2physical(&logical_col("a").eq(expr_lit(2)), &schema);
5733 let left_value = left_expr.evaluate(&batch).unwrap();
5734 assert!(matches!(
5735 check_short_circuit(&left_value, &Operator::And),
5736 ShortCircuitStrategy::ReturnLeft
5737 ));
5738
5739 let left_expr = logical2physical(&logical_col("a").eq(expr_lit(3)), &schema);
5741 let left_value = left_expr.evaluate(&batch).unwrap();
5742 let ColumnarValue::Array(array) = &left_value else {
5743 panic!("Expected ColumnarValue::Array");
5744 };
5745 let ShortCircuitStrategy::PreSelection { mask, fill_value } =
5746 check_short_circuit(&left_value, &Operator::And)
5747 else {
5748 panic!("Expected ShortCircuitStrategy::PreSelection");
5749 };
5750 assert!(!fill_value);
5753 let expected_boolean_arr: Vec<_> =
5754 as_boolean_array(array).unwrap().iter().collect();
5755 let boolean_arr: Vec<_> = mask.iter().collect();
5756 assert_eq!(expected_boolean_arr, boolean_arr);
5757
5758 let left_expr = logical2physical(&logical_col("a").gt(expr_lit(0)), &schema);
5760 let left_value = left_expr.evaluate(&batch).unwrap();
5761 assert!(matches!(
5762 check_short_circuit(&left_value, &Operator::Or),
5763 ShortCircuitStrategy::ReturnLeft
5764 ));
5765
5766 let left_expr: Arc<dyn PhysicalExpr> =
5768 logical2physical(&logical_col("a").gt(expr_lit(2)), &schema);
5769 let left_value = left_expr.evaluate(&batch).unwrap();
5770 let ColumnarValue::Array(array) = &left_value else {
5771 panic!("Expected ColumnarValue::Array");
5772 };
5773 let ShortCircuitStrategy::PreSelection { mask, fill_value } =
5774 check_short_circuit(&left_value, &Operator::Or)
5775 else {
5776 panic!("Expected ShortCircuitStrategy::PreSelection");
5777 };
5778 assert!(fill_value);
5781 let negated_lhs: Vec<_> = as_boolean_array(array)
5782 .unwrap()
5783 .iter()
5784 .map(|v| v.map(|b| !b))
5785 .collect();
5786 let boolean_arr: Vec<_> = mask.iter().collect();
5787 assert_eq!(negated_lhs, boolean_arr);
5788
5789 let left_expr: Arc<dyn PhysicalExpr> =
5791 logical2physical(&logical_col("a").gt(expr_lit(4)), &schema);
5792 let left_value = left_expr.evaluate(&batch).unwrap();
5793 assert!(matches!(
5794 check_short_circuit(&left_value, &Operator::Or),
5795 ShortCircuitStrategy::None
5796 ));
5797
5798 let schema_nullable = Arc::new(Schema::new(vec![
5800 Field::new("c", DataType::Boolean, true),
5801 Field::new("d", DataType::Boolean, true),
5802 ]));
5803
5804 let c_array = Arc::new(BooleanArray::from(vec![
5806 Some(true),
5807 Some(false),
5808 None,
5809 Some(true),
5810 None,
5811 ])) as ArrayRef;
5812 let d_array = Arc::new(BooleanArray::from(vec![
5813 Some(false),
5814 Some(true),
5815 Some(false),
5816 None,
5817 Some(true),
5818 ])) as ArrayRef;
5819
5820 let batch_nullable = RecordBatch::try_new(
5821 Arc::clone(&schema_nullable),
5822 vec![Arc::clone(&c_array), Arc::clone(&d_array)],
5823 )
5824 .unwrap();
5825
5826 let mixed_nulls = logical2physical(&logical_col("c"), &schema_nullable);
5828 let mixed_nulls_value = mixed_nulls.evaluate(&batch_nullable).unwrap();
5829 assert!(matches!(
5830 check_short_circuit(&mixed_nulls_value, &Operator::And),
5831 ShortCircuitStrategy::None
5832 ));
5833
5834 assert!(matches!(
5836 check_short_circuit(&mixed_nulls_value, &Operator::Or),
5837 ShortCircuitStrategy::None
5838 ));
5839
5840 let all_nulls = Arc::new(BooleanArray::from(vec![None, None, None])) as ArrayRef;
5842 let null_batch = RecordBatch::try_new(
5843 Arc::new(Schema::new(vec![Field::new("e", DataType::Boolean, true)])),
5844 vec![all_nulls],
5845 )
5846 .unwrap();
5847
5848 let null_expr = logical2physical(&logical_col("e"), &null_batch.schema());
5849 let null_value = null_expr.evaluate(&null_batch).unwrap();
5850
5851 assert!(matches!(
5853 check_short_circuit(&null_value, &Operator::And),
5854 ShortCircuitStrategy::None
5855 ));
5856 assert!(matches!(
5857 check_short_circuit(&null_value, &Operator::Or),
5858 ShortCircuitStrategy::None
5859 ));
5860
5861 let scalar_true = ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)));
5864 assert!(matches!(
5865 check_short_circuit(&scalar_true, &Operator::Or),
5866 ShortCircuitStrategy::ReturnLeft
5867 )); assert!(matches!(
5869 check_short_circuit(&scalar_true, &Operator::And),
5870 ShortCircuitStrategy::ReturnRight
5871 )); let scalar_false = ColumnarValue::Scalar(ScalarValue::Boolean(Some(false)));
5875 assert!(matches!(
5876 check_short_circuit(&scalar_false, &Operator::And),
5877 ShortCircuitStrategy::ReturnLeft
5878 )); assert!(matches!(
5880 check_short_circuit(&scalar_false, &Operator::Or),
5881 ShortCircuitStrategy::ReturnRight
5882 )); let scalar_null = ColumnarValue::Scalar(ScalarValue::Boolean(None));
5886 assert!(matches!(
5887 check_short_circuit(&scalar_null, &Operator::And),
5888 ShortCircuitStrategy::None
5889 ));
5890 assert!(matches!(
5891 check_short_circuit(&scalar_null, &Operator::Or),
5892 ShortCircuitStrategy::None
5893 ));
5894 }
5895
5896 #[test]
5901 fn test_pre_selection_scatter() {
5902 fn create_bool_array(bools: Vec<bool>) -> BooleanArray {
5903 BooleanArray::from(bools.into_iter().map(Some).collect::<Vec<_>>())
5904 }
5905 {
5907 let left = create_bool_array(vec![true, false, true, false, true]);
5910 let right = create_bool_array(vec![false, true, false]);
5911
5912 let result = pre_selection_scatter(&left, Some(&right), false).unwrap();
5913 let result_arr = result.into_array(left.len()).unwrap();
5914
5915 let expected = create_bool_array(vec![false, false, true, false, false]);
5916 assert_eq!(&expected, result_arr.as_boolean());
5917 }
5918 {
5920 let left =
5923 create_bool_array(vec![false, true, true, false, true, true, true]);
5924 let right = create_bool_array(vec![true, false, false, true, false]);
5925
5926 let result = pre_selection_scatter(&left, Some(&right), false).unwrap();
5927 let result_arr = result.into_array(left.len()).unwrap();
5928
5929 let expected =
5930 create_bool_array(vec![false, true, false, false, false, true, false]);
5931 assert_eq!(&expected, result_arr.as_boolean());
5932 }
5933 {
5935 let left = create_bool_array(vec![true, false, false]);
5938 let right = create_bool_array(vec![false]);
5939
5940 let result = pre_selection_scatter(&left, Some(&right), false).unwrap();
5941 let result_arr = result.into_array(left.len()).unwrap();
5942
5943 let expected = create_bool_array(vec![false, false, false]);
5944 assert_eq!(&expected, result_arr.as_boolean());
5945 }
5946 {
5948 let left = create_bool_array(vec![false, false, true]);
5951 let right = create_bool_array(vec![false]);
5952
5953 let result = pre_selection_scatter(&left, Some(&right), false).unwrap();
5954 let result_arr = result.into_array(left.len()).unwrap();
5955
5956 let expected = create_bool_array(vec![false, false, false]);
5957 assert_eq!(&expected, result_arr.as_boolean());
5958 }
5959 {
5961 let left = create_bool_array(vec![false, true, false, true]);
5964 let right = BooleanArray::from(vec![None, Some(false)]);
5965
5966 let result = pre_selection_scatter(&left, Some(&right), false).unwrap();
5967 let result_arr = result.into_array(left.len()).unwrap();
5968
5969 let expected = BooleanArray::from(vec![
5970 Some(false),
5971 None, Some(false),
5973 Some(false),
5974 ]);
5975 assert_eq!(&expected, result_arr.as_boolean());
5976 }
5977 {
5979 let left = create_bool_array(vec![true, false, true, false, true]);
5982 let right = create_bool_array(vec![false, true, false]);
5983
5984 let result = pre_selection_scatter(&left, Some(&right), true).unwrap();
5985 let result_arr = result.into_array(left.len()).unwrap();
5986
5987 let expected = create_bool_array(vec![false, true, true, true, false]);
5989 assert_eq!(&expected, result_arr.as_boolean());
5990 }
5991 {
5993 let left = create_bool_array(vec![false, true, false, true]);
5996 let right = BooleanArray::from(vec![None, Some(false)]);
5997
5998 let result = pre_selection_scatter(&left, Some(&right), true).unwrap();
5999 let result_arr = result.into_array(left.len()).unwrap();
6000
6001 let expected = BooleanArray::from(vec![
6002 Some(true), None, Some(true), Some(false),
6006 ]);
6007 assert_eq!(&expected, result_arr.as_boolean());
6008 }
6009 }
6010
6011 #[test]
6012 fn test_and_true_preselection_returns_lhs() {
6013 let schema =
6014 Arc::new(Schema::new(vec![Field::new("c", DataType::Boolean, false)]));
6015 let c_array = Arc::new(BooleanArray::from(vec![false, true, false, false, false]))
6016 as ArrayRef;
6017 let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::clone(&c_array)])
6018 .unwrap();
6019
6020 let expr = logical2physical(&logical_col("c").and(expr_lit(true)), &schema);
6021
6022 let result = expr.evaluate(&batch).unwrap();
6023 let ColumnarValue::Array(result_arr) = result else {
6024 panic!("Expected ColumnarValue::Array");
6025 };
6026
6027 let expected: Vec<_> = c_array.as_boolean().iter().collect();
6028 let actual: Vec<_> = result_arr.as_boolean().iter().collect();
6029 assert_eq!(
6030 expected, actual,
6031 "AND with TRUE must equal LHS even with PreSelection"
6032 );
6033 }
6034
6035 #[test]
6036 fn test_or_false_preselection_returns_lhs() {
6037 let schema =
6040 Arc::new(Schema::new(vec![Field::new("c", DataType::Boolean, false)]));
6041 let c_array =
6042 Arc::new(BooleanArray::from(vec![true, false, true, true, true])) as ArrayRef;
6043 let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::clone(&c_array)])
6044 .unwrap();
6045
6046 let expr = logical2physical(&logical_col("c").or(expr_lit(false)), &schema);
6047
6048 let result = expr.evaluate(&batch).unwrap();
6049 let ColumnarValue::Array(result_arr) = result else {
6050 panic!("Expected ColumnarValue::Array");
6051 };
6052
6053 let expected: Vec<_> = c_array.as_boolean().iter().collect();
6054 let actual: Vec<_> = result_arr.as_boolean().iter().collect();
6055 assert_eq!(
6056 expected, actual,
6057 "OR with FALSE must equal LHS even with PreSelection"
6058 );
6059 }
6060
6061 #[test]
6062 fn test_or_preselection_matches_kleene() {
6063 use arrow::compute::kernels::boolean::or_kleene;
6065
6066 let schema = Arc::new(Schema::new(vec![
6067 Field::new("c", DataType::Boolean, true),
6068 Field::new("d", DataType::Boolean, true),
6069 ]));
6070
6071 let c = BooleanArray::from(vec![
6073 true, true, false, true, true, true, true, false, true, true,
6074 ]);
6075
6076 let d_cases = vec![
6077 BooleanArray::from(vec![
6079 Some(false),
6080 Some(true),
6081 Some(true),
6082 Some(false),
6083 Some(false),
6084 Some(true),
6085 Some(false),
6086 None,
6087 Some(true),
6088 None,
6089 ]),
6090 BooleanArray::from(vec![Some(true); 10]),
6092 BooleanArray::from(vec![Some(false); 10]),
6094 ];
6095
6096 for d in d_cases {
6097 let batch = RecordBatch::try_new(
6098 Arc::clone(&schema),
6099 vec![
6100 Arc::new(c.clone()) as ArrayRef,
6101 Arc::new(d.clone()) as ArrayRef,
6102 ],
6103 )
6104 .unwrap();
6105
6106 let expr = logical2physical(&logical_col("c").or(logical_col("d")), &schema);
6107 let result = expr.evaluate(&batch).unwrap().into_array(c.len()).unwrap();
6108
6109 let expected = or_kleene(&c, &d).unwrap();
6110 assert_eq!(
6111 expected,
6112 *result.as_boolean(),
6113 "OR pre-selection must match Kleene OR for d = {d:?}"
6114 );
6115 }
6116 }
6117
6118 #[test]
6119 fn test_evaluate_bounds_int32() {
6120 let schema = Schema::new(vec![
6121 Field::new("a", DataType::Int32, false),
6122 Field::new("b", DataType::Int32, false),
6123 ]);
6124
6125 let a = Arc::new(Column::new("a", 0)) as _;
6126 let b = Arc::new(Column::new("b", 1)) as _;
6127
6128 let add_expr =
6130 binary_expr(Arc::clone(&a), Operator::Plus, Arc::clone(&b), &schema).unwrap();
6131 let add_bounds = add_expr
6132 .evaluate_bounds(&[
6133 &Interval::make(Some(1), Some(10)).unwrap(),
6134 &Interval::make(Some(5), Some(15)).unwrap(),
6135 ])
6136 .unwrap();
6137 assert_eq!(add_bounds, Interval::make(Some(6), Some(25)).unwrap());
6138
6139 let sub_expr =
6141 binary_expr(Arc::clone(&a), Operator::Minus, Arc::clone(&b), &schema)
6142 .unwrap();
6143 let sub_bounds = sub_expr
6144 .evaluate_bounds(&[
6145 &Interval::make(Some(1), Some(10)).unwrap(),
6146 &Interval::make(Some(5), Some(15)).unwrap(),
6147 ])
6148 .unwrap();
6149 assert_eq!(sub_bounds, Interval::make(Some(-14), Some(5)).unwrap());
6150
6151 let mul_expr =
6153 binary_expr(Arc::clone(&a), Operator::Multiply, Arc::clone(&b), &schema)
6154 .unwrap();
6155 let mul_bounds = mul_expr
6156 .evaluate_bounds(&[
6157 &Interval::make(Some(1), Some(10)).unwrap(),
6158 &Interval::make(Some(5), Some(15)).unwrap(),
6159 ])
6160 .unwrap();
6161 assert_eq!(mul_bounds, Interval::make(Some(5), Some(150)).unwrap());
6162
6163 let div_expr =
6165 binary_expr(Arc::clone(&a), Operator::Divide, Arc::clone(&b), &schema)
6166 .unwrap();
6167 let div_bounds = div_expr
6168 .evaluate_bounds(&[
6169 &Interval::make(Some(10), Some(20)).unwrap(),
6170 &Interval::make(Some(2), Some(5)).unwrap(),
6171 ])
6172 .unwrap();
6173 assert_eq!(div_bounds, Interval::make(Some(2), Some(10)).unwrap());
6174 }
6175
6176 #[test]
6177 fn test_evaluate_bounds_bool() {
6178 let schema = Schema::new(vec![
6179 Field::new("a", DataType::Boolean, false),
6180 Field::new("b", DataType::Boolean, false),
6181 ]);
6182
6183 let a = Arc::new(Column::new("a", 0)) as _;
6184 let b = Arc::new(Column::new("b", 1)) as _;
6185
6186 let or_expr =
6188 binary_expr(Arc::clone(&a), Operator::Or, Arc::clone(&b), &schema).unwrap();
6189 let or_bounds = or_expr
6190 .evaluate_bounds(&[
6191 &Interval::make(Some(true), Some(true)).unwrap(),
6192 &Interval::make(Some(false), Some(false)).unwrap(),
6193 ])
6194 .unwrap();
6195 assert_eq!(or_bounds, Interval::make(Some(true), Some(true)).unwrap());
6196
6197 let and_expr =
6199 binary_expr(Arc::clone(&a), Operator::And, Arc::clone(&b), &schema).unwrap();
6200 let and_bounds = and_expr
6201 .evaluate_bounds(&[
6202 &Interval::make(Some(true), Some(true)).unwrap(),
6203 &Interval::make(Some(false), Some(false)).unwrap(),
6204 ])
6205 .unwrap();
6206 assert_eq!(
6207 and_bounds,
6208 Interval::make(Some(false), Some(false)).unwrap()
6209 );
6210 }
6211
6212 #[test]
6213 fn test_evaluate_nested_type() {
6214 let batch_schema = Arc::new(Schema::new(vec![
6215 Field::new(
6216 "a",
6217 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
6218 true,
6219 ),
6220 Field::new(
6221 "b",
6222 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
6223 true,
6224 ),
6225 ]));
6226
6227 let mut list_builder_a = ListBuilder::new(Int32Builder::new());
6228
6229 list_builder_a.append_value([Some(1)]);
6230 list_builder_a.append_value([Some(2)]);
6231 list_builder_a.append_value([]);
6232 list_builder_a.append_value([None]);
6233
6234 let list_array_a: ArrayRef = Arc::new(list_builder_a.finish());
6235
6236 let mut list_builder_b = ListBuilder::new(Int32Builder::new());
6237
6238 list_builder_b.append_value([Some(1)]);
6239 list_builder_b.append_value([Some(2)]);
6240 list_builder_b.append_value([]);
6241 list_builder_b.append_value([None]);
6242
6243 let list_array_b: ArrayRef = Arc::new(list_builder_b.finish());
6244
6245 let batch =
6246 RecordBatch::try_new(batch_schema, vec![list_array_a, list_array_b]).unwrap();
6247
6248 let schema = Arc::new(Schema::new(vec![
6249 Field::new(
6250 "a",
6251 DataType::List(Arc::new(Field::new("foo", DataType::Int32, true))),
6252 true,
6253 ),
6254 Field::new(
6255 "b",
6256 DataType::List(Arc::new(Field::new("bar", DataType::Int32, true))),
6257 true,
6258 ),
6259 ]));
6260
6261 let a = Arc::new(Column::new("a", 0)) as _;
6262 let b = Arc::new(Column::new("b", 1)) as _;
6263
6264 let eq_expr =
6265 binary_expr(Arc::clone(&a), Operator::Eq, Arc::clone(&b), &schema).unwrap();
6266
6267 let eq_result = eq_expr.evaluate(&batch).unwrap();
6268 let expected =
6269 BooleanArray::from_iter(vec![Some(true), Some(true), Some(true), Some(true)]);
6270 assert_eq!(eq_result.into_array(4).unwrap().as_boolean(), &expected);
6271 }
6272}