fix[rust]: scan_csv 'n_rows' should block predicate pushdown (#4611)

pola-rs · Aug 29, 2022 · 85eeaec · 85eeaec
1 parent 2a3514a
commit 85eeaec
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 7 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs b/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs
@@ -302,14 +302,31 @@ impl PredicatePushDown {
             } => {
                 let predicate = predicate_at_scan(acc_predicates, predicate, expr_arena);
 
-                let lp = CsvScan {
-                    path,
-                    schema,
-                    output_schema,
-                    options,
-                    predicate,
-                    aggregate,
+                let lp = if let (Some(predicate), Some(_)) = (predicate, options.n_rows) {
+                    let lp = CsvScan {
+                        path,
+                        schema,
+                        output_schema,
+                        options,
+                        predicate: None,
+                        aggregate,
+                    };
+                    let input = lp_arena.add(lp);
+                    Selection {
+                        input,
+                        predicate
+                    }
+                } else {
+                    CsvScan {
+                        path,
+                        schema,
+                        output_schema,
+                        options,
+                        predicate,
+                        aggregate,
+                    }
                 };
+
                 Ok(lp)
             }
             AnonymousScan {

diff --git a/py-polars/tests/io/test_lazy_csv.py b/py-polars/tests/io/test_lazy_csv.py
@@ -65,3 +65,18 @@ def test_scan_csv_schema_overwrite_and_dtypes_overwrite(foods_csv: str) -> None:
         .collect()
         .dtypes
     ) == [pl.Utf8, pl.Utf8, pl.Float32, pl.Int64]
+
+
+def test_lazy_n_rows(foods_csv: str) -> None:
+    df = (
+        pl.scan_csv(foods_csv, n_rows=4, row_count_name="idx")
+        .filter(pl.col("idx") > 2)
+        .collect()
+    )
+    assert df.to_dict(False) == {
+        "idx": [3],
+        "category": ["fruit"],
+        "calories": [60],
+        "fats_g": [0.0],
+        "sugars_g": [11],
+    }