register_pairs

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
 ............
 

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: ...

src/tri_map.c

@@ -310,6 +310,60 @@ 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_ISALIGNED(src_a)
+            || !PyArray_IS_C_CONTIGUOUS(dst_a) || !PyArray_ISALIGNED(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 +1454,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},

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)
+            )