Skip to content

POC Use fused types more, tempita less #22432

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 5 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
218 changes: 150 additions & 68 deletions pandas/_libs/algos_common_helper.pxi.in
Original file line number Diff line number Diff line change
Expand Up @@ -19,33 +19,44 @@ WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
# 1-d template
#----------------------------------------------------------------------

{{py:
ctypedef fused algos_t:
float64_t
float32_t
object
int32_t
int64_t
uint64_t
uint8_t

# name, c_type, dtype, can_hold_na, nogil
dtypes = [('float64', 'float64_t', 'np.float64', True, True),
('float32', 'float32_t', 'np.float32', True, True),
('object', 'object', 'object', True, False),
('int32', 'int32_t', 'np.int32', False, True),
('int64', 'int64_t', 'np.int64', False, True),
('uint64', 'uint64_t', 'np.uint64', False, True),
('bool', 'uint8_t', 'np.bool', False, True)]

def get_dispatch(dtypes):
@cython.wraparound(False)
@cython.boundscheck(False)
def arrmap(ndarray[algos_t] index, object func):
cdef:
Py_ssize_t length = index.shape[0]
Py_ssize_t i = 0
ndarray[object] result = np.empty(length, dtype=np.object_)

for name, c_type, dtype, can_hold_na, nogil in dtypes:
from pandas._libs.lib import maybe_convert_objects

nogil_str = 'with nogil:' if nogil else ''
tab = ' ' if nogil else ''
yield name, c_type, dtype, can_hold_na, nogil_str, tab
}}
for i in range(length):
result[i] = func(index[i])

return maybe_convert_objects(result)

{{for name, c_type, dtype, can_hold_na, nogil_str, tab
in get_dispatch(dtypes)}}

arrmap_float64 = arrmap["float64_t"]
arrmap_float32 = arrmap["float32_t"]
arrmap_object = arrmap["object"]
arrmap_int32 = arrmap["int32_t"]
arrmap_int64 = arrmap["int64_t"]
arrmap_uint64 = arrmap["uint64_t"]
arrmap_bool = arrmap["uint8_t"]


@cython.wraparound(False)
@cython.boundscheck(False)
cpdef map_indices_{{name}}(ndarray[{{c_type}}] index):
cpdef map_indices(ndarray[algos_t] index):
"""
Produce a dict mapping the values of the input array to their respective
locations.
Expand All @@ -55,8 +66,9 @@ cpdef map_indices_{{name}}(ndarray[{{c_type}}] index):

Better to do this with Cython because of the enormous speed boost.
"""
cdef Py_ssize_t i, length
cdef dict result = {}
cdef:
Py_ssize_t i, length
dict result = {}

length = len(index)

Expand All @@ -66,13 +78,22 @@ cpdef map_indices_{{name}}(ndarray[{{c_type}}] index):
return result


map_indices_float64 = map_indices["float64_t"]
map_indices_float32 = map_indices["float32_t"]
map_indices_object = map_indices["object"]
map_indices_int32 = map_indices["int32_t"]
map_indices_int64 = map_indices["int64_t"]
map_indices_uint64 = map_indices["uint64_t"]
map_indices_bool = map_indices["uint8_t"]


@cython.boundscheck(False)
@cython.wraparound(False)
def pad_{{name}}(ndarray[{{c_type}}] old, ndarray[{{c_type}}] new, limit=None):
def pad(ndarray[algos_t] old, ndarray[algos_t] new, limit=None):
cdef:
Py_ssize_t i, j, nleft, nright
ndarray[int64_t, ndim=1] indexer
{{c_type}} cur, next
algos_t cur, next
int lim, fill_count = 0

nleft = len(old)
Expand Down Expand Up @@ -129,20 +150,28 @@ def pad_{{name}}(ndarray[{{c_type}}] old, ndarray[{{c_type}}] new, limit=None):

return indexer

pad_float64 = pad["float64_t"]
pad_float32 = pad["float32_t"]
pad_object = pad["object"]
pad_int32 = pad["int32_t"]
pad_int64 = pad["int64_t"]
pad_uint64 = pad["uint64_t"]
pad_bool = pad["uint8_t"]


@cython.boundscheck(False)
@cython.wraparound(False)
def pad_inplace_{{name}}(ndarray[{{c_type}}] values,
ndarray[uint8_t, cast=True] mask,
limit=None):
def pad_inplace(ndarray[algos_t] values,
ndarray[uint8_t, cast=True] mask,
limit=None):
cdef:
Py_ssize_t i, N
{{c_type}} val
int lim, fill_count = 0
cdef Py_ssize_t i, N
cdef algos_t val
cdef int lim, fill_count = 0

N = len(values)

# GH 2778
# GH#2778
if N == 0:
return

Expand All @@ -167,19 +196,28 @@ def pad_inplace_{{name}}(ndarray[{{c_type}}] values,
val = values[i]


pad_inplace_float64 = pad_inplace["float64_t"]
pad_inplace_float32 = pad_inplace["float32_t"]
pad_inplace_object = pad_inplace["object"]
pad_inplace_int32 = pad_inplace["int32_t"]
pad_inplace_int64 = pad_inplace["int64_t"]
pad_inplace_uint64 = pad_inplace["uint64_t"]
pad_inplace_bool = pad_inplace["uint8_t"]


@cython.boundscheck(False)
@cython.wraparound(False)
def pad_2d_inplace_{{name}}(ndarray[{{c_type}}, ndim=2] values,
ndarray[uint8_t, ndim=2] mask,
limit=None):
def pad_2d_inplace(ndarray[algos_t, ndim=2] values,
ndarray[uint8_t, ndim=2] mask,
limit=None):
cdef:
Py_ssize_t i, j, N, K
{{c_type}} val
algos_t val
int lim, fill_count = 0

K, N = (<object> values).shape

# GH 2778
# GH#2778
if N == 0:
return

Expand All @@ -205,6 +243,16 @@ def pad_2d_inplace_{{name}}(ndarray[{{c_type}}, ndim=2] values,
fill_count = 0
val = values[j, i]


pad_2d_inplace_float64 = pad_2d_inplace["float64_t"]
pad_2d_inplace_float32 = pad_2d_inplace["float32_t"]
pad_2d_inplace_object = pad_2d_inplace["object"]
pad_2d_inplace_int32 = pad_2d_inplace["int32_t"]
pad_2d_inplace_int64 = pad_2d_inplace["int64_t"]
pad_2d_inplace_uint64 = pad_2d_inplace["uint64_t"]
pad_2d_inplace_bool = pad_2d_inplace["uint8_t"]


"""
Backfilling logic for generating fill vector

Expand Down Expand Up @@ -233,13 +281,12 @@ D

@cython.boundscheck(False)
@cython.wraparound(False)
def backfill_{{name}}(ndarray[{{c_type}}] old, ndarray[{{c_type}}] new,
limit=None):
def backfill(ndarray[algos_t] old, ndarray[algos_t] new,limit=None):
cdef:
Py_ssize_t i, j, nleft, nright
ndarray[int64_t, ndim=1] indexer
{{c_type}} cur, prev
int lim, fill_count = 0
cdef Py_ssize_t i, j, nleft, nright
cdef ndarray[int64_t, ndim=1] indexer
cdef algos_t cur, prev
cdef int lim, fill_count = 0

nleft = len(old)
nright = len(new)
Expand Down Expand Up @@ -297,19 +344,28 @@ def backfill_{{name}}(ndarray[{{c_type}}] old, ndarray[{{c_type}}] new,
return indexer


backfill_float64 = backfill["float64_t"]
backfill_float32 = backfill["float32_t"]
backfill_object = backfill["object"]
backfill_int32 = backfill["int32_t"]
backfill_int64 = backfill["int64_t"]
backfill_uint64 = backfill["uint64_t"]
backfill_bool = backfill["uint8_t"]


@cython.boundscheck(False)
@cython.wraparound(False)
def backfill_inplace_{{name}}(ndarray[{{c_type}}] values,
ndarray[uint8_t, cast=True] mask,
limit=None):
def backfill_inplace(ndarray[algos_t] values,
ndarray[uint8_t, cast=True] mask,
limit=None):
cdef:
Py_ssize_t i, N
{{c_type}} val
int lim, fill_count = 0
cdef Py_ssize_t i, N
cdef algos_t val
cdef int lim, fill_count = 0

N = len(values)

# GH 2778
# GH#2778
if N == 0:
return

Expand All @@ -334,19 +390,28 @@ def backfill_inplace_{{name}}(ndarray[{{c_type}}] values,
val = values[i]


backfill_inplace_float64 = backfill_inplace["float64_t"]
backfill_inplace_float32 = backfill_inplace["float32_t"]
backfill_inplace_object = backfill_inplace["object"]
backfill_inplace_int32 = backfill_inplace["int32_t"]
backfill_inplace_int64 = backfill_inplace["int64_t"]
backfill_inplace_uint64 = backfill_inplace["uint64_t"]
backfill_inplace_bool = backfill_inplace["uint8_t"]


@cython.boundscheck(False)
@cython.wraparound(False)
def backfill_2d_inplace_{{name}}(ndarray[{{c_type}}, ndim=2] values,
ndarray[uint8_t, ndim=2] mask,
limit=None):
def backfill_2d_inplace(ndarray[algos_t, ndim=2] values,
ndarray[uint8_t, ndim=2] mask,
limit=None):
cdef:
Py_ssize_t i, j, N, K
{{c_type}} val
algos_t val
int lim, fill_count = 0

K, N = (<object> values).shape

# GH 2778
# GH#2778
if N == 0:
return

Expand All @@ -373,6 +438,39 @@ def backfill_2d_inplace_{{name}}(ndarray[{{c_type}}, ndim=2] values,
val = values[j, i]


backfill_2d_inplace_float64 = backfill_2d_inplace["float64_t"]
backfill_2d_inplace_float32 = backfill_2d_inplace["float32_t"]
backfill_2d_inplace_object = backfill_2d_inplace["object"]
backfill_2d_inplace_int32 = backfill_2d_inplace["int32_t"]
backfill_2d_inplace_int64 = backfill_2d_inplace["int64_t"]
backfill_2d_inplace_uint64 = backfill_2d_inplace["uint64_t"]
backfill_2d_inplace_bool = backfill_2d_inplace["uint8_t"]


{{py:

# name, c_type, dtype, can_hold_na, nogil
dtypes = [('float64', 'float64_t', 'np.float64', True, True),
('float32', 'float32_t', 'np.float32', True, True),
('object', 'object', 'object', True, False),
('int32', 'int32_t', 'np.int32', False, True),
('int64', 'int64_t', 'np.int64', False, True),
('uint64', 'uint64_t', 'np.uint64', False, True),
('bool', 'uint8_t', 'np.bool', False, True)]

def get_dispatch(dtypes):

for name, c_type, dtype, can_hold_na, nogil in dtypes:

nogil_str = 'with nogil:' if nogil else ''
tab = ' ' if nogil else ''
yield name, c_type, dtype, can_hold_na, nogil_str, tab
}}

{{for name, c_type, dtype, can_hold_na, nogil_str, tab
in get_dispatch(dtypes)}}


@cython.boundscheck(False)
@cython.wraparound(False)
def is_monotonic_{{name}}(ndarray[{{c_type}}] arr, bint timelike):
Expand Down Expand Up @@ -429,22 +527,6 @@ def is_monotonic_{{name}}(ndarray[{{c_type}}] arr, bint timelike):
return is_monotonic_inc, is_monotonic_dec, \
is_unique and (is_monotonic_inc or is_monotonic_dec)


@cython.wraparound(False)
@cython.boundscheck(False)
def arrmap_{{name}}(ndarray[{{c_type}}] index, object func):
cdef:
Py_ssize_t length = index.shape[0]
Py_ssize_t i = 0
ndarray[object] result = np.empty(length, dtype=np.object_)

from pandas._libs.lib import maybe_convert_objects

for i in range(length):
result[i] = func(index[i])

return maybe_convert_objects(result)

{{endfor}}

#----------------------------------------------------------------------
Expand Down
31 changes: 18 additions & 13 deletions pandas/_libs/algos_take_helper.pxi.in
Original file line number Diff line number Diff line change
Expand Up @@ -264,29 +264,34 @@ def take_2d_multi_{{name}}_{{dest}}(ndarray[{{c_type_in}}, ndim=2] values,
# take_2d internal function
#----------------------------------------------------------------------

{{py:

# dtype, ctype, init_result
dtypes = [('float64', 'float64_t', 'np.empty_like(values)'),
('uint64', 'uint64_t', 'np.empty_like(values)'),
('object', 'object', 'values.copy()'),
('int64', 'int64_t', 'np.empty_like(values)')]
}}
ctypedef fused take_t:
float64_t
uint64_t
object
int64_t

{{for dtype, ctype, init_result in dtypes}}

cdef _take_2d_{{dtype}}(ndarray[{{ctype}}, ndim=2] values, object idx):
cdef _take_2d(ndarray[take_t, ndim=2] values, object idx):
cdef:
Py_ssize_t i, j, N, K
ndarray[Py_ssize_t, ndim=2, cast=True] indexer = idx
ndarray[{{ctype}}, ndim=2] result
ndarray[take_t, ndim=2] result
object val

N, K = (<object> values).shape
result = {{init_result}}
if take_t is object:
result = values.copy()
else:
result = np.empty_like(values)

for i in range(N):
for j in range(K):
result[i, j] = values[i, indexer[i, j]]
return result

{{endfor}}

# TODO: Are these treated as cdefs?
_take_2d_float64 = _take_2d[float64_t]
_take_2d_uint64 = _take_2d[uint64_t]
_take_2d_object = _take_2d[object]
_take_2d_int64 = _take_2d[int64_t]
Loading