Skip to content

Commit ff62d61

Browse files
committed
Added code for group_fillna
1 parent 699e84a commit ff62d61

File tree

1 file changed

+70
-1
lines changed

1 file changed

+70
-1
lines changed

pandas/_libs/groupby_helper.pxi.in

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,7 @@ def group_ohlc_{{name}}(ndarray[{{dest_type2}}, ndim=2] out,
273273
{{endfor}}
274274

275275
#----------------------------------------------------------------------
276-
# group_nth, group_last, group_rank
276+
# group_nth, group_last, group_rank, group_fillna
277277
#----------------------------------------------------------------------
278278

279279
{{py:
@@ -574,6 +574,75 @@ def group_rank_{{name}}(ndarray[float64_t, ndim=2] out,
574574
for i in range(N):
575575
out[i, 0] = out[i, 0] / grp_sizes[i, 0]
576576
{{endif}}
577+
578+
@cython.wraparound(False)
579+
@cython.boundscheck(False)
580+
def group_fillna_{{name}}(ndarray[{{dest_type2}}, ndim=2] out,
581+
ndarray[{{c_type}}, ndim=2] values,
582+
ndarray[int64_t] labels,
583+
object method,
584+
int64_t limit):
585+
"""Fills values forwards or backwards within a group
586+
587+
Parameters
588+
----------
589+
out : array of {{dest_type2}} values which this method will write its
590+
results to
591+
values : array of {{c_type}} values which may require filling
592+
labels : array containing unique label for each group, with its ordering
593+
matching up to the corresponding record in `values`
594+
method : {'ffill', 'bfill'}
595+
Direction for fill to be applied (forwards or backwards, respectively)
596+
limit : Consecutive values to fill before stopping, or -1 for no limit
597+
598+
Notes
599+
-----
600+
This method modifies the `out` parameter rather than returning an object
601+
"""
602+
cdef:
603+
Py_ssize_t i, N
604+
ndarray[uint8_t] mask
605+
ndarray[int64_t] sorted_labels
606+
{{dest_type2}} curr_fill_val = {{nan_val}}
607+
int64_t idx, filled_vals=0
608+
609+
N, K = (<object> values).shape
610+
611+
{{if name=='int64'}}
612+
mask = (values[:, 0] == {{nan_val}}).astype(np.uint8)
613+
{{elif name=='object'}}
614+
mask = np.array([x != x for x in values[:, 0]]).astype(np.uint8)
615+
{{else}}
616+
mask = np.isnan(values[:, 0]).astype(np.uint8)
617+
{{endif}}
618+
619+
sorted_labels = np.argsort(labels)
620+
if method == 'bfill':
621+
sorted_labels[::-1].sort()
622+
623+
{{if name == 'object'}}
624+
if True: # make templating happy
625+
{{else}}
626+
with nogil:
627+
{{endif}}
628+
for i in range(N):
629+
idx = sorted_labels[i]
630+
if mask[idx]: # is missing
631+
if limit == -1 or filled_vals < limit:
632+
out[idx, 0] = curr_fill_val
633+
else:
634+
out[idx, 0] == {{nan_val}}
635+
filled_vals += 1
636+
else: # reset items when not missing
637+
filled_vals = 0
638+
curr_fill_val = values[idx, 0]
639+
out[idx, 0] = values[idx, 0]
640+
641+
# If we move to the next group, reset
642+
# the fill_val and counter
643+
if i == N - 1 or labels[idx] != labels[sorted_labels[i+1]]:
644+
curr_fill_val = {{nan_val}}
645+
filled_vals = 0
577646
{{endfor}}
578647

579648

0 commit comments

Comments
 (0)