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 cc614cc3e3fc8..25f01435d9ea4 100644 --- a/src/common/io/src/bitmap.rs +++ b/src/common/io/src/bitmap.rs @@ -728,23 +728,292 @@ 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 +} + +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 +/// 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 { + // Fast path: empty bitmap contains nothing + if buf.is_empty() { + return Ok(false); + } + validate_serialized_bitmap_header_and_length(buf)?; + if let Some(roaring_buf) = as_roaring(buf) { + Ok(reader::bitmap_contains(roaring_buf, value)?) + } else { + 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 let Some(roaring_buf) = as_roaring(buf) { + Ok(reader::bitmap_min(roaring_buf)?) + } else { + 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 let Some(roaring_buf) = as_roaring(buf) { + Ok(reader::bitmap_max(roaring_buf)?) + } else { + 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(bitmap_len(rhs)? == 0); + } + validate_serialized_bitmap_header_and_length(lhs)?; + validate_serialized_bitmap_header_and_length(rhs)?; + + // 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, 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_roaring, lhs_small.len())? { + return Ok(false); + } + // rhs has <= lhs_small.len() values, but still need to check containment + let lhs_bm = deserialize_bitmap(lhs)?; + let rhs_bm = deserialize_bitmap(rhs)?; + 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)?; + + // 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)?); + } + + // 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)); + } + + // One side HybridLarge/Legacy, the other HybridSmall: probe + let (roaring_buf, small) = if let Some(rb) = as_roaring(lhs) { + (rb, SmallReader::new(rhs)?) + } else { + (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> { 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 +1037,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 => { @@ -1154,6 +1423,7 @@ fn small_intersection_len(lhs: &SmallBitmap, rhs: &SmallBitmap) -> u64 { #[cfg(test)] mod tests { + use smallvec::smallvec; use super::*; @@ -1429,4 +1699,627 @@ mod tests { let decoded = deserialize_bitmap(&legacy).unwrap(); assert_eq!(decoded.into_iter().collect::>(), vec![1, 5, 42]); } + + // 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 + } + + fn make_buf(tree: &RoaringTreemap) -> Vec { + let bm = HybridBitmap::from_iter(tree.iter()); + let mut buf = Vec::new(); + 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, + ); + // 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, + ); + // 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(), + ); + // 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(), + ); + // 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(), + ); + // lhs: empty | rhs: HE + f( + "format: Empty+HE", + &[], + &he_buf, + &RoaringTreemap::new(), + &he_tree, + ); + // lhs: HE | rhs: empty + f( + "format: HE+Empty", + &he_buf, + &[], + &he_tree, + &RoaringTreemap::new(), + ); + // lhs: HE | rhs: HE + f("format: HE+HE", &he_buf, &he_buf, &he_tree, &he_tree); + } + + #[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}" + ); + }); + } + + #[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_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}"); + }); + } + + #[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() { + 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}"); + }); + } + + /// 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); + } + } + } } diff --git a/src/common/io/src/bitmap/reader.rs b/src/common/io/src/bitmap/reader.rs index fbd2934a707b6..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; @@ -26,6 +27,14 @@ 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; +const CONTAINER_MAX: usize = WORD_BITS * BITMAP_WORDS; + #[derive(Clone)] pub struct TreemapReader<'a> { buf: &'a [u8], @@ -47,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> { @@ -113,15 +199,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 +225,6 @@ impl BitmapReader<'_> { self.containers as usize } - #[allow(dead_code)] pub fn prefix(&self) -> u32 { self.prefix } @@ -164,6 +246,65 @@ 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 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 { @@ -178,16 +319,616 @@ 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?; +pub struct ContainerReader<'a> { + key: u16, + cardinality: usize, + is_array: bool, + data: &'a [u8], +} + +impl<'a> ContainerReader<'a> { + pub fn key(&self) -> u16 { + self.key + } + + pub fn contains(&self, low16: u16) -> bool { + if self.is_array { + self.contains_array(low16) + } else { + self.contains_bitmap(low16) + } + } + + pub fn min(&self) -> Option { + if self.is_array { + self.array_first() + } else { + self.bitmap_first() + } + } + + pub fn max(&self) -> Option { + if self.is_array { + self.array_last() + } else { + self.bitmap_last() + } + } + + 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) + } + } + + 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) + } + } + + 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 + } + + 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(self.data[start..start + WORD_BYTES].try_into().unwrap()); + word & (1 << bit_index) != 0 + } + + fn array_first(&self) -> Option { + if self.cardinality == 0 { + return None; + } + Some(u16::from_le_bytes(self.data[0..2].try_into().unwrap())) + } + + 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 + } + + 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(), + )) + } + + 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 + } + + 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 + } + + 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 + } + + 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; + } + } + 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; + } + 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 + } + + 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 + } + + 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 + } +} + +/// 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; +} + +struct ContainerVisitor<'a, H> { + lhs_buf: &'a [u8], + rhs_buf: &'a [u8], + handler: H, +} + +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, + } + } +} + +impl ContainerVisitor<'_, H> { + fn visit(mut self) -> io::Result { + use std::cmp::Ordering::*; + + let lhs_tree = TreemapReader::new(self.lhs_buf)?; + let rhs_tree = TreemapReader::new(self.rhs_buf)?; + + 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()?; + + // 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; + } + } + } + + // 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()) + } + + fn visit_one_side_containers( + &mut self, + bitmap: &BitmapReader<'_>, + is_lhs: bool, + other_remaining: bool, + ) -> io::Result> { for i in 0..bitmap.containers() { - sum += bitmap.description(i)?.cardinality(); + 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 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; + } + } + } + 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(())) + } + + 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 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 { + self.handler.handle_rhs_only(container, other_remaining) } } - Ok(sum) +} + +struct HasAnyHandler; + +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(()) + } + + fn handle_rhs_only(&mut self, _: &ContainerReader, lhs_remaining: bool) -> ControlFlow { + if !lhs_remaining { + return ControlFlow::Break(false); + } + ControlFlow::Continue(()) + } + + 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(()) + } + } + + fn handle_none(&mut self) -> bool { + false + } +} + +struct HasAllHandler; + +impl ContainerHandler for HasAllHandler { + type Output = bool; + + fn handle_lhs_only(&mut self, _: &ContainerReader, _: bool) -> ControlFlow { + ControlFlow::Continue(()) + } + + fn handle_rhs_only(&mut self, _: &ContainerReader, _: bool) -> ControlFlow { + ControlFlow::Break(false) + } + + 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(()) + } + } + + fn handle_none(&mut self) -> bool { + true + } +} + +pub fn bitmap_len(buf: &[u8]) -> io::Result { + TreemapReader::new(buf)?.len() +} + +pub(crate) fn bitmap_len_above(buf: &[u8], threshold: usize) -> io::Result { + TreemapReader::new(buf)?.len_above(threshold) +} + +pub(crate) fn bitmap_contains(buf: &[u8], value: u64) -> io::Result { + TreemapReader::new(buf)?.contains(value) +} + +pub(crate) fn bitmap_min(buf: &[u8]) -> io::Result> { + TreemapReader::new(buf)?.min() +} + +pub(crate) fn bitmap_max(buf: &[u8]) -> io::Result> { + TreemapReader::new(buf)?.max() +} + +pub(crate) fn bitmap_has_any(lhs: &[u8], rhs: &[u8]) -> io::Result { + ContainerVisitor::new(lhs, rhs, HasAnyHandler).visit() +} + +pub(crate) fn bitmap_has_all(lhs: &[u8], rhs: &[u8]) -> io::Result { + ContainerVisitor::new(lhs, rhs, HasAllHandler).visit() } pub fn intersection_with_serialized(tree: &mut RoaringTreemap, buf: &[u8]) -> io::Result<()> { @@ -288,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 a3b3a8c571e0b..b0f13040b83c4 100644 --- a/src/common/io/src/lib.rs +++ b/src/common/io/src/lib.rs @@ -53,6 +53,11 @@ 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::bitmap_has_all; +pub use bitmap::bitmap_has_any; +pub use bitmap::bitmap_max; +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/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; diff --git a/src/query/functions/src/scalars/bitmap.rs b/src/query/functions/src/scalars/bitmap.rs index 1190e52a0c732..d97b671a06a88 100644 --- a/src/query/functions/src/scalars/bitmap.rs +++ b/src/query/functions/src/scalars/bitmap.rs @@ -37,7 +37,12 @@ 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_has_all; +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; use databend_common_io::deserialize_bitmap; use databend_common_io::parse_bitmap; use itertools::join; @@ -269,9 +274,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()); @@ -381,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)); + } }, ), ); @@ -413,23 +410,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); + } }, ), ); @@ -448,14 +437,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 @@ -480,14 +467,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