diff --git a/core/Cargo.toml b/core/Cargo.toml index 0a4d068..67176cf 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -30,3 +30,4 @@ zstd = "0.13" arrow-schema = "58" arrow-array = "58" arrow-buffer = "58" +twox-hash = { version = "1.6", default-features = false } diff --git a/core/src/bloom.rs b/core/src/bloom.rs new file mode 100644 index 0000000..460b51c --- /dev/null +++ b/core/src/bloom.rs @@ -0,0 +1,584 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; +use std::hash::Hasher; +use std::io; + +use arrow_array::*; +use arrow_schema::DataType; + +use crate::values::Value; +use crate::varint; +use twox_hash::XxHash64; + +const SBBF_SALT: [u32; 8] = [ + 0x47b6137b, 0x44974d91, 0x8824ad5b, 0xa2b7289d, 0x705495c7, 0x2df1424b, 0x9efc4947, 0x5c6bfb31, +]; + +pub const BLOOM_HEADER_VERSION: u8 = 1; +pub const BLOOM_ALGORITHM_SBBF: u8 = 1; +pub const BLOOM_HASH_XXHASH64: u8 = 1; +pub const BLOOM_COMPRESSION_NONE: u8 = 0; + +pub const BLOOM_BLOCK_BYTES: usize = 32; +pub const BLOOM_BLOCK_WORDS: usize = 8; + +pub const BLOOM_MIN_BLOCKS: usize = 1; +// Hard cap to keep one filter under ~512 MiB even on hostile configs. +pub const BLOOM_MAX_BLOCKS: usize = 1 << 24; + +#[derive(Debug, Clone)] +pub struct BloomFilterConfig { + pub column_name: String, + pub ndv: u64, + pub fpp: f64, +} + +#[derive(Debug, Clone)] +pub struct SplitBlockBloomFilter { + blocks: Vec<[u32; BLOOM_BLOCK_WORDS]>, +} + +impl SplitBlockBloomFilter { + pub fn with_num_blocks(num_blocks: usize) -> Self { + let n = num_blocks + .clamp(BLOOM_MIN_BLOCKS, BLOOM_MAX_BLOCKS) + .next_power_of_two(); + Self { + blocks: vec![[0u32; BLOOM_BLOCK_WORDS]; n], + } + } + + pub fn with_capacity(ndv: u64, fpp: f64) -> Self { + Self::with_num_blocks(num_blocks_for(ndv, fpp)) + } + + pub fn num_blocks(&self) -> usize { + self.blocks.len() + } + + pub fn bitset_num_bytes(&self) -> usize { + self.blocks.len() * BLOOM_BLOCK_BYTES + } + + #[inline] + pub fn insert_hash(&mut self, h: u64) { + let idx = block_index(h, self.blocks.len()); + let block = &mut self.blocks[idx]; + let mask = block_mask(h as u32); + for i in 0..BLOOM_BLOCK_WORDS { + block[i] |= mask[i]; + } + } + + #[inline] + pub fn contains_hash(&self, h: u64) -> bool { + let idx = block_index(h, self.blocks.len()); + let block = &self.blocks[idx]; + let mask = block_mask(h as u32); + for i in 0..BLOOM_BLOCK_WORDS { + if (block[i] & mask[i]) != mask[i] { + return false; + } + } + true + } + + pub fn write_to(&self, buf: &mut Vec) { + buf.push(BLOOM_HEADER_VERSION); + buf.push(BLOOM_ALGORITHM_SBBF); + buf.push(BLOOM_HASH_XXHASH64); + buf.push(BLOOM_COMPRESSION_NONE); + let num_bytes = self.bitset_num_bytes() as u32; + varint::encode(buf, num_bytes); + for block in &self.blocks { + for word in block { + buf.extend_from_slice(&word.to_le_bytes()); + } + } + } + + pub fn read_from(data: &[u8]) -> io::Result { + if data.len() < 4 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "bloom: header truncated", + )); + } + let header_version = data[0]; + let algorithm = data[1]; + let hash = data[2]; + let compression = data[3]; + if header_version != BLOOM_HEADER_VERSION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("bloom: unsupported header version {}", header_version), + )); + } + if algorithm != BLOOM_ALGORITHM_SBBF { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("bloom: unsupported algorithm {}", algorithm), + )); + } + if hash != BLOOM_HASH_XXHASH64 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("bloom: unsupported hash {}", hash), + )); + } + if compression != BLOOM_COMPRESSION_NONE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("bloom: unsupported compression {}", compression), + )); + } + let mut pos = 4; + let num_bytes = varint::decode(data, &mut pos)? as usize; + if !num_bytes.is_multiple_of(BLOOM_BLOCK_BYTES) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "bloom: numBytes {} not a multiple of block size {}", + num_bytes, BLOOM_BLOCK_BYTES + ), + )); + } + let num_blocks = num_bytes / BLOOM_BLOCK_BYTES; + if num_blocks == 0 || !num_blocks.is_power_of_two() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "bloom: numBlocks {} not a positive power of two", + num_blocks + ), + )); + } + if num_blocks > BLOOM_MAX_BLOCKS { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "bloom: numBlocks {} exceeds cap {}", + num_blocks, BLOOM_MAX_BLOCKS + ), + )); + } + if pos + num_bytes > data.len() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "bloom: bitset truncated", + )); + } + let mut blocks = Vec::with_capacity(num_blocks); + for b in 0..num_blocks { + let mut block = [0u32; BLOOM_BLOCK_WORDS]; + for (i, word) in block.iter_mut().enumerate() { + let off = pos + b * BLOOM_BLOCK_BYTES + i * 4; + *word = u32::from_le_bytes(data[off..off + 4].try_into().unwrap()); + } + blocks.push(block); + } + Ok(Self { blocks }) + } +} + +#[inline] +fn block_index(h: u64, num_blocks: usize) -> usize { + let top = h >> 32; + let z = num_blocks as u64; + ((top.wrapping_mul(z)) >> 32) as usize +} + +#[inline] +fn block_mask(x: u32) -> [u32; BLOOM_BLOCK_WORDS] { + let mut mask = [0u32; BLOOM_BLOCK_WORDS]; + for i in 0..BLOOM_BLOCK_WORDS { + let y = x.wrapping_mul(SBBF_SALT[i]); + let bit = y >> 27; + mask[i] = 1u32 << bit; + } + mask +} + +fn bits_per_insert_for_fpp(fpp: f64) -> f64 { + if fpp >= 0.10 { + 6.0 + } else if fpp >= 0.01 { + 10.5 + } else if fpp >= 0.001 { + 16.9 + } else if fpp >= 0.0001 { + 26.4 + } else { + 41.0 + } +} + +pub fn num_blocks_for(ndv: u64, fpp: f64) -> usize { + let bits_per_insert = bits_per_insert_for_fpp(fpp); + let ndv_for_calc = ndv.max(1); + let total_bits = ((ndv_for_calc as f64) * bits_per_insert).ceil() as u64; + let blocks = total_bits.div_ceil(256) as usize; + blocks.max(BLOOM_MIN_BLOCKS).next_power_of_two() +} + +pub fn hash_value(value: &Value) -> u64 { + let mut h = XxHash64::with_seed(0); + write_canonical_bytes(value, &mut h); + h.finish() +} + +fn write_canonical_bytes(value: &Value, h: &mut XxHash64) { + match value { + Value::Boolean(b) => h.write(&[*b as u8]), + Value::TinyInt(x) => h.write(&x.to_le_bytes()), + Value::SmallInt(x) => h.write(&x.to_le_bytes()), + Value::Integer(x) => h.write(&x.to_le_bytes()), + Value::BigInt(x) => h.write(&x.to_le_bytes()), + Value::Float(x) => h.write(&x.to_le_bytes()), + Value::Double(x) => h.write(&x.to_le_bytes()), + Value::Date(x) => h.write(&x.to_le_bytes()), + Value::String(bytes) => h.write(bytes), + _ => {} + } +} + +pub fn supports_bloom(dt: &DataType) -> bool { + matches!( + dt, + DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Float32 + | DataType::Float64 + | DataType::Date32 + | DataType::Utf8 + ) +} + +struct ColEntry { + column_index: usize, + batch_col_index: usize, + data_type: DataType, + ndv: u64, + fpp: f64, + filter: SplitBlockBloomFilter, +} + +pub type ResolvedBloomColumn = (usize, usize, DataType, u64, f64); + +pub struct BloomFilterCollector { + entries: Vec, +} + +impl BloomFilterCollector { + pub fn new(columns: &[ResolvedBloomColumn]) -> Self { + let entries = columns + .iter() + .map(|(idx, batch_idx, dt, ndv, fpp)| ColEntry { + column_index: *idx, + batch_col_index: *batch_idx, + data_type: dt.clone(), + ndv: *ndv, + fpp: *fpp, + filter: SplitBlockBloomFilter::with_capacity(*ndv, *fpp), + }) + .collect(); + Self { entries } + } + + pub fn update_batch(&mut self, batch: &RecordBatch) { + for entry in &mut self.entries { + let array = batch.column(entry.batch_col_index).as_ref(); + insert_array_hashes(array, &entry.data_type, &mut entry.filter); + } + } + + pub fn finish(&mut self) -> Vec<(usize, SplitBlockBloomFilter)> { + let mut out = Vec::with_capacity(self.entries.len()); + for entry in &mut self.entries { + let fresh = SplitBlockBloomFilter::with_capacity(entry.ndv, entry.fpp); + let prev = std::mem::replace(&mut entry.filter, fresh); + out.push((entry.column_index, prev)); + } + out + } +} + +fn insert_array_hashes(array: &dyn Array, dt: &DataType, filter: &mut SplitBlockBloomFilter) { + let len = array.len(); + for row in 0..len { + if array.is_null(row) { + continue; + } + if let Some(v) = extract_value(array, row, dt) { + filter.insert_hash(hash_value(&v)); + } + } +} + +fn extract_value(array: &dyn Array, row: usize, dt: &DataType) -> Option { + match dt { + DataType::Boolean => array + .as_any() + .downcast_ref::() + .map(|a| Value::Boolean(a.value(row))), + DataType::Int8 => array + .as_any() + .downcast_ref::() + .map(|a| Value::TinyInt(a.value(row))), + DataType::Int16 => array + .as_any() + .downcast_ref::() + .map(|a| Value::SmallInt(a.value(row))), + DataType::Int32 => array + .as_any() + .downcast_ref::() + .map(|a| Value::Integer(a.value(row))), + DataType::Int64 => array + .as_any() + .downcast_ref::() + .map(|a| Value::BigInt(a.value(row))), + DataType::Float32 => array + .as_any() + .downcast_ref::() + .map(|a| Value::Float(a.value(row))), + DataType::Float64 => array + .as_any() + .downcast_ref::() + .map(|a| Value::Double(a.value(row))), + DataType::Date32 => array + .as_any() + .downcast_ref::() + .map(|a| Value::Date(a.value(row))), + DataType::Utf8 => array + .as_any() + .downcast_ref::() + .map(|a| Value::String(a.value(row).as_bytes().to_vec())), + _ => None, + } +} + +#[derive(Debug, Clone)] +pub struct BloomEntryMeta { + pub column_index: usize, + pub offset: u64, + pub total_bytes: usize, +} + +pub fn serialize_index_tail(blooms: &[BloomEntryMeta]) -> Vec { + let mut buf = Vec::with_capacity(1 + blooms.len() * 12); + varint::encode(&mut buf, blooms.len() as u32); + for b in blooms { + varint::encode(&mut buf, b.column_index as u32); + buf.extend_from_slice(&b.offset.to_be_bytes()); + varint::encode(&mut buf, b.total_bytes as u32); + } + buf +} + +pub fn deserialize_index_tail(data: &[u8], pos: &mut usize) -> io::Result> { + let num = varint::decode(data, pos)? as usize; + let mut out = Vec::with_capacity(num); + for _ in 0..num { + let column_index = varint::decode(data, pos)? as usize; + if *pos + 8 > data.len() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "bloom index tail: offset truncated", + )); + } + let offset = u64::from_be_bytes(data[*pos..*pos + 8].try_into().unwrap()); + *pos += 8; + let total_bytes = varint::decode(data, pos)? as usize; + out.push(BloomEntryMeta { + column_index, + offset, + total_bytes, + }); + } + Ok(out) +} + +// Convenience helper: validate user-supplied configs against the schema, +// returning the per-collector tuples used by BloomFilterCollector::new. +// Bubbles up clear errors for missing names or unsupported types. +pub fn resolve_configs( + configs: &[BloomFilterConfig], + schema_columns: &[crate::schema::ColumnMeta], + batch_col_map: &[usize], +) -> io::Result> { + let mut name_to_idx: HashMap<&str, usize> = HashMap::with_capacity(schema_columns.len()); + for (i, c) in schema_columns.iter().enumerate() { + name_to_idx.insert(c.name.as_str(), i); + } + let mut seen: HashMap = HashMap::new(); + let mut out = Vec::with_capacity(configs.len()); + for c in configs { + let idx = *name_to_idx.get(c.column_name.as_str()).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "bloom_filter_columns: column '{}' not found in schema", + c.column_name + ), + ) + })?; + if seen.insert(idx, ()).is_some() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "bloom_filter_columns: duplicate entry for column '{}'", + c.column_name + ), + )); + } + let dt = &schema_columns[idx].data_type; + if !supports_bloom(dt) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "bloom_filter_columns: column '{}' has unsupported type {:?} for bloom filter", + c.column_name, dt + ), + )); + } + if !(0.0 < c.fpp && c.fpp < 1.0) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "bloom_filter_columns: column '{}' fpp must be in (0, 1), got {}", + c.column_name, c.fpp + ), + )); + } + out.push((idx, batch_col_map[idx], dt.clone(), c.ndv, c.fpp)); + } + out.sort_by_key(|(idx, _, _, _, _)| *idx); + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn block_mask_sets_eight_bits() { + let m = block_mask(0xdeadbeef); + let total_bits: u32 = m.iter().map(|w| w.count_ones()).sum(); + assert_eq!(total_bits, 8); + } + + #[test] + fn insert_then_contains() { + let mut f = SplitBlockBloomFilter::with_num_blocks(64); + for i in 0..1000u64 { + f.insert_hash(hash_value(&Value::BigInt(i as i64))); + } + for i in 0..1000u64 { + assert!(f.contains_hash(hash_value(&Value::BigInt(i as i64)))); + } + } + + #[test] + fn fpp_within_expectation() { + let n = 10_000u64; + let target_fpp = 0.01; + let mut f = SplitBlockBloomFilter::with_capacity(n, target_fpp); + for i in 0..n { + f.insert_hash(hash_value(&Value::BigInt(i as i64))); + } + let probe = 200_000u64; + let mut fp = 0u64; + for i in n..(n + probe) { + if f.contains_hash(hash_value(&Value::BigInt(i as i64))) { + fp += 1; + } + } + let observed = fp as f64 / probe as f64; + assert!( + observed < target_fpp * 3.0, + "observed fpp {} exceeded 3x target {}", + observed, + target_fpp + ); + } + + #[test] + fn write_read_roundtrip() { + let mut f = SplitBlockBloomFilter::with_num_blocks(8); + for i in 0..500u64 { + f.insert_hash(hash_value(&Value::BigInt(i as i64))); + } + let mut buf = Vec::new(); + f.write_to(&mut buf); + let g = SplitBlockBloomFilter::read_from(&buf).unwrap(); + assert_eq!(g.num_blocks(), f.num_blocks()); + for i in 0..500u64 { + assert!(g.contains_hash(hash_value(&Value::BigInt(i as i64)))); + } + } + + #[test] + fn read_rejects_bad_discriminants() { + let mut buf = vec![ + BLOOM_HEADER_VERSION, + 99, + BLOOM_HASH_XXHASH64, + BLOOM_COMPRESSION_NONE, + ]; + varint::encode(&mut buf, 32); + buf.extend(vec![0u8; 32]); + let err = SplitBlockBloomFilter::read_from(&buf).unwrap_err(); + assert!(format!("{}", err).contains("unsupported algorithm")); + } + + #[test] + fn num_blocks_for_known_sizes() { + assert!(num_blocks_for(0, 0.01) >= 1); + assert!(num_blocks_for(1_000_000, 0.01) > num_blocks_for(1_000, 0.01)); + assert!(num_blocks_for(1_000, 0.0001) > num_blocks_for(1_000, 0.01)); + } + + #[test] + fn index_tail_roundtrip() { + let blooms = vec![ + BloomEntryMeta { + column_index: 1, + offset: 1024, + total_bytes: 200, + }, + BloomEntryMeta { + column_index: 4, + offset: 1224, + total_bytes: 500, + }, + ]; + let buf = serialize_index_tail(&blooms); + let mut pos = 0; + let got = deserialize_index_tail(&buf, &mut pos).unwrap(); + assert_eq!(pos, buf.len()); + assert_eq!(got.len(), 2); + assert_eq!(got[0].column_index, 1); + assert_eq!(got[0].offset, 1024); + assert_eq!(got[0].total_bytes, 200); + assert_eq!(got[1].column_index, 4); + } +} diff --git a/core/src/lib.rs b/core/src/lib.rs index 3f28500..d06c1a6 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +pub mod bloom; pub mod bpe; pub mod bucket_reader; pub mod bucket_writer; diff --git a/core/src/reader.rs b/core/src/reader.rs index e111b9d..19012cf 100644 --- a/core/src/reader.rs +++ b/core/src/reader.rs @@ -21,6 +21,7 @@ use std::sync::Arc; use arrow_array::{ArrayRef, RecordBatch, RecordBatchOptions}; use arrow_schema::{DataType, Field, Schema}; +use crate::bloom::{self, BloomEntryMeta, SplitBlockBloomFilter}; use crate::bucket_reader::{read_typed_value, read_variable_value, BucketReader, ColumnPageReader}; use crate::schema::MosaicSchema; use crate::spec::*; @@ -184,6 +185,7 @@ pub struct RowGroupMeta { pub bucket_offsets: Vec, pub bucket_layouts: Vec, pub stats: Vec, + pub blooms: Vec, } pub trait ReaderAccess { @@ -220,6 +222,12 @@ pub trait ReaderAccess { fn project(&mut self, column_names: &[&str]) -> io::Result<()>; fn row_group_stats(&self, rg_index: usize) -> io::Result<&[ColumnStats]>; fn row_group_num_rows(&self, rg_index: usize) -> io::Result; + fn row_group_bloom_meta(&self, rg_index: usize) -> io::Result<&[BloomEntryMeta]>; + fn bloom_filter( + &self, + rg_index: usize, + column_index: usize, + ) -> io::Result>; } pub struct MosaicReader { @@ -264,7 +272,7 @@ impl MosaicReader { } let version = footer[25]; - if version != VERSION { + if !(VERSION_MIN_SUPPORTED..=VERSION).contains(&version) { return Err(io::Error::new( io::ErrorKind::InvalidData, format!("unsupported version: {}", version), @@ -398,11 +406,18 @@ impl MosaicReader { let rg_stats = stats::deserialize_stats(index_data, &mut pos, &schema.columns, num_rows)?; + let rg_blooms = if version >= 2 { + bloom::deserialize_index_tail(index_data, &mut pos)? + } else { + Vec::new() + }; + row_group_metas.push(RowGroupMeta { num_rows, bucket_offsets, bucket_layouts, stats: rg_stats, + blooms: rg_blooms, }); } @@ -537,6 +552,49 @@ impl ReaderAccess for MosaicReader { Ok(self.row_group_metas[rg_index].num_rows) } + fn row_group_bloom_meta(&self, rg_index: usize) -> io::Result<&[BloomEntryMeta]> { + if rg_index >= self.row_group_metas.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "row group index {} out of range (num_row_groups={})", + rg_index, + self.row_group_metas.len() + ), + )); + } + Ok(&self.row_group_metas[rg_index].blooms) + } + + fn bloom_filter( + &self, + rg_index: usize, + column_index: usize, + ) -> io::Result> { + if rg_index >= self.row_group_metas.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "row group index {} out of range (num_row_groups={})", + rg_index, + self.row_group_metas.len() + ), + )); + } + let meta = self.row_group_metas[rg_index] + .blooms + .iter() + .find(|b| b.column_index == column_index); + let meta = match meta { + Some(m) => m, + None => return Ok(None), + }; + let mut buf = vec![0u8; meta.total_bytes]; + self.input.read_at(meta.offset, &mut buf)?; + let filter = SplitBlockBloomFilter::read_from(&buf)?; + Ok(Some(filter)) + } + fn row_group_reader(&self, rg_index: usize) -> io::Result { match &self.projected_columns { Some(cols) => self.row_group_reader_projected(rg_index, cols), diff --git a/core/src/spec.rs b/core/src/spec.rs index 5957b91..8d338cd 100644 --- a/core/src/spec.rs +++ b/core/src/spec.rs @@ -16,9 +16,12 @@ // under the License. pub const MAGIC: [u8; 4] = [b'M', b'O', b'S', b'A']; -pub const VERSION: u8 = 1; +pub const VERSION: u8 = 2; +pub const VERSION_MIN_SUPPORTED: u8 = 1; pub const FOOTER_SIZE: usize = 32; +pub const BLOOM_DEFAULT_FPP: f64 = 0.01; + pub const COMPRESSION_NONE: u8 = 0; pub const COMPRESSION_ZSTD: u8 = 1; diff --git a/core/src/writer.rs b/core/src/writer.rs index cea1582..46ddcf9 100644 --- a/core/src/writer.rs +++ b/core/src/writer.rs @@ -20,6 +20,7 @@ use std::io; use arrow_array::*; use arrow_schema::Schema; +use crate::bloom::{self, BloomEntryMeta, BloomFilterCollector, BloomFilterConfig}; use crate::bucket_writer::{BucketWriter, PagedBucketOutput}; use crate::schema::MosaicSchema; use crate::spec::*; @@ -50,6 +51,7 @@ pub struct WriterOptions { pub max_dict_entries: usize, pub stats_columns: Vec, pub page_size_threshold: usize, + pub bloom_filter_columns: Vec, } impl Default for WriterOptions { @@ -63,6 +65,7 @@ impl Default for WriterOptions { max_dict_entries: DEFAULT_DICT_MAX_ENTRIES, stats_columns: Vec::new(), page_size_threshold: DEFAULT_PAGE_SIZE_THRESHOLD, + bloom_filter_columns: Vec::new(), } } } @@ -72,6 +75,7 @@ struct RowGroupMeta { bucket_offsets: Vec, bucket_layouts: Vec, stats: Vec, + blooms: Vec, } pub struct MosaicWriter { @@ -93,6 +97,7 @@ pub struct MosaicWriter { total_uncompressed: u64, total_compressed: u64, stats_collector: Option, + bloom_collector: Option, closed: bool, } @@ -175,6 +180,17 @@ impl MosaicWriter { Some(StatsCollector::new(&cols)) }; + let bloom_collector = if options.bloom_filter_columns.is_empty() { + None + } else { + let resolved = bloom::resolve_configs( + &options.bloom_filter_columns, + &schema.columns, + &batch_col_map, + )?; + Some(BloomFilterCollector::new(&resolved)) + }; + let active_buckets: Vec = bucket_writers .iter() .enumerate() @@ -205,6 +221,7 @@ impl MosaicWriter { total_uncompressed: 0, total_compressed: 0, stats_collector, + bloom_collector, closed: false, }) } @@ -285,6 +302,9 @@ impl MosaicWriter { if let Some(ref mut collector) = self.stats_collector { collector.update_batch(batch); } + if let Some(ref mut collector) = self.bloom_collector { + collector.update_batch(batch); + } self.current_row_group_rows += batch.num_rows(); self.current_buffered_size += size; @@ -383,11 +403,33 @@ impl MosaicWriter { None => Vec::new(), }; + let row_blooms = match &mut self.bloom_collector { + Some(collector) => { + let filters = collector.finish(); + let mut metas = Vec::with_capacity(filters.len()); + let mut filter_bytes = Vec::new(); + for (column_index, filter) in filters { + filter_bytes.clear(); + filter.write_to(&mut filter_bytes); + let offset = self.out.pos(); + self.out.write(&filter_bytes)?; + metas.push(BloomEntryMeta { + column_index, + offset, + total_bytes: filter_bytes.len(), + }); + } + metas + } + None => Vec::new(), + }; + self.row_group_metas.push(RowGroupMeta { num_rows: self.current_row_group_rows, bucket_offsets, bucket_layouts, stats: row_stats, + blooms: row_blooms, }); self.current_row_group_rows = 0; @@ -543,6 +585,8 @@ impl MosaicWriter { } let stats_bytes = stats::serialize_stats(&meta.stats, &self.schema.columns); index_buf.extend_from_slice(&stats_bytes); + let bloom_bytes = bloom::serialize_index_tail(&meta.blooms); + index_buf.extend_from_slice(&bloom_bytes); } self.out.write(&index_buf)?; diff --git a/core/tests/bloom_test.rs b/core/tests/bloom_test.rs new file mode 100644 index 0000000..25664b0 --- /dev/null +++ b/core/tests/bloom_test.rs @@ -0,0 +1,249 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::io; +use std::sync::Arc; + +use arrow_array::*; +use arrow_schema::{DataType, Field, Schema}; + +use paimon_mosaic_core::bloom::{hash_value, BloomFilterConfig}; +use paimon_mosaic_core::reader::{InputFile, MosaicReader, ReaderAccess}; +use paimon_mosaic_core::values::Value; +use paimon_mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions}; + +struct MemOutputFile { + buf: Vec, +} + +impl OutputFile for MemOutputFile { + fn write(&mut self, data: &[u8]) -> io::Result<()> { + self.buf.extend_from_slice(data); + Ok(()) + } + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + fn pos(&self) -> u64 { + self.buf.len() as u64 + } +} + +struct ByteArrayInputFile { + data: Vec, +} + +impl InputFile for ByteArrayInputFile { + fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> { + let s = offset as usize; + let e = s + buf.len(); + buf.copy_from_slice(&self.data[s..e]); + Ok(()) + } +} + +fn write_file( + schema: Arc, + batches: Vec, + bloom_columns: Vec, + num_buckets: usize, +) -> Vec { + let out = MemOutputFile { buf: Vec::new() }; + let mut writer = MosaicWriter::new( + out, + schema.as_ref(), + WriterOptions { + num_buckets, + bloom_filter_columns: bloom_columns, + ..Default::default() + }, + ) + .unwrap(); + for b in batches { + writer.write_batch(&b).unwrap(); + } + writer.close().unwrap(); + writer.output().buf.clone() +} + +#[test] +fn bloom_present_values_match() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("payload", DataType::Int32, true), + ])); + let n = 5000i64; + let id_array: Int64Array = (0..n).collect(); + let payload_array: Int32Array = (0..n).map(|i| Some(i as i32)).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(id_array), Arc::new(payload_array)], + ) + .unwrap(); + let bytes = write_file( + schema.clone(), + vec![batch], + vec![BloomFilterConfig { + column_name: "id".to_string(), + ndv: n as u64, + fpp: 0.01, + }], + 4, + ); + let len = bytes.len() as u64; + let reader = MosaicReader::new(ByteArrayInputFile { data: bytes }, len).expect("reader open"); + + assert_eq!(reader.num_row_groups(), 1); + let metas = reader.row_group_bloom_meta(0).unwrap(); + assert_eq!(metas.len(), 1, "exactly one bloom for column id"); + assert_eq!(metas[0].column_index, 0); + assert!(metas[0].total_bytes > 0); + + let filter = reader + .bloom_filter(0, 0) + .expect("bloom fetch ok") + .expect("filter present"); + + for i in 0..n { + let h = hash_value(&Value::BigInt(i)); + assert!(filter.contains_hash(h), "present value {} missing", i); + } + + let probe = 10_000i64; + let mut fp = 0u64; + for i in n..(n + probe) { + let h = hash_value(&Value::BigInt(i)); + if filter.contains_hash(h) { + fp += 1; + } + } + let observed = fp as f64 / probe as f64; + assert!( + observed < 0.05, + "observed fpp {} too high (target was 0.01)", + observed + ); +} + +#[test] +fn no_bloom_when_not_configured() { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let id_array: Int64Array = (0..100i64).collect(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(id_array)]).unwrap(); + let bytes = write_file(schema.clone(), vec![batch], Vec::new(), 1); + let len = bytes.len() as u64; + let reader = MosaicReader::new(ByteArrayInputFile { data: bytes }, len).unwrap(); + + let metas = reader.row_group_bloom_meta(0).unwrap(); + assert!(metas.is_empty()); + assert!(reader.bloom_filter(0, 0).unwrap().is_none()); +} + +#[test] +fn bloom_string_column_skips_absent_value() { + let schema = Arc::new(Schema::new(vec![Field::new("name", DataType::Utf8, false)])); + let names: Vec<&str> = vec!["alice", "bob", "carol", "dave", "eve"]; + let arr = StringArray::from(names.clone()); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap(); + let bytes = write_file( + schema.clone(), + vec![batch], + vec![BloomFilterConfig { + column_name: "name".to_string(), + ndv: 1024, + fpp: 0.001, + }], + 1, + ); + let len = bytes.len() as u64; + let reader = MosaicReader::new(ByteArrayInputFile { data: bytes }, len).unwrap(); + let filter = reader.bloom_filter(0, 0).unwrap().unwrap(); + + for n in &names { + let h = hash_value(&Value::String(n.as_bytes().to_vec())); + assert!(filter.contains_hash(h), "present name {} missing", n); + } + let absent = "zachary"; + let h = hash_value(&Value::String(absent.as_bytes().to_vec())); + assert!( + !filter.contains_hash(h), + "filter false-positive on small set" + ); +} + +#[test] +fn bloom_per_row_group_independent() { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + + let batches: Vec = vec![ + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from((0..1000i64).collect::>()))], + ) + .unwrap(), + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from( + (1_000_000..1_001_000i64).collect::>(), + ))], + ) + .unwrap(), + ]; + + let mut out = MemOutputFile { buf: Vec::new() }; + let mut writer = MosaicWriter::new( + out, + schema.as_ref(), + WriterOptions { + num_buckets: 1, + row_group_max_size: 16 * 1024, + bloom_filter_columns: vec![BloomFilterConfig { + column_name: "id".to_string(), + ndv: 4000, + fpp: 0.01, + }], + ..Default::default() + }, + ) + .unwrap(); + for b in &batches { + writer.write_batch(b).unwrap(); + // Force separate row groups by manually flushing buffer through tiny RG size. + } + writer.close().unwrap(); + let bytes = writer.output().buf.clone(); + drop(writer); + out = MemOutputFile { buf: Vec::new() }; + drop(out); + + let len = bytes.len() as u64; + let reader = MosaicReader::new(ByteArrayInputFile { data: bytes }, len).unwrap(); + assert!(reader.num_row_groups() >= 1); + + let f0 = reader.bloom_filter(0, 0).unwrap().unwrap(); + for i in 0..1000i64 { + assert!(f0.contains_hash(hash_value(&Value::BigInt(i)))); + } + if reader.num_row_groups() >= 2 { + let f1 = reader.bloom_filter(1, 0).unwrap().unwrap(); + for i in 1_000_000..1_001_000i64 { + assert!(f1.contains_hash(hash_value(&Value::BigInt(i)))); + } + let h = hash_value(&Value::BigInt(0)); + assert!(!f1.contains_hash(h)); + } +} diff --git a/cpp/test_mosaic.cpp b/cpp/test_mosaic.cpp index 1d08012..6e233a8 100644 --- a/cpp/test_mosaic.cpp +++ b/cpp/test_mosaic.cpp @@ -824,6 +824,92 @@ static void test_stats_empty_string_min() { printf(" PASS test_stats_empty_string_min\n"); } +static void test_bloom_int64_hits_and_misses() { + auto schema = arrow::schema({arrow::field("id", arrow::int64(), false)}); + arrow::Int64Builder id_b; + for (int64_t i = 0; i < 2000; i++) { + assert(id_b.Append(i).ok()); + } + auto batch = arrow::RecordBatch::Make(schema, 2000, {id_b.Finish().ValueUnsafe()}); + + mosaic::WriterOptions opts; + opts.num_buckets = 2; + opts.bloom_filter_columns.push_back({"id", 2000, 0.01}); + + MemBuffer write_buf; + struct ArrowSchema c_schema; + auto st = arrow::ExportSchema(*schema, &c_schema); + assert(st.ok()); + mosaic::Writer writer(make_output(write_buf), &c_schema, opts); + + struct ArrowArray c_array; + struct ArrowSchema c_batch_schema; + st = arrow::ExportRecordBatch(*batch, &c_array, &c_batch_schema); + assert(st.ok()); + writer.write(&c_array, &c_batch_schema); + writer.close(); + + auto reader = mosaic::make_reader(make_input(write_buf), write_buf.data.size()); + for (int64_t i = 0; i < 2000; i++) { + ASSERT_TRUE(reader.bloom_might_contain( + 0, 0, 4, reinterpret_cast(&i), sizeof(int64_t))); + } + uint64_t fp = 0; + const uint64_t probe = 5000; + for (int64_t i = 1000000; i < 1000000 + (int64_t)probe; i++) { + if (reader.bloom_might_contain(0, 0, 4, + reinterpret_cast(&i), + sizeof(int64_t))) { + fp++; + } + } + double observed = static_cast(fp) / probe; + ASSERT_TRUE(observed < 0.05); + printf(" PASS test_bloom_int64_hits_and_misses (observed fpp %f)\n", observed); +} + +static void test_bloom_string_rejects_absent() { + auto schema = arrow::schema({arrow::field("name", arrow::utf8(), false)}); + arrow::StringBuilder name_b; + const std::vector names = {"alice", "bob", "carol", "dave", "eve"}; + for (auto const& n : names) { + assert(name_b.Append(n).ok()); + } + auto batch = arrow::RecordBatch::Make(schema, names.size(), + {name_b.Finish().ValueUnsafe()}); + + mosaic::WriterOptions opts; + opts.num_buckets = 1; + opts.bloom_filter_columns.push_back({"name", 1024, 0.001}); + + MemBuffer write_buf; + struct ArrowSchema c_schema; + auto st = arrow::ExportSchema(*schema, &c_schema); + assert(st.ok()); + mosaic::Writer writer(make_output(write_buf), &c_schema, opts); + + struct ArrowArray c_array; + struct ArrowSchema c_batch_schema; + st = arrow::ExportRecordBatch(*batch, &c_array, &c_batch_schema); + assert(st.ok()); + writer.write(&c_array, &c_batch_schema); + writer.close(); + + auto reader = mosaic::make_reader(make_input(write_buf), write_buf.data.size()); + for (auto const& n : names) { + ASSERT_TRUE(reader.bloom_might_contain( + 0, 0, 10, + reinterpret_cast(n.data()), + static_cast(n.size()))); + } + std::string absent = "zachary"; + ASSERT_TRUE(!reader.bloom_might_contain( + 0, 0, 10, + reinterpret_cast(absent.data()), + static_cast(absent.size()))); + printf(" PASS test_bloom_string_rejects_absent\n"); +} + int main() { printf("Running Mosaic C++ tests...\n"); test_basic_roundtrip(); @@ -841,6 +927,8 @@ int main() { test_writer_stats_all_null(); test_writer_stats_matches_reader(); test_stats_empty_string_min(); - printf("All %d tests passed.\n", 15); + test_bloom_int64_hits_and_misses(); + test_bloom_string_rejects_absent(); + printf("All %d tests passed.\n", 17); return 0; } diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 789d07d..60f37e0 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -28,8 +28,10 @@ use arrow_array::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; use arrow_array::{RecordBatch, StructArray}; use arrow_schema::{Field, Schema}; +use mosaic_core::bloom::{self, BloomFilterConfig}; use mosaic_core::reader::{InputFile, MosaicReader, ReaderAccess}; use mosaic_core::spec::*; +use mosaic_core::values::Value; use mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions}; type MinMaxPair = (Option>, Option>); @@ -109,6 +111,13 @@ impl OutputFile for FfiOutputFile { // ======================== Writer Options ======================== +#[repr(C)] +pub struct MosaicBloomConfig { + pub column_name: *const c_char, + pub ndv: u64, + pub fpp: f64, +} + #[repr(C)] pub struct MosaicWriterOptions { pub compression: u8, @@ -120,6 +129,8 @@ pub struct MosaicWriterOptions { pub stats_columns: *const *const c_char, pub num_stats_columns: u32, pub page_size_threshold: u32, + pub bloom_filter_columns: *const MosaicBloomConfig, + pub num_bloom_filter_columns: u32, } /// Returns default writer options. @@ -135,6 +146,8 @@ pub extern "C" fn mosaic_writer_options_default() -> MosaicWriterOptions { stats_columns: ptr::null(), num_stats_columns: 0, page_size_threshold: DEFAULT_PAGE_SIZE_THRESHOLD as u32, + bloom_filter_columns: ptr::null(), + num_bloom_filter_columns: 0, } } @@ -201,6 +214,36 @@ pub unsafe extern "C" fn mosaic_writer_open( } else { options.num_buckets as usize }; + let bloom_cols = + if options.bloom_filter_columns.is_null() || options.num_bloom_filter_columns == 0 { + Vec::new() + } else { + let entries = std::slice::from_raw_parts( + options.bloom_filter_columns, + options.num_bloom_filter_columns as usize, + ); + let mut out = Vec::with_capacity(entries.len()); + for entry in entries { + if entry.column_name.is_null() { + set_error("bloom_filter_columns contains null column_name".into()); + return ptr::null_mut(); + } + let cstr = std::ffi::CStr::from_ptr(entry.column_name); + let name = match cstr.to_str() { + Ok(s) => s.to_owned(), + Err(_) => { + set_error("bloom_filter_columns contains invalid UTF-8 name".into()); + return ptr::null_mut(); + } + }; + out.push(BloomFilterConfig { + column_name: name, + ndv: entry.ndv, + fpp: entry.fpp, + }); + } + out + }; let opts = WriterOptions { compression: options.compression, zstd_level: options.zstd_level, @@ -210,6 +253,7 @@ pub unsafe extern "C" fn mosaic_writer_open( max_dict_entries: options.max_dict_entries as usize, stats_columns: stats_cols, page_size_threshold: options.page_size_threshold as usize, + bloom_filter_columns: bloom_cols, }; match MosaicWriter::new(ffi_stream, &arrow_schema, opts) { Ok(writer) => Box::into_raw(Box::new(MosaicWriterHandle { @@ -978,6 +1022,113 @@ pub unsafe extern "C" fn mosaic_record_batch_free(handle: *mut MosaicRecordBatch } } +#[no_mangle] +pub unsafe extern "C" fn mosaic_reader_bloom_might_contain( + handle: *const MosaicReaderHandle, + rg_index: u32, + column_index: u32, + type_byte: u8, + value_bytes: *const u8, + value_len: u32, + out_might_contain: *mut u8, +) -> c_int { + if handle.is_null() || out_might_contain.is_null() { + set_error("null pointer".into()); + return -1; + } + if value_bytes.is_null() && value_len > 0 { + set_error("value_bytes is null with non-zero length".into()); + return -1; + } + let bytes = if value_len == 0 { + &[][..] + } else { + std::slice::from_raw_parts(value_bytes, value_len as usize) + }; + let value = match value_from_type_byte(type_byte, bytes) { + Ok(v) => v, + Err(msg) => { + set_error(msg); + return -1; + } + }; + let h = &*handle; + let filter = match h + .reader + .bloom_filter(rg_index as usize, column_index as usize) + { + Ok(opt) => opt, + Err(e) => { + set_error(e.to_string()); + return -1; + } + }; + match filter { + None => { + *out_might_contain = 1; + 1 + } + Some(f) => { + let hv = bloom::hash_value(&value); + *out_might_contain = if f.contains_hash(hv) { 1 } else { 0 }; + 0 + } + } +} + +fn value_from_type_byte(type_byte: u8, bytes: &[u8]) -> Result { + let need = |n: usize| -> Result<(), String> { + if bytes.len() != n { + Err(format!( + "value_bytes length {} does not match expected {} for type {}", + bytes.len(), + n, + type_byte + )) + } else { + Ok(()) + } + }; + match type_byte { + 0 => { + need(1)?; + Ok(Value::Boolean(bytes[0] != 0)) + } + 1 => { + need(1)?; + Ok(Value::TinyInt(bytes[0] as i8)) + } + 2 => { + need(2)?; + Ok(Value::SmallInt(i16::from_le_bytes([bytes[0], bytes[1]]))) + } + 3 => { + need(4)?; + Ok(Value::Integer(i32::from_le_bytes( + bytes.try_into().unwrap(), + ))) + } + 4 => { + need(8)?; + Ok(Value::BigInt(i64::from_le_bytes(bytes.try_into().unwrap()))) + } + 5 => { + need(4)?; + Ok(Value::Float(f32::from_le_bytes(bytes.try_into().unwrap()))) + } + 6 => { + need(8)?; + Ok(Value::Double(f64::from_le_bytes(bytes.try_into().unwrap()))) + } + 7 => { + need(4)?; + Ok(Value::Date(i32::from_le_bytes(bytes.try_into().unwrap()))) + } + 10 => Ok(Value::String(bytes.to_vec())), + other => Err(format!("unsupported bloom value type_byte {}", other)), + } +} + // ======================== Error ======================== /// Get the last error message. Returns a NUL-terminated pointer to a thread-local string. diff --git a/include/mosaic.hpp b/include/mosaic.hpp index 6f80a40..9bb5d5c 100644 --- a/include/mosaic.hpp +++ b/include/mosaic.hpp @@ -82,6 +82,12 @@ inline int64_t stream_get_pos(void* ctx) noexcept { } // namespace detail +struct BloomFilterConfig { + std::string column_name; + uint64_t ndv = 0; + double fpp = 0.01; +}; + struct WriterOptions { uint8_t compression = 1; // ZSTD int zstd_level = 1; @@ -92,6 +98,7 @@ struct WriterOptions { const char* const* stats_columns = nullptr; uint32_t num_stats_columns = 0; uint32_t page_size_threshold = 32 * 1024; + std::vector bloom_filter_columns; }; // ======================== Statistics ======================== @@ -127,6 +134,20 @@ class Writer { c_opts.num_stats_columns = opts.num_stats_columns; c_opts.page_size_threshold = opts.page_size_threshold; + std::vector bloom_c_opts; + bloom_c_opts.reserve(opts.bloom_filter_columns.size()); + for (auto const& b : opts.bloom_filter_columns) { + MosaicBloomConfig bc{}; + bc.column_name = b.column_name.c_str(); + bc.ndv = b.ndv; + bc.fpp = b.fpp; + bloom_c_opts.push_back(bc); + } + if (!bloom_c_opts.empty()) { + c_opts.bloom_filter_columns = bloom_c_opts.data(); + c_opts.num_bloom_filter_columns = static_cast(bloom_c_opts.size()); + } + handle_ = mosaic_writer_open(stream, static_cast(arrow_schema), c_opts); if (!handle_) throw Error("failed to open writer"); } @@ -296,6 +317,20 @@ class Reader { return out; } + bool bloom_might_contain(uint32_t rg_index, uint32_t column_index, + uint8_t type_byte, + const uint8_t* value_bytes, uint32_t value_len) const { + uint8_t out = 1; + int rc = mosaic_reader_bloom_might_contain( + handle_, rg_index, column_index, type_byte, + value_bytes, value_len, &out); + if (rc < 0) { + const char* err = mosaic_last_error(); + throw Error(err ? err : "bloom_might_contain failed"); + } + return out != 0; + } + std::vector get_row_group_statistics(uint32_t rg_index) { uint32_t n = 0; check(mosaic_reader_row_group_num_stats(handle_, rg_index, &n)); diff --git a/java/src/main/java/org/apache/paimon/mosaic/BloomFilterConfig.java b/java/src/main/java/org/apache/paimon/mosaic/BloomFilterConfig.java new file mode 100644 index 0000000..e708b3a --- /dev/null +++ b/java/src/main/java/org/apache/paimon/mosaic/BloomFilterConfig.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.paimon.mosaic; + +import java.util.Objects; + +public final class BloomFilterConfig { + + public static final double DEFAULT_FPP = 0.01; + + private final String columnName; + private final long ndv; + private final double fpp; + + public BloomFilterConfig(String columnName, long ndv) { + this(columnName, ndv, DEFAULT_FPP); + } + + public BloomFilterConfig(String columnName, long ndv, double fpp) { + this.columnName = Objects.requireNonNull(columnName, "columnName"); + if (ndv < 0) { + throw new IllegalArgumentException("ndv must be non-negative"); + } + if (!(fpp > 0.0 && fpp < 1.0)) { + throw new IllegalArgumentException("fpp must be in (0, 1), got " + fpp); + } + this.ndv = ndv; + this.fpp = fpp; + } + + public String columnName() { return columnName; } + public long ndv() { return ndv; } + public double fpp() { return fpp; } +} diff --git a/java/src/main/java/org/apache/paimon/mosaic/MosaicReader.java b/java/src/main/java/org/apache/paimon/mosaic/MosaicReader.java index cab99db..3e49af5 100644 --- a/java/src/main/java/org/apache/paimon/mosaic/MosaicReader.java +++ b/java/src/main/java/org/apache/paimon/mosaic/MosaicReader.java @@ -19,8 +19,12 @@ package org.apache.paimon.mosaic; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import org.apache.arrow.c.ArrowArray; @@ -28,6 +32,8 @@ import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; public class MosaicReader implements AutoCloseable { @@ -121,6 +127,89 @@ public Map getRowGroupStatistics(int rgIndex) { return Collections.unmodifiableMap(result); } + public boolean bloomMightContain(int rgIndex, String columnName, Object value) { + List fields = schema.getFields(); + int columnIndex = -1; + ArrowType arrowType = null; + for (int i = 0; i < fields.size(); i++) { + if (fields.get(i).getName().equals(columnName)) { + columnIndex = i; + arrowType = fields.get(i).getType(); + break; + } + } + if (columnIndex < 0) { + throw new IllegalArgumentException("column not found: " + columnName); + } + int typeByte = arrowTypeByte(arrowType); + byte[] encoded = encodeValue(typeByte, value); + return NativeLib.nativeReaderBloomMightContain( + handle, rgIndex, columnIndex, typeByte, encoded); + } + + private static int arrowTypeByte(ArrowType type) { + if (type instanceof ArrowType.Bool) return 0; + if (type instanceof ArrowType.Int) { + int bw = ((ArrowType.Int) type).getBitWidth(); + switch (bw) { + case 8: return 1; + case 16: return 2; + case 32: return 3; + case 64: return 4; + default: throw new IllegalArgumentException("unsupported int width: " + bw); + } + } + if (type instanceof ArrowType.FloatingPoint) { + switch (((ArrowType.FloatingPoint) type).getPrecision()) { + case SINGLE: return 5; + case DOUBLE: return 6; + default: throw new IllegalArgumentException("unsupported float precision"); + } + } + if (type instanceof ArrowType.Date) return 7; + if (type instanceof ArrowType.Utf8) return 10; + throw new IllegalArgumentException("unsupported arrow type for bloom: " + type); + } + + private static byte[] encodeValue(int typeByte, Object value) { + switch (typeByte) { + case 0: + return new byte[] { (byte) (((Boolean) value) ? 1 : 0) }; + case 1: + return new byte[] { ((Number) value).byteValue() }; + case 2: { + ByteBuffer bb = ByteBuffer.allocate(2).order(ByteOrder.LITTLE_ENDIAN); + bb.putShort(((Number) value).shortValue()); + return bb.array(); + } + case 3: + case 7: { + ByteBuffer bb = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN); + bb.putInt(((Number) value).intValue()); + return bb.array(); + } + case 5: { + ByteBuffer bb = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN); + bb.putFloat(((Number) value).floatValue()); + return bb.array(); + } + case 4: { + ByteBuffer bb = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN); + bb.putLong(((Number) value).longValue()); + return bb.array(); + } + case 6: { + ByteBuffer bb = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN); + bb.putDouble(((Number) value).doubleValue()); + return bb.array(); + } + case 10: + return ((String) value).getBytes(StandardCharsets.UTF_8); + default: + throw new IllegalArgumentException("unsupported type byte for value encoding: " + typeByte); + } + } + @Override public void close() { if (handle != 0) { diff --git a/java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java b/java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java index 461a694..c7c444c 100644 --- a/java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java +++ b/java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java @@ -49,6 +49,16 @@ public MosaicWriter(OutputStream outputStream, Schema arrowSchema, WriterOptions try (ArrowSchema cSchema = ArrowSchema.allocateNew(allocator)) { try { Data.exportSchema(allocator, arrowSchema, null, cSchema); + java.util.List blooms = options.getBloomFilterColumns(); + String[] bloomNames = new String[blooms.size()]; + long[] bloomNdvs = new long[blooms.size()]; + double[] bloomFpps = new double[blooms.size()]; + for (int i = 0; i < blooms.size(); i++) { + BloomFilterConfig c = blooms.get(i); + bloomNames[i] = c.columnName(); + bloomNdvs[i] = c.ndv(); + bloomFpps[i] = c.fpp(); + } this.handle = NativeLib.nativeWriterOpen( outputStream, cSchema.memoryAddress(), @@ -59,7 +69,10 @@ public MosaicWriter(OutputStream outputStream, Schema arrowSchema, WriterOptions options.getMaxDictTotalBytes(), options.getMaxDictEntries(), options.getStatsColumns(), - options.getPageSizeThreshold()); + options.getPageSizeThreshold(), + bloomNames, + bloomNdvs, + bloomFpps); } finally { releaseExported(cSchema); } diff --git a/java/src/main/java/org/apache/paimon/mosaic/NativeLib.java b/java/src/main/java/org/apache/paimon/mosaic/NativeLib.java index 62ffd81..e4e504a 100644 --- a/java/src/main/java/org/apache/paimon/mosaic/NativeLib.java +++ b/java/src/main/java/org/apache/paimon/mosaic/NativeLib.java @@ -105,7 +105,9 @@ static native long nativeWriterOpen(OutputStream stream, long arrowSchemaAddr, int numBuckets, int compression, int zstdLevel, long rowGroupMaxSize, int maxDictTotalBytes, int maxDictEntries, String[] statsColumns, - int pageSizeThreshold); + int pageSizeThreshold, + String[] bloomColumns, long[] bloomNdvs, + double[] bloomFpps); static native void nativeWriterClose(long handle); static native void nativeWriterFree(long handle); static native long nativeWriterEstimatedSize(long handle); @@ -137,4 +139,9 @@ static native long nativeWriterOpen(OutputStream stream, long arrowSchemaAddr, static native long[] nativeReaderRowGroupStatNullCounts(long handle, int rgIndex); static native byte[][] nativeReaderRowGroupStatMins(long handle, int rgIndex); static native byte[][] nativeReaderRowGroupStatMaxs(long handle, int rgIndex); + + // Bloom filter + static native boolean nativeReaderBloomMightContain(long handle, int rgIndex, + int columnIndex, int typeByte, + byte[] valueBytes); } diff --git a/java/src/main/java/org/apache/paimon/mosaic/WriterOptions.java b/java/src/main/java/org/apache/paimon/mosaic/WriterOptions.java index d20791a..ca6e140 100644 --- a/java/src/main/java/org/apache/paimon/mosaic/WriterOptions.java +++ b/java/src/main/java/org/apache/paimon/mosaic/WriterOptions.java @@ -19,6 +19,10 @@ package org.apache.paimon.mosaic; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + public class WriterOptions { public static final int COMPRESSION_ZSTD = 1; @@ -31,6 +35,7 @@ public class WriterOptions { private int maxDictEntries = 255; private String[] statsColumns = new String[0]; private int pageSizeThreshold = 32 * 1024; + private List bloomFilterColumns = Collections.emptyList(); public WriterOptions() {} @@ -74,6 +79,26 @@ public WriterOptions pageSizeThreshold(int threshold) { return this; } + public WriterOptions bloomFilterColumns(List configs) { + this.bloomFilterColumns = configs == null + ? Collections.emptyList() + : new ArrayList<>(configs); + return this; + } + + public WriterOptions bloomFilterColumns(BloomFilterConfig... configs) { + if (configs == null || configs.length == 0) { + this.bloomFilterColumns = Collections.emptyList(); + } else { + List list = new ArrayList<>(configs.length); + for (BloomFilterConfig c : configs) { + list.add(c); + } + this.bloomFilterColumns = list; + } + return this; + } + int getCompression() { return compression; } int getZstdLevel() { return zstdLevel; } int getNumBuckets() { return numBuckets; } @@ -82,4 +107,5 @@ public WriterOptions pageSizeThreshold(int threshold) { int getMaxDictEntries() { return maxDictEntries; } String[] getStatsColumns() { return statsColumns; } int getPageSizeThreshold() { return pageSizeThreshold; } + List getBloomFilterColumns() { return bloomFilterColumns; } } diff --git a/java/src/test/java/org/apache/paimon/mosaic/BloomFilterTest.java b/java/src/test/java/org/apache/paimon/mosaic/BloomFilterTest.java new file mode 100644 index 0000000..2899bf4 --- /dev/null +++ b/java/src/test/java/org/apache/paimon/mosaic/BloomFilterTest.java @@ -0,0 +1,154 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.paimon.mosaic; + +import java.io.ByteArrayOutputStream; +import java.util.Arrays; +import java.util.Collections; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import static org.junit.Assert.*; + +public class BloomFilterTest { + + private BufferAllocator allocator; + + @Before + public void setUp() { + allocator = new RootAllocator(); + } + + @After + public void tearDown() { + allocator.close(); + } + + private MosaicReader readerFromBytes(byte[] data) { + InputFile inputFile = (position, buffer, offset, length) -> + System.arraycopy(data, (int) position, buffer, offset, length); + return MosaicReader.open(inputFile, data.length, allocator); + } + + @Test + public void bigIntColumnHitsAndMisses() { + Schema schema = new Schema(Collections.singletonList( + Field.notNullable("id", new ArrowType.Int(64, true)))); + WriterOptions opts = new WriterOptions() + .numBuckets(2) + .bloomFilterColumns(Collections.singletonList( + new BloomFilterConfig("id", 2000, 0.01))); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (MosaicWriter writer = new MosaicWriter(baos, schema, opts, allocator); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + BigIntVector id = (BigIntVector) root.getVector("id"); + int n = 2000; + id.allocateNew(n); + for (int i = 0; i < n; i++) { + id.set(i, i); + } + root.setRowCount(n); + writer.write(root); + } + byte[] bytes = baos.toByteArray(); + + try (MosaicReader reader = readerFromBytes(bytes)) { + for (long i = 0; i < 2000; i++) { + assertTrue("present value " + i + " missed", reader.bloomMightContain(0, "id", i)); + } + int falsePositives = 0; + int probe = 5000; + for (long i = 1_000_000; i < 1_000_000 + probe; i++) { + if (reader.bloomMightContain(0, "id", i)) { + falsePositives++; + } + } + double rate = (double) falsePositives / probe; + assertTrue("fpp " + rate + " above 5%", rate < 0.05); + } + } + + @Test + public void stringColumnRejectsAbsentValue() { + Schema schema = new Schema(Collections.singletonList( + Field.notNullable("name", ArrowType.Utf8.INSTANCE))); + WriterOptions opts = new WriterOptions() + .numBuckets(1) + .bloomFilterColumns(Collections.singletonList( + new BloomFilterConfig("name", 1024, 0.001))); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + String[] names = new String[] {"alice", "bob", "carol", "dave", "eve"}; + try (MosaicWriter writer = new MosaicWriter(baos, schema, opts, allocator); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VarCharVector vec = (VarCharVector) root.getVector("name"); + vec.allocateNew(names.length); + for (int i = 0; i < names.length; i++) { + vec.setSafe(i, names[i].getBytes(java.nio.charset.StandardCharsets.UTF_8)); + } + root.setRowCount(names.length); + writer.write(root); + } + byte[] bytes = baos.toByteArray(); + + try (MosaicReader reader = readerFromBytes(bytes)) { + for (String n : names) { + assertTrue("present name " + n + " missed", reader.bloomMightContain(0, "name", n)); + } + assertFalse("absent name unexpectedly present", reader.bloomMightContain(0, "name", "zachary")); + } + } + + @Test + public void noBloomReturnsTrueConservatively() { + Schema schema = new Schema(Collections.singletonList( + Field.notNullable("id", new ArrowType.Int(64, true)))); + WriterOptions opts = new WriterOptions().numBuckets(1); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (MosaicWriter writer = new MosaicWriter(baos, schema, opts, allocator); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + BigIntVector id = (BigIntVector) root.getVector("id"); + id.allocateNew(3); + id.set(0, 1L); + id.set(1, 2L); + id.set(2, 3L); + root.setRowCount(3); + writer.write(root); + } + byte[] bytes = baos.toByteArray(); + + try (MosaicReader reader = readerFromBytes(bytes)) { + assertTrue(reader.bloomMightContain(0, "id", 99999L)); + } + } +} diff --git a/jni/src/lib.rs b/jni/src/lib.rs index 98fe7c4..f11c9b7 100644 --- a/jni/src/lib.rs +++ b/jni/src/lib.rs @@ -21,9 +21,10 @@ use std::ptr; use std::sync::Arc; use jni::objects::{ - GlobalRef, JByteArray, JClass, JMethodID, JObject, JObjectArray, JString, JValue, + GlobalRef, JByteArray, JClass, JDoubleArray, JLongArray, JMethodID, JObject, JObjectArray, + JString, JValue, }; -use jni::sys::{jint, jlong, jlongArray}; +use jni::sys::{jboolean, jint, jlong, jlongArray, JNI_TRUE}; use jni::JNIEnv; use jni::JavaVM; @@ -31,8 +32,10 @@ use arrow_array::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; use arrow_array::{RecordBatch, StructArray}; use arrow_schema::Schema; +use mosaic_core::bloom::{self, BloomFilterConfig}; use mosaic_core::reader::{InputFile, MosaicReader, ReaderAccess, RowGroupReader}; use mosaic_core::spec::*; +use mosaic_core::values::Value; use mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions}; fn panic_message(e: &Box) -> String { @@ -209,6 +212,9 @@ pub extern "system" fn Java_org_apache_paimon_mosaic_NativeLib_nativeWriterOpen( max_dict_entries: jint, stats_columns: JObjectArray<'_>, page_size_threshold: jint, + bloom_columns: JObjectArray<'_>, + bloom_ndvs: JLongArray<'_>, + bloom_fpps: JDoubleArray<'_>, ) -> jlong { let raw_env = env.get_raw(); let result = panic::catch_unwind(AssertUnwindSafe(|| { @@ -303,6 +309,50 @@ pub extern "system" fn Java_org_apache_paimon_mosaic_NativeLib_nativeWriterOpen( num_buckets as usize }; + let bloom_cols: Vec = match env.get_array_length(&bloom_columns) { + Ok(len) if len > 0 => { + let count = len as usize; + let mut ndvs_buf = vec![0i64; count]; + if let Err(e) = env.get_long_array_region(&bloom_ndvs, 0, &mut ndvs_buf) { + throw(&mut env, &format!("failed to read bloom_ndvs: {}", e)); + return 0; + } + let mut fpps_buf = vec![0f64; count]; + if let Err(e) = env.get_double_array_region(&bloom_fpps, 0, &mut fpps_buf) { + throw(&mut env, &format!("failed to read bloom_fpps: {}", e)); + return 0; + } + let mut out = Vec::with_capacity(count); + for i in 0..count { + let obj = match env.get_object_array_element(&bloom_columns, i as jint) { + Ok(o) => o, + Err(_) => { + throw(&mut env, "failed to read bloom_columns element"); + return 0; + } + }; + let jstr = JString::from(obj); + let name: String = match env.get_string(&jstr) { + Ok(s) => s.into(), + Err(_) => { + throw( + &mut env, + "failed to convert bloom_columns element to string", + ); + return 0; + } + }; + out.push(BloomFilterConfig { + column_name: name, + ndv: ndvs_buf[i] as u64, + fpp: fpps_buf[i], + }); + } + out + } + _ => Vec::new(), + }; + let opts = WriterOptions { compression: compression as u8, zstd_level, @@ -312,6 +362,7 @@ pub extern "system" fn Java_org_apache_paimon_mosaic_NativeLib_nativeWriterOpen( max_dict_entries: max_dict_entries as usize, stats_columns: stats_cols, page_size_threshold: page_size_threshold as usize, + bloom_filter_columns: bloom_cols, }; let writer = match MosaicWriter::new(jni_stream, &arrow_schema, opts) { @@ -1005,3 +1056,130 @@ pub extern "system" fn Java_org_apache_paimon_mosaic_NativeLib_nativeRowGroupRea } } } + +#[no_mangle] +pub extern "system" fn Java_org_apache_paimon_mosaic_NativeLib_nativeReaderBloomMightContain( + mut env: JNIEnv, + _class: JClass, + handle: jlong, + rg_index: jint, + column_index: jint, + type_byte: jint, + value_bytes: JByteArray<'_>, +) -> jboolean { + let raw_env = env.get_raw(); + let result = panic::catch_unwind(AssertUnwindSafe(|| { + if handle == 0 { + throw(&mut env, "null reader handle"); + return JNI_TRUE; + } + let len = match env.get_array_length(&value_bytes) { + Ok(n) => n as usize, + Err(e) => { + throw( + &mut env, + &format!("failed to read value bytes length: {}", e), + ); + return JNI_TRUE; + } + }; + let mut buf = vec![0i8; len]; + if len > 0 + && env + .get_byte_array_region(&value_bytes, 0, &mut buf) + .is_err() + { + throw(&mut env, "failed to copy value bytes"); + return JNI_TRUE; + } + let bytes: Vec = buf.into_iter().map(|b| b as u8).collect(); + let value = match jni_value_from_type_byte(type_byte as u8, &bytes) { + Ok(v) => v, + Err(msg) => { + throw(&mut env, &msg); + return JNI_TRUE; + } + }; + let rh = unsafe { &*(handle as *const ReaderHandle) }; + let filter = match rh + .reader + .bloom_filter(rg_index as usize, column_index as usize) + { + Ok(opt) => opt, + Err(e) => { + throw(&mut env, &format!("bloom_filter fetch failed: {}", e)); + return JNI_TRUE; + } + }; + let result = match filter { + None => true, + Some(f) => f.contains_hash(bloom::hash_value(&value)), + }; + if result { + JNI_TRUE + } else { + 0u8 + } + })); + match result { + Ok(v) => v, + Err(e) => { + let mut env = unsafe { JNIEnv::from_raw(raw_env).unwrap() }; + throw(&mut env, &panic_message(&e)); + JNI_TRUE + } + } +} + +fn jni_value_from_type_byte(type_byte: u8, bytes: &[u8]) -> Result { + let need = |n: usize| -> Result<(), String> { + if bytes.len() != n { + Err(format!( + "value bytes length {} does not match expected {} for type {}", + bytes.len(), + n, + type_byte + )) + } else { + Ok(()) + } + }; + match type_byte { + 0 => { + need(1)?; + Ok(Value::Boolean(bytes[0] != 0)) + } + 1 => { + need(1)?; + Ok(Value::TinyInt(bytes[0] as i8)) + } + 2 => { + need(2)?; + Ok(Value::SmallInt(i16::from_le_bytes([bytes[0], bytes[1]]))) + } + 3 => { + need(4)?; + Ok(Value::Integer(i32::from_le_bytes( + bytes.try_into().unwrap(), + ))) + } + 4 => { + need(8)?; + Ok(Value::BigInt(i64::from_le_bytes(bytes.try_into().unwrap()))) + } + 5 => { + need(4)?; + Ok(Value::Float(f32::from_le_bytes(bytes.try_into().unwrap()))) + } + 6 => { + need(8)?; + Ok(Value::Double(f64::from_le_bytes(bytes.try_into().unwrap()))) + } + 7 => { + need(4)?; + Ok(Value::Date(i32::from_le_bytes(bytes.try_into().unwrap()))) + } + 10 => Ok(Value::String(bytes.to_vec())), + other => Err(format!("unsupported bloom value type_byte {}", other)), + } +} diff --git a/python/mosaic/__init__.py b/python/mosaic/__init__.py index f9eb3b8..0a4a4bf 100644 --- a/python/mosaic/__init__.py +++ b/python/mosaic/__init__.py @@ -16,6 +16,7 @@ # under the License. from .mosaic import ( + BloomFilterConfig, ColumnStatistics, MosaicReader, MosaicWriter, @@ -25,6 +26,7 @@ ) __all__ = [ + "BloomFilterConfig", "ColumnStatistics", "WriterOptions", "MosaicWriter", diff --git a/python/mosaic/_ffi.py b/python/mosaic/_ffi.py index 1b94438..d77118a 100644 --- a/python/mosaic/_ffi.py +++ b/python/mosaic/_ffi.py @@ -110,6 +110,14 @@ class MosaicInputFile(Structure): ] +class MosaicBloomConfig(Structure): + _fields_ = [ + ("column_name", c_char_p), + ("ndv", c_uint64), + ("fpp", c_double), + ] + + class MosaicWriterOptions(Structure): _fields_ = [ ("compression", c_uint8), @@ -121,6 +129,8 @@ class MosaicWriterOptions(Structure): ("stats_columns", POINTER(c_char_p)), ("num_stats_columns", c_uint32), ("page_size_threshold", c_uint32), + ("bloom_filter_columns", POINTER(MosaicBloomConfig)), + ("num_bloom_filter_columns", c_uint32), ] @@ -227,6 +237,12 @@ class MosaicWriterOptions(Structure): ] lib.mosaic_reader_row_group_stats.restype = c_int +lib.mosaic_reader_bloom_might_contain.argtypes = [ + c_void_p, c_uint32, c_uint32, c_uint8, + POINTER(c_uint8), c_uint32, POINTER(c_uint8), +] +lib.mosaic_reader_bloom_might_contain.restype = c_int + # ======================== Error ======================== lib.mosaic_last_error.argtypes = [] diff --git a/python/mosaic/mosaic.py b/python/mosaic/mosaic.py index b08096d..d2f00a5 100644 --- a/python/mosaic/mosaic.py +++ b/python/mosaic/mosaic.py @@ -85,6 +85,17 @@ def _fetch_rg_stats(num_stats_fn, stats_fn, handle, rg_index): return result +class BloomFilterConfig: + def __init__(self, column_name, ndv, fpp=0.01): + if ndv < 0: + raise ValueError("ndv must be non-negative") + if not (0.0 < fpp < 1.0): + raise ValueError(f"fpp must be in (0, 1), got {fpp}") + self.column_name = column_name + self.ndv = int(ndv) + self.fpp = float(fpp) + + class WriterOptions: COMPRESSION_NONE = 0 COMPRESSION_ZSTD = 1 @@ -99,6 +110,7 @@ def __init__( max_dict_entries=255, stats_columns=None, page_size_threshold=32 * 1024, + bloom_filter_columns=None, ): self.compression = compression self.zstd_level = zstd_level @@ -108,6 +120,7 @@ def __init__( self.max_dict_entries = max_dict_entries self.stats_columns = stats_columns or [] self.page_size_threshold = page_size_threshold + self.bloom_filter_columns = bloom_filter_columns or [] def _to_ffi(self): opts = _ffi.MosaicWriterOptions() @@ -129,6 +142,20 @@ def _to_ffi(self): opts.stats_columns = None opts.num_stats_columns = 0 opts.page_size_threshold = self.page_size_threshold + if self.bloom_filter_columns: + encoded_names = [c.column_name.encode("utf-8") for c in self.bloom_filter_columns] + bloom_arr = (_ffi.MosaicBloomConfig * len(self.bloom_filter_columns))() + for i, c in enumerate(self.bloom_filter_columns): + bloom_arr[i].column_name = encoded_names[i] + bloom_arr[i].ndv = c.ndv + bloom_arr[i].fpp = c.fpp + refs.append(bloom_arr) + refs.append(encoded_names) + opts.bloom_filter_columns = bloom_arr + opts.num_bloom_filter_columns = len(self.bloom_filter_columns) + else: + opts.bloom_filter_columns = None + opts.num_bloom_filter_columns = 0 return opts, refs @@ -396,6 +423,27 @@ def get_row_group_statistics(self, rg_index): lib.mosaic_reader_row_group_stats, self._handle, rg_index) + def bloom_might_contain(self, rg_index, column_name, value): + column_index = self._schema.get_field_index(column_name) + if column_index < 0: + raise ValueError(f"column not found: {column_name}") + field = self._schema.field(column_index) + type_byte, encoded = _encode_bloom_value(field.type, value) + buf = (ctypes.c_uint8 * len(encoded)).from_buffer_copy(encoded) + out = ctypes.c_uint8(0) + rc = lib.mosaic_reader_bloom_might_contain( + self._handle, + rg_index, + column_index, + type_byte, + buf, + len(encoded), + ctypes.byref(out), + ) + if rc < 0: + _check_error("bloom_might_contain failed") + return out.value != 0 + def close(self): if self._handle: lib.mosaic_reader_free(self._handle) @@ -414,6 +462,29 @@ def __del__(self): self.close() +def _encode_bloom_value(arrow_type, value): + import struct + if pa.types.is_boolean(arrow_type): + return 0, bytes([1 if value else 0]) + if pa.types.is_int8(arrow_type): + return 1, struct.pack("