linkrs_query/executor/streaming/operators/
visited_set.rs

1use bitvec::prelude::*;
2use std::collections::HashSet;
3
4use linkrs_core::types::storage_ids::VertexId;
5
6const DENSE_THRESHOLD: usize = 64;
7
8enum Inner {
9    Sparse(HashSet<VertexId>),
10    Dense {
11        ids: HashSet<VertexId>,
12        bitmap: BitVec,
13        offset: i64,
14    },
15}
16
17pub struct VisitedSet {
18    inner: Inner,
19    switch_count: usize,
20}
21
22impl VisitedSet {
23    pub fn new() -> Self {
24        Self {
25            inner: Inner::Sparse(HashSet::new()),
26            switch_count: 0,
27        }
28    }
29
30    pub fn with_capacity(cap: usize) -> Self {
31        Self {
32            inner: Inner::Sparse(HashSet::with_capacity(cap)),
33            switch_count: 0,
34        }
35    }
36
37    pub fn contains(&self, id: &VertexId) -> bool {
38        match &self.inner {
39            Inner::Sparse(set) => set.contains(id),
40            Inner::Dense {
41                ids,
42                bitmap,
43                offset,
44            } => {
45                if let Some(idx) = id_to_dense_index(id, *offset) {
46                    if idx < bitmap.len() && bitmap[idx] {
47                        return true;
48                    }
49                }
50                ids.contains(id)
51            }
52        }
53    }
54
55    pub fn insert(&mut self, id: VertexId) -> bool {
56        match &mut self.inner {
57            Inner::Sparse(set) => {
58                if set.len() >= DENSE_THRESHOLD && self.switch_count < 3 {
59                    if let Some((offset, max_id)) = try_switch_to_dense(set, &id) {
60                        let mut bitmap = bitvec![0; (max_id - offset + 1) as usize];
61                        for existing in set.iter() {
62                            if let Some(idx) = id_to_dense_index(existing, offset) {
63                                if idx < bitmap.len() {
64                                    bitmap.set(idx, true);
65                                }
66                            }
67                        }
68                        let new = if let Some(idx) = id_to_dense_index(&id, offset) {
69                            if idx < bitmap.len() {
70                                let already = bitmap[idx];
71                                bitmap.set(idx, true);
72                                already
73                            } else {
74                                false
75                            }
76                        } else {
77                            false
78                        };
79                        if !new {
80                            set.insert(id);
81                        }
82                        self.inner = Inner::Dense {
83                            ids: std::mem::take(set),
84                            bitmap,
85                            offset,
86                        };
87                        self.switch_count += 1;
88                        return !new;
89                    }
90                }
91                set.insert(id)
92            }
93            Inner::Dense {
94                ids,
95                bitmap,
96                offset,
97            } => {
98                if let Some(idx) = id_to_dense_index(&id, *offset) {
99                    if idx < bitmap.len() {
100                        if bitmap[idx] {
101                            return false;
102                        }
103                        bitmap.set(idx, true);
104                        return ids.insert(id);
105                    }
106                }
107                ids.insert(id)
108            }
109        }
110    }
111
112    pub fn len(&self) -> usize {
113        match &self.inner {
114            Inner::Sparse(set) => set.len(),
115            Inner::Dense { ids, .. } => ids.len(),
116        }
117    }
118
119    pub fn is_empty(&self) -> bool {
120        self.len() == 0
121    }
122
123    pub fn clear(&mut self) {
124        match &mut self.inner {
125            Inner::Sparse(set) => set.clear(),
126            Inner::Dense { ids, bitmap, .. } => {
127                ids.clear();
128                bitmap.clear();
129            }
130        }
131    }
132
133    pub fn iter(&self) -> Box<dyn Iterator<Item = &VertexId> + '_> {
134        match &self.inner {
135            Inner::Sparse(set) => Box::new(set.iter()),
136            Inner::Dense { ids, .. } => Box::new(ids.iter()),
137        }
138    }
139}
140
141fn id_to_dense_index(id: &VertexId, offset: i64) -> Option<usize> {
142    let raw = id.as_int64()?;
143    let idx = raw.checked_sub(offset)?;
144    if idx < 0 {
145        return None;
146    }
147    usize::try_from(idx).ok()
148}
149
150fn try_switch_to_dense(set: &HashSet<VertexId>, candidate: &VertexId) -> Option<(i64, i64)> {
151    let min_id = set.iter().filter_map(|id| id.as_int64()).min()?;
152    let max_id = set.iter().filter_map(|id| id.as_int64()).max()?;
153    let candidate_int = candidate.as_int64()?;
154
155    let overall_min = min_id.min(candidate_int);
156    let overall_max = max_id.max(candidate_int);
157
158    let range_len = (overall_max - overall_min + 1) as u64;
159    if range_len > 1_000_000 {
160        return None;
161    }
162
163    Some((overall_min, overall_max))
164}
165
166impl Default for VisitedSet {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172impl std::fmt::Debug for VisitedSet {
173    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174        match &self.inner {
175            Inner::Sparse(set) => write!(f, "VisitedSet(Sparse, len={})", set.len()),
176            Inner::Dense { ids, bitmap, .. } => {
177                write!(
178                    f,
179                    "VisitedSet(Dense, ids={}, bitmap_ones={})",
180                    ids.len(),
181                    bitmap.count_ones()
182                )
183            }
184        }
185    }
186}
linkrs_query/executor/streaming/operators/visited_set.rs

linkrs_query/executor/streaming/operators/
visited_set.rs
