arraykit 1.8.0__tar.gz → 1.9.0__tar.gz

{arraykit-1.8.0/arraykit.egg-info → arraykit-1.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraykit
-Version: 1.8.0
+Version: 1.9.0
 Summary: Array utilities for StaticFrame
 Author: Christopher Ariza, Brandt Bucher, Charles Burkland
 License: MIT
@@ -64,6 +64,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
 ............
 

{arraykit-1.8.0 → arraykit-1.9.0}/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
 ............
 

arraykit-1.9.0/VERSION

@@ -0,0 +1,2 @@
+1.9.0
+

{arraykit-1.8.0 → arraykit-1.9.0/arraykit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraykit
-Version: 1.8.0
+Version: 1.9.0
 Summary: Array utilities for StaticFrame
 Author: Christopher Ariza, Brandt Bucher, Charles Burkland
 License: MIT
@@ -64,6 +64,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
 ............
 

{arraykit-1.8.0 → arraykit-1.9.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: ...

{arraykit-1.8.0 → arraykit-1.9.0}/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},
 };

{arraykit-1.8.0 → arraykit-1.9.0}/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},

{arraykit-1.8.0 → arraykit-1.9.0}/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')
+
+
 # -------------------------------------------------------------------------------
 
 

{arraykit-1.8.0 → arraykit-1.9.0}/test/test_tri_map.py

@@ -1439,3 +1439,72 @@ class TestUnit(unittest.TestCase):
                 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))

arraykit-1.8.0/VERSION

@@ -1,2 +0,0 @@
-1.8.0
-