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

Code before:
    pub fn create() -> BoxedRule {
        Box::new(TopNOnIndexRule {})
    }

    fn try_on_index(
        &self,
        logical_top_n: &LogicalTopN,
        logical_scan: LogicalScan,
        required_order: &Order,
    ) -> Option<PlanRef> {
        let output_len = logical_scan.output_col_idx().len();
        let scan_for_index = if logical_scan.output_col_idx() == logical_scan.required_col_idx() {
            logical_scan
        } else {
            // `required_col_idx` is built by appending predicate columns to `output_col_idx`,
            // so `output_col_idx` is always a prefix of `required_col_idx`.
            logical_scan.clone_with_output_indices(logical_scan.required_col_idx().clone())
        };
        let scan_predicates = scan_for_index.predicate();
        let output_col_map = scan_for_index
            .output_col_idx()
            .iter()
            .cloned()
            .enumerate()
            .map(|(id, col)| (col, id))
            .collect::<BTreeMap<_, _>>();
        let input_refs = scan_predicates.get_eq_const_input_refs();
        let prefix = input_refs
            .into_iter()
            .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(),
                            },

Review comment (refactor):
The logic for building the output column map from scan output indices is duplicated in both methods (lines 66-72 and lines 138-144). This creates a BTreeMap that maps column indices to their output positions.

Consider extracting this into a helper method to improve maintainability and reduce code duplication.

Code after:
        &self,
        logical_top_n: &LogicalTopN,
        logical_scan: LogicalScan,
        required_order: &Order,
    ) -> Option<PlanRef> {
        let (scan_for_index, output_len, _output_col_map, prefix) =
            Self::prepare_scan_for_predicate(logical_scan);
        let scan_output_len = scan_for_index.output_col_idx().len();
        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) {
                    selected_index = Some(Self::finish_top_n(
                        logical_top_n,
                        index_scan.into(),
                        prefix,
                        output_len,
                        scan_output_len,
                    ));
                }
            }
        }
        selected_index
    }

    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,
                    )
                })