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

Code before:
            Some(p)
        } else {
            self.try_on_index(logical_top_n, logical_scan, order)
        }
    }
}

impl TopNOnIndexRule {
    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(),
                            },

Review comment (refactor):
The logic for handling the output column indices is duplicated in both `try_on_index` and `try_on_pk` methods (lines 57-64 and lines 129-136). Both methods check if `output_col_idx` equals `required_col_idx` and conditionally clone the scan with updated indices.

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

Code after:
}

impl TopNOnIndexRule {
    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 (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()