risingwave_hummock_sdk/
table_stats.rs

1// Copyright 2025 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 std::borrow::Borrow;
16use std::collections::HashMap;
17
18use risingwave_common::catalog::TableId;
19use risingwave_pb::hummock::PbTableStats;
20
21use crate::version::HummockVersion;
22
23pub type TableStatsMap = HashMap<u32, TableStats>;
24
25pub type PbTableStatsMap = HashMap<u32, PbTableStats>;
26
27#[derive(Default, Debug, Clone)]
28pub struct TableStats {
29    pub total_key_size: i64,
30    pub total_value_size: i64,
31    pub total_key_count: i64,
32
33    // `total_compressed_size`` represents the size that the table takes up in the output sst
34    //  and this field is only filled and used by CN flushes, not compactor compaction
35    pub total_compressed_size: u64,
36}
37
38impl From<&TableStats> for PbTableStats {
39    fn from(value: &TableStats) -> Self {
40        Self {
41            total_key_size: value.total_key_size,
42            total_value_size: value.total_value_size,
43            total_key_count: value.total_key_count,
44            total_compressed_size: value.total_compressed_size,
45        }
46    }
47}
48
49impl From<TableStats> for PbTableStats {
50    fn from(value: TableStats) -> Self {
51        (&value).into()
52    }
53}
54
55impl From<&PbTableStats> for TableStats {
56    fn from(value: &PbTableStats) -> Self {
57        Self {
58            total_key_size: value.total_key_size,
59            total_value_size: value.total_value_size,
60            total_key_count: value.total_key_count,
61            total_compressed_size: value.total_compressed_size,
62        }
63    }
64}
65
66impl TableStats {
67    pub fn add(&mut self, other: &TableStats) {
68        self.total_key_size += other.total_key_size;
69        self.total_value_size += other.total_value_size;
70        self.total_key_count += other.total_key_count;
71        self.total_compressed_size += other.total_compressed_size;
72    }
73}
74
75pub fn add_prost_table_stats(this: &mut PbTableStats, other: &PbTableStats) {
76    this.total_key_size += other.total_key_size;
77    this.total_value_size += other.total_value_size;
78    this.total_key_count += other.total_key_count;
79    this.total_compressed_size += other.total_compressed_size;
80}
81
82pub fn add_prost_table_stats_map(this: &mut PbTableStatsMap, other: &PbTableStatsMap) {
83    for (table_id, stats) in other {
84        add_prost_table_stats(this.entry(*table_id).or_default(), stats);
85    }
86}
87
88pub fn add_table_stats_map(this: &mut TableStatsMap, other: &TableStatsMap) {
89    for (table_id, stats) in other {
90        this.entry(*table_id).or_default().add(stats);
91    }
92}
93
94pub fn to_prost_table_stats_map(
95    table_stats: impl Borrow<TableStatsMap>,
96) -> HashMap<u32, PbTableStats> {
97    table_stats
98        .borrow()
99        .iter()
100        .map(|(t, s)| (*t, s.into()))
101        .collect()
102}
103
104pub fn from_prost_table_stats_map(
105    table_stats: impl Borrow<HashMap<u32, PbTableStats>>,
106) -> HashMap<u32, TableStats> {
107    table_stats
108        .borrow()
109        .iter()
110        .map(|(t, s)| (*t, s.into()))
111        .collect()
112}
113
114pub fn purge_prost_table_stats(
115    table_stats: &mut PbTableStatsMap,
116    hummock_version: &HummockVersion,
117) -> bool {
118    let prev_count = table_stats.len();
119    table_stats.retain(|table_id, _| {
120        hummock_version
121            .state_table_info
122            .info()
123            .contains_key(&TableId::new(*table_id))
124    });
125    prev_count != table_stats.len()
126}