dbsp/trace/filter/
batch.rs

1//! Key filters for exact seeks.
2//!
3//! These filters only answer "definitely not present" versus "might still be
4//! present", so callers can reject cheap cases before doing the exact lookup.
5
6use crate::{
7    dynamic::{DataTrait, DynVec},
8    storage::{
9        file::{
10            BatchKeyFilter, FilterKind, FilterStats, TrackingFilterStats, TrackingRoaringBitmap,
11            reader::FilteredKeys,
12        },
13        tracking_bloom_filter::TrackingBloomFilter,
14    },
15    trace::filter::key_range::KeyRange,
16};
17use size_of::SizeOf;
18use std::sync::Arc;
19
20/// A cheap, in-memory precheck used by `seek_key_exact`.
21///
22/// Each filter may only reject a key early. Returning `true` means "the key
23/// might still be present", so the caller must continue with the next filter or
24/// the exact lookup.
25pub(crate) trait BatchFilter<K>: Send + Sync
26where
27    K: DataTrait + ?Sized,
28{
29    /// Returns `false` only when the key is definitely absent.
30    ///
31    /// Filters that need a hash compute it through `hash`, so a chain of
32    /// filters pays that cost at most once.
33    fn maybe_contains_key(&self, key: &K, hash: &mut Option<u64>) -> bool;
34
35    /// Filter kind for observability.
36    fn kind(&self) -> FilterKind;
37
38    /// Statistics for this filter.
39    fn stats(&self) -> FilterStats;
40}
41
42/// Ordered key filters for exact-key probes against file-backed batches.
43///
44/// The key range is kept separately from the trait objects because batches
45/// thread typed bounds through file-writer finalization and rebuild the range
46/// when they create a new reader. All other filters are stored behind a trait
47/// object so we can swap bloom for bitmap or another single post-range filter
48/// without changing the batch structs again.
49pub struct BatchFilters<K>
50where
51    K: DataTrait + ?Sized,
52{
53    range_filter: Arc<TrackedRangeFilter<K>>,
54    membership_filter: Option<Arc<dyn BatchFilter<K>>>,
55}
56
57/// Runtime statistics for the range and membership filters inside
58/// [`BatchFilters`].
59#[derive(Clone, Copy, Debug, Default, PartialEq)]
60pub struct BatchFilterStats {
61    pub range_filter: FilterStats,
62    pub membership_filter: FilterStats,
63}
64
65/// A range filter with statistics.
66#[derive(SizeOf)]
67struct TrackedRangeFilter<K>
68where
69    K: DataTrait + ?Sized,
70{
71    range: Option<KeyRange<K>>,
72    #[size_of(skip)]
73    tracking: TrackingFilterStats,
74}
75
76impl<K> TrackedRangeFilter<K>
77where
78    K: DataTrait + ?Sized,
79{
80    fn new(range: Option<KeyRange<K>>) -> Self {
81        let size_byte = range
82            .as_ref()
83            .map(|range| range.size_of().total_bytes())
84            .unwrap_or_default();
85        Self {
86            range,
87            tracking: TrackingFilterStats::new(size_byte),
88        }
89    }
90
91    fn stats(&self) -> FilterStats {
92        self.tracking.stats()
93    }
94}
95
96impl<K> BatchFilters<K>
97where
98    K: DataTrait + ?Sized,
99{
100    pub(crate) fn new(
101        range_filter: Option<KeyRange<K>>,
102        membership_filter: Option<Arc<dyn BatchFilter<K>>>,
103    ) -> Self {
104        Self {
105            range_filter: Arc::new(TrackedRangeFilter::new(range_filter)),
106            membership_filter,
107        }
108    }
109
110    /// Rebuilds the filter chain from persisted key bounds and an optional
111    /// membership filter.
112    ///
113    /// This is the public constructor for callers that read filter state back
114    /// from disk and want the same range-then-membership behavior that trace
115    /// batches use for `seek_key_exact`.
116    pub fn from_parts(
117        key_bounds: Option<(Box<K>, Box<K>)>,
118        membership_filter: Option<BatchKeyFilter>,
119    ) -> Self {
120        Self::from_file(key_bounds.map(KeyRange::from), membership_filter)
121    }
122
123    /// Builds the current file-backed filter chain.
124    ///
125    /// The range check runs before the bloom filter so out-of-range keys never
126    /// pay the hash or bloom lookup cost.
127    pub(crate) fn from_file(
128        key_range: Option<KeyRange<K>>,
129        membership_filter: Option<BatchKeyFilter>,
130    ) -> Self {
131        let membership_filter = membership_filter.map(Arc::<dyn BatchFilter<K>>::from);
132        Self::new(key_range, membership_filter)
133    }
134
135    /// Returns cumulative statistics for the range and membership filters.
136    pub fn stats(&self) -> BatchFilterStats {
137        BatchFilterStats {
138            range_filter: self.range_filter.stats(),
139            membership_filter: self
140                .membership_filter
141                .as_ref()
142                .map(|filter| filter.stats())
143                .unwrap_or_default(),
144        }
145    }
146
147    pub fn membership_filter_kind(&self) -> FilterKind {
148        self.membership_filter
149            .as_ref()
150            .map(|filter| filter.kind())
151            .unwrap_or(FilterKind::None)
152    }
153
154    /// Returns the cached key bounds, when available.
155    pub fn key_bounds(&self) -> Option<(&K, &K)> {
156        self.range_filter.range.as_ref().map(|range| range.bounds())
157    }
158
159    /// Returns `keys`, optionally narrowed to the indexes that pass the filter
160    /// chain when that is cheap enough to be worthwhile.
161    pub(crate) fn filtered_keys<'a>(&self, keys: &'a DynVec<K>) -> FilteredKeys<'a, K> {
162        debug_assert!(keys.is_sorted_by(&|a, b| a.cmp(b)));
163
164        let mut filter_pass_keys = Vec::with_capacity(keys.len().min(50));
165        for (index, key) in keys.dyn_iter().enumerate() {
166            if self.maybe_contains_key(key, None) {
167                filter_pass_keys.push(index);
168                if filter_pass_keys.len() >= keys.len() / 300 {
169                    return FilteredKeys::all(keys);
170                }
171            }
172        }
173
174        FilteredKeys::with_filter_pass_keys(keys, Some(filter_pass_keys))
175    }
176
177    /// Returns `false` only when `key` is definitely not present.
178    ///
179    /// Passing a cached `hash` avoids recomputing it when the caller already
180    /// has one available.
181    pub fn maybe_contains_key(&self, key: &K, mut hash: Option<u64>) -> bool {
182        // The range filter runs first so later filters only see keys inside
183        // the batch bounds. That matters for roaring-backed batches: once the
184        // range check has accepted the key, min-offset conversion is known to
185        // fit in `u32` because roaring filters are only built for batches
186        // whose full key range fits in `u32`.
187        if !self.range_filter.maybe_contains_key(key, &mut hash) {
188            return false;
189        }
190
191        self.membership_filter
192            .as_ref()
193            .is_none_or(|filter| filter.maybe_contains_key(key, &mut hash))
194    }
195}
196
197impl<K> Clone for BatchFilters<K>
198where
199    K: DataTrait + ?Sized,
200{
201    fn clone(&self) -> Self {
202        Self {
203            range_filter: self.range_filter.clone(),
204            membership_filter: self.membership_filter.clone(),
205        }
206    }
207}
208
209impl<K> SizeOf for BatchFilters<K>
210where
211    K: DataTrait + ?Sized,
212{
213    fn size_of_children(&self, context: &mut size_of::Context) {
214        self.range_filter.size_of_with_context(context);
215        context.add(
216            self.membership_filter
217                .as_ref()
218                .map(|filter| filter.stats().size_byte)
219                .unwrap_or_default(),
220        );
221    }
222}
223
224impl<K> BatchFilter<K> for Arc<TrackedRangeFilter<K>>
225where
226    K: DataTrait + ?Sized,
227{
228    fn maybe_contains_key(&self, key: &K, _hash: &mut Option<u64>) -> bool {
229        let is_hit = self.range.as_ref().is_some_and(|range| range.contains(key));
230        self.tracking.record(is_hit);
231        is_hit
232    }
233
234    fn kind(&self) -> FilterKind {
235        FilterKind::Range
236    }
237
238    fn stats(&self) -> FilterStats {
239        self.as_ref().stats()
240    }
241}
242
243impl<K> BatchFilter<K> for TrackingBloomFilter
244where
245    K: DataTrait + ?Sized,
246{
247    fn maybe_contains_key(&self, key: &K, hash: &mut Option<u64>) -> bool {
248        let hash = hash.get_or_insert_with(|| key.default_hash());
249        self.contains_hash(*hash)
250    }
251
252    fn kind(&self) -> FilterKind {
253        FilterKind::Bloom
254    }
255
256    fn stats(&self) -> FilterStats {
257        TrackingBloomFilter::stats(self)
258    }
259}
260
261impl<K> BatchFilter<K> for TrackingRoaringBitmap
262where
263    K: DataTrait + ?Sized,
264{
265    fn maybe_contains_key(&self, key: &K, _hash: &mut Option<u64>) -> bool {
266        self.maybe_contains_key(key)
267    }
268
269    fn kind(&self) -> FilterKind {
270        FilterKind::Roaring
271    }
272
273    fn stats(&self) -> FilterStats {
274        TrackingRoaringBitmap::stats(self)
275    }
276}
277
278impl<K> From<BatchKeyFilter> for Arc<dyn BatchFilter<K>>
279where
280    K: DataTrait + ?Sized,
281{
282    fn from(filter: BatchKeyFilter) -> Self {
283        match filter {
284            BatchKeyFilter::Bloom(filter) => Arc::new(filter),
285            BatchKeyFilter::RoaringU32(filter) => Arc::new(filter),
286        }
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::{BatchFilter, TrackedRangeFilter};
293    use crate::{dynamic::DynData, storage::file::FilterStats, trace::filter::key_range::KeyRange};
294    use std::sync::Arc;
295
296    #[test]
297    fn tracked_range_filter_stats() {
298        let filter = Arc::new(TrackedRangeFilter::new(Some(KeyRange::from_refs(
299            (&1i32) as &DynData,
300            (&10i32) as &DynData,
301        ))));
302
303        assert!(filter.maybe_contains_key((&5i32) as &DynData, &mut None));
304        assert!(!filter.maybe_contains_key((&11i32) as &DynData, &mut None));
305
306        let stats = filter.stats();
307        assert!(stats.size_byte > 0);
308        assert_eq!(
309            stats,
310            FilterStats {
311                size_byte: stats.size_byte,
312                hits: 1,
313                misses: 1,
314            }
315        );
316    }
317}
dbsp/trace/filter/batch.rs

dbsp/trace/filter/
batch.rs
