Add memory resources to reduce, column, column_factories, and contiguous_split (#20135)

Contributes to #15170

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

Approvers:
  - Bradley Dice (https://github.com/bdice)

URL: #20135

Author: vyasr

python/pylibcudf/pylibcudf/column.pxd

@@ -87,7 +87,7 @@ cdef class Column:
     cpdef gpumemoryview data(self)
     cpdef gpumemoryview null_mask(self)
     cpdef list children(self)
-    cpdef Column copy(self, Stream stream=*)
+    cpdef Column copy(self, Stream stream=*, DeviceMemoryResource mr=*)
     cpdef uint64_t device_buffer_size(self)
     cpdef Column with_mask(self, gpumemoryview, size_type)
 

python/pylibcudf/pylibcudf/column.pyi

@@ -52,15 +52,22 @@ class Column:
     def data(self) -> gpumemoryview | None: ...
     def null_mask(self) -> gpumemoryview | None: ...
     def children(self) -> list[Column]: ...
-    def copy(self, stream: Stream | None = None) -> Column: ...
+    def copy(
+        self,
+        stream: Stream | None = None,
+        mr: DeviceMemoryResource | None = None,
+    ) -> Column: ...
     def device_buffer_size(self) -> int: ...
     def with_mask(
         self, mask: gpumemoryview | None, null_count: int
     ) -> Column: ...
     def list_view(self) -> ListColumnView: ...
     @staticmethod
     def from_scalar(
-        scalar: Scalar, size: int, stream: Stream | None = None
+        scalar: Scalar,
+        size: int,
+        stream: Stream | None = None,
+        mr: DeviceMemoryResource | None = None,
     ) -> Column: ...
     def to_scalar(
         self,
@@ -69,7 +76,10 @@ class Column:
     ) -> Scalar: ...
     @staticmethod
     def all_null_like(
-        like: Column, size: int, stream: Stream | None = None
+        like: Column,
+        size: int,
+        stream: Stream | None = None,
+        mr: DeviceMemoryResource | None = None,
     ) -> Column: ...
     @staticmethod
     def from_rmm_buffer(

python/pylibcudf/pylibcudf/column.pyx

@@ -747,7 +747,12 @@ cdef class Column:
         )
 
     @staticmethod
-    def from_scalar(Scalar slr, size_type size, Stream stream=None):
+    def from_scalar(
+        Scalar slr,
+        size_type size,
+        Stream stream=None,
+        DeviceMemoryResource mr=None,
+    ):
         """Create a Column from a Scalar.
 
         Parameters
@@ -767,13 +772,15 @@ cdef class Column:
         cdef const scalar* c_scalar = slr.get()
         cdef unique_ptr[column] c_result
         stream = _get_stream(stream)
+        mr = _get_memory_resource(mr)
         with nogil:
             c_result = make_column_from_scalar(
                 dereference(c_scalar),
                 size,
-                stream.view()
+                stream.view(),
+                mr.get_mr()
             )
-        return Column.from_libcudf(move(c_result), stream)
+        return Column.from_libcudf(move(c_result), stream, mr)
 
     cpdef Scalar to_scalar(self, Stream stream=None, DeviceMemoryResource mr=None):
         """
@@ -807,7 +814,12 @@ cdef class Column:
         return Scalar.from_libcudf(move(result))
 
     @staticmethod
-    def all_null_like(Column like, size_type size, Stream stream=None):
+    def all_null_like(
+        Column like,
+        size_type size,
+        Stream stream=None,
+        DeviceMemoryResource mr=None,
+    ):
         """Create an all null column from a template.
 
         Parameters
@@ -827,13 +839,15 @@ cdef class Column:
         cdef Scalar slr = Scalar.empty_like(like)
         cdef unique_ptr[column] c_result
         stream = _get_stream(stream)
+        mr = _get_memory_resource(mr)
         with nogil:
             c_result = make_column_from_scalar(
                 dereference(slr.get()),
                 size,
-                stream.view()
+                stream.view(),
+                mr.get_mr()
             )
-        return Column.from_libcudf(move(c_result), stream)
+        return Column.from_libcudf(move(c_result), stream, mr)
 
     @staticmethod
     cdef Column _wrap_nested_list_column(
@@ -1250,13 +1264,14 @@ cdef class Column:
         """The children of the column."""
         return self._children
 
-    cpdef Column copy(self, Stream stream=None):
+    cpdef Column copy(self, Stream stream=None, DeviceMemoryResource mr=None):
         """Create a copy of the column."""
         cdef unique_ptr[column] c_result
         stream = _get_stream(stream)
+        mr = _get_memory_resource(mr)
         with nogil:
             c_result = make_unique[column](self.view(), stream.view())
-        return Column.from_libcudf(move(c_result), stream)
+        return Column.from_libcudf(move(c_result), stream, mr)
 
     cpdef uint64_t device_buffer_size(self):
         """

python/pylibcudf/pylibcudf/column_factories.pxd

@@ -1,5 +1,6 @@
 # Copyright (c) 2024-2025, NVIDIA CORPORATION.
 from pylibcudf.libcudf.types cimport mask_state
+from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource
 from rmm.pylibrmm.stream cimport Stream
 
 from .column cimport Column
@@ -18,39 +19,45 @@ ctypedef fused MaskArg:
 cpdef Column make_empty_column(
     MakeEmptyColumnOperand type_or_id,
     Stream stream=*,
+    DeviceMemoryResource mr=*,
 )
 
 cpdef Column make_numeric_column(
     DataType type_,
     size_type size,
     MaskArg mask,
     Stream stream = *,
+    DeviceMemoryResource mr = *,
 )
 
 cpdef Column make_fixed_point_column(
     DataType type_,
     size_type size,
     MaskArg mask,
     Stream stream = *,
+    DeviceMemoryResource mr = *,
 )
 
 cpdef Column make_timestamp_column(
     DataType type_,
     size_type size,
     MaskArg mask,
     Stream stream = *,
+    DeviceMemoryResource mr = *,
 )
 
 cpdef Column make_duration_column(
     DataType type_,
     size_type size,
     MaskArg mask,
     Stream stream = *,
+    DeviceMemoryResource mr = *,
 )
 
 cpdef Column make_fixed_width_column(
     DataType type_,
     size_type size,
     MaskArg mask,
     Stream stream = *,
+    DeviceMemoryResource mr = *,
 )

python/pylibcudf/pylibcudf/column_factories.pyi

@@ -1,4 +1,5 @@
 # Copyright (c) 2024, NVIDIA CORPORATION.
+from rmm.pylibrmm.memory_resource import DeviceMemoryResource
 from rmm.pylibrmm.stream import Stream
 
 from pylibcudf.column import Column
@@ -7,19 +8,40 @@ from pylibcudf.types import DataType, MaskState, TypeId
 def make_empty_column(
     type_or_id: DataType | TypeId,
     stream: Stream | None = None,
+    mr: DeviceMemoryResource | None = None,
 ) -> Column: ...
 def make_numeric_column(
-    type_: DataType, size: int, mstate: MaskState, stream: Stream | None = None
+    type_: DataType,
+    size: int,
+    mstate: MaskState,
+    stream: Stream | None = None,
+    mr: DeviceMemoryResource | None = None,
 ) -> Column: ...
 def make_fixed_point_column(
-    type_: DataType, size: int, mstate: MaskState, stream: Stream | None = None
+    type_: DataType,
+    size: int,
+    mstate: MaskState,
+    stream: Stream | None = None,
+    mr: DeviceMemoryResource | None = None,
 ) -> Column: ...
 def make_timestamp_column(
-    type_: DataType, size: int, mstate: MaskState, stream: Stream | None = None
+    type_: DataType,
+    size: int,
+    mstate: MaskState,
+    stream: Stream | None = None,
+    mr: DeviceMemoryResource | None = None,
 ) -> Column: ...
 def make_duration_column(
-    type_: DataType, size: int, mstate: MaskState, stream: Stream | None = None
+    type_: DataType,
+    size: int,
+    mstate: MaskState,
+    stream: Stream | None = None,
+    mr: DeviceMemoryResource | None = None,
 ) -> Column: ...
 def make_fixed_width_column(
-    type_: DataType, size: int, mstate: MaskState, stream: Stream | None = None
+    type_: DataType,
+    size: int,
+    mstate: MaskState,
+    stream: Stream | None = None,
+    mr: DeviceMemoryResource | None = None,
 ) -> Column: ...