This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11953-89fc173a933d78e5cee08a47c901209c2de59521 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit e2c48ef0dd83fa464d643e25f8e9f9177ca7316e Author: Ed Seidl <[email protected]> AuthorDate: Sun Nov 1 15:00:01 2026 +0000 feat(parquet): Add async `Sbbf::read_from_column_chunk_async` Bloom filter API (#11953) # Which issue does this PR close? - Closes #9067 # Rationale for this change Applications with existing Parquet column metadata and an `AsyncFileReader` cannot currently load a Bloom filter through a public API without constructing a `ParquetRecordBatchStreamBuilder` or duplicating internal header parsing. Expose the async equivalent of `Sbbf::read_from_column_chunk` so metadata-only pruning code can read a selected physical column's Bloom filter directly. # What changes are included in this PR? Add `Sbbf::read_from_column_chunk_async<R: AsyncFileReader + ?Sized>` and delegate `ParquetRecordBatchStreamBuilder::get_row_group_column_bloom_filter` to it. # Are these changes tested? Yes, new tests are added. # Are there any significant user-facing changes? Adds to the public API, but changes no public signatures. Malformed lengths or incomplete bitsets will now return an error. # AI disclosure AI assistance was used to author the implementation and tests. --- parquet/src/arrow/async_reader/mod.rs | 63 +---- parquet/src/bloom_filter/mod.rs | 112 ++++++++ parquet/tests/arrow_reader/bloom_filter/async.rs | 333 ++++++++++++++++++++++- 3 files changed, 444 insertions(+), 64 deletions(-) diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 7a1e78d962..8a46c25cb6 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -40,10 +40,7 @@ use crate::arrow::arrow_reader::{ ArrowReaderBuilder, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, }; -use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash}; -use crate::bloom_filter::{ - SBBF_HEADER_SIZE_ESTIMATE, Sbbf, chunk_read_bloom_filter_header_and_offset, -}; +use crate::bloom_filter::Sbbf; use crate::errors::{ParquetError, Result}; use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; @@ -593,63 +590,7 @@ impl<T: AsyncFileReader + Send + 'static> ParquetRecordBatchStreamBuilder<T> { let metadata = self.metadata.row_group(row_group_idx); let column_metadata = metadata.column(column_idx); - let offset: u64 = if let Some(offset) = column_metadata.bloom_filter_offset() { - offset - .try_into() - .map_err(|_| ParquetError::General("Bloom filter offset is invalid".to_string()))? - } else { - return Ok(None); - }; - - let buffer = match column_metadata.bloom_filter_length() { - Some(length) => self.input.0.get_bytes(offset..offset + length as u64), - None => self - .input - .0 - .get_bytes(offset..offset + SBBF_HEADER_SIZE_ESTIMATE as u64), - } - .await?; - - let (header, bitset_offset) = - chunk_read_bloom_filter_header_and_offset(offset, buffer.clone())?; - - match header.algorithm { - BloomFilterAlgorithm::BLOCK => { - // this match exists to future proof the singleton algorithm enum - } - } - match header.compression { - BloomFilterCompression::UNCOMPRESSED => { - // this match exists to future proof the singleton compression enum - } - } - match header.hash { - BloomFilterHash::XXHASH => { - // this match exists to future proof the singleton hash enum - } - } - - let bitset = match column_metadata.bloom_filter_length() { - Some(_) => { - let bitset_start = bitset_offset - .checked_sub(offset) - .and_then(|start| usize::try_from(start).ok()) - .ok_or_else(|| { - ParquetError::General("Bloom filter offset is invalid".to_string()) - })?; - buffer.slice(bitset_start..) - } - None => { - let bitset_length: u64 = header.num_bytes.try_into().map_err(|_| { - ParquetError::General("Bloom filter length is invalid".to_string()) - })?; - self.input - .0 - .get_bytes(bitset_offset..bitset_offset + bitset_length) - .await? - } - }; - Ok(Some(Sbbf::new(&bitset))) + Sbbf::read_from_column_chunk_async(column_metadata, &mut self.input.0).await } /// Select row groups and rows using row-group-local coordinates. diff --git a/parquet/src/bloom_filter/mod.rs b/parquet/src/bloom_filter/mod.rs index b936bcbdb2..ab09188186 100644 --- a/parquet/src/bloom_filter/mod.rs +++ b/parquet/src/bloom_filter/mod.rs @@ -112,6 +112,8 @@ //! [sbbf-paper]: https://arxiv.org/pdf/2101.01719 //! [bf-formulae]: http://tfk.mit.edu/pdf/bloom.pdf +#[cfg(all(feature = "arrow", feature = "async"))] +use crate::arrow::async_reader::AsyncFileReader; use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash}; use crate::data_type::AsBytes; use crate::errors::{ParquetError, Result}; @@ -520,6 +522,82 @@ impl Sbbf { Ok(Some(Self::new(&bitset))) } + /// Read a Bloom filter asynchronously from existing column-chunk metadata. + /// + /// The metadata and reader must refer to the same file. + /// + /// A missing Bloom filter offset returns `None` without performing I/O. Invalid + /// metadata, malformed filters, and read failures return an error. + /// + /// When the total filter length is present, one range is fetched. Otherwise, + /// this may perform additional fetches. + #[cfg(all(feature = "arrow", feature = "async"))] + pub async fn read_from_column_chunk_async<R: AsyncFileReader + ?Sized>( + column_metadata: &ColumnChunkMetaData, + reader: &mut R, + ) -> Result<Option<Self>> { + let Some(offset) = column_metadata.bloom_filter_offset() else { + return Ok(None); + }; + let offset = + u64::try_from(offset).map_err(|_| general_err!("Bloom filter offset is invalid"))?; + let length = column_metadata + .bloom_filter_length() + .map(|length| { + usize::try_from(length) + .ok() + .filter(|length| *length > 0) + .ok_or_else(|| general_err!("Bloom filter length is invalid")) + }) + .transpose()?; + let buffer = read_bloom_filter_bytes_async( + reader, + offset, + length.unwrap_or(SBBF_HEADER_SIZE_ESTIMATE), + ) + .await?; + let (header, header_length) = read_bloom_filter_header_and_length(buffer.clone())?; + + // The following matches are present to force an explicit decision + // if supported variants change. + match header.algorithm { + BloomFilterAlgorithm::BLOCK => {} + } + match header.hash { + BloomFilterHash::XXHASH => {} + } + match header.compression { + BloomFilterCompression::UNCOMPRESSED => {} + } + + let bitset_length = usize::try_from(header.num_bytes) + .map_err(|_| general_err!("Bloom filter bitset length is invalid"))?; + if bitset_length < BITSET_MIN_LENGTH || !bitset_length.is_multiple_of(32) { + return Err(general_err!( + "Bloom filter bitset length must be a positive multiple of 32" + )); + } + let bitset_offset = offset + .checked_add(header_length) + .ok_or_else(|| general_err!("Bloom filter offset overflow"))?; + let bitset = if let Some(length) = length { + let header_length = usize::try_from(header_length) + .map_err(|_| general_err!("Bloom filter header length is invalid"))?; + let total = header_length + .checked_add(bitset_length) + .ok_or_else(|| general_err!("Bloom filter length overflow"))?; + if total != length { + return Err(general_err!( + "Bloom filter length does not match its header" + )); + } + buffer.slice(header_length..total) + } else { + read_bloom_filter_bytes_async(reader, bitset_offset, bitset_length).await? + }; + Ok(Some(Self::new(&bitset))) + } + /// Map a 64-bit hash to a block index in `[0, num_blocks)`. /// /// Uses the "multiply-and-shift" trick (a fast alternative to modulo): @@ -778,6 +856,25 @@ impl Sbbf { } } +#[cfg(all(feature = "arrow", feature = "async"))] +async fn read_bloom_filter_bytes_async<R: AsyncFileReader + ?Sized>( + reader: &mut R, + offset: u64, + length: usize, +) -> Result<Bytes> { + let end = offset + .checked_add(length as u64) + .ok_or_else(|| general_err!("Bloom filter read range overflow"))?; + let bytes = reader.get_bytes(offset..end).await?; + if bytes.len() != length { + return Err(general_err!( + "Bloom filter read length mismatch: expected {length} bytes, got {}", + bytes.len() + )); + } + Ok(bytes) +} + // per spec we use xxHash with seed=0 const SEED: u64 = 0; @@ -790,6 +887,21 @@ fn hash_as_bytes<A: AsBytes + ?Sized>(value: &A) -> u64 { mod tests { use super::*; + #[cfg(all(feature = "arrow", feature = "async"))] + #[tokio::test] + async fn test_async_bloom_filter_range_overflow_before_io() { + let mut reader = std::io::Cursor::new(Vec::<u8>::new()); + let error = read_bloom_filter_bytes_async(&mut reader, u64::MAX, 1) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("Bloom filter read range overflow") + ); + assert_eq!(reader.position(), 0); + } + #[test] fn test_hash_bytes() { assert_eq!(hash_as_bytes(""), 17241709254077376921); diff --git a/parquet/tests/arrow_reader/bloom_filter/async.rs b/parquet/tests/arrow_reader/bloom_filter/async.rs index e230b33d2d..d228fb14d8 100644 --- a/parquet/tests/arrow_reader/bloom_filter/async.rs +++ b/parquet/tests/arrow_reader/bloom_filter/async.rs @@ -16,12 +16,74 @@ // under the License. use crate::io::TestReader; +use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray, StructArray}; +use arrow_schema::{DataType, Field}; use bytes::Bytes; -use futures::TryStreamExt; +use futures::{ + TryStreamExt, + future::{BoxFuture, ready}, +}; use parquet::{ - arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder}, - file::properties::WriterProperties, + arrow::{ + ArrowWriter, ParquetRecordBatchStreamBuilder, arrow_reader::ArrowReaderOptions, + async_reader::AsyncFileReader, + }, + basic::Type as PhysicalType, + bloom_filter::Sbbf, + errors::{ParquetError, Result}, + file::{ + metadata::{ColumnChunkMetaData, ParquetMetaData}, + properties::WriterProperties, + reader::FileReader, + serialized_reader::SerializedFileReader, + }, + schema::types::{ColumnDescriptor, ColumnPath, Type}, }; +use std::collections::VecDeque; +use std::ops::Range; +use std::sync::Arc; + +#[tokio::test] +async fn test_standalone_bloom_filter_after_nested_columns() { + let nested = StructArray::from(vec![ + ( + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + ), + ( + Arc::new(Field::new("b", DataType::Int32, false)), + Arc::new(Int32Array::from(vec![3, 4])) as ArrayRef, + ), + ]); + let batch = RecordBatch::try_from_iter([ + ("nested", Arc::new(nested) as ArrayRef), + ( + "value", + Arc::new(StringArray::from(vec!["present", "other"])) as ArrayRef, + ), + ]) + .unwrap(); + let props = WriterProperties::builder() + .set_column_bloom_filter_enabled("value".into(), true) + .build(); + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(data); + let metadata_reader = SerializedFileReader::new(data.clone()).unwrap(); + let row_group = metadata_reader.metadata().row_group(0); + // The second Arrow field is the third physical Parquet leaf. + assert_eq!(row_group.column(2).column_path().string(), "value"); + let mut reader = TestReader::new(data); + let filter = Sbbf::read_from_column_chunk_async(row_group.column(2), &mut reader) + .await + .unwrap() + .unwrap(); + assert!(filter.check("present")); + assert!(filter.check("other")); + assert_eq!(reader.requests().lock().unwrap().len(), 1); +} #[tokio::test] async fn test_get_row_group_column_bloom_filter_without_length() { @@ -61,6 +123,7 @@ async fn test_get_row_group_column_bloom_filter_with_length() { async fn test_get_row_group_column_bloom_filter(data: Bytes, with_length: bool) { let async_reader = TestReader::new(data.clone()); + let builder_requests = async_reader.requests(); let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader) .await @@ -71,6 +134,18 @@ async fn test_get_row_group_column_bloom_filter(data: Bytes, with_length: bool) let row_group = metadata.row_group(0); let column = row_group.column(0); assert_eq!(column.bloom_filter_length().is_some(), with_length); + // Compare standalone reads with the existing sync reader and async builder. + let mut reader = TestReader::new(data.clone()); + let standalone = Sbbf::read_from_column_chunk_async(column, &mut reader) + .await + .unwrap() + .unwrap(); + let sync = Sbbf::read_from_column_chunk(column, &data) + .unwrap() + .unwrap(); + assert_eq!(standalone.num_blocks(), sync.num_blocks()); + assert!(standalone.check(&"Hello")); + assert!(!standalone.check(&"Hello_Not_Exists")); let sbbf = builder .get_row_group_column_bloom_filter(0, 0) @@ -79,4 +154,256 @@ async fn test_get_row_group_column_bloom_filter(data: Bytes, with_length: bool) .unwrap(); assert!(sbbf.check(&"Hello")); assert!(!sbbf.check(&"Hello_Not_Exists")); + assert_eq!(sbbf.num_blocks(), standalone.num_blocks()); + // Delegation must preserve the same Bloom-only ranges, with or without a length. + assert_eq!( + *builder_requests.lock().unwrap(), + *reader.requests().lock().unwrap() + ); + assert_eq!( + builder_requests.lock().unwrap().len(), + if with_length { 1 } else { 2 } + ); +} + +#[derive(Default)] +struct RecordingReader { + responses: VecDeque<Result<Bytes>>, + requests: Vec<Range<u64>>, +} + +impl RecordingReader { + fn new(responses: impl IntoIterator<Item = Bytes>) -> Self { + Self { + responses: responses.into_iter().map(Ok).collect(), + requests: Vec::new(), + } + } +} + +impl AsyncFileReader for RecordingReader { + fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> { + self.requests.push(range); + Box::pin(ready(self.responses.pop_front().expect("unexpected read"))) + } + + fn get_metadata<'a>( + &'a mut self, + _: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result<Arc<ParquetMetaData>>> { + panic!("Bloom reads must not load metadata") + } +} + +fn column(offset: Option<i64>, length: Option<i32>) -> ColumnChunkMetaData { + let descriptor = Arc::new(ColumnDescriptor::new( + Arc::new( + Type::primitive_type_builder("value", PhysicalType::BYTE_ARRAY) + .build() + .unwrap(), + ), + 0, + 0, + ColumnPath::new(vec!["value".to_owned()]), + )); + ColumnChunkMetaData::builder(descriptor) + .set_bloom_filter_offset(offset) + .set_bloom_filter_length(length) + .build() + .unwrap() +} + +fn encode(num_bytes: i32, bitset_length: usize) -> Bytes { + // Compact-Thrift zigzag varints for the sizes used by these malformed fixtures. + let encoded_size: &[u8] = match num_bytes { + -1 => &[0x01], + 0 => &[0x00], + 32 => &[0x40], + 33 => &[0x42], + 64 => &[0x80, 0x01], + _ => panic!("unsupported fixture size"), + }; + let mut bytes = vec![0x15]; // Field 1: i32 numBytes. + bytes.extend_from_slice(encoded_size); + // Fields 2–4: BLOCK, XXHASH, UNCOMPRESSED (each a union containing an empty struct), + // followed by the header's STOP byte. + bytes.extend_from_slice(&[0x1c, 0x1c, 0, 0, 0x1c, 0x1c, 0, 0, 0x1c, 0x1c, 0, 0, 0]); + bytes.resize(bytes.len() + bitset_length, 0); + bytes.into() +} + +#[tokio::test] +async fn standalone_filter_reads_use_exact_ranges() { + // Three blocks are valid even though the writer normally uses powers of two. + let mut filter = Sbbf::new(&[0; 96]); + filter.insert("present"); + let mut serialized = Vec::new(); + filter.write(&mut serialized).unwrap(); + let bytes = Bytes::from(serialized); + let header_length = (bytes.len() - filter.num_blocks() * 32) as u64; + let offset = 13; + + for length in [Some(bytes.len() as i32), None] { + let metadata = column(Some(offset as i64), length); + let responses = if length.is_some() { + vec![bytes.clone()] + } else { + vec![bytes.slice(..20), bytes.slice(header_length as usize..)] + }; + let mut reader = RecordingReader::new(responses); + let borrowed: &mut dyn AsyncFileReader = &mut reader; + let loaded = Sbbf::read_from_column_chunk_async(&metadata, borrowed) + .await + .unwrap() + .unwrap(); + assert_eq!(loaded.num_blocks(), 3); + assert!(loaded.check("present")); + // The legacy prefix overlaps the bitset; resume at the decoded header end. + let expected = if length.is_some() { + std::iter::once(offset..offset + bytes.len() as u64).collect() + } else { + vec![ + offset..offset + 20, + offset + header_length..offset + bytes.len() as u64, + ] + }; + assert_eq!(reader.requests, expected); + } +} + +#[tokio::test] +async fn missing_and_invalid_locations_do_not_read() { + let mut reader = RecordingReader::default(); + // An orphan length does not advertise a filter without an offset. + assert!( + Sbbf::read_from_column_chunk_async(&column(None, Some(-1)), &mut reader) + .await + .unwrap() + .is_none() + ); + // Reject invalid signed offsets and nonpositive lengths before issuing any read. + for metadata in [ + column(Some(-1), None), + column(Some(0), Some(-1)), + column(Some(0), Some(0)), + ] { + assert!( + Sbbf::read_from_column_chunk_async(&metadata, &mut reader) + .await + .is_err() + ); + } + assert!(reader.requests.is_empty()); +} + +#[tokio::test] +async fn invalid_headers_and_bitsets_remain_errors() { + let mut canonical = Vec::new(); + Sbbf::new(&[0; 32]).write(&mut canonical).unwrap(); + assert_eq!(encode(32, 32).as_ref(), canonical.as_slice()); + // Reject invalid bitset sizes, then both undersized and oversized payloads. + for (declared, actual) in [(-1, 32), (0, 0), (33, 33), (64, 32), (32, 64)] { + let bytes = encode(declared, actual); + let mut reader = RecordingReader::new([bytes.clone()]); + assert!( + Sbbf::read_from_column_chunk_async( + &column(Some(0), Some(bytes.len() as i32)), + &mut reader, + ) + .await + .is_err(), + "declared {declared}, actual {actual}" + ); + assert_eq!(reader.requests.len(), 1); + } + // Decode failures must remain errors, not be mistaken for an absent filter. + let mut unsupported = encode(32, 32).to_vec(); + unsupported[3] = 0x2c; // Unknown algorithm union variant. + for bytes in [Bytes::from_static(&[0]), Bytes::from(unsupported)] { + let mut reader = RecordingReader::new([bytes.clone()]); + assert!( + Sbbf::read_from_column_chunk_async( + &column(Some(0), Some(bytes.len() as i32)), + &mut reader, + ) + .await + .is_err() + ); + } +} + +#[tokio::test] +async fn exact_responses_and_transport_errors_are_required() { + let bytes = encode(32, 32); + // With a known total length, the entire response must match the requested range. + for response_length in [bytes.len() - 1, bytes.len() + 1] { + let mut response = bytes.to_vec(); + response.resize(response_length, 0); + let mut reader = RecordingReader::new([response.into()]); + assert!( + Sbbf::read_from_column_chunk_async( + &column(Some(0), Some(bytes.len() as i32)), + &mut reader, + ) + .await + .is_err() + ); + } + // Without a total length, validate the separate bitset response against the header. + for bitset_length in [31, 33] { + let mut reader = + RecordingReader::new([bytes.slice(..20), Bytes::from(vec![0; bitset_length])]); + assert!( + Sbbf::read_from_column_chunk_async(&column(Some(0), None), &mut reader) + .await + .is_err() + ); + assert_eq!(reader.requests.len(), 2); + } + // Propagate transport failures from either the header or bitset read. + for failure_at in [0, 1] { + let mut reader = RecordingReader::new([bytes.slice(..20)]); + reader.responses.insert( + failure_at, + Err(ParquetError::General("injected read failure".to_owned())), + ); + assert!( + Sbbf::read_from_column_chunk_async(&column(Some(0), None), &mut reader) + .await + .is_err() + ); + assert_eq!(reader.requests.len(), failure_at + 1); + } +} + +#[tokio::test] +async fn extended_headers_require_an_advertised_length() { + let mut serialized = Vec::new(); + Sbbf::new(&[0; 32]).write(&mut serialized).unwrap(); + let canonical = Bytes::from(serialized); + let header_length = canonical.len() - 32; + let mut bytes = canonical[..header_length - 1].to_vec(); + // Unknown binary field extends the header beyond the legacy prefix. + bytes.extend_from_slice(&[0x18, 24]); + bytes.extend_from_slice(&[0; 24]); + bytes.push(0); + bytes.extend_from_slice(&canonical[header_length..]); + let bytes = Bytes::from(bytes); + // An advertised length lets the parser skip the complete unknown field. + let mut known = RecordingReader::new([bytes.clone()]); + assert!( + Sbbf::read_from_column_chunk_async(&column(Some(0), Some(bytes.len() as i32)), &mut known,) + .await + .unwrap() + .is_some() + ); + // The legacy 20-byte prefix cannot decode this header; do not return a partial filter. + let mut legacy = RecordingReader::new([bytes.slice(..20)]); + assert!( + Sbbf::read_from_column_chunk_async(&column(Some(0), None), &mut legacy) + .await + .is_err() + ); + assert_eq!(legacy.requests.len(), 1); + assert_eq!(legacy.requests[0], 0..20); }
