From e2693e054566b7bd3c0892a8dfa07d7414967bfc Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyas.ramasubramani@gmail.com>
Date: Thu, 25 Mar 2021 22:06:29 -0700
Subject: [PATCH] Expose libcudf's label_bins function to cudf (#7724)

This PR is a follow-up to #7554 and exposes the feature implemented there via Cython for consumption in cudf's Python API.

Authors:
  - Vyas Ramasubramani (@vyasr)

Approvers:
  - Keith Kraus (@kkraus14)

URL: https://github.com/rapidsai/cudf/pull/7724
---
 python/cudf/cudf/_lib/cpp/labeling.pxd | 19 +++++++++++
 python/cudf/cudf/_lib/labeling.pyx     | 47 ++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 python/cudf/cudf/_lib/cpp/labeling.pxd
 create mode 100644 python/cudf/cudf/_lib/labeling.pyx

diff --git a/python/cudf/cudf/_lib/cpp/labeling.pxd b/python/cudf/cudf/_lib/cpp/labeling.pxd
new file mode 100644
index 00000000000..996ae4f9e38
--- /dev/null
+++ b/python/cudf/cudf/_lib/cpp/labeling.pxd
@@ -0,0 +1,19 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+from libcpp.memory cimport unique_ptr
+
+from cudf._lib.cpp.column.column cimport column
+from cudf._lib.cpp.column.column_view cimport column_view
+
+cdef extern from "cudf/labeling/label_bins.hpp" namespace "cudf" nogil:
+    ctypedef enum inclusive:
+        YES "cudf::inclusive::YES"
+        NO "cudf::inclusive::NO"
+
+    cdef unique_ptr[column] label_bins (
+        const column_view &input,
+        const column_view &left_edges,
+        inclusive left_inclusive,
+        const column_view &right_edges,
+        inclusive right_inclusive
+    ) except +
diff --git a/python/cudf/cudf/_lib/labeling.pyx b/python/cudf/cudf/_lib/labeling.pyx
new file mode 100644
index 00000000000..1b553024347
--- /dev/null
+++ b/python/cudf/cudf/_lib/labeling.pyx
@@ -0,0 +1,47 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+import numpy as np
+from enum import IntEnum
+
+from libc.stdint cimport uint32_t
+from libcpp cimport bool as cbool
+from libcpp.memory cimport unique_ptr
+from libcpp.utility cimport move
+
+from cudf._lib.column cimport Column
+from cudf._lib.replace import replace_nulls
+
+from cudf._lib.cpp.labeling cimport inclusive
+from cudf._lib.cpp.labeling cimport label_bins as cpp_label_bins
+from cudf._lib.cpp.column.column cimport column
+from cudf._lib.cpp.column.column_view cimport column_view
+
+
+# Note that the parameter input shadows a Python built-in in the local scope,
+# but I'm not too concerned about that since there's no use-case for actual
+# input in this context.
+def label_bins(Column input, Column left_edges, cbool left_inclusive,
+               Column right_edges, cbool right_inclusive):
+    cdef inclusive c_left_inclusive = \
+        inclusive.YES if left_inclusive else inclusive.NO
+    cdef inclusive c_right_inclusive = \
+        inclusive.YES if right_inclusive else inclusive.NO
+
+    cdef column_view input_view = input.view()
+    cdef column_view left_edges_view = left_edges.view()
+    cdef column_view right_edges_view = right_edges.view()
+
+    cdef unique_ptr[column] c_result
+
+    with nogil:
+        c_result = move(
+            cpp_label_bins(
+                input_view,
+                left_edges_view,
+                c_left_inclusive,
+                right_edges_view,
+                c_right_inclusive,
+            )
+        )
+
+    return Column.from_unique_ptr(move(c_result))