Require passing memory resources to from_libcudf methods (#20171)

Resolves  #14229

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

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

URL: #20171

Author: vyasr

python/cudf/cudf/_lib/strings_udf.pyx

@@ -17,6 +17,7 @@ from pylibcudf.libcudf.strings_udf cimport (
 from rmm.librmm.device_buffer cimport device_buffer
 from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 from rmm.pylibrmm.stream import DEFAULT_STREAM
+from rmm.mr import get_current_device_resource
 
 import numpy as np
 
@@ -38,7 +39,11 @@ def column_from_managed_udf_string_array(DeviceBuffer d_buffer):
     with nogil:
         c_result = move(cpp_column_from_managed_udf_string_array(data, size))
 
-    return plc_Column.from_libcudf(move(c_result), DEFAULT_STREAM)
+    return plc_Column.from_libcudf(
+        move(c_result),
+        DEFAULT_STREAM,
+        get_current_device_resource(),
+    )
 
 
 def get_character_flags_table_ptr():

python/pylibcudf/pylibcudf/column.pxd

@@ -60,7 +60,7 @@ cdef class Column:
     cdef Column from_libcudf(
         unique_ptr[column] libcudf_col,
         Stream stream,
-        DeviceMemoryResource mr=*
+        DeviceMemoryResource mr
     )
 
     @staticmethod

python/pylibcudf/pylibcudf/column.pyx

@@ -604,7 +604,7 @@ cdef class Column:
     cdef Column from_libcudf(
         unique_ptr[column] libcudf_col,
         Stream stream,
-        DeviceMemoryResource mr=None
+        DeviceMemoryResource mr
     ):
         """Create a Column from a libcudf column.
 
@@ -619,8 +619,6 @@ cdef class Column:
 
         cdef column_contents contents = libcudf_col.get().release()
 
-        stream = _get_stream(stream)
-        mr = _get_memory_resource(mr)
         # Note that when converting to cudf Column objects we'll need to pull
         # out the base object.
         cdef gpumemoryview data = gpumemoryview(

python/pylibcudf/pylibcudf/contiguous_split.pxd

@@ -50,5 +50,6 @@ cpdef PackedColumns pack(Table input)
 cpdef Table unpack(PackedColumns input, DeviceMemoryResource mr=*)
 
 cpdef Table unpack_from_memoryviews(
-    memoryview metadata, gpumemoryview gpu_data, DeviceMemoryResource mr=*
+    memoryview metadata,
+    gpumemoryview gpu_data,
 )

python/pylibcudf/pylibcudf/contiguous_split.pyx

@@ -374,7 +374,8 @@ cpdef Table unpack(PackedColumns input, DeviceMemoryResource mr=None):
 
 
 cpdef Table unpack_from_memoryviews(
-    memoryview metadata, gpumemoryview gpu_data, DeviceMemoryResource mr=None
+    memoryview metadata,
+    gpumemoryview gpu_data,
 ):
     """Deserialize the result of `pack`.
 
@@ -388,19 +389,22 @@ cpdef Table unpack_from_memoryviews(
         The packed metadata to unpack.
     gpu_data : gpumemoryview
         The packed gpu_data to unpack.
-    mr : DeviceMemoryResource, optional
-        Device memory resource used to allocate the returned table's device memory.
 
     Returns
     -------
     Table
         Copy of the packed columns.
     """
-    mr = _get_memory_resource(mr)
     if metadata.nbytes == 0:
         if gpu_data.__cuda_array_interface__["data"][0] != 0:
             raise ValueError("Expected an empty gpu_data from unpacking an empty table")
-        return Table.from_libcudf(make_unique[table](table_view()), stream=None, mr=mr)
+        # For an empty table we just attach the default stream and mr since neither will
+        # be used for any operations.
+        return Table.from_libcudf(
+            make_unique[table](table_view()),
+            _get_stream(),
+            _get_memory_resource(),
+        )
 
     # Extract the raw data pointers
     cdef const uint8_t[::1] _metadata = metadata

python/pylibcudf/pylibcudf/io/types.pxd

@@ -85,7 +85,7 @@ cdef class TableWithMetadata:
 
     @staticmethod
     cdef TableWithMetadata from_libcudf(
-        table_with_metadata& tbl, Stream stream=*, DeviceMemoryResource mr=*
+        table_with_metadata& tbl, Stream stream, DeviceMemoryResource mr
     )
 
 cdef class SourceInfo:

python/pylibcudf/pylibcudf/io/types.pyx

@@ -29,7 +29,6 @@ from pylibcudf.libcudf.io.types cimport (
 from pylibcudf.libcudf.types cimport size_type
 from pylibcudf.libcudf.utilities.span cimport host_span, device_span
 
-from pylibcudf.utils cimport _get_stream, _get_memory_resource
 from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 from rmm.pylibrmm.stream cimport Stream
 from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
@@ -395,13 +394,11 @@ cdef class TableWithMetadata:
     @staticmethod
     cdef TableWithMetadata from_libcudf(
         table_with_metadata& tbl_with_meta,
-        Stream stream = None,
-        DeviceMemoryResource mr = None
+        Stream stream,
+        DeviceMemoryResource mr
     ):
         """Create a Python TableWithMetadata from a libcudf table_with_metadata"""
         cdef TableWithMetadata out = TableWithMetadata.__new__(TableWithMetadata)
-        stream = _get_stream(stream)
-        mr = _get_memory_resource(mr)
         out.tbl = Table.from_libcudf(move(tbl_with_meta.tbl), stream, mr)
         out.metadata = tbl_with_meta.metadata
         return out

python/pylibcudf/pylibcudf/table.pxd

@@ -19,8 +19,8 @@ cdef class Table:
     @staticmethod
     cdef Table from_libcudf(
         unique_ptr[table] libcudf_tbl,
-        Stream stream=*,
-        DeviceMemoryResource mr=*
+        Stream stream,
+        DeviceMemoryResource mr
     )
 
     @staticmethod

python/pylibcudf/pylibcudf/table.pyx

@@ -33,7 +33,6 @@ from pylibcudf.libcudf.table.table cimport table
 
 from .column cimport Column
 from .types cimport DataType
-from .utils cimport _get_stream, _get_memory_resource
 from pylibcudf._interop_helpers cimport (
     _release_schema,
     _release_array,
@@ -202,8 +201,8 @@ cdef class Table:
     @staticmethod
     cdef Table from_libcudf(
         unique_ptr[table] libcudf_tbl,
-        Stream stream=None,
-        DeviceMemoryResource mr=None
+        Stream stream,
+        DeviceMemoryResource mr
     ):
         """Create a Table from a libcudf table.
 
@@ -214,8 +213,6 @@ cdef class Table:
         cdef vector[unique_ptr[column]] c_columns = dereference(libcudf_tbl).release()
 
         cdef vector[unique_ptr[column]].size_type i
-        stream = _get_stream(stream)
-        mr = _get_memory_resource(mr)
         return Table([
             Column.from_libcudf(move(c_columns[i]), stream, mr)
             for i in range(c_columns.size())