risingwave_common/hash/consistent_hash/
bitmap.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::ops::RangeInclusive;
16use std::sync::{Arc, LazyLock};
17
18use crate::bitmap::{Bitmap, BitmapBuilder};
19use crate::hash::VirtualNode;
20use crate::hash::table_distribution::SINGLETON_VNODE;
21
22/// An extension trait for `Bitmap` to support virtual node operations.
23#[easy_ext::ext(VnodeBitmapExt)]
24impl Bitmap {
25    /// Enumerates the virtual nodes set to 1 in the bitmap.
26    pub fn iter_vnodes(&self) -> impl Iterator<Item = VirtualNode> + '_ {
27        self.iter_ones().map(VirtualNode::from_index)
28    }
29
30    /// Enumerates the virtual nodes set to 1 in the bitmap.
31    pub fn iter_vnodes_scalar(&self) -> impl Iterator<Item = i16> + '_ {
32        self.iter_vnodes().map(|vnode| vnode.to_scalar())
33    }
34
35    /// Returns an iterator which yields the position ranges of continuous virtual nodes set to 1 in
36    /// the bitmap.
37    pub fn vnode_ranges(&self) -> impl Iterator<Item = RangeInclusive<VirtualNode>> + '_ {
38        self.high_ranges()
39            .map(|r| (VirtualNode::from_index(*r.start())..=VirtualNode::from_index(*r.end())))
40    }
41
42    /// Returns whether only the [`SINGLETON_VNODE`] is set in the bitmap.
43    pub fn is_singleton(&self) -> bool {
44        self.count_ones() == 1 && self.iter_vnodes().next().unwrap() == SINGLETON_VNODE
45    }
46
47    /// Get the reference to a vnode bitmap for singleton actor or table, i.e., with length
48    /// 1 and the only [`SINGLETON_VNODE`] set to true.
49    pub fn singleton() -> &'static Self {
50        Self::singleton_arc()
51    }
52
53    /// Get the reference to a vnode bitmap for singleton actor or table, i.e., with length
54    /// 1 and the only [`SINGLETON_VNODE`] set to true.
55    pub fn singleton_arc() -> &'static Arc<Self> {
56        static SINGLETON: LazyLock<Arc<Bitmap>> = LazyLock::new(|| {
57            let mut builder = BitmapBuilder::zeroed(1);
58            builder.set(SINGLETON_VNODE.to_index(), true);
59            builder.finish().into()
60        });
61        &SINGLETON
62    }
63}
risingwave_common/hash/consistent_hash/bitmap.rs

risingwave_common/hash/consistent_hash/
bitmap.rs
