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

python/ray/data/_internal/arrow_block.py

@@ -109,7 +109,7 @@ class ArrowBlockBuilder(TableBlockBuilder[T]):
     def __init__(self):
         if pyarrow is None:
             raise ImportError("Run `pip install pyarrow` for Arrow support")
-        super().__init__(pyarrow.Table)
+        super().__init__((pyarrow.Table, bytes))
 
     @staticmethod
     def _table_from_pydict(columns: Dict[str, List[Any]]) -> Block:

python/ray/data/_internal/delegating_block_builder.py

@@ -53,7 +53,7 @@ def add_block(self, block: Block):
             return
         if self._builder is None:
             self._builder = accessor.builder()
-        self._builder.add_block(block)
+        self._builder.add_block(accessor.to_block())
 
     def will_build_yield_copy(self) -> bool:
         if self._builder is None:

python/ray/data/tests/test_raydp_dataset.py

@@ -2,6 +2,7 @@
 import ray
 import raydp
 import torch
+import pandas
 
 
 # RayDP tests require Ray Java. Make sure ray jar is built before running this test.
@@ -51,6 +52,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