Skip to content

Commit

Permalink
Make RowSelection's from_consecutive_ranges public (#5848)
Browse files Browse the repository at this point in the history
* Make RowSelection's from_consecutive_ranges public

This constructor method should be easier to use.

* Address reviewers' comments

---------

Co-authored-by: Andrew Lamb <[email protected]>
  • Loading branch information
advancedxy and alamb authored Jun 11, 2024
1 parent b8c2741 commit a20116e
Showing 1 changed file with 16 additions and 3 deletions.
19 changes: 16 additions & 3 deletions parquet/src/arrow/arrow_reader/selection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,13 @@ impl RowSelector {
///
/// let actual: Vec<RowSelector> = selection.into();
/// assert_eq!(actual, expected);
///
/// // you can also create a selection from consecutive ranges
/// let ranges = vec![5..10, 10..15];
/// let selection =
/// RowSelection::from_consecutive_ranges(ranges.into_iter(), 20);
/// let actual: Vec<RowSelector> = selection.into();
/// assert_eq!(actual, expected);
/// ```
///
/// A [`RowSelection`] maintains the following invariants:
Expand Down Expand Up @@ -115,7 +122,7 @@ impl RowSelection {
}

/// Creates a [`RowSelection`] from an iterator of consecutive ranges to keep
pub(crate) fn from_consecutive_ranges<I: Iterator<Item = Range<usize>>>(
pub fn from_consecutive_ranges<I: Iterator<Item = Range<usize>>>(
ranges: I,
total_rows: usize,
) -> Self {
Expand Down Expand Up @@ -1136,7 +1143,7 @@ mod tests {
}

#[test]
fn test_empty_ranges() {
fn test_from_ranges() {
let ranges = [1..3, 4..6, 6..6, 8..8, 9..10];
let selection = RowSelection::from_consecutive_ranges(ranges.into_iter(), 10);
assert_eq!(
Expand All @@ -1149,7 +1156,13 @@ mod tests {
RowSelector::skip(3),
RowSelector::select(1)
]
)
);

let out_of_order_ranges = [1..3, 8..10, 4..7];
let result = std::panic::catch_unwind(|| {
RowSelection::from_consecutive_ranges(out_of_order_ranges.into_iter(), 10)
});
assert!(result.is_err());
}

#[test]
Expand Down

0 comments on commit a20116e

Please sign in to comment.