[Ray Data] fix problem: to_pandas failed on datasets returned by from_spark

python/ray/data/_internal/arrow_block.py

@@ -111,6 +111,23 @@ def __init__(self):
             raise ImportError("Run `pip install pyarrow` for Arrow support")
         super().__init__(pyarrow.Table)
 
+    def add_block(self, block: Any) -> None:
+        if not isinstance(block, pyarrow.Table) and not isinstance(block, bytes):
+            raise TypeError(
+                f"Got a block of type {type(block)}, expected {self._block_type}."
+                "If you are mapping a function, ensure it returns an "
+                "object with the expected type. Block:\n"
+                f"{block}"
+            )
+        accessor = BlockAccessor.for_block(block)
+        if isinstance(block, bytes):
+            reader = pyarrow.ipc.open_stream(block)
+            self._tables.append(reader.read_all())
+        else:
+            self._tables.append(block)
+        self._tables_size_bytes += accessor.size_bytes()
+        self._num_rows += accessor.num_rows()
+
     @staticmethod
     def _table_from_pydict(columns: Dict[str, List[Any]]) -> Block:
         for col_name, col in columns.items():

python/ray/data/tests/test_raydp_dataset.py

@@ -2,6 +2,7 @@
 import ray
 import raydp
 import torch
+import pandas
 
 
 @pytest.fixture(scope="function")
@@ -50,6 +51,14 @@ def test_raydp_to_torch_iter(spark):
     assert torch.equal(data_features, features) and torch.equal(data_labels, labels)
 
 
+def test_to_pandas(spark):
+    df = spark.range(100)
+    ds = ray.data.from_spark(df)
+    pdf = ds.to_pandas()
+    pdf2 = df.toPandas()
+    pandas.testing.assert_frame_equal(pdf, pdf2)
+
+
 if __name__ == "__main__":
     import sys