Skip to content

Commit

Permalink
PERF: lib.generate_slices (#42097)
Browse files Browse the repository at this point in the history
  • Loading branch information
mzeitlin11 authored Jun 18, 2021
1 parent f652994 commit 648eb40
Showing 1 changed file with 17 additions and 15 deletions.
32 changes: 17 additions & 15 deletions pandas/_libs/lib.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -906,12 +906,13 @@ def count_level_2d(ndarray[uint8_t, ndim=2, cast=True] mask,
return counts


@cython.wraparound(False)
@cython.boundscheck(False)
def generate_slices(const intp_t[:] labels, Py_ssize_t ngroups):
cdef:
Py_ssize_t i, group_size, n, start
intp_t lab
object slobj
ndarray[int64_t] starts, ends
int64_t[::1] starts, ends

n = len(labels)

Expand All @@ -920,19 +921,20 @@ def generate_slices(const intp_t[:] labels, Py_ssize_t ngroups):

start = 0
group_size = 0
for i in range(n):
lab = labels[i]
if lab < 0:
start += 1
else:
group_size += 1
if i == n - 1 or lab != labels[i + 1]:
starts[lab] = start
ends[lab] = start + group_size
start += group_size
group_size = 0

return starts, ends
with nogil:
for i in range(n):
lab = labels[i]
if lab < 0:
start += 1
else:
group_size += 1
if i == n - 1 or lab != labels[i + 1]:
starts[lab] = start
ends[lab] = start + group_size
start += group_size
group_size = 0

return np.asarray(starts), np.asarray(ends)


def indices_fast(ndarray[intp_t] index, const int64_t[:] labels, list keys,
Expand Down

0 comments on commit 648eb40

Please sign in to comment.