Skip to main content

risingwave_frontend/optimizer/plan_node/
logical_join.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::ops::Deref;
17
18use fixedbitset::FixedBitSet;
19use itertools::{EitherOrBoth, Itertools};
20use pretty_xmlish::{Pretty, XmlNode};
21use risingwave_expr::bail;
22use risingwave_pb::expr::expr_node::PbType;
23use risingwave_pb::plan_common::{AsOfJoinDesc, JoinType, PbAsOfJoinInequalityType};
24use risingwave_sqlparser::ast::AsOf;
25
26use super::generic::{
27    GenericPlanNode, GenericPlanRef, push_down_into_join, push_down_join_condition,
28};
29use super::utils::{Distill, childless_record};
30use super::{
31    BackfillType, BatchPlanRef, ColPrunable, ExprRewritable, Logical, LogicalPlanRef as PlanRef,
32    PlanBase, PlanTreeNodeBinary, PredicatePushdown, StreamHashJoin, StreamPlanRef, StreamProject,
33    ToBatch, ToStream, generic, try_enforce_locality_requirement,
34};
35use crate::error::{ErrorCode, Result, RwError};
36use crate::expr::{CollectInputRef, Expr, ExprImpl, ExprRewriter, ExprType, ExprVisitor, InputRef};
37use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
38use crate::optimizer::plan_node::generic::DynamicFilter;
39use crate::optimizer::plan_node::stream_asof_join::StreamAsOfJoin;
40use crate::optimizer::plan_node::utils::IndicesDisplay;
41use crate::optimizer::plan_node::{
42    BatchHashJoin, BatchLookupJoin, BatchNestedLoopJoin, ColumnPruningContext, EqJoinPredicate,
43    LogicalFilter, LogicalScan, PredicatePushdownContext, RewriteStreamContext,
44    StreamDynamicFilter, StreamFilter, StreamTableScan, StreamTemporalJoin, ToStreamContext,
45};
46use crate::optimizer::plan_visitor::LogicalCardinalityExt;
47use crate::optimizer::property::{Distribution, RequiredDist};
48use crate::utils::{ColIndexMapping, ColIndexMappingRewriteExt, Condition, ConditionDisplay};
49
50/// `LogicalJoin` combines two relations according to some condition.
51///
52/// Each output row has fields from the left and right inputs. The set of output rows is a subset
53/// of the cartesian product of the two inputs; precisely which subset depends on the join
54/// condition. In addition, the output columns are a subset of the columns of the left and
55/// right columns, dependent on the output indices provided. A repeat output index is illegal.
56#[derive(Debug, Clone, PartialEq, Eq, Hash)]
57pub struct LogicalJoin {
58    pub base: PlanBase<Logical>,
59    core: generic::Join<PlanRef>,
60}
61
62impl Distill for LogicalJoin {
63    fn distill<'a>(&self) -> XmlNode<'a> {
64        let verbose = self.base.ctx().is_explain_verbose();
65        let mut vec = Vec::with_capacity(if verbose { 3 } else { 2 });
66        vec.push(("type", Pretty::debug(&self.join_type())));
67
68        let concat_schema = self.core.concat_schema();
69        let cond = Pretty::debug(&ConditionDisplay {
70            condition: self.on(),
71            input_schema: &concat_schema,
72        });
73        vec.push(("on", cond));
74
75        if verbose {
76            let data = IndicesDisplay::from_join(&self.core, &concat_schema);
77            vec.push(("output", data));
78        }
79
80        childless_record("LogicalJoin", vec)
81    }
82}
83
84impl LogicalJoin {
85    pub(crate) fn new(left: PlanRef, right: PlanRef, join_type: JoinType, on: Condition) -> Self {
86        let core = generic::Join::with_full_output(left, right, join_type, on);
87        Self::with_core(core)
88    }
89
90    pub(crate) fn with_output_indices(
91        left: PlanRef,
92        right: PlanRef,
93        join_type: JoinType,
94        on: Condition,
95        output_indices: Vec<usize>,
96    ) -> Self {
97        let core = generic::Join::new(left, right, on, join_type, output_indices);
98        Self::with_core(core)
99    }
100
101    pub fn with_core(core: generic::Join<PlanRef>) -> Self {
102        let base = PlanBase::new_logical_with_core(&core);
103        LogicalJoin { base, core }
104    }
105
106    pub fn create(
107        left: PlanRef,
108        right: PlanRef,
109        join_type: JoinType,
110        on_clause: ExprImpl,
111    ) -> PlanRef {
112        Self::new(left, right, join_type, Condition::with_expr(on_clause)).into()
113    }
114
115    pub fn internal_column_num(&self) -> usize {
116        self.core.internal_column_num()
117    }
118
119    pub fn i2l_col_mapping_ignore_join_type(&self) -> ColIndexMapping {
120        self.core.i2l_col_mapping_ignore_join_type()
121    }
122
123    pub fn i2r_col_mapping_ignore_join_type(&self) -> ColIndexMapping {
124        self.core.i2r_col_mapping_ignore_join_type()
125    }
126
127    /// Get a reference to the logical join's on.
128    pub fn on(&self) -> &Condition {
129        self.core
130            .on
131            .as_condition_ref()
132            .expect("logical join should store predicate as Condition")
133    }
134
135    pub fn core(&self) -> &generic::Join<PlanRef> {
136        &self.core
137    }
138
139    /// Collect all input ref in the on condition. And separate them into left and right.
140    pub fn input_idx_on_condition(&self) -> (Vec<usize>, Vec<usize>) {
141        let input_refs = self
142            .core
143            .on
144            .as_condition_ref()
145            .expect("logical join should store predicate as Condition")
146            .collect_input_refs(self.core.left.schema().len() + self.core.right.schema().len());
147        let index_group = input_refs
148            .ones()
149            .chunk_by(|i| *i < self.core.left.schema().len());
150        let left_index = index_group
151            .into_iter()
152            .next()
153            .map_or(vec![], |group| group.1.collect_vec());
154        let right_index = index_group.into_iter().next().map_or(vec![], |group| {
155            group
156                .1
157                .map(|i| i - self.core.left.schema().len())
158                .collect_vec()
159        });
160        (left_index, right_index)
161    }
162
163    /// Get the join type of the logical join.
164    pub fn join_type(&self) -> JoinType {
165        self.core.join_type
166    }
167
168    /// Get the eq join key of the logical join.
169    pub fn eq_indexes(&self) -> Vec<(usize, usize)> {
170        self.core.eq_indexes()
171    }
172
173    /// Get the output indices of the logical join.
174    pub fn output_indices(&self) -> &Vec<usize> {
175        &self.core.output_indices
176    }
177
178    /// Clone with new output indices
179    pub fn clone_with_output_indices(&self, output_indices: Vec<usize>) -> Self {
180        Self::with_core(generic::Join {
181            output_indices,
182            ..self.core.clone()
183        })
184    }
185
186    /// Clone with new `on` condition
187    pub fn clone_with_cond(&self, on: Condition) -> Self {
188        Self::with_core(generic::Join {
189            on: generic::JoinOn::Condition(on),
190            ..self.core.clone()
191        })
192    }
193
194    pub fn is_left_join(&self) -> bool {
195        matches!(self.join_type(), JoinType::LeftSemi | JoinType::LeftAnti)
196    }
197
198    pub fn is_right_join(&self) -> bool {
199        matches!(self.join_type(), JoinType::RightSemi | JoinType::RightAnti)
200    }
201
202    pub fn is_full_out(&self) -> bool {
203        self.core.is_full_out()
204    }
205
206    pub fn is_asof_join(&self) -> bool {
207        self.join_type() == JoinType::AsofInner || self.join_type() == JoinType::AsofLeftOuter
208    }
209
210    pub fn output_indices_are_trivial(&self) -> bool {
211        itertools::equal(
212            self.output_indices().iter().cloned(),
213            0..self.internal_column_num(),
214        )
215    }
216
217    /// Try to simplify the outer join with the predicate on the top of the join
218    ///
219    /// now it is just a naive implementation for comparison expression, we can give a more general
220    /// implementation with constant folding in future
221    fn simplify_outer(predicate: &Condition, left_col_num: usize, join_type: JoinType) -> JoinType {
222        let (mut gen_null_in_left, mut gen_null_in_right) = match join_type {
223            JoinType::LeftOuter => (false, true),
224            JoinType::RightOuter => (true, false),
225            JoinType::FullOuter => (true, true),
226            _ => return join_type,
227        };
228
229        for expr in &predicate.conjunctions {
230            if let ExprImpl::FunctionCall(func) = expr {
231                match func.func_type() {
232                    ExprType::Equal
233                    | ExprType::NotEqual
234                    | ExprType::LessThan
235                    | ExprType::LessThanOrEqual
236                    | ExprType::GreaterThan
237                    | ExprType::GreaterThanOrEqual => {
238                        for input in func.inputs() {
239                            if let ExprImpl::InputRef(input) = input {
240                                let idx = input.index;
241                                if idx < left_col_num {
242                                    gen_null_in_left = false;
243                                } else {
244                                    gen_null_in_right = false;
245                                }
246                            }
247                        }
248                    }
249                    _ => {}
250                };
251            }
252        }
253
254        match (gen_null_in_left, gen_null_in_right) {
255            (true, true) => JoinType::FullOuter,
256            (true, false) => JoinType::RightOuter,
257            (false, true) => JoinType::LeftOuter,
258            (false, false) => JoinType::Inner,
259        }
260    }
261
262    /// Index Join:
263    /// Try to convert logical join into batch lookup join and meanwhile it will do
264    /// the index selection for the lookup table so that we can benefit from indexes.
265    fn to_batch_lookup_join_with_index_selection(
266        &self,
267        predicate: EqJoinPredicate,
268        batch_join: generic::Join<BatchPlanRef>,
269    ) -> Result<Option<BatchLookupJoin>> {
270        match batch_join.join_type {
271            JoinType::Inner
272            | JoinType::LeftOuter
273            | JoinType::LeftSemi
274            | JoinType::LeftAnti
275            | JoinType::AsofInner
276            | JoinType::AsofLeftOuter => {}
277            _ => return Ok(None),
278        };
279
280        // Index selection for index join.
281        let right = self.right();
282        // Lookup Join only supports basic tables on the join's right side.
283        let logical_scan: &LogicalScan = if let Some(logical_scan) = right.as_logical_scan() {
284            logical_scan
285        } else {
286            return Ok(None);
287        };
288
289        let mut result_plan = None;
290        // Lookup primary table.
291        if let Some(lookup_join) =
292            self.to_batch_lookup_join(predicate.clone(), batch_join.clone())?
293        {
294            result_plan = Some(lookup_join);
295        }
296
297        if self
298            .core
299            .ctx()
300            .session_ctx()
301            .config()
302            .enable_index_selection()
303        {
304            let indexes = logical_scan.table_indexes();
305            for index in indexes {
306                if let Some(index_scan) = logical_scan.to_index_scan_if_index_covered(index) {
307                    let index_scan: PlanRef = index_scan.into();
308                    let that = self.clone_with_left_right(self.left(), index_scan.clone());
309                    let mut new_batch_join = batch_join.clone();
310                    new_batch_join.right =
311                        index_scan.to_batch().expect("index scan failed to batch");
312
313                    // Lookup covered index.
314                    if let Some(lookup_join) =
315                        that.to_batch_lookup_join(predicate.clone(), new_batch_join)?
316                    {
317                        match &result_plan {
318                            None => result_plan = Some(lookup_join),
319                            Some(prev_lookup_join) => {
320                                // Prefer to choose lookup join with longer lookup prefix len.
321                                if prev_lookup_join.lookup_prefix_len()
322                                    < lookup_join.lookup_prefix_len()
323                                {
324                                    result_plan = Some(lookup_join)
325                                }
326                            }
327                        }
328                    }
329                }
330            }
331        }
332
333        Ok(result_plan)
334    }
335
336    /// Try to convert logical join into batch lookup join.
337    fn to_batch_lookup_join(
338        &self,
339        predicate: EqJoinPredicate,
340        logical_join: generic::Join<BatchPlanRef>,
341    ) -> Result<Option<BatchLookupJoin>> {
342        let logical_scan: &LogicalScan =
343            if let Some(logical_scan) = self.core.right.as_logical_scan() {
344                logical_scan
345            } else {
346                return Ok(None);
347            };
348        Self::gen_batch_lookup_join(logical_scan, predicate, logical_join, self.is_asof_join())
349    }
350
351    pub fn gen_batch_lookup_join(
352        logical_scan: &LogicalScan,
353        predicate: EqJoinPredicate,
354        logical_join: generic::Join<BatchPlanRef>,
355        is_as_of: bool,
356    ) -> Result<Option<BatchLookupJoin>> {
357        match logical_join.join_type {
358            JoinType::Inner
359            | JoinType::LeftOuter
360            | JoinType::LeftSemi
361            | JoinType::LeftAnti
362            | JoinType::AsofInner
363            | JoinType::AsofLeftOuter => {}
364            _ => return Ok(None),
365        };
366
367        let table = logical_scan.table();
368        let output_column_ids = logical_scan.output_column_ids();
369
370        // Verify that the right join key columns are the the prefix of the primary key and
371        // also contain the distribution key.
372        let order_col_ids = table.order_column_ids();
373        let dist_key = table.distribution_key.clone();
374        // The at least prefix of order key that contains distribution key.
375        let mut dist_key_in_order_key_pos = vec![];
376        for d in dist_key {
377            let pos = table
378                .order_column_indices()
379                .position(|x| x == d)
380                .expect("dist_key must in order_key");
381            dist_key_in_order_key_pos.push(pos);
382        }
383        // The shortest prefix of order key that contains distribution key.
384        //
385        // If the lookup table has a singleton distribution (i.e. an empty distribution key),
386        // any non-empty prefix of the order key can be used for the lookup, so we require a
387        // prefix of length at least 1.
388        let shortest_prefix_len = dist_key_in_order_key_pos
389            .iter()
390            .max()
391            .map_or(1, |pos| pos + 1);
392
393        // Reorder the join equal predicate to match the order key.
394        let mut reorder_idx = Vec::with_capacity(shortest_prefix_len);
395        for order_col_id in order_col_ids {
396            let mut found = false;
397            for (i, eq_idx) in predicate.right_eq_indexes().into_iter().enumerate() {
398                if order_col_id == output_column_ids[eq_idx] {
399                    reorder_idx.push(i);
400                    found = true;
401                    break;
402                }
403            }
404            if !found {
405                break;
406            }
407        }
408        if reorder_idx.len() < shortest_prefix_len {
409            return Ok(None);
410        }
411        let lookup_prefix_len = reorder_idx.len();
412        let predicate = predicate.reorder(&reorder_idx);
413
414        // Extract the predicate from logical scan. Only pure scan is supported.
415        let (new_scan, scan_predicate, project_expr) = logical_scan.predicate_pull_up();
416        // Construct output column to require column mapping
417        let o2r = if let Some(project_expr) = project_expr {
418            project_expr
419                .into_iter()
420                .map(|x| x.as_input_ref().unwrap().index)
421                .collect_vec()
422        } else {
423            (0..logical_scan.output_col_idx().len()).collect_vec()
424        };
425        let left_schema_len = logical_join.left.schema().len();
426
427        let mut join_predicate_rewriter = LookupJoinPredicateRewriter {
428            offset: left_schema_len,
429            mapping: o2r.clone(),
430        };
431
432        let new_eq_cond = predicate
433            .eq_cond()
434            .rewrite_expr(&mut join_predicate_rewriter);
435
436        let mut scan_predicate_rewriter = LookupJoinScanPredicateRewriter {
437            offset: left_schema_len,
438        };
439
440        let new_other_cond = predicate
441            .other_cond()
442            .clone()
443            .rewrite_expr(&mut join_predicate_rewriter)
444            .and(scan_predicate.rewrite_expr(&mut scan_predicate_rewriter));
445
446        let new_join_on = new_eq_cond.and(new_other_cond);
447        let new_predicate =
448            EqJoinPredicate::create(left_schema_len, new_scan.schema().len(), new_join_on);
449
450        // We discovered that we cannot use a lookup join after pulling up the predicate
451        // from one side and simplifying the condition. Let's use some other join instead.
452        if !new_predicate.has_eq() {
453            return Ok(None);
454        }
455
456        // Rewrite the join output indices and all output indices referred to the old scan need to
457        // rewrite.
458        let new_join_output_indices = logical_join
459            .output_indices
460            .iter()
461            .map(|&x| {
462                if x < left_schema_len {
463                    x
464                } else {
465                    o2r[x - left_schema_len] + left_schema_len
466                }
467            })
468            .collect_vec();
469
470        let new_scan_output_column_ids = new_scan.output_column_ids();
471        let as_of = new_scan.as_of.clone();
472        let new_logical_scan: LogicalScan = new_scan.into();
473
474        // Construct a new logical join, because we have change its RHS.
475        let new_logical_join = generic::Join::new_with_eq_predicate(
476            logical_join.left,
477            new_logical_scan.to_batch()?,
478            new_predicate,
479            logical_join.join_type,
480            new_join_output_indices,
481        );
482
483        let asof_desc = is_as_of
484            .then(|| {
485                Self::get_inequality_desc_from_predicate(
486                    predicate.other_cond().clone(),
487                    left_schema_len,
488                )
489            })
490            .transpose()?;
491
492        Ok(Some(BatchLookupJoin::new(
493            new_logical_join,
494            table.clone(),
495            new_scan_output_column_ids,
496            lookup_prefix_len,
497            false,
498            as_of,
499            asof_desc,
500        )))
501    }
502
503    pub fn decompose(self) -> (PlanRef, PlanRef, Condition, JoinType, Vec<usize>) {
504        self.core.decompose()
505    }
506
507    fn dynamic_filter_candidate(&self, predicate: &Condition) -> Option<(usize, PbType)> {
508        // Dynamic filter only supports `Inner`/`LeftSemi`.
509        if !matches!(self.join_type(), JoinType::Inner | JoinType::LeftSemi) {
510            return None;
511        }
512
513        // Dynamic filter requires right side to be a scalar with one column.
514        if !self.right().max_one_row() || self.right().schema().len() != 1 {
515            return None;
516        }
517
518        // Dynamic filter only supports a single comparison predicate.
519        if predicate.conjunctions.len() > 1 {
520            return None;
521        }
522        let expr: ExprImpl = predicate.clone().into();
523        let (left_ref, comparator, right_ref) = expr.as_comparison_cond()?;
524
525        // Comparison must cross inputs: left input ref vs right scalar input ref.
526        let left_len = self.left().schema().len();
527        let condition_cross_inputs = left_ref.index < left_len && right_ref.index == left_len;
528        if !condition_cross_inputs {
529            return None;
530        }
531
532        // Comparison keys must be type aligned.
533        if self.left().schema().fields()[left_ref.index].data_type
534            != self.right().schema().fields()[0].data_type
535        {
536            return None;
537        }
538
539        // Dynamic filter output can only come from the left side.
540        if !self.output_indices().iter().all(|i| *i < left_len) {
541            return None;
542        }
543
544        Some((left_ref.index, comparator))
545    }
546
547    /// Check whether this join can be treated as a temporal filter for locality optimization.
548    ///
549    /// This is intentionally stricter than dynamic filter:
550    /// - right side must be a `LogicalNow`,
551    /// - and all dynamic filter preconditions must hold.
552    fn temporal_filter_candidate(&self) -> bool {
553        self.right().as_logical_now().is_some()
554            && self.dynamic_filter_candidate(self.on()).is_some()
555    }
556}
557
558impl PlanTreeNodeBinary<Logical> for LogicalJoin {
559    fn left(&self) -> PlanRef {
560        self.core.left.clone()
561    }
562
563    fn right(&self) -> PlanRef {
564        self.core.right.clone()
565    }
566
567    fn clone_with_left_right(&self, left: PlanRef, right: PlanRef) -> Self {
568        Self::with_core(generic::Join {
569            left,
570            right,
571            ..self.core.clone()
572        })
573    }
574
575    fn rewrite_with_left_right(
576        &self,
577        left: PlanRef,
578        left_col_change: ColIndexMapping,
579        right: PlanRef,
580        right_col_change: ColIndexMapping,
581    ) -> (Self, ColIndexMapping) {
582        let (new_on, new_output_indices) = {
583            let (mut map, _) = left_col_change.clone().into_parts();
584            let (mut right_map, _) = right_col_change.clone().into_parts();
585            for i in right_map.iter_mut().flatten() {
586                *i += left.schema().len();
587            }
588            map.append(&mut right_map);
589            let mut mapping = ColIndexMapping::new(map, left.schema().len() + right.schema().len());
590
591            let new_output_indices = self
592                .output_indices()
593                .iter()
594                .map(|&i| mapping.map(i))
595                .collect::<Vec<_>>();
596            let new_on = self.on().clone().rewrite_expr(&mut mapping);
597            (new_on, new_output_indices)
598        };
599
600        let join = Self::with_output_indices(
601            left,
602            right,
603            self.join_type(),
604            new_on,
605            new_output_indices.clone(),
606        );
607
608        let new_i2o = ColIndexMapping::with_remaining_columns(
609            &new_output_indices,
610            join.internal_column_num(),
611        );
612
613        let old_o2i = self.core.o2i_col_mapping();
614
615        let old_o2l = old_o2i
616            .composite(&self.core.i2l_col_mapping())
617            .composite(&left_col_change);
618        let old_o2r = old_o2i
619            .composite(&self.core.i2r_col_mapping())
620            .composite(&right_col_change);
621        let new_l2o = join.core.l2i_col_mapping().composite(&new_i2o);
622        let new_r2o = join.core.r2i_col_mapping().composite(&new_i2o);
623
624        let out_col_change = old_o2l
625            .composite(&new_l2o)
626            .union(&old_o2r.composite(&new_r2o));
627        (join, out_col_change)
628    }
629}
630
631impl_plan_tree_node_for_binary! { Logical, LogicalJoin }
632
633impl ColPrunable for LogicalJoin {
634    fn prune_col(&self, required_cols: &[usize], ctx: &mut ColumnPruningContext) -> PlanRef {
635        // make `required_cols` point to internal table instead of output schema.
636        let required_cols = required_cols
637            .iter()
638            .map(|i| self.output_indices()[*i])
639            .collect_vec();
640        let left_len = self.left().schema().fields.len();
641
642        let total_len = self.left().schema().len() + self.right().schema().len();
643        let mut resized_required_cols = FixedBitSet::with_capacity(total_len);
644
645        required_cols.iter().for_each(|&i| {
646            if self.is_right_join() {
647                resized_required_cols.insert(left_len + i);
648            } else {
649                resized_required_cols.insert(i);
650            }
651        });
652
653        // add those columns which are required in the join condition to
654        // to those that are required in the output
655        let mut visitor = CollectInputRef::new(resized_required_cols);
656        self.on().visit_expr(&mut visitor);
657        let left_right_required_cols = FixedBitSet::from(visitor).ones().collect_vec();
658
659        let mut left_required_cols = Vec::new();
660        let mut right_required_cols = Vec::new();
661        left_right_required_cols.iter().for_each(|&i| {
662            if i < left_len {
663                left_required_cols.push(i);
664            } else {
665                right_required_cols.push(i - left_len);
666            }
667        });
668
669        let mut on = self.on().clone();
670        let mut mapping =
671            ColIndexMapping::with_remaining_columns(&left_right_required_cols, total_len);
672        on = on.rewrite_expr(&mut mapping);
673
674        let new_output_indices = {
675            let required_inputs_in_output = if self.is_left_join() {
676                &left_required_cols
677            } else if self.is_right_join() {
678                &right_required_cols
679            } else {
680                &left_right_required_cols
681            };
682
683            let mapping =
684                ColIndexMapping::with_remaining_columns(required_inputs_in_output, total_len);
685            required_cols.iter().map(|&i| mapping.map(i)).collect_vec()
686        };
687
688        LogicalJoin::with_output_indices(
689            self.left().prune_col(&left_required_cols, ctx),
690            self.right().prune_col(&right_required_cols, ctx),
691            self.join_type(),
692            on,
693            new_output_indices,
694        )
695        .into()
696    }
697}
698
699impl ExprRewritable<Logical> for LogicalJoin {
700    fn has_rewritable_expr(&self) -> bool {
701        true
702    }
703
704    fn rewrite_exprs(&self, r: &mut dyn ExprRewriter) -> PlanRef {
705        let mut core = self.core.clone();
706        core.rewrite_exprs(r);
707        Self {
708            base: self.base.clone_with_new_plan_id(),
709            core,
710        }
711        .into()
712    }
713}
714
715impl ExprVisitable for LogicalJoin {
716    fn visit_exprs(&self, v: &mut dyn ExprVisitor) {
717        self.core.visit_exprs(v);
718    }
719}
720
721/// We are trying to derive a predicate to apply to the other side of a join if all
722/// the `InputRef`s in the predicate are eq condition columns, and can hence be substituted
723/// with the corresponding eq condition columns of the other side.
724///
725/// Strategy:
726/// 1. If the function is pure except for any `InputRef` (which may refer to impure computation),
727///    then we proceed. Else abort.
728/// 2. Then, we collect `InputRef`s in the conjunction.
729/// 3. If they are all columns in the given side of join eq condition, then we proceed. Else abort.
730/// 4. We then rewrite the `ExprImpl`, by replacing `InputRef` column indices with the equivalent in
731///    the other side.
732///
733/// # Arguments
734///
735/// Suppose we derive a predicate from the left side to be pushed to the right side.
736/// * `expr`: An expr from the left side.
737/// * `col_num`: The number of columns in the left side.
738fn derive_predicate_from_eq_condition(
739    expr: &ExprImpl,
740    eq_condition: &EqJoinPredicate,
741    col_num: usize,
742    expr_is_left: bool,
743) -> Option<ExprImpl> {
744    if expr.is_impure() {
745        return None;
746    }
747    let eq_indices = eq_condition
748        .eq_indexes_typed()
749        .iter()
750        .filter_map(|(l, r)| {
751            if l.return_type() != r.return_type() {
752                None
753            } else if expr_is_left {
754                Some(l.index())
755            } else {
756                Some(r.index())
757            }
758        })
759        .collect_vec();
760    if expr
761        .collect_input_refs(col_num)
762        .ones()
763        .any(|index| !eq_indices.contains(&index))
764    {
765        // expr contains an InputRef not in eq_condition
766        return None;
767    }
768    // The function is pure except for `InputRef` and all `InputRef`s are `eq_condition` indices.
769    // Hence, we can substitute those `InputRef`s with indices from the other side.
770    let other_side_mapping = if expr_is_left {
771        eq_condition.eq_indexes_typed().into_iter().collect()
772    } else {
773        eq_condition
774            .eq_indexes_typed()
775            .into_iter()
776            .map(|(x, y)| (y, x))
777            .collect()
778    };
779    struct InputRefsRewriter {
780        mapping: HashMap<InputRef, InputRef>,
781    }
782    impl ExprRewriter for InputRefsRewriter {
783        fn rewrite_input_ref(&mut self, input_ref: InputRef) -> ExprImpl {
784            self.mapping[&input_ref].clone().into()
785        }
786    }
787    Some(
788        InputRefsRewriter {
789            mapping: other_side_mapping,
790        }
791        .rewrite_expr(expr.clone()),
792    )
793}
794
795/// Rewrite the join predicate and all columns referred to the scan side need to rewrite.
796struct LookupJoinPredicateRewriter {
797    offset: usize,
798    mapping: Vec<usize>,
799}
800impl ExprRewriter for LookupJoinPredicateRewriter {
801    fn rewrite_input_ref(&mut self, input_ref: InputRef) -> ExprImpl {
802        if input_ref.index() < self.offset {
803            input_ref.into()
804        } else {
805            InputRef::new(
806                self.mapping[input_ref.index() - self.offset] + self.offset,
807                input_ref.return_type(),
808            )
809            .into()
810        }
811    }
812}
813
814/// Rewrite the scan predicate so we can add it to the join predicate.
815struct LookupJoinScanPredicateRewriter {
816    offset: usize,
817}
818impl ExprRewriter for LookupJoinScanPredicateRewriter {
819    fn rewrite_input_ref(&mut self, input_ref: InputRef) -> ExprImpl {
820        InputRef::new(input_ref.index() + self.offset, input_ref.return_type()).into()
821    }
822}
823
824impl PredicatePushdown for LogicalJoin {
825    /// Pushes predicates above and within a join node into the join node and/or its children nodes.
826    ///
827    /// # Which predicates can be pushed
828    ///
829    /// For inner join, we can do all kinds of pushdown.
830    ///
831    /// For left/right semi join, we can push filter to left/right and on-clause,
832    /// and push on-clause to left/right.
833    ///
834    /// For left/right anti join, we can push filter to left/right, but on-clause can not be pushed
835    ///
836    /// ## Outer Join
837    ///
838    /// Preserved Row table
839    /// : The table in an Outer Join that must return all rows.
840    ///
841    /// Null Supplying table
842    /// : This is the table that has nulls filled in for its columns in unmatched rows.
843    ///
844    /// |                          | Preserved Row table | Null Supplying table |
845    /// |--------------------------|---------------------|----------------------|
846    /// | Join predicate (on)      | Not Pushed          | Pushed               |
847    /// | Where predicate (filter) | Pushed              | Not Pushed           |
848    fn predicate_pushdown(
849        &self,
850        predicate: Condition,
851        ctx: &mut PredicatePushdownContext,
852    ) -> PlanRef {
853        // rewrite output col referencing indices as internal cols
854        let mut predicate = {
855            let mut mapping = self.core.o2i_col_mapping();
856            predicate.rewrite_expr(&mut mapping)
857        };
858
859        let left_col_num = self.left().schema().len();
860        let right_col_num = self.right().schema().len();
861        let join_type = LogicalJoin::simplify_outer(&predicate, left_col_num, self.join_type());
862
863        let push_down_temporal_predicate = self.temporal_join_on().is_none();
864
865        let (left_from_filter, right_from_filter, on) = push_down_into_join(
866            &mut predicate,
867            left_col_num,
868            right_col_num,
869            join_type,
870            push_down_temporal_predicate,
871        );
872
873        let mut new_on = self.on().clone().and(on);
874        let (left_from_on, right_from_on) = push_down_join_condition(
875            &mut new_on,
876            left_col_num,
877            right_col_num,
878            join_type,
879            push_down_temporal_predicate,
880        );
881
882        let left_predicate = left_from_filter.and(left_from_on);
883        let right_predicate = right_from_filter.and(right_from_on);
884
885        // Derive conditions to push to the other side based on eq condition columns
886        let eq_condition = EqJoinPredicate::create(left_col_num, right_col_num, new_on.clone());
887
888        // Only push to RHS if RHS is inner side of a join (RHS requires match on LHS)
889        let right_from_left = if matches!(
890            join_type,
891            JoinType::Inner | JoinType::LeftOuter | JoinType::RightSemi | JoinType::LeftSemi
892        ) {
893            Condition {
894                conjunctions: left_predicate
895                    .conjunctions
896                    .iter()
897                    .filter_map(|expr| {
898                        derive_predicate_from_eq_condition(expr, &eq_condition, left_col_num, true)
899                    })
900                    .collect(),
901            }
902        } else {
903            Condition::true_cond()
904        };
905
906        // Only push to LHS if LHS is inner side of a join (LHS requires match on RHS)
907        let left_from_right = if matches!(
908            join_type,
909            JoinType::Inner | JoinType::RightOuter | JoinType::LeftSemi | JoinType::RightSemi
910        ) {
911            Condition {
912                conjunctions: right_predicate
913                    .conjunctions
914                    .iter()
915                    .filter_map(|expr| {
916                        derive_predicate_from_eq_condition(
917                            expr,
918                            &eq_condition,
919                            right_col_num,
920                            false,
921                        )
922                    })
923                    .collect(),
924            }
925        } else {
926            Condition::true_cond()
927        };
928
929        let left_predicate = left_predicate.and(left_from_right);
930        let right_predicate = right_predicate.and(right_from_left);
931
932        let new_left = self.left().predicate_pushdown(left_predicate, ctx);
933        let new_right = self.right().predicate_pushdown(right_predicate, ctx);
934        let new_join = LogicalJoin::with_output_indices(
935            new_left,
936            new_right,
937            join_type,
938            new_on,
939            self.output_indices().clone(),
940        );
941
942        let mut mapping = self.core.i2o_col_mapping();
943        predicate = predicate.rewrite_expr(&mut mapping);
944        LogicalFilter::create(new_join.into(), predicate)
945    }
946}
947
948#[derive(Clone, Copy)]
949struct TemporalJoinScan<'a>(&'a LogicalScan);
950
951impl<'a> Deref for TemporalJoinScan<'a> {
952    type Target = LogicalScan;
953
954    fn deref(&self) -> &Self::Target {
955        self.0
956    }
957}
958
959impl LogicalJoin {
960    fn get_stream_input_for_hash_join(
961        &self,
962        predicate: &EqJoinPredicate,
963        ctx: &mut ToStreamContext,
964    ) -> Result<(StreamPlanRef, StreamPlanRef)> {
965        use super::stream::prelude::*;
966
967        let mut right = self.right().to_stream_with_dist_required(
968            &RequiredDist::shard_by_key(self.right().schema().len(), &predicate.right_eq_indexes()),
969            ctx,
970        )?;
971        let r2l =
972            predicate.r2l_eq_columns_mapping(self.left().schema().len(), right.schema().len());
973        let l2r =
974            predicate.l2r_eq_columns_mapping(self.left().schema().len(), right.schema().len());
975        let mut left;
976        let right_dist = right.distribution();
977        match right_dist {
978            Distribution::HashShard(_) => {
979                let left_dist = r2l
980                    .rewrite_required_distribution(&RequiredDist::PhysicalDist(right_dist.clone()));
981                left = self.left().to_stream_with_dist_required(&left_dist, ctx)?;
982            }
983            Distribution::UpstreamHashShard(_, _) => {
984                left = self.left().to_stream_with_dist_required(
985                    &RequiredDist::shard_by_key(
986                        self.left().schema().len(),
987                        &predicate.left_eq_indexes(),
988                    ),
989                    ctx,
990                )?;
991                let left_dist = left.distribution();
992                match left_dist {
993                    Distribution::HashShard(_) => {
994                        let right_dist = l2r.rewrite_required_distribution(
995                            &RequiredDist::PhysicalDist(left_dist.clone()),
996                        );
997                        right = right_dist.streaming_enforce_if_not_satisfies(right)?
998                    }
999                    Distribution::UpstreamHashShard(_, _) => {
1000                        left = RequiredDist::hash_shard(&predicate.left_eq_indexes())
1001                            .streaming_enforce_if_not_satisfies(left)?;
1002                        right = RequiredDist::hash_shard(&predicate.right_eq_indexes())
1003                            .streaming_enforce_if_not_satisfies(right)?;
1004                    }
1005                    _ => unreachable!(),
1006                }
1007            }
1008            _ => unreachable!(),
1009        }
1010        Ok((left, right))
1011    }
1012
1013    fn to_stream_hash_join(
1014        &self,
1015        predicate: EqJoinPredicate,
1016        ctx: &mut ToStreamContext,
1017    ) -> Result<StreamPlanRef> {
1018        use super::stream::prelude::*;
1019
1020        assert!(predicate.has_eq());
1021        let (left, right) = self.get_stream_input_for_hash_join(&predicate, ctx)?;
1022
1023        let mut core = self.core.clone_with_inputs(left, right);
1024        core.on = generic::JoinOn::EqPredicate(predicate);
1025
1026        // Convert to Hash Join for equal joins
1027        // For inner joins, pull non-equal conditions to a filter operator on top of it by default.
1028        // We do so as the filter operator can apply the non-equal condition batch-wise (vectorized)
1029        // as opposed to the HashJoin, which applies the condition row-wise.
1030        // However, the default behavior of pulling up non-equal conditions can be overridden by the
1031        // session variable `streaming_force_filter_inside_join` as it can save unnecessary
1032        // materialization of rows only to be filtered later.
1033
1034        let stream_hash_join = StreamHashJoin::new(core.clone())?;
1035        let predicate = stream_hash_join.eq_join_predicate().clone();
1036
1037        let force_filter_inside_join = self
1038            .base
1039            .ctx()
1040            .session_ctx()
1041            .config()
1042            .streaming_force_filter_inside_join();
1043
1044        let pull_filter = self.join_type() == JoinType::Inner
1045            && stream_hash_join.eq_join_predicate().has_non_eq()
1046            && stream_hash_join.inequality_pairs().is_empty()
1047            && (!force_filter_inside_join);
1048        if pull_filter {
1049            let default_indices = (0..self.internal_column_num()).collect::<Vec<_>>();
1050
1051            let mut core = core;
1052            core.output_indices = default_indices.clone();
1053            // Temporarily remove output indices.
1054            let eq_cond = EqJoinPredicate::new(
1055                Condition::true_cond(),
1056                predicate.eq_keys().to_vec(),
1057                self.left().schema().len(),
1058                self.right().schema().len(),
1059            );
1060            core.on = generic::JoinOn::EqPredicate(eq_cond);
1061            let hash_join = StreamHashJoin::new(core)?.into();
1062            let logical_filter = generic::Filter::new(predicate.non_eq_cond(), hash_join);
1063            let plan = StreamFilter::new(logical_filter).into();
1064            if self.output_indices() != &default_indices {
1065                let logical_project = generic::Project::with_mapping(
1066                    plan,
1067                    ColIndexMapping::with_remaining_columns(
1068                        self.output_indices(),
1069                        self.internal_column_num(),
1070                    ),
1071                );
1072                Ok(StreamProject::new(logical_project).into())
1073            } else {
1074                Ok(plan)
1075            }
1076        } else {
1077            Ok(stream_hash_join.into())
1078        }
1079    }
1080
1081    pub fn should_be_temporal_join(&self) -> bool {
1082        self.temporal_join_on().is_some()
1083    }
1084
1085    fn temporal_join_on(&self) -> Option<TemporalJoinScan<'_>> {
1086        if let Some(logical_scan) = self.core.right.as_logical_scan() {
1087            matches!(logical_scan.as_of(), Some(AsOf::ProcessTime))
1088                .then_some(TemporalJoinScan(logical_scan))
1089        } else {
1090            None
1091        }
1092    }
1093
1094    fn should_be_stream_temporal_join<'a>(
1095        &'a self,
1096        ctx: &ToStreamContext,
1097    ) -> Result<Option<TemporalJoinScan<'a>>> {
1098        Ok(if let Some(scan) = self.temporal_join_on() {
1099            if ctx.backfill_type().is_snapshot_backfill() {
1100                return Err(RwError::from(ErrorCode::NotSupported(
1101                    "Temporal join with snapshot backfill not supported".into(),
1102                    "Please use arrangement backfill".into(),
1103                )));
1104            }
1105            if scan.cross_database() {
1106                return Err(RwError::from(ErrorCode::NotSupported(
1107                        "Temporal join requires the lookup table to be in the same database as the stream source table".into(),
1108                        "Please ensure both tables are in the same database".into(),
1109                    )));
1110            }
1111            Some(scan)
1112        } else {
1113            None
1114        })
1115    }
1116
1117    fn to_stream_temporal_join_with_index_selection(
1118        &self,
1119        logical_scan: TemporalJoinScan<'_>,
1120        predicate: EqJoinPredicate,
1121        ctx: &mut ToStreamContext,
1122    ) -> Result<StreamPlanRef> {
1123        // Use primary table.
1124        let mut result_plan: Result<StreamTemporalJoin> =
1125            self.to_stream_temporal_join(logical_scan, predicate.clone(), ctx);
1126        // Return directly if this temporal join can match the pk of its right table.
1127        if let Ok(temporal_join) = &result_plan
1128            && temporal_join.eq_join_predicate().eq_indexes().len()
1129                == logical_scan.primary_key().len()
1130        {
1131            return result_plan.map(|x| x.into());
1132        }
1133        if self
1134            .core
1135            .ctx()
1136            .session_ctx()
1137            .config()
1138            .enable_index_selection()
1139        {
1140            let indexes = logical_scan.table_indexes();
1141            for index in indexes {
1142                // Use index table
1143                if let Some(index_scan) = logical_scan.to_index_scan_if_index_covered(index) {
1144                    let index_scan: PlanRef = index_scan.into();
1145                    let that = self.clone_with_left_right(self.left(), index_scan.clone());
1146                    if let Ok(temporal_join) = that.to_stream_temporal_join(
1147                        that.temporal_join_on().expect(
1148                            "index scan created from temporal join scan must also be temporal join",
1149                        ),
1150                        predicate.clone(),
1151                        ctx,
1152                    ) {
1153                        match &result_plan {
1154                            Err(_) => result_plan = Ok(temporal_join),
1155                            Ok(prev_temporal_join) => {
1156                                // Prefer to the temporal join with a longer lookup prefix len.
1157                                if prev_temporal_join.eq_join_predicate().eq_indexes().len()
1158                                    < temporal_join.eq_join_predicate().eq_indexes().len()
1159                                {
1160                                    result_plan = Ok(temporal_join)
1161                                }
1162                            }
1163                        }
1164                    }
1165                }
1166            }
1167        }
1168
1169        result_plan.map(|x| x.into())
1170    }
1171
1172    fn temporal_join_scan_predicate_pull_up(
1173        logical_scan: TemporalJoinScan<'_>,
1174        predicate: EqJoinPredicate,
1175        output_indices: &[usize],
1176        left_schema_len: usize,
1177    ) -> Result<(StreamTableScan, EqJoinPredicate, Condition, Vec<usize>)> {
1178        // Extract the predicate from logical scan. Only pure scan is supported.
1179        let (new_scan, scan_predicate, project_expr) = logical_scan.predicate_pull_up();
1180        // Construct output column to require column mapping
1181        let o2r = if let Some(project_expr) = project_expr {
1182            project_expr
1183                .into_iter()
1184                .map(|x| x.as_input_ref().unwrap().index)
1185                .collect_vec()
1186        } else {
1187            (0..logical_scan.output_col_idx().len()).collect_vec()
1188        };
1189        let mut join_predicate_rewriter = LookupJoinPredicateRewriter {
1190            offset: left_schema_len,
1191            mapping: o2r.clone(),
1192        };
1193
1194        let new_eq_cond = predicate
1195            .eq_cond()
1196            .rewrite_expr(&mut join_predicate_rewriter);
1197
1198        let mut scan_predicate_rewriter = LookupJoinScanPredicateRewriter {
1199            offset: left_schema_len,
1200        };
1201
1202        let new_other_cond = predicate
1203            .other_cond()
1204            .clone()
1205            .rewrite_expr(&mut join_predicate_rewriter)
1206            .and(scan_predicate.rewrite_expr(&mut scan_predicate_rewriter));
1207
1208        let new_join_on = new_eq_cond.and(new_other_cond);
1209
1210        let new_predicate = EqJoinPredicate::create(
1211            left_schema_len,
1212            new_scan.schema().len(),
1213            new_join_on.clone(),
1214        );
1215
1216        // Rewrite the join output indices and all output indices referred to the old scan need to
1217        // rewrite.
1218        let new_join_output_indices = output_indices
1219            .iter()
1220            .map(|&x| {
1221                if x < left_schema_len {
1222                    x
1223                } else {
1224                    o2r[x - left_schema_len] + left_schema_len
1225                }
1226            })
1227            .collect_vec();
1228
1229        let new_stream_table_scan =
1230            StreamTableScan::new_with_backfill_type(new_scan, BackfillType::Replicated);
1231        Ok((
1232            new_stream_table_scan,
1233            new_predicate,
1234            new_join_on,
1235            new_join_output_indices,
1236        ))
1237    }
1238
1239    fn to_stream_temporal_join(
1240        &self,
1241        logical_scan: TemporalJoinScan<'_>,
1242        predicate: EqJoinPredicate,
1243        ctx: &mut ToStreamContext,
1244    ) -> Result<StreamTemporalJoin> {
1245        use super::stream::prelude::*;
1246
1247        assert!(predicate.has_eq());
1248
1249        let table = logical_scan.table();
1250        let output_column_ids = logical_scan.output_column_ids();
1251
1252        // Verify that the right join key columns are the the prefix of the primary key and
1253        // also contain the distribution key.
1254        let order_col_ids = table.order_column_ids();
1255        let dist_key = table.distribution_key.clone();
1256
1257        let mut dist_key_in_order_key_pos = vec![];
1258        for d in dist_key {
1259            let pos = table
1260                .order_column_indices()
1261                .position(|x| x == d)
1262                .expect("dist_key must in order_key");
1263            dist_key_in_order_key_pos.push(pos);
1264        }
1265        // The shortest prefix of order key that contains distribution key.
1266        let shortest_prefix_len = dist_key_in_order_key_pos
1267            .iter()
1268            .max()
1269            .map_or(0, |pos| pos + 1);
1270
1271        // Reorder the join equal predicate to match the order key.
1272        let mut reorder_idx = Vec::with_capacity(shortest_prefix_len);
1273        for order_col_id in order_col_ids {
1274            let mut found = false;
1275            for (i, eq_idx) in predicate.right_eq_indexes().into_iter().enumerate() {
1276                if order_col_id == output_column_ids[eq_idx] {
1277                    reorder_idx.push(i);
1278                    found = true;
1279                    break;
1280                }
1281            }
1282            if !found {
1283                break;
1284            }
1285        }
1286        if reorder_idx.len() < shortest_prefix_len {
1287            return Err(RwError::from(ErrorCode::NotSupported(
1288                "Temporal join requires the equivalence join condition includes the key columns that form the distribution key of the lookup table".into(),
1289                concat!(
1290                    "Use DESCRIBE <table_name> to view the table's key information.\n",
1291                    "You can create an index on the lookup table to facilitate the temporal join if necessary."
1292                ).into(),
1293            )));
1294        }
1295        let lookup_prefix_len = reorder_idx.len();
1296        let predicate = predicate.reorder(&reorder_idx);
1297
1298        let required_dist = if dist_key_in_order_key_pos.is_empty() {
1299            RequiredDist::single()
1300        } else {
1301            let left_eq_indexes = predicate.left_eq_indexes();
1302            let left_dist_key = dist_key_in_order_key_pos
1303                .iter()
1304                .map(|pos| left_eq_indexes[*pos])
1305                .collect_vec();
1306
1307            RequiredDist::hash_shard(&left_dist_key)
1308        };
1309
1310        let left = self.left().to_stream(ctx)?;
1311        // Enforce a shuffle for the temporal join LHS to let the scheduler be able to schedule the join fragment together with the RHS with a `no_shuffle` exchange.
1312        let left = required_dist.stream_enforce(left);
1313
1314        let (new_stream_table_scan, new_predicate, new_join_on, new_join_output_indices) =
1315            Self::temporal_join_scan_predicate_pull_up(
1316                logical_scan,
1317                predicate,
1318                self.output_indices(),
1319                self.left().schema().len(),
1320            )?;
1321
1322        let right = RequiredDist::no_shuffle(new_stream_table_scan.into());
1323        if !new_predicate.has_eq() {
1324            return Err(RwError::from(ErrorCode::NotSupported(
1325                "Temporal join requires a non trivial join condition".into(),
1326                "Please remove the false condition of the join".into(),
1327            )));
1328        }
1329
1330        // Construct a new logical join, because we have change its RHS.
1331        let new_logical_join = generic::Join::new(
1332            left,
1333            right,
1334            new_join_on,
1335            self.join_type(),
1336            new_join_output_indices,
1337        );
1338
1339        let new_predicate = new_predicate.retain_prefix_eq_key(lookup_prefix_len);
1340
1341        let mut new_logical_join = new_logical_join;
1342        new_logical_join.on = generic::JoinOn::EqPredicate(new_predicate);
1343        StreamTemporalJoin::new(new_logical_join, false)
1344    }
1345
1346    fn to_stream_nested_loop_temporal_join(
1347        &self,
1348        logical_scan: TemporalJoinScan<'_>,
1349        predicate: EqJoinPredicate,
1350        ctx: &mut ToStreamContext,
1351    ) -> Result<StreamPlanRef> {
1352        use super::stream::prelude::*;
1353        assert!(!predicate.has_eq());
1354
1355        let left = self.left().to_stream_with_dist_required(
1356            &RequiredDist::PhysicalDist(Distribution::Broadcast),
1357            ctx,
1358        )?;
1359        assert!(left.as_stream_exchange().is_some());
1360
1361        if self.join_type() != JoinType::Inner {
1362            return Err(RwError::from(ErrorCode::NotSupported(
1363                "Temporal join requires an inner join".into(),
1364                "Please use an inner join".into(),
1365            )));
1366        }
1367
1368        if !left.append_only() {
1369            return Err(RwError::from(ErrorCode::NotSupported(
1370                "Nested-loop Temporal join requires the left hash side to be append only".into(),
1371                "Please ensure the left hash side is append only".into(),
1372            )));
1373        }
1374
1375        let (new_stream_table_scan, new_predicate, new_join_on, new_join_output_indices) =
1376            Self::temporal_join_scan_predicate_pull_up(
1377                logical_scan,
1378                predicate,
1379                self.output_indices(),
1380                self.left().schema().len(),
1381            )?;
1382
1383        let right = RequiredDist::no_shuffle(new_stream_table_scan.into());
1384
1385        // Construct a new logical join, because we have change its RHS.
1386        let new_logical_join = generic::Join::new(
1387            left,
1388            right,
1389            new_join_on,
1390            self.join_type(),
1391            new_join_output_indices,
1392        );
1393
1394        let mut new_logical_join = new_logical_join;
1395        new_logical_join.on = generic::JoinOn::EqPredicate(new_predicate);
1396        Ok(StreamTemporalJoin::new(new_logical_join, true)?.into())
1397    }
1398
1399    fn to_stream_dynamic_filter(
1400        &self,
1401        predicate: Condition,
1402        ctx: &mut ToStreamContext,
1403    ) -> Result<Option<StreamPlanRef>> {
1404        use super::stream::prelude::*;
1405
1406        // If there is exactly one predicate, it is a comparison (<, <=, >, >=), and the
1407        // join is a `Inner` or `LeftSemi` join, we can convert the scalar subquery into a
1408        // `StreamDynamicFilter`
1409        let Some((left_key_idx, comparator)) = self.dynamic_filter_candidate(&predicate) else {
1410            return Ok(None);
1411        };
1412
1413        let left = self.left().to_stream(ctx)?.enforce_concrete_distribution();
1414        let right = self.right().to_stream_with_dist_required(
1415            &RequiredDist::PhysicalDist(Distribution::Broadcast),
1416            ctx,
1417        )?;
1418
1419        assert!(right.as_stream_exchange().is_some());
1420        assert_eq!(
1421            *Itertools::exactly_one(right.inputs().iter())
1422                .unwrap()
1423                .distribution(),
1424            Distribution::Single
1425        );
1426
1427        let core = DynamicFilter::new(comparator, left_key_idx, left, right);
1428        let plan = StreamDynamicFilter::new(core)?.into();
1429        // TODO: `DynamicFilterExecutor` should support `output_indices` in `ChunkBuilder`
1430        if self
1431            .output_indices()
1432            .iter()
1433            .copied()
1434            .ne(0..self.left().schema().len())
1435        {
1436            // The schema of dynamic filter is always the same as the left side now, and we have
1437            // checked that all output columns are from the left side before.
1438            let logical_project = generic::Project::with_mapping(
1439                plan,
1440                ColIndexMapping::with_remaining_columns(
1441                    self.output_indices(),
1442                    self.left().schema().len(),
1443                ),
1444            );
1445            Ok(Some(StreamProject::new(logical_project).into()))
1446        } else {
1447            Ok(Some(plan))
1448        }
1449    }
1450
1451    pub fn index_lookup_join_to_batch_lookup_join(&self) -> Result<Option<BatchPlanRef>> {
1452        let predicate = EqJoinPredicate::create(
1453            self.left().schema().len(),
1454            self.right().schema().len(),
1455            self.on().clone(),
1456        );
1457        assert!(predicate.has_eq());
1458
1459        let join = self
1460            .core
1461            .clone_with_inputs(self.core.left.to_batch()?, self.core.right.to_batch()?);
1462
1463        Ok(self.to_batch_lookup_join(predicate, join)?.map(Into::into))
1464    }
1465
1466    fn to_stream_asof_join(
1467        &self,
1468        predicate: EqJoinPredicate,
1469        ctx: &mut ToStreamContext,
1470    ) -> Result<StreamPlanRef> {
1471        use super::stream::prelude::*;
1472
1473        if predicate.eq_keys().is_empty() {
1474            return Err(ErrorCode::InvalidInputSyntax(
1475                "AsOf join requires at least 1 equal condition".to_owned(),
1476            )
1477            .into());
1478        }
1479
1480        let (left, right) = self.get_stream_input_for_hash_join(&predicate, ctx)?;
1481        let left_len = left.schema().len();
1482        let mut core = self.core.clone_with_inputs(left, right);
1483        core.on = generic::JoinOn::EqPredicate(predicate);
1484
1485        let inequality_desc = Self::get_inequality_desc_from_predicate(
1486            core.on
1487                .as_eq_predicate_ref()
1488                .expect("core predicate must exist")
1489                .other_cond()
1490                .clone(),
1491            left_len,
1492        )?;
1493
1494        Ok(StreamAsOfJoin::new(core, inequality_desc)?.into())
1495    }
1496
1497    /// Convert the logical join to a Hash join.
1498    fn to_batch_hash_join(
1499        &self,
1500        logical_join: generic::Join<BatchPlanRef>,
1501        predicate: EqJoinPredicate,
1502    ) -> Result<BatchPlanRef> {
1503        use super::batch::prelude::*;
1504
1505        let left_schema_len = logical_join.left.schema().len();
1506        let asof_desc = self
1507            .is_asof_join()
1508            .then(|| {
1509                Self::get_inequality_desc_from_predicate(
1510                    predicate.other_cond().clone(),
1511                    left_schema_len,
1512                )
1513            })
1514            .transpose()?;
1515
1516        let logical_join = generic::Join {
1517            on: generic::JoinOn::EqPredicate(predicate),
1518            ..logical_join
1519        };
1520        let batch_join = BatchHashJoin::new(logical_join, asof_desc);
1521        Ok(batch_join.into())
1522    }
1523
1524    pub fn get_inequality_desc_from_predicate(
1525        predicate: Condition,
1526        left_input_len: usize,
1527    ) -> Result<AsOfJoinDesc> {
1528        let expr: ExprImpl = predicate.into();
1529        if let Some((left_input_ref, expr_type, right_input_ref)) = expr.as_comparison_cond() {
1530            if left_input_ref.index() < left_input_len && right_input_ref.index() >= left_input_len
1531            {
1532                Ok(AsOfJoinDesc {
1533                    left_idx: left_input_ref.index() as u32,
1534                    right_idx: (right_input_ref.index() - left_input_len) as u32,
1535                    inequality_type: Self::expr_type_to_comparison_type(expr_type)?.into(),
1536                })
1537            } else {
1538                bail!("inequal condition from the same side should be push down in optimizer");
1539            }
1540        } else {
1541            Err(ErrorCode::InvalidInputSyntax(
1542                "AsOf join requires exactly 1 ineuquality condition".to_owned(),
1543            )
1544            .into())
1545        }
1546    }
1547
1548    fn expr_type_to_comparison_type(expr_type: PbType) -> Result<PbAsOfJoinInequalityType> {
1549        match expr_type {
1550            PbType::LessThan => Ok(PbAsOfJoinInequalityType::AsOfInequalityTypeLt),
1551            PbType::LessThanOrEqual => Ok(PbAsOfJoinInequalityType::AsOfInequalityTypeLe),
1552            PbType::GreaterThan => Ok(PbAsOfJoinInequalityType::AsOfInequalityTypeGt),
1553            PbType::GreaterThanOrEqual => Ok(PbAsOfJoinInequalityType::AsOfInequalityTypeGe),
1554            _ => Err(ErrorCode::InvalidInputSyntax(format!(
1555                "Invalid comparison type: {}",
1556                expr_type.as_str_name()
1557            ))
1558            .into()),
1559        }
1560    }
1561}
1562
1563impl ToBatch for LogicalJoin {
1564    fn to_batch(&self) -> Result<crate::optimizer::plan_node::BatchPlanRef> {
1565        let predicate = EqJoinPredicate::create(
1566            self.left().schema().len(),
1567            self.right().schema().len(),
1568            self.on().clone(),
1569        );
1570
1571        let batch_join = self
1572            .core
1573            .clone_with_inputs(self.core.left.to_batch()?, self.core.right.to_batch()?);
1574
1575        let ctx = self.base.ctx();
1576        let config = ctx.session_ctx().config();
1577
1578        if predicate.has_eq() {
1579            if !predicate.eq_keys_are_type_aligned() {
1580                return Err(ErrorCode::InternalError(format!(
1581                    "Join eq keys are not aligned for predicate: {predicate:?}"
1582                ))
1583                .into());
1584            }
1585            if config.batch_enable_lookup_join()
1586                && let Some(lookup_join) = self.to_batch_lookup_join_with_index_selection(
1587                    predicate.clone(),
1588                    batch_join.clone(),
1589                )?
1590            {
1591                return Ok(lookup_join.into());
1592            }
1593            self.to_batch_hash_join(batch_join, predicate)
1594        } else if self.is_asof_join() {
1595            Err(ErrorCode::InvalidInputSyntax(
1596                "AsOf join requires at least 1 equal condition".to_owned(),
1597            )
1598            .into())
1599        } else {
1600            // Convert to Nested-loop Join for non-equal joins
1601            Ok(BatchNestedLoopJoin::new(batch_join).into())
1602        }
1603    }
1604}
1605
1606impl ToStream for LogicalJoin {
1607    fn to_stream(
1608        &self,
1609        ctx: &mut ToStreamContext,
1610    ) -> Result<crate::optimizer::plan_node::StreamPlanRef> {
1611        if self
1612            .on()
1613            .conjunctions
1614            .iter()
1615            .any(|cond| cond.count_nows() > 0)
1616        {
1617            return Err(ErrorCode::NotSupported(
1618                "optimizer has tried to separate the temporal predicate(with now() expression) from the on condition, but it still reminded in on join's condition. Considering move it into WHERE clause?".to_owned(),
1619                 "please refer to https://docs.risingwave.com/processing/sql/temporal-filters for more information".to_owned()).into());
1620        }
1621
1622        let predicate = EqJoinPredicate::create(
1623            self.left().schema().len(),
1624            self.right().schema().len(),
1625            self.on().clone(),
1626        );
1627
1628        if self.join_type() == JoinType::AsofInner || self.join_type() == JoinType::AsofLeftOuter {
1629            self.to_stream_asof_join(predicate, ctx)
1630        } else if predicate.has_eq() {
1631            if !predicate.eq_keys_are_type_aligned() {
1632                return Err(ErrorCode::InternalError(format!(
1633                    "Join eq keys are not aligned for predicate: {predicate:?}"
1634                ))
1635                .into());
1636            }
1637
1638            if let Some(scan) = self.should_be_stream_temporal_join(ctx)? {
1639                self.to_stream_temporal_join_with_index_selection(scan, predicate, ctx)
1640            } else {
1641                self.to_stream_hash_join(predicate, ctx)
1642            }
1643        } else if let Some(scan) = self.should_be_stream_temporal_join(ctx)? {
1644            self.to_stream_nested_loop_temporal_join(scan, predicate, ctx)
1645        } else if let Some(dynamic_filter) =
1646            self.to_stream_dynamic_filter(self.on().clone(), ctx)?
1647        {
1648            Ok(dynamic_filter)
1649        } else {
1650            Err(RwError::from(ErrorCode::NotSupported(
1651                "streaming nested-loop join".to_owned(),
1652                "The non-equal join in the query requires a nested-loop join executor, which could be very expensive to run. \
1653                 Consider rewriting the query to use dynamic filter as a substitute if possible.\n\
1654                 See also: https://docs.risingwave.com/processing/sql/dynamic-filters".to_owned(),
1655            )))
1656        }
1657    }
1658
1659    fn logical_rewrite_for_stream(
1660        &self,
1661        ctx: &mut RewriteStreamContext,
1662    ) -> Result<(PlanRef, ColIndexMapping)> {
1663        let eq_indexes = self.eq_indexes();
1664        let (logical_left, logical_right) = if eq_indexes.is_empty() {
1665            (self.left(), self.right())
1666        } else {
1667            let lhs_join_key_idx = eq_indexes.iter().map(|(l, _)| *l).collect_vec();
1668            if self.should_be_temporal_join() {
1669                (
1670                    try_enforce_locality_requirement(
1671                        self.left(),
1672                        &lhs_join_key_idx,
1673                        ctx.locality_backfill_enabled(),
1674                    ),
1675                    self.right(),
1676                )
1677            } else {
1678                let rhs_join_key_idx = eq_indexes.iter().map(|(_, r)| *r).collect_vec();
1679                (
1680                    try_enforce_locality_requirement(
1681                        self.left(),
1682                        &lhs_join_key_idx,
1683                        ctx.locality_backfill_enabled(),
1684                    ),
1685                    try_enforce_locality_requirement(
1686                        self.right(),
1687                        &rhs_join_key_idx,
1688                        ctx.locality_backfill_enabled(),
1689                    ),
1690                )
1691            }
1692        };
1693
1694        let (left, left_col_change) = logical_left.logical_rewrite_for_stream(ctx)?;
1695        let left_len = left.schema().len();
1696        let (right, right_col_change) = logical_right.logical_rewrite_for_stream(ctx)?;
1697        let (join, out_col_change) = self.rewrite_with_left_right(
1698            left.clone(),
1699            left_col_change,
1700            right.clone(),
1701            right_col_change,
1702        );
1703
1704        let mapping = ColIndexMapping::with_remaining_columns(
1705            join.output_indices(),
1706            join.internal_column_num(),
1707        );
1708
1709        let l2o = join.core.l2i_col_mapping().composite(&mapping);
1710        let r2o = join.core.r2i_col_mapping().composite(&mapping);
1711
1712        // Add missing pk indices to the logical join
1713        let mut left_to_add = left
1714            .expect_stream_key()
1715            .iter()
1716            .cloned()
1717            .filter(|i| l2o.try_map(*i).is_none())
1718            .collect_vec();
1719
1720        let mut right_to_add = right
1721            .expect_stream_key()
1722            .iter()
1723            .filter(|&&i| r2o.try_map(i).is_none())
1724            .map(|&i| i + left_len)
1725            .collect_vec();
1726
1727        // NOTE(st1page): add join keys in the pk_indices a work around before we really have stream
1728        // key.
1729        let right_len = right.schema().len();
1730        let eq_predicate = EqJoinPredicate::create(left_len, right_len, join.on().clone());
1731
1732        let either_or_both = self.core.add_which_join_key_to_pk();
1733
1734        for (lk, rk) in eq_predicate.eq_indexes() {
1735            match either_or_both {
1736                EitherOrBoth::Left(_) => {
1737                    if l2o.try_map(lk).is_none() {
1738                        left_to_add.push(lk);
1739                    }
1740                }
1741                EitherOrBoth::Right(_) => {
1742                    if r2o.try_map(rk).is_none() {
1743                        right_to_add.push(rk + left_len)
1744                    }
1745                }
1746                EitherOrBoth::Both(_, _) => {
1747                    if l2o.try_map(lk).is_none() {
1748                        left_to_add.push(lk);
1749                    }
1750                    if r2o.try_map(rk).is_none() {
1751                        right_to_add.push(rk + left_len)
1752                    }
1753                }
1754            };
1755        }
1756        let left_to_add = left_to_add.into_iter().unique();
1757        let right_to_add = right_to_add.into_iter().unique();
1758        // NOTE(st1page) over
1759
1760        let mut new_output_indices = join.output_indices().clone();
1761        if !join.is_right_join() {
1762            new_output_indices.extend(left_to_add);
1763        }
1764        if !join.is_left_join() {
1765            new_output_indices.extend(right_to_add);
1766        }
1767
1768        let join_with_pk = join.clone_with_output_indices(new_output_indices);
1769
1770        let plan = if join_with_pk.join_type() == JoinType::FullOuter {
1771            // ignore the all NULL to maintain the stream key's uniqueness, see https://github.com/risingwavelabs/risingwave/issues/8084 for more information
1772
1773            let l2o = join_with_pk
1774                .core
1775                .l2i_col_mapping()
1776                .composite(&join_with_pk.core.i2o_col_mapping());
1777            let r2o = join_with_pk
1778                .core
1779                .r2i_col_mapping()
1780                .composite(&join_with_pk.core.i2o_col_mapping());
1781            let mut left_right_keys = join_with_pk
1782                .left()
1783                .expect_stream_key()
1784                .iter()
1785                .map(|i| l2o.map(*i))
1786                .collect_vec();
1787            left_right_keys.extend(
1788                join_with_pk
1789                    .right()
1790                    .expect_stream_key()
1791                    .iter()
1792                    .map(|i| r2o.map(*i)),
1793            );
1794            left_right_keys.extend(
1795                eq_predicate
1796                    .eq_indexes()
1797                    .iter()
1798                    .flat_map(|(lk, rk)| [l2o.map(*lk), r2o.map(*rk)]),
1799            );
1800            let left_right_keys = left_right_keys.into_iter().unique().collect_vec();
1801            let plan: PlanRef = join_with_pk.into();
1802            LogicalFilter::filter_out_all_null_keys(plan, &left_right_keys)
1803        } else {
1804            join_with_pk.into()
1805        };
1806
1807        // the added columns is at the end, so it will not change the exists column index
1808        Ok((plan, out_col_change))
1809    }
1810
1811    fn try_better_locality(&self, columns: &[usize]) -> Option<PlanRef> {
1812        // Only propagate locality for temporal-filter.
1813        if !self.temporal_filter_candidate() {
1814            return None;
1815        }
1816
1817        // Temporal filter only outputs columns from left input, so mapping is safe.
1818        let o2i_mapping = self.core.o2i_col_mapping();
1819        let left_input_columns = columns
1820            .iter()
1821            .map(|&col| o2i_mapping.try_map(col))
1822            .collect::<Option<Vec<usize>>>()?;
1823        if let Some(better_left_plan) = self.left().try_better_locality(&left_input_columns) {
1824            return Some(
1825                self.clone_with_left_right(better_left_plan, self.right())
1826                    .into(),
1827            );
1828        }
1829        None
1830    }
1831}
1832
1833#[cfg(test)]
1834mod tests {
1835
1836    use std::collections::HashSet;
1837
1838    use risingwave_common::catalog::{Field, Schema};
1839    use risingwave_common::types::{DataType, Datum};
1840    use risingwave_pb::expr::expr_node::Type;
1841
1842    use super::*;
1843    use crate::expr::{FunctionCall, Literal, assert_eq_input_ref};
1844    use crate::optimizer::optimizer_context::OptimizerContext;
1845    use crate::optimizer::plan_node::LogicalValues;
1846    use crate::optimizer::property::FunctionalDependency;
1847
1848    /// Pruning
1849    /// ```text
1850    /// Join(on: input_ref(1)=input_ref(3))
1851    ///   TableScan(v1, v2, v3)
1852    ///   TableScan(v4, v5, v6)
1853    /// ```
1854    /// with required columns [2,3] will result in
1855    /// ```text
1856    /// Project(input_ref(1), input_ref(2))
1857    ///   Join(on: input_ref(0)=input_ref(2))
1858    ///     TableScan(v2, v3)
1859    ///     TableScan(v4)
1860    /// ```
1861    #[tokio::test]
1862    async fn test_prune_join() {
1863        let ty = DataType::Int32;
1864        let ctx = OptimizerContext::mock();
1865        let fields: Vec<Field> = (1..7)
1866            .map(|i| Field::with_name(ty.clone(), format!("v{}", i)))
1867            .collect();
1868        let left = LogicalValues::new(
1869            vec![],
1870            Schema {
1871                fields: fields[0..3].to_vec(),
1872            },
1873            ctx.clone(),
1874        );
1875        let right = LogicalValues::new(
1876            vec![],
1877            Schema {
1878                fields: fields[3..6].to_vec(),
1879            },
1880            ctx,
1881        );
1882        let on: ExprImpl = ExprImpl::FunctionCall(Box::new(
1883            FunctionCall::new(
1884                Type::Equal,
1885                vec![
1886                    ExprImpl::InputRef(Box::new(InputRef::new(1, ty.clone()))),
1887                    ExprImpl::InputRef(Box::new(InputRef::new(3, ty))),
1888                ],
1889            )
1890            .unwrap(),
1891        ));
1892        let join_type = JoinType::Inner;
1893        let join: PlanRef = LogicalJoin::new(
1894            left.into(),
1895            right.into(),
1896            join_type,
1897            Condition::with_expr(on),
1898        )
1899        .into();
1900
1901        // Perform the prune
1902        let required_cols = vec![2, 3];
1903        let plan = join.prune_col(&required_cols, &mut ColumnPruningContext::new(join.clone()));
1904
1905        // Check the result
1906        let join = plan.as_logical_join().unwrap();
1907        assert_eq!(join.schema().fields().len(), 2);
1908        assert_eq!(join.schema().fields()[0], fields[2]);
1909        assert_eq!(join.schema().fields()[1], fields[3]);
1910
1911        let expr: ExprImpl = join.on().clone().into();
1912        let call = expr.as_function_call().unwrap();
1913        assert_eq_input_ref!(&call.inputs()[0], 0);
1914        assert_eq_input_ref!(&call.inputs()[1], 2);
1915
1916        let left = join.left();
1917        let left = left.as_logical_values().unwrap();
1918        assert_eq!(left.schema().fields(), &fields[1..3]);
1919        let right = join.right();
1920        let right = right.as_logical_values().unwrap();
1921        assert_eq!(right.schema().fields(), &fields[3..4]);
1922    }
1923
1924    /// Semi join panicked previously at `prune_col`. Add test to prevent regression.
1925    #[tokio::test]
1926    async fn test_prune_semi_join() {
1927        let ty = DataType::Int32;
1928        let ctx = OptimizerContext::mock();
1929        let fields: Vec<Field> = (1..7)
1930            .map(|i| Field::with_name(ty.clone(), format!("v{}", i)))
1931            .collect();
1932        let left = LogicalValues::new(
1933            vec![],
1934            Schema {
1935                fields: fields[0..3].to_vec(),
1936            },
1937            ctx.clone(),
1938        );
1939        let right = LogicalValues::new(
1940            vec![],
1941            Schema {
1942                fields: fields[3..6].to_vec(),
1943            },
1944            ctx,
1945        );
1946        let on: ExprImpl = ExprImpl::FunctionCall(Box::new(
1947            FunctionCall::new(
1948                Type::Equal,
1949                vec![
1950                    ExprImpl::InputRef(Box::new(InputRef::new(1, ty.clone()))),
1951                    ExprImpl::InputRef(Box::new(InputRef::new(4, ty))),
1952                ],
1953            )
1954            .unwrap(),
1955        ));
1956        for join_type in [
1957            JoinType::LeftSemi,
1958            JoinType::RightSemi,
1959            JoinType::LeftAnti,
1960            JoinType::RightAnti,
1961        ] {
1962            let join = LogicalJoin::new(
1963                left.clone().into(),
1964                right.clone().into(),
1965                join_type,
1966                Condition::with_expr(on.clone()),
1967            );
1968
1969            let offset = if join.is_right_join() { 3 } else { 0 };
1970            let join: PlanRef = join.into();
1971            // Perform the prune
1972            let required_cols = vec![0];
1973            // key 0 is never used in the join (always key 1)
1974            let plan = join.prune_col(&required_cols, &mut ColumnPruningContext::new(join.clone()));
1975            let as_plan = plan.as_logical_join().unwrap();
1976            // Check the result
1977            assert_eq!(as_plan.schema().fields().len(), 1);
1978            assert_eq!(as_plan.schema().fields()[0], fields[offset]);
1979
1980            // Perform the prune
1981            let required_cols = vec![0, 1, 2];
1982            // should not panic here
1983            let plan = join.prune_col(&required_cols, &mut ColumnPruningContext::new(join.clone()));
1984            let as_plan = plan.as_logical_join().unwrap();
1985            // Check the result
1986            assert_eq!(as_plan.schema().fields().len(), 3);
1987            assert_eq!(as_plan.schema().fields()[0], fields[offset]);
1988            assert_eq!(as_plan.schema().fields()[1], fields[offset + 1]);
1989            assert_eq!(as_plan.schema().fields()[2], fields[offset + 2]);
1990        }
1991    }
1992
1993    /// Pruning
1994    /// ```text
1995    /// Join(on: input_ref(1)=input_ref(3))
1996    ///   TableScan(v1, v2, v3)
1997    ///   TableScan(v4, v5, v6)
1998    /// ```
1999    /// with required columns [1, 3] will result in
2000    /// ```text
2001    /// Join(on: input_ref(0)=input_ref(1))
2002    ///   TableScan(v2)
2003    ///   TableScan(v4)
2004    /// ```
2005    #[tokio::test]
2006    async fn test_prune_join_no_project() {
2007        let ty = DataType::Int32;
2008        let ctx = OptimizerContext::mock();
2009        let fields: Vec<Field> = (1..7)
2010            .map(|i| Field::with_name(ty.clone(), format!("v{}", i)))
2011            .collect();
2012        let left = LogicalValues::new(
2013            vec![],
2014            Schema {
2015                fields: fields[0..3].to_vec(),
2016            },
2017            ctx.clone(),
2018        );
2019        let right = LogicalValues::new(
2020            vec![],
2021            Schema {
2022                fields: fields[3..6].to_vec(),
2023            },
2024            ctx,
2025        );
2026        let on: ExprImpl = ExprImpl::FunctionCall(Box::new(
2027            FunctionCall::new(
2028                Type::Equal,
2029                vec![
2030                    ExprImpl::InputRef(Box::new(InputRef::new(1, ty.clone()))),
2031                    ExprImpl::InputRef(Box::new(InputRef::new(3, ty))),
2032                ],
2033            )
2034            .unwrap(),
2035        ));
2036        let join_type = JoinType::Inner;
2037        let join: PlanRef = LogicalJoin::new(
2038            left.into(),
2039            right.into(),
2040            join_type,
2041            Condition::with_expr(on),
2042        )
2043        .into();
2044
2045        // Perform the prune
2046        let required_cols = vec![1, 3];
2047        let plan = join.prune_col(&required_cols, &mut ColumnPruningContext::new(join.clone()));
2048
2049        // Check the result
2050        let join = plan.as_logical_join().unwrap();
2051        assert_eq!(join.schema().fields().len(), 2);
2052        assert_eq!(join.schema().fields()[0], fields[1]);
2053        assert_eq!(join.schema().fields()[1], fields[3]);
2054
2055        let expr: ExprImpl = join.on().clone().into();
2056        let call = expr.as_function_call().unwrap();
2057        assert_eq_input_ref!(&call.inputs()[0], 0);
2058        assert_eq_input_ref!(&call.inputs()[1], 1);
2059
2060        let left = join.left();
2061        let left = left.as_logical_values().unwrap();
2062        assert_eq!(left.schema().fields(), &fields[1..2]);
2063        let right = join.right();
2064        let right = right.as_logical_values().unwrap();
2065        assert_eq!(right.schema().fields(), &fields[3..4]);
2066    }
2067
2068    /// Convert
2069    /// ```text
2070    /// Join(on: ($1 = $3) AND ($2 == 42))
2071    ///   TableScan(v1, v2, v3)
2072    ///   TableScan(v4, v5, v6)
2073    /// ```
2074    /// to
2075    /// ```text
2076    /// Filter($2 == 42)
2077    ///   HashJoin(on: $1 = $3)
2078    ///     TableScan(v1, v2, v3)
2079    ///     TableScan(v4, v5, v6)
2080    /// ```
2081    #[tokio::test]
2082    async fn test_join_to_batch() {
2083        let ctx = OptimizerContext::mock();
2084        let fields: Vec<Field> = (1..7)
2085            .map(|i| Field::with_name(DataType::Int32, format!("v{}", i)))
2086            .collect();
2087        let left = LogicalValues::new(
2088            vec![],
2089            Schema {
2090                fields: fields[0..3].to_vec(),
2091            },
2092            ctx.clone(),
2093        );
2094        let right = LogicalValues::new(
2095            vec![],
2096            Schema {
2097                fields: fields[3..6].to_vec(),
2098            },
2099            ctx,
2100        );
2101
2102        fn input_ref(i: usize) -> ExprImpl {
2103            ExprImpl::InputRef(Box::new(InputRef::new(i, DataType::Int32)))
2104        }
2105        let eq_cond = ExprImpl::FunctionCall(Box::new(
2106            FunctionCall::new(Type::Equal, vec![input_ref(1), input_ref(3)]).unwrap(),
2107        ));
2108        let non_eq_cond = ExprImpl::FunctionCall(Box::new(
2109            FunctionCall::new(
2110                Type::Equal,
2111                vec![
2112                    input_ref(2),
2113                    ExprImpl::Literal(Box::new(Literal::new(
2114                        Datum::Some(42_i32.into()),
2115                        DataType::Int32,
2116                    ))),
2117                ],
2118            )
2119            .unwrap(),
2120        ));
2121        // Condition: ($1 = $3) AND ($2 == 42)
2122        let on_cond = ExprImpl::FunctionCall(Box::new(
2123            FunctionCall::new(Type::And, vec![eq_cond.clone(), non_eq_cond.clone()]).unwrap(),
2124        ));
2125
2126        let join_type = JoinType::Inner;
2127        let logical_join = LogicalJoin::new(
2128            left.into(),
2129            right.into(),
2130            join_type,
2131            Condition::with_expr(on_cond),
2132        );
2133
2134        // Perform `to_batch`
2135        let result = logical_join.to_batch().unwrap();
2136
2137        // Expected plan:  HashJoin($1 = $3 AND $2 == 42)
2138        let hash_join = result.as_batch_hash_join().unwrap();
2139        assert_eq!(
2140            ExprImpl::from(hash_join.eq_join_predicate().eq_cond()),
2141            eq_cond
2142        );
2143        assert_eq!(
2144            *hash_join
2145                .eq_join_predicate()
2146                .non_eq_cond()
2147                .conjunctions
2148                .first()
2149                .unwrap(),
2150            non_eq_cond
2151        );
2152    }
2153
2154    /// Pruning
2155    /// ```text
2156    /// Join(on: input_ref(1)=input_ref(3))
2157    ///   TableScan(v1, v2, v3)
2158    ///   TableScan(v4, v5, v6)
2159    /// ```
2160    /// with required columns [3, 2] will result in
2161    /// ```text
2162    /// Project(input_ref(2), input_ref(1))
2163    ///   Join(on: input_ref(0)=input_ref(2))
2164    ///     TableScan(v2, v3)
2165    ///     TableScan(v4)
2166    /// ```
2167    #[tokio::test]
2168    async fn test_join_column_prune_with_order_required() {
2169        let ty = DataType::Int32;
2170        let ctx = OptimizerContext::mock();
2171        let fields: Vec<Field> = (1..7)
2172            .map(|i| Field::with_name(ty.clone(), format!("v{}", i)))
2173            .collect();
2174        let left = LogicalValues::new(
2175            vec![],
2176            Schema {
2177                fields: fields[0..3].to_vec(),
2178            },
2179            ctx.clone(),
2180        );
2181        let right = LogicalValues::new(
2182            vec![],
2183            Schema {
2184                fields: fields[3..6].to_vec(),
2185            },
2186            ctx,
2187        );
2188        let on: ExprImpl = ExprImpl::FunctionCall(Box::new(
2189            FunctionCall::new(
2190                Type::Equal,
2191                vec![
2192                    ExprImpl::InputRef(Box::new(InputRef::new(1, ty.clone()))),
2193                    ExprImpl::InputRef(Box::new(InputRef::new(3, ty))),
2194                ],
2195            )
2196            .unwrap(),
2197        ));
2198        let join_type = JoinType::Inner;
2199        let join: PlanRef = LogicalJoin::new(
2200            left.into(),
2201            right.into(),
2202            join_type,
2203            Condition::with_expr(on),
2204        )
2205        .into();
2206
2207        // Perform the prune
2208        let required_cols = vec![3, 2];
2209        let plan = join.prune_col(&required_cols, &mut ColumnPruningContext::new(join.clone()));
2210
2211        // Check the result
2212        let join = plan.as_logical_join().unwrap();
2213        assert_eq!(join.schema().fields().len(), 2);
2214        assert_eq!(join.schema().fields()[0], fields[3]);
2215        assert_eq!(join.schema().fields()[1], fields[2]);
2216
2217        let expr: ExprImpl = join.on().clone().into();
2218        let call = expr.as_function_call().unwrap();
2219        assert_eq_input_ref!(&call.inputs()[0], 0);
2220        assert_eq_input_ref!(&call.inputs()[1], 2);
2221
2222        let left = join.left();
2223        let left = left.as_logical_values().unwrap();
2224        assert_eq!(left.schema().fields(), &fields[1..3]);
2225        let right = join.right();
2226        let right = right.as_logical_values().unwrap();
2227        assert_eq!(right.schema().fields(), &fields[3..4]);
2228    }
2229
2230    #[tokio::test]
2231    async fn fd_derivation_inner_outer_join() {
2232        // left: [l0, l1], right: [r0, r1, r2]
2233        // FD: l0 --> l1, r0 --> { r1, r2 }
2234        // On: l0 = 0 AND l1 = r1
2235        //
2236        // Inner Join:
2237        //  Schema: [l0, l1, r0, r1, r2]
2238        //  FD: l0 --> l1, r0 --> { r1, r2 }, {} --> l0, l1 --> r1, r1 --> l1
2239        // Left Outer Join:
2240        //  Schema: [l0, l1, r0, r1, r2]
2241        //  FD: l0 --> l1
2242        // Right Outer Join:
2243        //  Schema: [l0, l1, r0, r1, r2]
2244        //  FD: r0 --> { r1, r2 }
2245        // Full Outer Join:
2246        //  Schema: [l0, l1, r0, r1, r2]
2247        //  FD: empty
2248        // Left Semi/Anti Join:
2249        //  Schema: [l0, l1]
2250        //  FD: l0 --> l1
2251        // Right Semi/Anti Join:
2252        //  Schema: [r0, r1, r2]
2253        //  FD: r0 --> {r1, r2}
2254        let ctx = OptimizerContext::mock();
2255        let left = {
2256            let fields: Vec<Field> = vec![
2257                Field::with_name(DataType::Int32, "l0"),
2258                Field::with_name(DataType::Int32, "l1"),
2259            ];
2260            let mut values = LogicalValues::new(vec![], Schema { fields }, ctx.clone());
2261            // 0 --> 1
2262            values
2263                .base
2264                .functional_dependency_mut()
2265                .add_functional_dependency_by_column_indices(&[0], &[1]);
2266            values
2267        };
2268        let right = {
2269            let fields: Vec<Field> = vec![
2270                Field::with_name(DataType::Int32, "r0"),
2271                Field::with_name(DataType::Int32, "r1"),
2272                Field::with_name(DataType::Int32, "r2"),
2273            ];
2274            let mut values = LogicalValues::new(vec![], Schema { fields }, ctx);
2275            // 0 --> 1, 2
2276            values
2277                .base
2278                .functional_dependency_mut()
2279                .add_functional_dependency_by_column_indices(&[0], &[1, 2]);
2280            values
2281        };
2282        // l0 = 0 AND l1 = r1
2283        let on: ExprImpl = FunctionCall::new(
2284            Type::And,
2285            vec![
2286                FunctionCall::new(
2287                    Type::Equal,
2288                    vec![
2289                        InputRef::new(0, DataType::Int32).into(),
2290                        ExprImpl::literal_int(0),
2291                    ],
2292                )
2293                .unwrap()
2294                .into(),
2295                FunctionCall::new(
2296                    Type::Equal,
2297                    vec![
2298                        InputRef::new(1, DataType::Int32).into(),
2299                        InputRef::new(3, DataType::Int32).into(),
2300                    ],
2301                )
2302                .unwrap()
2303                .into(),
2304            ],
2305        )
2306        .unwrap()
2307        .into();
2308        let expected_fd_set = [
2309            (
2310                JoinType::Inner,
2311                [
2312                    // inherit from left
2313                    FunctionalDependency::with_indices(5, &[0], &[1]),
2314                    // inherit from right
2315                    FunctionalDependency::with_indices(5, &[2], &[3, 4]),
2316                    // constant column in join condition
2317                    FunctionalDependency::with_indices(5, &[], &[0]),
2318                    // eq column in join condition
2319                    FunctionalDependency::with_indices(5, &[1], &[3]),
2320                    FunctionalDependency::with_indices(5, &[3], &[1]),
2321                ]
2322                .into_iter()
2323                .collect::<HashSet<_>>(),
2324            ),
2325            (JoinType::FullOuter, HashSet::new()),
2326            (
2327                JoinType::RightOuter,
2328                [
2329                    // inherit from right
2330                    FunctionalDependency::with_indices(5, &[2], &[3, 4]),
2331                ]
2332                .into_iter()
2333                .collect::<HashSet<_>>(),
2334            ),
2335            (
2336                JoinType::LeftOuter,
2337                [
2338                    // inherit from left
2339                    FunctionalDependency::with_indices(5, &[0], &[1]),
2340                ]
2341                .into_iter()
2342                .collect::<HashSet<_>>(),
2343            ),
2344            (
2345                JoinType::LeftSemi,
2346                [
2347                    // inherit from left
2348                    FunctionalDependency::with_indices(2, &[0], &[1]),
2349                ]
2350                .into_iter()
2351                .collect::<HashSet<_>>(),
2352            ),
2353            (
2354                JoinType::LeftAnti,
2355                [
2356                    // inherit from left
2357                    FunctionalDependency::with_indices(2, &[0], &[1]),
2358                ]
2359                .into_iter()
2360                .collect::<HashSet<_>>(),
2361            ),
2362            (
2363                JoinType::RightSemi,
2364                [
2365                    // inherit from right
2366                    FunctionalDependency::with_indices(3, &[0], &[1, 2]),
2367                ]
2368                .into_iter()
2369                .collect::<HashSet<_>>(),
2370            ),
2371            (
2372                JoinType::RightAnti,
2373                [
2374                    // inherit from right
2375                    FunctionalDependency::with_indices(3, &[0], &[1, 2]),
2376                ]
2377                .into_iter()
2378                .collect::<HashSet<_>>(),
2379            ),
2380        ];
2381
2382        for (join_type, expected_res) in expected_fd_set {
2383            let join = LogicalJoin::new(
2384                left.clone().into(),
2385                right.clone().into(),
2386                join_type,
2387                Condition::with_expr(on.clone()),
2388            );
2389            let fd_set = join
2390                .functional_dependency()
2391                .as_dependencies()
2392                .iter()
2393                .cloned()
2394                .collect::<HashSet<_>>();
2395            assert_eq!(fd_set, expected_res);
2396        }
2397    }
2398}