diff --git a/python/pylibcudf/pylibcudf/interop.pxd b/python/pylibcudf/pylibcudf/interop.pxd index a02261db74c..a46d720ae91 100644 --- a/python/pylibcudf/pylibcudf/interop.pxd +++ b/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=*) diff --git a/python/pylibcudf/pylibcudf/interop.pyi b/python/pylibcudf/pylibcudf/interop.pyi index decd0e412c4..a2e650b2dbd 100644 --- a/python/pylibcudf/pylibcudf/interop.pyi +++ b/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: ... diff --git a/python/pylibcudf/pylibcudf/interop.pyx b/python/pylibcudf/pylibcudf/interop.pyx index 75edb59b093..86e265e898b 100644 --- a/python/pylibcudf/pylibcudf/interop.pyx +++ b/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,6 +169,7 @@ 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 @@ -171,14 +177,14 @@ cpdef Table from_dlpack(object managed_tensor, Stream stream=None): # 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, diff --git a/python/pylibcudf/pylibcudf/join.pyx b/python/pylibcudf/pylibcudf/join.pyx index ed631d991cd..404fffa7033 100644 --- a/python/pylibcudf/pylibcudf/join.pyx +++ b/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) diff --git a/python/pylibcudf/pylibcudf/labeling.pyx b/python/pylibcudf/pylibcudf/labeling.pyx index 2c7694f06f8..0617806b225 100644 --- a/python/pylibcudf/pylibcudf/labeling.pyx +++ b/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) diff --git a/python/pylibcudf/pylibcudf/libcudf/interop.pxd b/python/pylibcudf/pylibcudf/libcudf/interop.pxd index 301ae41d667..b80431c960d 100644 --- a/python/pylibcudf/pylibcudf/libcudf/interop.pxd +++ b/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: diff --git a/python/pylibcudf/pylibcudf/libcudf/join.pxd b/python/pylibcudf/pylibcudf/libcudf/join.pxd index 111576ea1d9..a637270dce4 100644 --- a/python/pylibcudf/pylibcudf/libcudf/join.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/join.pxd @@ -13,6 +13,7 @@ from pylibcudf.libcudf.table.table cimport table from pylibcudf.libcudf.table.table_view cimport table_view from pylibcudf.libcudf.types cimport null_equality, size_type from rmm.librmm.cuda_stream_view cimport cuda_stream_view +from rmm.librmm.memory_resource cimport device_memory_resource from rmm.librmm.device_uvector cimport device_uvector from pylibcudf.libcudf.utilities.span cimport device_span @@ -25,80 +26,95 @@ cdef extern from "cudf/join/join.hpp" namespace "cudf" nogil: cdef gather_map_pair_type inner_join( const table_view left_keys, const table_view right_keys, - cuda_stream_view stream + null_equality nulls_equal, + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type left_join( const table_view left_keys, const table_view right_keys, - cuda_stream_view stream + null_equality nulls_equal, + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type full_join( const table_view left_keys, const table_view right_keys, - cuda_stream_view stream + null_equality nulls_equal, + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type left_semi_join( const table_view left_keys, const table_view right_keys, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type left_anti_join( const table_view left_keys, const table_view right_keys, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type inner_join( const table_view left_keys, const table_view right_keys, null_equality nulls_equal, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type left_join( const table_view left_keys, const table_view right_keys, null_equality nulls_equal, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type full_join( const table_view left_keys, const table_view right_keys, null_equality nulls_equal, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type left_semi_join( const table_view left_keys, const table_view right_keys, null_equality nulls_equal, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type left_anti_join( const table_view left_keys, const table_view right_keys, null_equality nulls_equal, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef unique_ptr[table] cross_join( const table_view left, const table_view right, - cuda_stream_view stream - ) except + + cuda_stream_view stream, + device_memory_resource* mr + ) except +libcudf_exception_handler cdef extern from "cudf/join/conditional_join.hpp" namespace "cudf" nogil: cdef gather_map_pair_type conditional_inner_join( const table_view left, const table_view right, const expression binary_predicate, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type conditional_inner_join( @@ -106,14 +122,16 @@ cdef extern from "cudf/join/conditional_join.hpp" namespace "cudf" nogil: const table_view right, const expression binary_predicate, optional[size_t] output_size, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type conditional_left_join( const table_view left, const table_view right, const expression binary_predicate, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type conditional_left_join( @@ -121,21 +139,24 @@ cdef extern from "cudf/join/conditional_join.hpp" namespace "cudf" nogil: const table_view right, const expression binary_predicate, optional[size_t] output_size, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type conditional_full_join( const table_view left, const table_view right, const expression binary_predicate, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type conditional_left_semi_join( const table_view left, const table_view right, const expression binary_predicate, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type conditional_left_semi_join( @@ -143,14 +164,16 @@ cdef extern from "cudf/join/conditional_join.hpp" namespace "cudf" nogil: const table_view right, const expression binary_predicate, optional[size_t] output_size, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type conditional_left_anti_join( const table_view left, const table_view right, const expression binary_predicate, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type conditional_left_anti_join( @@ -158,7 +181,8 @@ cdef extern from "cudf/join/conditional_join.hpp" namespace "cudf" nogil: const table_view right, const expression binary_predicate, optional[size_t] output_size, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef extern from "cudf/join/mixed_join.hpp" namespace "cudf" nogil: @@ -170,7 +194,8 @@ cdef extern from "cudf/join/mixed_join.hpp" namespace "cudf" nogil: const expression binary_predicate, null_equality compare_nulls, output_size_data_type output_size_data, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type mixed_left_join( @@ -181,7 +206,8 @@ cdef extern from "cudf/join/mixed_join.hpp" namespace "cudf" nogil: const expression binary_predicate, null_equality compare_nulls, output_size_data_type output_size_data, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_pair_type mixed_full_join( @@ -192,7 +218,8 @@ cdef extern from "cudf/join/mixed_join.hpp" namespace "cudf" nogil: const expression binary_predicate, null_equality compare_nulls, output_size_data_type output_size_data, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type mixed_left_semi_join( @@ -202,7 +229,8 @@ cdef extern from "cudf/join/mixed_join.hpp" namespace "cudf" nogil: const table_view right_conditional, const expression binary_predicate, null_equality compare_nulls, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler cdef gather_map_type mixed_left_anti_join( @@ -212,5 +240,6 @@ cdef extern from "cudf/join/mixed_join.hpp" namespace "cudf" nogil: const table_view right_conditional, const expression binary_predicate, null_equality compare_nulls, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler diff --git a/python/pylibcudf/pylibcudf/libcudf/labeling.pxd b/python/pylibcudf/pylibcudf/libcudf/labeling.pxd index d2681c46b15..70b61fc8d82 100644 --- a/python/pylibcudf/pylibcudf/libcudf/labeling.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/labeling.pxd @@ -6,6 +6,7 @@ from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view from rmm.librmm.cuda_stream_view cimport cuda_stream_view +from rmm.librmm.memory_resource cimport device_memory_resource cdef extern from "cudf/labeling/label_bins.hpp" namespace "cudf" nogil: @@ -19,5 +20,6 @@ cdef extern from "cudf/labeling/label_bins.hpp" namespace "cudf" nogil: inclusive left_inclusive, const column_view &right_edges, inclusive right_inclusive, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler diff --git a/python/pylibcudf/pylibcudf/libcudf/merge.pxd b/python/pylibcudf/pylibcudf/libcudf/merge.pxd index fa17e1000f9..e80f56a47ea 100644 --- a/python/pylibcudf/pylibcudf/libcudf/merge.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/merge.pxd @@ -7,6 +7,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 "cudf/merge.hpp" namespace "cudf" nogil: @@ -15,5 +16,6 @@ cdef extern from "cudf/merge.hpp" namespace "cudf" nogil: vector[libcudf_types.size_type] key_cols, vector[libcudf_types.order] column_order, vector[libcudf_types.null_order] null_precedence, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler diff --git a/python/pylibcudf/pylibcudf/libcudf/round.pxd b/python/pylibcudf/pylibcudf/libcudf/round.pxd index 58bc9b5a4dc..4c7812d967d 100644 --- a/python/pylibcudf/pylibcudf/libcudf/round.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/round.pxd @@ -6,6 +6,7 @@ from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view from rmm.librmm.cuda_stream_view cimport cuda_stream_view +from rmm.librmm.memory_resource cimport device_memory_resource cdef extern from "cudf/round.hpp" namespace "cudf" nogil: @@ -18,5 +19,6 @@ cdef extern from "cudf/round.hpp" namespace "cudf" nogil: const column_view& input, int32_t decimal_places, rounding_method method, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler diff --git a/python/pylibcudf/pylibcudf/libcudf/transpose.pxd b/python/pylibcudf/pylibcudf/libcudf/transpose.pxd index ab0944914f2..b428b6e61b7 100644 --- a/python/pylibcudf/pylibcudf/libcudf/transpose.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/transpose.pxd @@ -6,6 +6,7 @@ from pylibcudf.libcudf.column.column cimport column 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 "cudf/transpose.hpp" namespace "cudf" nogil: @@ -14,5 +15,6 @@ cdef extern from "cudf/transpose.hpp" namespace "cudf" nogil: table_view ] transpose( table_view input_table, - cuda_stream_view stream + cuda_stream_view stream, + device_memory_resource* mr ) except +libcudf_exception_handler diff --git a/python/pylibcudf/pylibcudf/merge.pxd b/python/pylibcudf/pylibcudf/merge.pxd index 0cfe40dbe2c..6c3bf0d13e4 100644 --- a/python/pylibcudf/pylibcudf/merge.pxd +++ b/python/pylibcudf/pylibcudf/merge.pxd @@ -3,6 +3,7 @@ from .table cimport Table from rmm.pylibrmm.stream cimport Stream +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource cpdef Table merge ( @@ -10,5 +11,6 @@ cpdef Table merge ( list key_cols, list column_order, list null_precedence, - Stream stream=* + Stream stream=*, + DeviceMemoryResource mr=* ) diff --git a/python/pylibcudf/pylibcudf/merge.pyi b/python/pylibcudf/pylibcudf/merge.pyi index 26cfdc6ea0c..009a1e83642 100644 --- a/python/pylibcudf/pylibcudf/merge.pyi +++ b/python/pylibcudf/pylibcudf/merge.pyi @@ -1,5 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. +from rmm.pylibrmm.memory_resource import DeviceMemoryResource from rmm.pylibrmm.stream import Stream from pylibcudf.table import Table @@ -11,4 +12,5 @@ def merge( column_order: list[Order], null_precedence: list[NullOrder], stream: Stream | None = None, + mr: DeviceMemoryResource | None = None, ) -> Table: ... diff --git a/python/pylibcudf/pylibcudf/merge.pyx b/python/pylibcudf/pylibcudf/merge.pyx index cc585b11cc5..87db65a195d 100644 --- a/python/pylibcudf/pylibcudf/merge.pyx +++ b/python/pylibcudf/pylibcudf/merge.pyx @@ -9,9 +9,10 @@ from pylibcudf.libcudf.table.table_view cimport table_view from pylibcudf.libcudf.types cimport null_order, order, size_type from rmm.pylibrmm.stream cimport Stream +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .table cimport Table -from .utils cimport _get_stream +from .utils cimport _get_stream, _get_memory_resource __all__ = ["merge"] @@ -20,7 +21,8 @@ cpdef Table merge ( list key_cols, list column_order, list null_precedence, - Stream stream=None + Stream stream=None, + DeviceMemoryResource mr=None ): """Merge a set of sorted tables. @@ -38,6 +40,8 @@ cpdef Table merge ( Whether nulls should come before or after non-nulls. 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 ------- @@ -54,6 +58,7 @@ cpdef Table merge ( cdef unique_ptr[table] c_result stream = _get_stream(stream) + mr = _get_memory_resource(mr) with nogil: c_result = cpp_merge.merge( @@ -61,6 +66,7 @@ cpdef Table merge ( c_key_cols, c_column_order, c_null_precedence, - stream.view() + stream.view(), + mr.get_mr() ) - return Table.from_libcudf(move(c_result), stream) + return Table.from_libcudf(move(c_result), stream, mr) diff --git a/python/pylibcudf/pylibcudf/round.pyx b/python/pylibcudf/pylibcudf/round.pyx index d7be3d14fe0..8edc675fd82 100644 --- a/python/pylibcudf/pylibcudf/round.pyx +++ b/python/pylibcudf/pylibcudf/round.pyx @@ -57,7 +57,8 @@ cpdef Column round( source.view(), decimal_places, round_method, - stream.view() + stream.view(), + mr.get_mr() ) return Column.from_libcudf(move(c_result), stream, mr) diff --git a/python/pylibcudf/pylibcudf/transpose.pxd b/python/pylibcudf/pylibcudf/transpose.pxd index 28db765e48f..eba8989c4c5 100644 --- a/python/pylibcudf/pylibcudf/transpose.pxd +++ b/python/pylibcudf/pylibcudf/transpose.pxd @@ -2,6 +2,7 @@ from .table cimport Table from rmm.pylibrmm.stream cimport Stream +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource -cpdef Table transpose(Table input_table, Stream stream=*) +cpdef Table transpose(Table input_table, Stream stream=*, DeviceMemoryResource mr=*) diff --git a/python/pylibcudf/pylibcudf/transpose.pyi b/python/pylibcudf/pylibcudf/transpose.pyi index 4acbac7ea52..7de87a96129 100644 --- a/python/pylibcudf/pylibcudf/transpose.pyi +++ b/python/pylibcudf/pylibcudf/transpose.pyi @@ -1,7 +1,12 @@ # Copyright (c) 2024, NVIDIA CORPORATION. +from rmm.pylibrmm.memory_resource import DeviceMemoryResource from rmm.pylibrmm.stream import Stream from pylibcudf.table import Table -def transpose(input_table: Table, stream: Stream | None = None) -> Table: ... +def transpose( + input_table: Table, + stream: Stream | None = None, + mr: DeviceMemoryResource | None = None, +) -> Table: ... diff --git a/python/pylibcudf/pylibcudf/transpose.pyx b/python/pylibcudf/pylibcudf/transpose.pyx index cbb23e2358e..b49bcadcde6 100644 --- a/python/pylibcudf/pylibcudf/transpose.pyx +++ b/python/pylibcudf/pylibcudf/transpose.pyx @@ -7,14 +7,17 @@ from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.table.table_view cimport table_view from rmm.pylibrmm.stream cimport Stream +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .column cimport Column from .table cimport Table -from .utils cimport _get_stream +from .utils cimport _get_stream, _get_memory_resource __all__ = ["transpose"] -cpdef Table transpose(Table input_table, Stream stream=None): +cpdef Table transpose( + Table input_table, Stream stream=None, DeviceMemoryResource mr=None +): """Transpose a Table. For details, see :cpp:func:`transpose`. @@ -25,6 +28,8 @@ cpdef Table transpose(Table input_table, Stream stream=None): Table to transpose 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 ------- @@ -34,12 +39,15 @@ cpdef Table transpose(Table input_table, Stream stream=None): cdef pair[unique_ptr[column], table_view] c_result cdef Table owner_table stream = _get_stream(stream) + mr = _get_memory_resource(mr) with nogil: - c_result = cpp_transpose.transpose(input_table.view(), stream.view()) + c_result = cpp_transpose.transpose( + input_table.view(), stream.view(), mr.get_mr() + ) owner_table = Table( - [Column.from_libcudf(move(c_result.first), stream)] * + [Column.from_libcudf(move(c_result.first), stream, mr)] * c_result.second.num_columns() )