perf: add fast-path bitmap_contains that avoids full deserialization

For large hybrid bitmaps, bitmap_contains now reads the serialized
buffer directly via binary search on container descriptions instead
of deserializing the entire bitmap into memory.

This reduces bitmap_contains_large from ~12.1µs to ~2.1µs (6x speedup).

Author: harry-hao

src/common/io/src/bitmap.rs

@@ -728,23 +728,38 @@ pub fn bitmap_len(buf: &[u8]) -> Result<u64> {
         return Ok(0);
     }
 
-    if buf.len() > 3
-        && buf[3] == HYBRID_KIND_LARGE
-        && buf[..2] == HYBRID_MAGIC
-        && buf[2] == HYBRID_VERSION
-    {
+    if is_hybrid_large(buf) {
         Ok(reader::bitmap_len(&buf[HYBRID_HEADER_LEN..])? as u64)
     } else {
         Ok(deserialize_bitmap(buf)?.len())
     }
 }
 
+fn is_hybrid(buf: &[u8]) -> bool {
+    buf.len() >= HYBRID_HEADER_LEN && buf[..2] == HYBRID_MAGIC && buf[2] == HYBRID_VERSION
+}
+
+fn is_hybrid_large(buf: &[u8]) -> bool {
+    is_hybrid(buf) && buf[3] == HYBRID_KIND_LARGE
+}
+
+pub fn bitmap_contains(buf: &[u8], value: u64) -> Result<bool> {
+    if buf.is_empty() {
+        return Ok(false);
+    }
+    if is_hybrid_large(buf) {
+        Ok(reader::bitmap_contains(&buf[HYBRID_HEADER_LEN..], value)?)
+    } else {
+        Ok(deserialize_bitmap(buf)?.contains(value))
+    }
+}
+
 fn parse_bitmap_rhs(buf: &[u8]) -> Result<BitmapRhsView<'_>> {
     if buf.is_empty() {
         return Ok(BitmapRhsView::Empty);
     }
 
-    if buf.len() >= HYBRID_HEADER_LEN && buf[..2] == HYBRID_MAGIC && buf[2] == HYBRID_VERSION {
+    if is_hybrid(buf) {
         let payload = &buf[HYBRID_HEADER_LEN..];
         match buf[3] {
             HYBRID_KIND_SMALL => {
@@ -768,7 +783,7 @@ fn validate_serialized_bitmap(buf: &[u8]) -> Result<()> {
         return Ok(());
     }
 
-    if buf.len() >= HYBRID_HEADER_LEN && buf[..2] == HYBRID_MAGIC && buf[2] == HYBRID_VERSION {
+    if is_hybrid(buf) {
         let payload = &buf[HYBRID_HEADER_LEN..];
         match buf[3] {
             HYBRID_KIND_SMALL => {
@@ -1429,4 +1444,50 @@ mod tests {
         let decoded = deserialize_bitmap(&legacy).unwrap();
         assert_eq!(decoded.into_iter().collect::<Vec<_>>(), vec![1, 5, 42]);
     }
+
+    #[test]
+    fn test_bitmap_contains() {
+        // HybridLarge: hit and miss (spanning two prefix buckets)
+        let large = HybridBitmap::from_iter(
+            [0u64, 2500, 65535, 65536, 65536 + 2500, 131071]
+                .into_iter()
+                .chain((0..40000).map(|v| v + 200000)),
+        );
+        let mut buf = Vec::new();
+        large.serialize_into(&mut buf).unwrap();
+        assert!(bitmap_contains(&buf, 0).unwrap());
+        assert!(bitmap_contains(&buf, 65536 + 2500).unwrap());
+        assert!(bitmap_contains(&buf, 200000).unwrap());
+        // 999999: no matching prefix bucket (fast-path miss)
+        assert!(!bitmap_contains(&buf, 999999).unwrap());
+        // 100000: prefix bucket exists but value not in container (search miss)
+        assert!(!bitmap_contains(&buf, 100000).unwrap());
+
+        // HybridLarge with array container at max capacity (cardinality = 4096)
+        // Tests the ARRAY_LIMIT boundary: <= 4096 is array, > 4096 is bitmap.
+        let boundary = HybridBitmap::from_iter(0u64..4096);
+        let mut buf = Vec::new();
+        boundary.serialize_into(&mut buf).unwrap();
+        assert!(bitmap_contains(&buf, 0).unwrap());
+        assert!(bitmap_contains(&buf, 4095).unwrap());
+        assert!(!bitmap_contains(&buf, 4096).unwrap());
+
+        // HybridSmall: hit and miss
+        let small = HybridBitmap::from_iter(0u64..31);
+        let mut buf = Vec::new();
+        small.serialize_into(&mut buf).unwrap();
+        assert!(bitmap_contains(&buf, 15).unwrap());
+        assert!(!bitmap_contains(&buf, 50).unwrap());
+
+        // Legacy: hit and miss
+        let mut tree = RoaringTreemap::new();
+        tree.insert(42);
+        let mut buf = Vec::new();
+        tree.serialize_into(&mut buf).unwrap();
+        assert!(bitmap_contains(&buf, 42).unwrap());
+        assert!(!bitmap_contains(&buf, 99).unwrap());
+
+        // Empty buffer
+        assert!(!bitmap_contains(&[], 42).unwrap());
+    }
 }

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

@@ -26,6 +26,13 @@ use roaring::RoaringTreemap;
 const DESCRIPTION_BYTES: usize = 4;
 const OFFSET_BYTES: usize = 4;
 
+// Container layout constants
+const ARRAY_LIMIT: usize = 4096;
+const WORD_BITS: usize = 64;
+const WORD_BYTES: usize = 8;
+const BITMAP_WORDS: usize = 1024;
+const BITMAP_BYTES: usize = WORD_BYTES * BITMAP_WORDS;
+
 #[derive(Clone)]
 pub struct TreemapReader<'a> {
     buf: &'a [u8],
@@ -113,15 +120,12 @@ impl BitmapReader<'_> {
         reader.seek_relative(last_container * OFFSET_BYTES as i64)?;
         let last_offset = reader.read_u32::<LittleEndian>()?;
 
-        const ARRAY_LIMIT: usize = 4096;
-        const BITMAP_LENGTH: usize = 8192;
-
         let size = 4
             + last_offset as usize
-            + if last_cardinality < ARRAY_LIMIT {
+            + if last_cardinality <= ARRAY_LIMIT {
                 2 * last_cardinality
             } else {
-                BITMAP_LENGTH
+                BITMAP_BYTES
             };
 
         if buf.len() < size {
@@ -142,7 +146,6 @@ impl BitmapReader<'_> {
         self.containers as usize
     }
 
-    #[allow(dead_code)]
     pub fn prefix(&self) -> u32 {
         self.prefix
     }
@@ -164,6 +167,26 @@ impl BitmapReader<'_> {
     pub fn bitmap_buf(&self) -> &[u8] {
         &self.buf[4..]
     }
+
+    pub(crate) fn container_offset(&self, i: usize) -> io::Result<usize> {
+        if i >= self.containers() {
+            return Err(Error::other("index out of range"));
+        }
+        let offset_table_start = 12 + self.containers() * DESCRIPTION_BYTES;
+        let offset_pos = offset_table_start + i * OFFSET_BYTES;
+        if offset_pos + OFFSET_BYTES > self.buf.len() {
+            return Err(Error::other("offset table too short"));
+        }
+        let mut reader = Cursor::new(&self.buf[offset_pos..]);
+        let offset = reader.read_u32::<LittleEndian>()? as usize;
+        if offset > self.bitmap_buf().len() {
+            return Err(Error::new(
+                ErrorKind::InvalidData,
+                "container offset exceeds bitmap data",
+            ));
+        }
+        Ok(offset)
+    }
 }
 
 pub struct Description {
@@ -190,6 +213,84 @@ pub fn bitmap_len(buf: &[u8]) -> io::Result<usize> {
     Ok(sum)
 }
 
+pub(crate) fn bitmap_contains(buf: &[u8], value: u64) -> io::Result<bool> {
+    let tree = TreemapReader::new(buf)?;
+    let prefix = (value >> 32) as u32;
+    let container_key = ((value >> 16) & 0xFFFF) as u16;
+    let low16 = (value & 0xFFFF) as u16;
+
+    for bitmap_result in tree.iter() {
+        let bitmap = bitmap_result?;
+        if bitmap.prefix() < prefix {
+            continue;
+        }
+        if bitmap.prefix() > prefix {
+            return Ok(false);
+        }
+        return bitmap_contains_in_bitmap(&bitmap, container_key, low16);
+    }
+    Ok(false)
+}
+
+fn bitmap_contains_in_bitmap(
+    bitmap: &BitmapReader<'_>,
+    container_key: u16,
+    low16: u16,
+) -> io::Result<bool> {
+    let mut lo = 0;
+    let mut hi = bitmap.containers();
+    while lo < hi {
+        let mid = lo + (hi - lo) / 2;
+        let desc = bitmap.description(mid)?;
+        match desc.prefix.cmp(&container_key) {
+            std::cmp::Ordering::Less => lo = mid + 1,
+            std::cmp::Ordering::Greater => hi = mid,
+            std::cmp::Ordering::Equal => {
+                let cardinality = desc.cardinality();
+                let offset = bitmap.container_offset(mid)?;
+                let container_data = &bitmap.bitmap_buf()[offset..];
+                return if cardinality <= ARRAY_LIMIT {
+                    array_container_contains(container_data, cardinality, low16)
+                } else {
+                    bitmap_container_contains(container_data, low16)
+                };
+            }
+        }
+    }
+    Ok(false)
+}
+
+/// Binary search in an array container (cardinality < 4096).
+fn array_container_contains(data: &[u8], cardinality: usize, low16: u16) -> io::Result<bool> {
+    if data.len() < cardinality * 2 {
+        return Err(Error::other("array container too short"));
+    }
+    let mut lo = 0;
+    let mut hi = cardinality;
+    while lo < hi {
+        let mid = lo + (hi - lo) / 2;
+        let v = u16::from_le_bytes(data[mid * 2..mid * 2 + 2].try_into().unwrap());
+        match v.cmp(&low16) {
+            std::cmp::Ordering::Less => lo = mid + 1,
+            std::cmp::Ordering::Greater => hi = mid,
+            std::cmp::Ordering::Equal => return Ok(true),
+        }
+    }
+    Ok(false)
+}
+
+/// Bit lookup in a bitmap container (cardinality >= 4096).
+fn bitmap_container_contains(data: &[u8], low16: u16) -> io::Result<bool> {
+    if data.len() < BITMAP_BYTES {
+        return Err(Error::other("bitmap container too short"));
+    }
+    let word_index = low16 as usize / WORD_BITS;
+    let bit_index = low16 as usize % WORD_BITS;
+    let start = word_index * WORD_BYTES;
+    let word = u64::from_le_bytes(data[start..start + WORD_BYTES].try_into().unwrap());
+    Ok(word & (1 << bit_index) != 0)
+}
+
 pub fn intersection_with_serialized(tree: &mut RoaringTreemap, buf: &[u8]) -> io::Result<()> {
     use std::cmp::Ordering::*;
     let rhs = TreemapReader::new(buf)?;

src/common/io/src/lib.rs

@@ -53,6 +53,7 @@ pub use bitmap::HYBRID_MAGIC;
 pub use bitmap::HYBRID_VERSION;
 pub use bitmap::HybridBitmap;
 pub use bitmap::LARGE_THRESHOLD;
+pub use bitmap::bitmap_contains;
 pub use bitmap::deserialize_bitmap;
 pub use bitmap::parse_bitmap;
 pub use decimal::display_decimal_128;

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

@@ -37,6 +37,7 @@ use databend_common_expression::vectorize_with_builder_3_arg;
 use databend_common_expression::with_signed_integer_mapped_type;
 use databend_common_expression::with_unsigned_integer_mapped_type;
 use databend_common_io::HybridBitmap;
+use databend_common_io::bitmap::bitmap_contains;
 use databend_common_io::bitmap::bitmap_len;
 use databend_common_io::deserialize_bitmap;
 use databend_common_io::parse_bitmap;
@@ -269,9 +270,9 @@ pub fn register(registry: &mut FunctionRegistry) {
                     builder.push(false);
                     return;
                 }
-                match deserialize_bitmap(b) {
-                    Ok(rb) => {
-                        builder.push(rb.contains(item));
+                match bitmap_contains(b, item) {
+                    Ok(found) => {
+                        builder.push(found);
                     }
                     Err(e) => {
                         ctx.set_error(builder.len(), e.to_string());