alamb commented on code in PR #11975:
URL: https://github.com/apache/arrow-rs/pull/11975#discussion_r4457558357
##########
parquet/tests/arrow_reader/bloom_filter/async.rs:
##########
@@ -166,111 +153,6 @@ async fn test_get_row_group_column_bloom_filter(data:
Bytes, with_length: bool)
);
}
-#[derive(Default)]
Review Comment:
❤️
##########
parquet/src/bloom_filter/mod.rs:
##########
@@ -536,66 +487,25 @@ impl Sbbf {
column_metadata: &ColumnChunkMetaData,
reader: &mut R,
) -> Result<Option<Self>> {
- let Some(offset) = column_metadata.bloom_filter_offset() else {
+ let Some(mut decoder) = SbbfPushDecoder::try_new(column_metadata)?
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"
- ));
+ loop {
+ let result = decoder.try_decode()?;
Review Comment:
this is so nice
##########
parquet/src/bloom_filter/mod.rs:
##########
@@ -875,31 +827,385 @@ async fn read_bloom_filter_bytes_async<R:
AsyncFileReader + ?Sized>(
Ok(bytes)
}
-// per spec we use xxHash with seed=0
-const SEED: u64 = 0;
+/// Parses a Bloom filter header and returns the number of bytes consumed.
+fn read_bloom_filter_header_and_length(buffer: &[u8]) ->
Result<(BloomFilterHeader, usize)> {
+ let mut prot = ThriftSliceInputProtocol::new(buffer);
+ let header = BloomFilterHeader::read_thrift(&mut prot)
+ .map_err(|e| ParquetError::General(format!("Could not read bloom
filter header: {e}")))?;
+ Ok((header, buffer.len() - prot.as_slice().len()))
+}
-#[inline]
-fn hash_as_bytes<A: AsBytes + ?Sized>(value: &A) -> u64 {
- XxHash64::oneshot(SEED, value.as_bytes())
+/// Decoding state machine
+#[derive(Debug)]
+enum DecodeState {
+ /// Finding the first Bloom filter range
+ FindingRange,
+ // Reading the SBBF header
+ ReadingSbbfHeader(Range<u64>),
+ // Reading the SBBF bitset bytes
+ ReadingSbbfBitset(Range<u64>),
+ // Decoding is complete
+ Finished,
+ /// State left during the `try_decode` method so something valid is
present.
+ /// This state should never be observed.
+ Intermediate,
+}
+
+/// Decodes a Bloom filter from caller-provided file ranges, without
performing I/O.
+///
+/// The metadata and pushed bytes must refer to the same file. Ranges use
absolute
+/// file offsets. Call [`Self::try_decode`] to discover which ranges are
needed,
+/// then supply them using [`Self::push_range`] or [`Self::push_ranges`].
Bytes may
+/// also be pushed in advance to avoid requests for already-buffered data.
+///
+/// Each requested range must be contained in a single pushed buffer; adjacent
+/// buffers are not coalesced. After an error from [`Self::try_decode`],
discard
+/// the decoder.
+///
+/// # Example
+///
+/// ```no_run
+/// use bytes::Bytes;
+/// use parquet::DecodeResult;
+/// use parquet::bloom_filter::{Sbbf, SbbfPushDecoder};
+/// use parquet::errors::Result;
+/// use parquet::file::metadata::ColumnChunkMetaData;
+///
+/// fn decode_filter(metadata: &ColumnChunkMetaData, file_bytes: Bytes) ->
Result<Option<Sbbf>> {
+/// let Some(mut decoder) = SbbfPushDecoder::try_new(metadata)? else {
+/// return Ok(None);
+/// };
+/// loop {
+/// match decoder.try_decode()? {
+/// DecodeResult::NeedsData(ranges) => {
+/// let buffers = ranges.iter().map(|range| {
+/// file_bytes.slice(range.start as usize..range.end as
usize)
+/// }).collect();
+/// decoder.push_ranges(ranges, buffers)?;
+/// }
+/// DecodeResult::Data(filter) => return Ok(Some(filter)),
+/// DecodeResult::Finished => unreachable!("filter was returned
above"),
+/// }
+/// }
+/// }
+/// ```
+pub struct SbbfPushDecoder {
+ /// Decoding state
+ state: DecodeState,
+ /// Offset of the Bloom filter in the file
+ offset: u64,
+ /// Length of the Bloom filter
+ length: Option<usize>,
+ /// Underlying buffers
+ buffers: PushBuffers,
+}
+
+impl SbbfPushDecoder {
+ /// Create a new Bloom filter decoder for the given column chunk.
+ ///
+ /// Returns `None` if the column chunk has no Bloom filter offset.
+ /// Invalid offsets and nonpositive advertised lengths return an error.
+ pub fn try_new(column_metadata: &ColumnChunkMetaData) ->
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()?;
+ Ok(Some(Self {
+ state: DecodeState::FindingRange,
+ offset,
+ length,
+ buffers: PushBuffers::default(),
+ }))
+ }
+
+ /// Pushes a file range into the decoder's buffer.
+ ///
+ /// Returns an error if the buffer length differs from the range length or
+ /// decoding has finished.
+ pub fn push_range(&mut self, range: Range<u64>, buffer: Bytes) ->
Result<()> {
+ if matches!(&self.state, DecodeState::Finished) {
+ return Err(general_err!(
+ "SbbfPushDecoder: cannot push data after decoding is finished"
+ ));
+ }
+ self.buffers.push_range(range, buffer)?;
+ Ok(())
+ }
+
+ /// Pushes multiple file ranges into the decoder's buffer.
+ ///
+ /// The number of ranges must match the number of buffers, and each buffer
+ /// must have the same length as its range. Pushing after decoding has
+ /// finished returns an error.
+ pub fn push_ranges(&mut self, ranges: Vec<Range<u64>>, buffers:
Vec<Bytes>) -> Result<()> {
+ if matches!(&self.state, DecodeState::Finished) {
+ return Err(general_err!(
+ "SbbfPushDecoder: cannot push data after decoding is finished"
+ ));
+ }
+ self.buffers.push_ranges(ranges, buffers)?;
+ Ok(())
+ }
+
+ /// Returns the bytes for the given range from the internal buffer
+ fn get_bytes(&self, range: &Range<u64>) -> Result<Bytes> {
+ let start = range.start;
+ let raw_len = range.end - range.start;
+ let len: usize = raw_len.try_into().map_err(|_| {
+ ParquetError::General(format!(
+ "SbbfPushDecoder: Range length too large to fit in usize:
{raw_len}",
+ ))
+ })?;
+ self.buffers.get_bytes(start, len)
+ }
+
+ /// Attempts to decode the filter, returning the file ranges needed if data
+ /// is missing. Repeated calls before supplying data return the same
ranges.
+ ///
+ /// Without prefetched data, requests one range when `bloom_filter_length`
is
+ /// known, or two otherwise (a 20-byte header estimate, then the bitset).
+ /// Headers larger than the estimate require an advertised length.
+ /// After returning [`DecodeResult::Data`], subsequent calls return
+ /// [`DecodeResult::Finished`].
+ ///
+ /// The state transitions can be visualized as:
+ /// ```text
+ /// ┌──────────────┐
+ /// │ FindingRange │
+ /// └──────────────┘
+ /// │
+ /// choose initial range
+ /// │
+ /// ▼
+ /// ┌───────────────────┐ ───┐ bytes missing:
+ /// │ ReadingSbbfHeader │ │ return NeedsData(range),
+ /// └───────────────────┘ ◀──┘ wait for push_range(s)
+ /// │ │
+ /// length known length unknown
+ /// │ │
+ /// │ ▼
+ /// │ ┌───────────────────┐ ───┐ bytes missing:
+ /// │ │ ReadingSbbfBitset │ │ return
NeedsData(range),
+ /// │ └───────────────────┘ ◀──┘ wait for push_range(s)
+ /// │ │
+ /// └───────┬───────┘
+ /// │
+ /// return Data(Sbbf)
+ /// │
+ /// ▼
+ /// ┌──────────┐ ───┐
+ /// │ Finished │ │ return Finished
+ /// └──────────┘ ◀──┘
+ /// ```
+ ///
+ /// In the above:
+ /// - The initial range is the whole filter when its length is known,
otherwise a 20-byte header prefix.
+ /// - Transitions proceed within the same `try_decode()` call while the
required bytes are buffered;
+ /// push_range(s) supplies bytes without changing state.
+ /// - Header parsing/validation must succeed before advancing. Any
decoding error leaves the decoder unusable.
+ pub fn try_decode(&mut self) -> Result<DecodeResult<Sbbf>> {
+ loop {
+ match std::mem::replace(&mut self.state,
DecodeState::Intermediate) {
+ DecodeState::FindingRange => {
+ let range = bloom_filter_read_range(
+ self.offset,
+ 0..self.length.unwrap_or(SBBF_HEADER_SIZE_ESTIMATE),
+ )?;
+ self.state = DecodeState::ReadingSbbfHeader(range);
+ }
+ DecodeState::ReadingSbbfHeader(range) => {
+ if !self.buffers.has_range(&range) {
+ self.state =
DecodeState::ReadingSbbfHeader(range.clone());
+ return Ok(DecodeResult::NeedsData(vec![range]));
+ }
+ let buffer = self.get_bytes(&range)?;
+ let bitset_range = bloom_filter_bitset_range(&buffer,
self.length)?;
+ if self.length.is_some() {
+ let bitset = buffer.slice(bitset_range);
+ self.state = DecodeState::Finished;
+ return Ok(DecodeResult::Data(Sbbf::new(&bitset)));
+ }
+ let range = bloom_filter_read_range(self.offset,
bitset_range)?;
+ self.state = DecodeState::ReadingSbbfBitset(range);
+ }
+ DecodeState::ReadingSbbfBitset(range) => {
+ if !self.buffers.has_range(&range) {
+ self.state =
DecodeState::ReadingSbbfBitset(range.clone());
+ return Ok(DecodeResult::NeedsData(vec![range]));
+ }
+ let bitset = self.get_bytes(&range)?;
+ self.state = DecodeState::Finished;
+ return Ok(DecodeResult::Data(Sbbf::new(&bitset)));
+ }
+ DecodeState::Finished => {
+ self.state = DecodeState::Finished;
+ return Ok(DecodeResult::Finished);
+ }
+ DecodeState::Intermediate => {
+ return Err(general_err!(
+ "SbbfPushDecoder: internal error, invalid state"
+ ));
+ }
+ }
+ }
+ }
}
#[cfg(test)]
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]
Review Comment:
it isn't clear to me why these tests are not with the others in
`parquet/tests/arrow_reader/bloom_filter/push.rs` -- as they all seem to be
testing the public APIs
I recommend consolidating to parquet/tests/arrow_reader/bloom_filter/push.rs
##########
parquet/src/bloom_filter/mod.rs:
##########
@@ -875,31 +827,349 @@ async fn read_bloom_filter_bytes_async<R:
AsyncFileReader + ?Sized>(
Ok(bytes)
}
-// per spec we use xxHash with seed=0
-const SEED: u64 = 0;
+/// Parses a Bloom filter header and returns the number of bytes consumed.
+fn read_bloom_filter_header_and_length(buffer: &[u8]) ->
Result<(BloomFilterHeader, usize)> {
+ let mut prot = ThriftSliceInputProtocol::new(buffer);
+ let header = BloomFilterHeader::read_thrift(&mut prot)
+ .map_err(|e| ParquetError::General(format!("Could not read bloom
filter header: {e}")))?;
+ Ok((header, buffer.len() - prot.as_slice().len()))
+}
-#[inline]
-fn hash_as_bytes<A: AsBytes + ?Sized>(value: &A) -> u64 {
- XxHash64::oneshot(SEED, value.as_bytes())
+/// Decoding state machine
+#[derive(Debug)]
+enum DecodeState {
+ /// Finding the first Bloom filter range
+ FindingRange,
+ // Reading the SBBF header
+ ReadingSbbfHeader(Range<u64>),
+ // Reading the SBBF bitset bytes
+ ReadingSbbfBitset(Range<u64>),
+ // Decoding is complete
+ Finished,
+ /// State left during the `try_decode` method so something valid is
present.
+ /// This state should never be observed.
+ Intermediate,
+}
+
+/// Decodes a Bloom filter from caller-provided file ranges, without
performing I/O.
+///
+/// The metadata and pushed bytes must refer to the same file. Ranges use
absolute
+/// file offsets. Call [`Self::try_decode`] to discover which ranges are
needed,
+/// then supply them using [`Self::push_range`] or [`Self::push_ranges`].
Bytes may
+/// also be pushed in advance to avoid requests for already-buffered data.
+///
+/// Each requested range must be contained in a single pushed buffer; adjacent
+/// buffers are not coalesced. After an error from [`Self::try_decode`],
discard
+/// the decoder.
+///
+/// # Example
+///
+/// ```no_run
+/// use bytes::Bytes;
+/// use parquet::DecodeResult;
+/// use parquet::bloom_filter::{Sbbf, SbbfPushDecoder};
+/// use parquet::errors::Result;
+/// use parquet::file::metadata::ColumnChunkMetaData;
+///
+/// fn decode_filter(metadata: &ColumnChunkMetaData, file_bytes: Bytes) ->
Result<Option<Sbbf>> {
+/// let Some(mut decoder) = SbbfPushDecoder::try_new(metadata)? else {
+/// return Ok(None);
+/// };
+/// loop {
+/// match decoder.try_decode()? {
+/// DecodeResult::NeedsData(ranges) => {
+/// let buffers = ranges.iter().map(|range| {
+/// file_bytes.slice(range.start as usize..range.end as
usize)
+/// }).collect();
+/// decoder.push_ranges(ranges, buffers)?;
+/// }
+/// DecodeResult::Data(filter) => return Ok(Some(filter)),
+/// DecodeResult::Finished => unreachable!("filter was returned
above"),
+/// }
+/// }
+/// }
+/// ```
+pub struct SbbfPushDecoder {
+ /// Decoding state
+ state: DecodeState,
+ /// Offset of the Bloom filter in the file
+ offset: u64,
+ /// Length of the Bloom filter
+ length: Option<usize>,
+ /// Underlying buffers
+ buffers: PushBuffers,
+}
+
+impl SbbfPushDecoder {
+ /// Create a new Bloom filter decoder for the given column chunk.
+ ///
+ /// Returns `None` if the column chunk has no Bloom filter offset.
+ /// Invalid offsets and nonpositive advertised lengths return an error.
+ pub fn try_new(column_metadata: &ColumnChunkMetaData) ->
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()?;
+ Ok(Some(Self {
+ state: DecodeState::FindingRange,
+ offset,
+ length,
+ buffers: PushBuffers::default(),
+ }))
+ }
+
+ /// Pushes a file range into the decoder's buffer.
+ ///
+ /// Returns an error if the buffer length differs from the range length or
+ /// decoding has finished.
+ pub fn push_range(&mut self, range: Range<u64>, buffer: Bytes) ->
Result<()> {
+ if matches!(&self.state, DecodeState::Finished) {
+ return Err(general_err!(
+ "SbbfPushDecoder: cannot push data after decoding is finished"
+ ));
+ }
+ self.buffers.push_range(range, buffer)?;
+ Ok(())
+ }
+
+ /// Pushes multiple file ranges into the decoder's buffer.
+ ///
+ /// The number of ranges must match the number of buffers, and each buffer
+ /// must have the same length as its range. Pushing after decoding has
+ /// finished returns an error.
+ pub fn push_ranges(&mut self, ranges: Vec<Range<u64>>, buffers:
Vec<Bytes>) -> Result<()> {
+ if matches!(&self.state, DecodeState::Finished) {
+ return Err(general_err!(
+ "SbbfPushDecoder: cannot push data after decoding is finished"
+ ));
+ }
+ self.buffers.push_ranges(ranges, buffers)?;
+ Ok(())
+ }
+
+ /// Returns the bytes for the given range from the internal buffer
+ fn get_bytes(&self, range: &Range<u64>) -> Result<Bytes> {
+ let start = range.start;
+ let raw_len = range.end - range.start;
+ let len: usize = raw_len.try_into().map_err(|_| {
+ ParquetError::General(format!(
+ "SbbfPushDecoder: Range length too large to fit in usize:
{raw_len}",
+ ))
+ })?;
+ self.buffers.get_bytes(start, len)
+ }
+
+ /// Attempts to decode the filter, returning the file ranges needed if data
+ /// is missing. Repeated calls before supplying data return the same
ranges.
+ ///
+ /// Without prefetched data, requests one range when `bloom_filter_length`
is
+ /// known, or two otherwise (a 20-byte header estimate, then the bitset).
+ /// Headers larger than the estimate require an advertised length.
+ /// After returning [`DecodeResult::Data`], subsequent calls return
+ /// [`DecodeResult::Finished`].
+ pub fn try_decode(&mut self) -> Result<DecodeResult<Sbbf>> {
Review Comment:
I do enjoy making ascii art. Though recently I had a 🤯 experience of having
my agent reverse engineer the monodraw file format so I can have it draft ascii
art diagrams and then tweak them visually ❤️
##########
parquet/src/bloom_filter/mod.rs:
##########
@@ -875,31 +827,385 @@ async fn read_bloom_filter_bytes_async<R:
AsyncFileReader + ?Sized>(
Ok(bytes)
}
-// per spec we use xxHash with seed=0
-const SEED: u64 = 0;
+/// Parses a Bloom filter header and returns the number of bytes consumed.
+fn read_bloom_filter_header_and_length(buffer: &[u8]) ->
Result<(BloomFilterHeader, usize)> {
+ let mut prot = ThriftSliceInputProtocol::new(buffer);
+ let header = BloomFilterHeader::read_thrift(&mut prot)
+ .map_err(|e| ParquetError::General(format!("Could not read bloom
filter header: {e}")))?;
+ Ok((header, buffer.len() - prot.as_slice().len()))
+}
-#[inline]
-fn hash_as_bytes<A: AsBytes + ?Sized>(value: &A) -> u64 {
- XxHash64::oneshot(SEED, value.as_bytes())
+/// Decoding state machine
Review Comment:
as a minor nit I might move this into its own mod like bloom_filter/push.rs
or something, but this is totally fine too
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
