Avoid private cudf DeviceScalar in favor of using pylibcudf & pyarrow

nv-morpheus · Dec 19, 2024 · 9c1f4a2 · 9c1f4a2
1 parent 0ba9ee0
commit 9c1f4a2
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 23 deletions.
diff --git a/python/morpheus/morpheus/_lib/cudf_helpers.pyx b/python/morpheus/morpheus/_lib/cudf_helpers.pyx
@@ -35,15 +35,10 @@ from cudf._lib.column cimport Column
 
 # isort: off
 
-# imports needed for get_element, which is required by from_column_view_with_fix
-cimport pylibcudf.libcudf.copying as cpp_copying
-from pylibcudf.libcudf.column.column_view cimport column_view
-from libcpp.memory cimport make_unique, unique_ptr
-from pylibcudf.libcudf.scalar.scalar cimport scalar
-from pylibcudf cimport Table as plc_Table
-from cudf._lib.scalar cimport DeviceScalar
-
 # imports needed for from_column_view_with_fix
+import pylibcudf as plc
+from pylibcudf cimport Column as plc_Column, Table as plc_Table
+from pylibcudf.libcudf.column.column_view cimport column_view
 import rmm
 from libc.stdint cimport uintptr_t
 from cudf.core.buffer import (
@@ -64,18 +59,6 @@ from cudf._lib.null_mask import bitmask_allocation_size_bytes
 
 # isort: on
 
-cdef get_element(column_view col_view, size_type index):
-
-    cdef unique_ptr[scalar] c_output
-    with nogil:
-        c_output = move(
-            cpp_copying.get_element(col_view, index)
-        )
-
-    return DeviceScalar.from_unique_ptr(
-        move(c_output), dtype=dtype_from_column_view(col_view)
-    )
-
 cdef Column from_column_view_with_fix(column_view cv, object owner):
     """
     Given a ``cudf::column_view``, constructs a ``cudf.Column`` from it,
@@ -118,9 +101,13 @@ cdef Column from_column_view_with_fix(column_view cv, object owner):
             if offset_child_column.size() == 0:
                 base_nbytes = 0
             else:
-                chars_size = get_element(
-                    offset_child_column, offset_child_column.size()-1).value
-                base_nbytes = chars_size
+                offset_child_column_owner = owner.children[offsets_column_index]
+                plc_owner = offset_child_column_owner.to_pylibcudf(mode="read")
+                plc_offsets_col = plc_Column.from_column_view(
+                    offset_child_column, plc_owner
+                )
+                plc_scalar = plc.copying.get_element(plc_offsets_col, offset_child_column.size()-1)
+                base_nbytes = plc.interop.to_arrow(plc_scalar).as_py()
 
     if data_ptr:
         if data_owner is None:

diff --git a/python/morpheus/morpheus/_lib/cudf_helpers/__init__.pyi b/python/morpheus/morpheus/_lib/cudf_helpers/__init__.pyi
@@ -8,6 +8,7 @@ from cudf.core.dtypes import StructDtype
 import _cython_3_0_11
 import cudf
 import itertools
+import pylibcudf
 import rmm
 
 __all__ = [
@@ -19,6 +20,7 @@ __all__ = [
     "bitmask_allocation_size_bytes",
     "cudf",
     "itertools",
+    "plc",
     "rmm"
 ]