Use BufferOutputStream in P2P (dask#7991)

crusaderky · Jul 12, 2023 · 7b21399 · 7b21399
1 parent 6c37907
commit 7b21399
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
@@ -79,14 +79,12 @@ def list_of_buffers_to_table(data: list[bytes]) -> pa.Table:
 
 
 def serialize_table(table: pa.Table) -> bytes:
-    import io
-
     import pyarrow as pa
 
-    stream = io.BytesIO()
+    stream = pa.BufferOutputStream()
     with pa.ipc.new_stream(stream, table.schema) as writer:
         writer.write_table(table)
-    return stream.getvalue()
+    return stream.getvalue().to_pybytes()
 
 
 def deserialize_table(buffer: bytes) -> pa.Table: