fix(python): Allow passing files opened by fsspec in read_parquet (p…

…ola-rs#15770)
rancomp · Apr 19, 2024 · db456e1 · db456e1
1 parent 0c89c71
commit db456e1
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/py-polars/polars/io/parquet/functions.py b/py-polars/polars/io/parquet/functions.py
@@ -162,8 +162,8 @@ def read_parquet(
             memory_map=memory_map,
         )
 
-    # Read binary types using `read_parquet`
-    elif isinstance(source, (io.BufferedIOBase, io.RawIOBase, bytes)):
+    # Read file and bytes inputs using `read_parquet`
+    elif isinstance(source, (io.IOBase, bytes)):
         return _read_parquet_binary(
             source,
             columns=columns,

diff --git a/py-polars/tests/unit/io/test_parquet.py b/py-polars/tests/unit/io/test_parquet.py
@@ -5,6 +5,7 @@
 from decimal import Decimal
 from typing import TYPE_CHECKING, cast
 
+import fsspec
 import numpy as np
 import pandas as pd
 import pyarrow as pa
@@ -680,6 +681,20 @@ def test_read_parquet_binary_file_io(tmp_path: Path) -> None:
     assert_frame_equal(out, df)
 
 
+# https://github.com/pola-rs/polars/issues/15760
+@pytest.mark.write_disk()
+def test_read_parquet_binary_fsspec(tmp_path: Path) -> None:
+    tmp_path.mkdir(exist_ok=True)
+
+    df = pl.DataFrame({"a": [1, 2, 3]})
+    file_path = tmp_path / "test.parquet"
+    df.write_parquet(file_path)
+
+    with fsspec.open(file_path) as f:
+        out = pl.read_parquet(f)
+    assert_frame_equal(out, df)
+
+
 def test_read_parquet_binary_bytes_io() -> None:
     df = pl.DataFrame({"a": [1, 2, 3]})
     f = io.BytesIO()