From 5c8794ac81829fa02eb974968d728462f7faa545 Mon Sep 17 00:00:00 2001 From: HarryHao Date: Sun, 26 Jul 2026 22:12:58 +0800 Subject: [PATCH 1/9] chore: add bitmap scalar function benchmarks Added a `bitmap_scalar` bench group covering: - bitmap_contains - bitmap_count - bitmap_min - bitmap_max - bitmap_and - bitmap_has_any - bitmap_has_all --- src/query/functions/benches/bench.rs | 305 +++++++++++++++++++++++++++ 1 file changed, 305 insertions(+) diff --git a/src/query/functions/benches/bench.rs b/src/query/functions/benches/bench.rs index 37109aa2b7097..92ffb32b3fb5a 100644 --- a/src/query/functions/benches/bench.rs +++ b/src/query/functions/benches/bench.rs @@ -308,6 +308,311 @@ mod bitmap { } } +#[divan::bench_group(max_time = 0.5)] +mod bitmap_scalar { + use databend_common_expression::BlockEntry; + use databend_common_expression::DataBlock; + use databend_common_expression::Evaluator; + use databend_common_expression::Expr; + use databend_common_expression::FromData; + use databend_common_expression::FunctionContext; + use databend_common_expression::type_check; + use databend_common_expression::types::BitmapType; + use databend_common_expression::types::DataType; + use databend_common_expression_test_support as parser; + use databend_common_functions::BUILTIN_FUNCTIONS; + use databend_common_io::HybridBitmap; + + fn serialize_bitmap(bitmap: &HybridBitmap) -> Vec { + let mut data = Vec::new(); + bitmap.serialize_into(&mut data).unwrap(); + data + } + + fn bitmap_entry(bitmap: &HybridBitmap) -> BlockEntry { + BitmapType::from_data(vec![serialize_bitmap(bitmap); 1]).into() + } + + fn build_expr(sql: &str, columns: &[(&str, DataType)]) -> Expr { + let raw_expr = parser::parse_raw_expr(sql, columns, &BUILTIN_FUNCTIONS); + type_check::check(&raw_expr, &BUILTIN_FUNCTIONS).unwrap() + } + + fn eval_block(expr: &Expr, block: &DataBlock) { + let func_ctx = FunctionContext::default(); + let evaluator = Evaluator::new(block, &func_ctx, &BUILTIN_FUNCTIONS); + let result = evaluator.run(expr).unwrap(); + divan::black_box(result); + } + + fn large_bitmap() -> HybridBitmap { + let mut bm = HybridBitmap::new(); + for i in 0..10u64 { + for v in i * 65536..i * 65536 + 5000 { + bm.insert(v); + } + } + for j in 10..50u64 { + for v in j * 65536..j * 65536 + 100 { + bm.insert(v); + } + } + bm + } + + fn overlap_large_bitmap() -> HybridBitmap { + let mut bm = HybridBitmap::new(); + for i in 5..10u64 { + for v in i * 65536 - 3000..i * 65536 + 5000 { + bm.insert(v); + } + } + for i in 10..50u64 { + for v in i * 65536 - 30..i * 65536 + 50 { + bm.insert(v); + } + } + bm + } + + fn subset_large_bitmap() -> HybridBitmap { + let mut bm = HybridBitmap::new(); + for i in 0..5u64 { + for v in i * 65536 + 1000..i * 65536 + 4000 { + bm.insert(v); + } + } + for j in 10..30u64 { + for v in j * 65536 + 10..j * 65536 + 90 { + bm.insert(v); + } + } + bm + } + + fn disjoint_large_bitmap() -> HybridBitmap { + let mut bm = HybridBitmap::new(); + for i in 100..110u64 { + for v in i * 65536..i * 65536 + 5000 { + bm.insert(v); + } + } + for j in 110..150u64 { + for v in j * 65536..j * 65536 + 100 { + bm.insert(v); + } + } + + bm + } + + fn small_bitmap() -> HybridBitmap { + HybridBitmap::from_iter(24 * 65536..24 * 65536 + 31) + } + + fn overlap_small_bitmap() -> HybridBitmap { + HybridBitmap::from_iter(24 * 65536 - 10..24 * 65536 + 21) + } + + fn subset_small_bitmap() -> HybridBitmap { + HybridBitmap::from_iter(24 * 65536 + 1..24 * 65536 + 30) + } + + const C1: &[(&str, DataType)] = &[("a", DataType::Bitmap)]; + const C2: &[(&str, DataType)] = &[("a", DataType::Bitmap), ("b", DataType::Bitmap)]; + + #[divan::bench] + fn bitmap_contains_large(bencher: divan::Bencher) { + let bm = large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&bm)], 1); + let expr = build_expr("bitmap_contains(a, 5*65536+2500)", C1); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_contains_small(bencher: divan::Bencher) { + let bm = small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&bm)], 1); + let expr = build_expr("bitmap_contains(a, 15)", C1); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_count_large(bencher: divan::Bencher) { + let bm = large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&bm)], 1); + let expr = build_expr("bitmap_count(a)", C1); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_count_small(bencher: divan::Bencher) { + let bm = small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&bm)], 1); + let expr = build_expr("bitmap_count(a)", C1); + bencher.bench(|| eval_block(&expr, &block)); + } + + // bitmap_min + #[divan::bench] + fn bitmap_min_large(bencher: divan::Bencher) { + let bm = large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&bm)], 1); + let expr = build_expr("bitmap_min(a)", C1); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_min_small(bencher: divan::Bencher) { + let bm = small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&bm)], 1); + let expr = build_expr("bitmap_min(a)", C1); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_max_large(bencher: divan::Bencher) { + let bm = large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&bm)], 1); + let expr = build_expr("bitmap_max(a)", C1); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_max_small(bencher: divan::Bencher) { + let bm = small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&bm)], 1); + let expr = build_expr("bitmap_max(a)", C1); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_and_large_large(bencher: divan::Bencher) { + let a = large_bitmap(); + let b = overlap_large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_and(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_and_small_small(bencher: divan::Bencher) { + let a = small_bitmap(); + let b = overlap_small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_and(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_and_large_small(bencher: divan::Bencher) { + let a = large_bitmap(); + let b = small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_and(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_and_small_large(bencher: divan::Bencher) { + let a = small_bitmap(); + let b = large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_and(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_any_large_large(bencher: divan::Bencher) { + let a = large_bitmap(); + let b = overlap_large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_any(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_any_small_small(bencher: divan::Bencher) { + let a = small_bitmap(); + let b = overlap_small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_any(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_any_large_small(bencher: divan::Bencher) { + let a = large_bitmap(); + let b = small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_any(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_any_small_large(bencher: divan::Bencher) { + let a = small_bitmap(); + let b = large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_any(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_any_large_large_disjoint(bencher: divan::Bencher) { + let a = large_bitmap(); + let b = disjoint_large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_any(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_all_large_large(bencher: divan::Bencher) { + let a = large_bitmap(); + let b = subset_large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_all(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_all_small_small(bencher: divan::Bencher) { + let a = small_bitmap(); + let b = subset_small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_all(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_all_large_small(bencher: divan::Bencher) { + let a = large_bitmap(); + let b = small_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_all(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_all_small_large(bencher: divan::Bencher) { + let a = small_bitmap(); + let b = large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_all(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } + + #[divan::bench] + fn bitmap_has_all_large_large_disjoint(bencher: divan::Bencher) { + let a = large_bitmap(); + let b = disjoint_large_bitmap(); + let block = DataBlock::new(vec![bitmap_entry(&a), bitmap_entry(&b)], 1); + let expr = build_expr("bitmap_has_all(a, b)", C2); + bencher.bench(|| eval_block(&expr, &block)); + } +} + #[divan::bench_group(max_time = 0.5)] mod datetime_fast_path { use std::sync::LazyLock; From 4f90c09ed4dbd005f6fc8c3d5bf5664aa64ee10a Mon Sep 17 00:00:00 2001 From: HarryHao Date: Mon, 27 Jul 2026 00:05:52 +0800 Subject: [PATCH 2/9] perf: add fast-path bitmap_contains that avoids full deserialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- src/common/io/src/bitmap.rs | 113 ++++++++++++++++++++-- src/common/io/src/bitmap/reader.rs | 113 ++++++++++++++++++++-- src/common/io/src/lib.rs | 1 + src/query/functions/src/scalars/bitmap.rs | 7 +- 4 files changed, 218 insertions(+), 16 deletions(-) diff --git a/src/common/io/src/bitmap.rs b/src/common/io/src/bitmap.rs index cc614cc3e3fc8..3bc73c611f6f1 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -728,23 +728,76 @@ pub fn bitmap_len(buf: &[u8]) -> Result { 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 +} + +/// Validate that a serialized bitmap buffer has minimum length. +/// +/// Different from [`validate_serialized_bitmap`], this one does not check payload to be +/// lightweight. +pub(crate) fn validate_serialized_bitmap_header_and_length(buf: &[u8]) -> Result<()> { + if buf.is_empty() { + return Ok(()); + } else if is_hybrid(buf) { + match buf[3] { + HYBRID_KIND_SMALL => { + // HybridSmall payload: at least 1 byte (length byte) + if buf.len() < HYBRID_HEADER_LEN + 1 { + return Err(ErrorCode::BadBytes("hybrid small bitmap: buffer too short")); + } + } + HYBRID_KIND_LARGE => { + // HybridLarge payload: at least 8 bytes (u64 prefix bucket count) + if buf.len() < HYBRID_HEADER_LEN + 8 { + return Err(ErrorCode::BadBytes("hybrid large bitmap: buffer too short")); + } + } + kind => { + return Err(ErrorCode::BadBytes(format!( + "hybrid bitmap: invalid kind {kind}" + ))); + } + } + } else if buf.len() < 8 { + // Legacy RoaringTreemap: minimum 8 bytes (u64 prefix bucket count) + return Err(ErrorCode::BadBytes("bitmap: buffer too short")); + } + + Ok(()) +} + +pub fn bitmap_contains(buf: &[u8], value: u64) -> Result { + if buf.is_empty() { + return Ok(false); + } + + validate_serialized_bitmap_header_and_length(buf)?; + + 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> { 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 +821,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 +1482,50 @@ mod tests { let decoded = deserialize_bitmap(&legacy).unwrap(); assert_eq!(decoded.into_iter().collect::>(), 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()); + } } diff --git a/src/common/io/src/bitmap/reader.rs b/src/common/io/src/bitmap/reader.rs index fbd2934a707b6..95b7118914b2c 100644 --- a/src/common/io/src/bitmap/reader.rs +++ b/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::()?; - 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 { + 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::()? 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 { Ok(sum) } +pub(crate) fn bitmap_contains(buf: &[u8], value: u64) -> io::Result { + 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 { + 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 { + 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 { + 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)?; diff --git a/src/common/io/src/lib.rs b/src/common/io/src/lib.rs index a3b3a8c571e0b..d7f6867f23f58 100644 --- a/src/common/io/src/lib.rs +++ b/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; diff --git a/src/query/functions/src/scalars/bitmap.rs b/src/query/functions/src/scalars/bitmap.rs index 1190e52a0c732..7c9105a9c92ad 100644 --- a/src/query/functions/src/scalars/bitmap.rs +++ b/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()); From 39c851fb46e3ec5cf1bd1342021d7acc20ccc1c1 Mon Sep 17 00:00:00 2001 From: HarryHao Date: Mon, 27 Jul 2026 01:44:28 +0800 Subject: [PATCH 3/9] perf: add fast-path bitmap_min that avoids full deserialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For large hybrid bitmaps, bitmap_min now reads the minimum value directly from the serialized buffer (first element of the first container) instead of deserializing the entire roaring treemap. This reduces bitmap_min_large from ~10.3µs to ~562ns (18x speedup). --- src/common/io/src/bitmap.rs | 59 +++++++++++++++++++++++ src/common/io/src/bitmap/reader.rs | 46 ++++++++++++++++++ src/common/io/src/lib.rs | 1 + src/query/functions/src/scalars/bitmap.rs | 15 +++--- 4 files changed, 113 insertions(+), 8 deletions(-) diff --git a/src/common/io/src/bitmap.rs b/src/common/io/src/bitmap.rs index 3bc73c611f6f1..ec180aad5fc2c 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -792,6 +792,18 @@ pub fn bitmap_contains(buf: &[u8], value: u64) -> Result { } } +pub fn bitmap_min(buf: &[u8]) -> Result> { + if buf.is_empty() { + return Ok(None); + } + validate_serialized_bitmap_header_and_length(buf)?; + if is_hybrid_large(buf) { + Ok(reader::bitmap_min(&buf[HYBRID_HEADER_LEN..])?) + } else { + Ok(deserialize_bitmap(buf)?.min()) + } +} + fn parse_bitmap_rhs(buf: &[u8]) -> Result> { if buf.is_empty() { return Ok(BitmapRhsView::Empty); @@ -1528,4 +1540,51 @@ mod tests { // Empty buffer assert!(!bitmap_contains(&[], 42).unwrap()); } + + #[test] + fn test_bitmap_min() { + // HybridLarge: spanning multiple containers + 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_eq!(bitmap_min(&buf).unwrap(), Some(0)); + + // HybridLarge with array container at max capacity (cardinality = 4096) + let boundary = HybridBitmap::from_iter(0u64..4096); + let mut buf = Vec::new(); + boundary.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_min(&buf).unwrap(), Some(0)); + + // HybridLarge with multiple containers: array (4096) + bitmap (>4096) + let mixed = HybridBitmap::from_iter((0u64..4096).chain(65536..106496)); + let mut buf = Vec::new(); + mixed.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_min(&buf).unwrap(), Some(0)); + + // HybridSmall + let small = HybridBitmap::from_iter(0u64..31); + let mut buf = Vec::new(); + small.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_min(&buf).unwrap(), Some(0)); + + // HybridSmall with non-zero start + let small = HybridBitmap::from_iter(100u64..131); + let mut buf = Vec::new(); + small.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_min(&buf).unwrap(), Some(100)); + + // Legacy + let mut tree = RoaringTreemap::new(); + tree.insert(42); + let mut buf = Vec::new(); + tree.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_min(&buf).unwrap(), Some(42)); + + // Empty buffer + assert_eq!(bitmap_min(&[]).unwrap(), None); + } } diff --git a/src/common/io/src/bitmap/reader.rs b/src/common/io/src/bitmap/reader.rs index 95b7118914b2c..038a5744cdb4a 100644 --- a/src/common/io/src/bitmap/reader.rs +++ b/src/common/io/src/bitmap/reader.rs @@ -291,6 +291,52 @@ fn bitmap_container_contains(data: &[u8], low16: u16) -> io::Result { Ok(word & (1 << bit_index) != 0) } +pub(crate) fn bitmap_min(buf: &[u8]) -> io::Result> { + let tree = TreemapReader::new(buf)?; + let bitmap = match tree.iter().next() { + None => return Ok(None), + Some(b) => b?, + }; + if bitmap.containers() == 0 { + return Ok(None); + } + let desc = bitmap.description(0)?; + let offset = bitmap.container_offset(0)?; + let container_data = &bitmap.bitmap_buf()[offset..]; + let prefix = bitmap.prefix() as u64; + let container_key = desc.prefix as u64; + let cardinality = desc.cardinality(); + let low16 = if cardinality <= ARRAY_LIMIT { + array_container_first(container_data, cardinality)? + } else { + bitmap_container_first(container_data)? + }; + Ok(Some(prefix << 32 | container_key << 16 | low16 as u64)) +} + +fn array_container_first(data: &[u8], cardinality: usize) -> io::Result { + if data.len() < cardinality * 2 { + return Err(Error::other("array container too short")); + } + Ok(u16::from_le_bytes(data[0..2].try_into().unwrap())) +} + +fn bitmap_container_first(data: &[u8]) -> io::Result { + if data.len() < BITMAP_BYTES { + return Err(Error::other("bitmap container too short")); + } + // Find the lowest set bit in the 1024-word bitmap + for word_index in 0..BITMAP_WORDS { + let start = word_index * WORD_BYTES; + let word = u64::from_le_bytes(data[start..start + WORD_BYTES].try_into().unwrap()); + if word != 0 { + return Ok((word_index * WORD_BITS + word.trailing_zeros() as usize) as u16); + } + } + // All zeros — shouldn't happen for a valid bitmap container + Err(Error::other("bitmap container has no set bits")) +} + pub fn intersection_with_serialized(tree: &mut RoaringTreemap, buf: &[u8]) -> io::Result<()> { use std::cmp::Ordering::*; let rhs = TreemapReader::new(buf)?; diff --git a/src/common/io/src/lib.rs b/src/common/io/src/lib.rs index d7f6867f23f58..29e451816d4bb 100644 --- a/src/common/io/src/lib.rs +++ b/src/common/io/src/lib.rs @@ -54,6 +54,7 @@ pub use bitmap::HYBRID_VERSION; pub use bitmap::HybridBitmap; pub use bitmap::LARGE_THRESHOLD; pub use bitmap::bitmap_contains; +pub use bitmap::bitmap_min; pub use bitmap::deserialize_bitmap; pub use bitmap::parse_bitmap; pub use decimal::display_decimal_128; diff --git a/src/query/functions/src/scalars/bitmap.rs b/src/query/functions/src/scalars/bitmap.rs index 7c9105a9c92ad..a14ef98737d3e 100644 --- a/src/query/functions/src/scalars/bitmap.rs +++ b/src/query/functions/src/scalars/bitmap.rs @@ -39,6 +39,7 @@ 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::bitmap::bitmap_min; use databend_common_io::deserialize_bitmap; use databend_common_io::parse_bitmap; use itertools::join; @@ -481,14 +482,12 @@ pub fn register(registry: &mut FunctionRegistry) { builder.push(0); return; } - let val = match deserialize_bitmap(b) { - Ok(rb) => match rb.min() { - Some(val) => val, - None => { - ctx.set_error(builder.len(), "The bitmap is empty"); - 0 - } - }, + let val = match bitmap_min(b) { + Ok(Some(val)) => val, + Ok(None) => { + ctx.set_error(builder.len(), "The bitmap is empty"); + 0 + } Err(e) => { ctx.set_error(builder.len(), e.to_string()); 0 From 38510ebb5da88dea605e07f097688105d1de6383 Mon Sep 17 00:00:00 2001 From: HarryHao Date: Mon, 27 Jul 2026 08:04:10 +0800 Subject: [PATCH 4/9] perf: add fast-path bitmap_max that avoids full deserialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- src/common/io/src/bitmap.rs | 60 +++++++++++++++++++++++ src/common/io/src/bitmap/reader.rs | 55 +++++++++++++++++++++ src/common/io/src/lib.rs | 1 + src/query/functions/src/scalars/bitmap.rs | 15 +++--- 4 files changed, 123 insertions(+), 8 deletions(-) diff --git a/src/common/io/src/bitmap.rs b/src/common/io/src/bitmap.rs index ec180aad5fc2c..6e6eb6c9a7eaa 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -804,6 +804,18 @@ pub fn bitmap_min(buf: &[u8]) -> Result> { } } +pub fn bitmap_max(buf: &[u8]) -> Result> { + if buf.is_empty() { + return Ok(None); + } + validate_serialized_bitmap_header_and_length(buf)?; + if is_hybrid_large(buf) { + Ok(reader::bitmap_max(&buf[HYBRID_HEADER_LEN..])?) + } else { + Ok(deserialize_bitmap(buf)?.max()) + } +} + fn parse_bitmap_rhs(buf: &[u8]) -> Result> { if buf.is_empty() { return Ok(BitmapRhsView::Empty); @@ -1587,4 +1599,52 @@ mod tests { // Empty buffer assert_eq!(bitmap_min(&[]).unwrap(), None); } + + #[test] + fn test_bitmap_max() { + // HybridLarge: spanning multiple containers + 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_eq!(bitmap_max(&buf).unwrap(), Some(239999)); + + // HybridLarge with array container at max capacity (cardinality = 4096) + let boundary = HybridBitmap::from_iter(0u64..4096); + let mut buf = Vec::new(); + boundary.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_max(&buf).unwrap(), Some(4095)); + + // HybridLarge with multiple containers: array (4096) + bitmap (>4096) + let mixed = HybridBitmap::from_iter((0u64..4096).chain(65536..106496)); + let mut buf = Vec::new(); + mixed.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_max(&buf).unwrap(), Some(106495)); + + // HybridSmall + let small = HybridBitmap::from_iter(0u64..31); + let mut buf = Vec::new(); + small.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_max(&buf).unwrap(), Some(30)); + + // HybridSmall with different range + let small = HybridBitmap::from_iter(100u64..131); + let mut buf = Vec::new(); + small.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_max(&buf).unwrap(), Some(130)); + + // Legacy + let mut tree = RoaringTreemap::new(); + tree.insert(42); + tree.insert(100); + let mut buf = Vec::new(); + tree.serialize_into(&mut buf).unwrap(); + assert_eq!(bitmap_max(&buf).unwrap(), Some(100)); + + // Empty buffer + assert_eq!(bitmap_max(&[]).unwrap(), None); + } } diff --git a/src/common/io/src/bitmap/reader.rs b/src/common/io/src/bitmap/reader.rs index 038a5744cdb4a..dfd5d79e7848c 100644 --- a/src/common/io/src/bitmap/reader.rs +++ b/src/common/io/src/bitmap/reader.rs @@ -337,6 +337,61 @@ fn bitmap_container_first(data: &[u8]) -> io::Result { Err(Error::other("bitmap container has no set bits")) } +pub(crate) fn bitmap_max(buf: &[u8]) -> io::Result> { + let tree = TreemapReader::new(buf)?; + // Find the last prefix bucket + let mut last_bitmap = None; + for bitmap_result in tree.iter() { + last_bitmap = Some(bitmap_result?); + } + let bitmap = match last_bitmap { + None => return Ok(None), + Some(b) => b, + }; + if bitmap.containers() == 0 { + return Ok(None); + } + let last_idx = bitmap.containers() - 1; + let desc = bitmap.description(last_idx)?; + let offset = bitmap.container_offset(last_idx)?; + let container_data = &bitmap.bitmap_buf()[offset..]; + let prefix = bitmap.prefix() as u64; + let container_key = desc.prefix as u64; + let cardinality = desc.cardinality(); + let low16 = if cardinality <= ARRAY_LIMIT { + array_container_last(container_data, cardinality)? + } else { + bitmap_container_last(container_data)? + }; + Ok(Some(prefix << 32 | container_key << 16 | low16 as u64)) +} + +fn array_container_last(data: &[u8], cardinality: usize) -> io::Result { + if data.len() < cardinality * 2 { + return Err(Error::other("array container too short")); + } + let offset = (cardinality - 1) * 2; + Ok(u16::from_le_bytes( + data[offset..offset + 2].try_into().unwrap(), + )) +} + +fn bitmap_container_last(data: &[u8]) -> io::Result { + if data.len() < BITMAP_BYTES { + return Err(Error::other("bitmap container too short")); + } + // Find the highest set bit by scanning from the last word backwards + for word_index in (0..BITMAP_WORDS).rev() { + let start = word_index * WORD_BYTES; + let word = u64::from_le_bytes(data[start..start + WORD_BYTES].try_into().unwrap()); + if word != 0 { + let bit_index = WORD_BITS - 1 - word.leading_zeros() as usize; + return Ok((word_index * WORD_BITS + bit_index) as u16); + } + } + Err(Error::other("bitmap container has no set bits")) +} + pub fn intersection_with_serialized(tree: &mut RoaringTreemap, buf: &[u8]) -> io::Result<()> { use std::cmp::Ordering::*; let rhs = TreemapReader::new(buf)?; diff --git a/src/common/io/src/lib.rs b/src/common/io/src/lib.rs index 29e451816d4bb..4cffe99c4e924 100644 --- a/src/common/io/src/lib.rs +++ b/src/common/io/src/lib.rs @@ -54,6 +54,7 @@ pub use bitmap::HYBRID_VERSION; pub use bitmap::HybridBitmap; pub use bitmap::LARGE_THRESHOLD; pub use bitmap::bitmap_contains; +pub use bitmap::bitmap_max; pub use bitmap::bitmap_min; pub use bitmap::deserialize_bitmap; pub use bitmap::parse_bitmap; diff --git a/src/query/functions/src/scalars/bitmap.rs b/src/query/functions/src/scalars/bitmap.rs index a14ef98737d3e..2157f41b456c0 100644 --- a/src/query/functions/src/scalars/bitmap.rs +++ b/src/query/functions/src/scalars/bitmap.rs @@ -39,6 +39,7 @@ 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::bitmap::bitmap_max; use databend_common_io::bitmap::bitmap_min; use databend_common_io::deserialize_bitmap; use databend_common_io::parse_bitmap; @@ -450,14 +451,12 @@ pub fn register(registry: &mut FunctionRegistry) { builder.push(0); return; } - let val = match deserialize_bitmap(b) { - Ok(rb) => match rb.max() { - Some(val) => val, - None => { - ctx.set_error(builder.len(), "The bitmap is empty"); - 0 - } - }, + let val = match bitmap_max(b) { + Ok(Some(val)) => val, + Ok(None) => { + ctx.set_error(builder.len(), "The bitmap is empty"); + 0 + } Err(e) => { ctx.set_error(builder.len(), e.to_string()); 0 From 4bfba50f7b8570a61acc88b924d8720748ad87a2 Mon Sep 17 00:00:00 2001 From: HarryHao Date: Mon, 27 Jul 2026 08:43:32 +0800 Subject: [PATCH 5/9] perf: add fast-path bitmap_has_any that avoids full deserialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For large hybrid bitmaps, bitmap_has_any now checks intersection directly on the serialized buffers using container-level binary search instead of deserializing both bitmaps and computing full intersection. This reduces bitmap_has_any_large_large from ~26µs to ~2.5µs (10x), and disjoint cases from ~20µs to ~963ns (21x). --- src/common/io/src/bitmap.rs | 289 ++++++++++++++++++++++ src/common/io/src/bitmap/reader.rs | 231 +++++++++++++++++ src/common/io/src/lib.rs | 3 + src/query/functions/src/scalars/bitmap.rs | 17 +- 4 files changed, 528 insertions(+), 12 deletions(-) diff --git a/src/common/io/src/bitmap.rs b/src/common/io/src/bitmap.rs index 6e6eb6c9a7eaa..95cfaddbe76fe 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -28,6 +28,8 @@ use roaring::RoaringTreemap; use roaring::treemap::Iter; use smallvec::SmallVec; +pub use crate::bitmap::reader::BitmapStats; + mod reader; // https://github.com/ClickHouse/ClickHouse/blob/516a6ed6f8bd8c5f6eed3a10e9037580b2fb6152/src/AggregateFunctions/AggregateFunctionGroupBitmapData.h#L914 @@ -816,6 +818,102 @@ pub fn bitmap_max(buf: &[u8]) -> Result> { } } +pub fn bitmap_stats(buf: &[u8]) -> Result { + if buf.is_empty() { + return Ok(BitmapStats { + len: 0, + min: None, + max: None, + }); + } + if is_hybrid_large(buf) { + Ok(reader::bitmap_stats(&buf[HYBRID_HEADER_LEN..])?) + } else { + let bm = deserialize_bitmap(buf)?; + Ok(BitmapStats { + len: bm.len() as u64, + min: bm.min(), + max: bm.max(), + }) + } +} + +pub fn bitmap_has_any(lhs: &[u8], rhs: &[u8]) -> Result { + if lhs.is_empty() || rhs.is_empty() { + return Ok(false); + } + + validate_serialized_bitmap_header_and_length(lhs)?; + validate_serialized_bitmap_header_and_length(rhs)?; + + // Stats fast path: range rejection + // If lhs.max < rhs.min OR rhs.max < lhs.min, no overlap possible + if is_hybrid_large(lhs) && is_hybrid_large(rhs) { + let lhs_stats = reader::bitmap_stats(&lhs[HYBRID_HEADER_LEN..])?; + let rhs_stats = reader::bitmap_stats(&rhs[HYBRID_HEADER_LEN..])?; + if let (Some(lhs_min), Some(rhs_max)) = (lhs_stats.min, rhs_stats.max) + && lhs_min > rhs_max + { + return Ok(false); + } + if let (Some(rhs_min), Some(lhs_max)) = (rhs_stats.min, lhs_stats.max) + && rhs_min > lhs_max + { + return Ok(false); + } + Ok(reader::bitmap_has_any( + &lhs[HYBRID_HEADER_LEN..], + &rhs[HYBRID_HEADER_LEN..], + )?) + } else if is_hybrid_large(lhs) || is_hybrid_large(rhs) { + // Normalize: large is the HybridLarge side, other is the non-HybridLarge side + let (large, other) = if is_hybrid_large(lhs) { + (lhs, rhs) + } else { + (rhs, lhs) + }; + + // Stats fast path: range rejection + let large_stats = reader::bitmap_stats(&large[HYBRID_HEADER_LEN..])?; + let other_stats = bitmap_stats(other)?; + if let (Some(large_min), Some(other_max)) = (large_stats.min, other_stats.max) + && large_min > other_max + { + return Ok(false); + } + if let (Some(other_min), Some(large_max)) = (other_stats.min, large_stats.max) + && other_min > large_max + { + return Ok(false); + } + + // other is HybridSmall: iterate its (≤31) values through bitmap_contains on large. + // Each bitmap_contains call scans prefix buckets from the start, causing redundant + // work across iterations. This could be improved with a hint/cursor mechanism that + // tracks the last scanned position, but since HybridSmall holds at most 31 values + // the overhead is bounded (~μs-level). Revisit if this path becomes a bottleneck. + if is_hybrid(other) && !is_hybrid_large(other) { + let other_bm = deserialize_bitmap(other)?; + for value in other_bm.iter() { + if bitmap_contains(large, value)? { + return Ok(true); + } + } + return Ok(false); + } + + // other is Legacy (potentially large): fallback to full deserialization + let lhs_bm = deserialize_bitmap(lhs)?; + let rhs_bm = deserialize_bitmap(rhs)?; + Ok(lhs_bm.intersection_len(&rhs_bm) != 0) + } else { + // Both else: deserialize both + let lhs_bm = deserialize_bitmap(lhs)?; + let rhs_bm = deserialize_bitmap(rhs)?; + Ok(lhs_bm.intersection_len(&rhs_bm) != 0) + } +} + fn parse_bitmap_rhs(buf: &[u8]) -> Result> { if buf.is_empty() { return Ok(BitmapRhsView::Empty); @@ -1647,4 +1745,195 @@ mod tests { // Empty buffer assert_eq!(bitmap_max(&[]).unwrap(), None); } + + fn ground_truth_has_any(lhs: &[u8], rhs: &[u8]) -> bool { + let lhs_bm = deserialize_bitmap(lhs).unwrap(); + let rhs_bm = deserialize_bitmap(rhs).unwrap(); + lhs_bm.intersection_len(&rhs_bm) != 0 + } + + #[test] + fn test_bitmap_has_any() { + // Both HybridLarge with overlap: 0..50000 and 40000..90000 -> overlap at 40000-49999 + let lhs = HybridBitmap::from_iter(0u64..50000); + let rhs = HybridBitmap::from_iter(40000u64..90000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + + // Both HybridLarge disjoint: 0..50000 and 100000..150000 + let lhs = HybridBitmap::from_iter(0u64..50000); + let rhs = HybridBitmap::from_iter(100000u64..150000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(!bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + + // Both HybridSmall with overlap: 0..31 and 20..51 -> overlap at 20-30 + let lhs = HybridBitmap::from_iter(0u64..31); + let rhs = HybridBitmap::from_iter(20u64..51); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + + // Both HybridSmall disjoint: 0..15 and 20..35 + let lhs = HybridBitmap::from_iter(0u64..15); + let rhs = HybridBitmap::from_iter(20u64..35); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(!bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + + // Mixed (HybridLarge lhs, HybridSmall rhs) with overlap: 0..50000 and 10..41 + let lhs = HybridBitmap::from_iter(0u64..50000); + let rhs = HybridBitmap::from_iter(10u64..41); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + + // Mixed (HybridSmall lhs, HybridLarge rhs) with overlap: 10..41 and 0..50000 + let lhs = HybridBitmap::from_iter(10u64..41); + let rhs = HybridBitmap::from_iter(0u64..50000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + + // Empty lhs + let rhs = HybridBitmap::from_iter(0u64..50000); + let mut rhs_buf = Vec::new(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert!(!bitmap_has_any(&[], &rhs_buf).unwrap()); + + // Empty rhs + let lhs = HybridBitmap::from_iter(0u64..50000); + let mut lhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + assert!(!bitmap_has_any(&lhs_buf, &[]).unwrap()); + + // Both empty + assert!(!bitmap_has_any(&[], &[]).unwrap()); + + // Mixed (Legacy lhs, HybridLarge rhs) with overlap + let mut lhs_tree = RoaringTreemap::new(); + for v in [0u64, 2500, 40000] { + lhs_tree.insert(v); + } + let rhs = HybridBitmap::from_iter(0u64..50000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs_tree.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + + // Mixed (HybridLarge lhs, Legacy rhs) with overlap + let lhs = HybridBitmap::from_iter(0u64..50000); + let mut rhs_tree = RoaringTreemap::new(); + for v in [0u64, 2500, 40000] { + rhs_tree.insert(v); + } + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs_tree.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + + // Both Legacy with overlap + let mut lhs_tree = RoaringTreemap::new(); + for v in 0u64..100 { + lhs_tree.insert(v); + } + let mut rhs_tree = RoaringTreemap::new(); + for v in 50u64..150 { + rhs_tree.insert(v); + } + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs_tree.serialize_into(&mut lhs_buf).unwrap(); + rhs_tree.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_any(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + } + + #[test] + fn test_bitmap_stats() { + // HybridLarge stats + let large = HybridBitmap::from_iter(0u64..50000); + let mut buf = Vec::new(); + large.serialize_into(&mut buf).unwrap(); + let stats = bitmap_stats(&buf).unwrap(); + assert_eq!(stats.len, 50000); + assert_eq!(stats.min, Some(0)); + assert_eq!(stats.max, Some(49999)); + + // HybridSmall stats + let small = HybridBitmap::from_iter(100u64..131); + let mut buf = Vec::new(); + small.serialize_into(&mut buf).unwrap(); + let stats = bitmap_stats(&buf).unwrap(); + assert_eq!(stats.len, 31); + assert_eq!(stats.min, Some(100)); + assert_eq!(stats.max, Some(130)); + + // Empty stats + let stats = bitmap_stats(&[]).unwrap(); + assert_eq!(stats.len, 0); + assert_eq!(stats.min, None); + assert_eq!(stats.max, None); + + // Legacy stats + let mut tree = RoaringTreemap::new(); + tree.insert(42); + tree.insert(100); + let mut buf = Vec::new(); + tree.serialize_into(&mut buf).unwrap(); + let stats = bitmap_stats(&buf).unwrap(); + assert_eq!(stats.len, 2); + assert_eq!(stats.min, Some(42)); + assert_eq!(stats.max, Some(100)); + } } diff --git a/src/common/io/src/bitmap/reader.rs b/src/common/io/src/bitmap/reader.rs index dfd5d79e7848c..75d4b8a6680a8 100644 --- a/src/common/io/src/bitmap/reader.rs +++ b/src/common/io/src/bitmap/reader.rs @@ -427,6 +427,237 @@ pub fn intersection_with_serialized(tree: &mut RoaringTreemap, buf: &[u8]) -> io Ok(()) } +pub struct BitmapStats { + pub len: u64, + pub min: Option, + pub max: Option, +} + +pub(crate) fn bitmap_stats(buf: &[u8]) -> io::Result { + let tree = TreemapReader::new(buf)?; + + let mut total_len = 0u64; + let mut first_bitmap = None; + let mut last_bitmap = None; + + for bitmap_result in tree.iter() { + let bitmap = bitmap_result?; + for i in 0..bitmap.containers() { + total_len += bitmap.description(i)?.cardinality() as u64; + } + if first_bitmap.is_none() { + first_bitmap = Some(bitmap.clone()); + } + last_bitmap = Some(bitmap); + } + + if first_bitmap.is_none() || total_len == 0 { + return Ok(BitmapStats { + len: 0, + min: None, + max: None, + }); + } + + // Compute min from the first container of the first prefix bucket + let first = first_bitmap.unwrap(); + let first_desc = first.description(0)?; + let first_offset = first.container_offset(0)?; + let first_container_data = &first.bitmap_buf()[first_offset..]; + let first_prefix = first.prefix() as u64; + let first_container_key = first_desc.prefix as u64; + let first_cardinality = first_desc.cardinality(); + let first_low16 = if first_cardinality <= ARRAY_LIMIT { + array_container_first(first_container_data, first_cardinality)? + } else { + bitmap_container_first(first_container_data)? + }; + let min_val = first_prefix << 32 | first_container_key << 16 | first_low16 as u64; + + // Compute max from the last container of the last prefix bucket + let last = last_bitmap.unwrap(); + let last_idx = last.containers() - 1; + let last_desc = last.description(last_idx)?; + let last_offset = last.container_offset(last_idx)?; + let last_container_data = &last.bitmap_buf()[last_offset..]; + let last_prefix = last.prefix() as u64; + let last_container_key = last_desc.prefix as u64; + let last_cardinality = last_desc.cardinality(); + let last_low16 = if last_cardinality <= ARRAY_LIMIT { + array_container_last(last_container_data, last_cardinality)? + } else { + bitmap_container_last(last_container_data)? + }; + let max_val = last_prefix << 32 | last_container_key << 16 | last_low16 as u64; + + Ok(BitmapStats { + len: total_len, + min: Some(min_val), + max: Some(max_val), + }) +} + +pub(crate) fn bitmap_has_any(lhs: &[u8], rhs: &[u8]) -> io::Result { + use std::cmp::Ordering::*; + let lhs_tree = TreemapReader::new(lhs)?; + let rhs_tree = TreemapReader::new(rhs)?; + + let mut lhs_iter = lhs_tree.iter(); + let mut rhs_iter = rhs_tree.iter(); + + let mut lhs_curr = lhs_iter.next().transpose()?; + let mut rhs_curr = rhs_iter.next().transpose()?; + + while let (Some(lhs_bitmap), Some(rhs_bitmap)) = (lhs_curr.as_ref(), rhs_curr.as_ref()) { + match lhs_bitmap.prefix().cmp(&rhs_bitmap.prefix()) { + Less => { + lhs_curr = lhs_iter.next().transpose()?; + } + Greater => { + rhs_curr = rhs_iter.next().transpose()?; + } + Equal => { + if containers_has_any(lhs_bitmap, rhs_bitmap)? { + return Ok(true); + } + lhs_curr = lhs_iter.next().transpose()?; + rhs_curr = rhs_iter.next().transpose()?; + } + } + } + + Ok(false) +} + +fn containers_has_any(lhs: &BitmapReader<'_>, rhs: &BitmapReader<'_>) -> io::Result { + let lhs_count = lhs.containers(); + let rhs_count = rhs.containers(); + let mut i = 0; + let mut j = 0; + + while i < lhs_count && j < rhs_count { + let lhs_desc = lhs.description(i)?; + let rhs_desc = rhs.description(j)?; + match lhs_desc.prefix.cmp(&rhs_desc.prefix) { + std::cmp::Ordering::Less => i += 1, + std::cmp::Ordering::Greater => j += 1, + std::cmp::Ordering::Equal => { + let lhs_offset = lhs.container_offset(i)?; + let rhs_offset = rhs.container_offset(j)?; + let lhs_data = &lhs.bitmap_buf()[lhs_offset..]; + let rhs_data = &rhs.bitmap_buf()[rhs_offset..]; + let lhs_cardinality = lhs_desc.cardinality(); + let rhs_cardinality = rhs_desc.cardinality(); + if container_has_any(lhs_data, lhs_cardinality, rhs_data, rhs_cardinality)? { + return Ok(true); + } + i += 1; + j += 1; + } + } + } + + Ok(false) +} + +fn container_has_any( + lhs_data: &[u8], + lhs_cardinality: usize, + rhs_data: &[u8], + rhs_cardinality: usize, +) -> io::Result { + // Pigeonhole: two subsets of a 65536-element space with combined size > 65536 must overlap + if lhs_cardinality + rhs_cardinality > 65536 { + return Ok(true); + } + + let lhs_is_array = lhs_cardinality <= ARRAY_LIMIT; + let rhs_is_array = rhs_cardinality <= ARRAY_LIMIT; + + if lhs_is_array && rhs_is_array { + array_container_has_any(lhs_data, lhs_cardinality, rhs_data, rhs_cardinality) + } else if !lhs_is_array && !rhs_is_array { + bitmap_container_has_any(lhs_data, rhs_data) + } else if lhs_is_array { + array_bitmap_container_has_any(lhs_data, lhs_cardinality, rhs_data) + } else { + array_bitmap_container_has_any(rhs_data, rhs_cardinality, lhs_data) + } +} + +fn array_container_has_any( + lhs_data: &[u8], + lhs_cardinality: usize, + rhs_data: &[u8], + rhs_cardinality: usize, +) -> io::Result { + if lhs_data.len() < lhs_cardinality * 2 { + return Err(Error::other("lhs array container too short")); + } + if rhs_data.len() < rhs_cardinality * 2 { + return Err(Error::other("rhs array container too short")); + } + + let mut i = 0; + let mut j = 0; + while i < lhs_cardinality && j < rhs_cardinality { + let lv = u16::from_le_bytes(lhs_data[i * 2..i * 2 + 2].try_into().unwrap()); + let rv = u16::from_le_bytes(rhs_data[j * 2..j * 2 + 2].try_into().unwrap()); + if lv < rv { + i += 1; + } else if rv < lv { + j += 1; + } else { + return Ok(true); + } + } + Ok(false) +} + +fn bitmap_container_has_any(lhs_data: &[u8], rhs_data: &[u8]) -> io::Result { + if lhs_data.len() < BITMAP_BYTES { + return Err(Error::other("lhs bitmap container too short")); + } + if rhs_data.len() < BITMAP_BYTES { + return Err(Error::other("rhs bitmap container too short")); + } + + for word_index in 0..BITMAP_WORDS { + let start = word_index * WORD_BYTES; + let lhs_word = u64::from_le_bytes(lhs_data[start..start + WORD_BYTES].try_into().unwrap()); + let rhs_word = u64::from_le_bytes(rhs_data[start..start + WORD_BYTES].try_into().unwrap()); + if lhs_word & rhs_word != 0 { + return Ok(true); + } + } + Ok(false) +} + +fn array_bitmap_container_has_any( + array_data: &[u8], + array_cardinality: usize, + bitmap_data: &[u8], +) -> io::Result { + if array_data.len() < array_cardinality * 2 { + return Err(Error::other("array container too short")); + } + if bitmap_data.len() < BITMAP_BYTES { + return Err(Error::other("bitmap container too short")); + } + + for i in 0..array_cardinality { + let value = u16::from_le_bytes(array_data[i * 2..i * 2 + 2].try_into().unwrap()); + let word_index = value as usize / WORD_BITS; + let bit_index = value as usize % WORD_BITS; + let start = word_index * WORD_BYTES; + let word = u64::from_le_bytes(bitmap_data[start..start + WORD_BYTES].try_into().unwrap()); + if word & (1 << bit_index) != 0 { + return Ok(true); + } + } + Ok(false) +} + #[cfg(test)] mod tests { diff --git a/src/common/io/src/lib.rs b/src/common/io/src/lib.rs index 4cffe99c4e924..371f248fee9e4 100644 --- a/src/common/io/src/lib.rs +++ b/src/common/io/src/lib.rs @@ -46,6 +46,7 @@ mod stat_buffer; pub mod interval; pub mod wkb; +pub use bitmap::BitmapStats; pub use bitmap::HYBRID_HEADER_LEN; pub use bitmap::HYBRID_KIND_LARGE; pub use bitmap::HYBRID_KIND_SMALL; @@ -54,8 +55,10 @@ pub use bitmap::HYBRID_VERSION; pub use bitmap::HybridBitmap; pub use bitmap::LARGE_THRESHOLD; pub use bitmap::bitmap_contains; +pub use bitmap::bitmap_has_any; pub use bitmap::bitmap_max; pub use bitmap::bitmap_min; +pub use bitmap::bitmap_stats; pub use bitmap::deserialize_bitmap; pub use bitmap::parse_bitmap; pub use decimal::display_decimal_128; diff --git a/src/query/functions/src/scalars/bitmap.rs b/src/query/functions/src/scalars/bitmap.rs index 2157f41b456c0..a9fe63c84c322 100644 --- a/src/query/functions/src/scalars/bitmap.rs +++ b/src/query/functions/src/scalars/bitmap.rs @@ -38,6 +38,7 @@ 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_has_any; use databend_common_io::bitmap::bitmap_len; use databend_common_io::bitmap::bitmap_max; use databend_common_io::bitmap::bitmap_min; @@ -416,23 +417,15 @@ pub fn register(registry: &mut FunctionRegistry) { builder.push(false); return; } - let rb = match deserialize_bitmap(b) { - Ok(rb) => rb, - Err(e) => { - ctx.set_error(builder.len(), e.to_string()); - builder.push(false); - return; + match bitmap_has_any(b, items) { + Ok(has_any) => { + builder.push(has_any); } - }; - let rb2 = match deserialize_bitmap(items) { - Ok(rb) => rb, Err(e) => { ctx.set_error(builder.len(), e.to_string()); builder.push(false); - return; } - }; - builder.push(rb.intersection_len(&rb2) != 0); + } }, ), ); From 4ba28b9f08cd06232b1aec09ab5b07ced07ce7be Mon Sep 17 00:00:00 2001 From: HarryHao Date: Mon, 27 Jul 2026 10:27:15 +0800 Subject: [PATCH 6/9] perf: add fast-path bitmap_has_all that avoids full deserialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For large hybrid bitmaps, bitmap_has_all now checks superset relationships directly on serialized buffers using container-level comparison instead of deserializing both bitmaps. This reduces bitmap_has_all_large_large from ~39µs to ~11µs (3.5x), and disjoint cases from ~20µs to ~932ns (22x). --- src/common/io/src/bitmap.rs | 276 ++++++++++++++++++++++ src/common/io/src/bitmap/reader.rs | 194 +++++++++++++++ src/common/io/src/lib.rs | 1 + src/query/functions/src/scalars/bitmap.rs | 17 +- 4 files changed, 476 insertions(+), 12 deletions(-) diff --git a/src/common/io/src/bitmap.rs b/src/common/io/src/bitmap.rs index 95cfaddbe76fe..026ce37537d8f 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -838,6 +838,82 @@ pub fn bitmap_stats(buf: &[u8]) -> Result { } } +pub fn bitmap_has_all(lhs: &[u8], rhs: &[u8]) -> Result { + if rhs.is_empty() { + return Ok(true); + } + if lhs.is_empty() { + return Ok(false); + } + validate_serialized_bitmap_header_and_length(lhs)?; + validate_serialized_bitmap_header_and_length(rhs)?; + // Both HybridLarge: use stats for early rejection, then check directly on serialized buffers. + if is_hybrid_large(lhs) && is_hybrid_large(rhs) { + let lhs_stats = reader::bitmap_stats(&lhs[HYBRID_HEADER_LEN..])?; + let rhs_stats = reader::bitmap_stats(&rhs[HYBRID_HEADER_LEN..])?; + if let (Some(lhs_min), Some(rhs_min)) = (lhs_stats.min, rhs_stats.min) + && lhs_min > rhs_min + { + return Ok(false); + } + if let (Some(rhs_max), Some(lhs_max)) = (rhs_stats.max, lhs_stats.max) + && rhs_max > lhs_max + { + return Ok(false); + } + if lhs_stats.len < rhs_stats.len { + return Ok(false); + } + Ok(reader::bitmap_has_all( + &lhs[HYBRID_HEADER_LEN..], + &rhs[HYBRID_HEADER_LEN..], + )?) + } else if is_hybrid_large(lhs) || is_hybrid_large(rhs) { + // Mixed path: one side is HybridLarge, the other is HybridSmall or Legacy. + // Use stats for early rejection; if rhs is HybridSmall, iterate and probe lhs; + // otherwise fall back to full deserialization. + let lhs_stats = bitmap_stats(lhs)?; + let rhs_stats = bitmap_stats(rhs)?; + if let (Some(lhs_min), Some(rhs_min)) = (lhs_stats.min, rhs_stats.min) + && lhs_min > rhs_min + { + return Ok(false); + } + if let (Some(rhs_max), Some(lhs_max)) = (rhs_stats.max, lhs_stats.max) + && rhs_max > lhs_max + { + return Ok(false); + } + // lhs.len < rhs.len -> lhs cannot contain all of rhs + if lhs_stats.len < rhs_stats.len { + return Ok(false); + } + // rhs is HybridSmall: iterate its (≤31) values and probe lhs via bitmap_contains. + // Each bitmap_contains call scans prefix buckets from the start, causing redundant + // work across iterations. This could be improved with a hint/cursor mechanism that + // tracks the last scanned position, but since HybridSmall holds at most 31 values + // the overhead is bounded (~μs-level). Revisit if this path becomes a bottleneck. + if is_hybrid_large(lhs) && is_hybrid(rhs) && !is_hybrid_large(rhs) { + let rhs_bm = deserialize_bitmap(rhs)?; + for value in rhs_bm.iter() { + if !bitmap_contains(lhs, value)? { + return Ok(false); + } + } + return Ok(true); + } + // Fallback to deserialization and is_superset + let lhs_bm = deserialize_bitmap(lhs)?; + let rhs_bm = deserialize_bitmap(rhs)?; + Ok(lhs_bm.is_superset(&rhs_bm)) + } else { + // Both else: deserialize both + let lhs_bm = deserialize_bitmap(lhs)?; + let rhs_bm = deserialize_bitmap(rhs)?; + Ok(lhs_bm.is_superset(&rhs_bm)) + } +} + pub fn bitmap_has_any(lhs: &[u8], rhs: &[u8]) -> Result { if lhs.is_empty() || rhs.is_empty() { return Ok(false); @@ -1936,4 +2012,204 @@ mod tests { assert_eq!(stats.min, Some(42)); assert_eq!(stats.max, Some(100)); } + + fn ground_truth_has_all(lhs: &[u8], rhs: &[u8]) -> bool { + let lhs_bm = deserialize_bitmap(lhs).unwrap(); + let rhs_bm = deserialize_bitmap(rhs).unwrap(); + lhs_bm.is_superset(&rhs_bm) + } + + #[test] + fn test_bitmap_has_all() { + // Both HybridLarge, rhs is subset: 0..50000 contains 0..10000 + let lhs = HybridBitmap::from_iter(0u64..50000); + let rhs = HybridBitmap::from_iter(0u64..10000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Both HybridLarge, rhs is not subset: 0..50000 vs 40000..90000 + // rhs has values 50000-89999 that are not in lhs + let lhs = HybridBitmap::from_iter(0u64..50000); + let rhs = HybridBitmap::from_iter(40000u64..90000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Both HybridLarge, disjoint: 0..50000 vs 100000..150000 + let lhs = HybridBitmap::from_iter(0u64..50000); + let rhs = HybridBitmap::from_iter(100000u64..150000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Both HybridSmall, rhs is subset: 0..31 contains 5..15 + let lhs = HybridBitmap::from_iter(0u64..31); + let rhs = HybridBitmap::from_iter(5u64..15); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Both HybridSmall, rhs is not subset: 0..15 vs 10..25 + // rhs has values 15-24 that are not in lhs + let lhs = HybridBitmap::from_iter(0u64..15); + let rhs = HybridBitmap::from_iter(10u64..25); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Mixed (HybridLarge lhs, HybridSmall rhs) with superset: 0..50000 contains 5..15 + let lhs = HybridBitmap::from_iter(0u64..50000); + let rhs = HybridBitmap::from_iter(5u64..15); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Mixed (HybridLarge lhs, HybridSmall rhs not subset): 0..50000 vs 49990..50010 + // rhs has values 50000-50009 that are not in lhs + let lhs = HybridBitmap::from_iter(0u64..50000); + let rhs = HybridBitmap::from_iter(49990u64..50010); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Mixed (HybridSmall lhs, HybridLarge rhs): already rejected by len fast path + let lhs = HybridBitmap::from_iter(0u64..31); + let rhs = HybridBitmap::from_iter(0u64..50000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Empty rhs: always true + let lhs = HybridBitmap::from_iter(0u64..50000); + let mut lhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + assert!(bitmap_has_all(&lhs_buf, &[]).unwrap()); + + // Empty lhs with non-empty rhs: always false + let rhs = HybridBitmap::from_iter(0u64..50000); + let mut rhs_buf = Vec::new(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert!(!bitmap_has_all(&[], &rhs_buf).unwrap()); + + // Both empty: true + assert!(bitmap_has_all(&[], &[]).unwrap()); + + // Mixed (Legacy lhs, HybridLarge rhs) + let mut lhs_tree = RoaringTreemap::new(); + for v in 0u64..50000 { + lhs_tree.insert(v); + } + let rhs = HybridBitmap::from_iter(0u64..10000); + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs_tree.serialize_into(&mut lhs_buf).unwrap(); + rhs.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + + // Mixed (HybridLarge lhs, Legacy rhs) + let lhs = HybridBitmap::from_iter(0u64..10000); + let mut rhs_tree = RoaringTreemap::new(); + for v in 0u64..50000 { + rhs_tree.insert(v); + } + + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs.serialize_into(&mut lhs_buf).unwrap(); + rhs_tree.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + + // Both Legacy + let mut lhs_tree = RoaringTreemap::new(); + for v in 0u64..100 { + lhs_tree.insert(v); + } + let mut rhs_tree = RoaringTreemap::new(); + for v in 50u64..80 { + rhs_tree.insert(v); + } + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs_tree.serialize_into(&mut lhs_buf).unwrap(); + rhs_tree.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + + // Both Legacy, rhs not subset + let mut lhs_tree = RoaringTreemap::new(); + for v in 0u64..100 { + lhs_tree.insert(v); + } + let mut rhs_tree = RoaringTreemap::new(); + for v in 50u64..150 { + rhs_tree.insert(v); + } + let mut lhs_buf = Vec::new(); + let mut rhs_buf = Vec::new(); + lhs_tree.serialize_into(&mut lhs_buf).unwrap(); + rhs_tree.serialize_into(&mut rhs_buf).unwrap(); + assert_eq!( + bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), + ground_truth_has_all(&lhs_buf, &rhs_buf) + ); + assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + } } diff --git a/src/common/io/src/bitmap/reader.rs b/src/common/io/src/bitmap/reader.rs index 75d4b8a6680a8..5d7cfc3d41372 100644 --- a/src/common/io/src/bitmap/reader.rs +++ b/src/common/io/src/bitmap/reader.rs @@ -658,6 +658,200 @@ fn array_bitmap_container_has_any( Ok(false) } +pub(crate) fn bitmap_has_all(lhs: &[u8], rhs: &[u8]) -> io::Result { + use std::cmp::Ordering::*; + let lhs_tree = TreemapReader::new(lhs)?; + let rhs_tree = TreemapReader::new(rhs)?; + + let mut lhs_iter = lhs_tree.iter(); + let mut rhs_iter = rhs_tree.iter(); + + let mut lhs_curr = lhs_iter.next().transpose()?; + let mut rhs_curr = rhs_iter.next().transpose()?; + + while let (Some(lhs_bitmap), Some(rhs_bitmap)) = (lhs_curr.as_ref(), rhs_curr.as_ref()) { + match lhs_bitmap.prefix().cmp(&rhs_bitmap.prefix()) { + Less => { + lhs_curr = lhs_iter.next().transpose()?; + } + Greater => { + // rhs has a prefix bucket that lhs doesn't have — lhs cannot contain all rhs values + return Ok(false); + } + Equal => { + if !containers_has_all(lhs_bitmap, rhs_bitmap)? { + return Ok(false); + } + lhs_curr = lhs_iter.next().transpose()?; + rhs_curr = rhs_iter.next().transpose()?; + } + } + } + + // If rhs still has remaining prefix buckets, lhs doesn't cover them + if rhs_curr.is_some() { + return Ok(false); + } + + Ok(true) +} + +fn containers_has_all(lhs: &BitmapReader<'_>, rhs: &BitmapReader<'_>) -> io::Result { + let lhs_count = lhs.containers(); + let rhs_count = rhs.containers(); + let mut i = 0; + let mut j = 0; + + while j < rhs_count { + // If lhs is exhausted but rhs still has containers, lhs can't contain them + if i >= lhs_count { + return Ok(false); + } + let lhs_desc = lhs.description(i)?; + let rhs_desc = rhs.description(j)?; + match lhs_desc.prefix.cmp(&rhs_desc.prefix) { + std::cmp::Ordering::Less => i += 1, + std::cmp::Ordering::Greater => { + // rhs has a container key that lhs doesn't have + return Ok(false); + } + std::cmp::Ordering::Equal => { + let lhs_offset = lhs.container_offset(i)?; + let rhs_offset = rhs.container_offset(j)?; + let lhs_data = &lhs.bitmap_buf()[lhs_offset..]; + let rhs_data = &rhs.bitmap_buf()[rhs_offset..]; + let lhs_cardinality = lhs_desc.cardinality(); + let rhs_cardinality = rhs_desc.cardinality(); + if !container_has_all(lhs_data, lhs_cardinality, rhs_data, rhs_cardinality)? { + return Ok(false); + } + i += 1; + j += 1; + } + } + } + + Ok(true) +} + +fn container_has_all( + lhs_data: &[u8], + lhs_cardinality: usize, + rhs_data: &[u8], + rhs_cardinality: usize, +) -> io::Result { + // lhs cannot contain all of rhs if lhs has fewer values + if lhs_cardinality < rhs_cardinality { + return Ok(false); + } + // If lhs is the full bitmap (65536 values), it trivially contains any rhs subset + if lhs_cardinality == 65536 { + return Ok(true); + } + + let lhs_is_array = lhs_cardinality <= ARRAY_LIMIT; + let rhs_is_array = rhs_cardinality <= ARRAY_LIMIT; + + if lhs_is_array && rhs_is_array { + array_container_has_all_array(lhs_data, lhs_cardinality, rhs_data, rhs_cardinality) + } else if !lhs_is_array && !rhs_is_array { + bitmap_container_has_all_bitmap(lhs_data, rhs_data) + } else if lhs_is_array { + // lhs is array, rhs is bitmap: array container cardinality < 4096, + // bitmap container cardinality >= 4096, so lhs can never be superset of rhs + // Fast path: rhs has more values than lhs can contain + Ok(false) + } else { + // lhs is bitmap, rhs is array: check each rhs value is in lhs bitmap + array_values_all_in_bitmap_container(rhs_data, rhs_cardinality, lhs_data) + } +} + +fn array_container_has_all_array( + lhs_data: &[u8], + lhs_cardinality: usize, + rhs_data: &[u8], + rhs_cardinality: usize, +) -> io::Result { + if lhs_data.len() < lhs_cardinality * 2 { + return Err(Error::other("lhs array container too short")); + } + if rhs_data.len() < rhs_cardinality * 2 { + return Err(Error::other("rhs array container too short")); + } + + // Fast rejection: lhs must have at least as many values as rhs + if lhs_cardinality < rhs_cardinality { + return Ok(false); + } + + let mut i = 0; + let mut j = 0; + while j < rhs_cardinality { + if i >= lhs_cardinality { + return Ok(false); + } + let lv = u16::from_le_bytes(lhs_data[i * 2..i * 2 + 2].try_into().unwrap()); + let rv = u16::from_le_bytes(rhs_data[j * 2..j * 2 + 2].try_into().unwrap()); + if lv < rv { + i += 1; + } else if lv == rv { + i += 1; + j += 1; + } else { + return Ok(false); + } + } + + Ok(true) +} + +fn bitmap_container_has_all_bitmap(lhs_data: &[u8], rhs_data: &[u8]) -> io::Result { + if lhs_data.len() < BITMAP_BYTES { + return Err(Error::other("lhs bitmap container too short")); + } + if rhs_data.len() < BITMAP_BYTES { + return Err(Error::other("rhs bitmap container too short")); + } + + for word_index in 0..BITMAP_WORDS { + let start = word_index * WORD_BYTES; + let lhs_word = u64::from_le_bytes(lhs_data[start..start + WORD_BYTES].try_into().unwrap()); + let rhs_word = u64::from_le_bytes(rhs_data[start..start + WORD_BYTES].try_into().unwrap()); + // Check if all bits set in rhs are also set in lhs + // (rhs_word & !lhs_word) means value exists in rhs but not in lhs + if rhs_word & !lhs_word != 0 { + return Ok(false); + } + } + Ok(true) +} + +fn array_values_all_in_bitmap_container( + array_data: &[u8], + array_cardinality: usize, + bitmap_data: &[u8], +) -> io::Result { + if array_data.len() < array_cardinality * 2 { + return Err(Error::other("array container too short")); + } + if bitmap_data.len() < BITMAP_BYTES { + return Err(Error::other("bitmap container too short")); + } + + for i in 0..array_cardinality { + let value = u16::from_le_bytes(array_data[i * 2..i * 2 + 2].try_into().unwrap()); + let word_index = value as usize / WORD_BITS; + let bit_index = value as usize % WORD_BITS; + let start = word_index * WORD_BYTES; + let word = u64::from_le_bytes(bitmap_data[start..start + WORD_BYTES].try_into().unwrap()); + if word & (1 << bit_index) == 0 { + return Ok(false); + } + } + Ok(true) +} + #[cfg(test)] mod tests { diff --git a/src/common/io/src/lib.rs b/src/common/io/src/lib.rs index 371f248fee9e4..69c9af96b9a09 100644 --- a/src/common/io/src/lib.rs +++ b/src/common/io/src/lib.rs @@ -55,6 +55,7 @@ pub use bitmap::HYBRID_VERSION; pub use bitmap::HybridBitmap; pub use bitmap::LARGE_THRESHOLD; pub use bitmap::bitmap_contains; +pub use bitmap::bitmap_has_all; pub use bitmap::bitmap_has_any; pub use bitmap::bitmap_max; pub use bitmap::bitmap_min; diff --git a/src/query/functions/src/scalars/bitmap.rs b/src/query/functions/src/scalars/bitmap.rs index a9fe63c84c322..d97b671a06a88 100644 --- a/src/query/functions/src/scalars/bitmap.rs +++ b/src/query/functions/src/scalars/bitmap.rs @@ -38,6 +38,7 @@ 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_has_all; use databend_common_io::bitmap::bitmap_has_any; use databend_common_io::bitmap::bitmap_len; use databend_common_io::bitmap::bitmap_max; @@ -385,23 +386,15 @@ pub fn register(registry: &mut FunctionRegistry) { builder.push(false); return; } - let rb = match deserialize_bitmap(b) { - Ok(rb) => rb, - Err(e) => { - ctx.set_error(builder.len(), e.to_string()); - builder.push(false); - return; + match bitmap_has_all(b, items) { + Ok(has_all) => { + builder.push(has_all); } - }; - let rb2 = match deserialize_bitmap(items) { - Ok(rb) => rb, Err(e) => { ctx.set_error(builder.len(), e.to_string()); builder.push(false); - return; } - }; - builder.push(rb.is_superset(&rb2)); + } }, ), ); From 8db20994f69bdb7f1d14c12e9f031693cee5b3ef Mon Sep 17 00:00:00 2001 From: HarryHao Date: Wed, 29 Jul 2026 14:42:41 +0800 Subject: [PATCH 7/9] refactor: bitmap functions to improve maintainability Refactor the 5 bitmap operation functions to improve maintainability: - bitmap_contains - bitmap_min - bitmap_max - bitmap_has_any - bitmap_has_all Key changes: - ContainerReader: read and operate on a roaring container - SmallReader: same as ContainerReader but for HybridBitmap::Small - ContainerVisitor: merge traversal of two serialized treemaps - ContainerHandler: trait abstract container handle logic, used by ContainerVisitor - HasAnyHandler/HasAllHandler: impl ContainerHandler for bitmap_has_any/all - Tests restructured to compare against RoaringTreemap as ground truth, sharing fixtures across all tests, covering combinations of: format, container type, set relationship, and key alignment. --- src/common/io/src/bitmap.rs | 1251 +++++++++++----------- src/common/io/src/bitmap/reader.rs | 1566 ++++++++++++++++++---------- src/common/io/src/lib.rs | 2 - 3 files changed, 1602 insertions(+), 1217 deletions(-) diff --git a/src/common/io/src/bitmap.rs b/src/common/io/src/bitmap.rs index 026ce37537d8f..6d29be51e4932 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -28,8 +28,6 @@ use roaring::RoaringTreemap; use roaring::treemap::Iter; use smallvec::SmallVec; -pub use crate::bitmap::reader::BitmapStats; - mod reader; // https://github.com/ClickHouse/ClickHouse/blob/516a6ed6f8bd8c5f6eed3a10e9037580b2fb6152/src/AggregateFunctions/AggregateFunctionGroupBitmapData.h#L914 @@ -745,6 +743,112 @@ fn is_hybrid_large(buf: &[u8]) -> bool { is_hybrid(buf) && buf[3] == HYBRID_KIND_LARGE } +fn as_roaring(buf: &[u8]) -> Option<&[u8]> { + if is_hybrid_large(buf) { + Some(&buf[HYBRID_HEADER_LEN..]) + } else if !is_hybrid(buf) { + Some(buf) // Legacy + } else { + None // HybridSmall + } +} + +struct SmallReader<'a> { + len: usize, + values: &'a [u8], +} + +impl<'a> SmallReader<'a> { + fn new(buf: &'a [u8]) -> Result { + let payload = &buf[HYBRID_HEADER_LEN..]; + let (len, bytes) = decode_small_payload(payload)?; + Ok(Self { len, values: bytes }) + } + + fn len(&self) -> usize { + self.len + } + + fn is_empty(&self) -> bool { + self.len == 0 + } + + fn get(&self, i: usize) -> u64 { + read_u64_le(&self.values[i * 8..i * 8 + 8]) + } + + fn min(&self) -> Option { + if self.is_empty() { + None + } else { + Some(self.get(0)) + } + } + + fn max(&self) -> Option { + if self.is_empty() { + None + } else { + Some(self.get(self.len - 1)) + } + } + + fn contains(&self, value: u64) -> bool { + let mut lo = 0; + let mut hi = self.len; + while lo < hi { + let mid = lo + (hi - lo) / 2; + match self.get(mid).cmp(&value) { + std::cmp::Ordering::Less => lo = mid + 1, + std::cmp::Ordering::Greater => hi = mid, + std::cmp::Ordering::Equal => return true, + } + } + false + } + + fn has_any_with(&self, other: &SmallReader) -> bool { + let mut i = 0; + let mut j = 0; + while i < self.len && j < other.len { + let lv = self.get(i); + let rv = other.get(j); + if lv < rv { + i += 1; + } else if rv < lv { + j += 1; + } else { + return true; + } + } + false + } + + fn has_all_with(&self, other: &SmallReader) -> bool { + if self.len < other.len { + return false; + } + let mut i = 0; + let mut j = 0; + while j < other.len { + if i >= self.len { + return false; + } + let lv = self.get(i); + let rv = other.get(j); + if lv < rv { + i += 1; + } else if lv == rv { + i += 1; + j += 1; + } else { + return false; + } + } + true + } +} + /// Validate that a serialized bitmap buffer has minimum length. /// /// Different from [`validate_serialized_bitmap`], this one does not check payload to be @@ -781,213 +885,129 @@ pub(crate) fn validate_serialized_bitmap_header_and_length(buf: &[u8]) -> Result } pub fn bitmap_contains(buf: &[u8], value: u64) -> Result { + // Fast path: empty bitmap contains nothing if buf.is_empty() { return Ok(false); } - validate_serialized_bitmap_header_and_length(buf)?; - - if is_hybrid_large(buf) { - Ok(reader::bitmap_contains(&buf[HYBRID_HEADER_LEN..], value)?) + if let Some(roaring_buf) = as_roaring(buf) { + Ok(reader::bitmap_contains(roaring_buf, value)?) } else { - Ok(deserialize_bitmap(buf)?.contains(value)) + Ok(SmallReader::new(buf)?.contains(value)) } } pub fn bitmap_min(buf: &[u8]) -> Result> { + // Fast path: empty bitmap has no minimum if buf.is_empty() { return Ok(None); } validate_serialized_bitmap_header_and_length(buf)?; - if is_hybrid_large(buf) { - Ok(reader::bitmap_min(&buf[HYBRID_HEADER_LEN..])?) + if let Some(roaring_buf) = as_roaring(buf) { + Ok(reader::bitmap_min(roaring_buf)?) } else { - Ok(deserialize_bitmap(buf)?.min()) + Ok(SmallReader::new(buf)?.min()) } } pub fn bitmap_max(buf: &[u8]) -> Result> { + // Fast path: empty bitmap has no maximum if buf.is_empty() { return Ok(None); } validate_serialized_bitmap_header_and_length(buf)?; - if is_hybrid_large(buf) { - Ok(reader::bitmap_max(&buf[HYBRID_HEADER_LEN..])?) + if let Some(roaring_buf) = as_roaring(buf) { + Ok(reader::bitmap_max(roaring_buf)?) } else { - Ok(deserialize_bitmap(buf)?.max()) - } -} - -pub fn bitmap_stats(buf: &[u8]) -> Result { - if buf.is_empty() { - return Ok(BitmapStats { - len: 0, - min: None, - max: None, - }); - } - if is_hybrid_large(buf) { - Ok(reader::bitmap_stats(&buf[HYBRID_HEADER_LEN..])?) - } else { - let bm = deserialize_bitmap(buf)?; - Ok(BitmapStats { - len: bm.len() as u64, - min: bm.min(), - max: bm.max(), - }) + Ok(SmallReader::new(buf)?.max()) } } pub fn bitmap_has_all(lhs: &[u8], rhs: &[u8]) -> Result { + // Fast path: empty rhs is subset of anything if rhs.is_empty() { return Ok(true); } + // Fast path: empty lhs cannot contain non-empty rhs if lhs.is_empty() { - return Ok(false); + return Ok(bitmap_len(rhs)? == 0); } validate_serialized_bitmap_header_and_length(lhs)?; validate_serialized_bitmap_header_and_length(rhs)?; - // Both HybridLarge: use stats for early rejection, then check directly on serialized buffers. - if is_hybrid_large(lhs) && is_hybrid_large(rhs) { - let lhs_stats = reader::bitmap_stats(&lhs[HYBRID_HEADER_LEN..])?; - let rhs_stats = reader::bitmap_stats(&rhs[HYBRID_HEADER_LEN..])?; - if let (Some(lhs_min), Some(rhs_min)) = (lhs_stats.min, rhs_stats.min) - && lhs_min > rhs_min - { - return Ok(false); - } - if let (Some(rhs_max), Some(lhs_max)) = (rhs_stats.max, lhs_stats.max) - && rhs_max > lhs_max - { - return Ok(false); - } - if lhs_stats.len < rhs_stats.len { - return Ok(false); - } - Ok(reader::bitmap_has_all( - &lhs[HYBRID_HEADER_LEN..], - &rhs[HYBRID_HEADER_LEN..], - )?) - } else if is_hybrid_large(lhs) || is_hybrid_large(rhs) { - // Mixed path: one side is HybridLarge, the other is HybridSmall or Legacy. - // Use stats for early rejection; if rhs is HybridSmall, iterate and probe lhs; - // otherwise fall back to full deserialization. - let lhs_stats = bitmap_stats(lhs)?; - let rhs_stats = bitmap_stats(rhs)?; - if let (Some(lhs_min), Some(rhs_min)) = (lhs_stats.min, rhs_stats.min) - && lhs_min > rhs_min - { - return Ok(false); - } - if let (Some(rhs_max), Some(lhs_max)) = (rhs_stats.max, lhs_stats.max) - && rhs_max > lhs_max - { - return Ok(false); - } - // lhs.len < rhs.len -> lhs cannot contain all of rhs - if lhs_stats.len < rhs_stats.len { + + // Both HybridLarge or Legacy: use visitor (zero-copy) + if let (Some(lhs_buf), Some(rhs_buf)) = (as_roaring(lhs), as_roaring(rhs)) { + return Ok(reader::bitmap_has_all(lhs_buf, rhs_buf)?); + } + + // Both HybridSmall: two-pointer subset check + if as_roaring(lhs).is_none() && as_roaring(rhs).is_none() { + let lhs = SmallReader::new(lhs)?; + let rhs = SmallReader::new(rhs)?; + return Ok(lhs.has_all_with(&rhs)); + } + + // Fast path: HybridSmall lhs cannot contain HybridLarge rhs + if as_roaring(lhs).is_none() && is_hybrid_large(rhs) { + return Ok(false); + } + + // Fast path: HybridSmall lhs, Legacy rhs — cardinality check + if as_roaring(lhs).is_none() && as_roaring(rhs).is_some() { + let lhs_small = SmallReader::new(lhs)?; + // Fast path: rhs has more values than lhs, impossible to contain + if reader::bitmap_len_above(rhs, lhs_small.len())? { return Ok(false); } - // rhs is HybridSmall: iterate its (≤31) values and probe lhs via bitmap_contains. - // Each bitmap_contains call scans prefix buckets from the start, causing redundant - // work across iterations. This could be improved with a hint/cursor mechanism that - // tracks the last scanned position, but since HybridSmall holds at most 31 values - // the overhead is bounded (~μs-level). Revisit if this path becomes a bottleneck. - if is_hybrid_large(lhs) && is_hybrid(rhs) && !is_hybrid_large(rhs) { - let rhs_bm = deserialize_bitmap(rhs)?; - for value in rhs_bm.iter() { - if !bitmap_contains(lhs, value)? { - return Ok(false); - } - } - return Ok(true); - } - // Fallback to deserialization and is_superset + // rhs has <= lhs_small.len() values, but still need to check containment let lhs_bm = deserialize_bitmap(lhs)?; let rhs_bm = deserialize_bitmap(rhs)?; - Ok(lhs_bm.is_superset(&rhs_bm)) - } else { - // Both else: deserialize both - let lhs_bm = deserialize_bitmap(lhs)?; - let rhs_bm = deserialize_bitmap(rhs)?; - Ok(lhs_bm.is_superset(&rhs_bm)) + return Ok(lhs_bm.is_superset(&rhs_bm)); } + + // lhs is HybridLarge/Legacy, rhs is HybridSmall: probe + let roaring_buf = as_roaring(lhs).unwrap(); + let rhs_small = SmallReader::new(rhs)?; + for i in 0..rhs_small.len() { + if !reader::bitmap_contains(roaring_buf, rhs_small.get(i))? { + return Ok(false); + } + } + Ok(true) } pub fn bitmap_has_any(lhs: &[u8], rhs: &[u8]) -> Result { + // Fast path: empty bitmap has no intersection if lhs.is_empty() || rhs.is_empty() { return Ok(false); } - validate_serialized_bitmap_header_and_length(lhs)?; validate_serialized_bitmap_header_and_length(rhs)?; - // Stats fast path: range rejection - // If lhs.max < rhs.min OR rhs.max < lhs.min, no overlap possible - if is_hybrid_large(lhs) && is_hybrid_large(rhs) { - let lhs_stats = reader::bitmap_stats(&lhs[HYBRID_HEADER_LEN..])?; - let rhs_stats = reader::bitmap_stats(&rhs[HYBRID_HEADER_LEN..])?; - if let (Some(lhs_min), Some(rhs_max)) = (lhs_stats.min, rhs_stats.max) - && lhs_min > rhs_max - { - return Ok(false); - } - if let (Some(rhs_min), Some(lhs_max)) = (rhs_stats.min, lhs_stats.max) - && rhs_min > lhs_max - { - return Ok(false); - } - Ok(reader::bitmap_has_any( - &lhs[HYBRID_HEADER_LEN..], - &rhs[HYBRID_HEADER_LEN..], - )?) - } else if is_hybrid_large(lhs) || is_hybrid_large(rhs) { - // Normalize: large is the HybridLarge side, other is the non-HybridLarge side - let (large, other) = if is_hybrid_large(lhs) { - (lhs, rhs) - } else { - (rhs, lhs) - }; - - // Stats fast path: range rejection - let large_stats = reader::bitmap_stats(&large[HYBRID_HEADER_LEN..])?; - let other_stats = bitmap_stats(other)?; - if let (Some(large_min), Some(other_max)) = (large_stats.min, other_stats.max) - && large_min > other_max - { - return Ok(false); - } - if let (Some(other_min), Some(large_max)) = (other_stats.min, large_stats.max) - && other_min > large_max - { - return Ok(false); - } + // Both HybridLarge or Legacy: use visitor (zero-copy) + if let (Some(lhs_buf), Some(rhs_buf)) = (as_roaring(lhs), as_roaring(rhs)) { + return Ok(reader::bitmap_has_any(lhs_buf, rhs_buf)?); + } - // other is HybridSmall: iterate its (≤31) values through bitmap_contains on large. - // Each bitmap_contains call scans prefix buckets from the start, causing redundant - // work across iterations. This could be improved with a hint/cursor mechanism that - // tracks the last scanned position, but since HybridSmall holds at most 31 values - // the overhead is bounded (~μs-level). Revisit if this path becomes a bottleneck. - if is_hybrid(other) && !is_hybrid_large(other) { - let other_bm = deserialize_bitmap(other)?; - for value in other_bm.iter() { - if bitmap_contains(large, value)? { - return Ok(true); - } - } - return Ok(false); - } + // Both HybridSmall: two-pointer intersection + if as_roaring(lhs).is_none() && as_roaring(rhs).is_none() { + let lhs = SmallReader::new(lhs)?; + let rhs = SmallReader::new(rhs)?; + return Ok(lhs.has_any_with(&rhs)); + } - // other is Legacy (potentially large): fallback to full deserialization - let lhs_bm = deserialize_bitmap(lhs)?; - let rhs_bm = deserialize_bitmap(rhs)?; - Ok(lhs_bm.intersection_len(&rhs_bm) != 0) + // One side HybridLarge/Legacy, the other HybridSmall: probe + let (roaring_buf, small) = if let Some(rb) = as_roaring(lhs) { + (rb, SmallReader::new(rhs)?) } else { - // Both else: deserialize both - let lhs_bm = deserialize_bitmap(lhs)?; - let rhs_bm = deserialize_bitmap(rhs)?; - Ok(lhs_bm.intersection_len(&rhs_bm) != 0) + (as_roaring(rhs).unwrap(), SmallReader::new(lhs)?) + }; + for i in 0..small.len() { + if reader::bitmap_contains(roaring_buf, small.get(i))? { + return Ok(true); + } } + Ok(false) } fn parse_bitmap_rhs(buf: &[u8]) -> Result> { @@ -1405,6 +1425,7 @@ fn small_intersection_len(lhs: &SmallBitmap, rhs: &SmallBitmap) -> u64 { #[cfg(test)] mod tests { + use smallvec::smallvec; use super::*; @@ -1681,535 +1702,433 @@ mod tests { assert_eq!(decoded.into_iter().collect::>(), 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()); - } - - #[test] - fn test_bitmap_min() { - // HybridLarge: spanning multiple containers - 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_eq!(bitmap_min(&buf).unwrap(), Some(0)); - - // HybridLarge with array container at max capacity (cardinality = 4096) - let boundary = HybridBitmap::from_iter(0u64..4096); - let mut buf = Vec::new(); - boundary.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_min(&buf).unwrap(), Some(0)); - - // HybridLarge with multiple containers: array (4096) + bitmap (>4096) - let mixed = HybridBitmap::from_iter((0u64..4096).chain(65536..106496)); - let mut buf = Vec::new(); - mixed.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_min(&buf).unwrap(), Some(0)); - - // HybridSmall - let small = HybridBitmap::from_iter(0u64..31); - let mut buf = Vec::new(); - small.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_min(&buf).unwrap(), Some(0)); - - // HybridSmall with non-zero start - let small = HybridBitmap::from_iter(100u64..131); - let mut buf = Vec::new(); - small.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_min(&buf).unwrap(), Some(100)); - - // Legacy - let mut tree = RoaringTreemap::new(); - tree.insert(42); - let mut buf = Vec::new(); - tree.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_min(&buf).unwrap(), Some(42)); - - // Empty buffer - assert_eq!(bitmap_min(&[]).unwrap(), None); + // Tests for minimum deserialize bitmap functions. + // + // Cover public functions by comparing against RoaringTreemap: + // - [x] bitmap_contains + // - [x] bitmap_min + // - [x] bitmap_max + // - [x] bitmap_has_any + // - [x] bitmap_has_all + // + // Fixtures: + // - `for_each_fixture` for single-bitmap tests (contains, min, max) + // - `for_each_fixture_pair` for two-bitmap tests (has_any, has_all) + use rand::Rng; + use rand::SeedableRng; + use rand::rngs::SmallRng; + + fn create_bitmap(seed: u64) -> RoaringTreemap { + let mut rng = SmallRng::seed_from_u64(seed); + let mut bitmap = RoaringTreemap::new(); + for _ in 0..50 { + let v = rng.r#gen::(); + bitmap.insert(v); + } + for _ in 0..50 { + let v = rng.r#gen::(); + bitmap.insert(v & u32::MAX as u64); + } + for _ in 0..50 { + let v = rng.r#gen::(); + bitmap.insert(v & u16::MAX as u64); + } + bitmap } - #[test] - fn test_bitmap_max() { - // HybridLarge: spanning multiple containers - 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_eq!(bitmap_max(&buf).unwrap(), Some(239999)); - - // HybridLarge with array container at max capacity (cardinality = 4096) - let boundary = HybridBitmap::from_iter(0u64..4096); - let mut buf = Vec::new(); - boundary.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_max(&buf).unwrap(), Some(4095)); - - // HybridLarge with multiple containers: array (4096) + bitmap (>4096) - let mixed = HybridBitmap::from_iter((0u64..4096).chain(65536..106496)); - let mut buf = Vec::new(); - mixed.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_max(&buf).unwrap(), Some(106495)); - - // HybridSmall - let small = HybridBitmap::from_iter(0u64..31); - let mut buf = Vec::new(); - small.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_max(&buf).unwrap(), Some(30)); - - // HybridSmall with different range - let small = HybridBitmap::from_iter(100u64..131); + fn make_buf(tree: &RoaringTreemap) -> Vec { + let bm = HybridBitmap::from_iter(tree.iter()); let mut buf = Vec::new(); - small.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_max(&buf).unwrap(), Some(130)); - - // Legacy - let mut tree = RoaringTreemap::new(); - tree.insert(42); - tree.insert(100); - let mut buf = Vec::new(); - tree.serialize_into(&mut buf).unwrap(); - assert_eq!(bitmap_max(&buf).unwrap(), Some(100)); - - // Empty buffer - assert_eq!(bitmap_max(&[]).unwrap(), None); - } - - fn ground_truth_has_any(lhs: &[u8], rhs: &[u8]) -> bool { - let lhs_bm = deserialize_bitmap(lhs).unwrap(); - let rhs_bm = deserialize_bitmap(rhs).unwrap(); - lhs_bm.intersection_len(&rhs_bm) != 0 - } - - #[test] - fn test_bitmap_has_any() { - // Both HybridLarge with overlap: 0..50000 and 40000..90000 -> overlap at 40000-49999 - let lhs = HybridBitmap::from_iter(0u64..50000); - let rhs = HybridBitmap::from_iter(40000u64..90000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) + bm.serialize_into(&mut buf).unwrap(); + buf + } + + fn for_each_fixture(mut f: F) + where F: FnMut(&str, &[u8], &RoaringTreemap, u64) { + let fixtures: Vec<(&str, RoaringTreemap, u64)> = vec![ + // HL multi-prefix random + ("random", create_bitmap(123), u64::MAX), + // HS single-prefix + ("format: HS", (0..31u64).collect(), 50), + // HS multi-prefix sparse + ( + "format: HS multi-prefix", + RoaringTreemap::from_iter([0u64, 65535, 65536, (1u64 << 32) + 500]), + 1, + ), + // boundary: 4095 (array threshold-) + ("boundary: 4095", (0..4095u64).collect(), 4095), + // boundary: 4096 (array threshold) + ("boundary: 4096", (0..4096u64).collect(), 4096), + // boundary: 4097 (array threshold+, becomes bitmap) + ("boundary: 4097", (0..4097u64).collect(), 4097), + // boundary: 65535 (near-full container) + ("boundary: 65535", (0..65535u64).collect(), 65535), + // boundary: 65536 (full container) + ("boundary: 65536", (0..65536u64).collect(), 65536), + // boundary: 65537 (full+ container, splits to 2 containers) + ("boundary: 65537", (0..65537u64).collect(), 65537), + ]; + for (name, tree, miss) in fixtures { + let buf = make_buf(&tree); + f(name, &buf, &tree, miss); + } + // empty buffer + f("format: empty", &[], &RoaringTreemap::new(), 0); + // HE serialized empty + let he_tree = RoaringTreemap::new(); + let he_buf = make_buf(&he_tree); + f("format: HE", &he_buf, &he_tree, 0); + } + + fn for_each_fixture_pair(mut f: F) + where F: FnMut(&str, &[u8], &[u8], &RoaringTreemap, &RoaringTreemap) { + let scenarios: Vec<(&str, RoaringTreemap, RoaringTreemap)> = vec![ + // lhs: HL | rhs: HL overlap + ( + "format: HL+HL overlap", + (0..50000u64).collect(), + (30000..80000u64).collect(), + ), + // lhs: HL | rhs: HL disjoint + ( + "format: HL+HL disjoint", + (0..50000u64).collect(), + (100000..150000u64).collect(), + ), + // lhs: HL | rhs: HL superset + ( + "format: HL+HL superset", + (0..50000u64).collect(), + (0..100u64).collect(), + ), + // lhs: HL | rhs: HL not superset + ( + "format: HL+HL not superset", + (0..50000u64).collect(), + (40000..90000u64).collect(), + ), + // lhs: HL | rhs: HL self + ( + "format: HL+HL self", + (0..50000u64).collect(), + (0..50000u64).collect(), + ), + // lhs: HL | rhs: HS subset + ( + "format: HL+HS subset", + (0..50000u64).collect(), + (5..15u64).collect(), + ), + // lhs: HL | rhs: HS not subset + ( + "format: HL+HS not subset", + (0..50000u64).collect(), + (49990..50010u64).collect(), + ), + // lhs: HS | rhs: HL overlap + ( + "format: HS+HL overlap", + (0..31u64).collect(), + (0..50000u64).collect(), + ), + // lhs: HS | rhs: HL disjoint + ( + "format: HS+HL disjoint", + (0..31u64).collect(), + (100000..150000u64).collect(), + ), + // lhs: HS | rhs: HS overlap + ( + "format: HS+HS overlap", + (0..31u64).collect(), + (20..51u64).collect(), + ), + // lhs: HS | rhs: HS disjoint + ( + "format: HS+HS disjoint", + (0..15u64).collect(), + (20..35u64).collect(), + ), + // lhs: HS | rhs: HS subset + ( + "format: HS+HS subset", + (0..31u64).collect(), + (5..15u64).collect(), + ), + // lhs: HS | rhs: HS not subset + ( + "format: HS+HS not subset", + (0..15u64).collect(), + (10..25u64).collect(), + ), + // lhs: p0 | rhs: p0, p1 rhs-only prefix + ( + "prefix: rhs-only", + (0u64..0x10000).step_by(7).collect(), + (0u64..0x10000) + .step_by(5) + .chain(((1u64 << 32)..(1u64 << 32) + 0x10000).step_by(5)) + .collect(), + ), + // lhs: p0, p1 | rhs: p0 lhs-only prefix + ( + "prefix: lhs-only", + (0u64..0x10000) + .step_by(7) + .chain(((1u64 << 32)..(1u64 << 32) + 0x10000).step_by(7)) + .collect(), + (0u64..0x10000).step_by(5).collect(), + ), + // lhs: p0 | rhs: p1 disjoint prefixes + ( + "prefix: disjoint", + (0u64..0x10000).step_by(7).collect(), + ((1u64 << 32)..(1u64 << 32) + 0x10000).step_by(5).collect(), + ), + // lhs: p0, p1 | rhs: p1, p2 overlap prefixes + ( + "prefix: overlap", + (0u64..0x10000) + .step_by(7) + .chain(((1u64 << 32)..(1u64 << 32) + 0x10000).step_by(7)) + .collect(), + ((1u64 << 32)..(1u64 << 32) + 0x10000) + .step_by(5) + .chain(((2u64 << 32)..(2u64 << 32) + 0x10000).step_by(5)) + .collect(), + ), + // boundary: 4095 + 4096 + ( + "boundary: 4095+4096", + (0..4095u64).collect(), + (0..4096u64).collect(), + ), + // boundary: 4096 + 4097 + ( + "boundary: 4096+4097", + (0..4096u64).collect(), + (0..4097u64).collect(), + ), + // boundary: 4097 + 4097 + ( + "boundary: 4097+4097", + (0..4097u64).collect(), + (0..4097u64).collect(), + ), + // boundary: 65535 + 65536 + ( + "boundary: 65535+65536", + (0..65535u64).collect(), + (0..65536u64).collect(), + ), + // boundary: 65536 + 65536 + ( + "boundary: 65536+65536", + (0..65536u64).collect(), + (0..65536u64).collect(), + ), + // boundary: 65537 + 65537 + ( + "boundary: 65537+65537", + (0..65537u64).collect(), + (0..65537u64).collect(), + ), + // boundary: 65535 + 65537 + ( + "boundary: 65535+65537", + (0..65535u64).collect(), + (0..65537u64).collect(), + ), + // boundary: 4095 + 65535 + ( + "boundary: 4095+65535", + (0..4095u64).collect(), + (0..65535u64).collect(), + ), + // boundary: 4095 + 65536 + ( + "boundary: 4095+65536", + (0..4095u64).collect(), + (0..65536u64).collect(), + ), + // boundary: 4095 + 65537 + ( + "boundary: 4095+65537", + (0..4095u64).collect(), + (0..65537u64).collect(), + ), + // boundary: 4096 + 65535 + ( + "boundary: 4096+65535", + (0..4096u64).collect(), + (0..65535u64).collect(), + ), + // boundary: 4096 + 65536 + ( + "boundary: 4096+65536", + (0..4096u64).collect(), + (0..65536u64).collect(), + ), + // boundary: 4096 + 65537 + ( + "boundary: 4096+65537", + (0..4096u64).collect(), + (0..65537u64).collect(), + ), + // boundary: 4097 + 65535 + ( + "boundary: 4097+65535", + (0..4097u64).collect(), + (0..65535u64).collect(), + ), + // boundary: 4097 + 65536 + ( + "boundary: 4097+65536", + (0..4097u64).collect(), + (0..65536u64).collect(), + ), + // boundary: 4097 + 65537 + ( + "boundary: 4097+65537", + (0..4097u64).collect(), + (0..65537u64).collect(), + ), + ]; + for (name, lhs_tree, rhs_tree) in scenarios { + let lhs_buf = make_buf(&lhs_tree); + let rhs_buf = make_buf(&rhs_tree); + f(name, &lhs_buf, &rhs_buf, &lhs_tree, &rhs_tree); + } + // Empty scenarios + let hl_tree: RoaringTreemap = (0..50000u64).collect(); + let hl_buf = make_buf(&hl_tree); + let hs_tree: RoaringTreemap = (0..31u64).collect(); + let hs_buf = make_buf(&hs_tree); + let he_tree: RoaringTreemap = RoaringTreemap::new(); + let he_buf = make_buf(&he_tree); + // lhs: empty | rhs: HL + f( + "format: Empty+HL", + &[], + &hl_buf, + &RoaringTreemap::new(), + &hl_tree, ); - assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); - - // Both HybridLarge disjoint: 0..50000 and 100000..150000 - let lhs = HybridBitmap::from_iter(0u64..50000); - let rhs = HybridBitmap::from_iter(100000u64..150000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) + // lhs: HE | rhs: HL + f("format: HE+HL", &he_buf, &hl_buf, &he_tree, &hl_tree); + // lhs: empty | rhs: HS + f( + "format: Empty+HS", + &[], + &hs_buf, + &RoaringTreemap::new(), + &hs_tree, ); - assert!(!bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); - - // Both HybridSmall with overlap: 0..31 and 20..51 -> overlap at 20-30 - let lhs = HybridBitmap::from_iter(0u64..31); - let rhs = HybridBitmap::from_iter(20u64..51); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) + // lhs: HE | rhs: HS + f("format: HE+HS", &he_buf, &hs_buf, &he_tree, &hs_tree); + // lhs: HL | rhs: empty + f( + "format: HL+Empty", + &hl_buf, + &[], + &hl_tree, + &RoaringTreemap::new(), ); - assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); - - // Both HybridSmall disjoint: 0..15 and 20..35 - let lhs = HybridBitmap::from_iter(0u64..15); - let rhs = HybridBitmap::from_iter(20u64..35); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) + // lhs: HL | rhs: HE + f("format: HL+HE", &hl_buf, &he_buf, &hl_tree, &he_tree); + // lhs: HS | rhs: empty + f( + "format: HS+Empty", + &hs_buf, + &[], + &hs_tree, + &RoaringTreemap::new(), ); - assert!(!bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); - - // Mixed (HybridLarge lhs, HybridSmall rhs) with overlap: 0..50000 and 10..41 - let lhs = HybridBitmap::from_iter(0u64..50000); - let rhs = HybridBitmap::from_iter(10u64..41); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) + // lhs: HS | rhs: HE + f("format: HS+HE", &hs_buf, &he_buf, &hs_tree, &he_tree); + // lhs: empty | rhs: empty + f( + "format: Empty+Empty", + &[], + &[], + &RoaringTreemap::new(), + &RoaringTreemap::new(), ); - assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); - - // Mixed (HybridSmall lhs, HybridLarge rhs) with overlap: 10..41 and 0..50000 - let lhs = HybridBitmap::from_iter(10u64..41); - let rhs = HybridBitmap::from_iter(0u64..50000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) + // lhs: empty | rhs: HE + f( + "format: Empty+HE", + &[], + &he_buf, + &RoaringTreemap::new(), + &he_tree, ); - assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); - - // Empty lhs - let rhs = HybridBitmap::from_iter(0u64..50000); - let mut rhs_buf = Vec::new(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert!(!bitmap_has_any(&[], &rhs_buf).unwrap()); - - // Empty rhs - let lhs = HybridBitmap::from_iter(0u64..50000); - let mut lhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - assert!(!bitmap_has_any(&lhs_buf, &[]).unwrap()); - - // Both empty - assert!(!bitmap_has_any(&[], &[]).unwrap()); - - // Mixed (Legacy lhs, HybridLarge rhs) with overlap - let mut lhs_tree = RoaringTreemap::new(); - for v in [0u64, 2500, 40000] { - lhs_tree.insert(v); - } - let rhs = HybridBitmap::from_iter(0u64..50000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs_tree.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) + // lhs: HE | rhs: empty + f( + "format: HE+Empty", + &he_buf, + &[], + &he_tree, + &RoaringTreemap::new(), ); - assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + // lhs: HE | rhs: HE + f("format: HE+HE", &he_buf, &he_buf, &he_tree, &he_tree); + } - // Mixed (HybridLarge lhs, Legacy rhs) with overlap - let lhs = HybridBitmap::from_iter(0u64..50000); - let mut rhs_tree = RoaringTreemap::new(); - for v in [0u64, 2500, 40000] { - rhs_tree.insert(v); - } - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs_tree.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) - ); - assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + #[test] + fn test_bitmap_contains() { + for_each_fixture(|name, buf, tree, miss_value| { + if let Some(hit) = tree.min() { + let expected = tree.contains(hit); + let actual = bitmap_contains(buf, hit).unwrap(); + assert_eq!( + actual, expected, + "bitmap_contains hit: fixture={name}, val={hit}" + ); + } + let miss = miss_value; + let expected = tree.contains(miss); + let actual = bitmap_contains(buf, miss).unwrap(); + assert_eq!( + actual, expected, + "bitmap_contains miss: fixture={name}, val={miss}" + ); + }); + } - // Both Legacy with overlap - let mut lhs_tree = RoaringTreemap::new(); - for v in 0u64..100 { - lhs_tree.insert(v); - } - let mut rhs_tree = RoaringTreemap::new(); - for v in 50u64..150 { - rhs_tree.insert(v); - } - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs_tree.serialize_into(&mut lhs_buf).unwrap(); - rhs_tree.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_any(&lhs_buf, &rhs_buf) - ); - assert!(bitmap_has_any(&lhs_buf, &rhs_buf).unwrap()); + #[test] + fn test_bitmap_min() { + for_each_fixture(|name, buf, tree, _miss_value| { + let expected = tree.min(); + let actual = bitmap_min(buf).unwrap(); + assert_eq!(actual, expected, "bitmap_min: fixture={name}"); + }); } #[test] - fn test_bitmap_stats() { - // HybridLarge stats - let large = HybridBitmap::from_iter(0u64..50000); - let mut buf = Vec::new(); - large.serialize_into(&mut buf).unwrap(); - let stats = bitmap_stats(&buf).unwrap(); - assert_eq!(stats.len, 50000); - assert_eq!(stats.min, Some(0)); - assert_eq!(stats.max, Some(49999)); - - // HybridSmall stats - let small = HybridBitmap::from_iter(100u64..131); - let mut buf = Vec::new(); - small.serialize_into(&mut buf).unwrap(); - let stats = bitmap_stats(&buf).unwrap(); - assert_eq!(stats.len, 31); - assert_eq!(stats.min, Some(100)); - assert_eq!(stats.max, Some(130)); - - // Empty stats - let stats = bitmap_stats(&[]).unwrap(); - assert_eq!(stats.len, 0); - assert_eq!(stats.min, None); - assert_eq!(stats.max, None); - - // Legacy stats - let mut tree = RoaringTreemap::new(); - tree.insert(42); - tree.insert(100); - let mut buf = Vec::new(); - tree.serialize_into(&mut buf).unwrap(); - let stats = bitmap_stats(&buf).unwrap(); - assert_eq!(stats.len, 2); - assert_eq!(stats.min, Some(42)); - assert_eq!(stats.max, Some(100)); + fn test_bitmap_max() { + for_each_fixture(|name, buf, tree, _miss_value| { + let expected = tree.max(); + let actual = bitmap_max(buf).unwrap(); + assert_eq!(actual, expected, "bitmap_max: fixture={name}"); + }); } - fn ground_truth_has_all(lhs: &[u8], rhs: &[u8]) -> bool { - let lhs_bm = deserialize_bitmap(lhs).unwrap(); - let rhs_bm = deserialize_bitmap(rhs).unwrap(); - lhs_bm.is_superset(&rhs_bm) + #[test] + fn test_bitmap_has_any() { + for_each_fixture_pair(|name, lhs_buf, rhs_buf, lhs_tree, rhs_tree| { + let expected = !(lhs_tree & rhs_tree).is_empty(); + let actual = bitmap_has_any(lhs_buf, rhs_buf).unwrap(); + assert_eq!(actual, expected, "bitmap_has_any: fixture={name}"); + }); } #[test] fn test_bitmap_has_all() { - // Both HybridLarge, rhs is subset: 0..50000 contains 0..10000 - let lhs = HybridBitmap::from_iter(0u64..50000); - let rhs = HybridBitmap::from_iter(0u64..10000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Both HybridLarge, rhs is not subset: 0..50000 vs 40000..90000 - // rhs has values 50000-89999 that are not in lhs - let lhs = HybridBitmap::from_iter(0u64..50000); - let rhs = HybridBitmap::from_iter(40000u64..90000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Both HybridLarge, disjoint: 0..50000 vs 100000..150000 - let lhs = HybridBitmap::from_iter(0u64..50000); - let rhs = HybridBitmap::from_iter(100000u64..150000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Both HybridSmall, rhs is subset: 0..31 contains 5..15 - let lhs = HybridBitmap::from_iter(0u64..31); - let rhs = HybridBitmap::from_iter(5u64..15); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Both HybridSmall, rhs is not subset: 0..15 vs 10..25 - // rhs has values 15-24 that are not in lhs - let lhs = HybridBitmap::from_iter(0u64..15); - let rhs = HybridBitmap::from_iter(10u64..25); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Mixed (HybridLarge lhs, HybridSmall rhs) with superset: 0..50000 contains 5..15 - let lhs = HybridBitmap::from_iter(0u64..50000); - let rhs = HybridBitmap::from_iter(5u64..15); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Mixed (HybridLarge lhs, HybridSmall rhs not subset): 0..50000 vs 49990..50010 - // rhs has values 50000-50009 that are not in lhs - let lhs = HybridBitmap::from_iter(0u64..50000); - let rhs = HybridBitmap::from_iter(49990u64..50010); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Mixed (HybridSmall lhs, HybridLarge rhs): already rejected by len fast path - let lhs = HybridBitmap::from_iter(0u64..31); - let rhs = HybridBitmap::from_iter(0u64..50000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Empty rhs: always true - let lhs = HybridBitmap::from_iter(0u64..50000); - let mut lhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - assert!(bitmap_has_all(&lhs_buf, &[]).unwrap()); - - // Empty lhs with non-empty rhs: always false - let rhs = HybridBitmap::from_iter(0u64..50000); - let mut rhs_buf = Vec::new(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert!(!bitmap_has_all(&[], &rhs_buf).unwrap()); - - // Both empty: true - assert!(bitmap_has_all(&[], &[]).unwrap()); - - // Mixed (Legacy lhs, HybridLarge rhs) - let mut lhs_tree = RoaringTreemap::new(); - for v in 0u64..50000 { - lhs_tree.insert(v); - } - let rhs = HybridBitmap::from_iter(0u64..10000); - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs_tree.serialize_into(&mut lhs_buf).unwrap(); - rhs.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - - // Mixed (HybridLarge lhs, Legacy rhs) - let lhs = HybridBitmap::from_iter(0u64..10000); - let mut rhs_tree = RoaringTreemap::new(); - for v in 0u64..50000 { - rhs_tree.insert(v); - } - - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs.serialize_into(&mut lhs_buf).unwrap(); - rhs_tree.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - - // Both Legacy - let mut lhs_tree = RoaringTreemap::new(); - for v in 0u64..100 { - lhs_tree.insert(v); - } - let mut rhs_tree = RoaringTreemap::new(); - for v in 50u64..80 { - rhs_tree.insert(v); - } - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs_tree.serialize_into(&mut lhs_buf).unwrap(); - rhs_tree.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); - - // Both Legacy, rhs not subset - let mut lhs_tree = RoaringTreemap::new(); - for v in 0u64..100 { - lhs_tree.insert(v); - } - let mut rhs_tree = RoaringTreemap::new(); - for v in 50u64..150 { - rhs_tree.insert(v); - } - let mut lhs_buf = Vec::new(); - let mut rhs_buf = Vec::new(); - lhs_tree.serialize_into(&mut lhs_buf).unwrap(); - rhs_tree.serialize_into(&mut rhs_buf).unwrap(); - assert_eq!( - bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(), - ground_truth_has_all(&lhs_buf, &rhs_buf) - ); - assert!(!bitmap_has_all(&lhs_buf, &rhs_buf).unwrap()); + for_each_fixture_pair(|name, lhs_buf, rhs_buf, lhs_tree, rhs_tree| { + let expected = lhs_tree.is_superset(rhs_tree); + let actual = bitmap_has_all(lhs_buf, rhs_buf).unwrap(); + assert_eq!(actual, expected, "bitmap_has_all: fixture={name}"); + }); } } diff --git a/src/common/io/src/bitmap/reader.rs b/src/common/io/src/bitmap/reader.rs index 5d7cfc3d41372..0b2eddafd676f 100644 --- a/src/common/io/src/bitmap/reader.rs +++ b/src/common/io/src/bitmap/reader.rs @@ -17,6 +17,7 @@ use std::io::Cursor; use std::io::Error; use std::io::ErrorKind; use std::io::Seek; +use std::ops::ControlFlow; use byteorder::LittleEndian; use byteorder::ReadBytesExt; @@ -32,6 +33,7 @@ const WORD_BITS: usize = 64; const WORD_BYTES: usize = 8; const BITMAP_WORDS: usize = 1024; const BITMAP_BYTES: usize = WORD_BYTES * BITMAP_WORDS; +const CONTAINER_MAX: usize = WORD_BITS * BITMAP_WORDS; #[derive(Clone)] pub struct TreemapReader<'a> { @@ -54,6 +56,83 @@ impl<'a> TreemapReader<'a> { offset: 0, } } + + pub fn contains(&self, value: u64) -> io::Result { + let prefix = (value >> 32) as u32; + let container_key = ((value >> 16) & 0xFFFF) as u16; + let low16 = (value & 0xFFFF) as u16; + + for bitmap_result in self.iter() { + let bitmap = bitmap_result?; + if bitmap.prefix() < prefix { + continue; + } + if bitmap.prefix() > prefix { + return Ok(false); + } + return Ok(bitmap + .find_container(container_key)? + .is_some_and(|c| c.contains(low16))); + } + Ok(false) + } + + pub fn min(&self) -> io::Result> { + for bitmap_result in self.iter() { + let bitmap = bitmap_result?; + if bitmap.containers() == 0 { + continue; + } + let container = bitmap.container(0)?; + if let Some(low16) = container.min() { + let prefix = bitmap.prefix() as u64; + return Ok(Some( + (prefix << 32) | ((container.key() as u64) << 16) | low16 as u64, + )); + } + } + Ok(None) + } + + pub fn max(&self) -> io::Result> { + let mut last_result = None; + for bitmap_result in self.iter() { + let bitmap = bitmap_result?; + let last_idx = bitmap.containers() - 1; + let container = bitmap.container(last_idx)?; + if let Some(low16) = container.max() { + let prefix = bitmap.prefix() as u64; + last_result = + Some((prefix << 32) | ((container.key() as u64) << 16) | low16 as u64); + } + } + Ok(last_result) + } + + pub fn len(&self) -> io::Result { + let mut sum = 0; + for bitmap in self.iter() { + let bitmap = bitmap?; + for i in 0..bitmap.containers() { + sum += bitmap.description(i)?.cardinality(); + } + } + Ok(sum) + } + + pub fn len_above(&self, threshold: usize) -> io::Result { + let mut sum = 0; + for bitmap in self.iter() { + let bitmap = bitmap?; + for i in 0..bitmap.containers() { + sum += bitmap.description(i)?.cardinality(); + if sum > threshold { + return Ok(true); + } + } + } + Ok(false) + } } pub struct TreeMapIter<'a> { @@ -187,6 +266,45 @@ impl BitmapReader<'_> { } Ok(offset) } + + pub fn container(&self, index: usize) -> io::Result> { + let desc = self.description(index)?; + let offset = self.container_offset(index)?; + let cardinality = desc.cardinality(); + let data = &self.bitmap_buf()[offset..]; + + // Validate container data length + let required_len = if cardinality <= ARRAY_LIMIT { + cardinality * 2 + } else { + BITMAP_BYTES + }; + if data.len() < required_len { + return Err(Error::other("container data too short")); + } + + Ok(ContainerReader { + key: desc.prefix, + cardinality, + is_array: cardinality <= ARRAY_LIMIT, + data, + }) + } + + pub fn find_container(&self, key: u16) -> io::Result>> { + let mut lo = 0; + let mut hi = self.containers(); + while lo < hi { + let mid = lo + (hi - lo) / 2; + let desc = self.description(mid)?; + match desc.prefix.cmp(&key) { + std::cmp::Ordering::Less => lo = mid + 1, + std::cmp::Ordering::Greater => hi = mid, + std::cmp::Ordering::Equal => return Ok(Some(self.container(mid)?)), + } + } + Ok(None) + } } pub struct Description { @@ -201,655 +319,651 @@ impl Description { } } -pub fn bitmap_len(buf: &[u8]) -> io::Result { - let tree = TreemapReader::new(buf)?; - let mut sum = 0; - for bitmap in tree.iter() { - let bitmap = bitmap?; - for i in 0..bitmap.containers() { - sum += bitmap.description(i)?.cardinality(); - } - } - Ok(sum) +pub struct ContainerReader<'a> { + key: u16, + cardinality: usize, + is_array: bool, + data: &'a [u8], } -pub(crate) fn bitmap_contains(buf: &[u8], value: u64) -> io::Result { - 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; +impl<'a> ContainerReader<'a> { + pub fn key(&self) -> u16 { + self.key + } - for bitmap_result in tree.iter() { - let bitmap = bitmap_result?; - if bitmap.prefix() < prefix { - continue; - } - if bitmap.prefix() > prefix { - return Ok(false); + pub fn contains(&self, low16: u16) -> bool { + if self.is_array { + self.contains_array(low16) + } else { + self.contains_bitmap(low16) } - 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 { - 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) - }; - } + pub fn min(&self) -> Option { + if self.is_array { + self.array_first() + } else { + self.bitmap_first() } } - Ok(false) -} -/// Binary search in an array container (cardinality < 4096). -fn array_container_contains(data: &[u8], cardinality: usize, low16: u16) -> io::Result { - 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), + pub fn max(&self) -> Option { + if self.is_array { + self.array_last() + } else { + self.bitmap_last() } } - Ok(false) -} -/// Bit lookup in a bitmap container (cardinality >= 4096). -fn bitmap_container_contains(data: &[u8], low16: u16) -> io::Result { - if data.len() < BITMAP_BYTES { - return Err(Error::other("bitmap container too short")); + pub(crate) fn has_any_with(&self, other: &ContainerReader) -> bool { + // Fast path: pigeonhole principle + if self.cardinality + other.cardinality > CONTAINER_MAX { + return true; + } + // Fast path: range non-overlapping + if let (Some(self_max), Some(other_min)) = (self.max(), other.min()) + && self_max < other_min + { + return false; + } + if let (Some(self_min), Some(other_max)) = (self.min(), other.max()) + && self_min > other_max + { + return false; + } + // Type dispatch + if self.is_array && other.is_array { + self.array_has_any_array(other) + } else if !self.is_array && !other.is_array { + self.bitmap_has_any_bitmap(other) + } else if self.is_array { + self.array_has_any_bitmap(other) + } else { + other.array_has_any_bitmap(self) + } } - 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(crate) fn bitmap_min(buf: &[u8]) -> io::Result> { - let tree = TreemapReader::new(buf)?; - let bitmap = match tree.iter().next() { - None => return Ok(None), - Some(b) => b?, - }; - if bitmap.containers() == 0 { - return Ok(None); - } - let desc = bitmap.description(0)?; - let offset = bitmap.container_offset(0)?; - let container_data = &bitmap.bitmap_buf()[offset..]; - let prefix = bitmap.prefix() as u64; - let container_key = desc.prefix as u64; - let cardinality = desc.cardinality(); - let low16 = if cardinality <= ARRAY_LIMIT { - array_container_first(container_data, cardinality)? - } else { - bitmap_container_first(container_data)? - }; - Ok(Some(prefix << 32 | container_key << 16 | low16 as u64)) -} - -fn array_container_first(data: &[u8], cardinality: usize) -> io::Result { - if data.len() < cardinality * 2 { - return Err(Error::other("array container too short")); + pub(crate) fn has_all_with(&self, other: &ContainerReader) -> bool { + // Fast path: cardinality check + if self.cardinality < other.cardinality { + return false; + } + // Fast path: full bitmap + if self.cardinality == CONTAINER_MAX { + return true; + } + // Fast path: type check (array capacity < bitmap) + if self.is_array && !other.is_array { + return false; + } + // Fast path: range non-covering + if let (Some(self_min), Some(other_min)) = (self.min(), other.min()) + && self_min > other_min + { + return false; + } + if let (Some(self_max), Some(other_max)) = (self.max(), other.max()) + && self_max < other_max + { + return false; + } + // Type dispatch + if self.is_array && other.is_array { + self.array_has_all_array(other) + } else if !self.is_array && !other.is_array { + self.bitmap_has_all_bitmap(other) + } else { + // self is bitmap, other is array + self.bitmap_has_all_array(other) + } } - Ok(u16::from_le_bytes(data[0..2].try_into().unwrap())) -} -fn bitmap_container_first(data: &[u8]) -> io::Result { - if data.len() < BITMAP_BYTES { - return Err(Error::other("bitmap container too short")); + fn contains_array(&self, low16: u16) -> bool { + let mut lo = 0; + let mut hi = self.cardinality; + while lo < hi { + let mid = lo + (hi - lo) / 2; + let v = u16::from_le_bytes(self.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 true, + } + } + false } - // Find the lowest set bit in the 1024-word bitmap - for word_index in 0..BITMAP_WORDS { + + fn contains_bitmap(&self, low16: u16) -> bool { + 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()); - if word != 0 { - return Ok((word_index * WORD_BITS + word.trailing_zeros() as usize) as u16); - } + let word = u64::from_le_bytes(self.data[start..start + WORD_BYTES].try_into().unwrap()); + word & (1 << bit_index) != 0 } - // All zeros — shouldn't happen for a valid bitmap container - Err(Error::other("bitmap container has no set bits")) -} -pub(crate) fn bitmap_max(buf: &[u8]) -> io::Result> { - let tree = TreemapReader::new(buf)?; - // Find the last prefix bucket - let mut last_bitmap = None; - for bitmap_result in tree.iter() { - last_bitmap = Some(bitmap_result?); - } - let bitmap = match last_bitmap { - None => return Ok(None), - Some(b) => b, - }; - if bitmap.containers() == 0 { - return Ok(None); - } - let last_idx = bitmap.containers() - 1; - let desc = bitmap.description(last_idx)?; - let offset = bitmap.container_offset(last_idx)?; - let container_data = &bitmap.bitmap_buf()[offset..]; - let prefix = bitmap.prefix() as u64; - let container_key = desc.prefix as u64; - let cardinality = desc.cardinality(); - let low16 = if cardinality <= ARRAY_LIMIT { - array_container_last(container_data, cardinality)? - } else { - bitmap_container_last(container_data)? - }; - Ok(Some(prefix << 32 | container_key << 16 | low16 as u64)) -} + fn array_first(&self) -> Option { + if self.cardinality == 0 { + return None; + } + Some(u16::from_le_bytes(self.data[0..2].try_into().unwrap())) + } -fn array_container_last(data: &[u8], cardinality: usize) -> io::Result { - if data.len() < cardinality * 2 { - return Err(Error::other("array container too short")); + fn bitmap_first(&self) -> Option { + for word_index in 0..BITMAP_WORDS { + let start = word_index * WORD_BYTES; + let word = u64::from_le_bytes(self.data[start..start + WORD_BYTES].try_into().unwrap()); + if word != 0 { + return Some((word_index * WORD_BITS + word.trailing_zeros() as usize) as u16); + } + } + None } - let offset = (cardinality - 1) * 2; - Ok(u16::from_le_bytes( - data[offset..offset + 2].try_into().unwrap(), - )) -} -fn bitmap_container_last(data: &[u8]) -> io::Result { - if data.len() < BITMAP_BYTES { - return Err(Error::other("bitmap container too short")); + fn array_last(&self) -> Option { + if self.cardinality == 0 { + return None; + } + let offset = (self.cardinality - 1) * 2; + Some(u16::from_le_bytes( + self.data[offset..offset + 2].try_into().unwrap(), + )) } - // Find the highest set bit by scanning from the last word backwards - for word_index in (0..BITMAP_WORDS).rev() { - let start = word_index * WORD_BYTES; - let word = u64::from_le_bytes(data[start..start + WORD_BYTES].try_into().unwrap()); - if word != 0 { - let bit_index = WORD_BITS - 1 - word.leading_zeros() as usize; - return Ok((word_index * WORD_BITS + bit_index) as u16); + + fn bitmap_last(&self) -> Option { + for word_index in (0..BITMAP_WORDS).rev() { + let start = word_index * WORD_BYTES; + let word = u64::from_le_bytes(self.data[start..start + WORD_BYTES].try_into().unwrap()); + if word != 0 { + let bit_index = WORD_BITS - 1 - word.leading_zeros() as usize; + return Some((word_index * WORD_BITS + bit_index) as u16); + } } + None } - Err(Error::other("bitmap container has no set bits")) -} -pub fn intersection_with_serialized(tree: &mut RoaringTreemap, buf: &[u8]) -> io::Result<()> { - use std::cmp::Ordering::*; - let rhs = TreemapReader::new(buf)?; - let mut bitmaps = Vec::new(); - let mut lhs_iter = tree.bitmaps(); - let mut rhs_iter = rhs.iter(); + fn array_has_any_array(&self, other: &ContainerReader) -> bool { + let mut i = 0; + let mut j = 0; + while i < self.cardinality && j < other.cardinality { + let lv = u16::from_le_bytes(self.data[i * 2..i * 2 + 2].try_into().unwrap()); + let rv = u16::from_le_bytes(other.data[j * 2..j * 2 + 2].try_into().unwrap()); + if lv < rv { + i += 1; + } else if rv < lv { + j += 1; + } else { + return true; + } + } + false + } - let mut lhs_curr = lhs_iter.next(); - let mut rhs_curr = rhs_iter.next().transpose()?; + fn bitmap_has_any_bitmap(&self, other: &ContainerReader) -> bool { + for word_index in 0..BITMAP_WORDS { + let start = word_index * WORD_BYTES; + let lhs_word = + u64::from_le_bytes(self.data[start..start + WORD_BYTES].try_into().unwrap()); + let rhs_word = + u64::from_le_bytes(other.data[start..start + WORD_BYTES].try_into().unwrap()); + if lhs_word & rhs_word != 0 { + return true; + } + } + false + } - while let (Some((lhs_prefix, lhs_bitmap)), Some(rhs_bitmap)) = (lhs_curr, rhs_curr.as_ref()) { - match lhs_prefix.cmp(&rhs_bitmap.prefix()) { - Less => { - lhs_curr = lhs_iter.next(); + fn array_has_any_bitmap(&self, other: &ContainerReader) -> bool { + for i in 0..self.cardinality { + let value = u16::from_le_bytes(self.data[i * 2..i * 2 + 2].try_into().unwrap()); + let word_index = value as usize / WORD_BITS; + let bit_index = value as usize % WORD_BITS; + let start = word_index * WORD_BYTES; + let word = + u64::from_le_bytes(other.data[start..start + WORD_BYTES].try_into().unwrap()); + if word & (1 << bit_index) != 0 { + return true; } - Greater => { - rhs_curr = rhs_iter.next().transpose()?; + } + false + } + + fn array_has_all_array(&self, other: &ContainerReader) -> bool { + let mut i = 0; + let mut j = 0; + while j < other.cardinality { + if i >= self.cardinality { + return false; } - Equal => { - let intersection = lhs_bitmap - .intersection_with_serialized_unchecked(Cursor::new(rhs_bitmap.bitmap_buf()))?; - if !intersection.is_empty() { - bitmaps.push((lhs_prefix, intersection)); - } - lhs_curr = lhs_iter.next(); - rhs_curr = rhs_iter.next().transpose()?; + let lv = u16::from_le_bytes(self.data[i * 2..i * 2 + 2].try_into().unwrap()); + let rv = u16::from_le_bytes(other.data[j * 2..j * 2 + 2].try_into().unwrap()); + if lv < rv { + i += 1; + } else if lv == rv { + i += 1; + j += 1; + } else { + return false; } } + true } - *tree = RoaringTreemap::from_bitmaps(bitmaps); + fn bitmap_has_all_bitmap(&self, other: &ContainerReader) -> bool { + for word_index in 0..BITMAP_WORDS { + let start = word_index * WORD_BYTES; + let lhs_word = + u64::from_le_bytes(self.data[start..start + WORD_BYTES].try_into().unwrap()); + let rhs_word = + u64::from_le_bytes(other.data[start..start + WORD_BYTES].try_into().unwrap()); + // Check if all bits set in rhs are also set in lhs + if rhs_word & !lhs_word != 0 { + return false; + } + } + true + } - Ok(()) + fn bitmap_has_all_array(&self, other: &ContainerReader) -> bool { + for i in 0..other.cardinality { + let value = u16::from_le_bytes(other.data[i * 2..i * 2 + 2].try_into().unwrap()); + let word_index = value as usize / WORD_BITS; + let bit_index = value as usize % WORD_BITS; + let start = word_index * WORD_BYTES; + let word = u64::from_le_bytes(self.data[start..start + WORD_BYTES].try_into().unwrap()); + if word & (1 << bit_index) == 0 { + return false; + } + } + true + } } -pub struct BitmapStats { - pub len: u64, - pub min: Option, - pub max: Option, +/// Handler trait for zero-copy container-pair traversal. +/// +/// RoaringTreemap has two levels: prefix buckets (u32) and containers (u16 key). +/// `ContainerVisitor` walks both levels in sorted order and dispatches to the +/// handler at four kinds of events: +/// +/// - `handle_lhs_only`: a container exists in lhs but not in rhs at the same key. +/// - `handle_rhs_only`: a container exists in rhs but not in lhs at the same key. +/// - `handle_matched`: both sides have a container at the same key. +/// - `handle_none`: both bitmaps are empty. +/// +/// The `remaining` parameters signal whether the other side has more data +/// beyond the current prefix bucket or container key. This lets handlers +/// decide early termination (e.g., `has_any` returns false when both sides +/// are exhausted with no match). +/// +/// Return `ControlFlow::Break(output)` to stop traversal immediately, +/// or `ControlFlow::Continue(())` to keep walking. +trait ContainerHandler { + type Output; + + fn handle_lhs_only( + &mut self, + container: &ContainerReader, + rhs_remaining: bool, + ) -> ControlFlow; + fn handle_rhs_only( + &mut self, + container: &ContainerReader, + lhs_remaining: bool, + ) -> ControlFlow; + fn handle_matched( + &mut self, + lhs: &ContainerReader, + rhs: &ContainerReader, + lhs_remaining: bool, + rhs_remaining: bool, + ) -> ControlFlow; + fn handle_none(&mut self) -> Self::Output; } -pub(crate) fn bitmap_stats(buf: &[u8]) -> io::Result { - let tree = TreemapReader::new(buf)?; - - let mut total_len = 0u64; - let mut first_bitmap = None; - let mut last_bitmap = None; +struct ContainerVisitor<'a, H> { + lhs_buf: &'a [u8], + rhs_buf: &'a [u8], + handler: H, +} - for bitmap_result in tree.iter() { - let bitmap = bitmap_result?; - for i in 0..bitmap.containers() { - total_len += bitmap.description(i)?.cardinality() as u64; - } - if first_bitmap.is_none() { - first_bitmap = Some(bitmap.clone()); +impl<'a, H> ContainerVisitor<'a, H> { + fn new(lhs_buf: &'a [u8], rhs_buf: &'a [u8], handler: H) -> Self { + Self { + lhs_buf, + rhs_buf, + handler, } - last_bitmap = Some(bitmap); - } - - if first_bitmap.is_none() || total_len == 0 { - return Ok(BitmapStats { - len: 0, - min: None, - max: None, - }); } - - // Compute min from the first container of the first prefix bucket - let first = first_bitmap.unwrap(); - let first_desc = first.description(0)?; - let first_offset = first.container_offset(0)?; - let first_container_data = &first.bitmap_buf()[first_offset..]; - let first_prefix = first.prefix() as u64; - let first_container_key = first_desc.prefix as u64; - let first_cardinality = first_desc.cardinality(); - let first_low16 = if first_cardinality <= ARRAY_LIMIT { - array_container_first(first_container_data, first_cardinality)? - } else { - bitmap_container_first(first_container_data)? - }; - let min_val = first_prefix << 32 | first_container_key << 16 | first_low16 as u64; - - // Compute max from the last container of the last prefix bucket - let last = last_bitmap.unwrap(); - let last_idx = last.containers() - 1; - let last_desc = last.description(last_idx)?; - let last_offset = last.container_offset(last_idx)?; - let last_container_data = &last.bitmap_buf()[last_offset..]; - let last_prefix = last.prefix() as u64; - let last_container_key = last_desc.prefix as u64; - let last_cardinality = last_desc.cardinality(); - let last_low16 = if last_cardinality <= ARRAY_LIMIT { - array_container_last(last_container_data, last_cardinality)? - } else { - bitmap_container_last(last_container_data)? - }; - let max_val = last_prefix << 32 | last_container_key << 16 | last_low16 as u64; - - Ok(BitmapStats { - len: total_len, - min: Some(min_val), - max: Some(max_val), - }) } -pub(crate) fn bitmap_has_any(lhs: &[u8], rhs: &[u8]) -> io::Result { - use std::cmp::Ordering::*; - let lhs_tree = TreemapReader::new(lhs)?; - let rhs_tree = TreemapReader::new(rhs)?; +impl ContainerVisitor<'_, H> { + fn visit(mut self) -> io::Result { + use std::cmp::Ordering::*; - let mut lhs_iter = lhs_tree.iter(); - let mut rhs_iter = rhs_tree.iter(); + let lhs_tree = TreemapReader::new(self.lhs_buf)?; + let rhs_tree = TreemapReader::new(self.rhs_buf)?; - let mut lhs_curr = lhs_iter.next().transpose()?; - let mut rhs_curr = rhs_iter.next().transpose()?; + let mut lhs_iter = lhs_tree.iter(); + let mut rhs_iter = rhs_tree.iter(); + let mut lhs_cur = lhs_iter.next().transpose()?; + let mut rhs_cur = rhs_iter.next().transpose()?; - while let (Some(lhs_bitmap), Some(rhs_bitmap)) = (lhs_curr.as_ref(), rhs_curr.as_ref()) { - match lhs_bitmap.prefix().cmp(&rhs_bitmap.prefix()) { - Less => { - lhs_curr = lhs_iter.next().transpose()?; - } - Greater => { - rhs_curr = rhs_iter.next().transpose()?; - } - Equal => { - if containers_has_any(lhs_bitmap, rhs_bitmap)? { - return Ok(true); + // Both empty + if lhs_cur.is_none() && rhs_cur.is_none() { + return Ok(self.handler.handle_none()); + } + + // Treemap-level traversal + while let (Some(lhs_bitmap), Some(rhs_bitmap)) = (lhs_cur.as_ref(), rhs_cur.as_ref()) { + match lhs_bitmap.prefix().cmp(&rhs_bitmap.prefix()) { + Less => { + if let ControlFlow::Break(r) = + self.visit_one_side_containers(lhs_bitmap, true, true)? + { + return Ok(r); + } + lhs_cur = lhs_iter.next().transpose()?; + } + Greater => { + if let ControlFlow::Break(r) = + self.visit_one_side_containers(rhs_bitmap, false, true)? + { + return Ok(r); + } + rhs_cur = rhs_iter.next().transpose()?; + } + Equal => { + let next_lhs = lhs_iter.next().transpose()?; + let next_rhs = rhs_iter.next().transpose()?; + let lhs_has_more_prefixes = next_lhs.is_some(); + let rhs_has_more_prefixes = next_rhs.is_some(); + + if let ControlFlow::Break(r) = self.visit_both_sides_containers( + lhs_bitmap, + rhs_bitmap, + lhs_has_more_prefixes, + rhs_has_more_prefixes, + )? { + return Ok(r); + } + + lhs_cur = next_lhs; + rhs_cur = next_rhs; } - lhs_curr = lhs_iter.next().transpose()?; - rhs_curr = rhs_iter.next().transpose()?; } } + + // Handle remaining containers if any + if let ControlFlow::Break(r) = + self.visit_one_side_remaining_containers(lhs_cur.as_ref(), &mut lhs_iter, true)? + { + return Ok(r); + } + if let ControlFlow::Break(r) = + self.visit_one_side_remaining_containers(rhs_cur.as_ref(), &mut rhs_iter, false)? + { + return Ok(r); + } + + Ok(self.handler.handle_none()) } - Ok(false) -} + fn visit_one_side_containers( + &mut self, + bitmap: &BitmapReader<'_>, + is_lhs: bool, + other_remaining: bool, + ) -> io::Result> { + for i in 0..bitmap.containers() { + let container = bitmap.container(i)?; + if let ControlFlow::Break(r) = + self.visit_one_side_container(&container, is_lhs, other_remaining) + { + return Ok(ControlFlow::Break(r)); + } + } + Ok(ControlFlow::Continue(())) + } -fn containers_has_any(lhs: &BitmapReader<'_>, rhs: &BitmapReader<'_>) -> io::Result { - let lhs_count = lhs.containers(); - let rhs_count = rhs.containers(); - let mut i = 0; - let mut j = 0; - - while i < lhs_count && j < rhs_count { - let lhs_desc = lhs.description(i)?; - let rhs_desc = rhs.description(j)?; - match lhs_desc.prefix.cmp(&rhs_desc.prefix) { - std::cmp::Ordering::Less => i += 1, - std::cmp::Ordering::Greater => j += 1, - std::cmp::Ordering::Equal => { - let lhs_offset = lhs.container_offset(i)?; - let rhs_offset = rhs.container_offset(j)?; - let lhs_data = &lhs.bitmap_buf()[lhs_offset..]; - let rhs_data = &rhs.bitmap_buf()[rhs_offset..]; - let lhs_cardinality = lhs_desc.cardinality(); - let rhs_cardinality = rhs_desc.cardinality(); - if container_has_any(lhs_data, lhs_cardinality, rhs_data, rhs_cardinality)? { - return Ok(true); + fn visit_both_sides_containers( + &mut self, + lhs_bitmap: &BitmapReader<'_>, + rhs_bitmap: &BitmapReader<'_>, + lhs_has_more_prefixes: bool, + rhs_has_more_prefixes: bool, + ) -> io::Result> { + use std::cmp::Ordering::*; + + let mut i = 0; + let mut j = 0; + let lhs_count = lhs_bitmap.containers(); + let rhs_count = rhs_bitmap.containers(); + while i < lhs_count && j < rhs_count { + let lhs_container = lhs_bitmap.container(i)?; + let rhs_container = rhs_bitmap.container(j)?; + match lhs_container.key().cmp(&rhs_container.key()) { + Less => { + if let ControlFlow::Break(r) = self.visit_one_side_container( + &lhs_container, + true, // left + true, // rhs has more + ) { + return Ok(ControlFlow::Break(r)); + } + i += 1; + } + Greater => { + if let ControlFlow::Break(r) = self.visit_one_side_container( + &rhs_container, + false, // right + true, // lhs has more + ) { + return Ok(ControlFlow::Break(r)); + } + j += 1; + } + Equal => { + let lhs_remaining = (i + 1 < lhs_count) || lhs_has_more_prefixes; + let rhs_remaining = (j + 1 < rhs_count) || rhs_has_more_prefixes; + if let ControlFlow::Break(r) = self.handler.handle_matched( + &lhs_container, + &rhs_container, + lhs_remaining, + rhs_remaining, + ) { + return Ok(ControlFlow::Break(r)); + } + i += 1; + j += 1; } - i += 1; - j += 1; } } + while i < lhs_count { + let container = lhs_bitmap.container(i)?; + if let ControlFlow::Break(r) = + self.visit_one_side_container(&container, true, rhs_has_more_prefixes) + { + return Ok(ControlFlow::Break(r)); + } + i += 1; + } + while j < rhs_count { + let container = rhs_bitmap.container(j)?; + if let ControlFlow::Break(r) = + self.visit_one_side_container(&container, false, lhs_has_more_prefixes) + { + return Ok(ControlFlow::Break(r)); + } + j += 1; + } + Ok(ControlFlow::Continue(())) } - Ok(false) -} - -fn container_has_any( - lhs_data: &[u8], - lhs_cardinality: usize, - rhs_data: &[u8], - rhs_cardinality: usize, -) -> io::Result { - // Pigeonhole: two subsets of a 65536-element space with combined size > 65536 must overlap - if lhs_cardinality + rhs_cardinality > 65536 { - return Ok(true); - } - - let lhs_is_array = lhs_cardinality <= ARRAY_LIMIT; - let rhs_is_array = rhs_cardinality <= ARRAY_LIMIT; - - if lhs_is_array && rhs_is_array { - array_container_has_any(lhs_data, lhs_cardinality, rhs_data, rhs_cardinality) - } else if !lhs_is_array && !rhs_is_array { - bitmap_container_has_any(lhs_data, rhs_data) - } else if lhs_is_array { - array_bitmap_container_has_any(lhs_data, lhs_cardinality, rhs_data) - } else { - array_bitmap_container_has_any(rhs_data, rhs_cardinality, lhs_data) + fn visit_one_side_remaining_containers( + &mut self, + cur: Option<&BitmapReader<'_>>, + remaining: &mut TreeMapIter<'_>, + is_lhs: bool, + ) -> io::Result> { + if let Some(bitmap) = cur + && let ControlFlow::Break(r) = self.visit_one_side_containers(bitmap, is_lhs, false)? + { + return Ok(ControlFlow::Break(r)); + } + while let Some(bitmap) = remaining.next().transpose()? { + if let ControlFlow::Break(r) = self.visit_one_side_containers(&bitmap, is_lhs, false)? { + return Ok(ControlFlow::Break(r)); + } + } + Ok(ControlFlow::Continue(())) } -} -fn array_container_has_any( - lhs_data: &[u8], - lhs_cardinality: usize, - rhs_data: &[u8], - rhs_cardinality: usize, -) -> io::Result { - if lhs_data.len() < lhs_cardinality * 2 { - return Err(Error::other("lhs array container too short")); - } - if rhs_data.len() < rhs_cardinality * 2 { - return Err(Error::other("rhs array container too short")); - } - - let mut i = 0; - let mut j = 0; - while i < lhs_cardinality && j < rhs_cardinality { - let lv = u16::from_le_bytes(lhs_data[i * 2..i * 2 + 2].try_into().unwrap()); - let rv = u16::from_le_bytes(rhs_data[j * 2..j * 2 + 2].try_into().unwrap()); - if lv < rv { - i += 1; - } else if rv < lv { - j += 1; + fn visit_one_side_container( + &mut self, + container: &ContainerReader, + is_lhs: bool, + other_remaining: bool, + ) -> ControlFlow { + if is_lhs { + self.handler.handle_lhs_only(container, other_remaining) } else { - return Ok(true); + self.handler.handle_rhs_only(container, other_remaining) } } - Ok(false) } -fn bitmap_container_has_any(lhs_data: &[u8], rhs_data: &[u8]) -> io::Result { - if lhs_data.len() < BITMAP_BYTES { - return Err(Error::other("lhs bitmap container too short")); - } - if rhs_data.len() < BITMAP_BYTES { - return Err(Error::other("rhs bitmap container too short")); - } +struct HasAnyHandler; - for word_index in 0..BITMAP_WORDS { - let start = word_index * WORD_BYTES; - let lhs_word = u64::from_le_bytes(lhs_data[start..start + WORD_BYTES].try_into().unwrap()); - let rhs_word = u64::from_le_bytes(rhs_data[start..start + WORD_BYTES].try_into().unwrap()); - if lhs_word & rhs_word != 0 { - return Ok(true); +impl ContainerHandler for HasAnyHandler { + type Output = bool; + + fn handle_lhs_only(&mut self, _: &ContainerReader, rhs_remaining: bool) -> ControlFlow { + if !rhs_remaining { + return ControlFlow::Break(false); } + ControlFlow::Continue(()) } - Ok(false) -} -fn array_bitmap_container_has_any( - array_data: &[u8], - array_cardinality: usize, - bitmap_data: &[u8], -) -> io::Result { - if array_data.len() < array_cardinality * 2 { - return Err(Error::other("array container too short")); - } - if bitmap_data.len() < BITMAP_BYTES { - return Err(Error::other("bitmap container too short")); + fn handle_rhs_only(&mut self, _: &ContainerReader, lhs_remaining: bool) -> ControlFlow { + if !lhs_remaining { + return ControlFlow::Break(false); + } + ControlFlow::Continue(()) } - for i in 0..array_cardinality { - let value = u16::from_le_bytes(array_data[i * 2..i * 2 + 2].try_into().unwrap()); - let word_index = value as usize / WORD_BITS; - let bit_index = value as usize % WORD_BITS; - let start = word_index * WORD_BYTES; - let word = u64::from_le_bytes(bitmap_data[start..start + WORD_BYTES].try_into().unwrap()); - if word & (1 << bit_index) != 0 { - return Ok(true); + fn handle_matched( + &mut self, + lhs: &ContainerReader, + rhs: &ContainerReader, + lhs_remaining: bool, + rhs_remaining: bool, + ) -> ControlFlow { + if lhs.has_any_with(rhs) { + ControlFlow::Break(true) + } else if !lhs_remaining || !rhs_remaining { + ControlFlow::Break(false) + } else { + ControlFlow::Continue(()) } } - Ok(false) -} -pub(crate) fn bitmap_has_all(lhs: &[u8], rhs: &[u8]) -> io::Result { - use std::cmp::Ordering::*; - let lhs_tree = TreemapReader::new(lhs)?; - let rhs_tree = TreemapReader::new(rhs)?; + fn handle_none(&mut self) -> bool { + false + } +} - let mut lhs_iter = lhs_tree.iter(); - let mut rhs_iter = rhs_tree.iter(); +struct HasAllHandler; - let mut lhs_curr = lhs_iter.next().transpose()?; - let mut rhs_curr = rhs_iter.next().transpose()?; +impl ContainerHandler for HasAllHandler { + type Output = bool; - while let (Some(lhs_bitmap), Some(rhs_bitmap)) = (lhs_curr.as_ref(), rhs_curr.as_ref()) { - match lhs_bitmap.prefix().cmp(&rhs_bitmap.prefix()) { - Less => { - lhs_curr = lhs_iter.next().transpose()?; - } - Greater => { - // rhs has a prefix bucket that lhs doesn't have — lhs cannot contain all rhs values - return Ok(false); - } - Equal => { - if !containers_has_all(lhs_bitmap, rhs_bitmap)? { - return Ok(false); - } - lhs_curr = lhs_iter.next().transpose()?; - rhs_curr = rhs_iter.next().transpose()?; - } - } + fn handle_lhs_only(&mut self, _: &ContainerReader, _: bool) -> ControlFlow { + ControlFlow::Continue(()) } - // If rhs still has remaining prefix buckets, lhs doesn't cover them - if rhs_curr.is_some() { - return Ok(false); + fn handle_rhs_only(&mut self, _: &ContainerReader, _: bool) -> ControlFlow { + ControlFlow::Break(false) } - Ok(true) -} - -fn containers_has_all(lhs: &BitmapReader<'_>, rhs: &BitmapReader<'_>) -> io::Result { - let lhs_count = lhs.containers(); - let rhs_count = rhs.containers(); - let mut i = 0; - let mut j = 0; - - while j < rhs_count { - // If lhs is exhausted but rhs still has containers, lhs can't contain them - if i >= lhs_count { - return Ok(false); - } - let lhs_desc = lhs.description(i)?; - let rhs_desc = rhs.description(j)?; - match lhs_desc.prefix.cmp(&rhs_desc.prefix) { - std::cmp::Ordering::Less => i += 1, - std::cmp::Ordering::Greater => { - // rhs has a container key that lhs doesn't have - return Ok(false); - } - std::cmp::Ordering::Equal => { - let lhs_offset = lhs.container_offset(i)?; - let rhs_offset = rhs.container_offset(j)?; - let lhs_data = &lhs.bitmap_buf()[lhs_offset..]; - let rhs_data = &rhs.bitmap_buf()[rhs_offset..]; - let lhs_cardinality = lhs_desc.cardinality(); - let rhs_cardinality = rhs_desc.cardinality(); - if !container_has_all(lhs_data, lhs_cardinality, rhs_data, rhs_cardinality)? { - return Ok(false); - } - i += 1; - j += 1; - } + fn handle_matched( + &mut self, + lhs: &ContainerReader, + rhs: &ContainerReader, + lhs_remaining: bool, + rhs_remaining: bool, + ) -> ControlFlow { + if !lhs.has_all_with(rhs) || (!lhs_remaining && rhs_remaining) { + ControlFlow::Break(false) + } else { + ControlFlow::Continue(()) } } - Ok(true) + fn handle_none(&mut self) -> bool { + true + } } -fn container_has_all( - lhs_data: &[u8], - lhs_cardinality: usize, - rhs_data: &[u8], - rhs_cardinality: usize, -) -> io::Result { - // lhs cannot contain all of rhs if lhs has fewer values - if lhs_cardinality < rhs_cardinality { - return Ok(false); - } - // If lhs is the full bitmap (65536 values), it trivially contains any rhs subset - if lhs_cardinality == 65536 { - return Ok(true); - } - - let lhs_is_array = lhs_cardinality <= ARRAY_LIMIT; - let rhs_is_array = rhs_cardinality <= ARRAY_LIMIT; - - if lhs_is_array && rhs_is_array { - array_container_has_all_array(lhs_data, lhs_cardinality, rhs_data, rhs_cardinality) - } else if !lhs_is_array && !rhs_is_array { - bitmap_container_has_all_bitmap(lhs_data, rhs_data) - } else if lhs_is_array { - // lhs is array, rhs is bitmap: array container cardinality < 4096, - // bitmap container cardinality >= 4096, so lhs can never be superset of rhs - // Fast path: rhs has more values than lhs can contain - Ok(false) - } else { - // lhs is bitmap, rhs is array: check each rhs value is in lhs bitmap - array_values_all_in_bitmap_container(rhs_data, rhs_cardinality, lhs_data) - } +pub fn bitmap_len(buf: &[u8]) -> io::Result { + TreemapReader::new(buf)?.len() } -fn array_container_has_all_array( - lhs_data: &[u8], - lhs_cardinality: usize, - rhs_data: &[u8], - rhs_cardinality: usize, -) -> io::Result { - if lhs_data.len() < lhs_cardinality * 2 { - return Err(Error::other("lhs array container too short")); - } - if rhs_data.len() < rhs_cardinality * 2 { - return Err(Error::other("rhs array container too short")); - } +pub(crate) fn bitmap_len_above(buf: &[u8], threshold: usize) -> io::Result { + TreemapReader::new(buf)?.len_above(threshold) +} - // Fast rejection: lhs must have at least as many values as rhs - if lhs_cardinality < rhs_cardinality { - return Ok(false); - } +pub(crate) fn bitmap_contains(buf: &[u8], value: u64) -> io::Result { + TreemapReader::new(buf)?.contains(value) +} - let mut i = 0; - let mut j = 0; - while j < rhs_cardinality { - if i >= lhs_cardinality { - return Ok(false); - } - let lv = u16::from_le_bytes(lhs_data[i * 2..i * 2 + 2].try_into().unwrap()); - let rv = u16::from_le_bytes(rhs_data[j * 2..j * 2 + 2].try_into().unwrap()); - if lv < rv { - i += 1; - } else if lv == rv { - i += 1; - j += 1; - } else { - return Ok(false); - } - } +pub(crate) fn bitmap_min(buf: &[u8]) -> io::Result> { + TreemapReader::new(buf)?.min() +} - Ok(true) +pub(crate) fn bitmap_max(buf: &[u8]) -> io::Result> { + TreemapReader::new(buf)?.max() } -fn bitmap_container_has_all_bitmap(lhs_data: &[u8], rhs_data: &[u8]) -> io::Result { - if lhs_data.len() < BITMAP_BYTES { - return Err(Error::other("lhs bitmap container too short")); - } - if rhs_data.len() < BITMAP_BYTES { - return Err(Error::other("rhs bitmap container too short")); - } +pub(crate) fn bitmap_has_any(lhs: &[u8], rhs: &[u8]) -> io::Result { + ContainerVisitor::new(lhs, rhs, HasAnyHandler).visit() +} - for word_index in 0..BITMAP_WORDS { - let start = word_index * WORD_BYTES; - let lhs_word = u64::from_le_bytes(lhs_data[start..start + WORD_BYTES].try_into().unwrap()); - let rhs_word = u64::from_le_bytes(rhs_data[start..start + WORD_BYTES].try_into().unwrap()); - // Check if all bits set in rhs are also set in lhs - // (rhs_word & !lhs_word) means value exists in rhs but not in lhs - if rhs_word & !lhs_word != 0 { - return Ok(false); - } - } - Ok(true) +pub(crate) fn bitmap_has_all(lhs: &[u8], rhs: &[u8]) -> io::Result { + ContainerVisitor::new(lhs, rhs, HasAllHandler).visit() } -fn array_values_all_in_bitmap_container( - array_data: &[u8], - array_cardinality: usize, - bitmap_data: &[u8], -) -> io::Result { - if array_data.len() < array_cardinality * 2 { - return Err(Error::other("array container too short")); - } - if bitmap_data.len() < BITMAP_BYTES { - return Err(Error::other("bitmap container too short")); - } +pub fn intersection_with_serialized(tree: &mut RoaringTreemap, buf: &[u8]) -> io::Result<()> { + use std::cmp::Ordering::*; + let rhs = TreemapReader::new(buf)?; + let mut bitmaps = Vec::new(); + let mut lhs_iter = tree.bitmaps(); + let mut rhs_iter = rhs.iter(); - for i in 0..array_cardinality { - let value = u16::from_le_bytes(array_data[i * 2..i * 2 + 2].try_into().unwrap()); - let word_index = value as usize / WORD_BITS; - let bit_index = value as usize % WORD_BITS; - let start = word_index * WORD_BYTES; - let word = u64::from_le_bytes(bitmap_data[start..start + WORD_BYTES].try_into().unwrap()); - if word & (1 << bit_index) == 0 { - return Ok(false); + let mut lhs_curr = lhs_iter.next(); + let mut rhs_curr = rhs_iter.next().transpose()?; + + while let (Some((lhs_prefix, lhs_bitmap)), Some(rhs_bitmap)) = (lhs_curr, rhs_curr.as_ref()) { + match lhs_prefix.cmp(&rhs_bitmap.prefix()) { + Less => { + lhs_curr = lhs_iter.next(); + } + Greater => { + rhs_curr = rhs_iter.next().transpose()?; + } + Equal => { + let intersection = lhs_bitmap + .intersection_with_serialized_unchecked(Cursor::new(rhs_bitmap.bitmap_buf()))?; + if !intersection.is_empty() { + bitmaps.push((lhs_prefix, intersection)); + } + lhs_curr = lhs_iter.next(); + rhs_curr = rhs_iter.next().transpose()?; + } } } - Ok(true) + + *tree = RoaringTreemap::from_bitmaps(bitmaps); + + Ok(()) } #[cfg(test)] @@ -915,4 +1029,358 @@ mod tests { Ok(()) } + + // Tests for minimum deserialize bitmap functions. + // + // Cover public functions by comparing against RoaringTreemap: + // - [x] bitmap_contains + // - [x] bitmap_min + // - [x] bitmap_max + // - [x] bitmap_has_any + // - [x] bitmap_has_all + // + // Fixtures: + // - `for_each_fixture` for single-bitmap tests (contains, min, max) + // - `for_each_fixture_pair` for two-bitmap tests (has_any, has_all) + fn make_buf(tree: &RoaringTreemap) -> Vec { + let mut buf = Vec::new(); + tree.serialize_into(&mut buf).unwrap(); + buf + } + + fn for_each_fixture(mut f: F) + where F: FnMut(&str, &[u8], &RoaringTreemap, u64) { + let fixtures: Vec<(&str, RoaringTreemap, u64)> = vec![ + ("random", create_bitmap(123), u64::MAX), + ( + "sparse", + RoaringTreemap::from_iter([0u64, 65535, 65536, (1u64 << 32) + 500]), + 1, + ), + // p0 keys 0,1,2 (bitmap) multi-key bitmap + ( + "key: multi-array", + (0u64..0x10000) + .step_by(7) + .chain((0x10000u64..0x20000).step_by(7)) + .chain((0x20000u64..0x30000).step_by(7)) + .collect(), + 0x40000, + ), + // p0 keys 0,1 (bitmap) multi-key bitmap + ( + "key: multi-bitmap", + (0u64..50000).chain(0x10000u64..0x10000 + 50000).collect(), + 0x40000, + ), + // p0, p1 (bitmap) multi-prefix + ( + "prefix: multi", + (0u64..0x10000) + .step_by(7) + .chain(((1u64 << 32)..(1u64 << 32) + 0x10000).step_by(7)) + .collect(), + (2u64 << 32), + ), + ]; + for (name, tree, miss) in fixtures { + let buf = make_buf(&tree); + f(name, &buf, &tree, miss); + } + let empty_tree = RoaringTreemap::new(); + let empty_buf = make_buf(&empty_tree); + f("empty", &empty_buf, &empty_tree, 0); + } + + fn for_each_fixture_pair(mut f: F) + where F: FnMut(&str, &[u8], &[u8], &RoaringTreemap, &RoaringTreemap) { + let scenarios: Vec<(&str, RoaringTreemap, RoaringTreemap)> = vec![ + // lhs: p0 key 0 (bitmap) | rhs: p0 key 0 (bitmap) overlap + ( + "basic: overlap", + (0..50000u64).collect(), + (30000..80000u64).collect(), + ), + // lhs: p0 key 0 (bitmap) | rhs: p0 keys 1,2 (bitmap) disjoint + ( + "basic: disjoint", + (0..50000u64).collect(), + (100000..150000u64).collect(), + ), + // lhs: p0 key 0 (bitmap) | rhs: p0 key 0 (array) superset + ( + "basic: superset", + (0..50000u64).collect(), + (0..100u64).collect(), + ), + // lhs: p0 key 0 (bitmap) | rhs: p0 keys 0,1 (bitmap) not superset + ( + "basic: not superset", + (0..50000u64).collect(), + (40000..90000u64).collect(), + ), + // lhs: p0 key 0 (bitmap) | rhs: p0 key 0 (bitmap) self + ( + "basic: self", + (0..50000u64).collect(), + (0..50000u64).collect(), + ), + // lhs: p0 key 2 (bitmap) | rhs: p0 keys 1,2 (bitmap) rhs-only key before lhs + ( + "key: rhs-only before lhs", + (0x20000u64..0x30000).step_by(3).collect(), + (0x10000u64..0x30000).step_by(2).collect(), + ), + // lhs: p0 keys 0,1,2 (bitmap) | rhs: p0 keys 1,2 (bitmap) lhs-only key before rhs + ( + "key: lhs-only before rhs", + (0u64..0x10000) + .step_by(7) + .chain((0x10000u64..0x20000).step_by(7)) + .chain((0x20000u64..0x30000).step_by(7)) + .collect(), + (0x10000u64..0x20000) + .step_by(5) + .chain((0x20000u64..0x30000).step_by(5)) + .collect(), + ), + // lhs: p0 key 1 (bitmap) | rhs: p0 keys 1,2 (bitmap) first key shared, rhs has more + ( + "key: lhs first shared", + (0x10000u64..0x20000).step_by(3).collect(), + (0x10000u64..0x30000).step_by(2).collect(), + ), + // lhs: p0 keys 1,2 (bitmap) | rhs: p0 key 1 (bitmap) first key shared, lhs has more + ( + "key: rhs first shared", + (0x10000u64..0x30000).step_by(2).collect(), + (0x10000u64..0x20000).step_by(3).collect(), + ), + // lhs: p0 keys 0,2 (bitmap) | rhs: p0 keys 0,1,2 (bitmap) rhs-only key between lhs keys + ( + "key: lhs skips middle", + (0u64..0x10000) + .step_by(7) + .chain((0x20000u64..0x30000).step_by(7)) + .collect(), + (0u64..0x10000) + .step_by(5) + .chain((0x10000u64..0x20000).step_by(5)) + .chain((0x20000u64..0x30000).step_by(5)) + .collect(), + ), + // lhs: p0 keys 0,1,2 (bitmap) | rhs: p0 keys 0,2 (bitmap) lhs-only key between rhs keys + ( + "key: rhs skips middle", + (0u64..0x10000) + .step_by(7) + .chain((0x10000u64..0x20000).step_by(7)) + .chain((0x20000u64..0x30000).step_by(7)) + .collect(), + (0u64..0x10000) + .step_by(5) + .chain((0x20000u64..0x30000).step_by(5)) + .collect(), + ), + // lhs: p0 keys 0,1,2 (bitmap) | rhs: p0 keys 0,1 (bitmap) lhs-only key after rhs last + ( + "key: lhs-only after rhs", + (0u64..0x10000) + .step_by(7) + .chain((0x10000u64..0x20000).step_by(7)) + .chain((0x20000u64..0x30000).step_by(7)) + .collect(), + (0u64..0x10000) + .step_by(5) + .chain((0x10000u64..0x20000).step_by(5)) + .collect(), + ), + // lhs: p0 keys 0,1 (bitmap) | rhs: p0 keys 0,1,2 (bitmap) rhs-only key after lhs last + ( + "key: rhs-only after lhs", + (0u64..0x10000) + .step_by(7) + .chain((0x10000u64..0x20000).step_by(7)) + .collect(), + (0u64..0x10000) + .step_by(5) + .chain((0x10000u64..0x20000).step_by(5)) + .chain((0x20000u64..0x30000).step_by(5)) + .collect(), + ), + // lhs: p0 keys 0,2 (bitmap) | rhs: p0 keys 1,3 (bitmap) no matching keys + ( + "key: interleaved", + (0u64..0x10000) + .step_by(7) + .chain((0x20000u64..0x30000).step_by(7)) + .collect(), + (0x10000u64..0x20000) + .step_by(5) + .chain((0x30000u64..0x40000).step_by(5)) + .collect(), + ), + // lhs: p0 (bitmap) | rhs: p0, p1 (bitmap) rhs-only prefix + ( + "prefix: rhs-only", + (0u64..0x10000).step_by(7).collect(), + (0u64..0x10000) + .step_by(5) + .chain(((1u64 << 32)..(1u64 << 32) + 0x10000).step_by(5)) + .collect(), + ), + // lhs: p0, p1 (bitmap) | rhs: p0 (bitmap) lhs-only prefix + ( + "prefix: lhs-only", + (0u64..0x10000) + .step_by(7) + .chain(((1u64 << 32)..(1u64 << 32) + 0x10000).step_by(7)) + .collect(), + (0u64..0x10000).step_by(5).collect(), + ), + // lhs: p0 (bitmap) | rhs: p1 (bitmap) disjoint prefixes + ( + "prefix: disjoint", + (0u64..0x10000).step_by(7).collect(), + ((1u64 << 32)..(1u64 << 32) + 0x10000).step_by(5).collect(), + ), + // lhs: p0 key 0 (array) | rhs: p0 key 0 (bitmap) matched array+bitmap + ( + "store: array+bitmap", + (0u64..1000).collect(), + (0u64..50000).collect(), + ), + // lhs: p0 key 0 (bitmap) | rhs: p0 key 0 (array) matched bitmap+array + ( + "store: bitmap+array", + (0u64..50000).collect(), + (0u64..1000).collect(), + ), + // lhs: p0 key 0 (array), key 1 (bitmap) | rhs: p0 key 0 (bitmap), key 1 (array) mixed store across keys + ( + "store: mixed", + (0u64..1000).chain(0x10000u64..0x10000 + 50000).collect(), + (0u64..50000).chain(0x10000u64..0x10000 + 1000).collect(), + ), + ]; + for (name, lhs_tree, rhs_tree) in scenarios { + let lhs_buf = make_buf(&lhs_tree); + let rhs_buf = make_buf(&rhs_tree); + f(name, &lhs_buf, &rhs_buf, &lhs_tree, &rhs_tree); + } + // Empty scenarios + let empty_tree = RoaringTreemap::new(); + let empty_buf = make_buf(&empty_tree); + let non_empty: RoaringTreemap = (0..50000u64).collect(); + let non_empty_buf = make_buf(&non_empty); + f( + "empty+non-empty", + &empty_buf, + &non_empty_buf, + &empty_tree, + &non_empty, + ); + f( + "non-empty+empty", + &non_empty_buf, + &empty_buf, + &non_empty, + &empty_tree, + ); + f( + "empty+empty", + &empty_buf, + &empty_buf, + &empty_tree, + &empty_tree, + ); + } + + #[test] + fn test_bitmap_contains() -> io::Result<()> { + for_each_fixture(|name, buf, tree, miss_value| { + // Test a known-present value (min if exists) + if let Some(hit) = tree.min() { + let expected = tree.contains(hit); + let actual = bitmap_contains(buf, hit).unwrap(); + assert_eq!( + actual, expected, + "bitmap_contains hit: fixture={name}, val={hit}" + ); + } + // Test a known-absent value + let miss = miss_value; + let expected = tree.contains(miss); + let actual = bitmap_contains(buf, miss).unwrap(); + assert_eq!( + actual, expected, + "bitmap_contains miss: fixture={name}, val={miss}" + ); + }); + Ok(()) + } + + #[test] + fn test_bitmap_min() -> io::Result<()> { + for_each_fixture(|name, buf, tree, _miss_value| { + let expected = tree.min(); + let actual = bitmap_min(buf).unwrap(); + assert_eq!(actual, expected, "bitmap_min: fixture={name}"); + }); + Ok(()) + } + + #[test] + fn test_bitmap_max() -> io::Result<()> { + for_each_fixture(|name, buf, tree, _miss_value| { + let expected = tree.max(); + let actual = bitmap_max(buf).unwrap(); + assert_eq!(actual, expected, "bitmap_max: fixture={name}"); + }); + Ok(()) + } + + #[test] + fn test_bitmap_has_any() -> io::Result<()> { + for_each_fixture_pair(|name, lhs_buf, rhs_buf, lhs_tree, rhs_tree| { + let expected = !(lhs_tree & rhs_tree).is_empty(); + let actual = bitmap_has_any(lhs_buf, rhs_buf).unwrap(); + assert_eq!(actual, expected, "bitmap_has_any: fixture={name}"); + }); + Ok(()) + } + + #[test] + fn test_bitmap_has_all() -> io::Result<()> { + for_each_fixture_pair(|name, lhs_buf, rhs_buf, lhs_tree, rhs_tree| { + let expected = lhs_tree.is_superset(rhs_tree); + let actual = bitmap_has_all(lhs_buf, rhs_buf).unwrap(); + assert_eq!(actual, expected, "bitmap_has_all: fixture={name}"); + }); + Ok(()) + } + + #[test] + fn test_bitmap_len_above() -> io::Result<()> { + let bitmap = create_bitmap(123); + let mut buf = Vec::new(); + bitmap.serialize_into(&mut buf)?; + + let len = bitmap_len(&buf)?; + + // Threshold below actual length -> true + assert!(bitmap_len_above(&buf, len - 1)?); + // Threshold at actual length -> false + assert!(!bitmap_len_above(&buf, len)?); + // Threshold above actual length -> false + assert!(!bitmap_len_above(&buf, len + 1)?); + + // Empty bitmap + let empty = RoaringTreemap::new(); + let mut empty_buf = Vec::new(); + empty.serialize_into(&mut empty_buf)?; + assert!(!bitmap_len_above(&empty_buf, 0)?); + + Ok(()) + } } diff --git a/src/common/io/src/lib.rs b/src/common/io/src/lib.rs index 69c9af96b9a09..b0f13040b83c4 100644 --- a/src/common/io/src/lib.rs +++ b/src/common/io/src/lib.rs @@ -46,7 +46,6 @@ mod stat_buffer; pub mod interval; pub mod wkb; -pub use bitmap::BitmapStats; pub use bitmap::HYBRID_HEADER_LEN; pub use bitmap::HYBRID_KIND_LARGE; pub use bitmap::HYBRID_KIND_SMALL; @@ -59,7 +58,6 @@ pub use bitmap::bitmap_has_all; pub use bitmap::bitmap_has_any; pub use bitmap::bitmap_max; pub use bitmap::bitmap_min; -pub use bitmap::bitmap_stats; pub use bitmap::deserialize_bitmap; pub use bitmap::parse_bitmap; pub use decimal::display_decimal_128; From 33fabbd826b20d57da1dbc3468815f9ab29eb921 Mon Sep 17 00:00:00 2001 From: HarryHao Date: Thu, 30 Jul 2026 13:05:48 +0800 Subject: [PATCH 8/9] chore: add proptest for newly introduced bitmap functions This commit introduced [proptests](https://docs.rs/proptest) for newly introduced: - bitmap_contains - bitmap_min - bitmap_max - bitmap_has_any - bitmap_has_all Each compares the result against RoaringTreemap as ground truth. Strategy design inspired by roaring as: - Store (random bits) - Container (50% Array, 50% Bitmap) - Bitmap (0 to 16 containers) - Tree (0 to 16 Bitmaps) - Serialization (80% Hybrid, 10% Legacy, 10% Empty) --- Cargo.lock | 1 + src/common/io/Cargo.toml | 1 + src/common/io/src/bitmap.rs | 193 ++++++++++++++++++++++++++++++++++++ 3 files changed, 195 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index e0dbea911b197..e095dfe305e88 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4391,6 +4391,7 @@ dependencies = [ "jiff", "lexical-core", "micromarshal", + "proptest", "rand 0.8.5", "rmp-serde", "roaring 0.10.12", diff --git a/src/common/io/Cargo.toml b/src/common/io/Cargo.toml index 5ca0b339c76f7..c7e07c18ab08b 100644 --- a/src/common/io/Cargo.toml +++ b/src/common/io/Cargo.toml @@ -34,6 +34,7 @@ wkt = { workspace = true } [dev-dependencies] aho-corasick = { workspace = true } anyhow = { workspace = true } +proptest = { workspace = true } rand = { workspace = true } rmp-serde = { workspace = true } diff --git a/src/common/io/src/bitmap.rs b/src/common/io/src/bitmap.rs index 6d29be51e4932..4aa53c400fd77 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -2131,4 +2131,197 @@ mod tests { assert_eq!(actual, expected, "bitmap_has_all: fixture={name}"); }); } + + /// proptests for bitmap functions: + /// + /// - [x] bitmap_contains + /// - [x] bitmap_min + /// - [x] bitmap_max + /// - [x] bitmap_has_any + /// - [x] bitmap_has_all + /// + /// According to comment in src/common/column/tests/it/bitmap/assign_ops.rs, + /// following prop tests are ignored when using `miri`. + #[cfg_attr(miri, ignore)] + mod proptests { + use proptest::bits::BitSetLike; + use proptest::bits::SampledBitSetStrategy; + use proptest::collection::SizeRange; + use proptest::collection::btree_map; + use proptest::prelude::*; + use roaring::RoaringBitmap; + + use super::*; + + /// The random bits strategy. + /// + /// Generate random bits with [`Store::sampled`]. + #[derive(Clone, Debug)] + struct Store(Vec); + + impl Store { + fn sampled( + size: impl Into, + bits: impl Into, + ) -> SampledBitSetStrategy { + SampledBitSetStrategy::new(size.into(), bits.into()) + } + } + + /// Implement BitSetLike as required by SampledBitSetStrategy. + impl BitSetLike for Store { + fn new_bitset(max: usize) -> Self { + assert!(max <= u16::MAX as usize + 1); + Store(Vec::new()) + } + fn len(&self) -> usize { + u16::MAX as usize + 1 + } + fn test(&self, bit: usize) -> bool { + self.0.binary_search(&(bit as u16)).is_ok() + } + fn set(&mut self, bit: usize) { + let v = bit as u16; + if let Err(pos) = self.0.binary_search(&v) { + self.0.insert(pos, v); + } + } + fn clear(&mut self, bit: usize) { + let v = bit as u16; + if let Ok(pos) = self.0.binary_search(&v) { + self.0.remove(pos); + } + } + fn count(&self) -> usize { + self.0.len() + } + } + + /// The container strategy. + /// + /// Generates: + /// 50% array container with 1 to 4096 values + /// 50% bitmap container with 4097 to 65536 values + fn container_strategy() -> impl Strategy> { + prop_oneof![ + Store::sampled(1..=4096, ..=u16::MAX as usize).prop_map(|bs| bs.0), + Store::sampled(4097..u16::MAX as usize, ..=u16::MAX as usize).prop_map(|bs| bs.0), + ] + } + + /// The RoaringBitmap strategy. + /// + /// Generate 0 to 16 containers, with containers from [`container_strategy`]. + fn bitmap_strategy() -> impl Strategy { + btree_map(0u16..=16, container_strategy(), 0usize..=16).prop_map(|map| { + let mut bitmap = RoaringBitmap::new(); + for (key, values) in map { + for v in values { + bitmap.insert((key as u32) << 16 | v as u32); + } + } + bitmap + }) + } + + /// The RoaringTreemap strategy. + /// + /// Generate 0 to 16 RoaringBitmaps, each generated by [`bitmap_strategy`]. + fn tree_strategy() -> impl Strategy { + btree_map(0u32..=16, bitmap_strategy(), 0usize..=16).prop_map(|map| { + let mut treemap = RoaringTreemap::new(); + for (key, bitmap) in map { + if !bitmap.is_empty() { + for v in bitmap.iter() { + treemap.insert((key as u64) << 32 | v as u64); + } + } + } + treemap + }) + } + + /// The serialization strategy. + /// + /// Serialize RoaringTreemap from [`tree_strategy`] into `Vec`: + /// 80% HybridBitmap + /// 10% Legacy + /// 10% Empty + fn serialization_strategy() -> impl Strategy, RoaringTreemap)> { + prop_oneof![ + 8 => tree_strategy().prop_map(|tree| { + let bm = HybridBitmap::from_iter(tree.iter()); + let mut buf = Vec::new(); + bm.serialize_into(&mut buf).unwrap(); + (buf, tree) + }), + 1 => tree_strategy().prop_map(|tree| { + let mut buf = Vec::new(); + tree.serialize_into(&mut buf).unwrap(); + (buf, tree) + }), + 1 => Just((Vec::new(), RoaringTreemap::new())), + ] + } + + /// The probe strategy, picks random value for bitmap_contains to probe. + /// + /// 50% hit: probe = random value from tree (via nth) + /// 50% miss: probe = random u64 + fn probe_strategy() -> impl Strategy, RoaringTreemap, u64)> { + (serialization_strategy(), any::(), any::()).prop_map( + |((buf, tree), is_hit, random_value)| { + let probe = if is_hit && !tree.is_empty() { + let idx = (random_value % tree.len()) as usize; + tree.iter().nth(idx).unwrap_or(random_value) + } else { + random_value + }; + (buf, tree, probe) + }, + ) + } + + proptest! { + // Make the test run faster by limiting the number of cases, running 32 cases + // took ~16 seconds in release, and ~160 in debug. + // One can override this by setting the `PROPTEST_CASES` environment variable. + #![proptest_config(ProptestConfig::with_cases(32))] + + #[test] + fn prop_bitmap_contains((buf, tree, value) in probe_strategy()) { + let expected = tree.contains(value); + let actual = bitmap_contains(&buf, value).unwrap(); + assert_eq!(actual, expected); + } + + #[test] + fn prop_bitmap_min((buf, tree) in serialization_strategy()) { + assert_eq!(bitmap_min(&buf).unwrap(), tree.min()); + } + + #[test] + fn prop_bitmap_max((buf, tree) in serialization_strategy()) { + assert_eq!(bitmap_max(&buf).unwrap(), tree.max()); + } + + #[test] + fn prop_bitmap_has_any( + ((lhs_buf, lhs_tree), (rhs_buf, rhs_tree)) in (serialization_strategy(), serialization_strategy()) + ) { + let expected = !(lhs_tree & rhs_tree).is_empty(); + let actual = bitmap_has_any(&lhs_buf, &rhs_buf).unwrap(); + assert_eq!(actual, expected); + } + + #[test] + fn prop_bitmap_has_all( + ((lhs_buf, lhs_tree), (rhs_buf, rhs_tree)) in (serialization_strategy(), serialization_strategy()) + ) { + let expected = lhs_tree.is_superset(&rhs_tree); + let actual = bitmap_has_all(&lhs_buf, &rhs_buf).unwrap(); + assert_eq!(actual, expected); + } + } + } } From 29bbf638f6f4593e0d704943d4488f5004f6f9a2 Mon Sep 17 00:00:00 2001 From: HarryHao Date: Thu, 30 Jul 2026 18:44:48 +0800 Subject: [PATCH 9/9] fix: stop assuming HybridBitmap::Large having higher cardinality than Small Previously, when bitmap_has_all got a Small lhs and a Large rhs, it simply return false, since a set with fewer cardinality can't has_all another one with higher cardinality. But due to we don't always demote Large to Small after cardinality reducing operations, that's not true anymore. As a defense, we removed this branch, and the doubting rhs will fall into branch below, where rhs cardinality is actually checked before returning false. --- src/common/io/src/bitmap.rs | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/common/io/src/bitmap.rs b/src/common/io/src/bitmap.rs index 4aa53c400fd77..25f01435d9ea4 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -947,16 +947,14 @@ pub fn bitmap_has_all(lhs: &[u8], rhs: &[u8]) -> Result { return Ok(lhs.has_all_with(&rhs)); } - // Fast path: HybridSmall lhs cannot contain HybridLarge rhs - if as_roaring(lhs).is_none() && is_hybrid_large(rhs) { - return Ok(false); - } - - // Fast path: HybridSmall lhs, Legacy rhs — cardinality check + // Fast path: HybridSmall lhs, HybridLarge or Legacy rhs cardinality check + // We need to check cardinality here because: (1) it might be a legacy tree, + // and (2) we don't always demote HybridBitmap after reducing its cardinality immediately. if as_roaring(lhs).is_none() && as_roaring(rhs).is_some() { let lhs_small = SmallReader::new(lhs)?; + let rhs_roaring = as_roaring(rhs).unwrap(); // Fast path: rhs has more values than lhs, impossible to contain - if reader::bitmap_len_above(rhs, lhs_small.len())? { + if reader::bitmap_len_above(rhs_roaring, lhs_small.len())? { return Ok(false); } // rhs has <= lhs_small.len() values, but still need to check containment