Skip to content

Commit 5f6953e

Browse files
pad_2d_inplace
1 parent 3bf5e02 commit 5f6953e

File tree

3 files changed

+42
-26
lines changed

3 files changed

+42
-26
lines changed

pandas/_libs_numba/algos.py

Lines changed: 39 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -620,34 +620,49 @@ def _pad_inplace_with_limit(values: np.ndarray, mask: np.ndarray, limit: int) ->
620620
val = values[i]
621621

622622

623-
# @cython.boundscheck(False)
624-
# @cython.wraparound(False)
625-
# def pad_2d_inplace(algos_t[:, :] values, const uint8_t[:, :] mask, limit=None):
626-
# cdef:
627-
# Py_ssize_t i, j, N, K
628-
# algos_t val
629-
# int lim, fill_count = 0
623+
def pad_2d_inplace(
624+
values: np.ndarray, mask: np.ndarray, limit: int | None = None
625+
) -> None:
626+
_validate_limit(limit)
627+
_pad_2d_inplace(values, mask, limit)
630628

631-
# K, N = (<object>values).shape
632629

633-
# # GH#2778
634-
# if N == 0:
635-
# return
630+
@numba.jit
631+
def _pad_2d_inplace(values, mask, limit=None):
632+
if values.shape[1]:
633+
if limit is None:
634+
_pad_2d_inplace_no_limit(values, mask)
635+
else:
636+
_pad_2d_inplace_with_limit(values, mask, limit)
636637

637-
# lim = validate_limit(N, limit)
638638

639-
# for j in range(K):
640-
# fill_count = 0
641-
# val = values[j, 0]
642-
# for i in range(N):
643-
# if mask[j, i]:
644-
# if fill_count >= lim:
645-
# continue
646-
# fill_count += 1
647-
# values[j, i] = val
648-
# else:
649-
# fill_count = 0
650-
# val = values[j, i]
639+
@numba.jit
640+
def _pad_2d_inplace_no_limit(values, mask):
641+
K, N = values.shape
642+
for j in range(K):
643+
val = values[j, 0]
644+
for i in range(N):
645+
if mask[j, i]:
646+
values[j, i] = val
647+
else:
648+
val = values[j, i]
649+
650+
651+
@numba.jit
652+
def _pad_2d_inplace_with_limit(values, mask, limit):
653+
K, N = values.shape
654+
for j in range(K):
655+
fill_count = 0
656+
val = values[j, 0]
657+
for i in range(N):
658+
if mask[j, i]:
659+
if fill_count >= limit:
660+
continue
661+
fill_count += 1
662+
values[j, i] = val
663+
else:
664+
fill_count = 0
665+
val = values[j, i]
651666

652667

653668
# """

pandas/core/missing.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -710,7 +710,7 @@ def _pad_2d(values, limit=None, mask=None):
710710
values, mask = _fillna_prep(values, mask)
711711

712712
if np.all(values.shape):
713-
algos.pad_2d_inplace(values, mask, limit=limit)
713+
algos_numba.pad_2d_inplace(values, mask, limit=limit)
714714
else:
715715
# for test coverage
716716
pass

pandas/tests/apply/test_frame_transform.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import operator
22
import re
33

4+
from numba import NumbaWarning
45
import numpy as np
56
import pytest
67

@@ -205,7 +206,7 @@ def test_transform_bad_dtype(op, frame_or_series):
205206
msg = "Transform function failed"
206207

207208
# tshift is deprecated
208-
warn = None if op != "tshift" else FutureWarning
209+
warn = NumbaWarning if op != "tshift" else FutureWarning
209210
with tm.assert_produces_warning(warn):
210211
with pytest.raises(ValueError, match=msg):
211212
obj.transform(op)

0 commit comments

Comments
 (0)