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

risingwave_frontend/optimizer/plan_node/
convert.rs
