Skip to content

Commit 1762aaa

Browse files
committed
perf: add fast-path bitmap_max that avoids full deserialization
For large hybrid bitmaps, bitmap_max now reads the maximum value directly from the serialized buffer (last element of the last container) instead of deserializing the entire roaring treemap. This reduces bitmap_max_large from ~10.7µs to ~531ns (20x speedup).
1 parent 1402306 commit 1762aaa

4 files changed

Lines changed: 122 additions & 8 deletions

File tree

src/common/io/src/bitmap.rs

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -765,6 +765,17 @@ pub fn bitmap_min(buf: &[u8]) -> Result<Option<u64>> {
765765
}
766766
}
767767

768+
pub fn bitmap_max(buf: &[u8]) -> Result<Option<u64>> {
769+
if buf.is_empty() {
770+
return Ok(None);
771+
}
772+
if is_hybrid_large(buf) {
773+
Ok(reader::bitmap_max(&buf[HYBRID_HEADER_LEN..])?)
774+
} else {
775+
Ok(deserialize_bitmap(buf)?.max())
776+
}
777+
}
778+
768779
fn parse_bitmap_rhs(buf: &[u8]) -> Result<BitmapRhsView<'_>> {
769780
if buf.is_empty() {
770781
return Ok(BitmapRhsView::Empty);
@@ -1548,4 +1559,52 @@ mod tests {
15481559
// Empty buffer
15491560
assert_eq!(bitmap_min(&[]).unwrap(), None);
15501561
}
1562+
1563+
#[test]
1564+
fn test_bitmap_max() {
1565+
// HybridLarge: spanning multiple containers
1566+
let large = HybridBitmap::from_iter(
1567+
[0u64, 2500, 65535, 65536, 65536 + 2500, 131071]
1568+
.into_iter()
1569+
.chain((0..40000).map(|v| v + 200000)),
1570+
);
1571+
let mut buf = Vec::new();
1572+
large.serialize_into(&mut buf).unwrap();
1573+
assert_eq!(bitmap_max(&buf).unwrap(), Some(239999));
1574+
1575+
// HybridLarge with array container at max capacity (cardinality = 4096)
1576+
let boundary = HybridBitmap::from_iter(0u64..4096);
1577+
let mut buf = Vec::new();
1578+
boundary.serialize_into(&mut buf).unwrap();
1579+
assert_eq!(bitmap_max(&buf).unwrap(), Some(4095));
1580+
1581+
// HybridLarge with multiple containers: array (4096) + bitmap (>4096)
1582+
let mixed = HybridBitmap::from_iter((0u64..4096).chain(65536..106496));
1583+
let mut buf = Vec::new();
1584+
mixed.serialize_into(&mut buf).unwrap();
1585+
assert_eq!(bitmap_max(&buf).unwrap(), Some(106495));
1586+
1587+
// HybridSmall
1588+
let small = HybridBitmap::from_iter(0u64..31);
1589+
let mut buf = Vec::new();
1590+
small.serialize_into(&mut buf).unwrap();
1591+
assert_eq!(bitmap_max(&buf).unwrap(), Some(30));
1592+
1593+
// HybridSmall with different range
1594+
let small = HybridBitmap::from_iter(100u64..131);
1595+
let mut buf = Vec::new();
1596+
small.serialize_into(&mut buf).unwrap();
1597+
assert_eq!(bitmap_max(&buf).unwrap(), Some(130));
1598+
1599+
// Legacy
1600+
let mut tree = RoaringTreemap::new();
1601+
tree.insert(42);
1602+
tree.insert(100);
1603+
let mut buf = Vec::new();
1604+
tree.serialize_into(&mut buf).unwrap();
1605+
assert_eq!(bitmap_max(&buf).unwrap(), Some(100));
1606+
1607+
// Empty buffer
1608+
assert_eq!(bitmap_max(&[]).unwrap(), None);
1609+
}
15511610
}

src/common/io/src/bitmap/reader.rs

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -337,6 +337,61 @@ fn bitmap_container_first(data: &[u8]) -> io::Result<u16> {
337337
Err(Error::other("bitmap container has no set bits"))
338338
}
339339

340+
pub(crate) fn bitmap_max(buf: &[u8]) -> io::Result<Option<u64>> {
341+
let tree = TreemapReader::new(buf)?;
342+
// Find the last prefix bucket
343+
let mut last_bitmap = None;
344+
for bitmap_result in tree.iter() {
345+
last_bitmap = Some(bitmap_result?);
346+
}
347+
let bitmap = match last_bitmap {
348+
None => return Ok(None),
349+
Some(b) => b,
350+
};
351+
if bitmap.containers() == 0 {
352+
return Ok(None);
353+
}
354+
let last_idx = bitmap.containers() - 1;
355+
let desc = bitmap.description(last_idx)?;
356+
let offset = bitmap.container_offset(last_idx)?;
357+
let container_data = &bitmap.bitmap_buf()[offset..];
358+
let prefix = bitmap.prefix() as u64;
359+
let container_key = desc.prefix as u64;
360+
let cardinality = desc.cardinality();
361+
let low16 = if cardinality <= ARRAY_LIMIT {
362+
array_container_last(container_data, cardinality)?
363+
} else {
364+
bitmap_container_last(container_data)?
365+
};
366+
Ok(Some(prefix << 32 | container_key << 16 | low16 as u64))
367+
}
368+
369+
fn array_container_last(data: &[u8], cardinality: usize) -> io::Result<u16> {
370+
if data.len() < cardinality * 2 {
371+
return Err(Error::other("array container too short"));
372+
}
373+
let offset = (cardinality - 1) * 2;
374+
Ok(u16::from_le_bytes(
375+
data[offset..offset + 2].try_into().unwrap(),
376+
))
377+
}
378+
379+
fn bitmap_container_last(data: &[u8]) -> io::Result<u16> {
380+
if data.len() < BITMAP_BYTES {
381+
return Err(Error::other("bitmap container too short"));
382+
}
383+
// Find the highest set bit by scanning from the last word backwards
384+
for word_index in (0..BITMAP_WORDS).rev() {
385+
let start = word_index * WORD_BYTES;
386+
let word = u64::from_le_bytes(data[start..start + WORD_BYTES].try_into().unwrap());
387+
if word != 0 {
388+
let bit_index = WORD_BITS - 1 - word.leading_zeros() as usize;
389+
return Ok((word_index * WORD_BITS + bit_index) as u16);
390+
}
391+
}
392+
Err(Error::other("bitmap container has no set bits"))
393+
}
394+
340395
pub fn intersection_with_serialized(tree: &mut RoaringTreemap, buf: &[u8]) -> io::Result<()> {
341396
use std::cmp::Ordering::*;
342397
let rhs = TreemapReader::new(buf)?;

src/common/io/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ pub use bitmap::HYBRID_VERSION;
5454
pub use bitmap::HybridBitmap;
5555
pub use bitmap::LARGE_THRESHOLD;
5656
pub use bitmap::bitmap_contains;
57+
pub use bitmap::bitmap_max;
5758
pub use bitmap::bitmap_min;
5859
pub use bitmap::deserialize_bitmap;
5960
pub use bitmap::parse_bitmap;

src/query/functions/src/scalars/bitmap.rs

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ use databend_common_expression::with_unsigned_integer_mapped_type;
3939
use databend_common_io::HybridBitmap;
4040
use databend_common_io::bitmap::bitmap_contains;
4141
use databend_common_io::bitmap::bitmap_len;
42+
use databend_common_io::bitmap::bitmap_max;
4243
use databend_common_io::bitmap::bitmap_min;
4344
use databend_common_io::deserialize_bitmap;
4445
use databend_common_io::parse_bitmap;
@@ -450,14 +451,12 @@ pub fn register(registry: &mut FunctionRegistry) {
450451
builder.push(0);
451452
return;
452453
}
453-
let val = match deserialize_bitmap(b) {
454-
Ok(rb) => match rb.max() {
455-
Some(val) => val,
456-
None => {
457-
ctx.set_error(builder.len(), "The bitmap is empty");
458-
0
459-
}
460-
},
454+
let val = match bitmap_max(b) {
455+
Ok(Some(val)) => val,
456+
Ok(None) => {
457+
ctx.set_error(builder.len(), "The bitmap is empty");
458+
0
459+
}
461460
Err(e) => {
462461
ctx.set_error(builder.len(), e.to_string());
463462
0

0 commit comments

Comments
 (0)