refactor(optimizer): change stream key to option

src/frontend/src/optimizer/plan_node/derive.rs

@@ -22,6 +22,7 @@ use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
 
 use super::PlanRef;
 use crate::optimizer::property::Order;
+use crate::Explain;
 
 pub(crate) fn derive_columns(
     input_schema: &Schema,
@@ -82,7 +83,18 @@ pub(crate) fn derive_pk(
     columns: &[ColumnCatalog],
 ) -> (Vec<ColumnOrder>, Vec<usize>) {
     // Note(congyi): avoid pk duplication
-    let stream_key = input.stream_key().iter().copied().unique().collect_vec();
+    let stream_key = input
+        .stream_key()
+        .unwrap_or_else(|| {
+            panic!(
+                "should always have a stream key on the top of the stream plan but not, plan: {}",
+                input.explain_to_string()
+            )
+        })
+        .iter()
+        .copied()
+        .unique()
+        .collect_vec();
     let schema = input.schema();
 
     // Assert the uniqueness of column names and IDs, including hidden columns.

src/frontend/src/optimizer/plan_node/generic/agg.rs

@@ -344,7 +344,7 @@ impl<PlanRef: stream::StreamPlanRef> Agg<PlanRef> {
         window_col_idx: Option<usize>,
     ) -> Vec<AggCallState> {
         let in_fields = self.input.schema().fields().to_vec();
-        let in_pks = self.input.stream_key().to_vec();
+        let in_pks = self.input.stream_key().unwrap().to_vec();
         let in_append_only = self.input.append_only();
         let in_dist_key = self.input.distribution().dist_column_indices().to_vec();
 

src/frontend/src/optimizer/plan_node/generic/delete.rs

@@ -62,7 +62,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Delete<PlanRef> {
 
     fn stream_key(&self) -> Option<Vec<usize>> {
         if self.returning {
-            Some(self.input.stream_key().to_vec())
+            Some(self.input.stream_key()?.to_vec())
         } else {
             Some(vec![])
         }

src/frontend/src/optimizer/plan_node/generic/dynamic_filter.rs

@@ -59,7 +59,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for DynamicFilter<PlanRef> {
     }
 
     fn stream_key(&self) -> Option<Vec<usize>> {
-        Some(self.left.stream_key().to_vec())
+        Some(self.left.stream_key()?.to_vec())
     }
 
     fn ctx(&self) -> OptimizerContextRef {
@@ -151,7 +151,7 @@ pub fn infer_left_internal_table_catalog(
     let mut pk_indices = vec![left_key_index];
     let read_prefix_len_hint = pk_indices.len();
 
-    for i in me.stream_key() {
+    for i in me.stream_key().unwrap() {
         if *i != left_key_index {
             pk_indices.push(*i);
         }

src/frontend/src/optimizer/plan_node/generic/except.rs

@@ -34,7 +34,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Except<PlanRef> {
     }
 
     fn stream_key(&self) -> Option<Vec<usize>> {
-        Some(self.inputs[0].stream_key().to_vec())
+        Some(self.inputs[0].stream_key()?.to_vec())
     }
 
     fn ctx(&self) -> OptimizerContextRef {

src/frontend/src/optimizer/plan_node/generic/expand.rs

@@ -61,7 +61,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Expand<PlanRef> {
         let input_schema_len = self.input.schema().len();
         let mut pk_indices = self
             .input
-            .stream_key()
+            .stream_key()?
             .iter()
             .map(|&pk| pk + input_schema_len)
             .collect_vec();

src/frontend/src/optimizer/plan_node/generic/filter.rs

@@ -54,7 +54,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Filter<PlanRef> {
     }
 
     fn stream_key(&self) -> Option<Vec<usize>> {
-        Some(self.input.stream_key().to_vec())
+        Some(self.input.stream_key()?.to_vec())
     }
 
     fn ctx(&self) -> OptimizerContextRef {

src/frontend/src/optimizer/plan_node/generic/hop_window.rs

@@ -77,7 +77,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for HopWindow<PlanRef> {
         } else {
             let mut pk = self
                 .input
-                .stream_key()
+                .stream_key()?
                 .iter()
                 .filter_map(|&pk_idx| self.output_indices.iter().position(|&idx| idx == pk_idx))
                 .collect_vec();

src/frontend/src/optimizer/plan_node/generic/intersect.rs

@@ -33,7 +33,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Intersect<PlanRef> {
     }
 
     fn stream_key(&self) -> Option<Vec<usize>> {
-        Some(self.inputs[0].stream_key().to_vec())
+        Some(self.inputs[0].stream_key()?.to_vec())
     }
 
     fn ctx(&self) -> OptimizerContextRef {

src/frontend/src/optimizer/plan_node/generic/join.rs

@@ -93,8 +93,8 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Join<PlanRef> {
     fn stream_key(&self) -> Option<Vec<usize>> {
         let _left_len = self.left.schema().len();
         let _right_len = self.right.schema().len();
-        let left_pk = self.left.stream_key();
-        let right_pk = self.right.stream_key();
+        let left_pk = self.left.stream_key()?;
+        let right_pk = self.right.stream_key()?;
         let l2i = self.l2i_col_mapping();
         let r2i = self.r2i_col_mapping();
         let full_out_col_num = self.internal_column_num();
@@ -110,7 +110,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Join<PlanRef> {
 
         // NOTE(st1page): add join keys in the pk_indices a work around before we really have stream
         // key.
-        pk_indices.and_then(|mut pk_indices| {
+        pk_indices.and_then(|mut pk_indices: Vec<usize>| {
             let left_len = self.left.schema().len();
             let right_len = self.right.schema().len();
             let eq_predicate = EqJoinPredicate::create(left_len, right_len, self.on.clone());

src/frontend/src/optimizer/plan_node/generic/limit.rs

@@ -43,7 +43,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Limit<PlanRef> {
     }
 
     fn stream_key(&self) -> Option<Vec<usize>> {
-        Some(self.input.stream_key().to_vec())
+        Some(self.input.stream_key()?.to_vec())
     }
 }
 impl<PlanRef> Limit<PlanRef> {

src/frontend/src/optimizer/plan_node/generic/mod.rs

@@ -86,7 +86,7 @@ pub(super) use impl_distill_unit_from_fields;
 
 pub trait GenericPlanRef: Eq + Hash {
     fn schema(&self) -> &Schema;
-    fn stream_key(&self) -> &[usize];
+    fn stream_key(&self) -> Option<&[usize]>;
     fn functional_dependency(&self) -> &FunctionalDependencySet;
     fn ctx(&self) -> OptimizerContextRef;
 }

src/frontend/src/optimizer/plan_node/generic/over_window.rs

@@ -219,7 +219,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for OverWindow<PlanRef> {
     }
 
     fn stream_key(&self) -> Option<Vec<usize>> {
-        let mut output_pk = self.input.stream_key().to_vec();
+        let mut output_pk = self.input.stream_key()?.to_vec();
         for part_key_idx in self
             .window_functions
             .iter()

src/frontend/src/optimizer/plan_node/generic/project.rs

@@ -100,7 +100,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Project<PlanRef> {
     fn stream_key(&self) -> Option<Vec<usize>> {
         let i2o = self.i2o_col_mapping();
         self.input
-            .stream_key()
+            .stream_key()?
             .iter()
             .map(|pk_col| i2o.try_map(*pk_col))
             .collect::<Option<Vec<_>>>()

src/frontend/src/optimizer/plan_node/generic/project_set.rs

@@ -92,7 +92,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for ProjectSet<PlanRef> {
         let i2o = self.i2o_col_mapping();
         let mut pk = self
             .input
-            .stream_key()
+            .stream_key()?
             .iter()
             .map(|pk_col| i2o.try_map(*pk_col))
             .collect::<Option<Vec<_>>>()

src/frontend/src/optimizer/plan_node/generic/share.rs

@@ -44,7 +44,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Share<PlanRef> {
     }
 
     fn stream_key(&self) -> Option<Vec<usize>> {
-        Some(self.input.borrow().stream_key().to_vec())
+        Some(self.input.borrow().stream_key()?.to_vec())
     }
 
     fn ctx(&self) -> OptimizerContextRef {

src/frontend/src/optimizer/plan_node/generic/top_n.rs

@@ -42,7 +42,7 @@ impl<PlanRef: stream::StreamPlanRef> TopN<PlanRef> {
         &self,
         schema: &Schema,
         ctx: OptimizerContextRef,
-        stream_key: &[usize],
+        input_stream_key: &[usize],
         vnode_col_idx: Option<usize>,
     ) -> TableCatalog {
         let columns_fields = schema.fields().to_vec();
@@ -71,7 +71,7 @@ impl<PlanRef: stream::StreamPlanRef> TopN<PlanRef> {
             order_cols.insert(order.column_index);
         });
 
-        stream_key.iter().for_each(|idx| {
+        input_stream_key.iter().for_each(|idx| {
             if !order_cols.contains(idx) {
                 internal_table_catalog_builder.add_order_column(*idx, OrderType::ascending());
                 order_cols.insert(*idx);
@@ -176,7 +176,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for TopN<PlanRef> {
         if self.limit_attr.max_one_row() {
             Some(self.group_key.clone())
         } else {
-            let mut pk = self.input.stream_key().to_vec();
+            let mut pk = self.input.stream_key()?.to_vec();
             for i in &self.group_key {
                 if !pk.contains(i) {
                     pk.push(*i);

src/frontend/src/optimizer/plan_node/generic/union.rs

@@ -40,7 +40,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Union<PlanRef> {
         // Union all its inputs pks + source_col if exists
         let mut pk_indices = vec![];
         for input in &self.inputs {
-            for pk in input.stream_key() {
+            for pk in input.stream_key()? {
                 if !pk_indices.contains(pk) {
                     pk_indices.push(*pk);
                 }

src/frontend/src/optimizer/plan_node/generic/update.rs

@@ -64,7 +64,7 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Update<PlanRef> {
 
     fn stream_key(&self) -> Option<Vec<usize>> {
         if self.returning {
-            Some(self.input.stream_key().to_vec())
+            Some(self.input.stream_key()?.to_vec())
         } else {
             Some(vec![])
         }

src/frontend/src/optimizer/plan_node/logical_agg.rs

@@ -39,6 +39,7 @@ use crate::optimizer::property::{Distribution, Order, RequiredDist};
 use crate::utils::{
     ColIndexMapping, ColIndexMappingRewriteExt, Condition, GroupBy, IndexSet, Substitute,
 };
+use crate::Explain;
 
 /// `LogicalAgg` groups input data by their group key and computes aggregation functions.
 ///
@@ -222,8 +223,16 @@ impl LogicalAgg {
         // so it obeys consistent hash strategy via [`Distribution::HashShard`].
         let stream_input =
             if *input_dist == Distribution::SomeShard && self.core.must_try_two_phase_agg() {
-                RequiredDist::shard_by_key(stream_input.schema().len(), stream_input.stream_key())
-                    .enforce_if_not_satisfies(stream_input, &Order::any())?
+                RequiredDist::shard_by_key(
+                    stream_input.schema().len(),
+                    stream_input.stream_key().unwrap_or_else(|| {
+                        panic!(
+                        "should always have a stream key in the stream plan but not, sub plan: {}",
+                        stream_input.explain_to_string()
+                    )
+                    }),
+                )
+                .enforce_if_not_satisfies(stream_input, &Order::any())?
             } else {
                 stream_input
             };

src/frontend/src/optimizer/plan_node/logical_apply.rs

@@ -86,16 +86,13 @@ impl LogicalApply {
         let ctx = left.ctx();
         let join_core = generic::Join::with_full_output(left, right, join_type, on);
         let schema = join_core.schema();
-        let pk_indices = join_core.stream_key();
-        let (functional_dependency, pk_indices) = match pk_indices {
-            Some(pk_indices) => (
-                FunctionalDependencySet::with_key(schema.len(), &pk_indices),
-                pk_indices,
-            ),
-            None => (FunctionalDependencySet::new(schema.len()), vec![]),
+        let stream_key = join_core.stream_key();
+        let functional_dependency = match &stream_key {
+            Some(stream_key) => FunctionalDependencySet::with_key(schema.len(), stream_key),
+            None => FunctionalDependencySet::new(schema.len()),
         };
         let (left, right, on, join_type, _output_indices) = join_core.decompose();
-        let base = PlanBase::new_logical(ctx, schema, pk_indices, functional_dependency);
+        let base = PlanBase::new_logical(ctx, schema, stream_key, functional_dependency);
         LogicalApply {
             base,
             left,

src/frontend/src/optimizer/plan_node/logical_dedup.rs

@@ -38,13 +38,8 @@ pub struct LogicalDedup {
 
 impl LogicalDedup {
     pub fn new(input: PlanRef, dedup_cols: Vec<usize>) -> Self {
-        let base = PlanBase::new_logical(
-            input.ctx(),
-            input.schema().clone(),
-            dedup_cols.clone(),
-            input.functional_dependency().clone(),
-        );
         let core = generic::Dedup { input, dedup_cols };
+        let base = PlanBase::new_logical_with_core(&core);
         LogicalDedup { base, core }
     }
 

src/frontend/src/optimizer/plan_node/logical_hop_window.rs

@@ -28,6 +28,7 @@ use crate::optimizer::plan_node::{
     ColumnPruningContext, PredicatePushdownContext, RewriteStreamContext, ToStreamContext,
 };
 use crate::utils::{ColIndexMapping, ColIndexMappingRewriteExt, Condition};
+use crate::Explain;
 
 /// `LogicalHopWindow` implements Hop Table Function.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -63,23 +64,12 @@ impl LogicalHopWindow {
             output_indices,
         };
 
-        let _schema = core.schema();
-        let _pk_indices = core.stream_key();
         let ctx = core.ctx();
 
-        // NOTE(st1page): add join keys in the pk_indices a work around before we really have stream
-        // key.
-        // let pk_indices = match pk_indices {
-        //     Some(pk_indices) if functional_dependency.is_key(&pk_indices) => {
-        //         functional_dependency.minimize_key(&pk_indices)
-        //     }
-        //     _ => pk_indices.unwrap_or_default(),
-        // };
-
         let base = PlanBase::new_logical(
             ctx,
             core.schema(),
-            core.stream_key().unwrap_or_default(),
+            core.stream_key(),
             core.functional_dependency(),
         );
 
@@ -345,6 +335,12 @@ impl ToStream for LogicalHopWindow {
         output_indices.extend(
             input
                 .stream_key()
+                .unwrap_or_else(|| {
+                    panic!(
+                        "should always have a stream key in the stream plan but not, sub plan: {}",
+                        input.explain_to_string()
+                    )
+                })
                 .iter()
                 .cloned()
                 .filter(|i| i2o.try_map(*i).is_none()),

src/frontend/src/optimizer/plan_node/logical_join.rs

@@ -41,6 +41,7 @@ use crate::optimizer::plan_node::{
 use crate::optimizer::plan_visitor::LogicalCardinalityExt;
 use crate::optimizer::property::{Distribution, Order, RequiredDist};
 use crate::utils::{ColIndexMapping, ColIndexMappingRewriteExt, Condition, ConditionDisplay};
+use crate::Explain;
 
 /// `LogicalJoin` combines two relations according to some condition.
 ///
@@ -1396,13 +1397,25 @@ impl ToStream for LogicalJoin {
         // Add missing pk indices to the logical join
         let mut left_to_add = left
             .stream_key()
+            .unwrap_or_else(|| {
+                panic!(
+                    "should always have a stream key in the stream plan but not, sub plan: {}",
+                    left.explain_to_string()
+                )
+            })
             .iter()
             .cloned()
             .filter(|i| l2o.try_map(*i).is_none())
             .collect_vec();
 
         let mut right_to_add = right
             .stream_key()
+            .unwrap_or_else(|| {
+                panic!(
+                    "should always have a stream key in the stream plan but not, sub plan: {}",
+                    right.explain_to_string()
+                )
+            })
             .iter()
             .filter(|&&i| r2o.try_map(i).is_none())
             .map(|&i| i + left_len)
@@ -1465,12 +1478,15 @@ impl ToStream for LogicalJoin {
             let left_right_stream_keys = join_with_pk
                 .left()
                 .stream_key()
+                .unwrap_or_else(|| panic!("should always have a stream key in the stream plan but not, sub plan: {}",
+                    left.explain_to_string()))
                 .iter()
                 .map(|i| l2o.map(*i))
                 .chain(
                     join_with_pk
                         .right()
-                        .stream_key()
+                        .stream_key()                .unwrap_or_else(|| panic!("should always have a stream key in the stream plan but not, sub plan: {}",
+                            join_with_pk.right().explain_to_string()))
                         .iter()
                         .map(|i| r2o.map(*i)),
                 )