Use pylibcudf enums in cudf Python quantile (#17287)

Shouldn't need to use the "private" `pylibcudf.libcudf` types anymore now that the Python side enums are exposed Authors: - Matthew Roeschke (https://github.com/mroeschke) Approvers: - Matthew Murray (https://github.com/Matt711) URL: #17287
rapidsai · Nov 9, 2024 · e399e95 · e399e95
1 parent 0fc5fab
commit e399e95
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 79 deletions.
diff --git a/python/cudf/cudf/_lib/quantiles.pyx b/python/cudf/cudf/_lib/quantiles.pyx
@@ -6,14 +6,6 @@ from libcpp cimport bool
 from libcpp.vector cimport vector
 
 from cudf._lib.column cimport Column
-from cudf._lib.types cimport (
-    underlying_type_t_interpolation,
-    underlying_type_t_sorted,
-)
-
-from cudf._lib.types import Interpolation
-
-from pylibcudf.libcudf.types cimport interpolation, sorted
 
 from cudf._lib.utils cimport columns_from_pylibcudf_table
 
@@ -28,17 +20,13 @@ def quantile(
     Column ordered_indices,
     bool exact,
 ):
-    cdef interpolation c_interp = <interpolation>(
-        <underlying_type_t_interpolation> Interpolation[interp.upper()]
-    )
-
     return Column.from_pylibcudf(
         plc.quantiles.quantile(
             input.to_pylibcudf(mode="read"),
             q,
-            c_interp,
+            plc.types.Interpolation[interp.upper()],
             ordered_indices.to_pylibcudf(mode="read"),
-            <bool>exact
+            exact
         )
     )
 
@@ -51,22 +39,14 @@ def quantile_table(
     list column_order,
     list null_precedence,
 ):
-
-    cdef interpolation c_interp = <interpolation>(
-        <underlying_type_t_interpolation> interp
-    )
-    cdef sorted c_is_input_sorted = <sorted>(
-        <underlying_type_t_sorted> is_input_sorted
-    )
-
     return columns_from_pylibcudf_table(
         plc.quantiles.quantiles(
             plc.Table([
                 c.to_pylibcudf(mode="read") for c in source_columns
             ]),
             q,
-            c_interp,
-            c_is_input_sorted,
+            interp,
+            is_input_sorted,
             column_order,
             null_precedence
         )

diff --git a/python/cudf/cudf/_lib/types.pxd b/python/cudf/cudf/_lib/types.pxd
@@ -7,12 +7,7 @@ cimport pylibcudf.libcudf.types as libcudf_types
 from pylibcudf.libcudf.column.column_view cimport column_view
 from pylibcudf.libcudf.lists.lists_column_view cimport lists_column_view
 
-ctypedef bool underlying_type_t_order
-ctypedef bool underlying_type_t_null_order
-ctypedef bool underlying_type_t_sorted
-ctypedef int32_t underlying_type_t_interpolation
 ctypedef int32_t underlying_type_t_type_id
-ctypedef bool underlying_type_t_null_policy
 
 cdef dtype_from_column_view(column_view cv)
 

diff --git a/python/cudf/cudf/_lib/types.pyx b/python/cudf/cudf/_lib/types.pyx
@@ -11,12 +11,6 @@ cimport pylibcudf.libcudf.types as libcudf_types
 from pylibcudf.libcudf.column.column_view cimport column_view
 from pylibcudf.libcudf.lists.lists_column_view cimport lists_column_view
 
-from cudf._lib.types cimport (
-    underlying_type_t_interpolation,
-    underlying_type_t_order,
-    underlying_type_t_sorted,
-)
-
 import pylibcudf
 
 import cudf
@@ -151,44 +145,6 @@ datetime_unit_map = {
 size_type_dtype = LIBCUDF_TO_SUPPORTED_NUMPY_TYPES[pylibcudf.types.SIZE_TYPE_ID]
 
 
-class Interpolation(IntEnum):
-    LINEAR = (
-        <underlying_type_t_interpolation> libcudf_types.interpolation.LINEAR
-    )
-    LOWER = (
-        <underlying_type_t_interpolation> libcudf_types.interpolation.LOWER
-    )
-    HIGHER = (
-        <underlying_type_t_interpolation> libcudf_types.interpolation.HIGHER
-    )
-    MIDPOINT = (
-        <underlying_type_t_interpolation> libcudf_types.interpolation.MIDPOINT
-    )
-    NEAREST = (
-        <underlying_type_t_interpolation> libcudf_types.interpolation.NEAREST
-    )
-
-
-class Order(IntEnum):
-    ASCENDING = <underlying_type_t_order> libcudf_types.order.ASCENDING
-    DESCENDING = <underlying_type_t_order> libcudf_types.order.DESCENDING
-
-
-class Sorted(IntEnum):
-    YES = <underlying_type_t_sorted> libcudf_types.sorted.YES
-    NO = <underlying_type_t_sorted> libcudf_types.sorted.NO
-
-
-class NullOrder(IntEnum):
-    BEFORE = <underlying_type_t_order> libcudf_types.null_order.BEFORE
-    AFTER = <underlying_type_t_order> libcudf_types.null_order.AFTER
-
-
-class NullHandling(IntEnum):
-    INCLUDE = <underlying_type_t_null_policy> libcudf_types.null_policy.INCLUDE
-    EXCLUDE = <underlying_type_t_null_policy> libcudf_types.null_policy.EXCLUDE
-
-
 cdef dtype_from_lists_column_view(column_view cv):
     # lists_column_view have no default constructor, so we heap
     # allocate it to get around Cython's limitation of requiring

diff --git a/python/cudf/cudf/core/frame.py b/python/cudf/cudf/core/frame.py
@@ -16,6 +16,8 @@
 import pyarrow as pa
 from typing_extensions import Self
 
+import pylibcudf as plc
+
 import cudf
 from cudf import _lib as libcudf
 from cudf.api.types import is_dtype_equal, is_scalar
@@ -789,15 +791,13 @@ def _quantile_table(
         column_order=(),
         null_precedence=(),
     ):
-        interpolation = libcudf.types.Interpolation[interpolation]
+        interpolation = plc.types.Interpolation[interpolation]
 
-        is_sorted = libcudf.types.Sorted["YES" if is_sorted else "NO"]
+        is_sorted = plc.types.Sorted["YES" if is_sorted else "NO"]
 
-        column_order = [libcudf.types.Order[key] for key in column_order]
+        column_order = [plc.types.Order[key] for key in column_order]
 
-        null_precedence = [
-            libcudf.types.NullOrder[key] for key in null_precedence
-        ]
+        null_precedence = [plc.types.NullOrder[key] for key in null_precedence]
 
         return self._from_columns_like_self(
             libcudf.quantiles.quantile_table(