Adjust search_sorted output return

rapidsai · Jan 27, 2025 · b3aa5cd · b3aa5cd
1 parent 622b4ef
commit b3aa5cd
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 32 deletions.
diff --git a/python/cudf/cudf/core/_internals/search.py b/python/cudf/cudf/core/_internals/search.py
@@ -1,12 +1,14 @@
 # Copyright (c) 2020-2025, NVIDIA CORPORATION.
 from __future__ import annotations
 
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 import pylibcudf as plc
 
 from cudf.core.buffer import acquire_spill_lock
-from cudf.core.column import ColumnBase
+
+if TYPE_CHECKING:
+    from cudf.core.column import ColumnBase
 
 
 @acquire_spill_lock()
@@ -16,7 +18,7 @@ def search_sorted(
     side: Literal["left", "right"],
     ascending: bool = True,
     na_position: Literal["first", "last"] = "last",
-) -> ColumnBase:
+) -> plc.Column:
     """Find indices where elements should be inserted to maintain order
 
     Parameters
@@ -43,11 +45,9 @@ def search_sorted(
         plc.search,
         "lower_bound" if side == "left" else "upper_bound",
     )
-    return ColumnBase.from_pylibcudf(
-        func(
-            plc.Table([col.to_pylibcudf(mode="read") for col in source]),
-            plc.Table([col.to_pylibcudf(mode="read") for col in values]),
-            column_order,
-            null_precedence,
-        )
+    return func(
+        plc.Table([col.to_pylibcudf(mode="read") for col in source]),
+        plc.Table([col.to_pylibcudf(mode="read") for col in values]),
+        column_order,
+        null_precedence,
     )
diff --git a/python/cudf/cudf/core/column/column.py b/python/cudf/cudf/core/column/column.py
@@ -1812,12 +1812,14 @@ def searchsorted(
             raise ValueError(
                 "Column searchsorted expects values to be column of same dtype"
             )
-        return search.search_sorted(  # type: ignore[return-value]
-            [self],
-            [value],
-            side=side,
-            ascending=ascending,
-            na_position=na_position,
+        return ColumnBase.from_pylibcudf(
+            search.search_sorted(  # type: ignore[return-value]
+                [self],
+                [value],
+                side=side,
+                ascending=ascending,
+                na_position=na_position,
+            )
         )
 
     def unique(self) -> Self:

diff --git a/python/cudf/cudf/core/frame.py b/python/cudf/cudf/core/frame.py
@@ -1348,12 +1348,14 @@ def searchsorted(
             for val, common_dtype in zip(values, common_dtype_list)
         ]
 
-        outcol = search.search_sorted(
-            sources,
-            values,
-            side,
-            ascending=ascending,
-            na_position=na_position,
+        outcol = ColumnBase.from_pylibcudf(
+            search.search_sorted(
+                sources,
+                values,
+                side,
+                ascending=ascending,
+                na_position=na_position,
+            )
         )
 
         # Return result as cupy array if the values is non-scalar

diff --git a/python/cudf/cudf/core/index.py b/python/cudf/cudf/core/index.py
@@ -121,17 +121,21 @@ def _lexsorted_equal_range(
     else:
         sort_inds = None
         sort_vals = idx
-    lower_bound = search.search_sorted(
-        list(sort_vals._columns),
-        keys,
-        side="left",
-        ascending=sort_vals.is_monotonic_increasing,
+    lower_bound = ColumnBase.from_pylibcudf(
+        search.search_sorted(
+            list(sort_vals._columns),
+            keys,
+            side="left",
+            ascending=sort_vals.is_monotonic_increasing,
+        )
     ).element_indexing(0)
-    upper_bound = search.search_sorted(
-        list(sort_vals._columns),
-        keys,
-        side="right",
-        ascending=sort_vals.is_monotonic_increasing,
+    upper_bound = ColumnBase.from_pylibcudf(
+        search.search_sorted(
+            list(sort_vals._columns),
+            keys,
+            side="right",
+            ascending=sort_vals.is_monotonic_increasing,
+        )
     ).element_indexing(0)
 
     return lower_bound, upper_bound, sort_inds