Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add a 'prefetch' option to ParquetRecordBatchStream to load the next row group while decoding #6676

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Add a 'prefetch' option to ParquetRecordBatchStream to load the next row group while decoding #6676

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4f682b1
Squashed commit of the following:
masonh22 Oct 27, 2024
d9c71b9
Fix future not waking due to noop_cx
masonh22 Oct 30, 2024
9eb6570
rustfmt and clippy fixes
masonh22 Oct 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions parquet/src/arrow/arrow_reader/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,8 @@ pub struct ArrowReaderBuilder<T> {
pub(crate) limit: Option<usize>,

pub(crate) offset: Option<usize>,

pub(crate) prefetch: bool,
}

impl<T> ArrowReaderBuilder<T> {
Expand All @@ -88,6 +90,7 @@ impl<T> ArrowReaderBuilder<T> {
selection: None,
limit: None,
offset: None,
prefetch: false,
}
}

Expand Down
217 changes: 195 additions & 22 deletions parquet/src/arrow/async_reader/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -413,6 +413,14 @@ impl<T: AsyncFileReader + Send + 'static> ParquetRecordBatchStreamBuilder<T> {
Ok(Some(Sbbf::new(&bitset)))
}

/// For async readers, load data for the next row group while decoding the
/// current row group.
///
/// Defaults to `false`
pub fn with_prefetch(self, prefetch: bool) -> Self {
Self { prefetch, ..self }
}

/// Build a new [`ParquetRecordBatchStream`]
pub fn build(self) -> Result<ParquetRecordBatchStream<T>> {
let num_row_groups = self.metadata.row_groups().len();
Expand Down Expand Up @@ -461,6 +469,8 @@ impl<T: AsyncFileReader + Send + 'static> ParquetRecordBatchStreamBuilder<T> {
row_groups,
projection: self.projection,
selection: self.selection,
prefetch_row_groups: self.prefetch,
next_reader: None,
schema,
reader: Some(reader),
state: StreamState::Init,
Expand Down Expand Up @@ -591,6 +601,8 @@ enum StreamState<T> {
Init,
/// Decoding a batch
Decoding(ParquetRecordBatchReader),
/// Decoding a batch while fetching another row group
Prefetch(ParquetRecordBatchReader, BoxFuture<'static, ReadResult<T>>),
/// Reading data from input
Reading(BoxFuture<'static, ReadResult<T>>),
/// Error
Expand All @@ -602,6 +614,7 @@ impl<T> std::fmt::Debug for StreamState<T> {
match self {
StreamState::Init => write!(f, "StreamState::Init"),
StreamState::Decoding(_) => write!(f, "StreamState::Decoding"),
StreamState::Prefetch(..) => write!(f, "StreamState::Prefetch"),
StreamState::Reading(_) => write!(f, "StreamState::Reading"),
StreamState::Error => write!(f, "StreamState::Error"),
}
Expand All @@ -623,6 +636,11 @@ pub struct ParquetRecordBatchStream<T> {

selection: Option<RowSelection>,

prefetch_row_groups: bool,

/// The next row group to decode if we are prefetching.
next_reader: Option<ParquetRecordBatchReader>,

/// This is an option so it can be moved into a future
reader: Option<ReaderFactory<T>>,

Expand Down Expand Up @@ -651,6 +669,33 @@ impl<T> ParquetRecordBatchStream<T> {
}
}

impl<T> ParquetRecordBatchStream<T>
where
T: AsyncFileReader + 'static,
{
/// Returns a future for reading row group `row_group_idx`.
///
/// Note: this function should only be called in [`StreamState::Init`] and
/// [`StreamState::Decoding`] as this takes [`ParquetRecordBatchStream::reader`]
/// and panics if it does not exist.
fn read_row_group(&mut self, row_group_idx: usize) -> BoxFuture<'static, ReadResult<T>> {
let reader = self.reader.take().expect("lost reader");

let row_count = self.metadata.row_group(row_group_idx).num_rows() as usize;

let selection = self.selection.as_mut().map(|s| s.split_off(row_count));

reader
.read_row_group(
row_group_idx,
selection,
self.projection.clone(),
self.batch_size,
)
.boxed()
}
}

impl<T> Stream for ParquetRecordBatchStream<T>
where
T: AsyncFileReader + Unpin + Send + 'static,
Expand All @@ -660,36 +705,89 @@ where
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
loop {
match &mut self.state {
StreamState::Decoding(batch_reader) => match batch_reader.next() {
Some(Ok(batch)) => {
return Poll::Ready(Some(Ok(batch)));
StreamState::Decoding(batch_reader) => {
let res: Self::Item = match batch_reader.next() {
Some(Ok(batch)) => Ok(batch),
Some(Err(e)) => {
self.state = StreamState::Error;
return Poll::Ready(Some(Err(ParquetError::ArrowError(e.to_string()))));
}
None => {
self.state = StreamState::Init;
continue;
}
};

if !self.prefetch_row_groups
|| self.row_groups.is_empty()
|| self.next_reader.is_some()
{
return Poll::Ready(Some(res));
}
Some(Err(e)) => {
self.state = StreamState::Error;
return Poll::Ready(Some(Err(ParquetError::ArrowError(e.to_string()))));

let old_state = std::mem::replace(&mut self.state, StreamState::Init);

let row_group_idx = self.row_groups.pop_front().unwrap(); // already checked that row_groups is not empty

let fut = self.read_row_group(row_group_idx);

if let StreamState::Decoding(batch_reader) = old_state {
self.state = StreamState::Prefetch(batch_reader, fut);
return Poll::Ready(Some(res));
} else {
unreachable!()
}
None => self.state = StreamState::Init,
},
}
StreamState::Prefetch(batch_reader, f) => {
let mut noop_cx = Context::from_waker(futures::task::noop_waker_ref());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the rationale for doing this?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wanted to avoid any potential overhead from using the real context when polling the future here. Since we're always returning Poll::Ready out of this state (or transitioning to another state), we don't need to rely on the real context to wake the main stream future.

I'm not an expert at async rust code though so if it would make more sense to do something else here I'm happy to make that change.

match f.poll_unpin(&mut noop_cx) {
Poll::Pending => (),
Poll::Ready(Ok((reader_factory, maybe_reader))) => {
let old_state = std::mem::replace(&mut self.state, StreamState::Init);
if let StreamState::Prefetch(batch_reader, _) = old_state {
self.state = StreamState::Decoding(batch_reader);
} else {
unreachable!()
}
self.reader = Some(reader_factory);
self.next_reader = maybe_reader;
continue;
}
Poll::Ready(Err(e)) => {
self.state = StreamState::Error;
return Poll::Ready(Some(Err(e)));
}
}

match batch_reader.next() {
Some(Ok(batch)) => return Poll::Ready(Some(Ok(batch))),
Some(Err(e)) => {
self.state = StreamState::Error;
return Poll::Ready(Some(Err(ParquetError::ArrowError(e.to_string()))));
}
None => {
let old_state = std::mem::replace(&mut self.state, StreamState::Init);
if let StreamState::Prefetch(_, f) = old_state {
self.state = StreamState::Reading(f);
continue;
} else {
unreachable!()
}
}
}
}
StreamState::Init => {
if let Some(batch_reader) = self.next_reader.take() {
self.state = StreamState::Decoding(batch_reader);
continue;
}

let row_group_idx = match self.row_groups.pop_front() {
Some(idx) => idx,
None => return Poll::Ready(None),
};

let reader = self.reader.take().expect("lost reader");

let row_count = self.metadata.row_group(row_group_idx).num_rows() as usize;

let selection = self.selection.as_mut().map(|s| s.split_off(row_count));

let fut = reader
.read_row_group(
row_group_idx,
selection,
self.projection.clone(),
self.batch_size,
)
.boxed();
let fut = self.read_row_group(row_group_idx);

self.state = StreamState::Reading(fut)
}
Expand Down Expand Up @@ -2037,4 +2135,79 @@ mod tests {
// Should only have made 3 requests
assert_eq!(requests.lock().unwrap().len(), 3);
}

#[tokio::test]
async fn test_reader_prefetch() {
let a = StringArray::from_iter_values(["a", "b", "b", "b", "c", "c"]);
let b = StringArray::from_iter_values(["1", "2", "3", "4", "5", "6"]);
let c = Int32Array::from_iter(0..6);
let data = RecordBatch::try_from_iter([
("a", Arc::new(a) as ArrayRef),
("b", Arc::new(b) as ArrayRef),
("c", Arc::new(c) as ArrayRef),
])
.unwrap();

let mut buf = Vec::with_capacity(1024);
let props = WriterProperties::builder()
.set_max_row_group_size(3)
.build();
let mut writer = ArrowWriter::try_new(&mut buf, data.schema(), Some(props)).unwrap();
writer.write(&data).unwrap();
writer.close().unwrap();

let data: Bytes = buf.into();
let metadata = ParquetMetaDataReader::new()
.parse_and_finish(&data)
.unwrap();

assert_eq!(metadata.num_row_groups(), 2);

let test = TestReader {
data,
metadata: Arc::new(metadata),
requests: Default::default(),
};

let mut stream = ParquetRecordBatchStreamBuilder::new(test.clone())
.await
.unwrap()
.with_batch_size(1)
.with_limit(5)
.with_prefetch(true)
.build()
.unwrap();

let batch1 = stream.try_next().await.unwrap().unwrap();
// Each batch should only have one row
assert_eq!(batch1.num_rows(), 1);
// Make sure we are pre-fetching
assert!(matches!(stream.state, StreamState::Prefetch(..)));

let batch2 = stream.try_next().await.unwrap().unwrap();
assert_eq!(batch2.num_rows(), 1);
// We should no longer be prefetching...
assert!(matches!(stream.state, StreamState::Decoding(..)));
// because we already fetched the next row group.
assert!(stream.next_reader.is_some());

let batch3 = stream.try_next().await.unwrap().unwrap();
assert_eq!(batch3.num_rows(), 1);
assert!(matches!(stream.state, StreamState::Decoding(..)));
assert!(stream.next_reader.is_some());

let batch4 = stream.try_next().await.unwrap().unwrap();
assert_eq!(batch4.num_rows(), 1);
assert!(matches!(stream.state, StreamState::Decoding(..)));
assert!(stream.next_reader.is_none());

let batch5 = stream.try_next().await.unwrap().unwrap();
assert_eq!(batch5.num_rows(), 1);
assert!(matches!(stream.state, StreamState::Decoding(..)));
assert!(stream.next_reader.is_none());

let batch6 = stream.try_next().await.unwrap();
assert!(batch6.is_none());
assert!(stream.next_reader.is_none());
}
}