risingwave_frontend/optimizer/plan_node/
stream_temporal_join.rs

1// Copyright 2023 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 itertools::Itertools;
16use pretty_xmlish::{Pretty, XmlNode};
17use risingwave_common::util::sort_util::OrderType;
18use risingwave_pb::plan_common::JoinType;
19use risingwave_pb::stream_plan::TemporalJoinNode;
20use risingwave_pb::stream_plan::stream_node::NodeBody;
21use risingwave_sqlparser::ast::AsOf;
22
23use super::stream::prelude::*;
24use super::utils::{Distill, childless_record, watermark_pretty};
25use super::{ExprRewritable, PlanBase, PlanTreeNodeBinary, StreamPlanRef as PlanRef, generic};
26use crate::TableCatalog;
27use crate::expr::{Expr, ExprRewriter, ExprVisitor};
28use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
29use crate::optimizer::plan_node::generic::GenericPlanNode;
30use crate::optimizer::plan_node::plan_tree_node::PlanTreeNodeUnary;
31use crate::optimizer::plan_node::utils::{IndicesDisplay, TableCatalogBuilder};
32use crate::optimizer::plan_node::{
33    EqJoinPredicate, EqJoinPredicateDisplay, StreamExchange, StreamTableScan, TryToStreamPb,
34};
35use crate::optimizer::property::Distribution;
36use crate::scheduler::SchedulerResult;
37use crate::stream_fragmenter::BuildFragmentGraphState;
38use crate::utils::ColIndexMappingRewriteExt;
39
40#[derive(Debug, Clone, PartialEq, Eq, Hash)]
41pub struct StreamTemporalJoin {
42    pub base: PlanBase<Stream>,
43    core: generic::Join<PlanRef>,
44    append_only: bool,
45    is_nested_loop: bool,
46    is_broadcast: bool,
47}
48
49impl StreamTemporalJoin {
50    pub fn new(core: generic::Join<PlanRef>, is_nested_loop: bool) -> Result<Self> {
51        core.on
52            .as_eq_predicate_ref()
53            .expect("StreamTemporalJoin requires JoinOn::EqPredicate in core");
54        assert!(core.join_type == JoinType::Inner || core.join_type == JoinType::LeftOuter);
55        // TODO(kind): theoretically, the impl can handle upsert stream.
56        let stream_kind = reject_upsert_input!(core.left);
57        let append_only = stream_kind.is_append_only();
58        assert!(!is_nested_loop || append_only);
59
60        let right = core.right.clone();
61        let exchange: &StreamExchange = right
62            .as_stream_exchange()
63            .expect("temporal join lookup side should have an exchange");
64        let exchange_input = exchange.input();
65        let scan: &StreamTableScan = exchange_input
66            .as_stream_table_scan()
67            .expect("should be a stream table scan");
68        let is_broadcast = matches!(scan.core().as_of, Some(AsOf::ProcessTimeBroadcast));
69        assert!(matches!(
70            scan.core().as_of,
71            Some(AsOf::ProcessTime | AsOf::ProcessTimeBroadcast)
72        ));
73        if is_broadcast {
74            assert!(!exchange.no_shuffle());
75            assert_eq!(exchange.distribution(), &Distribution::Broadcast);
76            assert!(!is_nested_loop);
77        } else {
78            assert!(exchange.no_shuffle());
79        }
80
81        let dist = if is_nested_loop {
82            // Use right side distribution directly if it's nested loop temporal join.
83            let r2o = core.r2i_col_mapping().composite(&core.i2o_col_mapping());
84            r2o.rewrite_provided_distribution(core.right.distribution())
85        } else {
86            // Use left side distribution directly if it's hash temporal join.
87            // https://github.com/risingwavelabs/risingwave/pull/19201#discussion_r1824031780
88            let l2o = core.l2i_col_mapping().composite(&core.i2o_col_mapping());
89            l2o.rewrite_provided_distribution(core.left.distribution())
90        };
91
92        // Use left side watermark directly.
93        let watermark_columns = core
94            .left
95            .watermark_columns()
96            .map_clone(&core.l2i_col_mapping())
97            .map_clone(&core.i2o_col_mapping());
98
99        let columns_monotonicity = core.i2o_col_mapping().rewrite_monotonicity_map(
100            &core
101                .l2i_col_mapping()
102                .rewrite_monotonicity_map(core.left.columns_monotonicity()),
103        );
104
105        let base = PlanBase::new_stream_with_core(
106            &core,
107            dist,
108            stream_kind,
109            false, // TODO(rc): derive EOWC property from input
110            watermark_columns,
111            columns_monotonicity,
112        );
113
114        Ok(Self {
115            base,
116            core,
117            append_only,
118            is_nested_loop,
119            is_broadcast,
120        })
121    }
122
123    /// Get join type
124    pub fn join_type(&self) -> JoinType {
125        self.core.join_type
126    }
127
128    pub fn eq_join_predicate(&self) -> &EqJoinPredicate {
129        self.core
130            .on
131            .as_eq_predicate_ref()
132            .expect("StreamTemporalJoin should store predicate as EqJoinPredicate")
133    }
134
135    pub fn append_only(&self) -> bool {
136        self.append_only
137    }
138
139    pub fn is_nested_loop(&self) -> bool {
140        self.is_nested_loop
141    }
142
143    pub fn is_broadcast(&self) -> bool {
144        self.is_broadcast
145    }
146
147    /// Return the memo-table catalog.
148    ///
149    /// The memo prefix is `join_key + left_stream_key`. A regular temporal join is distributed by
150    /// the lookup key, while a broadcast temporal join shuffles the left input by its stream key
151    /// (or keeps it singleton) and maps that distribution key to its copy in the
152    /// `left_stream_key` part of the prefix.
153    ///
154    /// Write pattern:
155    ///   for each left input row (with insert op), persist the matched right row followed by the
156    ///   memo prefix.
157    ///
158    /// Read pattern:
159    ///   for each left input row (with delete op), use the memo prefix to fetch and delete all
160    ///   right rows that matched when the left row was inserted.
161    pub fn infer_memo_table_catalog(&self, right_scan: &StreamTableScan) -> TableCatalog {
162        let left_eq_indexes = self.eq_join_predicate().left_eq_indexes();
163        let left_stream_key = self.core.left.expect_stream_key();
164        let memo_prefix_indices = left_eq_indexes
165            .iter()
166            .chain(left_stream_key)
167            .copied()
168            .collect_vec();
169        let read_prefix_len_hint = memo_prefix_indices.len();
170
171        // Build internal table
172        let mut internal_table_catalog_builder = TableCatalogBuilder::default();
173        // Add right table fields
174        let right_scan_schema = right_scan.core().schema();
175        for field in right_scan_schema.fields() {
176            internal_table_catalog_builder.add_column(field);
177        }
178        // Add the columns that identify and route a left row.
179        for field in memo_prefix_indices
180            .iter()
181            .map(|idx| &self.core.left.schema().fields()[*idx])
182        {
183            internal_table_catalog_builder.add_column(field);
184        }
185
186        let mut pk_indices = vec![];
187        pk_indices
188            .extend(right_scan_schema.len()..(right_scan_schema.len() + read_prefix_len_hint));
189        pk_indices.extend(right_scan.stream_key().unwrap());
190
191        pk_indices.iter().for_each(|idx| {
192            internal_table_catalog_builder.add_order_column(*idx, OrderType::ascending())
193        });
194
195        let internal_table_dist_keys = if self.is_broadcast {
196            let left_dist_key = self
197                .core
198                .left
199                .distribution()
200                .dist_column_indices_opt()
201                .expect("broadcast temporal join left input must have a concrete distribution");
202            left_dist_key
203                .iter()
204                .map(|dist_idx| {
205                    let position = left_stream_key
206                        .iter()
207                        .position(|idx| idx == dist_idx)
208                        .expect("left distribution key must be part of the left stream key");
209                    right_scan_schema.len() + left_eq_indexes.len() + position
210                })
211                .collect()
212        } else {
213            let dist_key_len = right_scan
214                .core()
215                .distribution_key()
216                .map(|keys| keys.len())
217                .unwrap_or(0);
218            (right_scan_schema.len()..(right_scan_schema.len() + dist_key_len)).collect()
219        };
220        internal_table_catalog_builder.build(internal_table_dist_keys, read_prefix_len_hint)
221    }
222}
223
224impl Distill for StreamTemporalJoin {
225    fn distill<'a>(&self) -> XmlNode<'a> {
226        let verbose = self.base.ctx().is_explain_verbose();
227        let mut vec = Vec::with_capacity(if verbose { 3 } else { 2 });
228        vec.push(("type", Pretty::debug(&self.core.join_type)));
229        vec.push(("append_only", Pretty::debug(&self.append_only)));
230
231        let concat_schema = self.core.concat_schema();
232        vec.push((
233            "predicate",
234            Pretty::debug(&EqJoinPredicateDisplay {
235                eq_join_predicate: self.eq_join_predicate(),
236                input_schema: &concat_schema,
237            }),
238        ));
239
240        vec.push(("nested_loop", Pretty::debug(&self.is_nested_loop)));
241        if self.is_broadcast {
242            vec.push(("broadcast", Pretty::debug(&self.is_broadcast)));
243        }
244
245        if let Some(ow) = watermark_pretty(self.base.watermark_columns(), self.schema()) {
246            vec.push(("output_watermarks", ow));
247        }
248
249        if verbose {
250            let data = IndicesDisplay::from_join(&self.core, &concat_schema);
251            vec.push(("output", data));
252        }
253
254        childless_record("StreamTemporalJoin", vec)
255    }
256}
257
258impl PlanTreeNodeBinary<Stream> for StreamTemporalJoin {
259    fn left(&self) -> PlanRef {
260        self.core.left.clone()
261    }
262
263    fn right(&self) -> PlanRef {
264        self.core.right.clone()
265    }
266
267    fn clone_with_left_right(&self, left: PlanRef, right: PlanRef) -> Self {
268        let mut core = self.core.clone();
269        core.left = left;
270        core.right = right;
271        Self::new(core, self.is_nested_loop).unwrap()
272    }
273}
274
275impl_plan_tree_node_for_binary! { Stream, StreamTemporalJoin }
276
277impl TryToStreamPb for StreamTemporalJoin {
278    fn try_to_stream_prost_body(
279        &self,
280        state: &mut BuildFragmentGraphState,
281    ) -> SchedulerResult<NodeBody> {
282        let left_jk_indices = self.eq_join_predicate().left_eq_indexes();
283        let right_jk_indices = self.eq_join_predicate().right_eq_indexes();
284        let left_jk_indices_prost = left_jk_indices.iter().map(|idx| *idx as i32).collect_vec();
285        let right_jk_indices_prost = right_jk_indices.iter().map(|idx| *idx as i32).collect_vec();
286
287        let null_safe_prost = self.eq_join_predicate().null_safes().into_iter().collect();
288
289        let right = self.right();
290        let exchange: &StreamExchange = right
291            .as_stream_exchange()
292            .expect("temporal join lookup side should have an exchange");
293        if self.is_broadcast {
294            assert!(!exchange.no_shuffle());
295            assert_eq!(exchange.distribution(), &Distribution::Broadcast);
296        } else {
297            assert!(exchange.no_shuffle());
298        }
299        let exchange_input = exchange.input();
300        let scan: &StreamTableScan = exchange_input
301            .as_stream_table_scan()
302            .expect("should be a stream table scan");
303
304        Ok(NodeBody::TemporalJoin(Box::new(TemporalJoinNode {
305            join_type: self.core.join_type as i32,
306            left_key: left_jk_indices_prost,
307            right_key: right_jk_indices_prost,
308            null_safe: null_safe_prost,
309            condition: self
310                .eq_join_predicate()
311                .other_cond()
312                .as_expr_unless_true()
313                .map(|expr| {
314                    expr.to_expr_proto_checked_pure(
315                        self.left().stream_kind().is_retract()
316                            || self.right().stream_kind().is_retract(),
317                        "JOIN condition",
318                    )
319                })
320                .transpose()?,
321            output_indices: self.core.output_indices.iter().map(|&x| x as u32).collect(),
322            table_desc: Some(scan.core().table_catalog.table_desc().try_to_protobuf()?),
323            table_output_indices: scan.core().output_col_idx.iter().map(|&i| i as _).collect(),
324            memo_table: if self.append_only {
325                None
326            } else {
327                let mut memo_table = self.infer_memo_table_catalog(scan);
328                memo_table = memo_table.with_id(state.gen_table_id_wrapped());
329                Some(memo_table.to_internal_table_prost())
330            },
331            is_nested_loop: self.is_nested_loop,
332            is_broadcast: self.is_broadcast,
333        })))
334    }
335}
336
337impl ExprRewritable<Stream> for StreamTemporalJoin {
338    fn has_rewritable_expr(&self) -> bool {
339        true
340    }
341
342    fn rewrite_exprs(&self, r: &mut dyn ExprRewriter) -> PlanRef {
343        let mut core = self.core.clone();
344        core.rewrite_exprs(r);
345        Self::new(core, self.is_nested_loop).unwrap().into()
346    }
347}
348
349impl ExprVisitable for StreamTemporalJoin {
350    fn visit_exprs(&self, v: &mut dyn ExprVisitor) {
351        self.core.visit_exprs(v);
352    }
353}
risingwave_frontend/optimizer/plan_node/stream_temporal_join.rs

risingwave_frontend/optimizer/plan_node/
stream_temporal_join.rs
