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

Code before:
                                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,
        order: &Order,
    ) -> Option<PlanRef> {
        let output_len = logical_scan.output_col_idx().len();
        let scan_for_pk = 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_pk.predicate();
        let output_col_map = scan_for_pk
            .output_col_idx()
            .iter()
            .cloned()

Review comment (refactor):
The logic for wrapping the TopN result with a LogicalProject when necessary is duplicated in both methods (lines 112-116 and lines 213-217). Both use the same pattern to decide whether to add a project based on the output column length.

Consider extracting this logic into a helper method that takes the TopN node and output length as parameters to eliminate this duplication.

Code after:
                    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,
            scan_output_len,
        ))
    }

    fn prepare_scan_for_predicate(
        logical_scan: LogicalScan,
    ) -> (
        LogicalScan,
        usize,
        BTreeMap<usize, usize>,
        HashSet<ColumnOrder>,
    ) {
        let output_len = logical_scan.output_col_idx().len();
        let scan = if logical_scan.output_col_idx() == logical_scan.required_col_idx() {
            logical_scan