feature(pytest): Update testing

Author: jakepenzak

ibis/backends/pyspark/tests/test_import_export.py

@@ -5,6 +5,7 @@
 
 import pandas as pd
 import pytest
+from pandas.testing import assert_frame_equal
 
 from ibis.backends.pyspark.datatypes import PySparkSchema
 
@@ -83,8 +84,11 @@ def test_to_parquet_read_parquet(con, tmp_path):
 
     t_in = con.read_parquet(tmp_path / "out_np")
 
-    assert t_out.to_pandas().shape == t_in.to_pandas().shape
-    assert sorted(t_out.columns) == sorted(t_in.columns)
+    cols = list(t_out.columns)
+    expected = t_out.to_pandas()[cols].sort_values(cols).reset_index(drop=True)
+    result = t_in.to_pandas()[cols].sort_values(cols).reset_index(drop=True)
+
+    assert_frame_equal(expected, result)
 
     # Partitions
     t_out = con.table("awards_players")
@@ -99,5 +103,8 @@ def test_to_parquet_read_parquet(con, tmp_path):
 
     t_in = con.read_parquet(tmp_path / "out_p")
 
-    assert t_out.to_pandas().shape == t_in.to_pandas().shape
-    assert sorted(t_out.columns) == sorted(t_in.columns)
+    cols = list(t_out.columns)
+    expected = t_out.to_pandas()[cols].sort_values(cols).reset_index(drop=True)
+    result = t_in.to_pandas()[cols].sort_values(cols).reset_index(drop=True)
+
+    assert_frame_equal(expected, result)

ibis/backends/tests/test_export.py

@@ -257,15 +257,30 @@ def test_table_to_parquet_writer_kwargs(version, tmp_path, backend, awards_playe
     outparquet = tmp_path / "out.parquet"
     awards_players.to_parquet(outparquet, version=version)
 
-    df = pd.read_parquet(outparquet)
+    if backend.name() == "pyspark":
+        # Pyspark will write more than one parquet file under outparquet as directory
+        parquet_files = sorted(outparquet.glob("*.parquet"))
+        df = (
+            pd.concat(map(pd.read_parquet, parquet_files))
+            .sort_values(list(awards_players.columns))
+            .reset_index(drop=True)
+        )
+        result = (
+            awards_players.to_pandas()
+            .sort_values(list(awards_players.columns))
+            .reset_index(drop=True)
+        )
+        backend.assert_frame_equal(result, df)
+    else:
+        df = pd.read_parquet(outparquet)
 
-    backend.assert_frame_equal(
-        awards_players.to_pandas().fillna(pd.NA), df.fillna(pd.NA)
-    )
+        backend.assert_frame_equal(
+            awards_players.to_pandas().fillna(pd.NA), df.fillna(pd.NA)
+        )
 
-    md = pa.parquet.read_metadata(outparquet)
+        md = pa.parquet.read_metadata(outparquet)
 
-    assert md.format_version == version
+        assert md.format_version == version
 
 
 @pytest.mark.notimpl(
@@ -333,7 +348,10 @@ def test_memtable_to_file(tmp_path, con, ftype, monkeypatch):
 
     getattr(con, f"to_{ftype}")(memtable, outfile)
 
-    assert outfile.is_file()
+    if con.name == "pyspark" and ftype == "parquet":
+        assert outfile.is_dir()
+    else:
+        assert outfile.is_file()
 
 
 def test_table_to_csv(tmp_path, backend, awards_players):