From cc71de4d8c70da4e3b5d74ad4be2e7a809820816 Mon Sep 17 00:00:00 2001 From: Christopher Ariza Date: Fri, 31 Jul 2026 18:21:53 -0700 Subject: [PATCH 1/2] rendering --- README.rst | 6 ++ src/__init__.pyi | 1 + src/tri_map.c | 54 +++++++++++++++ test/test_tri_map.py | 154 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 215 insertions(+) diff --git a/README.rst b/README.rst index 6be72869..71b94254 100644 --- a/README.rst +++ b/README.rst @@ -35,6 +35,12 @@ ArrayKit requires the following: What is New in ArrayKit ------------------------- +1.10.0 +............ + +Added ``TriMap.register_pairs()``. + + 1.9.0 ............ diff --git a/src/__init__.pyi b/src/__init__.pyi index e856161e..04559c28 100644 --- a/src/__init__.pyi +++ b/src/__init__.pyi @@ -54,6 +54,7 @@ class TriMap: def __repr__(self) -> str: ... def register_one(self, /, src_from: int, dst_from: int) -> None: ... def register_many_from_one(self, __dst_pos: np.ndarray) -> None: ... + def register_pairs(self, __src_pos: np.ndarray, __dst_pos: np.ndarray) -> None: ... def register_unmatched_dst(self) -> None: ... def register_many(self, /, src_from: int, dst_from: np.ndarray) -> None: ... def finalize(self) -> None: ... diff --git a/src/tri_map.c b/src/tri_map.c index 2168bb01..003c3d5d 100644 --- a/src/tri_map.c +++ b/src/tri_map.c @@ -310,6 +310,59 @@ TriMap_register_many_from_one(TriMapObject *self, PyObject *arg) { Py_RETURN_NONE; } +// Bulk pair registration: given two equal-length int64 arrays `src_pos` and `dst_pos`, +// register the pair (src_pos[i], dst_pos[i]) for each i in a single C loop -- equivalent +// to calling register_one(src_pos[i], dst_pos[i]) for each i, but without per-element +// Python overhead. Either position may be -1 to indicate an unmatched (fill) side. This +// generalizes register_many_from_one (where src is the implicit range 0..src_len) to +// arbitrary src positions, as needed for many-to-many joins. +PyObject * +TriMap_register_pairs(TriMapObject *self, PyObject *args) { + PyObject* src_arg; + PyObject* dst_arg; + if (!PyArg_ParseTuple(args, + "OO:register_pairs", + &src_arg, + &dst_arg)) { + return NULL; + } + if (self->finalized) { + PyErr_SetString(PyExc_RuntimeError, "Cannot register post finalization"); + return NULL; + } + if (!PyArray_Check(src_arg) || !PyArray_Check(dst_arg)) { + PyErr_SetString(PyExc_TypeError, "Must provide arrays"); + return NULL; + } + PyArrayObject* src_a = (PyArrayObject*)src_arg; + PyArrayObject* dst_a = (PyArrayObject*)dst_arg; + if (PyArray_TYPE(src_a) != NPY_INT64 || PyArray_TYPE(dst_a) != NPY_INT64) { + PyErr_SetString(PyExc_ValueError, "Arrays must be of type int64"); + return NULL; + } + if (PyArray_NDIM(src_a) != 1 || PyArray_NDIM(dst_a) != 1) { + PyErr_SetString(PyExc_ValueError, "Arrays must be 1-dimensional"); + return NULL; + } + if (!PyArray_IS_C_CONTIGUOUS(src_a) || !PyArray_IS_C_CONTIGUOUS(dst_a)) { + PyErr_SetString(PyExc_ValueError, "Arrays must be contiguous"); + return NULL; + } + npy_intp n = PyArray_SIZE(src_a); + if (n != PyArray_SIZE(dst_a)) { + PyErr_SetString(PyExc_ValueError, "Arrays must be the same length"); + return NULL; + } + const npy_int64* s = (npy_int64*)PyArray_DATA(src_a); + const npy_int64* d = (npy_int64*)PyArray_DATA(dst_a); + for (npy_intp i = 0; i < n; i++) { + if (AK_TM_register_one(self, (Py_ssize_t)s[i], (Py_ssize_t)d[i])) { + return NULL; + } + } + Py_RETURN_NONE; +} + PyObject * TriMap_register_unmatched_dst(TriMapObject *self) { if (self->finalized) { @@ -1400,6 +1453,7 @@ TriMap_map_dst_fill(TriMapObject *self, PyObject *args) { static PyMethodDef TriMap_methods[] = { {"register_one", (PyCFunction)TriMap_register_one, METH_VARARGS, NULL}, {"register_many_from_one", (PyCFunction)TriMap_register_many_from_one, METH_O, NULL}, + {"register_pairs", (PyCFunction)TriMap_register_pairs, METH_VARARGS, NULL}, {"register_unmatched_dst", (PyCFunction)TriMap_register_unmatched_dst, METH_NOARGS, NULL}, {"register_many", (PyCFunction)TriMap_register_many, METH_VARARGS, NULL}, {"finalize", (PyCFunction)TriMap_finalize, METH_NOARGS, NULL}, diff --git a/test/test_tri_map.py b/test/test_tri_map.py index cb73fcbe..fcf21c1a 100644 --- a/test/test_tri_map.py +++ b/test/test_tri_map.py @@ -1508,3 +1508,157 @@ def test_tri_map_register_many_from_one_errors(self) -> None: tm.finalize() with self.assertRaises(RuntimeError): # post-finalize tm.register_many_from_one(np.array([0, 1], dtype=np.int64)) + + # --------------------------------------------------------------------------- + # register_pairs (bulk arbitrary (src, dst) pairs) + + def test_tri_map_register_pairs_a(self) -> None: + # a many-to-many mapping: src 0 -> dst 0 and 1; src 1 -> dst 1 + tm = TriMap(2, 2) + tm.register_pairs( + np.array([0, 0, 1], dtype=np.int64), + np.array([0, 1, 1], dtype=np.int64), + ) + tm.finalize() + self.assertTrue(tm.is_many()) + self.assertTrue(tm.src_no_fill()) + self.assertTrue(tm.dst_no_fill()) + src = np.array([10, 20]) + dst = np.array([100, 200]) + self.assertEqual(tm.map_src_no_fill(src).tolist(), [10, 10, 20]) + self.assertEqual(tm.map_dst_no_fill(dst).tolist(), [100, 200, 200]) + + def test_tri_map_register_pairs_unmatched(self) -> None: + # -1 on either side marks that side as a fill row + tm = TriMap(3, 2) + tm.register_pairs( + np.array([0, 1, 2], dtype=np.int64), + np.array([0, 1, -1], dtype=np.int64), + ) + tm.finalize() + self.assertTrue(tm.src_no_fill()) # every row has a src + self.assertFalse(tm.dst_no_fill()) # row 2 has no dst + dst = np.array([100, 200]) + self.assertEqual( + tm.map_dst_fill(dst, -9, np.dtype(np.int64)).tolist(), [100, 200, -9] + ) + + def test_tri_map_register_pairs_unmatched_src(self) -> None: + # -1 on the src side (as register_unmatched_dst would produce for OUTER) + tm = TriMap(2, 3) + tm.register_pairs( + np.array([0, 1, -1], dtype=np.int64), + np.array([0, 1, 2], dtype=np.int64), + ) + tm.finalize() + self.assertFalse(tm.src_no_fill()) # row 2 has no src + self.assertTrue(tm.dst_no_fill()) # every row has a dst + src = np.array([10, 20]) + self.assertEqual( + tm.map_src_fill(src, -9, np.dtype(np.int64)).tolist(), [10, 20, -9] + ) + + def test_tri_map_register_pairs_is_many(self) -> None: + # a single output row per src/dst -> not many + tm = TriMap(2, 2) + tm.register_pairs( + np.array([0, 1], dtype=np.int64), + np.array([1, 0], dtype=np.int64), + ) + tm.finalize() + self.assertFalse(tm.is_many()) + + def test_tri_map_register_pairs_equivalence_loop(self) -> None: + # register_pairs == a register_one loop over the same pairs + rng = np.random.RandomState(0) + for _ in range(50): + src_len = int(rng.randint(1, 12)) + dst_len = int(rng.randint(1, 10)) + n = int(rng.randint(1, 20)) + src_pos = rng.randint(-1, src_len, size=n).astype(np.int64) + dst_pos = rng.randint(-1, dst_len, size=n).astype(np.int64) + tb = TriMap(src_len, dst_len) + tb.register_pairs(src_pos, dst_pos) + tb.finalize() + tl = TriMap(src_len, dst_len) + for i in range(n): + tl.register_one(int(src_pos[i]), int(dst_pos[i])) + tl.finalize() + src = np.arange(100, 100 + src_len) + dst = np.arange(200, 200 + dst_len) + self.assertEqual(tb.is_many(), tl.is_many()) + self.assertEqual( + tb.map_src_fill(src, -1, np.dtype(np.int64)).tolist(), + tl.map_src_fill(src, -1, np.dtype(np.int64)).tolist(), + ) + self.assertEqual( + tb.map_dst_fill(dst, -1, np.dtype(np.int64)).tolist(), + tl.map_dst_fill(dst, -1, np.dtype(np.int64)).tolist(), + ) + + def test_tri_map_register_pairs_equivalence_many_from_one(self) -> None: + # with src = arange(src_len), register_pairs matches register_many_from_one + rng = np.random.RandomState(1) + for _ in range(20): + src_len = int(rng.randint(1, 12)) + dst_len = int(rng.randint(1, 10)) + dst_pos = rng.randint(-1, dst_len, size=src_len).astype(np.int64) + src_pos = np.arange(src_len, dtype=np.int64) + tp = TriMap(src_len, dst_len) + tp.register_pairs(src_pos, dst_pos) + tp.finalize() + tm = TriMap(src_len, dst_len) + tm.register_many_from_one(dst_pos) + tm.finalize() + src = np.arange(100, 100 + src_len) + dst = np.arange(200, 200 + dst_len) + self.assertEqual(tp.is_many(), tm.is_many()) + self.assertEqual( + tp.map_src_fill(src, -1, np.dtype(np.int64)).tolist(), + tm.map_src_fill(src, -1, np.dtype(np.int64)).tolist(), + ) + self.assertEqual( + tp.map_dst_fill(dst, -1, np.dtype(np.int64)).tolist(), + tm.map_dst_fill(dst, -1, np.dtype(np.int64)).tolist(), + ) + + def test_tri_map_register_pairs_empty(self) -> None: + tm = TriMap(3, 3) + tm.register_pairs( + np.array([], dtype=np.int64), np.array([], dtype=np.int64) + ) + tm.finalize() + self.assertEqual(tm.map_src_no_fill(np.arange(3)).tolist(), []) + + def test_tri_map_register_pairs_errors(self) -> None: + with self.assertRaises(ValueError): # length mismatch + TriMap(3, 3).register_pairs( + np.array([0, 1], dtype=np.int64), np.array([0], dtype=np.int64) + ) + with self.assertRaises(ValueError): # wrong dtype + TriMap(3, 3).register_pairs( + np.array([0, 1], dtype=np.int32), np.array([0, 1], dtype=np.int64) + ) + with self.assertRaises(ValueError): # 2d + TriMap(3, 3).register_pairs( + np.array([[0, 1]], dtype=np.int64), np.array([[0, 1]], dtype=np.int64) + ) + with self.assertRaises(ValueError): # out of bounds src + TriMap(3, 3).register_pairs( + np.array([9], dtype=np.int64), np.array([0], dtype=np.int64) + ) + with self.assertRaises(ValueError): # out of bounds dst + TriMap(3, 3).register_pairs( + np.array([0], dtype=np.int64), np.array([9], dtype=np.int64) + ) + with self.assertRaises(TypeError): # not an array + TriMap(3, 3).register_pairs([0, 1], np.array([0, 1], dtype=np.int64)) + tm = TriMap(2, 2) + tm.register_pairs( + np.array([0, 1], dtype=np.int64), np.array([0, 1], dtype=np.int64) + ) + tm.finalize() + with self.assertRaises(RuntimeError): # post-finalize + tm.register_pairs( + np.array([0], dtype=np.int64), np.array([0], dtype=np.int64) + ) From 1dd7848a713cb87d71497aead2a5cad9b4836c0b Mon Sep 17 00:00:00 2001 From: Christopher Ariza Date: Fri, 31 Jul 2026 18:30:08 -0700 Subject: [PATCH 2/2] Enforce alignment Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/tri_map.c | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/tri_map.c b/src/tri_map.c index 003c3d5d..2ff204af 100644 --- a/src/tri_map.c +++ b/src/tri_map.c @@ -344,7 +344,8 @@ TriMap_register_pairs(TriMapObject *self, PyObject *args) { PyErr_SetString(PyExc_ValueError, "Arrays must be 1-dimensional"); return NULL; } - if (!PyArray_IS_C_CONTIGUOUS(src_a) || !PyArray_IS_C_CONTIGUOUS(dst_a)) { + if (!PyArray_IS_C_CONTIGUOUS(src_a) || !PyArray_ISALIGNED(src_a) + || !PyArray_IS_C_CONTIGUOUS(dst_a) || !PyArray_ISALIGNED(dst_a)) { PyErr_SetString(PyExc_ValueError, "Arrays must be contiguous"); return NULL; }