Skip to content

Commit

Permalink
don't overwrite skip_rows in slice_pushdown opotimization (#2556)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Feb 6, 2022
1 parent 9c0c31f commit b77c963
Show file tree
Hide file tree
Showing 4 changed files with 15 additions and 2 deletions.
2 changes: 1 addition & 1 deletion polars/polars-lazy/src/frame/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -938,7 +938,7 @@ impl LazyFrame {
self.slice(-1, 1)
}

/// Get the n last rows
/// Get the last `n` rows
pub fn tail(self, n: u32) -> LazyFrame {
let neg_tail = -(n as i64);
self.slice(neg_tail, n)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ impl SlicePushDown {
aggregate,
}, Some(state)) if state.offset >= 0 && predicate.is_none() => {
let mut options = options;
options.skip_rows = state.offset as usize;
options.skip_rows += state.offset as usize;
options.n_rows = Some(state.len as usize);

let lp = CsvScan {
Expand Down
12 changes: 12 additions & 0 deletions polars/polars-lazy/src/tests/io.rs
Original file line number Diff line number Diff line change
Expand Up @@ -316,3 +316,15 @@ fn test_slice_filter() -> Result<()> {

Ok(())
}

#[test]
fn skip_rows_and_slice() -> Result<()> {
let out = LazyCsvReader::new(FOODS_CSV.to_string())
.with_skip_rows(4)
.finish()?
.limit(1)
.collect()?;
assert_eq!(out.column("fruit")?.get(0), AnyValue::Utf8("seafood"));
assert_eq!(out.shape(), (1, 4));
Ok(())
}
1 change: 1 addition & 0 deletions polars/polars-lazy/src/tests/projection_queries.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ fn test_join_suffix_and_drop() -> Result<()> {
}

#[test]
#[cfg(feature = "cross_join")]
fn test_cross_join_pd() -> Result<()> {
let food = df![
"name"=> ["Omelette", "Fried Egg"],
Expand Down

0 comments on commit b77c963

Please sign in to comment.