risingwave_frontend/optimizer/
optimizer_context.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 core::fmt::Formatter;
16use std::cell::{Cell, RefCell, RefMut};
17use std::collections::HashMap;
18use std::marker::PhantomData;
19use std::rc::{Rc, Weak};
20use std::sync::Arc;
21
22use risingwave_common::id::SourceId;
23use risingwave_sqlparser::ast::{ExplainFormat, ExplainOptions, ExplainType};
24
25use super::property::WatermarkGroupId;
26use crate::expr::{CorrelatedId, SessionTimezone};
27use crate::handler::HandlerArgs;
28use crate::optimizer::plan_node::generic::Share;
29use crate::optimizer::plan_node::{LogicalPlanRef, PlanNodeId, StreamPlanRef};
30use crate::session::SessionImpl;
31use crate::utils::{OverwriteOptions, WithOptions};
32use crate::{Explain, TableCatalog};
33
34const RESERVED_ID_NUM: u16 = 10000;
35
36type PhantomUnsend = PhantomData<Rc<()>>;
37
38/// The stable identity of a shared subplan.
39///
40/// Unlike [`PlanNodeId`], a `ShareId` survives rebuilding the wrapper plan node around a share.
41/// [`OptimizerContext`] tracks the current input used to rebuild wrappers after a DAG pass.
42#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
43pub struct ShareId(u32);
44
45#[derive(Debug)]
46pub(in crate::optimizer) struct ShareEntry<P> {
47    current: RefCell<P>,
48    plan_node_id: PlanNodeId,
49}
50
51impl<P> ShareEntry<P> {
52    fn update_current(&self, current: P) {
53        *self.current.borrow_mut() = current;
54    }
55
56    pub(in crate::optimizer) fn plan_node_id(&self) -> PlanNodeId {
57        self.plan_node_id
58    }
59}
60
61impl<P: Clone> ShareEntry<P> {
62    fn current(&self) -> P {
63        self.current.borrow().clone()
64    }
65}
66
67pub(in crate::optimizer) type ShareEntryRef<P> = Rc<ShareEntry<P>>;
68
69struct ShareTable<P> {
70    /// Weak entries avoid a reference cycle through `PlanBase::ctx` in the stored plans. A live
71    /// share handle owns the corresponding strong entry.
72    entries: HashMap<ShareId, Weak<ShareEntry<P>>>,
73}
74
75impl<P> Default for ShareTable<P> {
76    fn default() -> Self {
77        Self {
78            entries: HashMap::new(),
79        }
80    }
81}
82
83pub struct OptimizerContext {
84    session_ctx: Arc<SessionImpl>,
85    /// The original SQL string, used for debugging.
86    sql: Arc<str>,
87    /// Normalized SQL string. See [`HandlerArgs::normalize_sql`].
88    normalized_sql: String,
89    /// Explain options
90    explain_options: ExplainOptions,
91    /// Store the trace of optimizer
92    optimizer_trace: RefCell<Vec<String>>,
93    /// Store the optimized logical plan of optimizer
94    logical_explain: RefCell<Option<String>>,
95    /// Store options or properties from the `with` clause
96    with_options: WithOptions,
97    /// Store the Session Timezone and whether it was used.
98    session_timezone: RefCell<SessionTimezone>,
99    /// Total number of optimization rules have been applied.
100    total_rule_applied: RefCell<usize>,
101    /// Store the configs can be overwritten in with clause
102    /// if not specified, use the value from session variable.
103    overwrite_options: OverwriteOptions,
104    /// Mapping from Iceberg table identifier to the current snapshot and its RisingWave commit
105    /// boundary. Used to keep multiple scans of the same table consistent within one query.
106    iceberg_snapshot_info_map: RefCell<HashMap<SourceId, Option<IcebergSnapshotInfo>>>,
107    /// Batch materialized view candidates for exact-match rewriting.
108    batch_mview_candidates: RefCell<Vec<MaterializedViewCandidate>>,
109
110    /// Last assigned plan node ID.
111    last_plan_node_id: Cell<i32>,
112    /// Last assigned share ID.
113    last_share_id: Cell<u32>,
114    /// Last assigned correlated ID.
115    last_correlated_id: Cell<u32>,
116    /// Last assigned expr display ID.
117    last_expr_display_id: Cell<usize>,
118    /// Last assigned watermark group ID.
119    last_watermark_group_id: Cell<u32>,
120
121    /// Tracks the current input of each live share while DAG-aware passes rebuild wrappers.
122    logical_share_table: RefCell<ShareTable<LogicalPlanRef>>,
123    stream_share_table: RefCell<ShareTable<StreamPlanRef>>,
124
125    _phantom: PhantomUnsend,
126}
127
128#[derive(Clone, Debug)]
129pub struct MaterializedViewCandidate {
130    pub plan: LogicalPlanRef,
131    pub table: Arc<TableCatalog>,
132}
133
134#[derive(Clone, Copy, Debug, PartialEq, Eq)]
135pub(crate) struct IcebergSnapshotInfo {
136    pub(crate) snapshot_id: i64,
137    pub(crate) commit_epoch: Option<u64>,
138}
139
140pub(in crate::optimizer) struct LastAssignedIds {
141    last_plan_node_id: i32,
142    last_correlated_id: u32,
143    last_expr_display_id: usize,
144    last_watermark_group_id: u32,
145}
146
147pub type OptimizerContextRef = Rc<OptimizerContext>;
148
149impl OptimizerContext {
150    /// Create a new [`OptimizerContext`] from the given [`HandlerArgs`], with empty
151    /// [`ExplainOptions`].
152    pub fn from_handler_args(handler_args: HandlerArgs) -> Self {
153        Self::new(handler_args, ExplainOptions::default())
154    }
155
156    /// Create a new [`OptimizerContext`] from the given [`HandlerArgs`] and [`ExplainOptions`].
157    pub fn new(mut handler_args: HandlerArgs, explain_options: ExplainOptions) -> Self {
158        let session_timezone = RefCell::new(SessionTimezone::new(
159            handler_args.session.config().timezone(),
160        ));
161        let overwrite_options = OverwriteOptions::new(&mut handler_args);
162        Self {
163            session_ctx: handler_args.session,
164            sql: handler_args.sql,
165            normalized_sql: handler_args.normalized_sql,
166            explain_options,
167            optimizer_trace: RefCell::new(vec![]),
168            logical_explain: RefCell::new(None),
169            with_options: handler_args.with_options,
170            session_timezone,
171            total_rule_applied: RefCell::new(0),
172            overwrite_options,
173            iceberg_snapshot_info_map: RefCell::new(HashMap::new()),
174            batch_mview_candidates: RefCell::new(Vec::new()),
175
176            last_plan_node_id: Cell::new(RESERVED_ID_NUM.into()),
177            last_share_id: Cell::new(0),
178            last_correlated_id: Cell::new(0),
179            last_expr_display_id: Cell::new(RESERVED_ID_NUM.into()),
180            last_watermark_group_id: Cell::new(RESERVED_ID_NUM.into()),
181
182            logical_share_table: RefCell::new(ShareTable::default()),
183            stream_share_table: RefCell::new(ShareTable::default()),
184
185            _phantom: Default::default(),
186        }
187    }
188
189    #[cfg(test)]
190    pub fn mock() -> OptimizerContextRef {
191        Self {
192            session_ctx: Arc::new(SessionImpl::mock()),
193            sql: Arc::from(""),
194            normalized_sql: "".to_owned(),
195            explain_options: ExplainOptions::default(),
196            optimizer_trace: RefCell::new(vec![]),
197            logical_explain: RefCell::new(None),
198            with_options: Default::default(),
199            session_timezone: RefCell::new(SessionTimezone::new("UTC".into())),
200            total_rule_applied: RefCell::new(0),
201            overwrite_options: OverwriteOptions::default(),
202            iceberg_snapshot_info_map: RefCell::new(HashMap::new()),
203            batch_mview_candidates: RefCell::new(Vec::new()),
204
205            last_plan_node_id: Cell::new(0),
206            last_share_id: Cell::new(0),
207            last_correlated_id: Cell::new(0),
208            last_expr_display_id: Cell::new(0),
209            last_watermark_group_id: Cell::new(0),
210
211            logical_share_table: RefCell::new(ShareTable::default()),
212            stream_share_table: RefCell::new(ShareTable::default()),
213
214            _phantom: Default::default(),
215        }
216        .into()
217    }
218
219    pub fn next_plan_node_id(&self) -> PlanNodeId {
220        self.last_plan_node_id.update(|id| id + 1);
221        PlanNodeId(self.last_plan_node_id.get())
222    }
223
224    fn next_share_id(&self) -> ShareId {
225        self.last_share_id.update(|id| id + 1);
226        ShareId(self.last_share_id.get())
227    }
228
229    fn share_entry<P>(
230        table: &RefCell<ShareTable<P>>,
231        share_id: ShareId,
232        convention: &str,
233    ) -> ShareEntryRef<P> {
234        table
235            .borrow()
236            .entries
237            .get(&share_id)
238            .unwrap_or_else(|| panic!("{convention} share {share_id:?} is not registered"))
239            .upgrade()
240            .unwrap_or_else(|| panic!("{convention} share {share_id:?} is no longer live"))
241    }
242
243    fn register_share<P: Clone>(&self, table: &RefCell<ShareTable<P>>, input: P) -> Share<P> {
244        let share_id = self.next_share_id();
245        let entry = Rc::new(ShareEntry {
246            current: RefCell::new(input.clone()),
247            plan_node_id: self.next_plan_node_id(),
248        });
249        table
250            .borrow_mut()
251            .entries
252            .try_insert(share_id, Rc::downgrade(&entry))
253            .expect("share id must be unique");
254        Share::new(share_id, input, entry)
255    }
256
257    fn share<P: Clone>(
258        &self,
259        table: &RefCell<ShareTable<P>>,
260        share_id: ShareId,
261        convention: &str,
262    ) -> Share<P> {
263        let entry = Self::share_entry(table, share_id, convention);
264        let input = entry.current();
265        Share::new(share_id, input, entry)
266    }
267
268    fn update_share<P>(
269        table: &RefCell<ShareTable<P>>,
270        share_id: ShareId,
271        new_input: P,
272        convention: &str,
273    ) {
274        let entry = Self::share_entry(table, share_id, convention);
275        entry.update_current(new_input);
276    }
277
278    pub(in crate::optimizer) fn register_logical_share(
279        &self,
280        input: LogicalPlanRef,
281    ) -> Share<LogicalPlanRef> {
282        self.register_share(&self.logical_share_table, input)
283    }
284
285    pub(in crate::optimizer) fn logical_share(&self, share_id: ShareId) -> Share<LogicalPlanRef> {
286        self.share(&self.logical_share_table, share_id, "logical")
287    }
288
289    pub(in crate::optimizer) fn update_logical_share(
290        &self,
291        share_id: ShareId,
292        new_input: LogicalPlanRef,
293    ) {
294        Self::update_share(&self.logical_share_table, share_id, new_input, "logical");
295    }
296
297    pub(in crate::optimizer) fn register_stream_share(
298        &self,
299        input: StreamPlanRef,
300    ) -> Share<StreamPlanRef> {
301        self.register_share(&self.stream_share_table, input)
302    }
303
304    pub(in crate::optimizer) fn update_stream_share(
305        &self,
306        share_id: ShareId,
307        new_input: StreamPlanRef,
308    ) {
309        Self::update_share(&self.stream_share_table, share_id, new_input, "stream");
310    }
311
312    pub fn next_correlated_id(&self) -> CorrelatedId {
313        self.last_correlated_id.update(|id| id + 1);
314        self.last_correlated_id.get()
315    }
316
317    pub fn next_expr_display_id(&self) -> usize {
318        self.last_expr_display_id.update(|id| id + 1);
319        self.last_expr_display_id.get()
320    }
321
322    pub fn next_watermark_group_id(&self) -> WatermarkGroupId {
323        self.last_watermark_group_id.update(|id| id + 1);
324        self.last_watermark_group_id.get()
325    }
326
327    pub(in crate::optimizer) fn backup_elem_ids(&self) -> LastAssignedIds {
328        LastAssignedIds {
329            last_plan_node_id: self.last_plan_node_id.get(),
330            last_correlated_id: self.last_correlated_id.get(),
331            last_expr_display_id: self.last_expr_display_id.get(),
332            last_watermark_group_id: self.last_watermark_group_id.get(),
333        }
334    }
335
336    /// This should only be called in [`crate::optimizer::plan_node::reorganize_elements_id`].
337    pub(in crate::optimizer) fn reset_elem_ids(&self) {
338        self.last_plan_node_id.set(0);
339        self.last_correlated_id.set(0);
340        self.last_expr_display_id.set(0);
341        self.last_watermark_group_id.set(0);
342    }
343
344    pub(in crate::optimizer) fn restore_elem_ids(&self, backup: LastAssignedIds) {
345        self.last_plan_node_id.set(backup.last_plan_node_id);
346        self.last_correlated_id.set(backup.last_correlated_id);
347        self.last_expr_display_id.set(backup.last_expr_display_id);
348        self.last_watermark_group_id
349            .set(backup.last_watermark_group_id);
350    }
351
352    pub fn add_rule_applied(&self, num: usize) {
353        *self.total_rule_applied.borrow_mut() += num;
354    }
355
356    pub fn total_rule_applied(&self) -> usize {
357        *self.total_rule_applied.borrow()
358    }
359
360    pub fn is_explain_verbose(&self) -> bool {
361        self.explain_options.verbose
362    }
363
364    pub fn is_explain_trace(&self) -> bool {
365        self.explain_options.trace
366    }
367
368    fn is_explain_logical(&self) -> bool {
369        self.explain_options.explain_type == ExplainType::Logical
370    }
371
372    pub fn trace(&self, str: impl Into<String>) {
373        // If explain type is logical, do not store the trace for any optimizations beyond logical.
374        if self.is_explain_logical() && self.logical_explain.borrow().is_some() {
375            return;
376        }
377        let mut optimizer_trace = self.optimizer_trace.borrow_mut();
378        let string = str.into();
379        tracing::info!(target: "explain_trace", "\n{}", string);
380        optimizer_trace.push(string);
381        optimizer_trace.push("\n".to_owned());
382    }
383
384    pub fn warn_to_user(&self, str: impl Into<String>) {
385        self.session_ctx().notice_to_user(str);
386    }
387
388    fn explain_plan_impl(&self, plan: &impl Explain) -> String {
389        match self.explain_options.explain_format {
390            ExplainFormat::Text => plan.explain_to_string(),
391            ExplainFormat::Json => plan.explain_to_json(),
392            ExplainFormat::Xml => plan.explain_to_xml(),
393            ExplainFormat::Yaml => plan.explain_to_yaml(),
394            ExplainFormat::Dot => plan.explain_to_dot(),
395        }
396    }
397
398    pub fn may_store_explain_logical(&self, plan: &LogicalPlanRef) {
399        if self.is_explain_logical() {
400            let str = self.explain_plan_impl(plan);
401            *self.logical_explain.borrow_mut() = Some(str);
402        }
403    }
404
405    pub fn take_logical(&self) -> Option<String> {
406        self.logical_explain.borrow_mut().take()
407    }
408
409    pub fn take_trace(&self) -> Vec<String> {
410        self.optimizer_trace.borrow_mut().drain(..).collect()
411    }
412
413    pub fn with_options(&self) -> &WithOptions {
414        &self.with_options
415    }
416
417    pub fn overwrite_options(&self) -> &OverwriteOptions {
418        &self.overwrite_options
419    }
420
421    pub fn add_batch_mview_candidate(&self, table: Arc<TableCatalog>, plan: LogicalPlanRef) {
422        self.batch_mview_candidates
423            .borrow_mut()
424            .push(MaterializedViewCandidate { plan, table });
425    }
426
427    pub fn batch_mview_candidates(&self) -> std::cell::Ref<'_, Vec<MaterializedViewCandidate>> {
428        self.batch_mview_candidates.borrow()
429    }
430
431    pub fn session_ctx(&self) -> &Arc<SessionImpl> {
432        &self.session_ctx
433    }
434
435    pub fn batch_plan_dml_wait_persistence(&self) -> bool {
436        let session_config = self.session_ctx.config();
437        !session_config.implicit_flush() && session_config.dml_wait_persistence()
438    }
439
440    /// Return the original SQL.
441    pub fn sql(&self) -> &str {
442        &self.sql
443    }
444
445    /// Return the normalized SQL.
446    pub fn normalized_sql(&self) -> &str {
447        &self.normalized_sql
448    }
449
450    pub fn session_timezone(&self) -> RefMut<'_, SessionTimezone> {
451        self.session_timezone.borrow_mut()
452    }
453
454    pub fn get_session_timezone(&self) -> String {
455        self.session_timezone.borrow().timezone()
456    }
457
458    pub(crate) fn iceberg_snapshot_info_map(
459        &self,
460    ) -> RefMut<'_, HashMap<SourceId, Option<IcebergSnapshotInfo>>> {
461        self.iceberg_snapshot_info_map.borrow_mut()
462    }
463}
464
465impl std::fmt::Debug for OptimizerContext {
466    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
467        write!(
468            f,
469            "QueryContext {{ sql = {}, explain_options = {}, with_options = {:?}, last_plan_node_id = {}, last_correlated_id = {} }}",
470            self.sql,
471            self.explain_options,
472            self.with_options,
473            self.last_plan_node_id.get(),
474            self.last_correlated_id.get(),
475        )
476    }
477}
risingwave_frontend/optimizer/optimizer_context.rs

risingwave_frontend/optimizer/
optimizer_context.rs
