1use std::cmp::Ordering;
16use std::collections::{BTreeMap, HashSet};
17use std::fmt::{self, Debug};
18use std::ops::Bound;
19use std::rc::Rc;
20use std::sync::LazyLock;
21
22use fixedbitset::FixedBitSet;
23use itertools::Itertools;
24use risingwave_common::catalog::{Schema, TableDesc};
25use risingwave_common::types::{DataType, DefaultOrd, ScalarImpl};
26use risingwave_common::util::iter_util::ZipEqFast;
27use risingwave_common::util::scan_range::{ScanRange, is_full_range};
28use risingwave_common::util::sort_util::{OrderType, cmp_rows};
29
30use crate::error::Result;
31use crate::expr::{
32 ExprDisplay, ExprImpl, ExprMutator, ExprRewriter, ExprType, ExprVisitor, FunctionCall,
33 InequalityInputPair, InputRef, collect_input_refs, column_self_eq_eliminate,
34 factorization_expr, fold_boolean_constant, push_down_not, to_conjunctions,
35 try_get_bool_constant,
36};
37use crate::utils::condition::cast_compare::{ResultForCmp, ResultForEq};
38
39#[derive(Debug, Clone, PartialEq, Eq, Hash)]
40pub struct Condition {
41 pub conjunctions: Vec<ExprImpl>,
43}
44
45impl IntoIterator for Condition {
46 type IntoIter = std::vec::IntoIter<ExprImpl>;
47 type Item = ExprImpl;
48
49 fn into_iter(self) -> Self::IntoIter {
50 self.conjunctions.into_iter()
51 }
52}
53
54impl fmt::Display for Condition {
55 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56 let mut conjunctions = self.conjunctions.iter();
57 if let Some(expr) = conjunctions.next() {
58 write!(f, "{:?}", expr)?;
59 }
60 if self.always_true() {
61 write!(f, "true")?;
62 } else {
63 for expr in conjunctions {
64 write!(f, " AND {:?}", expr)?;
65 }
66 }
67 Ok(())
68 }
69}
70
71impl Condition {
72 pub fn with_expr(expr: ExprImpl) -> Self {
73 let conjunctions = to_conjunctions(expr);
74
75 Self { conjunctions }.simplify()
76 }
77
78 pub fn true_cond() -> Self {
79 Self {
80 conjunctions: vec![],
81 }
82 }
83
84 pub fn false_cond() -> Self {
85 Self {
86 conjunctions: vec![ExprImpl::literal_bool(false)],
87 }
88 }
89
90 pub fn always_true(&self) -> bool {
91 self.conjunctions.is_empty()
92 }
93
94 pub fn always_false(&self) -> bool {
95 static FALSE: LazyLock<ExprImpl> = LazyLock::new(|| ExprImpl::literal_bool(false));
96 !self.conjunctions.is_empty() && self.conjunctions.contains(&*FALSE)
98 }
99
100 pub fn as_expr_unless_true(&self) -> Option<ExprImpl> {
102 if self.always_true() {
103 None
104 } else {
105 Some(self.clone().into())
106 }
107 }
108
109 #[must_use]
110 pub fn and(self, other: Self) -> Self {
111 let mut ret = self;
112 ret.conjunctions.extend(other.conjunctions);
113 ret.simplify()
114 }
115
116 #[must_use]
117 pub fn or(self, other: Self) -> Self {
118 let or_expr = ExprImpl::FunctionCall(
119 FunctionCall::new_unchecked(
120 ExprType::Or,
121 vec![self.into(), other.into()],
122 DataType::Boolean,
123 )
124 .into(),
125 );
126 let ret = Self::with_expr(or_expr);
127 ret.simplify()
128 }
129
130 #[must_use]
132 pub fn split(self, left_col_num: usize, right_col_num: usize) -> (Self, Self, Self) {
133 let left_bit_map = FixedBitSet::from_iter(0..left_col_num);
134 let right_bit_map = FixedBitSet::from_iter(left_col_num..left_col_num + right_col_num);
135
136 self.group_by::<_, 3>(|expr| {
137 let input_bits = expr.collect_input_refs(left_col_num + right_col_num);
138 if input_bits.is_subset(&left_bit_map) {
139 0
140 } else if input_bits.is_subset(&right_bit_map) {
141 1
142 } else {
143 2
144 }
145 })
146 .into_iter()
147 .next_tuple()
148 .unwrap()
149 }
150
151 pub fn collect_input_refs(&self, input_col_num: usize) -> FixedBitSet {
156 collect_input_refs(input_col_num, &self.conjunctions)
157 }
158
159 #[must_use]
173 pub fn split_by_input_col_nums(
174 self,
175 input_col_nums: &[usize],
176 only_eq: bool,
177 ) -> (BTreeMap<(usize, usize), Self>, Self) {
178 let mut bitmaps = Vec::with_capacity(input_col_nums.len());
179 let mut cols_seen = 0;
180 for cols in input_col_nums {
181 bitmaps.push(FixedBitSet::from_iter(cols_seen..cols_seen + cols));
182 cols_seen += cols;
183 }
184
185 let mut pairwise_conditions = BTreeMap::new();
186 let mut non_eq_join = vec![];
187
188 for expr in self.conjunctions {
189 let input_bits = expr.collect_input_refs(cols_seen);
190 let mut subset_indices = Vec::with_capacity(input_col_nums.len());
191 for (idx, bitmap) in bitmaps.iter().enumerate() {
192 if !input_bits.is_disjoint(bitmap) {
193 subset_indices.push(idx);
194 }
195 }
196 if subset_indices.len() != 2 || (only_eq && expr.as_eq_cond().is_none()) {
197 non_eq_join.push(expr);
198 } else {
199 let key = if subset_indices[0] < subset_indices[1] {
201 (subset_indices[0], subset_indices[1])
202 } else {
203 (subset_indices[1], subset_indices[0])
204 };
205 let e = pairwise_conditions
206 .entry(key)
207 .or_insert_with(Condition::true_cond);
208 e.conjunctions.push(expr);
209 }
210 }
211 (
212 pairwise_conditions,
213 Condition {
214 conjunctions: non_eq_join,
215 },
216 )
217 }
218
219 #[must_use]
220 pub fn split_eq_keys(
227 self,
228 left_col_num: usize,
229 right_col_num: usize,
230 ) -> (Vec<(InputRef, InputRef, bool)>, Self) {
231 let left_bit_map = FixedBitSet::from_iter(0..left_col_num);
232 let right_bit_map = FixedBitSet::from_iter(left_col_num..left_col_num + right_col_num);
233
234 let (mut eq_keys, mut others) = (vec![], vec![]);
235 self.conjunctions.into_iter().for_each(|expr| {
236 let input_bits = expr.collect_input_refs(left_col_num + right_col_num);
237 if input_bits.is_disjoint(&left_bit_map) || input_bits.is_disjoint(&right_bit_map) {
238 others.push(expr)
239 } else if let Some(columns) = expr.as_eq_cond() {
240 eq_keys.push((columns.0, columns.1, false));
241 } else if let Some(columns) = expr.as_is_not_distinct_from_cond() {
242 eq_keys.push((columns.0, columns.1, true));
243 } else {
244 others.push(expr)
245 }
246 });
247
248 (
249 eq_keys,
250 Condition {
251 conjunctions: others,
252 },
253 )
254 }
255
256 pub(crate) fn extract_inequality_keys(
263 &self,
264 left_col_num: usize,
265 right_col_num: usize,
266 ) -> Vec<(usize, InequalityInputPair)> {
267 let left_bit_map = FixedBitSet::from_iter(0..left_col_num);
268 let right_bit_map = FixedBitSet::from_iter(left_col_num..left_col_num + right_col_num);
269
270 self.conjunctions
271 .iter()
272 .enumerate()
273 .filter_map(|(conjunction_idx, expr)| {
274 let input_bits = expr.collect_input_refs(left_col_num + right_col_num);
275 if input_bits.is_disjoint(&left_bit_map) || input_bits.is_disjoint(&right_bit_map) {
276 None
277 } else {
278 expr.as_input_comparison_cond()
279 .map(|inequality_pair| (conjunction_idx, inequality_pair))
280 }
281 })
282 .collect_vec()
283 }
284
285 #[must_use]
288 pub fn split_disjoint(self, columns: &FixedBitSet) -> (Self, Self) {
289 self.group_by::<_, 2>(|expr| {
290 let input_bits = expr.collect_input_refs(columns.len());
291 input_bits.is_disjoint(columns) as usize
292 })
293 .into_iter()
294 .next_tuple()
295 .unwrap()
296 }
297
298 fn disjunctions_to_scan_ranges(
302 table_desc: Rc<TableDesc>,
303 max_split_range_gap: u64,
304 disjunctions: Vec<ExprImpl>,
305 ) -> Result<Option<(Vec<ScanRange>, bool)>> {
306 let disjunctions_result: Result<Vec<(Vec<ScanRange>, Self)>> = disjunctions
307 .into_iter()
308 .map(|x| {
309 Condition {
310 conjunctions: to_conjunctions(x),
311 }
312 .split_to_scan_ranges(table_desc.clone(), max_split_range_gap)
313 })
314 .collect();
315
316 let disjunctions_result = disjunctions_result?;
318
319 let all_equal = disjunctions_result
322 .iter()
323 .all(|(scan_ranges, other_condition)| {
324 other_condition.always_true()
325 && scan_ranges
326 .iter()
327 .all(|x| !x.eq_conds.is_empty() && is_full_range(&x.range))
328 });
329
330 if all_equal {
331 let scan_ranges = disjunctions_result
335 .into_iter()
336 .flat_map(|(scan_ranges, _)| scan_ranges)
337 .sorted_by(|a, b| a.eq_conds.len().cmp(&b.eq_conds.len()))
339 .collect_vec();
340 let mut non_overlap_scan_ranges: Vec<ScanRange> = vec![];
342 for s1 in &scan_ranges {
343 let overlap = non_overlap_scan_ranges.iter().any(|s2| {
344 #[allow(clippy::disallowed_methods)]
345 s1.eq_conds
346 .iter()
347 .zip(s2.eq_conds.iter())
348 .all(|(a, b)| a == b)
349 });
350 if !overlap {
354 non_overlap_scan_ranges.push(s1.clone());
355 }
356 }
357
358 Ok(Some((non_overlap_scan_ranges, false)))
359 } else {
360 let mut scan_ranges = vec![];
361 for (scan_ranges_chunk, _) in disjunctions_result {
362 if scan_ranges_chunk.is_empty() {
363 return Ok(None);
365 }
366
367 scan_ranges.extend(scan_ranges_chunk);
368 }
369
370 let order_types = table_desc
371 .pk
372 .iter()
373 .cloned()
374 .map(|x| {
375 if x.order_type.is_descending() {
376 x.order_type.reverse()
377 } else {
378 x.order_type
379 }
380 })
381 .collect_vec();
382 scan_ranges.sort_by(|left, right| {
383 let (left_start, _left_end) = &left.convert_to_range();
384 let (right_start, _right_end) = &right.convert_to_range();
385
386 let left_start_vec = match &left_start {
387 Bound::Included(vec) | Bound::Excluded(vec) => vec,
388 _ => &vec![],
389 };
390 let right_start_vec = match &right_start {
391 Bound::Included(vec) | Bound::Excluded(vec) => vec,
392 _ => &vec![],
393 };
394
395 if left_start_vec.is_empty() && right_start_vec.is_empty() {
396 return Ordering::Less;
397 }
398
399 if left_start_vec.is_empty() {
400 return Ordering::Less;
401 }
402
403 if right_start_vec.is_empty() {
404 return Ordering::Greater;
405 }
406
407 let cmp_column_len = left_start_vec.len().min(right_start_vec.len());
408 cmp_rows(
409 &left_start_vec[0..cmp_column_len],
410 &right_start_vec[0..cmp_column_len],
411 &order_types[0..cmp_column_len],
412 )
413 });
414
415 if scan_ranges.is_empty() {
416 return Ok(None);
417 }
418
419 if scan_ranges.len() == 1 {
420 return Ok(Some((scan_ranges, true)));
421 }
422
423 let mut output_scan_ranges: Vec<ScanRange> = vec![];
424 output_scan_ranges.push(scan_ranges[0].clone());
425 let mut idx = 1;
426 loop {
427 if idx >= scan_ranges.len() {
428 break;
429 }
430
431 let scan_range_left = output_scan_ranges.last_mut().unwrap();
432 let scan_range_right = &scan_ranges[idx];
433
434 if scan_range_left.eq_conds == scan_range_right.eq_conds {
435 if !ScanRange::is_overlap(scan_range_left, scan_range_right, &order_types) {
438 output_scan_ranges.push(scan_range_right.clone());
440 idx += 1;
441 continue;
442 }
443
444 fn merge_bound(
446 left_scan_range: &Bound<Vec<Option<ScalarImpl>>>,
447 right_scan_range: &Bound<Vec<Option<ScalarImpl>>>,
448 order_types: &[OrderType],
449 left_bound: bool,
450 ) -> Bound<Vec<Option<ScalarImpl>>> {
451 let left_scan_range = match left_scan_range {
452 Bound::Included(vec) | Bound::Excluded(vec) => vec,
453 Bound::Unbounded => return Bound::Unbounded,
454 };
455
456 let right_scan_range = match right_scan_range {
457 Bound::Included(vec) | Bound::Excluded(vec) => vec,
458 Bound::Unbounded => return Bound::Unbounded,
459 };
460
461 let cmp_len = left_scan_range.len().min(right_scan_range.len());
462
463 let cmp = cmp_rows(
464 &left_scan_range[..cmp_len],
465 &right_scan_range[..cmp_len],
466 &order_types[..cmp_len],
467 );
468
469 let bound = {
470 if (cmp.is_le() && left_bound) || (cmp.is_ge() && !left_bound) {
471 left_scan_range.to_vec()
472 } else {
473 right_scan_range.to_vec()
474 }
475 };
476
477 Bound::Included(bound)
479 }
480
481 scan_range_left.range.0 = merge_bound(
482 &scan_range_left.range.0,
483 &scan_range_right.range.0,
484 &order_types,
485 true,
486 );
487
488 scan_range_left.range.1 = merge_bound(
489 &scan_range_left.range.1,
490 &scan_range_right.range.1,
491 &order_types,
492 false,
493 );
494
495 if scan_range_left.is_full_table_scan() {
496 return Ok(None);
497 }
498 } else {
499 output_scan_ranges.push(scan_range_right.clone());
500 }
501
502 idx += 1;
503 }
504
505 Ok(Some((output_scan_ranges, true)))
506 }
507 }
508
509 fn split_row_cmp_to_scan_ranges(
510 &self,
511 table_desc: Rc<TableDesc>,
512 ) -> Result<Option<(Vec<ScanRange>, Self)>> {
513 let (mut row_conjunctions, row_conjunctions_without_struct): (Vec<_>, Vec<_>) =
514 self.conjunctions.clone().into_iter().partition(|expr| {
515 if let Some(f) = expr.as_function_call() {
516 if let Some(left_input) = f.inputs().get(0)
517 && let Some(left_input) = left_input.as_function_call()
518 && matches!(left_input.func_type(), ExprType::Row)
519 && left_input.inputs().iter().all(|x| x.is_input_ref())
520 && let Some(right_input) = f.inputs().get(1)
521 && right_input.is_literal()
522 {
523 true
524 } else {
525 false
526 }
527 } else {
528 false
529 }
530 });
531 if row_conjunctions.len() == 1 {
536 let row_conjunction = row_conjunctions.pop().unwrap();
537 let row_left_inputs = row_conjunction
538 .as_function_call()
539 .unwrap()
540 .inputs()
541 .get(0)
542 .unwrap()
543 .as_function_call()
544 .unwrap()
545 .inputs();
546 let row_right_literal = row_conjunction
547 .as_function_call()
548 .unwrap()
549 .inputs()
550 .get(1)
551 .unwrap()
552 .as_literal()
553 .unwrap();
554 if !matches!(row_right_literal.get_data(), Some(ScalarImpl::Struct(_))) {
555 return Ok(None);
556 }
557 let row_right_literal_data = row_right_literal.get_data().clone().unwrap();
558 let right_iter = row_right_literal_data.as_struct().fields();
559 let func_type = row_conjunction.as_function_call().unwrap().func_type();
560 if row_left_inputs.len() > 1
561 && (matches!(func_type, ExprType::LessThan)
562 || matches!(func_type, ExprType::GreaterThan))
563 {
564 let mut pk_struct = vec![];
565 let mut order_type = None;
566 let mut all_added = true;
567 let mut iter = row_left_inputs.iter().zip_eq_fast(right_iter);
568 for column_order in &table_desc.pk {
569 if let Some((left_expr, right_expr)) = iter.next() {
570 if left_expr.as_input_ref().unwrap().index != column_order.column_index {
571 all_added = false;
572 break;
573 }
574 match order_type {
575 Some(o) => {
576 if o != column_order.order_type {
577 all_added = false;
578 break;
579 }
580 }
581 None => order_type = Some(column_order.order_type),
582 }
583 pk_struct.push(right_expr.clone());
584 }
585 }
586
587 if !pk_struct.is_empty() {
589 if !all_added {
590 let scan_range = ScanRange {
591 eq_conds: vec![],
592 range: match func_type {
593 ExprType::GreaterThan => {
594 (Bound::Included(pk_struct), Bound::Unbounded)
595 }
596 ExprType::LessThan => {
597 (Bound::Unbounded, Bound::Included(pk_struct))
598 }
599 _ => unreachable!(),
600 },
601 };
602 return Ok(Some((
603 vec![scan_range],
604 Condition {
605 conjunctions: self.conjunctions.clone(),
606 },
607 )));
608 } else {
609 let scan_range = ScanRange {
610 eq_conds: vec![],
611 range: match func_type {
612 ExprType::GreaterThan => {
613 (Bound::Excluded(pk_struct), Bound::Unbounded)
614 }
615 ExprType::LessThan => {
616 (Bound::Unbounded, Bound::Excluded(pk_struct))
617 }
618 _ => unreachable!(),
619 },
620 };
621 return Ok(Some((
622 vec![scan_range],
623 Condition {
624 conjunctions: row_conjunctions_without_struct,
625 },
626 )));
627 }
628 }
629 }
630 }
631 Ok(None)
632 }
633
634 pub fn get_eq_const_input_refs(&self) -> Vec<InputRef> {
636 self.conjunctions
637 .iter()
638 .filter_map(|expr| expr.as_eq_const().map(|(input_ref, _)| input_ref))
639 .collect()
640 }
641
642 pub fn split_to_scan_ranges(
644 self,
645 table_desc: Rc<TableDesc>,
646 max_split_range_gap: u64,
647 ) -> Result<(Vec<ScanRange>, Self)> {
648 fn false_cond() -> (Vec<ScanRange>, Condition) {
649 (vec![], Condition::false_cond())
650 }
651
652 if self.conjunctions.len() == 1 {
654 if let Some(disjunctions) = self.conjunctions[0].as_or_disjunctions() {
655 if let Some((scan_ranges, maintaining_condition)) =
656 Self::disjunctions_to_scan_ranges(
657 table_desc,
658 max_split_range_gap,
659 disjunctions,
660 )?
661 {
662 if maintaining_condition {
663 return Ok((scan_ranges, self));
664 } else {
665 return Ok((scan_ranges, Condition::true_cond()));
666 }
667 } else {
668 return Ok((vec![], self));
669 }
670 }
671 }
672 if let Some((scan_ranges, other_condition)) =
673 self.split_row_cmp_to_scan_ranges(table_desc.clone())?
674 {
675 return Ok((scan_ranges, other_condition));
676 }
677
678 let mut groups = Self::classify_conjunctions_by_pk(self.conjunctions, &table_desc);
679 let mut other_conds = groups.pop().unwrap();
680
681 let mut scan_range = ScanRange::full_table_scan();
683 for i in 0..table_desc.order_column_indices().len() {
684 let group = std::mem::take(&mut groups[i]);
685 if group.is_empty() {
686 groups.push(other_conds);
687 return Ok((
688 if scan_range.is_full_table_scan() {
689 vec![]
690 } else {
691 vec![scan_range]
692 },
693 Self {
694 conjunctions: groups[i + 1..].concat(),
695 },
696 ));
697 }
698
699 let Some((
700 lower_bound_conjunctions,
701 upper_bound_conjunctions,
702 eq_conds,
703 part_of_other_conds,
704 )) = Self::analyze_group(group)?
705 else {
706 return Ok(false_cond());
707 };
708 other_conds.extend(part_of_other_conds.into_iter());
709
710 let lower_bound = Self::merge_lower_bound_conjunctions(lower_bound_conjunctions);
711 let upper_bound = Self::merge_upper_bound_conjunctions(upper_bound_conjunctions);
712
713 if Self::is_invalid_range(&lower_bound, &upper_bound) {
714 return Ok(false_cond());
715 }
716
717 match eq_conds.len() {
719 1 => {
720 let eq_conds =
721 Self::extract_eq_conds_within_range(eq_conds, &upper_bound, &lower_bound);
722 if eq_conds.is_empty() {
723 return Ok(false_cond());
724 }
725 scan_range.eq_conds.extend(eq_conds.into_iter());
726 }
727 0 => {
728 let convert = |bound| match bound {
729 Bound::Included(l) => Bound::Included(vec![Some(l)]),
730 Bound::Excluded(l) => Bound::Excluded(vec![Some(l)]),
731 Bound::Unbounded => Bound::Unbounded,
732 };
733 scan_range.range = (convert(lower_bound), convert(upper_bound));
734 other_conds.extend(groups[i + 1..].iter().flatten().cloned());
735 break;
736 }
737 _ => {
738 let eq_conds =
745 Self::extract_eq_conds_within_range(eq_conds, &upper_bound, &lower_bound);
746 if eq_conds.is_empty() {
747 return Ok(false_cond());
748 }
749 other_conds.extend(groups[i + 1..].iter().flatten().cloned());
750 let scan_ranges = eq_conds
751 .into_iter()
752 .map(|lit| {
753 let mut scan_range = scan_range.clone();
754 scan_range.eq_conds.push(lit);
755 scan_range
756 })
757 .collect();
758 return Ok((
759 scan_ranges,
760 Self {
761 conjunctions: other_conds,
762 },
763 ));
764 }
765 }
766 }
767
768 Ok((
769 if scan_range.is_full_table_scan() {
770 vec![]
771 } else if table_desc.columns[table_desc.order_column_indices()[0]]
772 .data_type
773 .is_int()
774 {
775 match scan_range.split_small_range(max_split_range_gap) {
776 Some(scan_ranges) => scan_ranges,
777 None => vec![scan_range],
778 }
779 } else {
780 vec![scan_range]
781 },
782 Self {
783 conjunctions: other_conds,
784 },
785 ))
786 }
787
788 fn classify_conjunctions_by_pk(
792 conjunctions: Vec<ExprImpl>,
793 table_desc: &Rc<TableDesc>,
794 ) -> Vec<Vec<ExprImpl>> {
795 let pk_column_ids = &table_desc.order_column_indices();
796 let pk_cols_num = pk_column_ids.len();
797 let cols_num = table_desc.columns.len();
798
799 let mut col_idx_to_pk_idx = vec![None; cols_num];
800 pk_column_ids.iter().enumerate().for_each(|(idx, pk_idx)| {
801 col_idx_to_pk_idx[*pk_idx] = Some(idx);
802 });
803
804 let mut groups = vec![vec![]; pk_cols_num + 1];
805 for (key, group) in &conjunctions.into_iter().chunk_by(|expr| {
806 let input_bits = expr.collect_input_refs(cols_num);
807 if input_bits.count_ones(..) == 1 {
808 let col_idx = input_bits.ones().next().unwrap();
809 col_idx_to_pk_idx[col_idx].unwrap_or(pk_cols_num)
810 } else {
811 pk_cols_num
812 }
813 }) {
814 groups[key].extend(group);
815 }
816
817 groups
818 }
819
820 #[allow(clippy::type_complexity)]
828 fn analyze_group(
829 group: Vec<ExprImpl>,
830 ) -> Result<
831 Option<(
832 Vec<Bound<ScalarImpl>>,
833 Vec<Bound<ScalarImpl>>,
834 Vec<Option<ScalarImpl>>,
835 Vec<ExprImpl>,
836 )>,
837 > {
838 let mut lower_bound_conjunctions = vec![];
839 let mut upper_bound_conjunctions = vec![];
840 let mut eq_conds = vec![];
842 let mut other_conds = vec![];
843
844 for expr in group {
846 if let Some((input_ref, const_expr)) = expr.as_eq_const() {
847 let new_expr = if let Ok(expr) = const_expr
848 .clone()
849 .cast_implicit(input_ref.data_type.clone())
850 {
851 expr
852 } else {
853 match self::cast_compare::cast_compare_for_eq(const_expr, input_ref.data_type) {
854 Ok(ResultForEq::Success(expr)) => expr,
855 Ok(ResultForEq::NeverEqual) => {
856 return Ok(None);
857 }
858 Err(_) => {
859 other_conds.push(expr);
860 continue;
861 }
862 }
863 };
864
865 let Some(new_cond) = new_expr.fold_const()? else {
866 return Ok(None);
868 };
869 if Self::mutual_exclusive_with_eq_conds(&new_cond, &eq_conds) {
870 return Ok(None);
871 }
872 eq_conds = vec![Some(new_cond)];
873 } else if expr.as_is_null().is_some() {
874 if !eq_conds.is_empty() && eq_conds.into_iter().all(|l| l.is_some()) {
875 return Ok(None);
876 }
877 eq_conds = vec![None];
878 } else if let Some((input_ref, in_const_list)) = expr.as_in_const_list() {
879 let mut scalars = HashSet::new();
880 for const_expr in in_const_list {
881 let const_expr = const_expr
884 .cast_implicit(input_ref.data_type.clone())
885 .unwrap();
886 let value = const_expr.fold_const()?;
887 let Some(value) = value else {
888 continue;
889 };
890 scalars.insert(Some(value));
891 }
892 if scalars.is_empty() {
893 return Ok(None);
895 }
896 if !eq_conds.is_empty() {
897 scalars = scalars
898 .intersection(&HashSet::from_iter(eq_conds))
899 .cloned()
900 .collect();
901 if scalars.is_empty() {
902 return Ok(None);
903 }
904 }
905 eq_conds = scalars
907 .into_iter()
908 .sorted_by(DefaultOrd::default_cmp)
909 .collect();
910 } else if let Some((input_ref, op, const_expr)) = expr.as_comparison_const() {
911 let new_expr = if let Ok(expr) = const_expr
912 .clone()
913 .cast_implicit(input_ref.data_type.clone())
914 {
915 expr
916 } else {
917 match self::cast_compare::cast_compare_for_cmp(
918 const_expr,
919 input_ref.data_type,
920 op,
921 ) {
922 Ok(ResultForCmp::Success(expr)) => expr,
923 _ => {
924 other_conds.push(expr);
925 continue;
926 }
927 }
928 };
929 let Some(value) = new_expr.fold_const()? else {
930 return Ok(None);
932 };
933 match op {
934 ExprType::LessThan => {
935 upper_bound_conjunctions.push(Bound::Excluded(value));
936 }
937 ExprType::LessThanOrEqual => {
938 upper_bound_conjunctions.push(Bound::Included(value));
939 }
940 ExprType::GreaterThan => {
941 lower_bound_conjunctions.push(Bound::Excluded(value));
942 }
943 ExprType::GreaterThanOrEqual => {
944 lower_bound_conjunctions.push(Bound::Included(value));
945 }
946 _ => unreachable!(),
947 }
948 } else {
949 other_conds.push(expr);
950 }
951 }
952 Ok(Some((
953 lower_bound_conjunctions,
954 upper_bound_conjunctions,
955 eq_conds,
956 other_conds,
957 )))
958 }
959
960 fn mutual_exclusive_with_eq_conds(
961 new_conds: &ScalarImpl,
962 eq_conds: &[Option<ScalarImpl>],
963 ) -> bool {
964 !eq_conds.is_empty()
965 && eq_conds.iter().all(|l| {
966 if let Some(l) = l {
967 l != new_conds
968 } else {
969 true
970 }
971 })
972 }
973
974 fn merge_lower_bound_conjunctions(lb: Vec<Bound<ScalarImpl>>) -> Bound<ScalarImpl> {
975 lb.into_iter()
976 .max_by(|a, b| {
977 match (a, b) {
979 (Bound::Included(_), Bound::Unbounded) => std::cmp::Ordering::Greater,
980 (Bound::Excluded(_), Bound::Unbounded) => std::cmp::Ordering::Greater,
981 (Bound::Unbounded, Bound::Included(_)) => std::cmp::Ordering::Less,
982 (Bound::Unbounded, Bound::Excluded(_)) => std::cmp::Ordering::Less,
983 (Bound::Unbounded, Bound::Unbounded) => std::cmp::Ordering::Equal,
984 (Bound::Included(a), Bound::Included(b)) => a.default_cmp(b),
985 (Bound::Excluded(a), Bound::Excluded(b)) => a.default_cmp(b),
986 (Bound::Included(a), Bound::Excluded(b)) => match a.default_cmp(b) {
988 std::cmp::Ordering::Equal => std::cmp::Ordering::Less,
989 other => other,
990 },
991 (Bound::Excluded(a), Bound::Included(b)) => match a.default_cmp(b) {
992 std::cmp::Ordering::Equal => std::cmp::Ordering::Greater,
993 other => other,
994 },
995 }
996 })
997 .unwrap_or(Bound::Unbounded)
998 }
999
1000 fn merge_upper_bound_conjunctions(ub: Vec<Bound<ScalarImpl>>) -> Bound<ScalarImpl> {
1001 ub.into_iter()
1002 .min_by(|a, b| {
1003 match (a, b) {
1005 (Bound::Included(_), Bound::Unbounded) => std::cmp::Ordering::Less,
1006 (Bound::Excluded(_), Bound::Unbounded) => std::cmp::Ordering::Less,
1007 (Bound::Unbounded, Bound::Included(_)) => std::cmp::Ordering::Greater,
1008 (Bound::Unbounded, Bound::Excluded(_)) => std::cmp::Ordering::Greater,
1009 (Bound::Unbounded, Bound::Unbounded) => std::cmp::Ordering::Equal,
1010 (Bound::Included(a), Bound::Included(b)) => a.default_cmp(b),
1011 (Bound::Excluded(a), Bound::Excluded(b)) => a.default_cmp(b),
1012 (Bound::Included(a), Bound::Excluded(b)) => match a.default_cmp(b) {
1014 std::cmp::Ordering::Equal => std::cmp::Ordering::Greater,
1015 other => other,
1016 },
1017 (Bound::Excluded(a), Bound::Included(b)) => match a.default_cmp(b) {
1018 std::cmp::Ordering::Equal => std::cmp::Ordering::Less,
1019 other => other,
1020 },
1021 }
1022 })
1023 .unwrap_or(Bound::Unbounded)
1024 }
1025
1026 fn is_invalid_range(lower_bound: &Bound<ScalarImpl>, upper_bound: &Bound<ScalarImpl>) -> bool {
1027 match (lower_bound, upper_bound) {
1028 (Bound::Included(l), Bound::Included(u)) => l.default_cmp(u).is_gt(), (Bound::Included(l), Bound::Excluded(u)) => l.default_cmp(u).is_ge(), (Bound::Excluded(l), Bound::Included(u)) => l.default_cmp(u).is_ge(), (Bound::Excluded(l), Bound::Excluded(u)) => l.default_cmp(u).is_ge(), _ => false,
1033 }
1034 }
1035
1036 fn extract_eq_conds_within_range(
1037 eq_conds: Vec<Option<ScalarImpl>>,
1038 upper_bound: &Bound<ScalarImpl>,
1039 lower_bound: &Bound<ScalarImpl>,
1040 ) -> Vec<Option<ScalarImpl>> {
1041 if Self::is_invalid_range(lower_bound, upper_bound) {
1044 return vec![];
1045 }
1046
1047 let is_extract_null = upper_bound == &Bound::Unbounded && lower_bound == &Bound::Unbounded;
1048
1049 eq_conds
1050 .into_iter()
1051 .filter(|cond| {
1052 if let Some(cond) = cond {
1053 match lower_bound {
1054 Bound::Included(val) => {
1055 if cond.default_cmp(val).is_lt() {
1056 return false;
1058 }
1059 }
1060 Bound::Excluded(val) => {
1061 if cond.default_cmp(val).is_le() {
1062 return false;
1064 }
1065 }
1066 Bound::Unbounded => {}
1067 }
1068 match upper_bound {
1069 Bound::Included(val) => {
1070 if cond.default_cmp(val).is_gt() {
1071 return false;
1073 }
1074 }
1075 Bound::Excluded(val) => {
1076 if cond.default_cmp(val).is_ge() {
1077 return false;
1079 }
1080 }
1081 Bound::Unbounded => {}
1082 }
1083 true
1084 } else {
1085 is_extract_null
1086 }
1087 })
1088 .collect()
1089 }
1090
1091 #[must_use]
1097 pub fn group_by<F, const N: usize>(self, f: F) -> [Self; N]
1098 where
1099 F: Fn(&ExprImpl) -> usize,
1100 {
1101 const EMPTY: Vec<ExprImpl> = vec![];
1102 let mut groups = [EMPTY; N];
1103 for (key, group) in &self.conjunctions.into_iter().chunk_by(|expr| {
1104 let i = f(expr);
1106 assert!(i < N);
1107 i
1108 }) {
1109 groups[key].extend(group);
1110 }
1111
1112 groups.map(|group| Condition {
1113 conjunctions: group,
1114 })
1115 }
1116
1117 #[must_use]
1118 pub fn rewrite_expr(self, rewriter: &mut (impl ExprRewriter + ?Sized)) -> Self {
1119 Self {
1120 conjunctions: self
1121 .conjunctions
1122 .into_iter()
1123 .map(|expr| rewriter.rewrite_expr(expr))
1124 .collect(),
1125 }
1126 .simplify()
1127 }
1128
1129 pub fn visit_expr<V: ExprVisitor + ?Sized>(&self, visitor: &mut V) {
1130 self.conjunctions
1131 .iter()
1132 .for_each(|expr| visitor.visit_expr(expr));
1133 }
1134
1135 pub fn visit_expr_mut(&mut self, mutator: &mut (impl ExprMutator + ?Sized)) {
1136 self.conjunctions
1137 .iter_mut()
1138 .for_each(|expr| mutator.visit_expr(expr))
1139 }
1140
1141 fn simplify(self) -> Self {
1144 let conjunctions: Vec<_> = self
1146 .conjunctions
1147 .into_iter()
1148 .map(push_down_not)
1149 .map(fold_boolean_constant)
1150 .map(column_self_eq_eliminate)
1151 .flat_map(to_conjunctions)
1152 .collect();
1153 let mut res: Vec<ExprImpl> = Vec::new();
1154 let mut visited: HashSet<ExprImpl> = HashSet::new();
1155 for expr in conjunctions {
1156 if !expr.has_subquery() {
1158 let results_of_factorization = factorization_expr(expr);
1159 res.extend(
1160 results_of_factorization
1161 .clone()
1162 .into_iter()
1163 .filter(|expr| !visited.contains(expr)),
1164 );
1165 visited.extend(results_of_factorization);
1166 } else {
1167 res.push(expr);
1169 }
1170 }
1171 res.retain(|expr| {
1173 if let Some(v) = try_get_bool_constant(expr)
1174 && v
1175 {
1176 false
1177 } else {
1178 true
1179 }
1180 });
1181 for expr in &mut res {
1183 if let Some(v) = try_get_bool_constant(expr) {
1184 if !v {
1185 res.clear();
1186 res.push(ExprImpl::literal_bool(false));
1187 break;
1188 }
1189 }
1190 }
1191 Self { conjunctions: res }
1192 }
1193}
1194
1195pub struct ConditionDisplay<'a> {
1196 pub condition: &'a Condition,
1197 pub input_schema: &'a Schema,
1198}
1199
1200impl ConditionDisplay<'_> {
1201 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1202 if self.condition.always_true() {
1203 write!(f, "true")
1204 } else {
1205 write!(
1206 f,
1207 "{}",
1208 self.condition
1209 .conjunctions
1210 .iter()
1211 .format_with(" AND ", |expr, f| {
1212 f(&ExprDisplay {
1213 expr,
1214 input_schema: self.input_schema,
1215 })
1216 })
1217 )
1218 }
1219 }
1220}
1221
1222impl fmt::Display for ConditionDisplay<'_> {
1223 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1224 self.fmt(f)
1225 }
1226}
1227
1228impl fmt::Debug for ConditionDisplay<'_> {
1229 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1230 self.fmt(f)
1231 }
1232}
1233
1234mod cast_compare {
1239 use risingwave_common::types::DataType;
1240
1241 use crate::expr::{Expr, ExprImpl, ExprType};
1242
1243 enum ShrinkResult {
1244 OutUpperBound,
1245 OutLowerBound,
1246 InRange(ExprImpl),
1247 }
1248
1249 pub enum ResultForEq {
1250 Success(ExprImpl),
1251 NeverEqual,
1252 }
1253
1254 pub enum ResultForCmp {
1255 Success(ExprImpl),
1256 OutUpperBound,
1257 OutLowerBound,
1258 }
1259
1260 pub fn cast_compare_for_eq(const_expr: ExprImpl, target: DataType) -> Result<ResultForEq, ()> {
1261 match (const_expr.return_type(), &target) {
1262 (DataType::Int64, DataType::Int32)
1263 | (DataType::Int64, DataType::Int16)
1264 | (DataType::Int32, DataType::Int16) => match shrink_integral(const_expr, target)? {
1265 ShrinkResult::InRange(expr) => Ok(ResultForEq::Success(expr)),
1266 ShrinkResult::OutUpperBound | ShrinkResult::OutLowerBound => {
1267 Ok(ResultForEq::NeverEqual)
1268 }
1269 },
1270 _ => Err(()),
1271 }
1272 }
1273
1274 pub fn cast_compare_for_cmp(
1275 const_expr: ExprImpl,
1276 target: DataType,
1277 _op: ExprType,
1278 ) -> Result<ResultForCmp, ()> {
1279 match (const_expr.return_type(), &target) {
1280 (DataType::Int64, DataType::Int32)
1281 | (DataType::Int64, DataType::Int16)
1282 | (DataType::Int32, DataType::Int16) => match shrink_integral(const_expr, target)? {
1283 ShrinkResult::InRange(expr) => Ok(ResultForCmp::Success(expr)),
1284 ShrinkResult::OutUpperBound => Ok(ResultForCmp::OutUpperBound),
1285 ShrinkResult::OutLowerBound => Ok(ResultForCmp::OutLowerBound),
1286 },
1287 _ => Err(()),
1288 }
1289 }
1290
1291 fn shrink_integral(const_expr: ExprImpl, target: DataType) -> Result<ShrinkResult, ()> {
1292 let (upper_bound, lower_bound) = match (const_expr.return_type(), &target) {
1293 (DataType::Int64, DataType::Int32) => (i32::MAX as i64, i32::MIN as i64),
1294 (DataType::Int64, DataType::Int16) | (DataType::Int32, DataType::Int16) => {
1295 (i16::MAX as i64, i16::MIN as i64)
1296 }
1297 _ => unreachable!(),
1298 };
1299 match const_expr.fold_const().map_err(|_| ())? {
1300 Some(scalar) => {
1301 let value = scalar.as_integral();
1302 if value > upper_bound {
1303 Ok(ShrinkResult::OutUpperBound)
1304 } else if value < lower_bound {
1305 Ok(ShrinkResult::OutLowerBound)
1306 } else {
1307 Ok(ShrinkResult::InRange(
1308 const_expr.cast_explicit(target).unwrap(),
1309 ))
1310 }
1311 }
1312 None => Ok(ShrinkResult::InRange(
1313 const_expr.cast_explicit(target).unwrap(),
1314 )),
1315 }
1316 }
1317}
1318
1319#[cfg(test)]
1320mod tests {
1321 use rand::Rng;
1322
1323 use super::*;
1324
1325 #[test]
1326 fn test_split() {
1327 let left_col_num = 3;
1328 let right_col_num = 2;
1329
1330 let ty = DataType::Int32;
1331
1332 let mut rng = rand::rng();
1333
1334 let left: ExprImpl = FunctionCall::new(
1335 ExprType::LessThanOrEqual,
1336 vec![
1337 InputRef::new(rng.random_range(0..left_col_num), ty.clone()).into(),
1338 InputRef::new(rng.random_range(0..left_col_num), ty.clone()).into(),
1339 ],
1340 )
1341 .unwrap()
1342 .into();
1343
1344 let right: ExprImpl = FunctionCall::new(
1345 ExprType::LessThan,
1346 vec![
1347 InputRef::new(
1348 rng.random_range(left_col_num..left_col_num + right_col_num),
1349 ty.clone(),
1350 )
1351 .into(),
1352 InputRef::new(
1353 rng.random_range(left_col_num..left_col_num + right_col_num),
1354 ty.clone(),
1355 )
1356 .into(),
1357 ],
1358 )
1359 .unwrap()
1360 .into();
1361
1362 let other: ExprImpl = FunctionCall::new(
1363 ExprType::GreaterThan,
1364 vec![
1365 InputRef::new(rng.random_range(0..left_col_num), ty.clone()).into(),
1366 InputRef::new(
1367 rng.random_range(left_col_num..left_col_num + right_col_num),
1368 ty,
1369 )
1370 .into(),
1371 ],
1372 )
1373 .unwrap()
1374 .into();
1375
1376 let cond = Condition::with_expr(other.clone())
1377 .and(Condition::with_expr(right.clone()))
1378 .and(Condition::with_expr(left.clone()));
1379
1380 let res = cond.split(left_col_num, right_col_num);
1381
1382 assert_eq!(res.0.conjunctions, vec![left]);
1383 assert_eq!(res.1.conjunctions, vec![right]);
1384 assert_eq!(res.2.conjunctions, vec![other]);
1385 }
1386
1387 #[test]
1388 fn test_self_eq_eliminate() {
1389 let left_col_num = 3;
1390 let right_col_num = 2;
1391
1392 let ty = DataType::Int32;
1393
1394 let mut rng = rand::rng();
1395
1396 let x: ExprImpl = InputRef::new(rng.random_range(0..left_col_num), ty.clone()).into();
1397
1398 let left: ExprImpl = FunctionCall::new(ExprType::Equal, vec![x.clone(), x.clone()])
1399 .unwrap()
1400 .into();
1401
1402 let right: ExprImpl = FunctionCall::new(
1403 ExprType::LessThan,
1404 vec![
1405 InputRef::new(
1406 rng.random_range(left_col_num..left_col_num + right_col_num),
1407 ty.clone(),
1408 )
1409 .into(),
1410 InputRef::new(
1411 rng.random_range(left_col_num..left_col_num + right_col_num),
1412 ty.clone(),
1413 )
1414 .into(),
1415 ],
1416 )
1417 .unwrap()
1418 .into();
1419
1420 let other: ExprImpl = FunctionCall::new(
1421 ExprType::GreaterThan,
1422 vec![
1423 InputRef::new(rng.random_range(0..left_col_num), ty.clone()).into(),
1424 InputRef::new(
1425 rng.random_range(left_col_num..left_col_num + right_col_num),
1426 ty,
1427 )
1428 .into(),
1429 ],
1430 )
1431 .unwrap()
1432 .into();
1433
1434 let cond = Condition::with_expr(other.clone())
1435 .and(Condition::with_expr(right.clone()))
1436 .and(Condition::with_expr(left.clone()));
1437
1438 let res = cond.split(left_col_num, right_col_num);
1439
1440 let left_res = FunctionCall::new(ExprType::IsNotNull, vec![x])
1441 .unwrap()
1442 .into();
1443
1444 assert_eq!(res.0.conjunctions, vec![left_res]);
1445 assert_eq!(res.1.conjunctions, vec![right]);
1446 assert_eq!(res.2.conjunctions, vec![other]);
1447 }
1448}