fix bug in slice pushdown

pola-rs · Jan 11, 2022 · 41f63cb · 41f63cb
1 parent c82c8d5
commit 41f63cb
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/optimizer/slice_pushdown.rs b/polars/polars-lazy/src/logical_plan/optimizer/slice_pushdown.rs
@@ -76,7 +76,7 @@ impl SlicePushDown {
                 options,
                 predicate,
                 aggregate,
-            }, Some(state)) if state.offset > 0 => {
+            }, Some(state)) if state.offset >= 0 => {
                 let mut options = options;
                 options.skip_rows = state.offset as usize;
                 options.n_rows = Some(state.len as usize);
@@ -204,9 +204,17 @@ impl SlicePushDown {
                 Ok(lp.from_exprs_and_input(exprs, new_inputs))
             }
             (catch_all, state) => {
-                assert!(state.is_none());
-                Ok(catch_all)
-
+                match state {
+                    Some(state) => {
+                        let input = lp_arena.add(catch_all);
+                        Ok(Slice {
+                            input,
+                            offset: state.offset,
+                            len: state.len
+                        })
+                    }
+                    None => Ok(catch_all)
+                }
             }
 
         }

diff --git a/polars/polars-lazy/src/tests/io.rs b/polars/polars-lazy/src/tests/io.rs
@@ -124,3 +124,11 @@ fn test_csv_globbing() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+pub fn test_simple_slice() -> Result<()> {
+    let out = scan_foods_parquet(false).limit(3).collect()?;
+    assert_eq!(out.height(), 3);
+
+    Ok(())
+}