fix(optimizer): logical scan pk derive (risingwavelabs#1340)

* fix LogicalScan::logical_rewrite_for_stream

* add ut

* fix test

Signed-off-by: Alex Chi <iskyzh@gmail.com>

Co-authored-by: Alex Chi <iskyzh@gmail.com>

rust/frontend/src/optimizer/plan_node/logical_scan.rs

@@ -170,13 +170,16 @@ impl ToStream for LogicalScan {
                 for idx in &self.required_col_idx {
                     col_ids.insert(self.table_desc.columns[*idx].column_id);
                 }
+                let mut col_id_to_tb_idx = HashMap::new();
+                for (tb_idx, c) in self.table_desc().columns.iter().enumerate() {
+                    col_id_to_tb_idx.insert(c.column_id, tb_idx);
+                }
                 let col_need_to_add = self
                     .table_desc
                     .pk
                     .iter()
-                    .enumerate()
-                    .filter(|(_idx, c)| !col_ids.contains(&c.column_desc.column_id))
-                    .map(|(idx, _c)| idx)
+                    .filter(|c| !col_ids.contains(&c.column_desc.column_id))
+                    .map(|c| col_id_to_tb_idx.get(&c.column_desc.column_id).unwrap())
                     .collect_vec();
                 let mut required_col_idx = self.required_col_idx.clone();
                 required_col_idx.extend(col_need_to_add);

rust/frontend/test_runner/src/lib.rs

@@ -163,6 +163,16 @@ impl TestCase {
                     Statement::CreateTable { name, columns, .. } => {
                         create_table::handle_create_table(context, name, columns).await?;
                     }
+                    Statement::CreateView {
+                        materialized: true,
+                        or_replace: false,
+                        name,
+                        query,
+                        ..
+                    } => {
+                        create_mv::handle_create_mv(context, name, query).await?;
+                    }
+
                     Statement::Drop(drop_statement) => {
                         let table_object_name = ObjectName(vec![drop_statement.name]);
                         drop_table::handle_drop_table(context, table_object_name).await?;

rust/frontend/test_runner/tests/testdata/mv_on_mv.yaml

@@ -0,0 +1,19 @@
+- id: create_tables
+  sql: |
+    create table t1(v1 int not null, v2 int not null);
+    create table t2(v1 real not null, v2 int not null, v3 real not null);
+    create materialized view m1 as select v1, v2 from t1 where v1 = 1;
+    create materialized view m2 as select v1, v2 from t1 where v2 = 1;
+- id: mv_on_mv_join
+  before:
+    - create_tables
+  sql: |
+    select m1.v1 as m1v1, m1.v2 as m1v2, m2.v1 as m2v1, m2.v2 as m2v2 from m1 join m2 on m1.v1 = m2.v1;
+  stream_plan: |
+    StreamMaterialize { columns: [m1v1, m1v2, m2v1, m2v2, expr#4(hidden), expr#5(hidden)], pk_columns: [expr#4, expr#5] }
+      StreamProject { exprs: [$0, $1, $3, $4, $2, $5], expr_alias: [m1v1, m1v2, m2v1, m2v2,  ,  ] }
+        StreamHashJoin { type: Inner, predicate: $0 = $3 }
+          StreamExchange { dist: HashShard([0]) }
+            StreamTableScan { table: m1, columns: [v1, v2, _row_id], pk_indices: [2] }
+          StreamExchange { dist: HashShard([0]) }
+            StreamTableScan { table: m2, columns: [v1, v2, _row_id], pk_indices: [2] }