Skip to main content

risingwave_frontend/optimizer/
mod.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::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/// `PlanRoot` is used to describe a plan. planner will construct a `PlanRoot` with `LogicalNode`.
103/// and required distribution and order. And `PlanRoot` can generate corresponding streaming or
104/// batch plan with optimization. the required Order and Distribution columns might be more than the
105/// output columns. for example:
106/// ```sql
107///    select v1 from t order by id;
108/// ```
109/// the plan will return two columns (id, v1), and the required order column is id. the id
110/// column is required in optimization, but the final generated plan will remove the unnecessary
111/// column in the result.
112#[derive(Educe)]
113#[educe(Debug, Clone)]
114pub struct PlanRoot<P: PlanPhase> {
115    // The current plan node.
116    pub plan: PlanRef<P::Convention>,
117    // The phase of the plan.
118    #[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
126/// `PlanPhase` is used to track the phase of the `PlanRoot`.
127/// Usually, it begins from `Logical` and ends with `Batch` or `Stream`, unless we want to construct a `PlanRoot` from an intermediate phase.
128/// Typical phase transformation are:
129/// - `Logical` -> `OptimizedLogicalForBatch` -> `Batch`
130/// - `Logical` -> `OptimizedLogicalForStream` -> `Stream`
131pub 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    /// Set customized names of the output fields, used for `CREATE [MATERIALIZED VIEW | SINK] r(a,
218    /// b, ..)`.
219    ///
220    /// If the number of names does not match the number of output fields, an error is returned.
221    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    /// Get the plan root's schema, only including the fields to be output.
232    pub fn schema(&self) -> Schema {
233        // The schema can be derived from the `out_fields` and `out_names`, so we don't maintain it
234        // as a field and always construct one on demand here to keep it in sync.
235        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    /// Transform the [`PlanRoot`] back to a [`PlanRef`] suitable to be used as a subplan, for
252    /// example as insert source or subquery. This ignores Order but retains post-Order pruning
253    /// (`out_fields`).
254    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    /// Transform the [`PlanRoot`] wrapped in an array-construction subquery to a [`PlanRef`]
262    /// supported by `ARRAY_AGG`. Similar to the unordered version, this abstracts away internal
263    /// `self.plan` which is further modified by `self.required_order` then `self.out_fields`.
264    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    /// Apply logical optimization to the plan for stream.
311    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    /// Apply logical optimization to the plan for batch.
317    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    /// Rejects `VARIANT` wherever a batch plan would group, deduplicate or order by it. See
330    /// [`crate::optimizer::variant_key`] for why this is batch's last line of defense.
331    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    /// Optimize and generate a singleton batch physical plan without exchange nodes.
349    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        // Inline session timezone mainly for rewriting now()
361        let mut plan = inline_session_timezone_in_exprs(ctx.clone(), self.plan.clone())?;
362
363        // Const eval of exprs at the last minute, but before `to_batch` to make functional index selection happy.
364        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        // Convert to physical plan node
372        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        // Inline session timezone
385        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        // Inline session timezone mainly for rewriting now()
432        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    /// Optimize and generate a batch query plan for distributed execution.
476    pub fn gen_batch_distributed_plan(mut self) -> Result<BatchPlanRef> {
477        self.required_dist = RequiredDist::single();
478        let mut plan = self.plan;
479
480        // Convert to distributed plan
481        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        // Add Project if the any position of `self.out_fields` is set to zero.
494        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        // Both two phase limit and topn could generate limit on top of the scan, so we push limit here.
500        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    /// Optimize and generate a batch query plan for local execution.
510    pub fn gen_batch_local_plan(self) -> Result<BatchPlanRef> {
511        let mut plan = self.plan;
512
513        // Convert to local plan node
514        plan = plan.to_local_with_order_required(&self.required_order)?;
515
516        // We remark that since the `to_local_with_order_required` does not enforce single
517        // distribution, we enforce at the root if needed.
518        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        // Add Project if the any position of `self.out_fields` is set to zero.
528        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        // Both two phase limit and topn could generate limit on top of the scan, so we push limit here.
540        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    /// Generate optimized stream plan
552    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        // Add Logstore for Unaligned join
592        // Apply this BEFORE delta join rule, because delta join removes
593        // the join
594        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            // TODO: make it a logical optimization.
606            // Rewrite joins with index to delta join
607            plan = plan.optimize_by_rules(&OptimizationStage::new(
608                "To IndexDeltaJoin",
609                vec![IndexDeltaJoinRule::create()],
610                ApplyOrder::BottomUp,
611            ))?;
612        }
613        // Inline session timezone
614        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        // Const eval of exprs at the last minute
622        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            // LocalityProviderCounter is non-zero only when locality backfill is enabled.
648            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            // missed_locality_providers can only be non-zero when locality backfill is disabled.
658            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    /// Generate create index or create materialize view plan.
701    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                        // TODO(st1page): https://github.com/risingwavelabs/risingwave/issues/7234
739                        // assert_eq!(target_size, output_indices.len());
740                        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            // TODO: can be `plan.plan.explain_to_string()`, but should explicitly specify the type due to some limitation of rust compiler
784            ctx.trace(<PlanRef<Stream> as Explain>::explain_to_string(&plan.plan));
785        }
786        Ok(plan)
787    }
788
789    /// Visit the plan root and compute the cardinality.
790    ///
791    /// Panics if not called on a logical plan.
792    fn compute_cardinality(&self) -> Cardinality {
793        CardinalityVisitor.visit(self.plan.clone())
794    }
795
796    /// Optimize and generate a create table plan.
797    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        // Snapshot backfill is not allowed for create table
822        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()) // pk column id must exist in table columns.
835                .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            // Add generated columns.
868            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                    // Dummy DML source planning does not create stream table scans, so this
932                    // required context value is only a placeholder and is not used for backfill
933                    // selection.
934                    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(), // should always be the same as dist of `Union`
989            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        // Add RowIDGen before WatermarkFilter, so filtering always sees a valid row-id key.
1030        if let Some(row_id_index) = row_id_index {
1031            stream_plan = add_row_id_gen(stream_plan, row_id_index);
1032        }
1033
1034        // Add WatermarkFilter node.
1035        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(), &not_null_idxs);
1064        }
1065
1066        // Determine if the table should be refreshable based on the connector type
1067        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        // Validate that refreshable tables have a user-defined primary key (i.e., does not have rowid)
1082        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    /// Optimize and generate a create materialized view plan.
1113    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    /// Optimize and generate a create index plan.
1137    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    /// Optimize and generate a create sink plan.
1187    #[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            // Snapshot backfill on sink-into-table is not allowed
1236            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    /// used when the plan has a target relation such as DML and sink into table, return the mapping from table's columns to the plan's schema
1294    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
1432/// As we always run the root stage locally, for some plan in root stage which need to execute in
1433/// compute node we insert an additional exhchange before it to avoid to include it in the root
1434/// stage.
1435///
1436/// Returns `true` if we must insert an additional exchange to ensure this.
1437fn 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            // A lookup join on a singleton lookup table may be already in `Single`
1446            // distribution and reads from the state store, so it must not be executed
1447            // in the root stage on the frontend.
1448            || plan.is_lookup_join()
1449    })
1450}
1451
1452/// The purpose is same as `require_additional_exchange_on_root_in_distributed_mode`. We separate
1453/// them for the different requirement of plan node in different execute mode.
1454fn 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}