risingwave_frontend/optimizer/
optimizer_context.rs1use core::fmt::Formatter;
16use std::cell::{Cell, RefCell, RefMut};
17use std::collections::HashMap;
18use std::marker::PhantomData;
19use std::rc::Rc;
20use std::sync::Arc;
21
22use risingwave_sqlparser::ast::{ExplainFormat, ExplainOptions, ExplainType};
23
24use super::property::WatermarkGroupId;
25use crate::PlanRef;
26use crate::binder::ShareId;
27use crate::expr::{CorrelatedId, SessionTimezone};
28use crate::handler::HandlerArgs;
29use crate::optimizer::plan_node::PlanNodeId;
30use crate::session::SessionImpl;
31use crate::utils::{OverwriteOptions, WithOptions};
32
33const RESERVED_ID_NUM: u16 = 10000;
34
35type PhantomUnsend = PhantomData<Rc<()>>;
36
37pub struct OptimizerContext {
38 session_ctx: Arc<SessionImpl>,
39 sql: Arc<str>,
41 normalized_sql: String,
43 explain_options: ExplainOptions,
45 optimizer_trace: RefCell<Vec<String>>,
47 logical_explain: RefCell<Option<String>>,
49 with_options: WithOptions,
51 session_timezone: RefCell<SessionTimezone>,
53 total_rule_applied: RefCell<usize>,
55 overwrite_options: OverwriteOptions,
58 rcte_cache: RefCell<HashMap<ShareId, PlanRef>>,
61
62 last_plan_node_id: Cell<i32>,
64 last_correlated_id: Cell<u32>,
66 last_expr_display_id: Cell<usize>,
68 last_watermark_group_id: Cell<u32>,
70
71 _phantom: PhantomUnsend,
72}
73
74pub(in crate::optimizer) struct LastAssignedIds {
75 last_plan_node_id: i32,
76 last_correlated_id: u32,
77 last_expr_display_id: usize,
78 last_watermark_group_id: u32,
79}
80
81pub type OptimizerContextRef = Rc<OptimizerContext>;
82
83impl OptimizerContext {
84 pub fn from_handler_args(handler_args: HandlerArgs) -> Self {
87 Self::new(handler_args, ExplainOptions::default())
88 }
89
90 pub fn new(mut handler_args: HandlerArgs, explain_options: ExplainOptions) -> Self {
92 let session_timezone = RefCell::new(SessionTimezone::new(
93 handler_args.session.config().timezone().to_owned(),
94 ));
95 let overwrite_options = OverwriteOptions::new(&mut handler_args);
96 Self {
97 session_ctx: handler_args.session,
98 sql: handler_args.sql,
99 normalized_sql: handler_args.normalized_sql,
100 explain_options,
101 optimizer_trace: RefCell::new(vec![]),
102 logical_explain: RefCell::new(None),
103 with_options: handler_args.with_options,
104 session_timezone,
105 total_rule_applied: RefCell::new(0),
106 overwrite_options,
107 rcte_cache: RefCell::new(HashMap::new()),
108
109 last_plan_node_id: Cell::new(RESERVED_ID_NUM.into()),
110 last_correlated_id: Cell::new(0),
111 last_expr_display_id: Cell::new(RESERVED_ID_NUM.into()),
112 last_watermark_group_id: Cell::new(RESERVED_ID_NUM.into()),
113
114 _phantom: Default::default(),
115 }
116 }
117
118 #[cfg(test)]
120 #[expect(clippy::unused_async)]
121 pub async fn mock() -> OptimizerContextRef {
122 Self {
123 session_ctx: Arc::new(SessionImpl::mock()),
124 sql: Arc::from(""),
125 normalized_sql: "".to_owned(),
126 explain_options: ExplainOptions::default(),
127 optimizer_trace: RefCell::new(vec![]),
128 logical_explain: RefCell::new(None),
129 with_options: Default::default(),
130 session_timezone: RefCell::new(SessionTimezone::new("UTC".into())),
131 total_rule_applied: RefCell::new(0),
132 overwrite_options: OverwriteOptions::default(),
133 rcte_cache: RefCell::new(HashMap::new()),
134
135 last_plan_node_id: Cell::new(0),
136 last_correlated_id: Cell::new(0),
137 last_expr_display_id: Cell::new(0),
138 last_watermark_group_id: Cell::new(0),
139
140 _phantom: Default::default(),
141 }
142 .into()
143 }
144
145 pub fn next_plan_node_id(&self) -> PlanNodeId {
146 self.last_plan_node_id.update(|id| id + 1);
147 PlanNodeId(self.last_plan_node_id.get())
148 }
149
150 pub fn next_correlated_id(&self) -> CorrelatedId {
151 self.last_correlated_id.update(|id| id + 1);
152 self.last_correlated_id.get()
153 }
154
155 pub fn next_expr_display_id(&self) -> usize {
156 self.last_expr_display_id.update(|id| id + 1);
157 self.last_expr_display_id.get()
158 }
159
160 pub fn next_watermark_group_id(&self) -> WatermarkGroupId {
161 self.last_watermark_group_id.update(|id| id + 1);
162 self.last_watermark_group_id.get()
163 }
164
165 pub(in crate::optimizer) fn backup_elem_ids(&self) -> LastAssignedIds {
166 LastAssignedIds {
167 last_plan_node_id: self.last_plan_node_id.get(),
168 last_correlated_id: self.last_correlated_id.get(),
169 last_expr_display_id: self.last_expr_display_id.get(),
170 last_watermark_group_id: self.last_watermark_group_id.get(),
171 }
172 }
173
174 pub(in crate::optimizer) fn reset_elem_ids(&self) {
176 self.last_plan_node_id.set(0);
177 self.last_correlated_id.set(0);
178 self.last_expr_display_id.set(0);
179 self.last_watermark_group_id.set(0);
180 }
181
182 pub(in crate::optimizer) fn restore_elem_ids(&self, backup: LastAssignedIds) {
183 self.last_plan_node_id.set(backup.last_plan_node_id);
184 self.last_correlated_id.set(backup.last_correlated_id);
185 self.last_expr_display_id.set(backup.last_expr_display_id);
186 self.last_watermark_group_id
187 .set(backup.last_watermark_group_id);
188 }
189
190 pub fn add_rule_applied(&self, num: usize) {
191 *self.total_rule_applied.borrow_mut() += num;
192 }
193
194 pub fn total_rule_applied(&self) -> usize {
195 *self.total_rule_applied.borrow()
196 }
197
198 pub fn is_explain_verbose(&self) -> bool {
199 self.explain_options.verbose
200 }
201
202 pub fn is_explain_trace(&self) -> bool {
203 self.explain_options.trace
204 }
205
206 pub fn is_explain_backfill(&self) -> bool {
207 self.explain_options.backfill
208 }
209
210 pub fn explain_type(&self) -> ExplainType {
211 self.explain_options.explain_type.clone()
212 }
213
214 pub fn explain_format(&self) -> ExplainFormat {
215 self.explain_options.explain_format.clone()
216 }
217
218 pub fn is_explain_logical(&self) -> bool {
219 self.explain_type() == ExplainType::Logical
220 }
221
222 pub fn trace(&self, str: impl Into<String>) {
223 if self.is_explain_logical() && self.logical_explain.borrow().is_some() {
225 return;
226 }
227 let mut optimizer_trace = self.optimizer_trace.borrow_mut();
228 let string = str.into();
229 tracing::info!(target: "explain_trace", "\n{}", string);
230 optimizer_trace.push(string);
231 optimizer_trace.push("\n".to_owned());
232 }
233
234 pub fn warn_to_user(&self, str: impl Into<String>) {
235 self.session_ctx().notice_to_user(str);
236 }
237
238 pub fn store_logical(&self, str: impl Into<String>) {
239 *self.logical_explain.borrow_mut() = Some(str.into())
240 }
241
242 pub fn take_logical(&self) -> Option<String> {
243 self.logical_explain.borrow_mut().take()
244 }
245
246 pub fn take_trace(&self) -> Vec<String> {
247 self.optimizer_trace.borrow_mut().drain(..).collect()
248 }
249
250 pub fn with_options(&self) -> &WithOptions {
251 &self.with_options
252 }
253
254 pub fn overwrite_options(&self) -> &OverwriteOptions {
255 &self.overwrite_options
256 }
257
258 pub fn session_ctx(&self) -> &Arc<SessionImpl> {
259 &self.session_ctx
260 }
261
262 pub fn sql(&self) -> &str {
264 &self.sql
265 }
266
267 pub fn normalized_sql(&self) -> &str {
269 &self.normalized_sql
270 }
271
272 pub fn session_timezone(&self) -> RefMut<'_, SessionTimezone> {
273 self.session_timezone.borrow_mut()
274 }
275
276 pub fn get_session_timezone(&self) -> String {
277 self.session_timezone.borrow().timezone()
278 }
279
280 pub fn get_rcte_cache_plan(&self, id: &ShareId) -> Option<PlanRef> {
281 self.rcte_cache.borrow().get(id).cloned()
282 }
283
284 pub fn insert_rcte_cache_plan(&self, id: ShareId, plan: PlanRef) {
285 self.rcte_cache.borrow_mut().insert(id, plan);
286 }
287}
288
289impl std::fmt::Debug for OptimizerContext {
290 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
291 write!(
292 f,
293 "QueryContext {{ sql = {}, explain_options = {}, with_options = {:?}, last_plan_node_id = {}, last_correlated_id = {} }}",
294 self.sql,
295 self.explain_options,
296 self.with_options,
297 self.last_plan_node_id.get(),
298 self.last_correlated_id.get(),
299 )
300 }
301}