parquet/bloom_filter/push.rs
1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Push-based decoding of serialized Parquet Bloom filters.
6
7use std::ops::Range;
8
9use bytes::Bytes;
10
11use super::{SBBF_HEADER_SIZE_ESTIMATE, Sbbf, bloom_filter_bitset_range, bloom_filter_read_range};
12use crate::DecodeResult;
13use crate::errors::{ParquetError, Result};
14use crate::file::metadata::ColumnChunkMetaData;
15use crate::file::reader::ChunkReader;
16use crate::util::push_buffers::PushBuffers;
17
18/// Decoding state machine
19#[derive(Debug)]
20enum DecodeState {
21 /// Finding the first Bloom filter range
22 FindingRange,
23 // Reading the SBBF header
24 ReadingSbbfHeader(Range<u64>),
25 // Reading the SBBF bitset bytes
26 ReadingSbbfBitset(Range<u64>),
27 // Decoding is complete
28 Finished,
29 /// State left during the `try_decode` method so something valid is present.
30 /// This state should never be observed.
31 Intermediate,
32}
33
34/// Decodes a Bloom filter from caller-provided file ranges, without performing I/O.
35///
36/// The metadata and pushed bytes must refer to the same file. Ranges use absolute
37/// file offsets. Call [`Self::try_decode`] to discover which ranges are needed,
38/// then supply them using [`Self::push_range`] or [`Self::push_ranges`]. Bytes may
39/// also be pushed in advance to avoid requests for already-buffered data.
40///
41/// Each requested range must be contained in a single pushed buffer; adjacent
42/// buffers are not coalesced. After an error from [`Self::try_decode`], discard
43/// the decoder.
44///
45/// # Example
46///
47/// ```rust
48/// use bytes::Bytes;
49/// use parquet::DecodeResult;
50/// use parquet::bloom_filter::{Sbbf, SbbfPushDecoder};
51/// use parquet::errors::Result;
52/// use parquet::file::metadata::ColumnChunkMetaData;
53///
54/// fn decode_filter(metadata: &ColumnChunkMetaData, file_bytes: Bytes) -> Result<Option<Sbbf>> {
55/// let Some(mut decoder) = SbbfPushDecoder::try_new(metadata)? else {
56/// return Ok(None);
57/// };
58/// loop {
59/// match decoder.try_decode()? {
60/// DecodeResult::NeedsData(ranges) => {
61/// let buffers = ranges.iter().map(|range| {
62/// file_bytes.slice(range.start as usize..range.end as usize)
63/// }).collect();
64/// decoder.push_ranges(ranges, buffers)?;
65/// }
66/// DecodeResult::Data(filter) => return Ok(Some(filter)),
67/// DecodeResult::Finished => unreachable!("filter was returned above"),
68/// }
69/// }
70/// }
71///
72/// # fn main() -> Result<()> {
73/// # use std::sync::Arc;
74/// # use parquet::schema::parser::parse_message_type;
75/// # use parquet::schema::types::SchemaDescriptor;
76/// # let schema = Arc::new(parse_message_type("message schema { REQUIRED BYTE_ARRAY value; }")?);
77/// # let descriptor = SchemaDescriptor::new(schema).column(0);
78/// # let mut original = Sbbf::new_with_num_of_bytes(32);
79/// # original.insert("present");
80/// # let offset = 13;
81/// # let mut serialized = vec![0; offset];
82/// # original.write(&mut serialized)?;
83/// # let file_bytes = Bytes::from(serialized);
84/// # let metadata = ColumnChunkMetaData::builder(descriptor)
85/// # .set_bloom_filter_offset(Some(offset as i64))
86/// # .set_bloom_filter_length(Some((file_bytes.len() - offset) as i32))
87/// # .build()?;
88/// let filter = decode_filter(&metadata, file_bytes)?.unwrap();
89/// assert!(filter.check("present"));
90/// # Ok(())
91/// # }
92/// ```
93pub struct SbbfPushDecoder {
94 /// Decoding state
95 state: DecodeState,
96 /// Offset of the Bloom filter in the file
97 offset: u64,
98 /// Length of the Bloom filter
99 length: Option<usize>,
100 /// Underlying buffers
101 buffers: PushBuffers,
102}
103
104impl SbbfPushDecoder {
105 /// Create a new Bloom filter decoder for the given column chunk.
106 ///
107 /// Returns `None` if the column chunk has no Bloom filter offset.
108 /// Invalid offsets and nonpositive advertised lengths return an error.
109 pub fn try_new(column_metadata: &ColumnChunkMetaData) -> Result<Option<Self>> {
110 let Some(offset) = column_metadata.bloom_filter_offset() else {
111 return Ok(None);
112 };
113 let offset =
114 u64::try_from(offset).map_err(|_| general_err!("Bloom filter offset is invalid"))?;
115 let length = column_metadata
116 .bloom_filter_length()
117 .map(|length| {
118 usize::try_from(length)
119 .ok()
120 .filter(|length| *length > 0)
121 .ok_or_else(|| general_err!("Bloom filter length is invalid"))
122 })
123 .transpose()?;
124 Ok(Some(Self {
125 state: DecodeState::FindingRange,
126 offset,
127 length,
128 buffers: PushBuffers::default(),
129 }))
130 }
131
132 /// Pushes a file range into the decoder's buffer.
133 ///
134 /// Returns an error if the buffer length differs from the range length or
135 /// decoding has finished.
136 pub fn push_range(&mut self, range: Range<u64>, buffer: Bytes) -> Result<()> {
137 if matches!(&self.state, DecodeState::Finished) {
138 return Err(general_err!(
139 "SbbfPushDecoder: cannot push data after decoding is finished"
140 ));
141 }
142 self.buffers.push_range(range, buffer)?;
143 Ok(())
144 }
145
146 /// Pushes multiple file ranges into the decoder's buffer.
147 ///
148 /// The number of ranges must match the number of buffers, and each buffer
149 /// must have the same length as its range. Pushing after decoding has
150 /// finished returns an error.
151 pub fn push_ranges(&mut self, ranges: Vec<Range<u64>>, buffers: Vec<Bytes>) -> Result<()> {
152 if matches!(&self.state, DecodeState::Finished) {
153 return Err(general_err!(
154 "SbbfPushDecoder: cannot push data after decoding is finished"
155 ));
156 }
157 self.buffers.push_ranges(ranges, buffers)?;
158 Ok(())
159 }
160
161 /// Returns the bytes for the given range from the internal buffer
162 fn get_bytes(&self, range: &Range<u64>) -> Result<Bytes> {
163 let start = range.start;
164 let raw_len = range.end - range.start;
165 let len: usize = raw_len.try_into().map_err(|_| {
166 ParquetError::General(format!(
167 "SbbfPushDecoder: Range length too large to fit in usize: {raw_len}",
168 ))
169 })?;
170 self.buffers.get_bytes(start, len)
171 }
172
173 /// Attempts to decode the filter, returning the file ranges needed if data
174 /// is missing. Repeated calls before supplying data return the same ranges.
175 ///
176 /// Without prefetched data, requests one range when `bloom_filter_length` is
177 /// known, or two otherwise (a 20-byte header estimate, then the bitset).
178 /// Headers larger than the estimate require an advertised length.
179 /// After returning [`DecodeResult::Data`], subsequent calls return
180 /// [`DecodeResult::Finished`].
181 ///
182 /// The state transitions can be visualized as:
183 /// ```text
184 /// ┌──────────────┐
185 /// │ FindingRange │
186 /// └──────────────┘
187 /// │
188 /// choose initial range
189 /// │
190 /// ▼
191 /// ┌───────────────────┐ ───┐ bytes missing:
192 /// │ ReadingSbbfHeader │ │ return NeedsData(range),
193 /// └───────────────────┘ ◀──┘ wait for push_range(s)
194 /// │ │
195 /// length known length unknown
196 /// │ │
197 /// │ ▼
198 /// │ ┌───────────────────┐ ───┐ bytes missing:
199 /// │ │ ReadingSbbfBitset │ │ return NeedsData(range),
200 /// │ └───────────────────┘ ◀──┘ wait for push_range(s)
201 /// │ │
202 /// └───────┬───────┘
203 /// │
204 /// return Data(Sbbf)
205 /// │
206 /// ▼
207 /// ┌──────────┐ ───┐
208 /// │ Finished │ │ return Finished
209 /// └──────────┘ ◀──┘
210 /// ```
211 ///
212 /// In the above:
213 /// - The initial range is the whole filter when its length is known, otherwise a 20-byte header prefix.
214 /// - Transitions proceed within the same `try_decode()` call while the required bytes are buffered;
215 /// push_range(s) supplies bytes without changing state.
216 /// - Header parsing/validation must succeed before advancing. Any decoding error leaves the decoder unusable.
217 pub fn try_decode(&mut self) -> Result<DecodeResult<Sbbf>> {
218 loop {
219 match std::mem::replace(&mut self.state, DecodeState::Intermediate) {
220 DecodeState::FindingRange => {
221 let range = bloom_filter_read_range(
222 self.offset,
223 0..self.length.unwrap_or(SBBF_HEADER_SIZE_ESTIMATE),
224 )?;
225 self.state = DecodeState::ReadingSbbfHeader(range);
226 }
227 DecodeState::ReadingSbbfHeader(range) => {
228 if !self.buffers.has_range(&range) {
229 self.state = DecodeState::ReadingSbbfHeader(range.clone());
230 return Ok(DecodeResult::NeedsData(vec![range]));
231 }
232 let buffer = self.get_bytes(&range)?;
233 let bitset_range = bloom_filter_bitset_range(&buffer, self.length)?;
234 if self.length.is_some() {
235 let bitset = buffer.slice(bitset_range);
236 self.state = DecodeState::Finished;
237 return Ok(DecodeResult::Data(Sbbf::new(&bitset)));
238 }
239 let range = bloom_filter_read_range(self.offset, bitset_range)?;
240 self.state = DecodeState::ReadingSbbfBitset(range);
241 }
242 DecodeState::ReadingSbbfBitset(range) => {
243 if !self.buffers.has_range(&range) {
244 self.state = DecodeState::ReadingSbbfBitset(range.clone());
245 return Ok(DecodeResult::NeedsData(vec![range]));
246 }
247 let bitset = self.get_bytes(&range)?;
248 self.state = DecodeState::Finished;
249 return Ok(DecodeResult::Data(Sbbf::new(&bitset)));
250 }
251 DecodeState::Finished => {
252 self.state = DecodeState::Finished;
253 return Ok(DecodeResult::Finished);
254 }
255 DecodeState::Intermediate => {
256 return Err(general_err!(
257 "SbbfPushDecoder: internal error, invalid state"
258 ));
259 }
260 }
261 }
262 }
263}