Join utils

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

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

src/auto_map.c

@@ -325,17 +325,17 @@
     }
     x = 0;
     while (m) {
         x = ((x << 28) & HASH_MODULUS) | x >> (HASH_BITS - 28);
         m *= 268435456.0;  /* 2**28 */
         e -= 28;
         y = (Py_uhash_t)m;  /* pull out integer part */
         m -= y;
         x += y;
         if (x >= HASH_MODULUS)
             x -= HASH_MODULUS;
     }
     e = e >= 0 ? e % HASH_BITS : HASH_BITS-1-((-1-e) % HASH_BITS);
     x = ((x << e) & HASH_MODULUS) | x >> (HASH_BITS - e);
     x = x * sign;
     if (x == (Py_uhash_t)-1)
         x = (Py_uhash_t)-2;
@@ -1505,7 +1505,7 @@
 static int
 grow_table(FAMObject *self, Py_ssize_t keys_size)
 {
     Py_ssize_t keys_load = keys_size / LOAD;
     Py_ssize_t size_old = self->table_size;
     if (keys_load < size_old) {
         return 0;
@@ -1919,6 +1919,209 @@
 # 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 @@
     {"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},
 };

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},