split imports and reference rmm.DeviceBuffer correctly

rapidsai · Oct 7, 2024 · 0309aee · 0309aee
1 parent 38bdc9e
commit 0309aee
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 12 deletions.
diff --git a/docs/cudf/source/conf.py b/docs/cudf/source/conf.py
@@ -342,7 +342,7 @@ def clean_all_xml_files(path):
     "cudf.Series": ("cudf.core.series.Series", "cudf.Series"),
     "cudf.Index": ("cudf.core.index.Index", "cudf.Index"),
     "cupy.core.core.ndarray": ("cupy.ndarray", "cupy.ndarray"),
-    "DeviceBuffer": ("rmm.pylibrmm.device_buffer.DeviceBuffer", "rmm.DeviceBuffer"),
+    "DeviceBuffer": ("rmm.DeviceBuffer", "rmm.DeviceBuffer"),
 }
 
 

diff --git a/python/cudf/cudf/_lib/strings_udf.pyx b/python/cudf/cudf/_lib/strings_udf.pyx
@@ -23,7 +23,8 @@ from pylibcudf.libcudf.strings_udf cimport (
     to_string_view_array as cpp_to_string_view_array,
     udf_string,
 )
-from rmm.pylibrmm.device_buffer cimport DeviceBuffer, device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 
 from cudf._lib.column cimport Column
 

diff --git a/python/pylibcudf/pylibcudf/libcudf/strings_udf.pxd b/python/pylibcudf/pylibcudf/libcudf/strings_udf.pxd
@@ -8,7 +8,7 @@ from pylibcudf.libcudf.column.column cimport column
 from pylibcudf.libcudf.column.column_view cimport column_view
 from pylibcudf.libcudf.types cimport size_type
 
-from rmm.pylibrmm.device_buffer cimport DeviceBuffer, device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
 
 
 cdef extern from "cudf/strings/udf/udf_string.hpp" namespace \

diff --git a/python/pylibcudf/pylibcudf/null_mask.pyx b/python/pylibcudf/pylibcudf/null_mask.pyx
@@ -6,12 +6,11 @@ from libcpp.utility cimport move
 from pylibcudf.libcudf cimport null_mask as cpp_null_mask
 from pylibcudf.libcudf.types cimport mask_state, size_type
 
-from rmm.pylibrmm.device_buffer cimport DeviceBuffer, device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 
 from pylibcudf.libcudf.types import mask_state as MaskState  # no-cython-lint
 
-import rmm  # no-cython-lint
-
 from .column cimport Column
 from .table cimport Table
 
@@ -21,7 +20,7 @@ cdef DeviceBuffer buffer_to_python(device_buffer buf):
 
 
 cpdef DeviceBuffer copy_bitmask(Column col):
-    """Copies ``col``'s bitmask into a ``rmm.DeviceBuffer``.
+    """Copies ``col``'s bitmask into a ``DeviceBuffer``.
 
     For details, see :cpp:func:`copy_bitmask`.
 
@@ -33,8 +32,8 @@ cpdef DeviceBuffer copy_bitmask(Column col):
     Returns
     -------
     rmm.DeviceBuffer
-        A ``rmm.DeviceBuffer`` containing ``col``'s bitmask, or an empty
-        ``rmm.DeviceBuffer`` if ``col`` is not nullable
+        A ``DeviceBuffer`` containing ``col``'s bitmask, or an empty
+        ``DeviceBuffer`` if ``col`` is not nullable
     """
     cdef device_buffer db
 
@@ -68,7 +67,7 @@ cpdef DeviceBuffer create_null_mask(
     size_type size,
     mask_state state = mask_state.UNINITIALIZED
 ):
-    """Creates a ``rmm.DeviceBuffer`` for use as a null value indicator bitmask of a
+    """Creates a ``DeviceBuffer`` for use as a null value indicator bitmask of a
     ``Column``.
 
     For details, see :cpp:func:`create_null_mask`.
@@ -85,7 +84,7 @@ cpdef DeviceBuffer create_null_mask(
     Returns
     -------
     rmm.DeviceBuffer
-        A ``rmm.DeviceBuffer`` for use as a null bitmask satisfying the desired size and
+        A ``DeviceBuffer`` for use as a null bitmask satisfying the desired size and
         state
     """
     cdef device_buffer db

diff --git a/python/pylibcudf/pylibcudf/transform.pyx b/python/pylibcudf/pylibcudf/transform.pyx
@@ -9,7 +9,8 @@ from pylibcudf.libcudf.table.table cimport table
 from pylibcudf.libcudf.table.table_view cimport table_view
 from pylibcudf.libcudf.types cimport bitmask_type, size_type
 
-from rmm.pylibrmm.device_buffer cimport DeviceBuffer, device_buffer
+from rmm.librmm.device_buffer cimport device_buffer
+from rmm.pylibrmm.device_buffer cimport DeviceBuffer
 
 from .column cimport Column
 from .gpumemoryview cimport gpumemoryview