risingwave_frontend/utils/
index_set.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::{BitAnd, BitOr};
16
17use fixedbitset::FixedBitSet;
18
19/// A set of deduplicated column indices.
20#[derive(Debug, Clone, PartialEq, Eq, Hash)]
21pub struct IndexSet {
22    bitset: FixedBitSet,
23}
24
25impl IndexSet {
26    /// Create an empty set.
27    pub fn empty() -> Self {
28        Self {
29            bitset: FixedBitSet::new(),
30        }
31    }
32
33    /// Get the number of indices.
34    pub fn len(&self) -> usize {
35        self.bitset.count_ones(..)
36    }
37
38    /// Check if the set is empty.
39    pub fn is_empty(&self) -> bool {
40        self.bitset.is_clear()
41    }
42
43    /// Check if the given index is contained.
44    pub fn contains(&self, index: usize) -> bool {
45        self.bitset.contains(index)
46    }
47
48    /// Iterator over indices.
49    pub fn indices(&self) -> impl Iterator<Item = usize> + '_ {
50        self.bitset.ones()
51    }
52
53    /// Convert to `Vec<usize>`.
54    pub fn to_vec(&self) -> Vec<usize> {
55        self.indices().collect()
56    }
57
58    /// Convert to `FixedBitSet`.
59    pub fn to_bitset(&self) -> FixedBitSet {
60        self.bitset.clone()
61    }
62
63    /// Iterator over indices as `u32`s. Please only use this when converting to protobuf.
64    pub fn indices_as_u32(&self) -> impl Iterator<Item = u32> + '_ {
65        self.bitset.ones().map(|i| i as u32)
66    }
67
68    /// Convert to `Vec<u32>`. Please only use this when converting to protobuf.
69    pub fn to_vec_as_u32(&self) -> Vec<u32> {
70        self.indices_as_u32().collect()
71    }
72
73    /// Insert an index.
74    pub fn insert(&mut self, index: usize) {
75        self.bitset.extend_one(index)
76    }
77}
78
79impl Extend<usize> for IndexSet {
80    fn extend<T: IntoIterator<Item = usize>>(&mut self, iter: T) {
81        self.bitset.extend(iter)
82    }
83}
84
85impl From<Vec<usize>> for IndexSet {
86    fn from(vec: Vec<usize>) -> Self {
87        vec.into_iter().collect()
88    }
89}
90
91impl FromIterator<usize> for IndexSet {
92    fn from_iter<T: IntoIterator<Item = usize>>(iter: T) -> Self {
93        IndexSet {
94            bitset: iter.into_iter().collect(),
95        }
96    }
97}
98
99impl BitAnd for IndexSet {
100    type Output = IndexSet;
101
102    fn bitand(self, rhs: Self) -> Self::Output {
103        IndexSet {
104            bitset: &self.bitset & &rhs.bitset,
105        }
106    }
107}
108
109impl BitAnd for &IndexSet {
110    type Output = IndexSet;
111
112    fn bitand(self, rhs: Self) -> Self::Output {
113        IndexSet {
114            bitset: &self.bitset & &rhs.bitset,
115        }
116    }
117}
118
119impl BitOr for IndexSet {
120    type Output = IndexSet;
121
122    fn bitor(self, rhs: Self) -> Self::Output {
123        IndexSet {
124            bitset: &self.bitset | &rhs.bitset,
125        }
126    }
127}
128
129impl BitOr for &IndexSet {
130    type Output = IndexSet;
131
132    fn bitor(self, rhs: Self) -> Self::Output {
133        IndexSet {
134            bitset: &self.bitset | &rhs.bitset,
135        }
136    }
137}
risingwave_frontend/utils/index_set.rs

risingwave_frontend/utils/
index_set.rs
