Use stream in cudf_polars.DataFrame.to_polars (#20323)

This ensures that the device-to-host transfer used in `cudf_polars.DataFrame.to_polars` performed on the stream associated with the dataframe.


As a rough test, I updated our stream generation function to always return a new stream:

```diff
diff --git a/python/cudf_polars/cudf_polars/utils/cuda_stream.py b/python/cudf_polars/cudf_polars/utils/cuda_stream.py
index 4f8d540..a3cf6c9 100644
--- a/python/cudf_polars/cudf_polars/utils/cuda_stream.py
+++ b/python/cudf_polars/cudf_polars/utils/cuda_stream.py
@@ -7,7 +7,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from rmm.pylibrmm.stream import DEFAULT_STREAM
+from rmm.pylibrmm.stream import DEFAULT_STREAM, Stream
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -22,7 +22,7 @@ def get_dask_cuda_stream() -> Stream:
 
 def get_cuda_stream() -> Stream:
     """Get the default CUDA stream for the current thread."""
-    return DEFAULT_STREAM
+    return Stream()
 
 
 def join_cuda_streams(
```

And then ran that under nsys. Here's a screenshot of the device-to-host copy that was previously on the default stream.
<img width="1427" height="257" alt="Screenshot 2025-10-21 at 10 31 22 AM" src="https://github.com/user-attachments/assets/5a210e06-9e00-469c-a9e1-1000989b2bb0" />

It's not visible in the screenshot, but there are other `to_arrow_host` calls that still show up on the default stream. I'm tracking those down still.

Closes #20309

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

Approvers:
  - Matthew Roeschke (https://github.com/mroeschke)

URL: #20323

Author: TomAugspurger

python/cudf_polars/cudf_polars/containers/dataframe.py

@@ -56,15 +56,21 @@ def _create_polars_column_metadata(
 # This is also defined in pylibcudf.interop
 class _ObjectWithArrowMetadata:
     def __init__(
-        self, obj: plc.Table | plc.Column, metadata: list[plc.interop.ColumnMetadata]
+        self,
+        obj: plc.Table | plc.Column,
+        metadata: list[plc.interop.ColumnMetadata],
+        stream: Stream,
     ) -> None:
         self.obj = obj
         self.metadata = metadata
+        self.stream = stream
 
     def __arrow_c_array__(
         self, requested_schema: None = None
     ) -> tuple[CapsuleType, CapsuleType]:
-        return self.obj._to_schema(self.metadata), self.obj._to_host_array()
+        return self.obj._to_schema(self.metadata), self.obj._to_host_array(
+            stream=self.stream
+        )
 
 
 # Pacify the type checker. DataFrame init asserts that all the columns
@@ -108,7 +114,9 @@ def to_polars(self) -> pl.DataFrame:
             _create_polars_column_metadata(name, dtype.polars_type)
             for name, dtype in zip(name_map, self.dtypes, strict=True)
         ]
-        table_with_metadata = _ObjectWithArrowMetadata(self.table, metadata)
+        table_with_metadata = _ObjectWithArrowMetadata(
+            self.table, metadata, self.stream
+        )
         df = pl.DataFrame(table_with_metadata)
         return df.rename(name_map).with_columns(
             pl.col(c.name).set_sorted(descending=c.order == plc.types.Order.DESCENDING)

python/pylibcudf/pylibcudf/_interop_helpers.pyx

@@ -12,6 +12,7 @@ from pylibcudf.libcudf.interop cimport (
     release_arrow_device_array_raw,
     release_arrow_schema_raw,
 )
+from .utils cimport _get_stream
 
 from dataclasses import dataclass, field
 
@@ -33,12 +34,16 @@ class ArrowLike(metaclass=_ArrowLikeMeta):
 
 
 class _ObjectWithArrowMetadata:
-    def __init__(self, obj, metadata=None):
+    def __init__(self, obj, metadata=None, stream=None):
         self.obj = obj
         self.metadata = metadata
+        self.stream = _get_stream(stream)
 
     def __arrow_c_array__(self, requested_schema=None):
-        return self.obj._to_schema(self.metadata), self.obj._to_host_array()
+        return (
+            self.obj._to_schema(self.metadata),
+            self.obj._to_host_array(stream=self.stream),
+        )
 
 
 @dataclass

python/pylibcudf/pylibcudf/column.pyi

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 from collections.abc import Iterable, Sequence
@@ -89,11 +89,13 @@ class Column:
     def from_rmm_buffer(
         buff: DeviceBuffer, dtype: DataType, size: int, children: list[Column]
     ) -> Column: ...
-    def to_arrow(self, metadata: list | str | None = None) -> ArrowLike: ...
+    def to_arrow(
+        self, metadata: list | str | None = None, stream: Stream | None = None
+    ) -> ArrowLike: ...
     # Private methods below are included because polars is currently using them,
     # but we want to remove stubs for these private methods eventually
     def _to_schema(self, metadata: Any = None) -> Any: ...
-    def _to_host_array(self) -> Any: ...
+    def _to_host_array(self, stream: Stream) -> Any: ...
     @staticmethod
     def from_arrow(
         obj: ArrowLike,

python/pylibcudf/pylibcudf/column.pyx

@@ -349,14 +349,17 @@ cdef class Column:
 
     def to_arrow(
         self,
-        metadata: ColumnMetadata | str | None = None
+        metadata: ColumnMetadata | str | None = None,
+        stream: Stream = None,
     ) -> ArrowLike:
         """Create a pyarrow array from a pylibcudf column.
 
         Parameters
         ----------
         metadata : ColumnMetadata | str | None
             The metadata to attach to the column.
+        stream : Stream | None
+            CUDA stream on which to perform the operation.
 
         Returns
         -------
@@ -371,7 +374,7 @@ cdef class Column:
         # TODO: Once the arrow C device interface registers more
         # types that it supports, we can call pa.array(self) if
         # no metadata is passed.
-        return pa.array(_ObjectWithArrowMetadata(self, metadata))
+        return pa.array(_ObjectWithArrowMetadata(self, metadata, stream))
 
     @staticmethod
     def from_arrow(
@@ -1324,10 +1327,10 @@ cdef class Column:
 
         return PyCapsule_New(<void*>raw_schema_ptr, 'arrow_schema', _release_schema)
 
-    def _to_host_array(self):
+    def _to_host_array(self, Stream stream):
         cdef ArrowArray* raw_host_array_ptr
         with nogil:
-            raw_host_array_ptr = to_arrow_host_raw(self.view())
+            raw_host_array_ptr = to_arrow_host_raw(self.view(), stream.view())
 
         return PyCapsule_New(<void*>raw_host_array_ptr, "arrow_array", _release_array)
 
@@ -1346,7 +1349,7 @@ cdef class Column:
         if requested_schema is not None:
             raise ValueError("pylibcudf.Column does not support alternative schema")
 
-        return self._to_schema(), self._to_host_array()
+        return self._to_schema(), self._to_host_array(_get_stream(None))
 
     def __arrow_c_device_array__(self, requested_schema=None, **kwargs):
         if requested_schema is not None:

python/pylibcudf/pylibcudf/libcudf/interop.pxd

@@ -135,7 +135,7 @@ cdef extern from *:
     template <typename ViewType>
     ArrowArray* to_arrow_host_raw(
       ViewType const& obj,
-      rmm::cuda_stream_view stream       = cudf::get_default_stream(),
+      rmm::cuda_stream_view stream,
       rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) {
       ArrowArray *arr = new ArrowArray();
       auto device_arr = cudf::to_arrow_host(obj, stream, mr);
@@ -221,10 +221,12 @@ cdef extern from *:
         ArrowSchema *
     ) except +libcudf_exception_handler nogil
     cdef ArrowArray* to_arrow_host_raw(
-        const table_view& tbl
+        const table_view& tbl,
+        cuda_stream_view stream,
     ) except +libcudf_exception_handler nogil
     cdef ArrowArray* to_arrow_host_raw(
-        const column_view& tbl
+        const column_view& tbl,
+        cuda_stream_view stream,
     ) except +libcudf_exception_handler nogil
     cdef void release_arrow_array_raw(
         ArrowArray *

python/pylibcudf/pylibcudf/table.pyi

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 from typing import Any
@@ -20,7 +20,7 @@ class Table:
     # Private methods below are included because polars is currently using them,
     # but we want to remove stubs for these private methods eventually
     def _to_schema(self, metadata: Any = None) -> Any: ...
-    def _to_host_array(self) -> Any: ...
+    def _to_host_array(self, stream: Stream) -> Any: ...
     @staticmethod
     def from_arrow(
         arrow_like: ArrowLike,

python/pylibcudf/pylibcudf/table.pyx

@@ -74,14 +74,17 @@ cdef class Table:
 
     def to_arrow(
         self,
-        metadata: list[ColumnMetadata | str] | None = None
+        metadata: list[ColumnMetadata | str] | None = None,
+        stream: Stream = None,
     ) -> ArrowLike:
         """Create a pyarrow table from a pylibcudf table.
 
         Parameters
         ----------
         metadata : list[ColumnMetadata | str] | None
             The metadata to attach to the columns of the table.
+        stream : Stream | None
+            CUDA stream on which to perform the operation.
 
         Returns
         -------
@@ -96,7 +99,7 @@ cdef class Table:
         # TODO: Once the arrow C device interface registers more
         # types that it supports, we can call pa.table(self) if
         # no metadata is passed.
-        return pa.table(_ObjectWithArrowMetadata(self, metadata))
+        return pa.table(_ObjectWithArrowMetadata(self, metadata, stream))
 
     @staticmethod
     def from_arrow(
@@ -320,10 +323,11 @@ cdef class Table:
 
         return PyCapsule_New(<void*>raw_schema_ptr, "arrow_schema", _release_schema)
 
-    def _to_host_array(self):
+    def _to_host_array(self, Stream stream):
         cdef ArrowArray* raw_host_array_ptr
+
         with nogil:
-            raw_host_array_ptr = to_arrow_host_raw(self.view())
+            raw_host_array_ptr = to_arrow_host_raw(self.view(), stream.view())
 
         return PyCapsule_New(<void*>raw_host_array_ptr, "arrow_array", _release_array)
 
@@ -342,7 +346,7 @@ cdef class Table:
         if requested_schema is not None:
             raise ValueError("pylibcudf.Table does not support alternative schema")
 
-        return self._to_schema(), self._to_host_array()
+        return self._to_schema(), self._to_host_array(_get_stream(None))
 
     def __arrow_c_device_array__(self, requested_schema=None, **kwargs):
         if requested_schema is not None:

python/pylibcudf/tests/test_column_to_arrow.py

@@ -1,10 +1,14 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
+import pytest
 from utils import assert_column_eq
 
+from rmm.pylibrmm.stream import Stream
 
-def test_column_to_arrow(table_data):
+
+@pytest.mark.parametrize("stream", [None, Stream()])
+def test_column_to_arrow(table_data, stream):
     plc_tbl, _ = table_data
     for col in plc_tbl.tbl.columns():
-        assert_column_eq(col, col.to_arrow())
+        assert_column_eq(col, col.to_arrow(stream=stream))

python/pylibcudf/tests/test_table.py

@@ -5,6 +5,8 @@
 import pytest
 from utils import assert_table_eq
 
+from rmm.pylibrmm.stream import Stream
+
 import pylibcudf as plc
 
 
@@ -23,10 +25,11 @@ def test_table_shape(arrow_tbl):
     assert plc_tbl.shape() == arrow_tbl.shape
 
 
-def test_table_to_arrow(table_data):
+@pytest.mark.parametrize("stream", [None, Stream()])
+def test_table_to_arrow(table_data, stream):
     plc_tbl, _ = table_data
     expect = plc_tbl.tbl
-    got = expect.to_arrow()
+    got = expect.to_arrow(stream=stream)
     # The order of `got` and `expect` is reversed here
     # because in almost all pylibcudf tests the `expect`
     # is a pyarrow object while `got` is a pylibcudf object,