diff --git a/README.rst b/README.rst index 391b6f3e..6be72869 100644 --- a/README.rst +++ b/README.rst @@ -35,6 +35,12 @@ ArrayKit requires the following: What is New in ArrayKit ------------------------- +1.9.0 +............ + +Added ``FrozenAutoMap.get_all_fill()`` and ``TriMap.register_many_from_one()``. + + 1.8.0 ............ diff --git a/src/__init__.pyi b/src/__init__.pyi index cfd71b26..e856161e 100644 --- a/src/__init__.pyi +++ b/src/__init__.pyi @@ -53,6 +53,7 @@ class TriMap: def __init__(self, /, src_len: int, dst_len: int) -> None: ... 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_unmatched_dst(self) -> None: ... def register_many(self, /, src_from: int, dst_from: np.ndarray) -> None: ... def finalize(self) -> None: ... @@ -142,6 +143,7 @@ class FrozenAutoMap: def items(self) -> tp.Iterator[tuple[_TLabel, int]]: ... def values(self) -> tp.Iterator[int]: ... def get_all(self, __key: list[_TLabel] | np.ndarray) -> np.ndarray: ... + def get_all_fill(self, __key: list[_TLabel] | np.ndarray) -> np.ndarray: ... def get_any(self, __key: list[_TLabel] | np.ndarray) -> list[int]: ... def __iter__(self) -> tp.Iterator[_TLabel]: ... def __getitem__(self, __key: tp.Any) -> int: ... diff --git a/src/auto_map.c b/src/auto_map.c index 4dbe6b3a..6f3298ee 100644 --- a/src/auto_map.c +++ b/src/auto_map.c @@ -1919,6 +1919,209 @@ fam_get_all(FAMObject *self, PyObject *key) { # undef GET_ALL_FLEXIBLE +// Fill variants of the GET_ALL_* macros: on a miss, write -1 into the output slot and +// continue (instead of raising KeyError). Depend on self, key_size, key_array, i, b, array. +# define GET_ALL_FILL_SCALARS(npy_type_src, npy_type_dst, kat, lookup_func, hash_func, post_deref) \ +{ \ + npy_type_dst v; \ + Py_ssize_t table_pos; \ + for (; i < key_size; i++) { \ + v = post_deref(*(npy_type_src*)PyArray_GETPTR1(key_array, i)); \ + table_pos = lookup_func(self, v, hash_func(v), kat); \ + if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \ + if (PyErr_Occurred()) { \ + Py_DECREF(array); \ + return NULL; \ + } \ + b[i] = -1; \ + continue; \ + } \ + b[i] = (npy_int64)self->table[table_pos].keys_pos; \ + } \ +} \ + +# define GET_ALL_FILL_DT64(npy_type_src, npy_type_dst, kat, lookup_func, hash_func) \ +{ \ + npy_type_dst v; \ + Py_ssize_t table_pos; \ + for (; i < key_size; i++) { \ + v = *(npy_type_src*)PyArray_GETPTR1(key_array, i); \ + table_pos = lookup_func(self, v, hash_func(v), kat); \ + if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \ + if (PyErr_Occurred()) { \ + Py_DECREF(array); \ + return NULL; \ + } \ + b[i] = -1; \ + continue; \ + } \ + b[i] = (npy_int64)self->table[table_pos].keys_pos; \ + } \ +} \ + +# define GET_ALL_FILL_FLEXIBLE(char_type, get_end_func, lookup_func, hash_func) \ +{ \ + char_type* v; \ + Py_ssize_t dt_size = PyArray_ITEMSIZE(key_array) / sizeof(char_type); \ + Py_ssize_t k_size; \ + Py_ssize_t table_pos; \ + for (; i < key_size; i++) { \ + v = (char_type*)PyArray_GETPTR1(key_array, i); \ + k_size = get_end_func(v, dt_size) - v; \ + table_pos = lookup_func(self, v, k_size, hash_func(v, k_size)); \ + if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \ + if (PyErr_Occurred()) { \ + Py_DECREF(array); \ + return NULL; \ + } \ + b[i] = -1; \ + continue; \ + } \ + b[i] = (npy_int64)self->table[table_pos].keys_pos; \ + } \ +} \ + +// Given a list or array of keys, return an input-aligned int64 array of looked-up +// positions, with -1 for any key not found (never raises KeyError). This is the +// "vectorized get with default -1" needed by left/outer joins. Immutable output. +static PyObject * +fam_get_all_fill(FAMObject *self, PyObject *key) { + Py_ssize_t key_size = 0; + Py_ssize_t keys_pos = -1; + PyObject* k = NULL; + PyObject *array = NULL; + Py_ssize_t i = 0; + + int key_is_list; + if (PyList_CheckExact(key)) { + key_is_list = 1; + key_size = PyList_GET_SIZE(key); + } + else if (PyArray_Check(key)) { + key_is_list = 0; + key_size = PyArray_SIZE((PyArrayObject *)key); + } + else { + PyErr_SetString(PyExc_TypeError, "Must provide a list or array."); + return NULL; + } + + npy_intp dims[] = {key_size}; + array = PyArray_EMPTY(1, dims, NPY_INT64, 0); + if (array == NULL) { + return NULL; + } + npy_int64* b = (npy_int64*)PyArray_DATA((PyArrayObject*)array); + + if (key_is_list) { + for (; i < key_size; i++) { + k = PyList_GET_ITEM(key, i); // borrow + keys_pos = lookup(self, k); + if (keys_pos < 0) { + if (PyErr_Occurred()) { + Py_DECREF(array); + return NULL; + } + b[i] = -1; + continue; + } + b[i] = (npy_int64)keys_pos; + } + } + else { // key is an array + PyArrayObject* key_array = (PyArrayObject *)key; + int key_array_t = PyArray_TYPE(key_array); + int use_typed = kat_is_kind(self->keys_array_type, PyArray_DESCR(key_array)->kind); + if (use_typed && key_array_t == NPY_DATETIME) { + NPY_DATETIMEUNIT key_unit = AK_dt_unit_from_array(key_array); + if (!kat_is_datetime_unit(self->keys_array_type, key_unit)) { + // mismatched units: use the coercing scalar path (below) rather than + // raising, so equal instants still resolve (and true misses give -1) + use_typed = 0; + } + } + if (use_typed) { + switch (key_array_t) { + case NPY_INT64: + GET_ALL_FILL_SCALARS(npy_int64, npy_int64, KAT_INT64, lookup_hash_int, int_to_hash,); + break; + case NPY_INT32: + GET_ALL_FILL_SCALARS(npy_int32, npy_int64, KAT_INT32, lookup_hash_int, int_to_hash,); + break; + case NPY_INT16: + GET_ALL_FILL_SCALARS(npy_int16, npy_int64, KAT_INT16, lookup_hash_int, int_to_hash,); + break; + case NPY_INT8: + GET_ALL_FILL_SCALARS(npy_int8, npy_int64, KAT_INT8, lookup_hash_int, int_to_hash,); + break; + case NPY_UINT64: + GET_ALL_FILL_SCALARS(npy_uint64, npy_uint64, KAT_UINT64, lookup_hash_uint, uint_to_hash,); + break; + case NPY_UINT32: + GET_ALL_FILL_SCALARS(npy_uint32, npy_uint64, KAT_UINT32, lookup_hash_uint, uint_to_hash,); + break; + case NPY_UINT16: + GET_ALL_FILL_SCALARS(npy_uint16, npy_uint64, KAT_UINT16, lookup_hash_uint, uint_to_hash,); + break; + case NPY_UINT8: + GET_ALL_FILL_SCALARS(npy_uint8, npy_uint64, KAT_UINT8, lookup_hash_uint, uint_to_hash,); + break; + case NPY_FLOAT64: + GET_ALL_FILL_SCALARS(npy_double, npy_double, KAT_FLOAT64, lookup_hash_double, double_to_hash,); + break; + case NPY_FLOAT32: + GET_ALL_FILL_SCALARS(npy_float, npy_double, KAT_FLOAT32, lookup_hash_double, double_to_hash,); + break; + case NPY_FLOAT16: + GET_ALL_FILL_SCALARS(npy_half, npy_double, KAT_FLOAT16, lookup_hash_double, double_to_hash, npy_half_to_double); + break; + case NPY_UNICODE: + GET_ALL_FILL_FLEXIBLE(Py_UCS4, ucs4_get_end_p, lookup_hash_unicode, unicode_to_hash); + break; + case NPY_STRING: + GET_ALL_FILL_FLEXIBLE(char, char_get_end_p, lookup_hash_string, string_to_hash); + break; + case NPY_DATETIME: + GET_ALL_FILL_DT64(npy_int64, npy_int64, KAT_INT64, lookup_hash_int, int_to_hash); + break; + default: + use_typed = 0; // unhandled kind: fall to the scalar path below + break; + } + } + if (!use_typed) { + for (; i < key_size; i++) { + k = PyArray_ToScalar(PyArray_GETPTR1(key_array, i), key_array); + if (k == NULL) { + Py_DECREF(array); + return NULL; + } + keys_pos = lookup(self, k); + if (keys_pos < 0) { + if (PyErr_Occurred()) { + Py_DECREF(k); + Py_DECREF(array); + return NULL; + } + Py_DECREF(k); + b[i] = -1; + continue; + } + Py_DECREF(k); + b[i] = (npy_int64)keys_pos; + } + } + } + + PyArray_CLEARFLAGS((PyArrayObject *)array, NPY_ARRAY_WRITEABLE); + return array; +} + +# undef GET_ALL_FILL_SCALARS +# undef GET_ALL_FILL_DT64 +# undef GET_ALL_FILL_FLEXIBLE + + static inline int append_ssize_t( PyObject* list, @@ -2974,6 +3177,7 @@ static PyMethodDef fam_methods[] = { {"keys", (PyCFunction) fam_keys, METH_NOARGS, NULL}, {"values", (PyCFunction) fam_values, METH_NOARGS, NULL}, {"get_all", (PyCFunction) fam_get_all, METH_O, NULL}, + {"get_all_fill", (PyCFunction) fam_get_all_fill, METH_O, NULL}, {"get_any", (PyCFunction) fam_get_any, METH_O, NULL}, {NULL}, }; diff --git a/src/tri_map.c b/src/tri_map.c index 6eb0c7d1..2168bb01 100644 --- a/src/tri_map.c +++ b/src/tri_map.c @@ -269,6 +269,47 @@ TriMap_register_one(TriMapObject *self, PyObject *args) { Py_RETURN_NONE; } +// Bulk one-to-one registration: given an int64 array `dst_pos` of length `src_len`, +// register src position i to dst position dst_pos[i] (or -1 for an unmatched src) in a +// single C loop -- equivalent to calling register_one(i, dst_pos[i]) for each i, but +// without per-element Python overhead. Used by the hash-join fast path. +PyObject * +TriMap_register_many_from_one(TriMapObject *self, PyObject *arg) { + if (self->finalized) { + PyErr_SetString(PyExc_RuntimeError, "Cannot register post finalization"); + return NULL; + } + if (!PyArray_Check(arg)) { + PyErr_SetString(PyExc_TypeError, "Must provide an array"); + return NULL; + } + PyArrayObject* a = (PyArrayObject*)arg; + if (PyArray_TYPE(a) != NPY_INT64) { + PyErr_SetString(PyExc_ValueError, "Array must be of type int64"); + return NULL; + } + if (PyArray_NDIM(a) != 1) { + PyErr_SetString(PyExc_ValueError, "Array must be 1-dimensional"); + return NULL; + } + if (!PyArray_IS_C_CONTIGUOUS(a)) { + PyErr_SetString(PyExc_ValueError, "Array must be contiguous"); + return NULL; + } + npy_intp n = PyArray_SIZE(a); + if (n != self->src_len) { + PyErr_SetString(PyExc_ValueError, "Array length must equal src_len"); + return NULL; + } + const npy_int64* d = (npy_int64*)PyArray_DATA(a); + for (npy_intp i = 0; i < n; i++) { + if (AK_TM_register_one(self, (Py_ssize_t)i, (Py_ssize_t)d[i])) { + return NULL; + } + } + Py_RETURN_NONE; +} + PyObject * TriMap_register_unmatched_dst(TriMapObject *self) { if (self->finalized) { @@ -1358,6 +1399,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_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_auto_map.py b/test/test_auto_map.py index a674bd92..e6ebade4 100644 --- a/test/test_auto_map.py +++ b/test/test_auto_map.py @@ -916,6 +916,109 @@ def test_fam_array_get_all_m3(): assert post2.tolist() == [2, 1] +# ------------------------------------------------------------------------------- +# get_all_fill: input-aligned lookup, -1 for misses (never raises KeyError) + + +def test_fam_array_get_all_fill_int_a(): + a1 = np.array((10, 20, 30, 40), dtype=np.int64) + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + # all present -> identical to get_all + keys = np.array([30, 10, 40], dtype=np.int64) + keys.flags.writeable = False + assert fam.get_all_fill(keys).tolist() == [2, 0, 3] + assert fam.get_all_fill(keys).tolist() == fam.get_all(keys).tolist() + + +def test_fam_array_get_all_fill_int_missing(): + a1 = np.array((10, 20, 30, 40), dtype=np.int64) + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + keys = np.array([30, 99, 10, 77], dtype=np.int64) + keys.flags.writeable = False + assert fam.get_all_fill(keys).tolist() == [2, -1, 0, -1] + + +def test_fam_array_get_all_fill_immutable_and_dtype(): + a1 = np.array((10, 20, 30), dtype=np.int64) + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + keys = np.array([10], dtype=np.int64) + keys.flags.writeable = False + post = fam.get_all_fill(keys) + assert not post.flags.writeable + assert post.dtype == np.dtype(np.int64) + + +def test_fam_array_get_all_fill_empty(): + a1 = np.array((10, 20), dtype=np.int64) + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + keys = np.array([], dtype=np.int64) + keys.flags.writeable = False + assert fam.get_all_fill(keys).tolist() == [] + + +def test_fam_array_get_all_fill_all_missing(): + a1 = np.array((10, 20), dtype=np.int64) + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + keys = np.array([1, 2, 3], dtype=np.int64) + keys.flags.writeable = False + assert fam.get_all_fill(keys).tolist() == [-1, -1, -1] + + +def test_fam_array_get_all_fill_kind_mismatch(): + # int map queried with a float array -> scalar fallback, -1 for misses + a1 = np.array((10, 20, 30), dtype=np.int64) + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + keys = np.array([10.0, 15.0, 30.0], dtype=np.float64) + keys.flags.writeable = False + assert fam.get_all_fill(keys).tolist() == [0, -1, 2] + + +def test_fam_array_get_all_fill_unicode(): + a1 = np.array(['a', 'bb', 'c']) + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + keys = np.array(['c', 'zz', 'a']) + keys.flags.writeable = False + assert fam.get_all_fill(keys).tolist() == [2, -1, 0] + + +def test_fam_array_get_all_fill_bytes(): + a1 = np.array([b'a', b'bb', b'c'], dtype='S2') + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + keys = np.array([b'c', b'zz', b'a'], dtype='S2') + keys.flags.writeable = False + assert fam.get_all_fill(keys).tolist() == [2, -1, 0] + + +def test_fam_array_get_all_fill_datetime(): + a1 = np.array(['2020-01', '2021-06', '2022-03'], dtype='datetime64[M]') + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + keys = np.array(['2022-03', '1999-12', '2020-01'], dtype='datetime64[M]') + keys.flags.writeable = False + assert fam.get_all_fill(keys).tolist() == [2, -1, 0] + + +def test_fam_array_get_all_fill_object_list(): + fam = FrozenAutoMap([('x',), ('y',)]) + assert fam.get_all_fill([('y',), ('z',), ('x',)]).tolist() == [1, -1, 0] + + +def test_fam_array_get_all_fill_bad_type(): + a1 = np.array((10, 20), dtype=np.int64) + a1.flags.writeable = False + fam = FrozenAutoMap(a1) + with pytest.raises(TypeError): + fam.get_all_fill('a') + + # ------------------------------------------------------------------------------- diff --git a/test/test_tri_map.py b/test/test_tri_map.py index e1708722..cb73fcbe 100644 --- a/test/test_tri_map.py +++ b/test/test_tri_map.py @@ -1439,3 +1439,72 @@ def test_tri_map_merge_i(self) -> None: np.datetime64('2005-11'), ], ) + + # ------------------------------------------------------------------ + # register_many_from_one (bulk one-to-one) + + def test_tri_map_register_many_from_one_a(self) -> None: + tm = TriMap(3, 4) + tm.register_many_from_one(np.array([2, 0, 3], dtype=np.int64)) + tm.finalize() + self.assertFalse(tm.is_many()) + src = np.array([10, 20, 30]) + dst = np.array([100, 200, 300, 400]) + self.assertEqual(tm.map_src_no_fill(src).tolist(), [10, 20, 30]) + self.assertEqual(tm.map_dst_no_fill(dst).tolist(), [300, 100, 400]) + + def test_tri_map_register_many_from_one_unmatched(self) -> None: + # -1 marks an unmatched src -> a src-only row (dst side gets a fill there) + tm = TriMap(3, 4) + tm.register_many_from_one(np.array([2, -1, 0], dtype=np.int64)) + tm.finalize() + self.assertTrue(tm.src_no_fill()) # every output row has a src value + self.assertFalse(tm.dst_no_fill()) # row 1 has no dst -> needs fill + dst = np.array([100, 200, 300, 400]) + self.assertEqual(tm.map_dst_fill(dst, -9, np.dtype(np.int64)).tolist(), [300, -9, 100]) + + def test_tri_map_register_many_from_one_is_many(self) -> None: + # two src rows matching the same dst -> is_many + tm = TriMap(3, 3) + tm.register_many_from_one(np.array([0, 0, 1], dtype=np.int64)) + tm.finalize() + self.assertTrue(tm.is_many()) + + def test_tri_map_register_many_from_one_equivalence(self) -> None: + rng = np.random.RandomState(0) + for _ in range(50): + 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) + tb = TriMap(src_len, dst_len) + tb.register_many_from_one(dst_pos) + tb.finalize() + tl = TriMap(src_len, dst_len) + for i in range(src_len): + tl.register_one(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_no_fill(src).tolist(), tl.map_src_no_fill(src).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_many_from_one_errors(self) -> None: + with self.assertRaises(ValueError): # wrong length + TriMap(3, 3).register_many_from_one(np.array([0, 1], dtype=np.int64)) + with self.assertRaises(ValueError): # wrong dtype + TriMap(3, 3).register_many_from_one(np.array([0, 1, 2], dtype=np.int32)) + with self.assertRaises(ValueError): # out of bounds dst + TriMap(3, 3).register_many_from_one(np.array([0, 1, 9], dtype=np.int64)) + with self.assertRaises(TypeError): # not an array + TriMap(3, 3).register_many_from_one([0, 1, 2]) + tm = TriMap(2, 2) + tm.register_many_from_one(np.array([0, 1], dtype=np.int64)) + tm.finalize() + with self.assertRaises(RuntimeError): # post-finalize + tm.register_many_from_one(np.array([0, 1], dtype=np.int64))