1use std::num::NonZeroU32;
16use std::ops::DerefMut;
17use std::sync::Arc;
18
19use risingwave_pb::catalog::PbVectorIndexInfo;
20
21pub mod plan_node;
22
23use plan_node::StreamFilter;
24pub use plan_node::{Explain, LogicalPlanRef, PlanRef};
25
26pub mod property;
27
28mod delta_join_solver;
29mod heuristic_optimizer;
30mod plan_rewriter;
31
32mod plan_visitor;
33
34#[cfg(feature = "datafusion")]
35pub use plan_visitor::DataFusionExecuteCheckerExt;
36pub use plan_visitor::{
37 ExecutionModeDecider, PlanVisitor, RelationCollectorVisitor, SysTableVisitor,
38};
39use risingwave_pb::plan_common::source_refresh_mode::RefreshMode;
40
41pub mod backfill_order_strategy;
42mod logical_optimization;
43mod optimizer_context;
44pub mod plan_expr_rewriter;
45mod plan_expr_visitor;
46mod rule;
47pub mod variant_key;
48
49use std::collections::{BTreeMap, HashMap};
50use std::marker::PhantomData;
51
52use educe::Educe;
53use fixedbitset::FixedBitSet;
54use itertools::Itertools;
55pub use logical_optimization::*;
56pub use optimizer_context::*;
57use plan_expr_rewriter::ConstEvalRewriter;
58use property::Order;
59use risingwave_common::bail;
60use risingwave_common::catalog::{
61 ColumnCatalog, ColumnDesc, ConflictBehavior, Field, FieldDisplay, Schema,
62};
63use risingwave_common::types::DataType;
64use risingwave_common::util::column_index_mapping::ColIndexMapping;
65use risingwave_common::util::iter_util::ZipEqDebug;
66use risingwave_connector::WithPropertiesExt;
67use risingwave_connector::sink::catalog::SinkFormatDesc;
68
69use self::heuristic_optimizer::ApplyOrder;
70use self::plan_node::generic::{self, PhysicalPlanRef};
71use self::plan_node::{
72 BatchProject, LogicalProject, LogicalSource, PartitionComputeInfo, StreamDml,
73 StreamMaterialize, StreamProject, StreamRowIdGen, StreamSink, StreamWatermarkFilter,
74 ToStreamContext, stream_enforce_eowc_requirement,
75};
76#[cfg(debug_assertions)]
77use self::plan_visitor::InputRefValidator;
78use self::plan_visitor::{CardinalityVisitor, StreamKeyChecker, has_batch_exchange};
79use self::property::{Cardinality, RequiredDist};
80use self::rule::*;
81use self::variant_key::variant_key_error;
82use crate::TableCatalog;
83use crate::catalog::table_catalog::TableType;
84use crate::catalog::{DatabaseId, SchemaId};
85use crate::error::{ErrorCode, Result};
86use crate::expr::TimestamptzExprFinder;
87use crate::handler::create_table::{CreateTableInfo, CreateTableProps};
88use crate::optimizer::plan_node::generic::{GenericPlanRef, SourceNodeKind, Union};
89use crate::optimizer::plan_node::{
90 BackfillType, Batch, BatchExchange, BatchPlanNodeType, BatchPlanRef, ConventionMarker,
91 PlanTreeNode, RewriteStreamContext, Stream, StreamExchange, StreamPlanRef, StreamUnion,
92 StreamUpstreamSinkUnion, StreamVectorIndexWrite, ToStream, VisitExprsRecursive,
93};
94use crate::optimizer::plan_visitor::{
95 LocalityProviderCounter, RwTimestampValidator, TemporalJoinValidator,
96};
97use crate::optimizer::property::Distribution;
98use crate::utils::{
99 ColIndexMappingRewriteExt, MV_REFRESH_INTERVAL_SEC_KEY, WithOptionsSecResolved,
100};
101
102#[derive(Educe)]
113#[educe(Debug, Clone)]
114pub struct PlanRoot<P: PlanPhase> {
115 pub plan: PlanRef<P::Convention>,
117 #[educe(Debug(ignore), Clone(method(PhantomData::clone)))]
119 _phase: PhantomData<P>,
120 required_dist: RequiredDist,
121 required_order: Order,
122 out_fields: FixedBitSet,
123 out_names: Vec<String>,
124}
125
126pub trait PlanPhase {
132 type Convention: ConventionMarker;
133}
134
135macro_rules! for_all_phase {
136 () => {
137 for_all_phase! {
138 { Logical, $crate::optimizer::plan_node::Logical },
139 { BatchOptimizedLogical, $crate::optimizer::plan_node::Logical },
140 { StreamOptimizedLogical, $crate::optimizer::plan_node::Stream },
141 { Batch, $crate::optimizer::plan_node::Batch },
142 { Stream, $crate::optimizer::plan_node::Stream }
143 }
144 };
145 ($({$phase:ident, $convention:ty}),+ $(,)?) => {
146 $(
147 paste::paste! {
148 pub struct [< PlanPhase$phase >];
149 impl PlanPhase for [< PlanPhase$phase >] {
150 type Convention = $convention;
151 }
152 pub type [< $phase PlanRoot >] = PlanRoot<[< PlanPhase$phase >]>;
153 }
154 )+
155 }
156}
157
158for_all_phase!();
159
160impl LogicalPlanRoot {
161 pub fn new_with_logical_plan(
162 plan: LogicalPlanRef,
163 required_dist: RequiredDist,
164 required_order: Order,
165 out_fields: FixedBitSet,
166 out_names: Vec<String>,
167 ) -> Self {
168 Self::new_inner(plan, required_dist, required_order, out_fields, out_names)
169 }
170}
171
172impl BatchPlanRoot {
173 pub fn new_with_batch_plan(
174 plan: BatchPlanRef,
175 required_dist: RequiredDist,
176 required_order: Order,
177 out_fields: FixedBitSet,
178 out_names: Vec<String>,
179 ) -> Self {
180 Self::new_inner(plan, required_dist, required_order, out_fields, out_names)
181 }
182}
183
184impl<P: PlanPhase> PlanRoot<P> {
185 fn new_inner(
186 plan: PlanRef<P::Convention>,
187 required_dist: RequiredDist,
188 required_order: Order,
189 out_fields: FixedBitSet,
190 out_names: Vec<String>,
191 ) -> Self {
192 let input_schema = plan.schema();
193 assert_eq!(input_schema.fields().len(), out_fields.len());
194 assert_eq!(out_fields.count_ones(..), out_names.len());
195
196 Self {
197 plan,
198 _phase: PhantomData,
199 required_dist,
200 required_order,
201 out_fields,
202 out_names,
203 }
204 }
205
206 fn into_phase<P2: PlanPhase>(self, plan: PlanRef<P2::Convention>) -> PlanRoot<P2> {
207 PlanRoot {
208 plan,
209 _phase: PhantomData,
210 required_dist: self.required_dist,
211 required_order: self.required_order,
212 out_fields: self.out_fields,
213 out_names: self.out_names,
214 }
215 }
216
217 pub fn set_out_names(&mut self, out_names: Vec<String>) -> Result<()> {
222 if out_names.len() != self.out_fields.count_ones(..) {
223 Err(ErrorCode::InvalidInputSyntax(
224 "number of column names does not match number of columns".to_owned(),
225 ))?
226 }
227 self.out_names = out_names;
228 Ok(())
229 }
230
231 pub fn schema(&self) -> Schema {
233 Schema {
236 fields: self
237 .out_fields
238 .ones()
239 .map(|i| self.plan.schema().fields()[i].clone())
240 .zip_eq_debug(&self.out_names)
241 .map(|(field, name)| Field {
242 name: name.clone(),
243 ..field
244 })
245 .collect(),
246 }
247 }
248}
249
250impl LogicalPlanRoot {
251 pub fn into_unordered_subplan(self) -> LogicalPlanRef {
255 if self.out_fields.count_ones(..) == self.out_fields.len() {
256 return self.plan;
257 }
258 LogicalProject::with_out_fields(self.plan, &self.out_fields).into()
259 }
260
261 pub fn into_array_agg(self) -> Result<LogicalPlanRef> {
265 use generic::Agg;
266 use plan_node::PlanAggCall;
267 use risingwave_common::types::ListValue;
268 use risingwave_expr::aggregate::PbAggKind;
269
270 use crate::expr::{ExprImpl, ExprType, FunctionCall, InputRef};
271 use crate::utils::{Condition, IndexSet};
272
273 let Ok(select_idx) = Itertools::exactly_one(self.out_fields.ones()) else {
274 bail!("subquery must return only one column");
275 };
276 let input_column_type = self.plan.schema().fields()[select_idx].data_type();
277 let return_type = DataType::list(input_column_type.clone());
278 let agg = Agg::new(
279 vec![PlanAggCall {
280 agg_type: PbAggKind::ArrayAgg.into(),
281 return_type: return_type.clone(),
282 inputs: vec![InputRef::new(select_idx, input_column_type.clone())],
283 distinct: false,
284 order_by: self.required_order.column_orders,
285 filter: Condition::true_cond(),
286 direct_args: vec![],
287 }],
288 IndexSet::empty(),
289 self.plan,
290 );
291 Ok(LogicalProject::create(
292 agg.into(),
293 vec![
294 FunctionCall::new(
295 ExprType::Coalesce,
296 vec![
297 InputRef::new(0, return_type).into(),
298 ExprImpl::literal_list(
299 ListValue::empty(&input_column_type),
300 input_column_type,
301 ),
302 ],
303 )
304 .unwrap()
305 .into(),
306 ],
307 ))
308 }
309
310 pub fn gen_optimized_logical_plan_for_stream(mut self) -> Result<LogicalPlanRoot> {
312 self.plan = LogicalOptimizer::gen_optimized_logical_plan_for_stream(self.plan.clone())?;
313 Ok(self)
314 }
315
316 pub fn gen_optimized_logical_plan_for_batch(self) -> Result<BatchOptimizedLogicalPlanRoot> {
318 let plan = LogicalOptimizer::gen_optimized_logical_plan_for_batch(self.plan.clone())?;
319 Ok(self.into_phase(plan))
320 }
321
322 pub fn gen_batch_plan(self) -> Result<BatchPlanRoot> {
323 self.gen_optimized_logical_plan_for_batch()?
324 .gen_batch_plan()
325 }
326}
327
328impl BatchOptimizedLogicalPlanRoot {
329 fn reject_variant_keys(&self) -> Result<()> {
332 if let Some(err) = StreamKeyChecker::Variant.visit(self.plan.clone()) {
333 return Err(variant_key_error(err));
334 }
335 let schema = self.plan.schema();
336 for order in &self.required_order.column_orders {
337 let field = &schema[order.column_index];
338 if field.data_type().contains_variant() {
339 return Err(variant_key_error(format!(
340 "VARIANT column \"{}\" should not be in the ORDER BY.",
341 FieldDisplay(field)
342 )));
343 }
344 }
345 Ok(())
346 }
347
348 pub fn gen_batch_plan(self) -> Result<BatchPlanRoot> {
350 self.reject_variant_keys()?;
351 if TemporalJoinValidator::exist_dangling_temporal_scan(self.plan.clone()) {
352 return Err(ErrorCode::NotSupported(
353 "do not support temporal join for batch queries".to_owned(),
354 "please use temporal join in streaming queries".to_owned(),
355 )
356 .into());
357 }
358
359 let ctx = self.plan.ctx();
360 let mut plan = inline_session_timezone_in_exprs(ctx.clone(), self.plan.clone())?;
362
363 plan = const_eval_exprs(plan)?;
365
366 if ctx.is_explain_trace() {
367 ctx.trace("Const eval exprs:");
368 ctx.trace(plan.explain_to_string());
369 }
370
371 let mut plan = plan.to_batch_with_order_required(&self.required_order)?;
373 if ctx.is_explain_trace() {
374 ctx.trace("To Batch Plan:");
375 ctx.trace(plan.explain_to_string());
376 }
377
378 plan = plan.optimize_by_rules(&OptimizationStage::<Batch>::new(
379 "Merge BatchProject",
380 vec![BatchProjectMergeRule::create()],
381 ApplyOrder::BottomUp,
382 ))?;
383
384 plan = inline_session_timezone_in_exprs(ctx.clone(), plan)?;
386
387 if ctx.is_explain_trace() {
388 ctx.trace("Inline Session Timezone:");
389 ctx.trace(plan.explain_to_string());
390 }
391
392 #[cfg(debug_assertions)]
393 InputRefValidator.validate(plan.clone());
394 assert_eq!(
395 *plan.distribution(),
396 Distribution::Single,
397 "{}",
398 plan.explain_to_string()
399 );
400 assert!(
401 !has_batch_exchange(plan.clone()),
402 "{}",
403 plan.explain_to_string()
404 );
405
406 let ctx = plan.ctx();
407 if ctx.is_explain_trace() {
408 ctx.trace("To Batch Physical Plan:");
409 ctx.trace(plan.explain_to_string());
410 }
411
412 Ok(self.into_phase(plan))
413 }
414
415 #[cfg(feature = "datafusion")]
416 pub fn gen_datafusion_logical_plan(
417 &self,
418 ) -> Result<Arc<datafusion::logical_expr::LogicalPlan>> {
419 use datafusion::logical_expr::{Expr as DFExpr, LogicalPlan, Projection, Sort};
420 use datafusion_common::Column;
421 use plan_visitor::LogicalPlanToDataFusionExt;
422
423 use crate::datafusion::{InputColumns, convert_column_order};
424
425 tracing::debug!(
426 "Converting RisingWave logical plan to DataFusion plan:\nRisingWave Plan: {:?}",
427 self.plan
428 );
429
430 let ctx = self.plan.ctx();
431 let mut plan = inline_session_timezone_in_exprs(ctx, self.plan.clone())?;
433 plan = const_eval_exprs(plan)?;
434
435 let mut df_plan = plan.to_datafusion_logical_plan()?;
436
437 if !self.required_order.is_any() {
438 let input_columns = InputColumns::new(df_plan.schema().as_ref(), plan.schema());
439 let expr = self
440 .required_order
441 .column_orders
442 .iter()
443 .map(|column_order| convert_column_order(column_order, &input_columns))
444 .collect_vec();
445 df_plan = Arc::new(LogicalPlan::Sort(Sort {
446 expr,
447 input: df_plan,
448 fetch: None,
449 }));
450 }
451
452 if self.out_names.len() < df_plan.schema().fields().len() {
453 let df_schema = df_plan.schema().as_ref();
454 let projection_exprs = self
455 .out_fields
456 .ones()
457 .zip_eq_debug(self.out_names.iter())
458 .map(|(i, name)| {
459 DFExpr::Column(Column::from(df_schema.qualified_field(i))).alias(name)
460 })
461 .collect_vec();
462 df_plan = Arc::new(LogicalPlan::Projection(Projection::try_new(
463 projection_exprs,
464 df_plan,
465 )?));
466 }
467
468 tracing::debug!("Converted DataFusion plan:\nDataFusion Plan: {:?}", df_plan);
469
470 Ok(df_plan)
471 }
472}
473
474impl BatchPlanRoot {
475 pub fn gen_batch_distributed_plan(mut self) -> Result<BatchPlanRef> {
477 self.required_dist = RequiredDist::single();
478 let mut plan = self.plan;
479
480 plan = plan.to_distributed_with_required(&self.required_order, &self.required_dist)?;
482
483 let ctx = plan.ctx();
484 if ctx.is_explain_trace() {
485 ctx.trace("To Batch Distributed Plan:");
486 ctx.trace(plan.explain_to_string());
487 }
488 if require_additional_exchange_on_root_in_distributed_mode(plan.clone()) {
489 plan =
490 BatchExchange::new(plan, self.required_order.clone(), Distribution::Single).into();
491 }
492
493 if self.out_fields.count_ones(..) != self.out_fields.len() {
495 plan =
496 BatchProject::new(generic::Project::with_out_fields(plan, &self.out_fields)).into();
497 }
498
499 let plan = plan.optimize_by_rules(&OptimizationStage::new(
501 "Push Limit To Scan",
502 vec![BatchPushLimitToScanRule::create()],
503 ApplyOrder::BottomUp,
504 ))?;
505
506 Ok(plan)
507 }
508
509 pub fn gen_batch_local_plan(self) -> Result<BatchPlanRef> {
511 let mut plan = self.plan;
512
513 plan = plan.to_local_with_order_required(&self.required_order)?;
515
516 let insert_exchange = match plan.distribution() {
519 Distribution::Single => require_additional_exchange_on_root_in_local_mode(plan.clone()),
520 _ => true,
521 };
522 if insert_exchange {
523 plan =
524 BatchExchange::new(plan, self.required_order.clone(), Distribution::Single).into()
525 }
526
527 if self.out_fields.count_ones(..) != self.out_fields.len() {
529 plan =
530 BatchProject::new(generic::Project::with_out_fields(plan, &self.out_fields)).into();
531 }
532
533 let ctx = plan.ctx();
534 if ctx.is_explain_trace() {
535 ctx.trace("To Batch Local Plan:");
536 ctx.trace(plan.explain_to_string());
537 }
538
539 let plan = plan.optimize_by_rules(&OptimizationStage::new(
541 "Push Limit To Scan",
542 vec![BatchPushLimitToScanRule::create()],
543 ApplyOrder::BottomUp,
544 ))?;
545
546 Ok(plan)
547 }
548}
549
550impl LogicalPlanRoot {
551 pub(crate) fn derive_backfill_type(&self, allow_snapshot_backfill: bool) -> BackfillType {
553 if allow_snapshot_backfill && self.should_use_snapshot_backfill() {
554 BackfillType::SnapshotBackfill
555 } else {
556 BackfillType::ArrangementBackfill
557 }
558 }
559
560 fn gen_optimized_stream_plan(
561 self,
562 emit_on_window_close: bool,
563 backfill_type: BackfillType,
564 ) -> Result<StreamOptimizedLogicalPlanRoot> {
565 let ctx = self.plan.ctx();
566 let _explain_trace = ctx.is_explain_trace();
567
568 let optimized_plan = self.gen_stream_plan(emit_on_window_close, backfill_type)?;
569
570 let mut plan = optimized_plan
571 .plan
572 .clone()
573 .optimize_by_rules(&OptimizationStage::new(
574 "Merge StreamProject",
575 vec![StreamProjectMergeRule::create()],
576 ApplyOrder::BottomUp,
577 ))?;
578
579 if ctx
580 .session_ctx()
581 .config()
582 .streaming_separate_consecutive_join()
583 {
584 plan = plan.optimize_by_rules(&OptimizationStage::new(
585 "Separate consecutive StreamHashJoin by no-shuffle StreamExchange",
586 vec![SeparateConsecutiveJoinRule::create()],
587 ApplyOrder::BottomUp,
588 ))?;
589 }
590
591 if ctx.session_ctx().config().streaming_enable_unaligned_join() {
595 plan = plan.optimize_by_rules(&OptimizationStage::new(
596 "Add Logstore for Unaligned join",
597 vec![AddLogstoreRule::create()],
598 ApplyOrder::BottomUp,
599 ))?;
600 }
601
602 if ctx.session_ctx().config().streaming_enable_delta_join()
603 && ctx.session_ctx().config().enable_index_selection()
604 {
605 plan = plan.optimize_by_rules(&OptimizationStage::new(
608 "To IndexDeltaJoin",
609 vec![IndexDeltaJoinRule::create()],
610 ApplyOrder::BottomUp,
611 ))?;
612 }
613 plan = inline_session_timezone_in_exprs(ctx.clone(), plan)?;
615
616 if ctx.is_explain_trace() {
617 ctx.trace("Inline session timezone:");
618 ctx.trace(plan.explain_to_string());
619 }
620
621 plan = const_eval_exprs(plan)?;
623
624 if ctx.is_explain_trace() {
625 ctx.trace("Const eval exprs:");
626 ctx.trace(plan.explain_to_string());
627 }
628
629 #[cfg(debug_assertions)]
630 InputRefValidator.validate(plan.clone());
631
632 if TemporalJoinValidator::exist_dangling_temporal_scan(plan.clone()) {
633 return Err(ErrorCode::NotSupported(
634 "exist dangling temporal scan".to_owned(),
635 "please check your temporal join syntax e.g. consider removing the right outer join if it is being used.".to_owned(),
636 ).into());
637 }
638
639 if RwTimestampValidator::select_rw_timestamp_in_stream_query(plan.clone()) {
640 return Err(ErrorCode::NotSupported(
641 "selecting `_rw_timestamp` in a streaming query is not allowed".to_owned(),
642 "please run the sql in batch mode or remove the column `_rw_timestamp` from the streaming query".to_owned(),
643 ).into());
644 }
645
646 if LocalityProviderCounter::count(plan.clone()) > 5 {
647 assert!(ctx.session_ctx().config().enable_locality_backfill());
649 risingwave_common::license::Feature::LocalityBackfill.check_available()?;
650 }
651
652 if ctx.missed_locality_providers() > 1
653 && risingwave_common::license::Feature::LocalityBackfill
654 .check_available()
655 .is_ok()
656 {
657 assert!(!ctx.session_ctx().config().enable_locality_backfill());
659 ctx.warn_to_user(format!(
660 "This streaming job has {} operators that could benefit from locality backfill. \
661 Consider enabling it with `SET enable_locality_backfill = true` for potentially \
662 faster backfill performance, when existing data volume in upstream(s) is large.",
663 ctx.missed_locality_providers()
664 ));
665 }
666
667 Ok(optimized_plan.into_phase(plan))
668 }
669
670 pub(crate) fn require_snapshot_backfill_for_batch_refresh(&self) -> Result<()> {
671 let ctx = self.plan.ctx();
672 let session_ctx = ctx.session_ctx();
673 let snapshot_backfill_enabled = session_ctx
674 .env()
675 .streaming_config()
676 .developer
677 .enable_snapshot_backfill
678 && session_ctx.config().streaming_use_snapshot_backfill();
679 if !snapshot_backfill_enabled {
680 return Err(ErrorCode::NotSupported(
681 "Batch refresh materialized view requires snapshot backfill".to_owned(),
682 format!(
683 "Please enable snapshot backfill or remove `{}` from the WITH clause.",
684 MV_REFRESH_INTERVAL_SEC_KEY
685 ),
686 )
687 .into());
688 }
689 if let Some(reason) = self.plan.forbid_snapshot_backfill() {
690 return Err(ErrorCode::NotSupported(
691 format!("Batch refresh materialized view requires snapshot backfill, but {reason}"),
692 "Please rewrite the query to avoid operators that forbid snapshot backfill."
693 .to_owned(),
694 )
695 .into());
696 }
697 Ok(())
698 }
699
700 fn gen_stream_plan(
702 self,
703 emit_on_window_close: bool,
704 backfill_type: BackfillType,
705 ) -> Result<StreamOptimizedLogicalPlanRoot> {
706 let ctx = self.plan.ctx();
707 let explain_trace = ctx.is_explain_trace();
708
709 let plan = {
710 {
711 if let Some(err) = StreamKeyChecker::Variant.visit(self.plan.clone()) {
712 return Err(variant_key_error(err));
713 }
714 if !ctx
715 .session_ctx()
716 .config()
717 .streaming_allow_jsonb_in_stream_key()
718 && let Some(err) = StreamKeyChecker::Jsonb.visit(self.plan.clone())
719 {
720 return Err(ErrorCode::NotSupported(
721 err,
722 "Using JSONB columns as part of the join or aggregation keys can severely impair performance. \
723 If you intend to proceed, force to enable it with: `set rw_streaming_allow_jsonb_in_stream_key to true`".to_owned(),
724 ).into());
725 }
726 let mut optimized_plan = self.gen_optimized_logical_plan_for_stream()?;
727 let (plan, out_col_change) = {
728 let (plan, out_col_change) = optimized_plan.plan.logical_rewrite_for_stream(
729 &mut RewriteStreamContext::new_with_backfill_type(backfill_type),
730 )?;
731 if out_col_change.is_injective() {
732 (plan, out_col_change)
733 } else {
734 let mut output_indices = (0..plan.schema().len()).collect_vec();
735 #[expect(unused_assignments)]
736 let (mut map, mut target_size) = out_col_change.into_parts();
737
738 target_size = plan.schema().len();
741 let mut tar_exists = vec![false; target_size];
742 for i in map.iter_mut().flatten() {
743 if tar_exists[*i] {
744 output_indices.push(*i);
745 *i = target_size;
746 target_size += 1;
747 } else {
748 tar_exists[*i] = true;
749 }
750 }
751 let plan =
752 LogicalProject::with_out_col_idx(plan, output_indices.into_iter());
753 let out_col_change = ColIndexMapping::new(map, target_size);
754 (plan.into(), out_col_change)
755 }
756 };
757 if explain_trace {
758 ctx.trace("Logical Rewrite For Stream:");
759 ctx.trace(plan.explain_to_string());
760 }
761
762 optimized_plan.required_dist =
763 out_col_change.rewrite_required_distribution(&optimized_plan.required_dist);
764 optimized_plan.required_order = out_col_change
765 .rewrite_required_order(&optimized_plan.required_order)
766 .unwrap();
767 optimized_plan.out_fields =
768 out_col_change.rewrite_bitset(&optimized_plan.out_fields);
769 let mut plan = plan.to_stream_with_dist_required(
770 &optimized_plan.required_dist,
771 &mut ToStreamContext::new_with_backfill_type(
772 emit_on_window_close,
773 backfill_type,
774 ),
775 )?;
776 plan = stream_enforce_eowc_requirement(ctx.clone(), plan, emit_on_window_close)?;
777 optimized_plan.into_phase(plan)
778 }
779 };
780
781 if explain_trace {
782 ctx.trace("To Stream Plan:");
783 ctx.trace(<PlanRef<Stream> as Explain>::explain_to_string(&plan.plan));
785 }
786 Ok(plan)
787 }
788
789 fn compute_cardinality(&self) -> Cardinality {
793 CardinalityVisitor.visit(self.plan.clone())
794 }
795
796 pub fn gen_table_plan(
798 self,
799 context: OptimizerContextRef,
800 table_name: String,
801 database_id: DatabaseId,
802 schema_id: SchemaId,
803 CreateTableInfo {
804 columns,
805 pk_column_ids,
806 row_id_index,
807 watermark_descs,
808 source_catalog,
809 version,
810 }: CreateTableInfo,
811 CreateTableProps {
812 definition,
813 append_only,
814 on_conflict,
815 with_version_columns,
816 webhook_info,
817 engine,
818 }: CreateTableProps,
819 ) -> Result<StreamMaterialize> {
820 let backfill_type = self.derive_backfill_type(false);
821 let stream_plan = self.gen_optimized_stream_plan(false, backfill_type)?;
823
824 assert!(!pk_column_ids.is_empty() || row_id_index.is_some());
825
826 let pk_column_indices = {
827 let mut id_to_idx = HashMap::new();
828
829 columns.iter().enumerate().for_each(|(idx, c)| {
830 id_to_idx.insert(c.column_id(), idx);
831 });
832 pk_column_ids
833 .iter()
834 .map(|c| id_to_idx.get(c).copied().unwrap()) .collect_vec()
836 };
837
838 fn inject_project_for_generated_column_if_needed(
839 columns: &[ColumnCatalog],
840 node: StreamPlanRef,
841 ) -> Result<StreamPlanRef> {
842 let exprs = LogicalSource::derive_output_exprs_from_generated_columns(columns)?;
843 if let Some(exprs) = exprs {
844 let logical_project = generic::Project::new(exprs, node);
845 return Ok(StreamProject::new(logical_project).into());
846 }
847 Ok(node)
848 }
849
850 #[derive(PartialEq, Debug, Copy, Clone)]
851 enum PrimaryKeyKind {
852 UserDefinedPrimaryKey,
853 NonAppendOnlyRowIdPk,
854 AppendOnlyRowIdPk,
855 }
856
857 fn inject_dml_node(
858 columns: &[ColumnCatalog],
859 append_only: bool,
860 stream_plan: StreamPlanRef,
861 pk_column_indices: &[usize],
862 kind: PrimaryKeyKind,
863 column_descs: Vec<ColumnDesc>,
864 ) -> Result<StreamPlanRef> {
865 let mut dml_node = StreamDml::new(stream_plan, append_only, column_descs).into();
866
867 dml_node = inject_project_for_generated_column_if_needed(columns, dml_node)?;
869
870 dml_node = match kind {
871 PrimaryKeyKind::UserDefinedPrimaryKey | PrimaryKeyKind::NonAppendOnlyRowIdPk => {
872 RequiredDist::hash_shard(pk_column_indices)
873 .streaming_enforce_if_not_satisfies(dml_node)?
874 }
875 PrimaryKeyKind::AppendOnlyRowIdPk => {
876 StreamExchange::new_no_shuffle(dml_node).into()
877 }
878 };
879
880 Ok(dml_node)
881 }
882
883 let kind = if let Some(row_id_index) = row_id_index {
884 assert_eq!(
885 Itertools::exactly_one(pk_column_indices.iter())
886 .copied()
887 .unwrap(),
888 row_id_index
889 );
890 if append_only {
891 PrimaryKeyKind::AppendOnlyRowIdPk
892 } else {
893 PrimaryKeyKind::NonAppendOnlyRowIdPk
894 }
895 } else {
896 PrimaryKeyKind::UserDefinedPrimaryKey
897 };
898
899 let column_descs: Vec<ColumnDesc> = columns
900 .iter()
901 .filter(|&c| c.can_dml())
902 .map(|c| c.column_desc.clone())
903 .collect();
904
905 let mut not_null_idxs = vec![];
906 for (idx, column) in column_descs.iter().enumerate() {
907 if !column.nullable {
908 not_null_idxs.push(idx);
909 }
910 }
911
912 let version_column_indices = if !with_version_columns.is_empty() {
913 find_version_column_indices(&columns, with_version_columns)?
914 } else {
915 vec![]
916 };
917
918 let with_external_source = source_catalog.is_some();
919 let (dml_source_node, external_source_node) = if with_external_source {
920 let dummy_source_node = LogicalSource::new(
921 None,
922 columns.clone(),
923 row_id_index,
924 SourceNodeKind::CreateTable,
925 context.clone(),
926 None,
927 )
928 .and_then(|s| {
929 s.to_stream(&mut ToStreamContext::new_with_backfill_type(
930 false,
931 BackfillType::ArrangementBackfill,
935 ))
936 })?;
937 let mut external_source_node = stream_plan.plan;
938 external_source_node =
939 inject_project_for_generated_column_if_needed(&columns, external_source_node)?;
940 external_source_node = match kind {
941 PrimaryKeyKind::UserDefinedPrimaryKey => {
942 RequiredDist::hash_shard(&pk_column_indices)
943 .streaming_enforce_if_not_satisfies(external_source_node)?
944 }
945
946 PrimaryKeyKind::NonAppendOnlyRowIdPk | PrimaryKeyKind::AppendOnlyRowIdPk => {
947 StreamExchange::new_no_shuffle(external_source_node).into()
948 }
949 };
950 (dummy_source_node, Some(external_source_node))
951 } else {
952 (stream_plan.plan, None)
953 };
954
955 let dml_node = inject_dml_node(
956 &columns,
957 append_only,
958 dml_source_node,
959 &pk_column_indices,
960 kind,
961 column_descs,
962 )?;
963
964 let dists = external_source_node
965 .iter()
966 .map(|input| input.distribution())
967 .chain([dml_node.distribution()])
968 .unique()
969 .collect_vec();
970
971 let dist = match &dists[..] {
972 &[Distribution::SomeShard, Distribution::HashShard(_)]
973 | &[Distribution::HashShard(_), Distribution::SomeShard] => Distribution::SomeShard,
974 &[dist @ Distribution::SomeShard] | &[dist @ Distribution::HashShard(_)] => {
975 dist.clone()
976 }
977 _ => {
978 unreachable!()
979 }
980 };
981
982 let generated_column_exprs =
983 LogicalSource::derive_output_exprs_from_generated_columns(&columns)?;
984 let upstream_sink_union = StreamUpstreamSinkUnion::new(
985 context.clone(),
986 dml_node.schema(),
987 dml_node.stream_key(),
988 dist.clone(), append_only,
990 row_id_index.is_none(),
991 generated_column_exprs,
992 );
993
994 let union_inputs = external_source_node
995 .into_iter()
996 .chain([dml_node, upstream_sink_union.into()])
997 .collect_vec();
998
999 let mut stream_plan: StreamPlanRef = StreamUnion::new_with_dist(
1000 Union {
1001 all: true,
1002 inputs: union_inputs,
1003 source_col: None,
1004 },
1005 dist,
1006 )
1007 .into();
1008
1009 let ttl_watermark_indices = watermark_descs
1010 .iter()
1011 .filter(|d| d.with_ttl)
1012 .map(|d| d.watermark_idx as usize)
1013 .collect_vec();
1014
1015 let add_row_id_gen = |stream_plan: StreamPlanRef, row_id_index| match kind {
1016 PrimaryKeyKind::UserDefinedPrimaryKey => {
1017 unreachable!()
1018 }
1019 PrimaryKeyKind::NonAppendOnlyRowIdPk | PrimaryKeyKind::AppendOnlyRowIdPk => {
1020 StreamRowIdGen::new_with_dist(
1021 stream_plan,
1022 row_id_index,
1023 Distribution::HashShard(vec![row_id_index]),
1024 )
1025 .into()
1026 }
1027 };
1028
1029 if let Some(row_id_index) = row_id_index {
1031 stream_plan = add_row_id_gen(stream_plan, row_id_index);
1032 }
1033
1034 if !watermark_descs.is_empty() {
1036 stream_plan = StreamWatermarkFilter::new(stream_plan, watermark_descs).into();
1037 }
1038
1039 let conflict_behavior = on_conflict.to_behavior(append_only, row_id_index.is_some())?;
1040
1041 if let ConflictBehavior::IgnoreConflict = conflict_behavior
1042 && !version_column_indices.is_empty()
1043 {
1044 Err(ErrorCode::InvalidParameterValue(
1045 "The with version column syntax cannot be used with the ignore behavior of on conflict".to_owned(),
1046 ))?
1047 }
1048
1049 let retention_seconds = context.with_options().retention_seconds();
1050
1051 let table_required_dist = {
1052 let mut bitset = FixedBitSet::with_capacity(columns.len());
1053 for idx in &pk_column_indices {
1054 bitset.insert(*idx);
1055 }
1056 RequiredDist::ShardByKey(bitset)
1057 };
1058
1059 let mut stream_plan = inline_session_timezone_in_exprs(context, stream_plan)?;
1060
1061 if !not_null_idxs.is_empty() {
1062 stream_plan =
1063 StreamFilter::filter_out_any_null_rows(stream_plan.clone(), ¬_null_idxs);
1064 }
1065
1066 let refreshable = source_catalog
1068 .as_ref()
1069 .map(|catalog| {
1070 catalog.with_properties.supports_full_reload_refresh()
1071 && matches!(
1072 catalog
1073 .refresh_mode
1074 .as_ref()
1075 .map(|refresh_mode| refresh_mode.refresh_mode),
1076 Some(Some(RefreshMode::FullReload(_)))
1077 )
1078 })
1079 .unwrap_or(false);
1080
1081 if refreshable && row_id_index.is_some() {
1083 return Err(crate::error::ErrorCode::BindError(
1084 "Refreshable tables must have a PRIMARY KEY. Please define a primary key for the table."
1085 .to_owned(),
1086 )
1087 .into());
1088 }
1089
1090 StreamMaterialize::create_for_table(
1091 stream_plan,
1092 table_name,
1093 database_id,
1094 schema_id,
1095 table_required_dist,
1096 Order::any(),
1097 columns,
1098 definition,
1099 conflict_behavior,
1100 version_column_indices,
1101 pk_column_indices,
1102 ttl_watermark_indices,
1103 row_id_index,
1104 version,
1105 retention_seconds,
1106 webhook_info,
1107 engine,
1108 refreshable,
1109 )
1110 }
1111
1112 pub fn gen_materialize_plan(
1114 self,
1115 database_id: DatabaseId,
1116 schema_id: SchemaId,
1117 mv_name: String,
1118 definition: String,
1119 emit_on_window_close: bool,
1120 backfill_type: BackfillType,
1121 ) -> Result<StreamMaterialize> {
1122 let cardinality = self.compute_cardinality();
1123 let stream_plan = self.gen_optimized_stream_plan(emit_on_window_close, backfill_type)?;
1124 StreamMaterialize::create(
1125 stream_plan,
1126 mv_name,
1127 database_id,
1128 schema_id,
1129 definition,
1130 TableType::MaterializedView,
1131 cardinality,
1132 None,
1133 )
1134 }
1135
1136 pub fn gen_index_plan(
1138 self,
1139 index_name: String,
1140 database_id: DatabaseId,
1141 schema_id: SchemaId,
1142 definition: String,
1143 retention_seconds: Option<NonZeroU32>,
1144 ) -> Result<StreamMaterialize> {
1145 let cardinality = self.compute_cardinality();
1146 let backfill_type = self.derive_backfill_type(false);
1147 let stream_plan = self.gen_optimized_stream_plan(false, backfill_type)?;
1148
1149 StreamMaterialize::create(
1150 stream_plan,
1151 index_name,
1152 database_id,
1153 schema_id,
1154 definition,
1155 TableType::Index,
1156 cardinality,
1157 retention_seconds,
1158 )
1159 }
1160
1161 pub fn gen_vector_index_plan(
1162 self,
1163 index_name: String,
1164 database_id: DatabaseId,
1165 schema_id: SchemaId,
1166 definition: String,
1167 retention_seconds: Option<NonZeroU32>,
1168 vector_index_info: PbVectorIndexInfo,
1169 ) -> Result<StreamVectorIndexWrite> {
1170 let cardinality = self.compute_cardinality();
1171 let backfill_type = self.derive_backfill_type(false);
1172 let stream_plan = self.gen_optimized_stream_plan(false, backfill_type)?;
1173
1174 StreamVectorIndexWrite::create(
1175 stream_plan,
1176 index_name,
1177 database_id,
1178 schema_id,
1179 definition,
1180 cardinality,
1181 retention_seconds,
1182 vector_index_info,
1183 )
1184 }
1185
1186 #[expect(clippy::too_many_arguments)]
1188 pub fn gen_sink_plan(
1189 self,
1190 sink_name: String,
1191 definition: String,
1192 properties: WithOptionsSecResolved,
1193 emit_on_window_close: bool,
1194 db_name: String,
1195 sink_from_table_name: String,
1196 format_desc: Option<SinkFormatDesc>,
1197 without_snapshot: bool,
1198 since_timestamp: bool,
1199 is_iceberg_engine_internal: bool,
1200 target_table: Option<Arc<TableCatalog>>,
1201 partition_info: Option<PartitionComputeInfo>,
1202 user_specified_columns: bool,
1203 auto_refresh_schema_from_table: Option<Arc<TableCatalog>>,
1204 ) -> Result<StreamSink> {
1205 let backfill_type = if since_timestamp {
1206 assert!(
1207 target_table.is_none(),
1208 "should not allow since_timestamp for sink-into-table"
1209 );
1210 if is_iceberg_engine_internal {
1211 return Err(ErrorCode::InvalidInputSyntax(
1212 "since_timestamp is not allowed for this sink".to_owned(),
1213 )
1214 .into());
1215 }
1216 BackfillType::SnapshotBackfillSinceTimestamp
1217 } else if without_snapshot {
1218 BackfillType::UpstreamOnlySink
1219 } else if target_table.is_none()
1220 && !is_iceberg_engine_internal
1221 && self.should_use_snapshot_backfill()
1222 && {
1223 if auto_refresh_schema_from_table.is_some() {
1224 self.plan.ctx().session_ctx().notice_to_user("Auto schema change only support for ArrangementBackfill. Switched to use ArrangementBackfill");
1225 false
1226 } else {
1227 true
1228 }
1229 }
1230 {
1231 assert!(
1232 target_table.is_none(),
1233 "should not allow snapshot backfill for sink-into-table"
1234 );
1235 BackfillType::SnapshotBackfill
1237 } else {
1238 BackfillType::ArrangementBackfill
1239 };
1240 if auto_refresh_schema_from_table.is_some()
1241 && backfill_type != BackfillType::ArrangementBackfill
1242 {
1243 return Err(ErrorCode::InvalidInputSyntax(format!(
1244 "auto schema change only support for ArrangementBackfill, but got: {:?}",
1245 backfill_type
1246 ))
1247 .into());
1248 }
1249 let stream_plan = self.gen_optimized_stream_plan(emit_on_window_close, backfill_type)?;
1250 let target_columns_to_plan_mapping = target_table.as_ref().map(|t| {
1251 let columns = t.columns_without_rw_timestamp();
1252 stream_plan.target_columns_to_plan_mapping(&columns, user_specified_columns)
1253 });
1254
1255 StreamSink::create(
1256 stream_plan,
1257 sink_name,
1258 db_name,
1259 sink_from_table_name,
1260 target_table,
1261 target_columns_to_plan_mapping,
1262 definition,
1263 properties,
1264 format_desc,
1265 partition_info,
1266 auto_refresh_schema_from_table,
1267 )
1268 }
1269
1270 pub fn should_use_snapshot_backfill(&self) -> bool {
1271 let ctx = self.plan.ctx();
1272 let session_ctx = ctx.session_ctx();
1273 let use_snapshot_backfill = session_ctx
1274 .env()
1275 .streaming_config()
1276 .developer
1277 .enable_snapshot_backfill
1278 && session_ctx.config().streaming_use_snapshot_backfill();
1279 if use_snapshot_backfill {
1280 if let Some(warning_msg) = self.plan.forbid_snapshot_backfill() {
1281 self.plan.ctx().session_ctx().notice_to_user(warning_msg);
1282 false
1283 } else {
1284 true
1285 }
1286 } else {
1287 false
1288 }
1289 }
1290}
1291
1292impl<P: PlanPhase> PlanRoot<P> {
1293 pub fn target_columns_to_plan_mapping(
1295 &self,
1296 tar_cols: &[ColumnCatalog],
1297 user_specified_columns: bool,
1298 ) -> Vec<Option<usize>> {
1299 #[expect(clippy::disallowed_methods)]
1300 let visible_cols: Vec<(usize, String)> = self
1301 .out_fields
1302 .ones()
1303 .zip_eq(self.out_names.iter().cloned())
1304 .collect_vec();
1305
1306 let visible_col_idxes = visible_cols.iter().map(|(i, _)| *i).collect_vec();
1307 let visible_col_idxes_by_name = visible_cols
1308 .iter()
1309 .map(|(i, name)| (name.as_ref(), *i))
1310 .collect::<BTreeMap<_, _>>();
1311
1312 tar_cols
1313 .iter()
1314 .enumerate()
1315 .filter(|(_, tar_col)| tar_col.can_dml())
1316 .map(|(tar_i, tar_col)| {
1317 if user_specified_columns {
1318 visible_col_idxes_by_name.get(tar_col.name()).cloned()
1319 } else {
1320 (tar_i < visible_col_idxes.len()).then(|| visible_cols[tar_i].0)
1321 }
1322 })
1323 .collect()
1324 }
1325}
1326
1327fn find_version_column_indices(
1328 column_catalog: &Vec<ColumnCatalog>,
1329 version_column_names: Vec<String>,
1330) -> Result<Vec<usize>> {
1331 let mut indices = Vec::new();
1332 for version_column_name in version_column_names {
1333 let mut found = false;
1334 for (index, column) in column_catalog.iter().enumerate() {
1335 if column.column_desc.name == version_column_name {
1336 if let &DataType::Jsonb
1337 | &DataType::Variant
1338 | &DataType::List(_)
1339 | &DataType::Struct(_)
1340 | &DataType::Bytea
1341 | &DataType::Boolean = column.data_type()
1342 {
1343 return Err(ErrorCode::InvalidInputSyntax(format!(
1344 "Version column {} must be of a comparable data type",
1345 version_column_name
1346 ))
1347 .into());
1348 }
1349 indices.push(index);
1350 found = true;
1351 break;
1352 }
1353 }
1354 if !found {
1355 return Err(ErrorCode::InvalidInputSyntax(format!(
1356 "Version column {} not found",
1357 version_column_name
1358 ))
1359 .into());
1360 }
1361 }
1362 Ok(indices)
1363}
1364
1365fn const_eval_exprs<C: ConventionMarker>(plan: PlanRef<C>) -> Result<PlanRef<C>> {
1366 let mut const_eval_rewriter = ConstEvalRewriter { error: None };
1367
1368 let plan = plan.rewrite_exprs_recursive(&mut const_eval_rewriter);
1369 if let Some(error) = const_eval_rewriter.error {
1370 return Err(error);
1371 }
1372 Ok(plan)
1373}
1374
1375fn inline_session_timezone_in_exprs<C: ConventionMarker>(
1376 ctx: OptimizerContextRef,
1377 plan: PlanRef<C>,
1378) -> Result<PlanRef<C>> {
1379 let mut v = TimestamptzExprFinder::default();
1380 plan.visit_exprs_recursive(&mut v);
1381 if v.has() {
1382 Ok(plan.rewrite_exprs_recursive(ctx.session_timezone().deref_mut()))
1383 } else {
1384 Ok(plan)
1385 }
1386}
1387
1388fn exist_and_no_exchange_before(
1389 plan: &BatchPlanRef,
1390 is_candidate: fn(&BatchPlanRef) -> bool,
1391) -> bool {
1392 if plan.node_type() == BatchPlanNodeType::BatchExchange {
1393 return false;
1394 }
1395 is_candidate(plan)
1396 || plan
1397 .inputs()
1398 .iter()
1399 .any(|input| exist_and_no_exchange_before(input, is_candidate))
1400}
1401
1402impl BatchPlanRef {
1403 fn is_user_table_scan(&self) -> bool {
1404 self.node_type() == BatchPlanNodeType::BatchSeqScan
1405 || self.node_type() == BatchPlanNodeType::BatchLogSeqScan
1406 || self.node_type() == BatchPlanNodeType::BatchVectorSearch
1407 }
1408
1409 fn is_source_scan(&self) -> bool {
1410 self.node_type() == BatchPlanNodeType::BatchSource
1411 || self.node_type() == BatchPlanNodeType::BatchKafkaScan
1412 || self.node_type() == BatchPlanNodeType::BatchIcebergScan
1413 }
1414
1415 fn is_lookup_join(&self) -> bool {
1416 self.node_type() == BatchPlanNodeType::BatchLookupJoin
1417 }
1418
1419 fn is_insert(&self) -> bool {
1420 self.node_type() == BatchPlanNodeType::BatchInsert
1421 }
1422
1423 fn is_update(&self) -> bool {
1424 self.node_type() == BatchPlanNodeType::BatchUpdate
1425 }
1426
1427 fn is_delete(&self) -> bool {
1428 self.node_type() == BatchPlanNodeType::BatchDelete
1429 }
1430}
1431
1432fn require_additional_exchange_on_root_in_distributed_mode(plan: BatchPlanRef) -> bool {
1438 assert_eq!(plan.distribution(), &Distribution::Single);
1439 exist_and_no_exchange_before(&plan, |plan| {
1440 plan.is_user_table_scan()
1441 || plan.is_source_scan()
1442 || plan.is_insert()
1443 || plan.is_update()
1444 || plan.is_delete()
1445 || plan.is_lookup_join()
1449 })
1450}
1451
1452fn require_additional_exchange_on_root_in_local_mode(plan: BatchPlanRef) -> bool {
1455 assert_eq!(plan.distribution(), &Distribution::Single);
1456 exist_and_no_exchange_before(&plan, |plan| {
1457 plan.is_user_table_scan() || plan.is_source_scan() || plan.is_insert()
1458 })
1459}
1460
1461#[cfg(test)]
1462mod tests {
1463 use super::*;
1464 use crate::optimizer::plan_node::LogicalValues;
1465
1466 #[tokio::test]
1467 async fn test_as_subplan() {
1468 let ctx = OptimizerContext::mock();
1469 let values = LogicalValues::new(
1470 vec![],
1471 Schema::new(vec![
1472 Field::with_name(DataType::Int32, "v1"),
1473 Field::with_name(DataType::Varchar, "v2"),
1474 ]),
1475 ctx,
1476 )
1477 .into();
1478 let out_fields = FixedBitSet::with_capacity_and_blocks(2, [1]);
1479 let out_names = vec!["v1".into()];
1480 let root = PlanRoot::new_with_logical_plan(
1481 values,
1482 RequiredDist::Any,
1483 Order::any(),
1484 out_fields,
1485 out_names,
1486 );
1487 let subplan = root.into_unordered_subplan();
1488 assert_eq!(
1489 subplan.schema(),
1490 &Schema::new(vec![Field::with_name(DataType::Int32, "v1")])
1491 );
1492 }
1493}