fix functional index selection

src/frontend/planner_test/tests/testdata/input/index_selection.yaml

@@ -453,3 +453,10 @@
     select ts::timestamptz from t where ts > now();
   expected_outputs:
   - batch_plan
+- name: funtional index with timezone2
+  sql: |
+    create table t (a int, ts timestamp without time zone);
+    create index idx on t (CAST(ts AS timestamptz));
+    select * from t order by ts;
+  expected_outputs:
+    - batch_plan

src/frontend/planner_test/tests/testdata/output/index_selection.yaml

@@ -766,3 +766,12 @@
     BatchExchange { order: [], dist: Single }
     └─BatchProject { exprs: [AtTimeZone(i.ts, 'UTC':Varchar) as $expr1] }
       └─BatchScan { table: i, columns: [i.ts], scan_ranges: [i.AT_TIME_ZONE > Timestamptz(Timestamptz(1617235200000000))], distribution: SomeShard }
+- name: funtional index with timezone2
+  sql: |
+    create table t (a int, ts timestamp without time zone);
+    create index idx on t (CAST(ts AS timestamptz));
+    select * from t order by ts;
+  batch_plan: |-
+    BatchExchange { order: [t.ts ASC], dist: Single }
+    └─BatchSort { order: [t.ts ASC] }
+      └─BatchScan { table: t, columns: [t.a, t.ts], distribution: SomeShard }

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

@@ -163,16 +163,14 @@ impl LogicalScan {
                         .pk()
                         .iter()
                         .map(|idx_item| {
-                            ColumnOrder::new(
-                                *output_col_map
-                                    .get(
-                                        s2p_mapping
-                                            .get(&idx_item.column_index)
-                                            .expect("should be in s2p mapping"),
-                                    )
-                                    .unwrap_or(&unmatched_idx),
-                                idx_item.order_type,
-                            )
+                            let idx = match s2p_mapping.get(&idx_item.column_index) {
+                                Some(col_idx) => {
+                                    *output_col_map.get(col_idx).unwrap_or(&unmatched_idx)
+                                }
+                                // After we support index on expressions, we need to handle the case where the column is not in the `s2p_mapping`.
+                                None => unmatched_idx,
+                            };
+                            ColumnOrder::new(idx, idx_item.order_type)
                         })
                         .collect(),
                 }