risingwave_frontend/optimizer/plan_node/
batch_lookup_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::sync::Arc;
16
17use pretty_xmlish::{Pretty, XmlNode};
18use risingwave_common::catalog::ColumnId;
19use risingwave_pb::batch_plan::plan_node::NodeBody;
20use risingwave_pb::batch_plan::{DistributedLookupJoinNode, LocalLookupJoinNode};
21use risingwave_pb::plan_common::AsOfJoinDesc;
22use risingwave_sqlparser::ast::AsOf;
23
24use super::batch::prelude::*;
25use super::utils::{Distill, childless_record, to_batch_query_epoch};
26use super::{BatchPlanRef as PlanRef, BatchSeqScan, ExprRewritable, generic};
27use crate::TableCatalog;
28use crate::error::Result;
29use crate::expr::{Expr, ExprRewriter, ExprVisitor};
30use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
31use crate::optimizer::plan_node::utils::IndicesDisplay;
32use crate::optimizer::plan_node::{
33    EqJoinPredicate, EqJoinPredicateDisplay, PlanBase, PlanTreeNodeUnary, ToDistributedBatch,
34    ToLocalBatch, TryToBatchPb,
35};
36use crate::optimizer::property::{Distribution, Order, RequiredDist};
37use crate::scheduler::SchedulerResult;
38use crate::utils::ColIndexMappingRewriteExt;
39
40#[derive(Debug, Clone, PartialEq, Eq, Hash)]
41pub struct BatchLookupJoin {
42    pub base: PlanBase<Batch>,
43    core: generic::Join<PlanRef>,
44
45    /// Table description of the right side table
46    right_table: Arc<TableCatalog>,
47
48    /// Output column ids of the right side table
49    right_output_column_ids: Vec<ColumnId>,
50
51    /// The prefix length of the order key of right side table.
52    lookup_prefix_len: usize,
53
54    /// If `distributed_lookup` is true, it will generate `DistributedLookupJoinNode` for
55    /// `ToBatchPb`. Otherwise, it will generate `LookupJoinNode`.
56    distributed_lookup: bool,
57
58    as_of: Option<AsOf>,
59    // `AsOf` join description
60    asof_desc: Option<AsOfJoinDesc>,
61}
62
63impl BatchLookupJoin {
64    pub fn new(
65        core: generic::Join<PlanRef>,
66        right_table: Arc<TableCatalog>,
67        right_output_column_ids: Vec<ColumnId>,
68        lookup_prefix_len: usize,
69        distributed_lookup: bool,
70        as_of: Option<AsOf>,
71        asof_desc: Option<AsOfJoinDesc>,
72    ) -> Self {
73        // We cannot create a `BatchLookupJoin` without any eq keys. We require eq keys to do the
74        // lookup.
75        let eq_join_predicate = core
76            .on
77            .as_eq_predicate_ref()
78            .expect("BatchLookupJoin requires JoinOn::EqPredicate in core");
79        assert!(eq_join_predicate.has_eq());
80        assert!(eq_join_predicate.eq_keys_are_type_aligned());
81        let dist = Self::derive_dist(core.left.distribution(), &core);
82        let base = PlanBase::new_batch_with_core(&core, dist, Order::any());
83        Self {
84            base,
85            core,
86            right_table,
87            right_output_column_ids,
88            lookup_prefix_len,
89            distributed_lookup,
90            as_of,
91            asof_desc,
92        }
93    }
94
95    fn derive_dist(left: &Distribution, core: &generic::Join<PlanRef>) -> Distribution {
96        match left {
97            Distribution::Single => Distribution::Single,
98            Distribution::HashShard(_) | Distribution::UpstreamHashShard(_, _) => {
99                let l2o = core.l2i_col_mapping().composite(&core.i2o_col_mapping());
100                l2o.rewrite_provided_distribution(left)
101            }
102            _ => unreachable!(),
103        }
104    }
105
106    fn eq_join_predicate(&self) -> &EqJoinPredicate {
107        self.core
108            .on
109            .as_eq_predicate_ref()
110            .expect("BatchLookupJoin should store predicate as EqJoinPredicate")
111    }
112
113    pub fn right_table(&self) -> &TableCatalog {
114        &self.right_table
115    }
116
117    fn clone_with_distributed_lookup(&self, input: PlanRef, distributed_lookup: bool) -> Self {
118        let mut batch_lookup_join = self.clone_with_input(input);
119        batch_lookup_join.distributed_lookup = distributed_lookup;
120        batch_lookup_join
121    }
122
123    pub fn lookup_prefix_len(&self) -> usize {
124        self.lookup_prefix_len
125    }
126}
127
128impl Distill for BatchLookupJoin {
129    fn distill<'a>(&self) -> XmlNode<'a> {
130        let verbose = self.base.ctx().is_explain_verbose();
131        let mut vec = Vec::with_capacity(if verbose { 3 } else { 2 });
132        vec.push(("type", Pretty::debug(&self.core.join_type)));
133
134        let concat_schema = self.core.concat_schema();
135        vec.push((
136            "predicate",
137            Pretty::debug(&EqJoinPredicateDisplay {
138                eq_join_predicate: self.eq_join_predicate(),
139                input_schema: &concat_schema,
140            }),
141        ));
142
143        if verbose {
144            let data = IndicesDisplay::from_join(&self.core, &concat_schema);
145            vec.push(("output", data));
146        }
147
148        let scan: &BatchSeqScan = self.core.right.as_batch_seq_scan().unwrap();
149
150        vec.push(("lookup table", Pretty::display(&scan.core().table_name())));
151
152        if let Some(as_of) = &self.as_of {
153            vec.push(("as_of", Pretty::debug(as_of)));
154        }
155
156        childless_record("BatchLookupJoin", vec)
157    }
158}
159
160impl PlanTreeNodeUnary<Batch> for BatchLookupJoin {
161    fn input(&self) -> PlanRef {
162        self.core.left.clone()
163    }
164
165    // Only change left side
166    fn clone_with_input(&self, input: PlanRef) -> Self {
167        let mut core = self.core.clone();
168        core.left = input;
169        Self::new(
170            core,
171            self.right_table.clone(),
172            self.right_output_column_ids.clone(),
173            self.lookup_prefix_len,
174            self.distributed_lookup,
175            self.as_of.clone(),
176            self.asof_desc,
177        )
178    }
179}
180
181impl_plan_tree_node_for_unary! { Batch, BatchLookupJoin }
182
183impl ToDistributedBatch for BatchLookupJoin {
184    fn to_distributed(&self) -> Result<PlanRef> {
185        let right_table = &self.right_table;
186
187        // The lookup table has a singleton distribution, so there's no way to align the
188        // left side with the distribution key of the right table. Instead, gather the left
189        // side into a single task and perform all lookups from there. Note that this is
190        // still correct because the lookup executor scans the table with a full vnode
191        // bitmap, regardless of which worker the task is scheduled to.
192        if right_table.distribution_key.is_empty() {
193            let input = self
194                .input()
195                .to_distributed_with_required(&Order::any(), &RequiredDist::single())?;
196            return Ok(self.clone_with_distributed_lookup(input, true).into());
197        }
198
199        // Align left distribution keys with the right table.
200        let mut exchange_dist_keys = vec![];
201        let left_eq_indexes = self.eq_join_predicate().left_eq_indexes();
202        for dist_col_index in &right_table.distribution_key {
203            let dist_col_id = right_table.columns[*dist_col_index].column_desc.column_id;
204            let output_pos = self
205                .right_output_column_ids
206                .iter()
207                .position(|p| *p == dist_col_id)
208                .unwrap();
209            let dist_in_eq_indexes = self
210                .eq_join_predicate()
211                .right_eq_indexes()
212                .iter()
213                .position(|col| *col == output_pos)
214                .unwrap();
215            assert!(dist_in_eq_indexes < self.lookup_prefix_len);
216            exchange_dist_keys.push(left_eq_indexes[dist_in_eq_indexes]);
217        }
218
219        assert!(!exchange_dist_keys.is_empty());
220
221        let input = self.input().to_distributed_with_required(
222            &Order::any(),
223            &RequiredDist::PhysicalDist(Distribution::UpstreamHashShard(
224                exchange_dist_keys,
225                self.right_table.id,
226            )),
227        )?;
228
229        Ok(self.clone_with_distributed_lookup(input, true).into())
230    }
231}
232
233impl TryToBatchPb for BatchLookupJoin {
234    fn try_to_batch_prost_body(&self) -> SchedulerResult<NodeBody> {
235        let eq_join_predicate = self.eq_join_predicate();
236        Ok(if self.distributed_lookup {
237            NodeBody::DistributedLookupJoin(DistributedLookupJoinNode {
238                join_type: self.core.join_type as i32,
239                condition: self
240                    .eq_join_predicate()
241                    .other_cond()
242                    .as_expr_unless_true()
243                    .map(|x| x.to_expr_proto()),
244                outer_side_key: self
245                    .eq_join_predicate()
246                    .left_eq_indexes()
247                    .into_iter()
248                    .map(|a| a as _)
249                    .collect(),
250                inner_side_key: self
251                    .eq_join_predicate()
252                    .right_eq_indexes()
253                    .into_iter()
254                    .map(|a| a as _)
255                    .collect(),
256                inner_side_table_desc: Some(self.right_table.table_desc().try_to_protobuf()?),
257                inner_side_column_ids: self
258                    .right_output_column_ids
259                    .iter()
260                    .map(ColumnId::get_id)
261                    .collect(),
262                output_indices: self.core.output_indices.iter().map(|&x| x as u32).collect(),
263                null_safe: eq_join_predicate.null_safes(),
264                lookup_prefix_len: self.lookup_prefix_len as u32,
265                query_epoch: to_batch_query_epoch(&self.as_of)?,
266                asof_desc: self.asof_desc,
267            })
268        } else {
269            NodeBody::LocalLookupJoin(LocalLookupJoinNode {
270                join_type: self.core.join_type as i32,
271                condition: self
272                    .eq_join_predicate()
273                    .other_cond()
274                    .as_expr_unless_true()
275                    .map(|x| x.to_expr_proto()),
276                outer_side_key: self
277                    .eq_join_predicate()
278                    .left_eq_indexes()
279                    .into_iter()
280                    .map(|a| a as _)
281                    .collect(),
282                inner_side_key: self
283                    .eq_join_predicate()
284                    .right_eq_indexes()
285                    .into_iter()
286                    .map(|a| a as _)
287                    .collect(),
288                inner_side_table_desc: Some(self.right_table.table_desc().try_to_protobuf()?),
289                inner_side_vnode_mapping: vec![], // To be filled in at local.rs
290                inner_side_column_ids: self
291                    .right_output_column_ids
292                    .iter()
293                    .map(ColumnId::get_id)
294                    .collect(),
295                output_indices: self.core.output_indices.iter().map(|&x| x as u32).collect(),
296                worker_nodes: vec![], // To be filled in at local.rs
297                null_safe: eq_join_predicate.null_safes(),
298                lookup_prefix_len: self.lookup_prefix_len as u32,
299                query_epoch: to_batch_query_epoch(&self.as_of)?,
300                asof_desc: self.asof_desc,
301            })
302        })
303    }
304}
305
306impl ToLocalBatch for BatchLookupJoin {
307    fn to_local(&self) -> Result<PlanRef> {
308        let input = RequiredDist::single()
309            .batch_enforce_if_not_satisfies(self.input().to_local()?, &Order::any())?;
310
311        Ok(self.clone_with_distributed_lookup(input, false).into())
312    }
313}
314
315impl ExprRewritable<Batch> for BatchLookupJoin {
316    fn has_rewritable_expr(&self) -> bool {
317        true
318    }
319
320    fn rewrite_exprs(&self, r: &mut dyn ExprRewriter) -> PlanRef {
321        let base = self.base.clone_with_new_plan_id();
322        let mut core = self.core.clone();
323        core.rewrite_exprs(r);
324        let mut new = self.clone();
325        new.base = base;
326        new.core = core;
327        new.into()
328    }
329}
330
331impl ExprVisitable for BatchLookupJoin {
332    fn visit_exprs(&self, v: &mut dyn ExprVisitor) {
333        self.core.visit_exprs(v);
334    }
335}
risingwave_frontend/optimizer/plan_node/batch_lookup_join.rs

risingwave_frontend/optimizer/plan_node/
batch_lookup_join.rs
