From 36968f520df44853e6ce2ef15541a8f9549bd807 Mon Sep 17 00:00:00 2001 From: Christopher Ariza Date: Thu, 30 Jul 2026 18:29:19 -0700 Subject: [PATCH 1/2] first rendering --- README.rst | 6 + src/__init__.py | 1 + src/__init__.pyi | 8 ++ src/_arraykit.c | 4 + src/methods.c | 156 +++++++++++++++++++++++ src/methods.h | 3 + test/test_fill_directional.py | 229 ++++++++++++++++++++++++++++++++++ 7 files changed, 407 insertions(+) create mode 100644 test/test_fill_directional.py diff --git a/README.rst b/README.rst index 1161a488..391b6f3e 100644 --- a/README.rst +++ b/README.rst @@ -35,6 +35,12 @@ ArrayKit requires the following: What is New in ArrayKit ------------------------- +1.8.0 +............ + +Added ``fill_directional()``. + + 1.7.0 ............ diff --git a/src/__init__.py b/src/__init__.py index 1fc933f5..bca91283 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -27,6 +27,7 @@ from ._arraykit import write_array_to_file as write_array_to_file from ._arraykit import factorize as factorize from ._arraykit import group_ordering as group_ordering +from ._arraykit import fill_directional as fill_directional from ._arraykit import count_iteration as count_iteration from ._arraykit import first_true_1d as first_true_1d from ._arraykit import first_true_2d as first_true_2d diff --git a/src/__init__.pyi b/src/__init__.pyi index 85304f9c..cfd71b26 100644 --- a/src/__init__.pyi +++ b/src/__init__.pyi @@ -230,6 +230,14 @@ def factorize( def group_ordering( codes: np.ndarray, *, size: tp.Optional[int] = ... ) -> tp.Tuple[np.ndarray, np.ndarray]: ... +def fill_directional( + array: np.ndarray, + target: np.ndarray, + *, + forward: bool = ..., + axis: int = ..., + limit: int = ..., +) -> np.ndarray: ... def first_true_1d(__array: np.ndarray, *, forward: bool) -> int: ... def first_true_2d(__array: np.ndarray, *, forward: bool, axis: int) -> np.ndarray: ... def nonzero_1d(__array: np.ndarray, /) -> np.ndarray: ... diff --git a/src/_arraykit.c b/src/_arraykit.c index db80aca4..a3cba481 100644 --- a/src/_arraykit.c +++ b/src/_arraykit.c @@ -78,6 +78,10 @@ static PyMethodDef arraykit_methods[] = { (PyCFunction)group_ordering, METH_VARARGS | METH_KEYWORDS, NULL}, + {"fill_directional", + (PyCFunction)fill_directional, + METH_VARARGS | METH_KEYWORDS, + NULL}, {NULL}, }; diff --git a/src/methods.c b/src/methods.c index 746846d7..015f6f17 100644 --- a/src/methods.c +++ b/src/methods.c @@ -1128,6 +1128,162 @@ group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) return NULL; } +// Fill one strided lane in place: walk positions in the fill direction, carrying +// the most recent non-target value into each target position (subject to `limit` +// consecutive fills per run). `elem_base`/`elem_stride` address elements in bytes; +// `tgt_base`/`tgt_stride` address the aligned Boolean target lane. Object arrays +// duplicate the carried PyObject* with balanced refcounts; all other dtypes memcpy +// the itemsize bytes. Leading targets (no prior source) are left unchanged. +static inline void +AK_fill_lane( + char *elem_base, + npy_intp elem_stride, + const npy_bool *tgt_base, + npy_intp tgt_stride, + npy_intp length, + npy_intp itemsize, + int is_object, + int forward, + npy_intp limit) +{ + char *last_valid = NULL; + npy_intp count = 0; + for (npy_intp k = 0; k < length; k++) { + npy_intp pos = forward ? k : (length - 1 - k); + char *elem = elem_base + pos * elem_stride; + if (tgt_base[pos * tgt_stride]) { + if (last_valid != NULL && (limit == 0 || count < limit)) { + if (is_object) { + PyObject **dst = (PyObject**)elem; + PyObject **src = (PyObject**)last_valid; + Py_INCREF(*src); + Py_XDECREF(*dst); + *dst = *src; + } + else { + memcpy(elem, last_valid, (size_t)itemsize); + } + count++; + } + // else: leading target (no source yet) or limit reached; leave as-is + } + else { + last_valid = elem; + count = 0; + } + } +} + +static char *fill_directional_kwarg_names[] = { + "array", + "target", + "forward", + "axis", + "limit", + NULL +}; + +// Directional (forward or backward) fill of `array` along `axis`, replacing each +// position flagged True in the Boolean `target` mask with the last (forward) or +// next (backward) non-target value. A single O(n) pass per lane; `limit` caps the +// number of consecutive fills per run (0 == unlimited). The target mask is supplied +// by the caller (e.g. isna/isfalsy), keeping this dtype- and predicate-agnostic. +// Returns a new, immutable, C-contiguous array. +PyObject * +fill_directional(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) +{ + PyArrayObject *array = NULL; + PyArrayObject *target = NULL; + int forward = 1; + int axis = 0; + Py_ssize_t limit = 0; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, + "O!O!|$pin:fill_directional", + fill_directional_kwarg_names, + &PyArray_Type, &array, + &PyArray_Type, &target, + &forward, + &axis, + &limit + )) { + return NULL; + } + + int ndim = PyArray_NDIM(array); + if (ndim != 1 && ndim != 2) { + PyErr_SetString(PyExc_ValueError, "array must be 1- or 2-dimensional"); + return NULL; + } + if (PyArray_TYPE(target) != NPY_BOOL) { + PyErr_SetString(PyExc_ValueError, "target must be a Boolean array"); + return NULL; + } + if (PyArray_NDIM(target) != ndim) { + PyErr_SetString(PyExc_ValueError, + "target must match the dimensionality of array"); + return NULL; + } + for (int i = 0; i < ndim; i++) { + if (PyArray_DIM(target, i) != PyArray_DIM(array, i)) { + PyErr_SetString(PyExc_ValueError, + "target must match the shape of array"); + return NULL; + } + } + if (!PyArray_IS_C_CONTIGUOUS(target)) { + PyErr_SetString(PyExc_ValueError, "target must be C-contiguous"); + return NULL; + } + if (ndim == 2 && axis != 0 && axis != 1) { + PyErr_SetString(PyExc_ValueError, "axis must be 0 or 1"); + return NULL; + } + if (limit < 0) { + PyErr_SetString(PyExc_ValueError, "limit must be non-negative"); + return NULL; + } + + // A C-order copy is the result: non-target positions are already correct, and + // its layout aligns element-for-element with the C-contiguous target mask. + PyObject *out = PyArray_NewCopy(array, NPY_CORDER); + if (!out) { + return NULL; + } + + int is_object = PyArray_TYPE(array) == NPY_OBJECT; + npy_intp itemsize = PyArray_ITEMSIZE((PyArrayObject*)out); + char *out_data = (char*)PyArray_DATA((PyArrayObject*)out); + const npy_bool *tgt = (const npy_bool*)PyArray_DATA(target); + + if (ndim == 1) { + npy_intp n = PyArray_DIM((PyArrayObject*)out, 0); + AK_fill_lane(out_data, itemsize, tgt, 1, n, + itemsize, is_object, forward, limit); + } + else { + npy_intp rows = PyArray_DIM((PyArrayObject*)out, 0); + npy_intp cols = PyArray_DIM((PyArrayObject*)out, 1); + if (axis == 0) { + // fill down each column: elements stride by a full row (cols * itemsize) + for (npy_intp c = 0; c < cols; c++) { + AK_fill_lane(out_data + c * itemsize, cols * itemsize, + tgt + c, cols, rows, itemsize, is_object, forward, limit); + } + } + else { + // fill across each row: elements are contiguous + for (npy_intp r = 0; r < rows; r++) { + AK_fill_lane(out_data + r * cols * itemsize, itemsize, + tgt + r * cols, 1, cols, itemsize, is_object, forward, limit); + } + } + } + + PyArray_CLEARFLAGS((PyArrayObject*)out, NPY_ARRAY_WRITEABLE); + return out; +} + PyObject * dtype_from_element(PyObject *Py_UNUSED(m), PyObject *arg) { diff --git a/src/methods.h b/src/methods.h index 574c7bdd..7b425300 100644 --- a/src/methods.h +++ b/src/methods.h @@ -72,6 +72,9 @@ first_true_2d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); PyObject * group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); +PyObject * +fill_directional(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); + PyObject * dtype_from_element(PyObject *Py_UNUSED(m), PyObject *arg); diff --git a/test/test_fill_directional.py b/test/test_fill_directional.py new file mode 100644 index 00000000..7196865f --- /dev/null +++ b/test/test_fill_directional.py @@ -0,0 +1,229 @@ +import sys +import unittest + +import numpy as np +from arraykit import fill_directional + + +def reference(array, target, *, forward=True, axis=0, limit=0): + """Pure-Python reference: carry the last (forward) or next (backward) non-target + value across target runs, capped at `limit` consecutive fills.""" + out = array.copy() + + def fill_lane(vals, tgt): + last = None + count = 0 + order = range(len(vals)) if forward else range(len(vals) - 1, -1, -1) + for i in order: + if tgt[i]: + if last is not None and (limit == 0 or count < limit): + vals[i] = last + count += 1 + else: + last = vals[i] + count = 0 + + if array.ndim == 1: + fill_lane(out, target) + elif axis == 0: + for c in range(out.shape[1]): + col = out[:, c].copy() + fill_lane(col, target[:, c]) + out[:, c] = col + else: + for r in range(out.shape[0]): + row = out[r, :].copy() + fill_lane(row, target[r, :]) + out[r, :] = row + return out + + +class TestUnit(unittest.TestCase): + # ------------------------------------------------------------------ + # basic 1D behavior + + def test_forward_1d(self) -> None: + a = np.array([1.0, np.nan, np.nan, 4.0, np.nan]) + t = np.isnan(a) + post = fill_directional(a, t) + self.assertEqual(post.tolist(), [1.0, 1.0, 1.0, 4.0, 4.0]) + + def test_backward_1d(self) -> None: + a = np.array([1.0, np.nan, np.nan, 4.0, np.nan]) + t = np.isnan(a) + post = fill_directional(a, t, forward=False) + # trailing target has no source and stays NaN + self.assertTrue( + np.array_equal(post, [1.0, 4.0, 4.0, 4.0, np.nan], equal_nan=True) + ) + + def test_leading_target_unchanged(self) -> None: + a = np.array([np.nan, np.nan, 3.0, np.nan]) + post = fill_directional(a, np.isnan(a)) + self.assertTrue( + np.array_equal(post, [np.nan, np.nan, 3.0, 3.0], equal_nan=True) + ) + + def test_all_target(self) -> None: + a = np.array([np.nan, np.nan, np.nan]) + post = fill_directional(a, np.isnan(a)) + self.assertTrue(np.array_equal(post, a, equal_nan=True)) + + def test_no_target(self) -> None: + a = np.arange(5.0) + post = fill_directional(a, np.zeros(5, dtype=bool)) + self.assertEqual(post.tolist(), a.tolist()) + + def test_empty(self) -> None: + a = np.array([], dtype=float) + post = fill_directional(a, np.array([], dtype=bool)) + self.assertEqual(len(post), 0) + + # ------------------------------------------------------------------ + # limit + + def test_limit_forward(self) -> None: + a = np.array([1.0, np.nan, np.nan, np.nan, 5.0]) + post = fill_directional(a, np.isnan(a), limit=1) + self.assertTrue( + np.array_equal(post, [1.0, 1.0, np.nan, np.nan, 5.0], equal_nan=True) + ) + + def test_limit_resets_per_run(self) -> None: + a = np.array([1.0, np.nan, 3.0, np.nan, np.nan, 6.0]) + post = fill_directional(a, np.isnan(a), limit=1) + self.assertTrue( + np.array_equal(post, [1.0, 1.0, 3.0, 3.0, np.nan, 6.0], equal_nan=True) + ) + + # ------------------------------------------------------------------ + # dtypes + + def test_int(self) -> None: + a = np.array([5, 0, 0, 8, 0]) + t = np.array([False, True, True, False, True]) + self.assertEqual(fill_directional(a, t).tolist(), [5, 5, 5, 8, 8]) + + def test_object(self) -> None: + a = np.array([None, 'a', None, 'b', None, None], dtype=object) + t = np.array([v is None for v in a]) + post = fill_directional(a, t) + self.assertEqual(post.tolist(), [None, 'a', 'a', 'b', 'b', 'b']) + + def test_object_no_refcount_leak(self) -> None: + marker = object() + a = np.array([marker, None, None, marker, None], dtype=object) + t = np.array([False, True, True, False, True]) + base = sys.getrefcount(marker) + for _ in range(2000): + post = fill_directional(a, t) + del post + self.assertEqual(sys.getrefcount(marker), base) + + def test_datetime(self) -> None: + a = np.array(['2020-01-01', 'NaT', '2020-01-03'], dtype='datetime64[D]') + t = np.isnat(a) + post = fill_directional(a, t) + self.assertEqual( + post.tolist(), + np.array( + ['2020-01-01', '2020-01-01', '2020-01-03'], dtype='datetime64[D]' + ).tolist(), + ) + + # ------------------------------------------------------------------ + # 2D + + def test_2d_axis0(self) -> None: + a = np.array([[1.0, np.nan], [np.nan, 5.0], [3.0, np.nan]]) + post = fill_directional(a, np.isnan(a), axis=0) + self.assertTrue( + np.array_equal(post, [[1.0, np.nan], [1.0, 5.0], [3.0, 5.0]], equal_nan=True) + ) + + def test_2d_axis1(self) -> None: + a = np.array([[1.0, np.nan, 3.0], [np.nan, np.nan, 6.0]]) + post = fill_directional(a, np.isnan(a), axis=1) + self.assertTrue( + np.array_equal( + post, [[1.0, 1.0, 3.0], [np.nan, np.nan, 6.0]], equal_nan=True + ) + ) + + # ------------------------------------------------------------------ + # immutability + + def test_immutable(self) -> None: + a = np.array([1.0, np.nan, 3.0]) + post = fill_directional(a, np.isnan(a)) + self.assertFalse(post.flags.writeable) + + # ------------------------------------------------------------------ + # differential vs reference across a matrix of parameters + + def test_matches_reference(self) -> None: + rng = np.random.default_rng(42) + arrays_1d = ( + np.round(rng.random(300), 2), + np.where(rng.random(300) < 0.5, np.nan, rng.random(300)), + rng.integers(0, 4, 300), + ) + for a in arrays_1d: + t = np.isnan(a) if a.dtype.kind == 'f' else (a % 3 == 0) + for forward in (True, False): + for limit in (0, 1, 3): + post = fill_directional(a, t, forward=forward, limit=limit) + exp = reference(a, t, forward=forward, limit=limit) + self.assertTrue( + np.array_equal(post, exp, equal_nan=(a.dtype.kind == 'f')) + ) + for shape in ((8, 5), (5, 8), (1, 6), (6, 1), (10, 10)): + a = np.where(rng.random(shape) < 0.35, np.nan, rng.random(shape)) + t = np.isnan(a) + for axis in (0, 1): + for forward in (True, False): + for limit in (0, 2): + post = fill_directional( + a, t, forward=forward, axis=axis, limit=limit + ) + exp = reference( + a, t, forward=forward, axis=axis, limit=limit + ) + self.assertTrue(np.array_equal(post, exp, equal_nan=True)) + + # ------------------------------------------------------------------ + # errors + + def test_error_target_not_bool(self) -> None: + a = np.arange(3.0) + with self.assertRaises(ValueError): + fill_directional(a, np.zeros(3, dtype=int)) + + def test_error_shape_mismatch(self) -> None: + a = np.arange(3.0) + with self.assertRaises(ValueError): + fill_directional(a, np.zeros(4, dtype=bool)) + + def test_error_ndim_mismatch(self) -> None: + a = np.arange(6.0).reshape(2, 3) + with self.assertRaises(ValueError): + fill_directional(a, np.zeros(6, dtype=bool)) + + def test_error_bad_axis(self) -> None: + a = np.arange(6.0).reshape(2, 3) + with self.assertRaises(ValueError): + fill_directional(a, np.zeros((2, 3), dtype=bool), axis=2) + + def test_error_negative_limit(self) -> None: + a = np.arange(3.0) + with self.assertRaises(ValueError): + fill_directional(a, np.zeros(3, dtype=bool), limit=-1) + + def test_error_3d(self) -> None: + a = np.zeros((2, 2, 2)) + with self.assertRaises(ValueError): + fill_directional(a, np.zeros((2, 2, 2), dtype=bool)) + + +if __name__ == '__main__': + unittest.main() From 46e5896b969292f0335a8ec031c6fc7f36a3dd36 Mon Sep 17 00:00:00 2001 From: Christopher Ariza Date: Thu, 30 Jul 2026 18:40:01 -0700 Subject: [PATCH 2/2] cleanup --- src/methods.c | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/methods.c b/src/methods.c index 015f6f17..186984f8 100644 --- a/src/methods.c +++ b/src/methods.c @@ -1131,15 +1131,15 @@ group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) // Fill one strided lane in place: walk positions in the fill direction, carrying // the most recent non-target value into each target position (subject to `limit` // consecutive fills per run). `elem_base`/`elem_stride` address elements in bytes; -// `tgt_base`/`tgt_stride` address the aligned Boolean target lane. Object arrays +// `target_base`/`target_stride` address the aligned Boolean target lane. Object arrays // duplicate the carried PyObject* with balanced refcounts; all other dtypes memcpy // the itemsize bytes. Leading targets (no prior source) are left unchanged. static inline void AK_fill_lane( char *elem_base, npy_intp elem_stride, - const npy_bool *tgt_base, - npy_intp tgt_stride, + const npy_bool *target_base, + npy_intp target_stride, npy_intp length, npy_intp itemsize, int is_object, @@ -1151,7 +1151,7 @@ AK_fill_lane( for (npy_intp k = 0; k < length; k++) { npy_intp pos = forward ? k : (length - 1 - k); char *elem = elem_base + pos * elem_stride; - if (tgt_base[pos * tgt_stride]) { + if (target_base[pos * target_stride]) { if (last_valid != NULL && (limit == 0 || count < limit)) { if (is_object) { PyObject **dst = (PyObject**)elem; @@ -1253,12 +1253,13 @@ fill_directional(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) int is_object = PyArray_TYPE(array) == NPY_OBJECT; npy_intp itemsize = PyArray_ITEMSIZE((PyArrayObject*)out); + char *out_data = (char*)PyArray_DATA((PyArrayObject*)out); - const npy_bool *tgt = (const npy_bool*)PyArray_DATA(target); + const npy_bool *target_data = (const npy_bool*)PyArray_DATA(target); if (ndim == 1) { npy_intp n = PyArray_DIM((PyArrayObject*)out, 0); - AK_fill_lane(out_data, itemsize, tgt, 1, n, + AK_fill_lane(out_data, itemsize, target_data, 1, n, itemsize, is_object, forward, limit); } else { @@ -1268,14 +1269,14 @@ fill_directional(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) // fill down each column: elements stride by a full row (cols * itemsize) for (npy_intp c = 0; c < cols; c++) { AK_fill_lane(out_data + c * itemsize, cols * itemsize, - tgt + c, cols, rows, itemsize, is_object, forward, limit); + target_data + c, cols, rows, itemsize, is_object, forward, limit); } } else { // fill across each row: elements are contiguous for (npy_intp r = 0; r < rows; r++) { AK_fill_lane(out_data + r * cols * itemsize, itemsize, - tgt + r * cols, 1, cols, itemsize, is_object, forward, limit); + target_data + r * cols, 1, cols, itemsize, is_object, forward, limit); } } }