Skip to content

Commit 840d999

Browse files
ZeutschlerZeutschler
committed
feat(iter): Add next_many() for batch iteration
Add next_many() method to RoaringBitmap iterators (both Iter and IntoIter) for efficient batch extraction of values into a buffer. This method is significantly faster than calling next() repeatedly: - 2.5-3.1x speedup for dense bitmaps (bitmap storage) - 10.5x speedup for sparse arrays (array storage) Implementation details: - Store::Iter: Direct slice copy for Array/Vec, bit extraction for Bitmap, run expansion for Interval stores - Container::Iter: Uses u16 buffer internally, combines with container key - Bitmap::Iter/IntoIter: Handles front/back iterators and container transitions The API mirrors the next_many() method available in CRoaring and the Go implementation of RoaringBitmap. Benchmarks (1M dense values): next(): 19.02ms next_many(): 6.16ms (3.09x faster) Benchmarks (10K sparse values, every 100th): next(): 1.67ms next_many(): 159.58µs (10.46x faster)
1 parent a96d4df commit 840d999

6 files changed

Lines changed: 551 additions & 0 deletions

File tree

‎roaring/src/bitmap/container.rs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -473,6 +473,34 @@ impl Iter<'_> {
473473
.next_range_back()
474474
.map(|r| util::join(self.key, *r.start())..=util::join(self.key, *r.end()))
475475
}
476+
477+
/// Read multiple values from the iterator into `dst`.
478+
/// Returns the number of values read.
479+
///
480+
/// This can be significantly faster than calling `next()` repeatedly.
481+
pub(crate) fn next_many(&mut self, dst: &mut [u32]) -> usize {
482+
// Use a temporary u16 buffer for the inner iterator
483+
const BUF_SIZE: usize = 256;
484+
let mut buf = [0u16; BUF_SIZE];
485+
486+
let key = self.key;
487+
let mut count = 0;
488+
489+
while count < dst.len() {
490+
let remaining = dst.len() - count;
491+
let to_read = remaining.min(BUF_SIZE);
492+
let n = self.inner.next_many(&mut buf[..to_read]);
493+
if n == 0 {
494+
break;
495+
}
496+
for i in 0..n {
497+
dst[count + i] = util::join(key, buf[i]);
498+
}
499+
count += n;
500+
}
501+
502+
count
503+
}
476504
}
477505

478506
impl fmt::Debug for Container {

‎roaring/src/bitmap/iter.rs‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,82 @@ impl Iter<'_> {
331331
pub fn next_range_back(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
332332
next_range_back_impl(&mut self.front, &mut self.containers, &mut self.back)
333333
}
334+
335+
/// Retrieve the next `dst.len()` values from the iterator and write them into `dst`.
336+
///
337+
/// Returns the number of values written. This will be less than `dst.len()` only
338+
/// if the iterator is exhausted.
339+
///
340+
/// This method is significantly faster than calling `next()` repeatedly due to
341+
/// reduced per-element overhead and better CPU cache utilization.
342+
///
343+
/// # Examples
344+
///
345+
/// ```rust
346+
/// use roaring::RoaringBitmap;
347+
///
348+
/// let bitmap: RoaringBitmap = (0..100).collect();
349+
/// let mut iter = bitmap.iter();
350+
/// let mut buf = [0u32; 32];
351+
///
352+
/// let n = iter.next_many(&mut buf);
353+
/// assert_eq!(n, 32);
354+
/// assert_eq!(buf[0], 0);
355+
/// assert_eq!(buf[31], 31);
356+
///
357+
/// // Iterate remainder
358+
/// let n = iter.next_many(&mut buf);
359+
/// assert_eq!(n, 32);
360+
/// assert_eq!(buf[0], 32);
361+
/// ```
362+
pub fn next_many(&mut self, dst: &mut [u32]) -> usize {
363+
if dst.is_empty() {
364+
return 0;
365+
}
366+
367+
let mut count = 0;
368+
369+
// First drain from the front container iterator if present
370+
if let Some(ref mut front_iter) = self.front {
371+
let n = front_iter.next_many(&mut dst[count..]);
372+
count += n;
373+
if count >= dst.len() {
374+
return count;
375+
}
376+
// Front is exhausted
377+
self.front = None;
378+
}
379+
380+
// Process remaining containers
381+
while count < dst.len() {
382+
let Some(container) = self.containers.next() else {
383+
// No more containers in the middle, try the back
384+
break;
385+
};
386+
let mut container_iter = container.into_iter();
387+
let n = container_iter.next_many(&mut dst[count..]);
388+
count += n;
389+
390+
// If container still has values, save it as new front
391+
if n > 0 && container_iter.len() > 0 {
392+
self.front = Some(container_iter);
393+
return count;
394+
}
395+
}
396+
397+
// Finally, try draining from the back iterator if present
398+
if count < dst.len() {
399+
if let Some(ref mut back_iter) = self.back {
400+
let n = back_iter.next_many(&mut dst[count..]);
401+
count += n;
402+
if back_iter.len() == 0 {
403+
self.back = None;
404+
}
405+
}
406+
}
407+
408+
count
409+
}
334410
}
335411

336412
impl IntoIter {
@@ -419,6 +495,82 @@ impl IntoIter {
419495
pub fn next_range_back(&mut self) -> Option<core::ops::RangeInclusive<u32>> {
420496
next_range_back_impl(&mut self.front, &mut self.containers, &mut self.back)
421497
}
498+
499+
/// Retrieve the next `dst.len()` values from the iterator and write them into `dst`.
500+
///
501+
/// Returns the number of values written. This will be less than `dst.len()` only
502+
/// if the iterator is exhausted.
503+
///
504+
/// This method is significantly faster than calling `next()` repeatedly due to
505+
/// reduced per-element overhead and better CPU cache utilization.
506+
///
507+
/// # Examples
508+
///
509+
/// ```rust
510+
/// use roaring::RoaringBitmap;
511+
///
512+
/// let bitmap: RoaringBitmap = (0..100).collect();
513+
/// let mut iter = bitmap.into_iter();
514+
/// let mut buf = [0u32; 32];
515+
///
516+
/// let n = iter.next_many(&mut buf);
517+
/// assert_eq!(n, 32);
518+
/// assert_eq!(buf[0], 0);
519+
/// assert_eq!(buf[31], 31);
520+
///
521+
/// // Iterate remainder
522+
/// let n = iter.next_many(&mut buf);
523+
/// assert_eq!(n, 32);
524+
/// assert_eq!(buf[0], 32);
525+
/// ```
526+
pub fn next_many(&mut self, dst: &mut [u32]) -> usize {
527+
if dst.is_empty() {
528+
return 0;
529+
}
530+
531+
let mut count = 0;
532+
533+
// First drain from the front container iterator if present
534+
if let Some(ref mut front_iter) = self.front {
535+
let n = front_iter.next_many(&mut dst[count..]);
536+
count += n;
537+
if count >= dst.len() {
538+
return count;
539+
}
540+
// Front is exhausted
541+
self.front = None;
542+
}
543+
544+
// Process remaining containers
545+
while count < dst.len() {
546+
let Some(container) = self.containers.next() else {
547+
// No more containers in the middle, try the back
548+
break;
549+
};
550+
let mut container_iter = container.into_iter();
551+
let n = container_iter.next_many(&mut dst[count..]);
552+
count += n;
553+
554+
// If container still has values, save it as new front
555+
if n > 0 && container_iter.len() > 0 {
556+
self.front = Some(container_iter);
557+
return count;
558+
}
559+
}
560+
561+
// Finally, try draining from the back iterator if present
562+
if count < dst.len() {
563+
if let Some(ref mut back_iter) = self.back {
564+
let n = back_iter.next_many(&mut dst[count..]);
565+
count += n;
566+
if back_iter.len() == 0 {
567+
self.back = None;
568+
}
569+
}
570+
}
571+
572+
count
573+
}
422574
}
423575

424576
fn size_hint_impl(

‎roaring/src/bitmap/store/bitmap_store.rs‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -691,6 +691,55 @@ impl<B: Borrow<[u64; BITMAP_LENGTH]>> BitmapIter<B> {
691691
let index = 63 - index_from_left;
692692
Some(64 * key_back + index)
693693
}
694+
695+
/// Read multiple values from the iterator into `dst`.
696+
/// Returns the number of values read.
697+
///
698+
/// This can be significantly faster than calling `next()` repeatedly.
699+
pub fn next_many(&mut self, dst: &mut [u16]) -> usize {
700+
if dst.is_empty() {
701+
return 0;
702+
}
703+
704+
let mut count = 0;
705+
let bits = self.bits.borrow();
706+
707+
while count < dst.len() {
708+
// Advance to next non-zero word if current is empty
709+
if self.value == 0 {
710+
if self.key >= self.key_back {
711+
break;
712+
}
713+
loop {
714+
self.key += 1;
715+
if self.key == self.key_back {
716+
self.value = core::mem::replace(&mut self.value_back, 0);
717+
break;
718+
}
719+
// Safety: key is always in bounds
720+
self.value = unsafe { *bits.get_unchecked(self.key as usize) };
721+
if self.value != 0 {
722+
break;
723+
}
724+
}
725+
if self.value == 0 {
726+
break;
727+
}
728+
}
729+
730+
// Extract set bits from current word
731+
let base = self.key as u16 * 64;
732+
while self.value != 0 && count < dst.len() {
733+
let bit_pos = self.value.trailing_zeros() as u16;
734+
dst[count] = base + bit_pos;
735+
count += 1;
736+
// Clear the lowest set bit
737+
self.value &= self.value - 1;
738+
}
739+
}
740+
741+
count
742+
}
694743
}
695744

696745
fn advance_to_next_nonzero_word<'a>(

‎roaring/src/bitmap/store/interval_store.rs‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -834,6 +834,59 @@ impl<I: SliceIterator<Interval>> RunIter<I> {
834834
let result = self.intervals.as_slice().last()?.end - self.backward_offset;
835835
Some(result)
836836
}
837+
838+
/// Read multiple values from the iterator into `dst`.
839+
/// Returns the number of values read.
840+
///
841+
/// This can be significantly faster than calling `next()` repeatedly
842+
/// because it processes runs in bulk.
843+
pub fn next_many(&mut self, dst: &mut [u16]) -> usize {
844+
if dst.is_empty() {
845+
return 0;
846+
}
847+
848+
let mut count = 0;
849+
850+
while count < dst.len() {
851+
let Some(interval) = self.intervals.as_slice().first() else {
852+
break;
853+
};
854+
855+
let end_offset = if self.intervals.as_slice().len() == 1 {
856+
self.backward_offset
857+
} else {
858+
0
859+
};
860+
861+
let start = interval.start + self.forward_offset;
862+
let end = interval.end - end_offset;
863+
864+
// How many values can we emit from this interval?
865+
let available = (end - start + 1) as usize;
866+
let to_emit = available.min(dst.len() - count);
867+
868+
// Emit values
869+
for i in 0..to_emit {
870+
dst[count + i] = start + i as u16;
871+
}
872+
count += to_emit;
873+
874+
// Advance within or past this interval
875+
if to_emit == available {
876+
// Consumed entire interval
877+
_ = self.intervals.next();
878+
self.forward_offset = 0;
879+
if self.intervals.as_slice().is_empty() {
880+
self.backward_offset = 0;
881+
}
882+
} else {
883+
// Partial consumption
884+
self.forward_offset += to_emit as u16;
885+
}
886+
}
887+
888+
count
889+
}
837890
}
838891

839892
impl<I: SliceIterator<Interval>> Iterator for RunIter<I> {

‎roaring/src/bitmap/store/mod.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1038,6 +1038,36 @@ impl Iterator for Iter<'_> {
10381038
}
10391039
}
10401040

1041+
impl Iter<'_> {
1042+
/// Read multiple values from the iterator into `dst`.
1043+
/// Returns the number of values read.
1044+
///
1045+
/// This can be significantly faster than calling `next()` repeatedly.
1046+
pub fn next_many(&mut self, dst: &mut [u16]) -> usize {
1047+
match self {
1048+
Iter::Array(inner) => {
1049+
let remaining = inner.as_slice();
1050+
let n = remaining.len().min(dst.len());
1051+
dst[..n].copy_from_slice(&remaining[..n]);
1052+
if n > 0 { _ = inner.nth(n - 1); }
1053+
n
1054+
}
1055+
Iter::Vec(inner) => {
1056+
let remaining = inner.as_slice();
1057+
let n = remaining.len().min(dst.len());
1058+
dst[..n].copy_from_slice(&remaining[..n]);
1059+
if n > 0 { _ = inner.nth(n - 1); }
1060+
n
1061+
}
1062+
Iter::BitmapBorrowed(inner) => inner.next_many(dst),
1063+
Iter::BitmapOwned(inner) => inner.next_many(dst),
1064+
Iter::RunBorrowed(inner) => inner.next_many(dst),
1065+
Iter::RunOwned(inner) => inner.next_many(dst),
1066+
}
1067+
}
1068+
}
1069+
1070+
10411071
impl DoubleEndedIterator for Iter<'_> {
10421072
fn next_back(&mut self) -> Option<Self::Item> {
10431073
match self {

0 commit comments

Comments
 (0)