risingwave_frontend/optimizer/plan_node/
convert.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;
16
17use risingwave_common::catalog::FieldDisplay;
18use risingwave_pb::stream_plan::StreamScanType;
19
20use super::*;
21use crate::optimizer::ShareId;
22use crate::optimizer::property::RequiredDist;
23
24/// `ToStream` converts a logical plan node to streaming physical node
25/// with an optional required distribution.
26///
27/// when implement this trait you can choose the two ways
28/// - Implement `to_stream` and use the default implementation of `to_stream_with_dist_required`
29/// - Or, if the required distribution is given, there will be a better plan. For example a hash
30///   join with hash-key(a,b) and the plan is required hash-distributed by (a,b,c). you can
31///   implement `to_stream_with_dist_required`, and implement `to_stream` with
32///   `to_stream_with_dist_required(RequiredDist::Any)`. you can see [`LogicalProject`] as an
33///   example.
34pub trait ToStream {
35    /// `logical_rewrite_for_stream` will rewrite the logical node, and return (`new_plan_node`,
36    /// `col_mapping`), the `col_mapping` is for original columns have been changed into some other
37    /// position.
38    ///
39    /// Now it is used to:
40    /// 1. ensure every plan node's output having pk column
41    /// 2. add `row_count`() in every Agg
42    fn logical_rewrite_for_stream(
43        &self,
44        ctx: &mut RewriteStreamContext,
45    ) -> Result<(LogicalPlanRef, ColIndexMapping)>;
46
47    /// `to_stream` is equivalent to `to_stream_with_dist_required(RequiredDist::Any)`
48    fn to_stream(&self, ctx: &mut ToStreamContext) -> Result<StreamPlanRef>;
49
50    /// convert the plan to streaming physical plan and satisfy the required distribution
51    fn to_stream_with_dist_required(
52        &self,
53        required_dist: &RequiredDist,
54        ctx: &mut ToStreamContext,
55    ) -> Result<StreamPlanRef> {
56        let ret = self.to_stream(ctx)?;
57        required_dist.streaming_enforce_if_not_satisfies(ret)
58    }
59
60    fn try_better_locality(&self, _columns: &[usize]) -> Option<LogicalPlanRef> {
61        None
62    }
63}
64
65/// Try to enforce the locality requirement on the given columns.
66/// If a better plan can be found, return the better plan.
67/// If no better plan can be found, and locality backfill is enabled, wrap the plan
68/// with `LogicalLocalityProvider`.
69/// Otherwise, return the plan as is.
70pub fn try_enforce_locality_requirement(
71    plan: LogicalPlanRef,
72    columns: &[usize],
73    locality_backfill_enabled: bool,
74) -> LogicalPlanRef {
75    assert!(!columns.is_empty());
76    if let Some(better_plan) = plan.try_better_locality(columns) {
77        better_plan
78    } else if locality_backfill_enabled {
79        LogicalLocalityProvider::new(plan, columns.to_owned()).into()
80    } else {
81        plan
82    }
83}
84
85pub fn stream_enforce_eowc_requirement(
86    ctx: OptimizerContextRef,
87    plan: StreamPlanRef,
88    emit_on_window_close: bool,
89) -> Result<StreamPlanRef> {
90    if emit_on_window_close && !plan.emit_on_window_close() {
91        let watermark_groups = plan.watermark_columns().grouped();
92        let n_watermark_groups = watermark_groups.len();
93        if n_watermark_groups == 0 {
94            Err(ErrorCode::NotSupported(
95                "The query cannot be executed in Emit-On-Window-Close mode.".to_owned(),
96                "Try define a watermark column in the source, or avoid aggregation without GROUP BY".to_owned(),
97            )
98            .into())
99        } else {
100            let first_watermark_group = watermark_groups.values().next().unwrap();
101            let watermark_col_idx = first_watermark_group.indices().next().unwrap();
102            if n_watermark_groups > 1 {
103                ctx.warn_to_user(format!(
104                    "There are multiple unrelated watermark columns in the query, the first one `{}` is used.",
105                    FieldDisplay(&plan.schema()[watermark_col_idx])
106                ));
107            }
108            Ok(StreamEowcSort::new(plan, watermark_col_idx).into())
109        }
110    } else {
111        Ok(plan)
112    }
113}
114
115#[derive(Debug, Clone)]
116pub struct RewriteStreamContext {
117    share_rewrite_map: HashMap<ShareId, (LogicalPlanRef, ColIndexMapping)>,
118    // Snapshot backfill needs upstream table primary-key semantics during logical rewrite
119    // so operators above `LogicalScan` can preserve hidden primary-key columns before
120    // `StreamTableScan` is built. Other backfill types keep logical stream-key semantics.
121    backfill_type: BackfillType,
122    locality_backfill_enabled: bool,
123}
124
125impl RewriteStreamContext {
126    pub fn new_with_backfill_type(
127        backfill_type: BackfillType,
128        locality_backfill_enabled: bool,
129    ) -> Self {
130        Self {
131            share_rewrite_map: HashMap::new(),
132            backfill_type,
133            locality_backfill_enabled,
134        }
135    }
136
137    pub fn backfill_type(&self) -> BackfillType {
138        self.backfill_type
139    }
140
141    pub fn locality_backfill_enabled(&self) -> bool {
142        self.locality_backfill_enabled
143    }
144
145    pub fn add_rewrite_result(
146        &mut self,
147        share_id: ShareId,
148        plan_ref: LogicalPlanRef,
149        col_change: ColIndexMapping,
150    ) {
151        let prev = self
152            .share_rewrite_map
153            .insert(share_id, (plan_ref, col_change));
154        assert!(prev.is_none());
155    }
156
157    pub fn get_rewrite_result(
158        &self,
159        share_id: ShareId,
160    ) -> Option<&(LogicalPlanRef, ColIndexMapping)> {
161        self.share_rewrite_map.get(&share_id)
162    }
163}
164
165#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
166pub enum BackfillType {
167    Replicated,
168    /// Frontend-only variant for snapshot-free sinks. It is serialized as
169    /// `StreamScanType::UpstreamOnly`, but derives an upsert stream kind.
170    UpstreamOnlySink,
171    ArrangementBackfill,
172    SnapshotBackfill,
173    /// Frontend-only variant for sinks created with `since_timestamp`.
174    /// It is serialized as `StreamScanType::SnapshotBackfill`, but derives
175    /// the same upsert stream kind as upstream-only sinks.
176    SnapshotBackfillSinceTimestamp,
177}
178
179impl BackfillType {
180    pub fn without_snapshot(self) -> bool {
181        matches!(
182            self,
183            BackfillType::UpstreamOnlySink | BackfillType::SnapshotBackfillSinceTimestamp
184        )
185    }
186
187    pub fn is_snapshot_backfill(self) -> bool {
188        matches!(
189            self,
190            BackfillType::SnapshotBackfill | BackfillType::SnapshotBackfillSinceTimestamp
191        )
192    }
193
194    pub fn to_stream_scan_type(self, is_cross_db: bool) -> StreamScanType {
195        if is_cross_db {
196            return StreamScanType::CrossDbSnapshotBackfill;
197        }
198
199        match self {
200            BackfillType::Replicated | BackfillType::UpstreamOnlySink => {
201                StreamScanType::UpstreamOnly
202            }
203            BackfillType::ArrangementBackfill => StreamScanType::ArrangementBackfill,
204            BackfillType::SnapshotBackfill | BackfillType::SnapshotBackfillSinceTimestamp => {
205                StreamScanType::SnapshotBackfill
206            }
207        }
208    }
209}
210
211#[derive(Debug, Clone)]
212pub struct ToStreamContext {
213    share_to_stream_map: HashMap<ShareId, StreamPlanRef>,
214    emit_on_window_close: bool,
215    backfill_type: BackfillType,
216}
217
218impl ToStreamContext {
219    pub fn new_with_backfill_type(emit_on_window_close: bool, backfill_type: BackfillType) -> Self {
220        Self {
221            share_to_stream_map: HashMap::new(),
222            emit_on_window_close,
223            backfill_type,
224        }
225    }
226
227    pub fn backfill_type(&self) -> BackfillType {
228        self.backfill_type
229    }
230
231    pub fn add_to_stream_result(&mut self, share_id: ShareId, plan_ref: StreamPlanRef) {
232        self.share_to_stream_map
233            .try_insert(share_id, plan_ref)
234            .unwrap();
235    }
236
237    pub fn get_to_stream_result(&self, share_id: ShareId) -> Option<&StreamPlanRef> {
238        self.share_to_stream_map.get(&share_id)
239    }
240
241    pub fn emit_on_window_close(&self) -> bool {
242        self.emit_on_window_close
243    }
244}
245
246/// `ToBatch` allows to convert a logical plan node to batch physical node
247/// with an optional required order.
248///
249/// The generated plan has single distribution and doesn't have any exchange nodes inserted.
250/// Use either [`ToLocalBatch`] or [`ToDistributedBatch`] after `ToBatch` to get a distributed plan.
251///
252/// To implement this trait you can choose one of the two ways:
253/// - Implement `to_batch` and use the default implementation of `to_batch_with_order_required`
254/// - Or, if a better plan can be generated when a required order is given, you can implement
255///   `to_batch_with_order_required`, and implement `to_batch` with
256///   `to_batch_with_order_required(&Order::any())`.
257pub trait ToBatch {
258    /// `to_batch` is equivalent to `to_batch_with_order_required(&Order::any())`
259    fn to_batch(&self) -> Result<BatchPlanRef>;
260    /// convert the plan to batch physical plan and satisfy the required Order
261    fn to_batch_with_order_required(&self, required_order: &Order) -> Result<BatchPlanRef> {
262        let ret = self.to_batch()?;
263        required_order.enforce_if_not_satisfies(ret)
264    }
265}
266
267/// Converts a batch physical plan to local plan for local execution.
268///
269/// This is quite similar to `ToBatch`, but different in several ways. For example it converts
270/// scan to exchange + scan.
271pub trait ToLocalBatch {
272    fn to_local(&self) -> Result<BatchPlanRef>;
273
274    /// Convert the plan to batch local physical plan and satisfy the required Order
275    fn to_local_with_order_required(&self, required_order: &Order) -> Result<BatchPlanRef> {
276        let ret = self.to_local()?;
277        required_order.enforce_if_not_satisfies(ret)
278    }
279}
280
281/// `ToDistributedBatch` allows to convert a batch physical plan to distributed batch plan, by
282/// insert exchange node, with an optional required order and distributed.
283///
284/// To implement this trait you can choose one of the two ways:
285/// - Implement `to_distributed` and use the default implementation of
286///   `to_distributed_with_required`
287/// - Or, if a better plan can be generated when a required order is given, you can implement
288///   `to_distributed_with_required`, and implement `to_distributed` with
289///   `to_distributed_with_required(&Order::any(), &RequiredDist::Any)`
290pub trait ToDistributedBatch {
291    /// `to_distributed` is equivalent to `to_distributed_with_required(&Order::any(),
292    /// &RequiredDist::Any)`
293    fn to_distributed(&self) -> Result<BatchPlanRef>;
294    /// insert the exchange in batch physical plan to satisfy the required Distribution and Order.
295    fn to_distributed_with_required(
296        &self,
297        required_order: &Order,
298        required_dist: &RequiredDist,
299    ) -> Result<BatchPlanRef> {
300        let ret = self.to_distributed()?;
301        let ret = required_order.enforce_if_not_satisfies(ret)?;
302        required_dist.batch_enforce_if_not_satisfies(ret, required_order)
303    }
304}
risingwave_frontend/optimizer/plan_node/convert.rs

risingwave_frontend/optimizer/plan_node/
convert.rs
