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

Code before:
// Copyright 2026 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::optimizer::PlanRef;
use crate::optimizer::plan_node::{Logical, LogicalPlanRef, LogicalScan};
use crate::optimizer::rule::{BoxedRule, Rule};

pub struct MvSelectionRule;

impl MvSelectionRule {
    pub fn create() -> BoxedRule<Logical> {
        Box::new(Self)
    }
}

impl Rule<Logical> for MvSelectionRule {
    fn apply(&self, plan: PlanRef<Logical>) -> Option<PlanRef<Logical>> {
        let ctx = plan.ctx();
        for candidate in ctx.batch_mview_candidates().iter() {
            if plan == candidate.plan {
                let output_len = plan.schema().len();
                if output_len > candidate.table.columns().len() {
                    continue;
                }
                let output_col_idx = (0..output_len).collect();
                let scan = LogicalScan::create(candidate.table.clone(), ctx.clone(), None)
                    .clone_with_output_indices(output_col_idx);
                let rewritten: LogicalPlanRef = scan.into();
                return Some(rewritten);
            }
        }
        None
    }
}


Review comment (bug):
`candidate_scan` relies on `debug_assert!(output_len <= candidate.table.columns().len())`, but in release builds this won’t run and `clone_with_output_indices` can panic if the lengths ever diverge. Consider restoring a runtime guard (like the previous `if output_len > ... { return None; }`) so MV selection can fail gracefully instead of crashing the optimizer.

Code after:
        None
    }
}

impl MvSelectionRule {
    fn rewrite_with_candidate(
        plan: &PlanRef,
        candidate: &MaterializedViewCandidate,
    ) -> Option<PlanRef> {
        if plan == &candidate.plan {
            return Some(Self::candidate_scan(candidate, plan.ctx())?.into());
        }
        Self::agg_rollup_rewrite(plan, candidate)
    }

    fn candidate_scan(
        candidate: &MaterializedViewCandidate,
        ctx: crate::optimizer::optimizer_context::OptimizerContextRef,
    ) -> Option<LogicalScan> {
        let output_len = candidate.plan.schema().len();
        debug_assert!(output_len <= candidate.table.columns().len());
        let output_col_idx = (0..output_len).collect();
        Some(
            LogicalScan::create(candidate.table.clone(), ctx, None)
                .clone_with_output_indices(output_col_idx),
        )
    }

    /// Rewrite aggregate query by rolling up aggregate MV states.
    ///
    /// Graph (matching + rewrite):
    /// ```text
    /// 1) Normalize both sides to base-column lineage
    ///
    /// Query side                                  MV candidate side
    /// Agg_q(Gq, Aq)                               [Project_mv]?
    ///   └─ [Project_q]?                             └─ Agg_mv(Gm, Am)
    ///       └─ BaseInput                                  └─ [Project_mv_in]?
    ///                                                         └─ BaseInput
    ///
    /// 2) Match under lineage mapping
    /// - same BaseInput
    /// - map(Gq) ⊆ map(Gm)
    /// - for each aq in Aq, find semantically equivalent am in Am
    /// - aq must support partial_to_total
    ///
    /// 3) Build rewritten plan
    /// Agg_rollup(group = mapped Gq, calls = total(aq) on MV columns)
    ///   └─ Scan(MV table)
    /// [Project] (optional, restore query output order)
    /// ```