Repository: risingwavelabs/risingwave
File: src/frontend/src/optimizer/rule/top_n_on_index_rule.rs (Rust)

Code before:
            .flat_map(|input_ref| {
                output_col_map
                    .get(&input_ref.index)
                    .into_iter()
                    .flat_map(|&output_idx| {
                        [
                            ColumnOrder {
                                column_index: output_idx,
                                order_type: OrderType::ascending_nulls_first(),
                            },
                            ColumnOrder {
                                column_index: output_idx,
                                order_type: OrderType::ascending_nulls_last(),
                            },
                            ColumnOrder {
                                column_index: output_idx,
                                order_type: OrderType::descending_nulls_first(),
                            },
                            ColumnOrder {
                                column_index: output_idx,
                                order_type: OrderType::descending_nulls_last(),
                            },
                        ]
                    })
            })
            .collect();
        let order_satisfied_index =
            scan_for_index.indexes_satisfy_order_with_prefix(required_order, &prefix);
        let mut longest_prefix: Option<Order> = None;
        let mut selected_index = None;
        for (index, prefix) in order_satisfied_index {
            if prefix.len() >= longest_prefix.as_ref().map_or(0, |p| p.len()) {
                longest_prefix = Some(prefix.clone());
                if let Some(index_scan) = scan_for_index.to_index_scan_if_index_covered(index) {
                    let top_n =
                        logical_top_n.clone_with_input_and_prefix(index_scan.into(), prefix);
                    selected_index = Some(if scan_for_index.output_col_idx().len() == output_len {
                        top_n.into()
                    } else {
                        LogicalProject::with_out_col_idx(top_n.into(), 0..output_len).into()
                    });
                }
            }
        }
        selected_index
    }

    fn try_on_pk(
        &self,
        logical_top_n: &LogicalTopN,
        logical_scan: LogicalScan,

Review comment (refactor):
The code for building the prefix from scan predicates is duplicated in both `try_on_index` and `try_on_pk` methods (lines 74-101 and lines 146-173). This duplication makes the code harder to maintain and increases the risk of inconsistencies if changes are needed.

Consider extracting this logic into a helper method that takes the scan predicates and output column map as parameters and returns the prefix collection.

Code after:
    fn try_on_pk(
        &self,
        logical_top_n: &LogicalTopN,
        logical_scan: LogicalScan,
        order: &Order,
    ) -> Option<PlanRef> {
        let (scan_for_pk, output_len, output_col_map, prefix) =
            Self::prepare_scan_for_predicate(logical_scan);
        let scan_output_len = scan_for_pk.output_col_idx().len();
        let unmatched_idx = output_col_map.len();
        let primary_key = scan_for_pk.primary_key();
        let primary_key_order = Order {
            column_orders: primary_key
                .iter()
                .map(|o| {
                    ColumnOrder::new(
                        *output_col_map
                            .get(&o.column_index)
                            .unwrap_or(&unmatched_idx),
                        o.order_type,
                    )
                })
                .collect::<Vec<_>>(),
        };
        let mut pk_orders_iter = primary_key_order.column_orders.iter().cloned().peekable();
        let fixed_prefix = {
            let mut fixed_prefix = vec![];
            loop {
                match pk_orders_iter.peek() {
                    Some(order) if prefix.contains(order) => {
                        let order = pk_orders_iter.next().unwrap();
                        fixed_prefix.push(order);
                    }
                    _ => break,
                }
            }
            Order {
                column_orders: fixed_prefix,
            }
        };
        let remaining_orders = Order {
            column_orders: pk_orders_iter.collect(),
        };
        if !remaining_orders.satisfies(order) {
            return None;
        }
        Some(Self::finish_top_n(
            logical_top_n,
            scan_for_pk.into(),
            fixed_prefix,
            output_len,