Skip to content

Commit

Permalink
Expose libcudf's label_bins function to cudf (#7724)
Browse files Browse the repository at this point in the history
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: #7724
  • Loading branch information
vyasr authored Mar 26, 2021
1 parent 20509d0 commit e2693e0
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 0 deletions.
19 changes: 19 additions & 0 deletions python/cudf/cudf/_lib/cpp/labeling.pxd
Original file line number Diff line number Diff line change
@@ -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 +
47 changes: 47 additions & 0 deletions python/cudf/cudf/_lib/labeling.pyx
Original file line number Diff line number Diff line change
@@ -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))

0 comments on commit e2693e0

Please sign in to comment.