alamb commented on code in PR #10702:
URL: https://github.com/apache/arrow-rs/pull/10702#discussion_r3806804500
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -64,6 +64,109 @@ pub mod statistics;
/// Default batch size for reading parquet files
pub const DEFAULT_BATCH_SIZE: usize = 1024;
+/// A row group and its optional row-group-local [`RowSelection`].
+///
+/// A `None` selection reads the entire row group. Omitting a row group skips
+/// it. Entries are decoded in the supplied order.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RowGroupSelection {
+ pub(crate) row_group_index: usize,
+ pub(crate) selection: Option<RowSelection>,
+}
+
+impl RowGroupSelection {
+ /// Creates a row-group-local selection.
+ pub fn new(row_group_index: usize, selection: Option<RowSelection>) ->
Self {
+ Self {
+ row_group_index,
+ selection,
+ }
+ }
+
+ /// The index of the row group this selection applies to.
+ pub fn row_group_index(&self) -> usize {
+ self.row_group_index
+ }
+
+ /// The row-group-local selection, or `None` if the entire row group is
+ /// read.
+ pub fn selection(&self) -> Option<&RowSelection> {
+ self.selection.as_ref()
+ }
+}
+
+/// Row-selection configuration shared by the Arrow reader builders.
+///
+/// `Global` is the legacy configuration formed by `with_row_groups` and
Review Comment:
nit: it would be nice to put these docs directly on the enum variants, along
with a description of what they mean, not just where they came from / what they
are used for
Something like this, perhaps:
```rust
/// Row-selection configuration shared by the Arrow reader builders.
pub(crate) enum RowGroupPlan {
/// (IS THIS TRUE?) Selection: first select any row_groups, adn then
apply the RowSelection to
/// any remaining rows.
/// formed by `with_row_groups` and ....)
Global {
...
}
```
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -64,6 +64,109 @@ pub mod statistics;
/// Default batch size for reading parquet files
pub const DEFAULT_BATCH_SIZE: usize = 1024;
+/// A row group and its optional row-group-local [`RowSelection`].
+///
+/// A `None` selection reads the entire row group. Omitting a row group skips
+/// it. Entries are decoded in the supplied order.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RowGroupSelection {
+ pub(crate) row_group_index: usize,
+ pub(crate) selection: Option<RowSelection>,
+}
+
+impl RowGroupSelection {
+ /// Creates a row-group-local selection.
Review Comment:
it might be good to define what `row-group-local` means more precisely --
specifically something like "The selection is relative to the rows in that
selection (e.g. row 100 means the 100th row in the row group)"
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -130,14 +233,12 @@ pub struct ArrowReaderBuilder<T> {
pub(crate) batch_size: usize,
- pub(crate) row_groups: Option<Vec<usize>>,
+ pub(crate) row_group_plan: RowGroupPlan,
Review Comment:
👍
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -219,11 +349,15 @@ impl<T> ArrowReaderBuilder<T> {
/// Only read data from the provided row group indexes
///
/// This is also called row group filtering
- pub fn with_row_groups(self, row_groups: Vec<usize>) -> Self {
- Self {
- row_groups: Some(row_groups),
- ..self
- }
+ ///
Review Comment:
If we are going to keep this wording I suggest making `async stream
builders` a link too -- something like
```diff
- /// On [`ParquetPushDecoderBuilder`] and async stream builders, which
- /// additionally offer `with_row_group_selections`, this cannot be
combined
- /// with that method; attempting to do so returns an error from `build`.
+ /// On [`ParquetPushDecoderBuilder`] and
[`ParquetRecordBatchStreamBuilder`],
+ /// which additionally offer `with_row_group_selections`, this cannot be
+ /// combined with that method; attempting to do so returns an error from
+ /// `build`.
///
/// [`ParquetPushDecoderBuilder`]:
crate::arrow::push_decoder::ParquetPushDecoderBuilder
+ /// [`ParquetRecordBatchStreamBuilder`]:
crate::arrow::async_reader::ParquetRecordBatchStreamBuilder
```
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -151,16 +252,43 @@ pub struct ArrowReaderBuilder<T> {
impl<T: Debug> Debug for ArrowReaderBuilder<T> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
- f.debug_struct("ArrowReaderBuilder<T>")
+ // `row_groups` and `selection` straddle `projection`/`filter` so that
the
+ // legacy configuration keeps its historical field order; only the
+ // row-group-local fields are new.
+ let (row_groups, selection, row_group_selections) = match
&self.row_group_plan {
Review Comment:
this feels overly complicated. Why not derive `Debug` for RowGroupPlan and
display it as normal here
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -219,11 +349,15 @@ impl<T> ArrowReaderBuilder<T> {
/// Only read data from the provided row group indexes
///
/// This is also called row group filtering
- pub fn with_row_groups(self, row_groups: Vec<usize>) -> Self {
- Self {
- row_groups: Some(row_groups),
- ..self
- }
+ ///
+ /// On [`ParquetPushDecoderBuilder`] and async stream builders, which
+ /// additionally offer `with_row_group_selections`, this cannot be combined
+ /// with that method; attempting to do so returns an error from `build`.
Review Comment:
It may be simpler here to just focus on what is not compatible rather than
trying to spell out the cases when it might go wrong (which the compiler will
prevent)
Something like
```rust
/// Can not be combined with [`Self::with_row_group_selections`]. Doing so
/// will return an error at [`build()`]
```
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -219,11 +349,15 @@ impl<T> ArrowReaderBuilder<T> {
/// Only read data from the provided row group indexes
///
/// This is also called row group filtering
- pub fn with_row_groups(self, row_groups: Vec<usize>) -> Self {
- Self {
- row_groups: Some(row_groups),
- ..self
- }
+ ///
+ /// On [`ParquetPushDecoderBuilder`] and async stream builders, which
+ /// additionally offer `with_row_group_selections`, this cannot be combined
+ /// with that method; attempting to do so returns an error from `build`.
+ ///
+ /// [`ParquetPushDecoderBuilder`]:
crate::arrow::push_decoder::ParquetPushDecoderBuilder
+ pub fn with_row_groups(mut self, row_groups: Vec<usize>) -> Self {
Review Comment:
I wonder if it is too complicated to figure out which API a user should use.
Maybe it would be simpler if we directed users to the preferred api
(`with_group_selections`)
Maybe we should even deprecate the methods and direct everyone towards
`with_row_group_selections`. And we could add an example of how to convert row
groups and selections into row group selections 🤔
##########
parquet/src/arrow/async_reader/mod.rs:
##########
@@ -644,6 +648,29 @@ impl<T: AsyncFileReader + Send + 'static>
ParquetRecordBatchStreamBuilder<T> {
Ok(Some(Sbbf::new(&bitset)))
}
+ /// Select row groups and rows using row-group-local coordinates.
+ ///
+ /// Entries are decoded in the supplied order, omitted row groups are
+ /// skipped, and a `None` selection reads the whole row group. This is
+ /// mutually exclusive with [`ArrowReaderBuilder::with_row_groups`] and
+ /// [`ArrowReaderBuilder::with_row_selection`]; combining them returns an
+ /// error from [`Self::build`].
+ ///
+ /// See [`ParquetPushDecoderBuilder::with_row_group_selections`] for the
Review Comment:
same comment as above
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -42,21 +44,78 @@ enum QueuedRowGroupDecision {
struct NextRowGroup {
row_group_idx: usize,
row_count: usize,
- /// This row group's slice of the global selection, or `None` when all rows
- /// are selected.
+ /// This row group's selection, or `None` when all rows are selected.
selection: Option<RowSelection>,
/// Budget snapshot to apply while decoding this row group.
budget: RowBudget,
}
+#[derive(Debug, Clone)]
+enum QueuedRowGroups {
Review Comment:
Can we please add some comments about what this struct represents (aka the
remaining selections for the reader)
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -42,21 +44,78 @@ enum QueuedRowGroupDecision {
struct NextRowGroup {
row_group_idx: usize,
row_count: usize,
- /// This row group's slice of the global selection, or `None` when all rows
- /// are selected.
+ /// This row group's selection, or `None` when all rows are selected.
selection: Option<RowSelection>,
/// Budget snapshot to apply while decoding this row group.
budget: RowBudget,
}
+#[derive(Debug, Clone)]
+enum QueuedRowGroups {
+ Global {
+ row_groups: VecDeque<usize>,
+ selection: Option<RowSelection>,
+ },
+ PerRowGroup(VecDeque<RowGroupSelection>),
+}
+
+impl QueuedRowGroups {
+ fn front(&self) -> Option<usize> {
+ match self {
+ Self::Global { row_groups, .. } => row_groups.front().copied(),
+ Self::PerRowGroup(row_groups) => row_groups
+ .front()
+ .map(|row_group| row_group.row_group_index),
+ }
+ }
+
+ fn len(&self) -> usize {
+ match self {
+ Self::Global { row_groups, .. } => row_groups.len(),
+ Self::PerRowGroup(row_groups) => row_groups.len(),
+ }
+ }
+
+ fn clear(&mut self) {
+ match self {
+ Self::Global {
+ row_groups,
+ selection,
+ } => {
+ row_groups.clear();
+ *selection = None;
+ }
+ Self::PerRowGroup(row_groups) => row_groups.clear(),
+ }
+ }
+}
+
+/// Number of rows in `row_group_idx`, or an error if the index is out of
+/// bounds for the file.
+fn row_group_num_rows(
Review Comment:
Perhaps this could be added as a method on ParquetMetaData (it seems useful
for multiple users)
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -273,16 +370,26 @@ impl RemainingRowGroups {
// `has_predicates` is recomputed by `build()` from the filter.
let RowGroupFrontier {
parquet_metadata,
- row_groups,
- selection,
+ queued,
budget,
has_predicates: _,
} = frontier;
+ let row_group_plan = match queued {
+ QueuedRowGroups::Global {
Review Comment:
Can we also make this a method on QueuedRowGroups (like `queued.into_plan()`
or something like that?
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -142,35 +223,57 @@ impl RowGroupFrontier {
/// Advance queued row groups until one should be handed to the builder.
fn next_readable_row_group(&mut self) -> Result<Option<NextRowGroup>,
ParquetError> {
loop {
- let Some(&row_group_idx) = self.row_groups.front() else {
+ let Some(row_group_idx) = self.queued.front() else {
return Ok(None);
};
+ // A global selection can be exhausted before its row-group queue.
+ // Per-row-group selections have no shared cursor to exhaust; empty
+ // local selections are discarded by the `selected_rows == 0` path
below.
if self.budget.is_exhausted()
- || self
- .selection
- .as_ref()
- .is_some_and(|selection| selection.row_count() == 0)
+ || matches!(
+ &self.queued,
+ QueuedRowGroups::Global {
+ selection: Some(selection),
+ ..
+ } if selection.row_count() == 0
+ )
{
self.clear_remaining();
return Ok(None);
}
let row_count = self.row_group_num_rows(row_group_idx)?;
- let (selection, selected_rows) = match self.selection.as_mut() {
+ // Pop the front row group, resolving its selection to
+ // row-group-local coordinates (splitting off the global cursor's
+ // prefix, or taking the already-local selection as-is).
+ let selection = match &mut self.queued {
Review Comment:
can we please make this a method on QueuedRowGroups rather than inlining it
here -- this code is already long and complicated enough
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -259,6 +393,12 @@ impl<T> ArrowReaderBuilder<T> {
/// applying the row selection, and therefore rows from skipped row groups
/// should not be included in the [`RowSelection`] (see example below)
///
+ /// On [`ParquetPushDecoderBuilder`] and async stream builders, which
Review Comment:
same comments as above
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -67,26 +126,49 @@ struct RowGroupFrontier {
impl RowGroupFrontier {
fn new(
parquet_metadata: Arc<ParquetMetaData>,
- row_groups: Vec<usize>,
- selection: Option<RowSelection>,
+ row_group_plan: RowGroupPlan,
budget: RowBudget,
has_predicates: bool,
- ) -> Self {
- Self {
+ ) -> Result<Self, ParquetError> {
+ let queued = match row_group_plan {
Review Comment:
This looks like it would more naturally be a constructor / from impl on
RowGroupPlan -- to construct `QueuedRowGroups` from a `RowGroupPlan`
That would also encapsulate the compleixty more
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -249,6 +250,55 @@ impl ParquetPushDecoderBuilder {
}
}
+ /// Select row groups and rows using row-group-local coordinates.
+ ///
+ /// Entries are decoded in the supplied order, and omitted row groups are
+ /// skipped. A row group listed more than once is decoded once per entry.
+ /// A `None` selection reads the entire row group. A selection shorter
+ /// than its row group skips the trailing rows, while a selection longer
+ /// than its row group returns an error from [`Self::build`]. Each
+ /// selection is passed through without being re-partitioned, so no
+ /// conversion between the bitmap and selector representations is forced
Review Comment:
I think the mention of different representations of the selection is
confusing here (it seems like an irrelevant implementation detail)
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -142,35 +223,54 @@ impl RowGroupFrontier {
/// Advance queued row groups until one should be handed to the builder.
fn next_readable_row_group(&mut self) -> Result<Option<NextRowGroup>,
ParquetError> {
loop {
- let Some(&row_group_idx) = self.row_groups.front() else {
+ let Some(row_group_idx) = self.queued.front() else {
return Ok(None);
};
if self.budget.is_exhausted()
Review Comment:
the comments are nice -- maybe we can also add a debug assert too to show
that the invariant claimed is true
##########
parquet/src/arrow/push_decoder/mod.rs:
##########
@@ -249,6 +250,55 @@ impl ParquetPushDecoderBuilder {
}
}
+ /// Select row groups and rows using row-group-local coordinates.
+ ///
+ /// Entries are decoded in the supplied order, and omitted row groups are
+ /// skipped. A row group listed more than once is decoded once per entry.
+ /// A `None` selection reads the entire row group. A selection shorter
+ /// than its row group skips the trailing rows, while a selection longer
+ /// than its row group returns an error from [`Self::build`]. Each
+ /// selection is passed through without being re-partitioned, so no
+ /// conversion between the bitmap and selector representations is forced
+ /// (how the selection is then materialized while reading is still governed
+ /// by [`ArrowReaderBuilder::with_row_selection_policy`]).
+ ///
+ /// [`ArrowReaderBuilder::with_offset`] and
+ /// [`ArrowReaderBuilder::with_limit`] apply after the row-group-local
+ /// selections and any [`ArrowReaderBuilder::with_row_filter`], across the
+ /// plan as a whole and in the supplied order: the offset skips the first N
+ /// remaining rows and the limit caps the total rows emitted, regardless of
+ /// which row group they come from.
+ ///
+ /// This configuration is mutually exclusive with
+ /// [`ArrowReaderBuilder::with_row_groups`] and
+ /// [`ArrowReaderBuilder::with_row_selection`]. Combining them returns an
+ /// error from [`Self::build`]. Calling this method more than once replaces
+ /// the previous row-group-local configuration.
+ ///
+ /// ```no_run
Review Comment:
It took me some time to figure out (with some help) that we duplicated
`with_row_group_selections` twice because it isn't supported via the
serialized reader
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -42,21 +44,78 @@ enum QueuedRowGroupDecision {
struct NextRowGroup {
row_group_idx: usize,
row_count: usize,
- /// This row group's slice of the global selection, or `None` when all rows
- /// are selected.
+ /// This row group's selection, or `None` when all rows are selected.
selection: Option<RowSelection>,
/// Budget snapshot to apply while decoding this row group.
budget: RowBudget,
}
+#[derive(Debug, Clone)]
+enum QueuedRowGroups {
Review Comment:
As a follow on PR it seems like it might be simpler to just do the
translation into `VecDequeue<RowGroupSelection>` once at build time rather than
having to support both types of groups during decode
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -64,6 +64,109 @@ pub mod statistics;
/// Default batch size for reading parquet files
pub const DEFAULT_BATCH_SIZE: usize = 1024;
+/// A row group and its optional row-group-local [`RowSelection`].
+///
+/// A `None` selection reads the entire row group. Omitting a row group skips
+/// it. Entries are decoded in the supplied order.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RowGroupSelection {
+ pub(crate) row_group_index: usize,
+ pub(crate) selection: Option<RowSelection>,
+}
+
+impl RowGroupSelection {
+ /// Creates a row-group-local selection.
+ pub fn new(row_group_index: usize, selection: Option<RowSelection>) ->
Self {
+ Self {
+ row_group_index,
+ selection,
+ }
+ }
+
+ /// The index of the row group this selection applies to.
+ pub fn row_group_index(&self) -> usize {
+ self.row_group_index
+ }
+
+ /// The row-group-local selection, or `None` if the entire row group is
+ /// read.
+ pub fn selection(&self) -> Option<&RowSelection> {
+ self.selection.as_ref()
+ }
+}
+
+/// Row-selection configuration shared by the Arrow reader builders.
+///
+/// `Global` is the legacy configuration formed by `with_row_groups` and
Review Comment:
(I am happy to make these changes myself, but I wanted to double check first)
##########
parquet/src/arrow/push_decoder/remaining.rs:
##########
@@ -142,35 +223,54 @@ impl RowGroupFrontier {
/// Advance queued row groups until one should be handed to the builder.
fn next_readable_row_group(&mut self) -> Result<Option<NextRowGroup>,
ParquetError> {
loop {
- let Some(&row_group_idx) = self.row_groups.front() else {
+ let Some(row_group_idx) = self.queued.front() else {
return Ok(None);
};
if self.budget.is_exhausted()
Review Comment:
Also putting it as a method on QueuedRowGroups would help 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]
