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

Add limit to ArrowReaderBuilder to push limit down to parquet reader #3633

Merged
merged 4 commits into from
Jan 30, 2023
+282 −5
Merged

Add limit to ArrowReaderBuilder to push limit down to parquet reader #3633

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a6b65fb
Add limit to ArrowReaderBuilder to push limit down to parquet reader
thinkharderdev Jan 30, 2023
db9c2e7
Update parquet/src/arrow/arrow_reader/mod.rs
thinkharderdev Jan 30, 2023
59eccc5
pr comments
thinkharderdev Jan 30, 2023
6805164
Apply limit to entire file instead of each row group
thinkharderdev Jan 30, 2023
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
58 changes: 57 additions & 1 deletion parquet/src/arrow/arrow_reader/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@ pub struct ArrowReaderBuilder<T> {
pub(crate) filter: Option<RowFilter>,

pub(crate) selection: Option<RowSelection>,

pub(crate) limit: Option<usize>,
}

impl<T> ArrowReaderBuilder<T> {
Expand Down Expand Up @@ -98,6 +100,7 @@ impl<T> ArrowReaderBuilder<T> {
projection: ProjectionMask::all(),
filter: None,
selection: None,
limit: None,
})
}

Expand Down Expand Up @@ -167,6 +170,17 @@ impl<T> ArrowReaderBuilder<T> {
..self
}
}

/// Provide a limit to the number of rows to be read
///
/// The limit will be applied after any [`Self::with_row_selection`] and [`Self::with_row_filter`]
/// allowing it to limit the final set of rows decoded after any pushed down predicates
pub fn with_limit(self, limit: usize) -> Self {
Self {
limit: Some(limit),
..self
}
}
}

/// Arrow reader api.
Expand Down Expand Up @@ -453,6 +467,19 @@ impl<T: ChunkReader + 'static> ArrowReaderBuilder<SyncReader<T>> {
selection = Some(RowSelection::from(vec![]));
}

// If a limit is defined, apply it to the final `RowSelection`
Copy link
Contributor

Choose a reason for hiding this comment

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

We need to also do something similar for the async reader

if let Some(limit) = self.limit {
selection = Some(
selection
.map(|selection| selection.limit(limit))
.unwrap_or_else(|| {
RowSelection::from(vec![RowSelector::select(
limit.min(reader.num_rows()),
)])
}),
);
}

Ok(ParquetRecordBatchReader::new(
batch_size,
array_reader,
Expand Down Expand Up @@ -1215,6 +1242,8 @@ mod tests {
row_selections: Option<(RowSelection, usize)>,
/// row filter
row_filter: Option<Vec<bool>>,
/// limit
limit: Option<usize>,
}

/// Manually implement this to avoid printing entire contents of row_selections and row_filter
Expand All @@ -1233,6 +1262,7 @@ mod tests {
.field("encoding", &self.encoding)
.field("row_selections", &self.row_selections.is_some())
.field("row_filter", &self.row_filter.is_some())
.field("limit", &self.limit)
.finish()
}
}
Expand All @@ -1252,6 +1282,7 @@ mod tests {
encoding: Encoding::PLAIN,
row_selections: None,
row_filter: None,
limit: None,
}
}
}
Expand Down Expand Up @@ -1323,6 +1354,13 @@ mod tests {
}
}

fn with_limit(self, limit: usize) -> Self {
Self {
limit: Some(limit),
..self
}
}

fn writer_props(&self) -> WriterProperties {
let builder = WriterProperties::builder()
.set_data_pagesize_limit(self.max_data_page_size)
Expand Down Expand Up @@ -1381,6 +1419,14 @@ mod tests {
TestOptions::new(2, 256, 127).with_null_percent(0),
// Test optional with nulls
TestOptions::new(2, 256, 93).with_null_percent(25),
// Test with limit of 0
TestOptions::new(4, 100, 25).with_limit(0),
// Test with limit of 50
TestOptions::new(4, 100, 25).with_limit(50),
// Test with limit equal to number of rows
TestOptions::new(4, 100, 25).with_limit(10),
// Test with limit larger than number of rows
TestOptions::new(4, 100, 25).with_limit(101),
// Test with no page-level statistics
TestOptions::new(2, 256, 91)
.with_null_percent(25)
Expand Down Expand Up @@ -1423,6 +1469,11 @@ mod tests {
TestOptions::new(2, 256, 93)
.with_null_percent(25)
.with_row_selections(),
// Test optional with nulls
TestOptions::new(2, 256, 93)
.with_null_percent(25)
.with_row_selections()
.with_limit(10),
// Test filter

// Test with row filter
Expand Down Expand Up @@ -1592,7 +1643,7 @@ mod tests {
}
};

let expected_data = match opts.row_filter {
let mut expected_data = match opts.row_filter {
Some(filter) => {
let expected_data = expected_data
.into_iter()
Expand Down Expand Up @@ -1622,6 +1673,11 @@ mod tests {
None => expected_data,
};

if let Some(limit) = opts.limit {
builder = builder.with_limit(limit);
expected_data = expected_data.into_iter().take(limit).collect();
}

let mut record_reader = builder
.with_batch_size(opts.record_batch_size)
.build()
Expand Down
82 changes: 81 additions & 1 deletion parquet/src/arrow/arrow_reader/selection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ use arrow_array::{Array, BooleanArray};
use arrow_select::filter::SlicesIterator;
use std::cmp::Ordering;
use std::collections::VecDeque;
use std::mem;
use std::ops::Range;

/// [`RowSelection`] is a collection of [`RowSelector`] used to skip rows when
Expand Down Expand Up @@ -111,7 +112,7 @@ impl RowSelection {
}

/// Creates a [`RowSelection`] from an iterator of consecutive ranges to keep
fn from_consecutive_ranges<I: Iterator<Item = Range<usize>>>(
pub(crate) fn from_consecutive_ranges<I: Iterator<Item = Range<usize>>>(
ranges: I,
total_rows: usize,
) -> Self {
Expand Down Expand Up @@ -371,6 +372,32 @@ impl RowSelection {
self
}

/// Limit this [`RowSelection`] to only select `limit` rows
pub(crate) fn limit(mut self, mut limit: usize) -> Self {
let mut new_selectors = Vec::with_capacity(self.selectors.len());
for mut selection in mem::take(&mut self.selectors) {
if limit == 0 {
break;
}

if !selection.skip {
if selection.row_count >= limit {
selection.row_count = limit;
new_selectors.push(selection);
break;
} else {
limit -= selection.row_count;
new_selectors.push(selection);
}
} else {
new_selectors.push(selection);
}
}

self.selectors = new_selectors;
self
}

/// Returns an iterator over the [`RowSelector`]s for this
/// [`RowSelection`].
pub fn iter(&self) -> impl Iterator<Item = &RowSelector> {
Expand Down Expand Up @@ -841,6 +868,59 @@ mod tests {
assert_eq!(selectors, round_tripped);
}

#[test]
fn test_limit() {
// Limit to existing limit should no-op
let selection =
RowSelection::from(vec![RowSelector::select(10), RowSelector::skip(90)]);
let limited = selection.limit(10);
assert_eq!(RowSelection::from(vec![RowSelector::select(10)]), limited);

let selection = RowSelection::from(vec![
RowSelector::select(10),
RowSelector::skip(10),
RowSelector::select(10),
RowSelector::skip(10),
RowSelector::select(10),
]);

let limited = selection.clone().limit(5);
let expected = vec![RowSelector::select(5)];
assert_eq!(limited.selectors, expected);

let limited = selection.clone().limit(15);
let expected = vec![
RowSelector::select(10),
RowSelector::skip(10),
RowSelector::select(5),
];
assert_eq!(limited.selectors, expected);

let limited = selection.clone().limit(0);
let expected = vec![];
assert_eq!(limited.selectors, expected);

let limited = selection.clone().limit(30);
let expected = vec![
RowSelector::select(10),
RowSelector::skip(10),
RowSelector::select(10),
RowSelector::skip(10),
RowSelector::select(10),
];
assert_eq!(limited.selectors, expected);

let limited = selection.limit(100);
let expected = vec![
RowSelector::select(10),
RowSelector::skip(10),
RowSelector::select(10),
RowSelector::skip(10),
RowSelector::select(10),
];
assert_eq!(limited.selectors, expected);
}

#[test]
fn test_scan_ranges() {
let index = vec![
Expand Down
Loading