Add CUDA stream to cudf_polars.Column.deserialize (#20396)

Column.deserialize uses `plc.contiguous_split.unpack_from_memoryviews`, which accepts a stream argument.

This also fixes the type annotation for `unpack_from_memoryviews` and `PackedColumns.unpack` which were missing a `Stream` argument.

Authors:
  - Tom Augspurger (https://github.com/TomAugspurger)

Approvers:
  - Matthew Murray (https://github.com/Matt711)

URL: #20396

Author: TomAugspurger

python/cudf_polars/cudf_polars/containers/column.py

@@ -71,7 +71,10 @@ def __init__(
 
     @classmethod
     def deserialize(
-        cls, header: ColumnHeader, frames: tuple[memoryview[bytes], plc.gpumemoryview]
+        cls,
+        header: ColumnHeader,
+        frames: tuple[memoryview[bytes], plc.gpumemoryview],
+        stream: Stream,
     ) -> Self:
         """
         Create a Column from a serialized representation returned by `.serialize()`.
@@ -82,6 +85,10 @@ def deserialize(
             The (unpickled) metadata required to reconstruct the object.
         frames
             Two-tuple of frames (a memoryview and a gpumemoryview).
+        stream
+            CUDA stream used for device memory operations and kernel launches
+            on this column. The caller is responsible for ensuring that
+            the data in ``frames`` is valid on ``stream``.
 
         Returns
         -------
@@ -90,7 +97,7 @@ def deserialize(
         """
         packed_metadata, packed_gpu_data = frames
         (plc_column,) = plc.contiguous_split.unpack_from_memoryviews(
-            packed_metadata, packed_gpu_data
+            packed_metadata, packed_gpu_data, stream
         ).columns()
         return cls(plc_column, **cls.deserialize_ctor_kwargs(header["column_kwargs"]))
 

python/cudf_polars/cudf_polars/containers/dataframe.py

@@ -249,6 +249,7 @@ def deserialize(
         table = plc.contiguous_split.unpack_from_memoryviews(
             packed_metadata,
             packed_gpu_data,
+            stream,
         )
         return cls(
             (

python/cudf_polars/cudf_polars/experimental/dask_registers.py

@@ -100,7 +100,11 @@ def _(
     ) -> Column:
         with log_errors():
             metadata, gpudata = frames
-            return Column.deserialize(header, (metadata, plc.gpumemoryview(gpudata)))
+            return Column.deserialize(
+                header,
+                (metadata, plc.gpumemoryview(gpudata)),
+                stream=get_dask_cuda_stream(),
+            )
 
     @overload
     def dask_serialize_column_or_frame(
@@ -144,7 +148,7 @@ def _(header: ColumnHeader, frames: tuple[memoryview[bytes], memoryview]) -> Col
                 frames[0],
                 plc.gpumemoryview(rmm.DeviceBuffer.to_device(frames[1])),
             )
-            return Column.deserialize(header, new_frames)
+            return Column.deserialize(header, new_frames, stream=get_dask_cuda_stream())
 
     @dask_serialize.register(DataFrame)
     def _(

python/cudf_polars/tests/containers/test_column.py

@@ -230,7 +230,7 @@ def test_serialize_cache_miss():
     # the same hash, so they cache the same, but have some difference
     # in behavior (e.g. isinstance).
     cudf_polars.containers.datatype._from_polars.cache_clear()
-    result = Column.deserialize(header, frames)
+    result = Column.deserialize(header, frames, stream=stream)
     assert result.dtype == dtype
 
 

python/pylibcudf/pylibcudf/contiguous_split.pyi

@@ -15,9 +15,11 @@ class PackedColumns:
     ) -> tuple[memoryview[bytes], gpumemoryview]: ...
 
 def pack(input: Table, stream: Stream | None = None) -> PackedColumns: ...
-def unpack(input: PackedColumns) -> Table: ...
+def unpack(input: PackedColumns, stream: Stream | None = None) -> Table: ...
 def unpack_from_memoryviews(
-    metadata: memoryview[bytes], gpu_data: gpumemoryview
+    metadata: memoryview[bytes],
+    gpu_data: gpumemoryview,
+    stream: Stream | None = None,
 ) -> Table: ...
 
 class ChunkedPack: