python: fix pyarrow imports (#4025)

pola-rs · Jul 15, 2022 · e012ec7 · e012ec7
1 parent 394c738
commit e012ec7
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git a/py-polars/polars/internals/frame.py b/py-polars/polars/internals/frame.py
@@ -65,6 +65,10 @@
 try:
     import pyarrow as pa
 
+    # do not remove these
+    import pyarrow.compute
+    import pyarrow.parquet
+
     _PYARROW_AVAILABLE = True
 except ImportError:  # pragma: no cover
     _PYARROW_AVAILABLE = False

diff --git a/py-polars/polars/io.py b/py-polars/polars/io.py
@@ -9,6 +9,11 @@
 try:
     import pyarrow as pa
 
+    # do not remove these
+    import pyarrow.csv
+    import pyarrow.feather
+    import pyarrow.parquet
+
     _PYARROW_AVAILABLE = True
 except ImportError:  # pragma: no cover
     _PYARROW_AVAILABLE = False

diff --git a/py-polars/tests/io/test_parquet.py b/py-polars/tests/io/test_parquet.py
@@ -39,6 +39,13 @@ def test_to_from_buffer(df: pl.DataFrame, compressions: list[str]) -> None:
             read_df = pl.read_parquet(buf)
             assert df.frame_equal(read_df, null_equal=True)
 
+    for use_pyarrow in [True, False]:
+        buf = io.BytesIO()
+        df.write_parquet(buf, use_pyarrow=use_pyarrow)
+        buf.seek(0)
+        read_df = pl.read_parquet(buf, use_pyarrow=use_pyarrow)
+        assert df.frame_equal(read_df, null_equal=True)
+
 
 def test_to_from_file(
     io_test_dir: str, df: pl.DataFrame, compressions: list[str]