Skip to main content

risingwave_frontend/optimizer/plan_visitor/
locality_backfill_scan_estimator.rs

1// Copyright 2026 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use risingwave_pb::hummock::HummockVersionStats;
16
17use super::{DefaultBehavior, LogicalPlanVisitor, Merge};
18use crate::optimizer::plan_node::{LogicalPlanRef, LogicalScan};
19use crate::optimizer::plan_visitor::PlanVisitor;
20
21#[derive(Debug, Clone)]
22pub struct LocalityBackfillScanEstimator<'a> {
23    table_stats: &'a HummockVersionStats,
24}
25
26impl<'a> LocalityBackfillScanEstimator<'a> {
27    pub fn estimate(plan: LogicalPlanRef, table_stats: &'a HummockVersionStats) -> u64 {
28        let mut estimator = Self { table_stats };
29        estimator.visit(plan)
30    }
31}
32
33impl LocalityBackfillScanEstimator<'_> {
34    fn table_size(&self, scan: &LogicalScan) -> u64 {
35        self.table_stats
36            .table_stats
37            .get(&scan.table().id.as_raw_id())
38            .map(|stats| {
39                (stats.total_key_size.max(0) as u64)
40                    .saturating_add(stats.total_value_size.max(0) as u64)
41            })
42            .unwrap_or_default()
43    }
44}
45
46impl LogicalPlanVisitor for LocalityBackfillScanEstimator<'_> {
47    type Result = u64;
48
49    type DefaultBehavior = impl DefaultBehavior<Self::Result>;
50
51    fn default_behavior() -> Self::DefaultBehavior {
52        Merge(u64::saturating_add)
53    }
54
55    fn visit_logical_scan(&mut self, scan: &LogicalScan) -> Self::Result {
56        self.table_size(scan)
57    }
58}