1use std::ops::Range;
22use std::pin::Pin;
23
24use futures::Stream;
25use futures::StreamExt;
26use futures::stream::BoxStream;
27use futures::stream::Peekable;
28use mysten_common::zip_debug_eq::ZipDebugEqIteratorExt;
29use roaring::RoaringBitmap;
30
31use super::BitmapBucketSource;
32use super::BitmapQuery;
33use super::BitmapScanBudget;
34use super::BitmapScanMetrics;
35use super::BucketItem;
36use super::BucketStream;
37use super::DedupedQuery;
38use super::LeafHead;
39use super::LeafStop;
40use super::ObserveOnDrop;
41use super::ScanDirection;
42use super::ScanStop;
43use super::SkipPolicy;
44use super::WatermarkPacer;
45use super::WatermarkPacing;
46use super::Watermarked;
47use super::WatermarkedBucketStream;
48use super::advance_in_direction;
49use super::bound_in_direction;
50use super::bucket_edges;
51use super::build_term_specs;
52use super::collapse;
53use super::count_on_floor_refs;
54use super::eval_term_at_bucket;
55use super::frontier_advanced;
56use super::leaf_skip_targets;
57use super::recompute_unreferenced;
58use super::strictly_before;
59use super::take_snapshot_bitmap;
60
61pub fn eval_bitmap_query_stream<S, F>(
75 source: S,
76 query: BitmapQuery,
77 range: Range<u64>,
78 bucket_size: u64,
79 direction: ScanDirection,
80 budget: u64,
81 policy: SkipPolicy,
82 pacing: WatermarkPacing,
83 on_metrics: F,
84) -> BoxStream<'static, Result<Watermarked<u64>, ScanStop>>
85where
86 S: BitmapBucketSource,
87 F: FnOnce(BitmapScanMetrics) + Send + 'static,
88{
89 let leaves = query.unique_leaf_count();
90 if (budget as usize) < leaves {
91 return async_stream::stream! {
95 yield Err(ScanStop::Fault(anyhow::anyhow!(
97 "bitmap scan budget {budget} is insufficient for {leaves} leaf streams; \
98 server is misconfigured"
99 )));
100 }
101 .boxed();
102 }
103 let budget = BitmapScanBudget::new(budget);
104 let unpaced = eval_bitmap_query_bucket_stream(
105 source,
106 query,
107 range.clone(),
108 bucket_size,
109 direction,
110 budget.clone(),
111 policy,
112 );
113 let pacer = WatermarkPacer::new(pacing, policy, budget.clone(), range.clone(), bucket_size);
114 let paced = pace_watermarked_buckets(unpaced, pacer);
115 let inner = flatten_watermarked_buckets(paced, range.clone(), bucket_size, direction);
116 let guard = ObserveOnDrop {
120 budget,
121 callback: Some(on_metrics),
122 };
123 async_stream::stream! {
124 let _guard = guard;
125 futures::pin_mut!(inner);
126 while let Some(item) = inner.next().await {
127 yield item;
128 }
129 }
130 .boxed()
131}
132
133fn pace_watermarked_buckets(
138 unpaced: WatermarkedBucketStream,
139 pacer: WatermarkPacer,
140) -> WatermarkedBucketStream {
141 async_stream::stream! {
142 let mut unpaced = unpaced;
143 let mut pacer = pacer;
144 while let Some(output) = unpaced.next().await {
145 if let Some(paced) = pacer.pace(output) {
146 yield paced;
147 }
148 }
149 if let Some(position) = pacer.finish() {
150 yield Ok(Watermarked::Watermark(position));
151 }
152 }
153 .boxed()
154}
155
156async fn peek_leaf<S>(
159 mut leaf: Pin<&mut Peekable<S>>,
160 idx: usize,
161 bucket_size: u64,
162 range: &Range<u64>,
163 direction: ScanDirection,
164 terminus: u64,
165) -> (usize, LeafHead, Option<u64>)
166where
167 S: Stream<Item = BucketItem>,
168{
169 match leaf.as_mut().peek().await {
170 Some(Ok((bucket, _))) => {
171 let (pre, _post) = bucket_edges(*bucket, bucket_size, range, direction);
172 (idx, LeafHead::Bucket(*bucket), Some(pre))
173 }
174 None => (idx, LeafHead::Eof, Some(terminus)),
175 Some(Err(_)) => (idx, LeafHead::Error, None),
176 }
177}
178
179pub(crate) fn eval_bitmap_query_bucket_stream<S>(
186 source: S,
187 query: BitmapQuery,
188 range: Range<u64>,
189 bucket_size: u64,
190 direction: ScanDirection,
191 budget: BitmapScanBudget,
192 policy: SkipPolicy,
193) -> WatermarkedBucketStream
194where
195 S: BitmapBucketSource,
196{
197 let DedupedQuery {
198 keys: unique_keys,
199 mut terms,
200 } = build_term_specs(query.terms);
201 let leaf_count = unique_keys.len();
202 let terminus = if direction.is_ascending() {
203 range.end
204 } else {
205 range.start
206 };
207 let request_floor = if direction.is_ascending() {
208 range.start
209 } else {
210 range.end
211 };
212
213 async_stream::try_stream! {
214 let mut leaves: Vec<Peekable<BucketStream>> = Vec::with_capacity(leaf_count);
215 for key in &unique_keys {
216 let raw = source.scan_bucket_stream(key.clone(), range.clone(), direction);
217 leaves.push(
218 budgeted_bucket_stream(raw, budget.clone(), true)
219 .boxed()
220 .peekable(),
221 );
222 }
223 let mut unreferenced = vec![false; leaf_count];
224 let mut front = vec![request_floor; leaf_count];
225 let mut drained = vec![0u64; leaf_count];
226 let mut progress_frontier = Some(request_floor);
227
228 loop {
229 let mut peeks = Vec::new();
230 for (i, leaf) in leaves.iter_mut().enumerate() {
231 if !unreferenced[i] {
232 peeks.push(peek_leaf(
233 Pin::new(leaf),
234 i,
235 bucket_size,
236 &range,
237 direction,
238 terminus,
239 ));
240 }
241 }
242 let results = futures::future::join_all(peeks).await;
243 let mut class: Vec<Option<LeafHead>> = (0..leaf_count).map(|_| None).collect();
244 for (i, head, scanned_to) in results {
245 if let Some(position) = scanned_to {
246 front[i] = advance_in_direction(front[i], position, direction);
247 }
248 class[i] = Some(head);
249 }
250
251 for term in terms.iter_mut() {
252 if !term.unsatisfiable
253 && term
254 .includes
255 .iter()
256 .any(|&i| matches!(class[i], Some(LeafHead::Eof)))
257 {
258 term.unsatisfiable = true;
259 }
260 }
261 recompute_unreferenced(&terms, &class, &mut unreferenced);
262
263 let targets = leaf_skip_targets(&terms, &class, &unreferenced, direction);
264 for (i, target) in targets.iter().enumerate() {
265 if let Some(target) = target {
266 let (pre, _) = bucket_edges(*target, bucket_size, &range, direction);
267 front[i] = advance_in_direction(front[i], pre, direction);
268 }
269 }
270
271 let mut errors: Vec<LeafStop> = Vec::new();
272 for i in 0..leaf_count {
273 if !unreferenced[i] && matches!(class[i], Some(LeafHead::Error)) {
274 match Pin::new(&mut leaves[i]).next().await {
275 Some(Err(error)) => errors.push(error),
276 _ => unreachable!("peek classified Error"),
277 }
278 }
279 }
280
281 let active: Vec<usize> = (0..leaf_count).filter(|&i| !unreferenced[i]).collect();
282 if active.is_empty() {
283 return;
284 }
285
286 let floor_pos = active
287 .iter()
288 .map(|&i| front[i])
289 .reduce(|a, b| bound_in_direction(a, b, direction))
290 .expect("active non-empty");
291 let collapsed = (!errors.is_empty()).then(|| collapse(errors, floor_pos));
292 let scan_limited = matches!(collapsed, Some(ScanStop::ScanLimit { .. }));
293 if !scan_limited && frontier_advanced(progress_frontier, floor_pos, direction) {
294 yield Watermarked::Watermark(floor_pos);
295 progress_frontier = Some(floor_pos);
296 }
297 if let Some(stop) = collapsed {
298 Err(stop)?;
299 }
300
301 let lagging: Vec<usize> = active
305 .iter()
306 .copied()
307 .filter(|&i| targets[i].is_some())
308 .collect();
309 for i in 0..leaf_count {
313 if targets[i].is_none() {
314 drained[i] = 0;
315 }
316 }
317
318 let eval_bucket = active
322 .iter()
323 .filter(|&&i| targets[i].is_none())
324 .filter_map(|&i| match class[i] {
325 Some(LeafHead::Bucket(bucket)) => Some(bucket),
326 _ => None,
327 })
328 .reduce(|a, b| bound_in_direction(a, b, direction))
329 .expect("at least one active leaf has a ready bucket");
330 let lagging_target = lagging
335 .iter()
336 .filter_map(|&i| targets[i])
337 .reduce(|a, b| bound_in_direction(a, b, direction));
338 let evaluate = lagging_target
339 .is_none_or(|target| strictly_before(eval_bucket, target, direction));
340
341 if evaluate {
342 let (_, post) = bucket_edges(eval_bucket, bucket_size, &range, direction);
343 let mut snapshot: Vec<Option<RoaringBitmap>> =
347 (0..leaf_count).map(|_| None).collect();
348 let mut on_floor = vec![false; leaf_count];
349 for i in 0..leaf_count {
350 if !unreferenced[i]
351 && targets[i].is_none()
352 && matches!(class[i], Some(LeafHead::Bucket(b)) if b == eval_bucket)
353 {
354 on_floor[i] = true;
355 front[i] = advance_in_direction(front[i], post, direction);
356 snapshot[i] = match Pin::new(&mut leaves[i]).next().await {
357 Some(Ok((_, bitmap))) => Some(bitmap),
358 _ => None,
359 };
360 }
361 }
362 let mut remaining_refs = count_on_floor_refs(&terms, &on_floor);
365 let mut result: Option<RoaringBitmap> = None;
366 for term in &terms {
367 if term.unsatisfiable {
368 continue;
369 }
370 let includes = term
371 .includes
372 .iter()
373 .map(|&i| {
374 take_snapshot_bitmap(
375 &mut snapshot,
376 &mut remaining_refs,
377 &on_floor,
378 i,
379 )
380 })
381 .collect();
382 let excludes = term
383 .excludes
384 .iter()
385 .map(|&i| {
386 take_snapshot_bitmap(
387 &mut snapshot,
388 &mut remaining_refs,
389 &on_floor,
390 i,
391 )
392 })
393 .collect();
394 if let Some(bitmap) = eval_term_at_bucket(includes, excludes) {
395 result = Some(match result {
396 None => bitmap,
397 Some(acc) => acc | bitmap,
398 });
399 }
400 }
401 if let Some(bitmap) = result {
402 yield Watermarked::Item((eval_bucket, bitmap));
403 }
404 if frontier_advanced(progress_frontier, post, direction) {
405 yield Watermarked::Watermark(post);
406 progress_frontier = Some(post);
407 }
408 }
409
410 let catch_up = leaves.iter_mut().zip_debug_eq(drained.iter_mut())
414 .enumerate()
415 .filter_map(|(i, (leaf, drained))| {
416 let target = targets[i]?;
417 let source = source.clone();
418 let key = unique_keys[i].clone();
419 let range = range.clone();
420 let budget = budget.clone();
421 Some(async move {
422 loop {
426 let dead_row = matches!(
427 Pin::new(&mut *leaf).peek().await,
428 Some(Ok((bucket, _)))
429 if strictly_before(*bucket, target, direction)
430 );
431 if !dead_row {
432 break;
433 }
434 if policy
438 .drain_probe_rows
439 .is_some_and(|probe| *drained >= u64::from(probe.get()))
440 {
441 budget.note_discarded();
445 let narrowed = match direction {
448 ScanDirection::Ascending => {
449 target.saturating_mul(bucket_size).max(range.start)
450 ..range.end
451 }
452 ScanDirection::Descending => {
453 range.start
454 ..target
455 .saturating_add(1)
456 .saturating_mul(bucket_size)
457 .min(range.end)
458 }
459 };
460 let raw = source.scan_bucket_stream(key, narrowed, direction);
461 *leaf = budgeted_bucket_stream(raw, budget.clone(), false)
465 .boxed()
466 .peekable();
467 budget.note_seek();
468 break;
469 }
470 let discarded = Pin::new(&mut *leaf).next().await;
473 debug_assert!(matches!(discarded, Some(Ok(_))));
474 *drained += 1;
475 budget.note_discarded();
476 }
477 })
478 });
479 let seeks_before_catch_up = budget.seeks();
480 futures::future::join_all(catch_up).await;
481 if budget.seeks() > seeks_before_catch_up {
483 budget.note_reopen_step();
484 }
485 }
486 }
487 .boxed()
488}
489
490fn budgeted_bucket_stream<S>(
495 inner: S,
496 budget: BitmapScanBudget,
497 reserve_first: bool,
498) -> impl Stream<Item = BucketItem> + Send + 'static
499where
500 S: Stream<Item = BucketItem> + Send + 'static,
501{
502 async_stream::try_stream! {
503 futures::pin_mut!(inner);
504 let mut first = reserve_first;
505 while let Some(item) = inner.next().await {
506 let item = item?;
507 if first {
508 budget.take_first();
509 first = false;
510 } else if !budget.try_take() {
511 Err(LeafStop::BudgetExhausted)?;
512 }
513 yield item;
514 }
515 }
516}
517
518pub fn flatten_watermarked_buckets<S, E>(
522 stream: S,
523 range: Range<u64>,
524 bucket_size: u64,
525 direction: ScanDirection,
526) -> impl Stream<Item = Result<Watermarked<u64>, E>> + Send + 'static
527where
528 S: Stream<Item = Result<Watermarked<(u64, RoaringBitmap)>, E>> + Send + 'static,
529 E: Send + 'static,
530{
531 async_stream::try_stream! {
532 if range.is_empty() {
533 return;
534 }
535 let start_bucket = range.start / bucket_size;
536 let end_bucket = (range.end - 1) / bucket_size;
537 futures::pin_mut!(stream);
538 while let Some(item) = stream.next().await {
539 match item? {
540 Watermarked::Watermark(p) => yield Watermarked::Watermark(p),
541 Watermarked::Item((bucket_id, bitmap)) => {
542 let bucket_start = bucket_id * bucket_size;
543 let is_first = bucket_id == start_bucket;
544 let is_last = bucket_id == end_bucket;
545 let lo = if is_first {
546 (range.start - bucket_start) as u32
547 } else {
548 0
549 };
550 let hi = if is_last {
551 ((range.end - bucket_start).min(bucket_size)) as u32
552 } else {
553 bucket_size as u32
554 };
555
556 if direction.is_ascending() {
557 for bit in bitmap.iter() {
558 if bit >= lo && bit < hi {
559 yield Watermarked::Item(bucket_start + bit as u64);
560 }
561 }
562 } else {
563 for bit in bitmap.iter().rev() {
564 if bit >= lo && bit < hi {
565 yield Watermarked::Item(bucket_start + bit as u64);
566 }
567 }
568 }
569 }
570 }
571 }
572 }
573}
574
575#[cfg(test)]
576mod tests {
577 use std::collections::BTreeMap;
578 use std::sync::Arc;
579
580 use futures::TryStreamExt;
581 use futures::stream;
582
583 use super::*;
584 use crate::bitmap_query::BitmapLiteral;
585 use crate::bitmap_query::BitmapTerm;
586 use crate::bitmap_query::BucketStream;
587 use crate::bitmap_query::test_utils::EVERY_WATERMARK;
588
589 const BUCKET_SIZE: u64 = 100_000;
590 type TestBuckets = BTreeMap<Vec<u8>, Vec<(u64, Vec<u32>)>>;
591
592 #[derive(Clone)]
593 struct TestBucketSource {
594 buckets: Arc<TestBuckets>,
595 }
596
597 impl BitmapBucketSource for TestBucketSource {
598 fn scan_bucket_stream(
599 &self,
600 dimension_key: Vec<u8>,
601 range: Range<u64>,
602 direction: ScanDirection,
603 ) -> BucketStream {
604 let mut buckets = self
605 .buckets
606 .get(&dimension_key)
607 .cloned()
608 .unwrap_or_default();
609 if range.is_empty() {
610 buckets.clear();
611 } else {
612 let first_bucket = range.start / BUCKET_SIZE;
613 let last_bucket = (range.end - 1) / BUCKET_SIZE;
614 buckets.retain(|(bucket, _)| first_bucket <= *bucket && *bucket <= last_bucket);
615 }
616 if matches!(direction, ScanDirection::Descending) {
617 buckets.reverse();
618 }
619 make_bucket_stream(
620 buckets
621 .iter()
622 .map(|(bucket_id, bits)| (*bucket_id, bits.as_slice()))
623 .collect(),
624 )
625 }
626 }
627
628 fn make_bitmap(bits: &[u32]) -> RoaringBitmap {
629 let mut bm = RoaringBitmap::new();
630 for &b in bits {
631 bm.insert(b);
632 }
633 bm
634 }
635
636 fn make_bucket_stream(items: Vec<(u64, &[u32])>) -> BucketStream {
637 let items: Vec<BucketItem> = items
638 .into_iter()
639 .map(|(bid, bits)| Ok((bid, make_bitmap(bits))))
640 .collect();
641 stream::iter(items).boxed()
642 }
643
644 async fn drain_marked(
647 stream: BoxStream<'static, Result<Watermarked<u64>, ScanStop>>,
648 ) -> Result<(Vec<u64>, Vec<u64>), ScanStop> {
649 let all: Vec<Watermarked<u64>> = stream.try_collect().await?;
650 let mut items = Vec::new();
651 let mut watermarks = Vec::new();
652 for m in all {
653 match m {
654 Watermarked::Item(v) => items.push(v),
655 Watermarked::Watermark(f) => watermarks.push(f),
656 }
657 }
658 Ok((items, watermarks))
659 }
660
661 fn test_key(value: &[u8]) -> Vec<u8> {
662 crate::dimensions::encode_dimension_key(crate::dimensions::IndexDimension::Sender, value)
663 }
664
665 fn include(value: &[u8]) -> BitmapLiteral {
666 BitmapLiteral::include(test_key(value)).unwrap()
667 }
668
669 fn exclude(value: &[u8]) -> BitmapLiteral {
670 BitmapLiteral::exclude(test_key(value)).unwrap()
671 }
672
673 #[tokio::test]
679 async fn nested_term_starvation_bundles_frontier_in_scan_limit() {
680 let source = TestBucketSource {
681 buckets: Arc::new(BTreeMap::from([
682 (
683 test_key(b"a"),
684 vec![
685 (0, vec![1]),
686 (1, vec![1]),
687 (2, vec![1]),
688 (3, vec![1]),
689 (4, vec![1]),
690 ],
691 ),
692 (test_key(b"b"), vec![(50, vec![1])]),
693 (test_key(b"c"), vec![(40, vec![7])]),
694 ])),
695 };
696 let query = BitmapQuery::new(vec![
697 BitmapTerm::new(vec![include(b"a"), include(b"b")]).unwrap(),
698 BitmapTerm::new(vec![include(b"c")]).unwrap(),
699 ])
700 .unwrap();
701
702 let stream = eval_bitmap_query_stream(
704 source,
705 query,
706 0..(60 * BUCKET_SIZE),
707 BUCKET_SIZE,
708 ScanDirection::Ascending,
709 3,
710 SkipPolicy::DRAIN_ONLY,
711 EVERY_WATERMARK,
712 |_| {},
713 );
714 let all: Vec<Result<Watermarked<u64>, ScanStop>> = stream.collect().await;
715
716 let items: Vec<u64> = all
717 .iter()
718 .filter_map(|result| match result {
719 Ok(Watermarked::Item(item)) => Some(*item),
720 _ => None,
721 })
722 .collect();
723 assert_eq!(items, vec![40 * BUCKET_SIZE + 7]);
724
725 let watermarks: Vec<u64> = all
726 .iter()
727 .filter_map(|r| match r {
728 Ok(Watermarked::Watermark(position)) => Some(*position),
729 _ => None,
730 })
731 .collect();
732 assert_eq!(watermarks, vec![40 * BUCKET_SIZE, 41 * BUCKET_SIZE]);
733 assert_eq!(
734 all.len(),
735 items.len() + watermarks.len() + 1,
736 "the stopping round must add only its terminal, not another beacon"
737 );
738 let err = all
739 .last()
740 .expect("non-empty")
741 .as_ref()
742 .expect_err("scan must terminate with an error");
743 match err {
744 ScanStop::ScanLimit { scan_frontier } => assert_eq!(
745 *scan_frontier,
746 50 * BUCKET_SIZE,
747 "terminal must carry the logically advanced merged floor"
748 ),
749 other => panic!("expected ScanLimit, got {other:?}"),
750 }
751 }
752
753 #[tokio::test]
764 async fn flush_on_error_delivers_below_floor_sibling_result() {
765 let source = TestBucketSource {
766 buckets: Arc::new(BTreeMap::from([
767 (
768 test_key(b"a"),
769 vec![
770 (0, vec![1]),
771 (1, vec![1]),
772 (2, vec![1]),
773 (3, vec![1]),
774 (4, vec![1]),
775 ],
776 ),
777 (test_key(b"b"), vec![(50, vec![1])]),
778 (test_key(b"c"), vec![(0, vec![7])]),
779 ])),
780 };
781 let query = BitmapQuery::new(vec![
782 BitmapTerm::new(vec![include(b"a"), include(b"b")]).unwrap(),
783 BitmapTerm::new(vec![include(b"c")]).unwrap(),
784 ])
785 .unwrap();
786
787 let stream = eval_bitmap_query_stream(
788 source,
789 query,
790 0..(60 * BUCKET_SIZE),
791 BUCKET_SIZE,
792 ScanDirection::Ascending,
793 3,
794 SkipPolicy::DRAIN_ONLY,
795 EVERY_WATERMARK,
796 |_| {},
797 );
798 let all: Vec<Result<Watermarked<u64>, ScanStop>> = stream.collect().await;
799
800 let items: Vec<u64> = all
802 .iter()
803 .filter_map(|r| match r {
804 Ok(Watermarked::Item(v)) => Some(*v),
805 _ => None,
806 })
807 .collect();
808 assert_eq!(items, vec![7], "c's below-floor match must be delivered");
809
810 let watermarks: Vec<u64> = all
811 .iter()
812 .filter_map(|r| match r {
813 Ok(Watermarked::Watermark(position)) => Some(*position),
814 _ => None,
815 })
816 .collect();
817 assert_eq!(watermarks, vec![BUCKET_SIZE]);
818 assert_eq!(
819 all.len(),
820 items.len() + watermarks.len() + 1,
821 "the stopping round must add only its terminal, not another beacon"
822 );
823
824 let err = all
825 .last()
826 .expect("non-empty")
827 .as_ref()
828 .expect_err("scan must terminate with an error");
829 let scan_frontier = match err {
830 ScanStop::ScanLimit { scan_frontier } => *scan_frontier,
831 other => panic!("expected ScanLimit, got {other:?}"),
832 };
833 assert_eq!(
834 scan_frontier,
835 50 * BUCKET_SIZE,
836 "terminal frontier must include the logical gap jump"
837 );
838 assert!(
839 items.iter().all(|&i| i < scan_frontier),
840 "delivered items must be below the resume frontier"
841 );
842 }
843
844 #[test]
845 fn bitmap_query_validation_rejects_empty_shapes() {
846 assert!(BitmapQuery::new(Vec::new()).is_err());
847 assert!(BitmapLiteral::include(Vec::new()).is_err());
848 assert!(
849 BitmapLiteral::include(vec![crate::dimensions::IndexDimension::Sender.tag_byte()])
850 .is_err()
851 );
852 assert!(BitmapLiteral::include(vec![0xff, 0x00]).is_err());
853 assert!(BitmapTerm::new(vec![exclude(b"neg")]).is_err());
854 }
855
856 #[tokio::test]
861 async fn shared_include_across_terms_scans_dimension_once() {
862 use crate::bitmap_query::test_utils::CountingBucketSource;
863
864 let source = CountingBucketSource::new(BTreeMap::from([
865 (test_key(b"a"), vec![(0, vec![1, 2, 3])]),
866 (test_key(b"b"), vec![(0, vec![1])]),
867 (test_key(b"c"), vec![(0, vec![2])]),
868 ]));
869 let query = BitmapQuery::new(vec![
870 BitmapTerm::new(vec![include(b"a"), include(b"b")]).unwrap(),
871 BitmapTerm::new(vec![include(b"a"), include(b"c")]).unwrap(),
872 ])
873 .unwrap();
874
875 let stream = eval_bitmap_query_stream(
876 source.clone(),
877 query,
878 0..200_000,
879 BUCKET_SIZE,
880 ScanDirection::Ascending,
881 u64::MAX,
882 SkipPolicy::DRAIN_ONLY,
883 EVERY_WATERMARK,
884 |_| {},
885 );
886 let (items, _watermarks) = drain_marked(stream).await.unwrap();
887
888 assert_eq!(items, vec![1, 2]);
891 assert_eq!(source.scan_count(&test_key(b"a")), 1);
893 assert_eq!(source.scan_count(&test_key(b"b")), 1);
894 assert_eq!(source.scan_count(&test_key(b"c")), 1);
895 }
896
897 #[tokio::test]
901 async fn shared_key_across_include_and_exclude_terms_scans_once() {
902 use crate::bitmap_query::test_utils::CountingBucketSource;
903
904 let source = CountingBucketSource::new(BTreeMap::from([
905 (test_key(b"a"), vec![(0, vec![1, 2])]),
906 (test_key(b"b"), vec![(0, vec![1, 2, 3])]),
907 ]));
908 let query = BitmapQuery::new(vec![
912 BitmapTerm::new(vec![include(b"b"), exclude(b"a")]).unwrap(),
913 BitmapTerm::new(vec![include(b"b"), include(b"a")]).unwrap(),
914 ])
915 .unwrap();
916
917 let stream = eval_bitmap_query_stream(
918 source.clone(),
919 query,
920 0..100_000,
921 BUCKET_SIZE,
922 ScanDirection::Ascending,
923 u64::MAX,
924 SkipPolicy::DRAIN_ONLY,
925 EVERY_WATERMARK,
926 |_| {},
927 );
928 let (items, _watermarks) = drain_marked(stream).await.unwrap();
929
930 assert_eq!(items, vec![1, 2, 3]);
931 assert_eq!(source.scan_count(&test_key(b"a")), 1);
932 assert_eq!(source.scan_count(&test_key(b"b")), 1);
933 }
934
935 #[tokio::test]
936 async fn eval_bitmap_query_stream_uses_backend_source() {
937 let source = TestBucketSource {
938 buckets: Arc::new(BTreeMap::from([
939 (test_key(b"a"), vec![(0, vec![1, 2, 3]), (1, vec![5])]),
940 (test_key(b"b"), vec![(0, vec![2, 3]), (1, vec![5])]),
941 (test_key(b"c"), vec![(0, vec![3])]),
942 ])),
943 };
944 let query = BitmapQuery::new(vec![
945 BitmapTerm::new(vec![include(b"a"), include(b"b"), exclude(b"c")]).unwrap(),
946 ])
947 .unwrap();
948
949 let stream = eval_bitmap_query_stream(
950 source,
951 query,
952 0..200_000,
953 BUCKET_SIZE,
954 ScanDirection::Ascending,
955 u64::MAX,
956 SkipPolicy::DRAIN_ONLY,
957 EVERY_WATERMARK,
958 |_| {},
959 );
960 let (items, _watermarks) = drain_marked(stream).await.unwrap();
961
962 assert_eq!(items, vec![2, BUCKET_SIZE + 5]);
963 }
964
965 #[tokio::test]
966 async fn eval_bitmap_query_stream_descending() {
967 let source = TestBucketSource {
968 buckets: Arc::new(BTreeMap::from([
969 (test_key(b"a"), vec![(0, vec![1, 2, 3]), (1, vec![5])]),
970 (test_key(b"b"), vec![(0, vec![2, 3]), (1, vec![5])]),
971 (test_key(b"c"), vec![(0, vec![3])]),
972 ])),
973 };
974 let query = BitmapQuery::new(vec![
975 BitmapTerm::new(vec![include(b"a"), include(b"b"), exclude(b"c")]).unwrap(),
976 ])
977 .unwrap();
978
979 let stream = eval_bitmap_query_stream(
980 source,
981 query,
982 0..200_000,
983 BUCKET_SIZE,
984 ScanDirection::Descending,
985 u64::MAX,
986 SkipPolicy::DRAIN_ONLY,
987 EVERY_WATERMARK,
988 |_| {},
989 );
990 let (items, _watermarks) = drain_marked(stream).await.unwrap();
991
992 assert_eq!(items, vec![BUCKET_SIZE + 5, 2]);
993 }
994
995 #[tokio::test]
998 async fn flatten_watermarked_buckets_ascending() {
999 let range = 50u64..(2 * BUCKET_SIZE + 50_001);
1000 let marked_buckets = stream::iter(vec![
1001 Ok::<_, LeafStop>(Watermarked::Item((
1002 0,
1003 make_bitmap(&[10, 50, (BUCKET_SIZE - 1) as u32]),
1004 ))),
1005 Ok(Watermarked::Watermark(BUCKET_SIZE)),
1006 Ok(Watermarked::Item((
1007 1,
1008 make_bitmap(&[0, (BUCKET_SIZE - 1) as u32]),
1009 ))),
1010 Ok(Watermarked::Watermark(2 * BUCKET_SIZE)),
1011 Ok(Watermarked::Item((2, make_bitmap(&[0, 50_000, 50_001])))),
1012 Ok(Watermarked::Watermark(2 * BUCKET_SIZE + 50_001)),
1013 ]);
1014 let out: Vec<Watermarked<u64>> = flatten_watermarked_buckets(
1015 marked_buckets,
1016 range,
1017 BUCKET_SIZE,
1018 ScanDirection::Ascending,
1019 )
1020 .try_collect()
1021 .await
1022 .unwrap();
1023
1024 assert_eq!(
1025 out,
1026 vec![
1027 Watermarked::Item(50),
1028 Watermarked::Item(BUCKET_SIZE - 1),
1029 Watermarked::Watermark(BUCKET_SIZE),
1030 Watermarked::Item(BUCKET_SIZE),
1031 Watermarked::Item(2 * BUCKET_SIZE - 1),
1032 Watermarked::Watermark(2 * BUCKET_SIZE),
1033 Watermarked::Item(2 * BUCKET_SIZE),
1034 Watermarked::Item(2 * BUCKET_SIZE + 50_000),
1035 Watermarked::Watermark(2 * BUCKET_SIZE + 50_001),
1036 ],
1037 );
1038 }
1039
1040 #[tokio::test]
1041 async fn scan_budget_below_unique_leaf_count_yields_misconfig_error() {
1042 let source = TestBucketSource {
1049 buckets: Arc::new(BTreeMap::from([(
1050 test_key(b"a"),
1051 vec![(0, vec![1, 2]), (1, vec![3])],
1052 )])),
1053 };
1054 let query = BitmapQuery::new(vec![BitmapTerm::new(vec![include(b"a")]).unwrap()]).unwrap();
1055
1056 let metrics = std::sync::Arc::new(std::sync::Mutex::new(None));
1057 let metrics_sink = metrics.clone();
1058 let stream = eval_bitmap_query_stream(
1059 source,
1060 query,
1061 0..200_000,
1062 BUCKET_SIZE,
1063 ScanDirection::Ascending,
1064 0,
1065 SkipPolicy::DRAIN_ONLY,
1066 EVERY_WATERMARK,
1067 move |m| *metrics_sink.lock().unwrap() = Some(m),
1068 );
1069 let err = drain_marked(stream).await.unwrap_err();
1070
1071 assert!(
1072 matches!(err, ScanStop::Fault(_)),
1073 "must be a Fault (cursorless), never a clean ScanLimit end"
1074 );
1075 assert!(
1076 err.to_string().contains("insufficient for"),
1077 "expected misconfig error, got {err:?}"
1078 );
1079 assert!(
1082 metrics.lock().unwrap().is_none(),
1083 "misconfig early-out should not emit scan metrics"
1084 );
1085 }
1086
1087 #[tokio::test]
1088 async fn scan_budget_shared_across_dimensions() {
1089 let source = TestBucketSource {
1093 buckets: Arc::new(BTreeMap::from([
1094 (
1095 test_key(b"a"),
1096 vec![(0, vec![1]), (1, vec![2]), (2, vec![3])],
1097 ),
1098 (
1099 test_key(b"b"),
1100 vec![(0, vec![1]), (1, vec![2]), (2, vec![3])],
1101 ),
1102 (
1103 test_key(b"c"),
1104 vec![(0, vec![1]), (1, vec![2]), (2, vec![3])],
1105 ),
1106 ])),
1107 };
1108 let query = BitmapQuery::new(vec![
1109 BitmapTerm::new(vec![include(b"a"), include(b"b"), include(b"c")]).unwrap(),
1110 ])
1111 .unwrap();
1112
1113 let metrics = std::sync::Arc::new(std::sync::Mutex::new(None));
1114 let metrics_sink = metrics.clone();
1115 let stream = eval_bitmap_query_stream(
1116 source,
1117 query,
1118 0..300_000,
1119 BUCKET_SIZE,
1120 ScanDirection::Ascending,
1121 4,
1122 SkipPolicy::DRAIN_ONLY,
1123 EVERY_WATERMARK,
1124 move |m| *metrics_sink.lock().unwrap() = Some(m),
1125 );
1126 let err = drain_marked(stream).await.unwrap_err();
1127
1128 assert!(
1129 matches!(err, ScanStop::ScanLimit { .. }),
1130 "expected ScanLimit, got {err:?}"
1131 );
1132 assert_eq!(metrics.lock().unwrap().unwrap().buckets_evaluated, 4);
1135 }
1136
1137 #[tokio::test]
1142 async fn scan_budget_exclude_side_exhaustion_does_not_leak_includes() {
1143 let source = TestBucketSource {
1144 buckets: Arc::new(BTreeMap::from([
1145 (
1146 test_key(b"inc"),
1147 vec![(0, vec![1]), (1, vec![2]), (2, vec![3])],
1148 ),
1149 (
1150 test_key(b"exc"),
1151 vec![(0, vec![1]), (1, vec![2]), (2, vec![3])],
1152 ),
1153 ])),
1154 };
1155 let query = BitmapQuery::new(vec![
1156 BitmapTerm::new(vec![include(b"inc"), exclude(b"exc")]).unwrap(),
1157 ])
1158 .unwrap();
1159
1160 let stream = eval_bitmap_query_stream(
1167 source,
1168 query,
1169 0..300_000,
1170 BUCKET_SIZE,
1171 ScanDirection::Ascending,
1172 2,
1173 SkipPolicy::DRAIN_ONLY,
1174 EVERY_WATERMARK,
1175 |_| {},
1176 );
1177 let result = drain_marked(stream).await;
1178
1179 let err = result.expect_err("must surface scan-limit, not silently emit includes");
1181 assert!(
1182 matches!(err, ScanStop::ScanLimit { .. }),
1183 "expected ScanLimit, got {err:?}"
1184 );
1185 }
1186
1187 #[tokio::test]
1192 async fn sparse_intersect_bundles_frontier_in_scan_limit() {
1193 let source = TestBucketSource {
1194 buckets: Arc::new(BTreeMap::from([
1195 (
1196 test_key(b"a"),
1197 vec![
1198 (0, vec![1]),
1199 (1, vec![1]),
1200 (2, vec![1]),
1201 (3, vec![1]),
1202 (4, vec![1]),
1203 ],
1204 ),
1205 (test_key(b"b"), vec![(100, vec![1])]),
1206 ])),
1207 };
1208 let query = BitmapQuery::new(vec![
1209 BitmapTerm::new(vec![include(b"a"), include(b"b")]).unwrap(),
1210 ])
1211 .unwrap();
1212
1213 let stream = eval_bitmap_query_stream(
1214 source,
1215 query,
1216 0..(110 * BUCKET_SIZE),
1217 BUCKET_SIZE,
1218 ScanDirection::Ascending,
1219 4,
1220 SkipPolicy::DRAIN_ONLY,
1221 EVERY_WATERMARK,
1222 |_| {},
1223 );
1224 let all: Vec<Result<Watermarked<u64>, ScanStop>> = stream.collect().await;
1225
1226 assert!(
1227 all.iter().all(|r| !matches!(r, Ok(Watermarked::Item(_)))),
1228 "disjoint intersect must not emit items"
1229 );
1230 let err = all
1231 .last()
1232 .expect("non-empty")
1233 .as_ref()
1234 .expect_err("scan must terminate with an error");
1235 let scan_frontier = match err {
1236 ScanStop::ScanLimit { scan_frontier } => *scan_frontier,
1237 other => panic!("expected ScanLimit, got {other:?}"),
1238 };
1239 assert_eq!(
1240 scan_frontier,
1241 100 * BUCKET_SIZE,
1242 "terminal must carry the logically advanced merged floor"
1243 );
1244 let watermarks: Vec<u64> = all
1245 .iter()
1246 .filter_map(|r| match r {
1247 Ok(Watermarked::Watermark(position)) => Some(*position),
1248 _ => None,
1249 })
1250 .collect();
1251 assert_eq!(
1252 watermarks,
1253 vec![100 * BUCKET_SIZE],
1254 "the stopping round must not append another frontier beacon"
1255 );
1256 assert_eq!(all.len(), watermarks.len() + 1);
1257 }
1258
1259 #[tokio::test]
1260 async fn eval_emits_watermarks_at_bucket_boundaries_ascending() {
1261 let source = TestBucketSource {
1262 buckets: Arc::new(BTreeMap::from([(
1263 test_key(b"a"),
1264 vec![(0, vec![1]), (3, vec![2]), (7, vec![3])],
1265 )])),
1266 };
1267 let query = BitmapQuery::new(vec![BitmapTerm::new(vec![include(b"a")]).unwrap()]).unwrap();
1268
1269 let stream = eval_bitmap_query_stream(
1270 source,
1271 query,
1272 0..(8 * BUCKET_SIZE),
1273 BUCKET_SIZE,
1274 ScanDirection::Ascending,
1275 u64::MAX,
1276 SkipPolicy::DRAIN_ONLY,
1277 EVERY_WATERMARK,
1278 |_| {},
1279 );
1280 let (items, watermarks) = drain_marked(stream).await.unwrap();
1281
1282 assert_eq!(items, vec![1, 3 * BUCKET_SIZE + 2, 7 * BUCKET_SIZE + 3]);
1284 assert_eq!(
1288 watermarks,
1289 vec![
1290 BUCKET_SIZE,
1291 3 * BUCKET_SIZE,
1292 4 * BUCKET_SIZE,
1293 7 * BUCKET_SIZE,
1294 8 * BUCKET_SIZE,
1295 ]
1296 );
1297 }
1298
1299 #[tokio::test]
1300 async fn descending_exclusive_upper_bound_is_not_progress() {
1301 let source = TestBucketSource {
1302 buckets: Arc::new(BTreeMap::from([(
1303 test_key(b"a"),
1304 vec![(0, vec![1]), (3, vec![2]), (7, vec![3])],
1305 )])),
1306 };
1307 let query = BitmapQuery::new(vec![BitmapTerm::new(vec![include(b"a")]).unwrap()]).unwrap();
1308
1309 let stream = eval_bitmap_query_stream(
1310 source,
1311 query,
1312 0..(8 * BUCKET_SIZE),
1313 BUCKET_SIZE,
1314 ScanDirection::Descending,
1315 u64::MAX,
1316 SkipPolicy::DRAIN_ONLY,
1317 EVERY_WATERMARK,
1318 |_| {},
1319 );
1320 let (items, watermarks) = drain_marked(stream).await.unwrap();
1321
1322 assert_eq!(items, vec![7 * BUCKET_SIZE + 3, 3 * BUCKET_SIZE + 2, 1]);
1323 assert_eq!(
1326 watermarks,
1327 vec![
1328 7 * BUCKET_SIZE,
1329 4 * BUCKET_SIZE,
1330 3 * BUCKET_SIZE,
1331 BUCKET_SIZE,
1332 0,
1333 ]
1334 );
1335 }
1336
1337 #[tokio::test]
1338 async fn natural_completion_omits_terminus_but_retains_earned_progress() {
1339 let source = TestBucketSource {
1344 buckets: Arc::new(BTreeMap::from([
1345 (
1346 test_key(b"a"),
1347 vec![(0, vec![1]), (2, vec![3]), (4, vec![5])],
1348 ),
1349 (
1350 test_key(b"b"),
1351 vec![(1, vec![1]), (3, vec![3]), (5, vec![5])],
1352 ),
1353 ])),
1354 };
1355 let query = BitmapQuery::new(vec![
1356 BitmapTerm::new(vec![include(b"a"), include(b"b")]).unwrap(),
1357 ])
1358 .unwrap();
1359
1360 let stream = eval_bitmap_query_stream(
1361 source,
1362 query,
1363 0..(6 * BUCKET_SIZE),
1364 BUCKET_SIZE,
1365 ScanDirection::Ascending,
1366 u64::MAX,
1367 SkipPolicy::DRAIN_ONLY,
1368 EVERY_WATERMARK,
1369 |_| {},
1370 );
1371 let (items, watermarks) = drain_marked(stream).await.unwrap();
1372
1373 assert!(items.is_empty(), "disjoint intersect must not emit items");
1374 assert_eq!(
1375 watermarks,
1376 vec![
1377 BUCKET_SIZE,
1378 2 * BUCKET_SIZE,
1379 3 * BUCKET_SIZE,
1380 4 * BUCKET_SIZE,
1381 5 * BUCKET_SIZE,
1382 ],
1383 );
1384 }
1385
1386 #[tokio::test]
1387 async fn eval_watermark_ordering_invariant_item_then_watermark() {
1388 let source = TestBucketSource {
1393 buckets: Arc::new(BTreeMap::from([(
1394 test_key(b"a"),
1395 vec![(0, vec![10, 20, 30]), (1, vec![40, 50])],
1396 )])),
1397 };
1398 let query = BitmapQuery::new(vec![BitmapTerm::new(vec![include(b"a")]).unwrap()]).unwrap();
1399
1400 let stream = eval_bitmap_query_stream(
1401 source,
1402 query,
1403 0..(2 * BUCKET_SIZE),
1404 BUCKET_SIZE,
1405 ScanDirection::Ascending,
1406 u64::MAX,
1407 SkipPolicy::DRAIN_ONLY,
1408 EVERY_WATERMARK,
1409 |_| {},
1410 );
1411 let all: Vec<Watermarked<u64>> = stream.try_collect().await.unwrap();
1412
1413 assert_eq!(
1417 all,
1418 vec![
1419 Watermarked::Item(10),
1420 Watermarked::Item(20),
1421 Watermarked::Item(30),
1422 Watermarked::Watermark(BUCKET_SIZE),
1423 Watermarked::Item(BUCKET_SIZE + 40),
1424 Watermarked::Item(BUCKET_SIZE + 50),
1425 Watermarked::Watermark(2 * BUCKET_SIZE),
1426 ],
1427 );
1428 }
1429
1430 #[tokio::test]
1431 async fn flatten_watermarked_buckets_descending() {
1432 let range = 50u64..(2 * BUCKET_SIZE + 50_001);
1433 let marked_buckets = stream::iter(vec![
1434 Ok::<_, LeafStop>(Watermarked::Item((2, make_bitmap(&[0, 50_000, 50_001])))),
1435 Ok(Watermarked::Watermark(2 * BUCKET_SIZE)),
1436 Ok(Watermarked::Item((
1437 1,
1438 make_bitmap(&[0, (BUCKET_SIZE - 1) as u32]),
1439 ))),
1440 Ok(Watermarked::Watermark(BUCKET_SIZE)),
1441 Ok(Watermarked::Item((
1442 0,
1443 make_bitmap(&[10, 50, (BUCKET_SIZE - 1) as u32]),
1444 ))),
1445 Ok(Watermarked::Watermark(50)),
1446 ]);
1447 let out: Vec<Watermarked<u64>> = flatten_watermarked_buckets(
1448 marked_buckets,
1449 range,
1450 BUCKET_SIZE,
1451 ScanDirection::Descending,
1452 )
1453 .try_collect()
1454 .await
1455 .unwrap();
1456
1457 assert_eq!(
1458 out,
1459 vec![
1460 Watermarked::Item(2 * BUCKET_SIZE + 50_000),
1461 Watermarked::Item(2 * BUCKET_SIZE),
1462 Watermarked::Watermark(2 * BUCKET_SIZE),
1463 Watermarked::Item(2 * BUCKET_SIZE - 1),
1464 Watermarked::Item(BUCKET_SIZE),
1465 Watermarked::Watermark(BUCKET_SIZE),
1466 Watermarked::Item(BUCKET_SIZE - 1),
1467 Watermarked::Item(50),
1468 Watermarked::Watermark(50),
1469 ],
1470 );
1471 }
1472
1473 #[test]
1476 fn collapse_single_budget_stop_binds_frontier() {
1477 assert!(matches!(
1478 collapse(vec![LeafStop::BudgetExhausted], 17 * BUCKET_SIZE),
1479 ScanStop::ScanLimit { scan_frontier } if scan_frontier == 17 * BUCKET_SIZE
1480 ));
1481 }
1482
1483 #[test]
1486 fn collapse_all_budget_stops_bind_frontier() {
1487 assert!(matches!(
1488 collapse(
1489 vec![LeafStop::BudgetExhausted, LeafStop::BudgetExhausted],
1490 23 * BUCKET_SIZE
1491 ),
1492 ScanStop::ScanLimit { scan_frontier } if scan_frontier == 23 * BUCKET_SIZE
1493 ));
1494 }
1495
1496 fn gap_probe_policy() -> SkipPolicy {
1497 SkipPolicy {
1498 drain_probe_rows: std::num::NonZeroU32::new(2),
1499 }
1500 }
1501
1502 fn skewed_and_source() -> TestBucketSource {
1503 TestBucketSource {
1504 buckets: Arc::new(BTreeMap::from([
1505 (test_key(b"a"), vec![(0, vec![1]), (50, vec![1])]),
1506 (
1507 test_key(b"b"),
1508 (0..=50).map(|bucket| (bucket, vec![1])).collect(),
1509 ),
1510 ])),
1511 }
1512 }
1513
1514 #[tokio::test]
1515 async fn skewed_and_seeks_past_dead_gap() {
1516 let query = BitmapQuery::new(vec![
1517 BitmapTerm::new(vec![include(b"a"), include(b"b")]).unwrap(),
1518 ])
1519 .unwrap();
1520 let (metrics_tx, metrics_rx) = std::sync::mpsc::channel();
1521 let stream = eval_bitmap_query_stream(
1522 skewed_and_source(),
1523 query,
1524 0..(51 * BUCKET_SIZE),
1525 BUCKET_SIZE,
1526 ScanDirection::Ascending,
1527 1_000,
1528 gap_probe_policy(),
1529 EVERY_WATERMARK,
1530 move |observed| metrics_tx.send(observed).unwrap(),
1531 );
1532 let (items, watermarks) = drain_marked(stream).await.unwrap();
1533
1534 assert_eq!(items, vec![1, 50 * BUCKET_SIZE + 1]);
1535 assert_eq!(
1536 watermarks,
1537 vec![BUCKET_SIZE, 50 * BUCKET_SIZE, 51 * BUCKET_SIZE]
1538 );
1539 assert_eq!(
1540 metrics_rx.recv().unwrap(),
1541 BitmapScanMetrics {
1542 buckets_evaluated: 7,
1543 buckets_discarded: 3,
1544 leaf_seeks: 1,
1545 }
1546 );
1547 }
1548
1549 #[tokio::test]
1550 async fn skewed_and_seeks_past_dead_gap_descending() {
1551 let query = BitmapQuery::new(vec![
1552 BitmapTerm::new(vec![include(b"a"), include(b"b")]).unwrap(),
1553 ])
1554 .unwrap();
1555 let (metrics_tx, metrics_rx) = std::sync::mpsc::channel();
1556 let stream = eval_bitmap_query_stream(
1557 skewed_and_source(),
1558 query,
1559 0..(51 * BUCKET_SIZE),
1560 BUCKET_SIZE,
1561 ScanDirection::Descending,
1562 1_000,
1563 gap_probe_policy(),
1564 EVERY_WATERMARK,
1565 move |observed| metrics_tx.send(observed).unwrap(),
1566 );
1567 let (items, watermarks) = drain_marked(stream).await.unwrap();
1568
1569 assert_eq!(items, vec![50 * BUCKET_SIZE + 1, 1]);
1570 assert_eq!(watermarks, vec![50 * BUCKET_SIZE, BUCKET_SIZE, 0]);
1571 assert_eq!(
1572 metrics_rx.recv().unwrap(),
1573 BitmapScanMetrics {
1574 buckets_evaluated: 7,
1575 buckets_discarded: 3,
1576 leaf_seeks: 1,
1577 }
1578 );
1579 }
1580
1581 #[tokio::test]
1582 async fn budget_death_mid_gap_carries_jumped_frontier() {
1583 let query = BitmapQuery::new(vec![
1584 BitmapTerm::new(vec![include(b"a"), include(b"b")]).unwrap(),
1585 ])
1586 .unwrap();
1587 let first: Vec<_> = eval_bitmap_query_stream(
1588 skewed_and_source(),
1589 query.clone(),
1590 0..(51 * BUCKET_SIZE),
1591 BUCKET_SIZE,
1592 ScanDirection::Ascending,
1593 5,
1594 gap_probe_policy(),
1595 EVERY_WATERMARK,
1596 |_| {},
1597 )
1598 .collect()
1599 .await;
1600 let first_items: Vec<_> = first
1601 .iter()
1602 .filter_map(|result| match result {
1603 Ok(Watermarked::Item(item)) => Some(*item),
1604 _ => None,
1605 })
1606 .collect();
1607 assert_eq!(first_items, vec![1]);
1608 assert!(matches!(
1609 first.last(),
1610 Some(Err(ScanStop::ScanLimit { scan_frontier }))
1611 if *scan_frontier == 50 * BUCKET_SIZE
1612 ));
1613
1614 let resumed = eval_bitmap_query_stream(
1615 skewed_and_source(),
1616 query,
1617 (50 * BUCKET_SIZE)..(51 * BUCKET_SIZE),
1618 BUCKET_SIZE,
1619 ScanDirection::Ascending,
1620 100,
1621 gap_probe_policy(),
1622 EVERY_WATERMARK,
1623 |_| {},
1624 );
1625 let (resumed_items, _) = drain_marked(resumed).await.unwrap();
1626 assert_eq!(resumed_items, vec![50 * BUCKET_SIZE + 1]);
1627 }
1628
1629 #[tokio::test]
1630 async fn shared_leaf_never_skips_past_its_own_terms_candidate() {
1631 let source = TestBucketSource {
1632 buckets: Arc::new(BTreeMap::from([
1633 (test_key(b"a"), vec![(0, vec![1]), (50, vec![1])]),
1634 (test_key(b"e"), vec![(10, vec![3])]),
1635 ])),
1636 };
1637 let query = BitmapQuery::new(vec![
1638 BitmapTerm::new(vec![include(b"a"), exclude(b"e")]).unwrap(),
1639 BitmapTerm::new(vec![include(b"e")]).unwrap(),
1640 ])
1641 .unwrap();
1642 let (metrics_tx, metrics_rx) = std::sync::mpsc::channel();
1643 let stream = eval_bitmap_query_stream(
1644 source,
1645 query,
1646 0..(51 * BUCKET_SIZE),
1647 BUCKET_SIZE,
1648 ScanDirection::Ascending,
1649 100,
1650 gap_probe_policy(),
1651 EVERY_WATERMARK,
1652 move |observed| metrics_tx.send(observed).unwrap(),
1653 );
1654 let (items, _) = drain_marked(stream).await.unwrap();
1655
1656 assert_eq!(items, vec![1, 10 * BUCKET_SIZE + 3, 50 * BUCKET_SIZE + 1]);
1657 let metrics = metrics_rx.recv().expect("metrics callback ran");
1658 assert_eq!(metrics.buckets_discarded, 0);
1659 assert_eq!(metrics.leaf_seeks, 0);
1660 }
1661
1662 #[tokio::test]
1663 async fn exclude_leaf_drains_dead_rows_when_dragged() {
1664 let source = TestBucketSource {
1665 buckets: Arc::new(BTreeMap::from([
1666 (test_key(b"a"), vec![(0, vec![1]), (50, vec![1])]),
1667 (test_key(b"e"), vec![(10, vec![3])]),
1668 ])),
1669 };
1670 let query = BitmapQuery::new(vec![
1671 BitmapTerm::new(vec![include(b"a"), exclude(b"e")]).unwrap(),
1672 ])
1673 .unwrap();
1674 let (metrics_tx, metrics_rx) = std::sync::mpsc::channel();
1675 let stream = eval_bitmap_query_stream(
1676 source,
1677 query,
1678 0..(51 * BUCKET_SIZE),
1679 BUCKET_SIZE,
1680 ScanDirection::Ascending,
1681 100,
1682 gap_probe_policy(),
1683 EVERY_WATERMARK,
1684 move |observed| metrics_tx.send(observed).unwrap(),
1685 );
1686 let (items, _) = drain_marked(stream).await.unwrap();
1687
1688 assert_eq!(items, vec![1, 50 * BUCKET_SIZE + 1]);
1689 let metrics = metrics_rx.recv().expect("metrics callback ran");
1690 assert_eq!(metrics.buckets_discarded, 1);
1691 assert_eq!(metrics.leaf_seeks, 0);
1692 }
1693
1694 #[test]
1697 fn collapse_fault_outranks_budget_stop() {
1698 let collapsed = collapse(
1699 vec![
1700 LeafStop::BudgetExhausted,
1701 LeafStop::Fault(anyhow::anyhow!("storage boom")),
1702 ],
1703 7,
1704 );
1705 match collapsed {
1706 ScanStop::Fault(e) => assert!(e.to_string().contains("storage boom")),
1707 other => panic!("expected Fault to win, got {other:?}"),
1708 }
1709 }
1710
1711 #[test]
1713 fn collapse_fault_outranks_budget_stop_and_cancelled() {
1714 let collapsed = collapse(
1715 vec![
1716 LeafStop::BudgetExhausted,
1717 LeafStop::Cancelled,
1718 LeafStop::Fault(anyhow::anyhow!("storage boom")),
1719 ],
1720 11,
1721 );
1722 match collapsed {
1723 ScanStop::Fault(e) => assert!(e.to_string().contains("storage boom")),
1724 other => panic!("expected Fault to win, got {other:?}"),
1725 }
1726 }
1727
1728 #[test]
1731 fn collapse_budget_stop_outranks_cancelled() {
1732 assert!(matches!(
1733 collapse(vec![LeafStop::Cancelled, LeafStop::BudgetExhausted], 29),
1734 ScanStop::ScanLimit { scan_frontier: 29 }
1735 ));
1736 }
1737
1738 #[test]
1741 fn collapse_all_cancelled() {
1742 assert!(matches!(
1743 collapse(vec![LeafStop::Cancelled, LeafStop::Cancelled], 31),
1744 ScanStop::Cancelled
1745 ));
1746 }
1747
1748 #[test]
1751 fn collapse_combines_concurrent_faults() {
1752 let collapsed = collapse(
1753 vec![
1754 LeafStop::Fault(anyhow::anyhow!("boom one")),
1755 LeafStop::Fault(anyhow::anyhow!("boom two")),
1756 ],
1757 37,
1758 );
1759 match collapsed {
1760 ScanStop::Fault(e) => {
1761 let s = e.to_string();
1762 assert!(s.contains("boom one"), "missing first fault: {s}");
1763 assert!(s.contains("boom two"), "missing second fault: {s}");
1764 }
1765 other => panic!("expected combined Fault, got {other:?}"),
1766 }
1767 }
1768
1769 #[test]
1772 fn from_anyhow_funnels_to_leaf_fault() {
1773 match LeafStop::from(anyhow::anyhow!("leaf storage boom")) {
1774 LeafStop::Fault(e) => assert!(e.to_string().contains("leaf storage boom")),
1775 other => panic!("expected leaf Fault, got {other:?}"),
1776 }
1777 }
1778
1779 #[test]
1782 fn from_anyhow_funnels_to_scan_fault() {
1783 match ScanStop::from(anyhow::anyhow!("merged storage boom")) {
1784 ScanStop::Fault(e) => assert!(e.to_string().contains("merged storage boom")),
1785 other => panic!("expected scan Fault, got {other:?}"),
1786 }
1787 }
1788
1789 fn stream_items(all: &[Result<Watermarked<u64>, ScanStop>]) -> Vec<u64> {
1790 all.iter()
1791 .filter_map(|r| match r {
1792 Ok(Watermarked::Item(v)) => Some(*v),
1793 _ => None,
1794 })
1795 .collect()
1796 }
1797
1798 #[tokio::test]
1803 async fn absent_include_annihilates_term() {
1804 let source = TestBucketSource {
1805 buckets: Arc::new(BTreeMap::from([(test_key(b"a"), vec![(0, vec![1, 2])])])),
1806 };
1807 let query =
1808 BitmapQuery::new(vec![BitmapTerm::new(vec![include(b"ghost")]).unwrap()]).unwrap();
1809
1810 let stream = eval_bitmap_query_stream(
1811 source,
1812 query,
1813 0..(2 * BUCKET_SIZE),
1814 BUCKET_SIZE,
1815 ScanDirection::Ascending,
1816 3,
1817 SkipPolicy::DRAIN_ONLY,
1818 EVERY_WATERMARK,
1819 |_| {},
1820 );
1821 let all: Vec<Result<Watermarked<u64>, ScanStop>> = stream.collect().await;
1822 let items = stream_items(&all);
1823
1824 assert!(
1825 items.is_empty(),
1826 "absent include must annihilate: {items:?}"
1827 );
1828 }
1829
1830 #[tokio::test]
1833 async fn absent_include_annihilates_term_despite_present_include() {
1834 let source = TestBucketSource {
1835 buckets: Arc::new(BTreeMap::from([(test_key(b"a"), vec![(0, vec![1, 2])])])),
1836 };
1837 let query = BitmapQuery::new(vec![
1838 BitmapTerm::new(vec![include(b"a"), include(b"ghost")]).unwrap(),
1839 ])
1840 .unwrap();
1841
1842 let stream = eval_bitmap_query_stream(
1843 source,
1844 query,
1845 0..(2 * BUCKET_SIZE),
1846 BUCKET_SIZE,
1847 ScanDirection::Ascending,
1848 3,
1849 SkipPolicy::DRAIN_ONLY,
1850 EVERY_WATERMARK,
1851 |_| {},
1852 );
1853 let all: Vec<Result<Watermarked<u64>, ScanStop>> = stream.collect().await;
1854 let items = stream_items(&all);
1855
1856 assert!(
1857 items.is_empty(),
1858 "absent include must annihilate: {items:?}"
1859 );
1860 }
1861
1862 #[tokio::test]
1865 async fn absent_exclude_is_noop() {
1866 let source = TestBucketSource {
1867 buckets: Arc::new(BTreeMap::from([(test_key(b"a"), vec![(0, vec![1, 2])])])),
1868 };
1869 let query = BitmapQuery::new(vec![
1870 BitmapTerm::new(vec![include(b"a"), exclude(b"ghost")]).unwrap(),
1871 ])
1872 .unwrap();
1873
1874 let stream = eval_bitmap_query_stream(
1875 source,
1876 query,
1877 0..(2 * BUCKET_SIZE),
1878 BUCKET_SIZE,
1879 ScanDirection::Ascending,
1880 3,
1881 SkipPolicy::DRAIN_ONLY,
1882 EVERY_WATERMARK,
1883 |_| {},
1884 );
1885 let all: Vec<Result<Watermarked<u64>, ScanStop>> = stream.collect().await;
1886
1887 assert_eq!(stream_items(&all), vec![1, 2]);
1888 }
1889}
1890
1891#[cfg(test)]
1892mod pacing_tests {
1893 use super::*;
1894 use crate::bitmap_query::test_utils::*;
1895
1896 const BUCKET_COUNT: u64 = 10;
1897
1898 fn scan_range() -> Range<u64> {
1899 0..(BUCKET_COUNT * BUCKET_SIZE)
1900 }
1901
1902 async fn itemless_watermarks_with_charges(
1906 direction: ScanDirection,
1907 pacing: Option<WatermarkPacing>,
1908 ) -> Vec<(u64, u64)> {
1909 let budget = BitmapScanBudget::new(1_000_000);
1910 let unpaced = eval_bitmap_query_bucket_stream(
1911 parallel_leaves_source(BUCKET_COUNT, &[]),
1912 parallel_leaves_query(),
1913 scan_range(),
1914 BUCKET_SIZE,
1915 direction,
1916 budget.clone(),
1917 SkipPolicy::DRAIN_ONLY,
1918 );
1919 let mut outputs = match pacing {
1920 Some(pacing) => pace_watermarked_buckets(
1921 unpaced,
1922 WatermarkPacer::new(
1923 pacing,
1924 SkipPolicy::DRAIN_ONLY,
1925 budget.clone(),
1926 scan_range(),
1927 BUCKET_SIZE,
1928 ),
1929 ),
1930 None => unpaced,
1931 };
1932 let mut watermarks = Vec::new();
1933 while let Some(output) = outputs.next().await {
1934 match output {
1935 Ok(Watermarked::Watermark(position)) => {
1936 watermarks.push((position, budget.buckets_evaluated()))
1937 }
1938 other => panic!("item-less scan emitted {other:?}"),
1939 }
1940 }
1941 watermarks
1942 }
1943
1944 async fn member_outputs<S: BitmapBucketSource>(
1946 source: S,
1947 direction: ScanDirection,
1948 budget: u64,
1949 pacing: WatermarkPacing,
1950 ) -> Vec<ObservedMemberOutput> {
1951 observe_member_outputs(
1952 eval_bitmap_query_stream(
1953 source,
1954 parallel_leaves_query(),
1955 scan_range(),
1956 BUCKET_SIZE,
1957 direction,
1958 budget,
1959 SkipPolicy::DRAIN_ONLY,
1960 pacing,
1961 |_| {},
1962 )
1963 .collect()
1964 .await,
1965 )
1966 }
1967
1968 #[tokio::test]
1973 async fn charged_buckets_emits_after_min_buckets_charged_rows() {
1974 for direction in [ScanDirection::Ascending, ScanDirection::Descending] {
1975 let unpaced = itemless_watermarks_with_charges(direction, None).await;
1976 assert_eq!(unpaced.len() as u64, BUCKET_COUNT, "{direction:?}");
1977 for min_buckets in [1, 3, 5] {
1978 let context = format!("{direction:?}, min_buckets {min_buckets}");
1979 let paced = itemless_watermarks_with_charges(
1980 direction,
1981 Some(charged_buckets(min_buckets, false)),
1982 )
1983 .await;
1984 assert_eq!(
1985 paced.first(),
1986 unpaced.first(),
1987 "first watermark must not be held: {context}"
1988 );
1989 let mut last_emission = paced[0];
1990 let mut later_emissions = paced[1..].iter().copied().peekable();
1991 for &(position, charged) in &unpaced[1..] {
1992 let charged_since_emission = charged - last_emission.1;
1993 if later_emissions.peek() == Some(&(position, charged)) {
1994 later_emissions.next();
1995 assert!(
1996 charged_since_emission >= min_buckets,
1997 "watermark {position} emitted after {charged_since_emission} \
1998 charged rows: {context}"
1999 );
2000 last_emission = (position, charged);
2001 } else {
2002 assert!(
2003 charged_since_emission < min_buckets,
2004 "watermark {position} held after {charged_since_emission} \
2005 charged rows: {context}"
2006 );
2007 }
2008 }
2009 assert_eq!(
2010 later_emissions.next(),
2011 None,
2012 "paced {paced:?} is not a subsequence of unpaced {unpaced:?}: {context}"
2013 );
2014 }
2015 }
2016 }
2017
2018 #[tokio::test]
2025 async fn item_discards_held_watermark_and_restarts_count() {
2026 let source = parallel_leaves_source(BUCKET_COUNT, &[6]);
2027 let item: ObservedMemberOutput = Ok(Watermarked::Item(6 * BUCKET_SIZE + 1));
2028 let watermark = |bucket: u64| -> ObservedMemberOutput {
2029 Ok(Watermarked::Watermark(bucket * BUCKET_SIZE))
2030 };
2031 let cases: [(ScanDirection, Vec<ObservedMemberOutput>); 2] = [
2032 (
2033 ScanDirection::Ascending,
2034 vec![watermark(1), watermark(4), item.clone(), watermark(10)],
2035 ),
2036 (
2037 ScanDirection::Descending,
2038 vec![watermark(9), item.clone(), watermark(3), watermark(0)],
2039 ),
2040 ];
2041 for (direction, expected) in cases {
2042 let paced = member_outputs(
2043 source.clone(),
2044 direction,
2045 1_000_000,
2046 charged_buckets(5, false),
2047 )
2048 .await;
2049 assert_eq!(paced, expected, "{direction:?}");
2050 }
2051 }
2052
2053 #[tokio::test]
2057 async fn scan_limit_discards_held_watermark() {
2058 let source = parallel_leaves_source(BUCKET_COUNT, &[]);
2059 for direction in [ScanDirection::Ascending, ScanDirection::Descending] {
2060 let unpaced = member_outputs(source.clone(), direction, 7, EVERY_WATERMARK).await;
2064 assert_eq!(unpaced.len(), 4, "{direction:?}: {unpaced:?}");
2065 assert!(
2066 matches!(unpaced[3], Err(ObservedStop::ScanLimit(_))),
2067 "{direction:?}: {unpaced:?}"
2068 );
2069 let paced =
2070 member_outputs(source.clone(), direction, 7, charged_buckets(5, true)).await;
2071 assert_eq!(
2072 paced,
2073 vec![unpaced[0].clone(), unpaced[3].clone()],
2074 "{direction:?}"
2075 );
2076 }
2077 }
2078
2079 #[tokio::test]
2083 async fn fault_and_cancellation_discard_held_watermark() {
2084 let stops: [(fn() -> LeafStop, ObservedStop); 2] = [
2085 (
2086 || LeafStop::Fault(anyhow::anyhow!("leaf storage fault")),
2087 ObservedStop::Fault,
2088 ),
2089 (|| LeafStop::Cancelled, ObservedStop::Cancelled),
2090 ];
2091 for (stop, observed_stop) in stops {
2092 let source = StoppingBucketSource {
2094 inner: parallel_leaves_source(BUCKET_COUNT, &[]),
2095 stop_key: test_key(b"a"),
2096 rows_before_stop: 5,
2097 stop,
2098 };
2099 for (direction, emitted_buckets) in [
2100 (ScanDirection::Ascending, [1, 4]),
2101 (ScanDirection::Descending, [9, 6]),
2102 ] {
2103 let context = format!("{direction:?}, {observed_stop:?}");
2104 let unpaced =
2105 member_outputs(source.clone(), direction, 1_000_000, EVERY_WATERMARK).await;
2106 let newest_unpaced_watermark =
2107 unpaced.iter().rev().find_map(|output| match output {
2108 Ok(Watermarked::Watermark(position)) => Some(*position),
2109 _ => None,
2110 });
2111 assert_eq!(
2112 newest_unpaced_watermark,
2113 Some(5 * BUCKET_SIZE),
2114 "{context}: {unpaced:?}"
2115 );
2116 assert_eq!(
2117 unpaced.last(),
2118 Some(&Err(observed_stop.clone())),
2119 "{context}"
2120 );
2121
2122 let paced = member_outputs(
2123 source.clone(),
2124 direction,
2125 1_000_000,
2126 charged_buckets(5, false),
2127 )
2128 .await;
2129 let mut expected: Vec<ObservedMemberOutput> = emitted_buckets
2130 .iter()
2131 .map(|bucket| Ok(Watermarked::Watermark(bucket * BUCKET_SIZE)))
2132 .collect();
2133 expected.push(Err(observed_stop.clone()));
2134 assert_eq!(paced, expected, "{context}");
2135 }
2136 }
2137 }
2138
2139 #[tokio::test]
2143 async fn clean_end_emits_held_watermark_only_when_configured() {
2144 let source = parallel_leaves_source(BUCKET_COUNT, &[]);
2145 for direction in [ScanDirection::Ascending, ScanDirection::Descending] {
2146 let unpaced =
2147 member_outputs(source.clone(), direction, 1_000_000, EVERY_WATERMARK).await;
2148 let newest_unpaced = unpaced.last().cloned().expect("unpaced run emits");
2149 assert!(
2150 matches!(newest_unpaced, Ok(Watermarked::Watermark(_))),
2151 "{direction:?}: {unpaced:?}"
2152 );
2153
2154 let mut emit_held = member_outputs(
2155 source.clone(),
2156 direction,
2157 1_000_000,
2158 charged_buckets(4, true),
2159 )
2160 .await;
2161 let discard_held = member_outputs(
2162 source.clone(),
2163 direction,
2164 1_000_000,
2165 charged_buckets(4, false),
2166 )
2167 .await;
2168 assert!(
2169 !discard_held.contains(&newest_unpaced),
2170 "{direction:?}: {discard_held:?}"
2171 );
2172 assert_eq!(emit_held.pop(), Some(newest_unpaced), "{direction:?}");
2173 assert_eq!(emit_held, discard_held, "{direction:?}");
2174 }
2175 }
2176
2177 async fn assert_zero_min_buckets_matches_unpaced<S: BitmapBucketSource>(source: S) {
2180 for direction in [ScanDirection::Ascending, ScanDirection::Descending] {
2181 for budget in [3, 7, 1_000_000] {
2182 let unpaced = observe_member_outputs(
2183 flatten_watermarked_buckets(
2184 eval_bitmap_query_bucket_stream(
2185 source.clone(),
2186 parallel_leaves_query(),
2187 scan_range(),
2188 BUCKET_SIZE,
2189 direction,
2190 BitmapScanBudget::new(budget),
2191 SkipPolicy::DRAIN_ONLY,
2192 ),
2193 scan_range(),
2194 BUCKET_SIZE,
2195 direction,
2196 )
2197 .collect()
2198 .await,
2199 );
2200 for emit_held_at_end in [false, true] {
2201 let paced = member_outputs(
2202 source.clone(),
2203 direction,
2204 budget,
2205 charged_buckets(0, emit_held_at_end),
2206 )
2207 .await;
2208 assert_eq!(
2209 paced, unpaced,
2210 "{direction:?}, budget {budget}, emit_held_at_end {emit_held_at_end}"
2211 );
2212 }
2213 }
2214 }
2215 }
2216
2217 #[tokio::test]
2218 async fn zero_min_buckets_matches_unpaced_evaluation() {
2219 assert_zero_min_buckets_matches_unpaced(parallel_leaves_source(BUCKET_COUNT, &[3, 6]))
2220 .await;
2221 assert_zero_min_buckets_matches_unpaced(StoppingBucketSource {
2222 inner: parallel_leaves_source(BUCKET_COUNT, &[3]),
2223 stop_key: test_key(b"b"),
2224 rows_before_stop: 6,
2225 stop: || LeafStop::Cancelled,
2226 })
2227 .await;
2228 }
2229}