ENH: 2D support for MaskedArray by jbrockmendel · Pull Request #38992 · pandas-dev/pandas (original) (raw)

so fixing this on my numba branch results in...

@numba.njit def _pad_2d_inplace(values, mask, limit=None): if values.shape[1]: K, N = values.shape if limit is None: for j in range(K): val, prev_mask = values[j, 0], mask[j, 0] for i in range(N): if mask[j, i]: values[j, i], mask[j, i] = val, prev_mask else: val, prev_mask = values[j, i], mask[j, i] else: for j in range(K): fill_count = 0 val, prev_mask = values[j, 0], mask[j, 0] for i in range(N): if mask[j, i]: if fill_count >= limit: continue fill_count += 1 values[j, i], mask[j, i] = val, prev_mask else: fill_count = 0 val, prev_mask = values[j, i], mask[j, i]

I have some duplication here but a perf improvement for the common case of no limit, the duplication can probably be mitigated by reshaping a 1d array and removing the 1d version pad_inplace