Merge remote-tracking branch 'upstream/main' into business

pola-rs · Apr 8, 2024 · ca7075d · ca7075d
2 parents ff3bbd6 + c211fad
commit ca7075d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/py-polars/polars/io/csv/functions.py b/py-polars/polars/io/csv/functions.py
@@ -284,6 +284,7 @@ def read_csv(
                     data,
                     pa.csv.ReadOptions(
                         skip_rows=skip_rows,
+                        skip_rows_after_names=skip_rows_after_header,
                         autogenerate_column_names=not has_header,
                         encoding=encoding,
                     ),

diff --git a/py-polars/tests/unit/io/test_csv.py b/py-polars/tests/unit/io/test_csv.py
@@ -2037,3 +2037,19 @@ def test_csv_escape_cf_15349() -> None:
     df.write_csv(f)
     f.seek(0)
     assert f.read() == b'test\nnormal\n"with\rcr"\n'
+
+
+@pytest.mark.parametrize("use_pyarrow", [True, False])
+def test_skip_rows_after_header_pyarrow(use_pyarrow: bool) -> None:
+    csv = textwrap.dedent(
+        """\
+        foo,bar
+        1,2
+        3,4
+        5,6
+        """
+    )
+    f = io.StringIO(csv)
+    df = pl.read_csv(f, skip_rows_after_header=1, use_pyarrow=use_pyarrow)
+    expected = pl.DataFrame({"foo": [3, 5], "bar": [4, 6]})
+    assert_frame_equal(df, expected)