adriangb commented on code in PR #11975:
URL: https://github.com/apache/arrow-rs/pull/11975#discussion_r4448361076
##########
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:
Would it be helpful to have an agent generate a state transition diagram so
readers can easily visualize the code?
##########
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,
+}
Review Comment:
🥳
--
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]
