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

risingwave_frontend/optimizer/
optimizer_context.rs
