Merge pull request #20046 from rapidsai/branch-25.10

Forward-merge branch-25.10 into branch-25.12

Author: GPUtester

python/pylibcudf/pylibcudf/interop.pxd

@@ -2,7 +2,10 @@
 
 from pylibcudf.table cimport Table
 from rmm.pylibrmm.stream cimport Stream
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
-cpdef Table from_dlpack(object managed_tensor, Stream stream=*)
+cpdef Table from_dlpack(
+    object managed_tensor, Stream stream=*, DeviceMemoryResource mr=*
+)
 
-cpdef object to_dlpack(Table input, Stream stream=*)
+cpdef object to_dlpack(Table input, Stream stream=*, DeviceMemoryResource mr=*)

python/pylibcudf/pylibcudf/interop.pyi

@@ -6,6 +6,7 @@ from typing import Any, overload
 
 import pyarrow as pa
 
+from rmm.pylibrmm.memory_resource import DeviceMemoryResource
 from rmm.pylibrmm.stream import Stream
 
 from pylibcudf.column import Column
@@ -51,6 +52,12 @@ def to_arrow(
     obj: Scalar, metadata: ColumnMetadata | str | None = None
 ) -> pa.Scalar[Any]: ...
 def from_dlpack(
-    managed_tensor: Any, stream: Stream | None = None
+    managed_tensor: Any,
+    stream: Stream | None = None,
+    mr: DeviceMemoryResource | None = None,
 ) -> Table: ...
-def to_dlpack(input: Table, stream: Stream | None = None) -> Any: ...
+def to_dlpack(
+    input: Table,
+    stream: Stream | None = None,
+    mr: DeviceMemoryResource | None = None,
+) -> Any: ...

python/pylibcudf/pylibcudf/interop.pyx

@@ -19,12 +19,13 @@ from pylibcudf.libcudf.interop cimport (
 from pylibcudf.libcudf.table.table cimport table
 
 from rmm.pylibrmm.stream cimport Stream
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 
 from .column cimport Column
 from .scalar cimport Scalar
 from .table cimport Table
 from .types cimport DataType
-from .utils cimport _get_stream
+from .utils cimport _get_stream, _get_memory_resource
 from ._interop_helpers import ColumnMetadata
 
 try:
@@ -136,7 +137,9 @@ if pa is not None:
         return plc_object.to_arrow(metadata=metadata)
 
 
-cpdef Table from_dlpack(object managed_tensor, Stream stream=None):
+cpdef Table from_dlpack(
+    object managed_tensor, Stream stream=None, DeviceMemoryResource mr=None
+):
     """
     Convert a DLPack DLTensor into a cudf table.
 
@@ -148,6 +151,8 @@ cpdef Table from_dlpack(object managed_tensor, Stream stream=None):
         A 1D or 2D column-major (Fortran order) tensor.
     stream : Stream | None
         CUDA stream on which to perform the operation.
+    mr : DeviceMemoryResource | None
+        Device memory resource used to allocate the returned table's device memory.
 
     Returns
     -------
@@ -164,21 +169,22 @@ cpdef Table from_dlpack(object managed_tensor, Stream stream=None):
         raise ValueError("PyCapsule object contained a NULL pointer")
     PyCapsule_SetName(managed_tensor, "used_dltensor")
     stream = _get_stream(stream)
+    mr = _get_memory_resource(mr)
 
     # Note: A copy is always performed when converting the dlpack
     # data to a libcudf table. We also delete the dlpack_tensor pointer
     # as the pointer is not deleted by libcudf's from_dlpack function.
     # TODO: https://github.com/rapidsai/cudf/issues/10874
     # TODO: https://github.com/rapidsai/cudf/issues/10849
     with nogil:
-        c_result = cpp_from_dlpack(dlpack_tensor, stream.view())
+        c_result = cpp_from_dlpack(dlpack_tensor, stream.view(), mr.get_mr())
 
-    cdef Table result = Table.from_libcudf(move(c_result), stream)
+    cdef Table result = Table.from_libcudf(move(c_result), stream, mr)
     dlpack_tensor.deleter(dlpack_tensor)
     return result
 
 
-cpdef object to_dlpack(Table input, Stream stream=None):
+cpdef object to_dlpack(Table input, Stream stream=None, DeviceMemoryResource mr=None):
     """
     Convert a cudf table into a DLPack DLTensor.
 
@@ -190,6 +196,9 @@ cpdef object to_dlpack(Table input, Stream stream=None):
         A 1D or 2D column-major (Fortran order) tensor.
     stream : Stream | None
         CUDA stream on which to perform the operation.
+    mr : DeviceMemoryResource | None
+        Device memory resource used to allocate the returned DLPack tensor's device
+        memory.
 
     Returns
     -------
@@ -204,9 +213,10 @@ cpdef object to_dlpack(Table input, Stream stream=None):
             )
     cdef DLManagedTensor *dlpack_tensor
     stream = _get_stream(stream)
+    mr = _get_memory_resource(mr)
 
     with nogil:
-        dlpack_tensor = cpp_to_dlpack(input.view(), stream.view())
+        dlpack_tensor = cpp_to_dlpack(input.view(), stream.view(), mr.get_mr())
 
     return PyCapsule_New(
         dlpack_tensor,

python/pylibcudf/pylibcudf/join.pyx

@@ -89,7 +89,7 @@ cpdef tuple inner_join(
 
     with nogil:
         c_result = cpp_join.inner_join(
-            left_keys.view(), right_keys.view(), nulls_equal, stream.view()
+            left_keys.view(), right_keys.view(), nulls_equal, stream.view(), mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -130,7 +130,7 @@ cpdef tuple left_join(
 
     with nogil:
         c_result = cpp_join.left_join(
-            left_keys.view(), right_keys.view(), nulls_equal, stream.view()
+            left_keys.view(), right_keys.view(), nulls_equal, stream.view(), mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -171,7 +171,7 @@ cpdef tuple full_join(
 
     with nogil:
         c_result = cpp_join.full_join(
-            left_keys.view(), right_keys.view(), nulls_equal, stream.view()
+            left_keys.view(), right_keys.view(), nulls_equal, stream.view(), mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -214,7 +214,8 @@ cpdef Column left_semi_join(
             left_keys.view(),
             right_keys.view(),
             nulls_equal,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return _column_from_gather_map(move(c_result), stream, mr)
 
@@ -254,7 +255,8 @@ cpdef Column left_anti_join(
             left_keys.view(),
             right_keys.view(),
             nulls_equal,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return _column_from_gather_map(move(c_result), stream, mr)
 
@@ -288,7 +290,9 @@ cpdef Table cross_join(
     mr = _get_memory_resource(mr)
 
     with nogil:
-        result = cpp_join.cross_join(left.view(), right.view(), stream.view())
+        result = cpp_join.cross_join(
+            left.view(), right.view(), stream.view(), mr.get_mr()
+        )
     return Table.from_libcudf(move(result), stream, mr)
 
 
@@ -330,7 +334,8 @@ cpdef tuple conditional_inner_join(
             right.view(),
             dereference(binary_predicate.c_obj.get()),
             output_size,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -376,7 +381,8 @@ cpdef tuple conditional_left_join(
             right.view(),
             dereference(binary_predicate.c_obj.get()),
             output_size,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -420,7 +426,8 @@ cpdef tuple conditional_full_join(
             left.view(),
             right.view(),
             dereference(binary_predicate.c_obj.get()),
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -465,7 +472,8 @@ cpdef Column conditional_left_semi_join(
             right.view(),
             dereference(binary_predicate.c_obj.get()),
             output_size,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return _column_from_gather_map(move(c_result), stream, mr)
 
@@ -507,7 +515,8 @@ cpdef Column conditional_left_anti_join(
             right.view(),
             dereference(binary_predicate.c_obj.get()),
             output_size,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return _column_from_gather_map(move(c_result), stream, mr)
 
@@ -562,7 +571,8 @@ cpdef tuple mixed_inner_join(
             dereference(binary_predicate.c_obj.get()),
             nulls_equal,
             empty_optional,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -620,7 +630,8 @@ cpdef tuple mixed_left_join(
             dereference(binary_predicate.c_obj.get()),
             nulls_equal,
             empty_optional,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -678,7 +689,8 @@ cpdef tuple mixed_full_join(
             dereference(binary_predicate.c_obj.get()),
             nulls_equal,
             empty_optional,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return (
         _column_from_gather_map(move(c_result.first), stream, mr),
@@ -733,7 +745,8 @@ cpdef Column mixed_left_semi_join(
             right_conditional.view(),
             dereference(binary_predicate.c_obj.get()),
             nulls_equal,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return _column_from_gather_map(move(c_result), stream, mr)
 
@@ -785,6 +798,7 @@ cpdef Column mixed_left_anti_join(
             right_conditional.view(),
             dereference(binary_predicate.c_obj.get()),
             nulls_equal,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
     return _column_from_gather_map(move(c_result), stream, mr)

python/pylibcudf/pylibcudf/labeling.pyx

@@ -63,7 +63,8 @@ cpdef Column label_bins(
             left_inclusive,
             right_edges.view(),
             right_inclusive,
-            stream.view()
+            stream.view(),
+            mr.get_mr()
         )
 
     return Column.from_libcudf(move(c_result), stream, mr)

python/pylibcudf/pylibcudf/libcudf/interop.pxd

@@ -12,6 +12,7 @@ from pylibcudf.libcudf.table.table cimport table
 from pylibcudf.libcudf.table.table_view cimport table_view
 
 from rmm.librmm.cuda_stream_view cimport cuda_stream_view
+from rmm.librmm.memory_resource cimport device_memory_resource
 
 
 cdef extern from "dlpack/dlpack.h" nogil:
@@ -38,12 +39,14 @@ cdef extern from "cudf/interop.hpp" namespace "cudf" \
         nogil:
     cdef unique_ptr[table] from_dlpack(
         const DLManagedTensor* managed_tensor,
-        cuda_stream_view stream
+        cuda_stream_view stream,
+        device_memory_resource* mr
     ) except +libcudf_exception_handler
 
     DLManagedTensor* to_dlpack(
         const table_view& input,
-        cuda_stream_view stream
+        cuda_stream_view stream,
+        device_memory_resource* mr
     ) except +libcudf_exception_handler
 
     cdef cppclass column_metadata: