Passing bool/str arg to cython

rapidsai · Mar 23, 2021 · 8acf8f0 · 8acf8f0
1 parent e5ba902
commit 8acf8f0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 11 deletions.
diff --git a/python/cudf/cudf/_lib/lists.pyx b/python/cudf/cudf/_lib/lists.pyx
@@ -27,7 +27,6 @@ from cudf._lib.table cimport Table
 from cudf._lib.types cimport (
     underlying_type_t_null_order, underlying_type_t_order
 )
-from cudf._lib.types import Order, NullOrder
 from cudf.core.dtypes import ListDtype
 
 from cudf._lib.cpp.lists.extract cimport extract_list_element
@@ -67,13 +66,15 @@ def explode_outer(Table tbl, int explode_column_idx, bool ignore_index=False):
     )
 
 
-def sort_lists(Column col, object order_enum, object null_order_enum):
+def sort_lists(Column col, bool ascending, str na_position):
     cdef shared_ptr[lists_column_view] list_view = (
         make_shared[lists_column_view](col.view())
     )
-    cdef order c_sort_order = <order><underlying_type_t_order>order_enum.value
+    cdef order c_sort_order = (
+        order.ASCENDING if ascending else order.DESCENDING
+    )
     cdef null_order c_null_prec = (
-        <null_order><underlying_type_t_null_order>null_order_enum.value
+        null_order.BEFORE if na_position == "first" else null_order.AFTER
     )
 
     cdef unique_ptr[column] c_result

diff --git a/python/cudf/cudf/core/column/lists.py b/python/cudf/cudf/core/column/lists.py
@@ -8,7 +8,6 @@
 import cudf
 from cudf._lib.copying import segmented_gather
 from cudf._lib.lists import count_elements, extract_element, sort_lists
-from cudf._lib.types import NullOrder, Order
 from cudf.core.buffer import Buffer
 from cudf.core.column import ColumnBase, as_column, column
 from cudf.core.column.methods import ColumnMethodsMixin
@@ -368,12 +367,7 @@ def sort_values(
         if is_list_dtype(self._column.children[1].dtype):
             raise NotImplementedError("Nested lists sort is not supported.")
 
-        sort_order = Order.ASCENDING if ascending else Order.DESCENDING
-        null_order = (
-            NullOrder.BEFORE if na_position == "first" else NullOrder.AFTER
-        )
-
         return self._return_or_inplace(
-            sort_lists(self._column, sort_order, null_order),
+            sort_lists(self._column, ascending, na_position),
             retain_index=not ignore_index,
         )