Skip to content

Commit be87bfa

Browse files
flexatoneflexatone
authored
Merge pull request #220 from static-frame/219/join-util
Join utils
2 parents 4171f7c + 38d362e commit be87bfa

6 files changed

Lines changed: 426 additions & 0 deletions

File tree

‎README.rst‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,12 @@ ArrayKit requires the following:
3535
What is New in ArrayKit
3636
-------------------------
3737

38+
1.9.0
39+
............
40+
41+
Added ``FrozenAutoMap.get_all_fill()`` and ``TriMap.register_many_from_one()``.
42+
43+
3844
1.8.0
3945
............
4046

‎src/__init__.pyi‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ class TriMap:
5353
def __init__(self, /, src_len: int, dst_len: int) -> None: ...
5454
def __repr__(self) -> str: ...
5555
def register_one(self, /, src_from: int, dst_from: int) -> None: ...
56+
def register_many_from_one(self, __dst_pos: np.ndarray) -> None: ...
5657
def register_unmatched_dst(self) -> None: ...
5758
def register_many(self, /, src_from: int, dst_from: np.ndarray) -> None: ...
5859
def finalize(self) -> None: ...
@@ -142,6 +143,7 @@ class FrozenAutoMap:
142143
def items(self) -> tp.Iterator[tuple[_TLabel, int]]: ...
143144
def values(self) -> tp.Iterator[int]: ...
144145
def get_all(self, __key: list[_TLabel] | np.ndarray) -> np.ndarray: ...
146+
def get_all_fill(self, __key: list[_TLabel] | np.ndarray) -> np.ndarray: ...
145147
def get_any(self, __key: list[_TLabel] | np.ndarray) -> list[int]: ...
146148
def __iter__(self) -> tp.Iterator[_TLabel]: ...
147149
def __getitem__(self, __key: tp.Any) -> int: ...

‎src/auto_map.c‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1919,6 +1919,209 @@ fam_get_all(FAMObject *self, PyObject *key) {
19191919
# undef GET_ALL_FLEXIBLE
19201920

19211921

1922+
// Fill variants of the GET_ALL_* macros: on a miss, write -1 into the output slot and
1923+
// continue (instead of raising KeyError). Depend on self, key_size, key_array, i, b, array.
1924+
# define GET_ALL_FILL_SCALARS(npy_type_src, npy_type_dst, kat, lookup_func, hash_func, post_deref) \
1925+
{ \
1926+
npy_type_dst v; \
1927+
Py_ssize_t table_pos; \
1928+
for (; i < key_size; i++) { \
1929+
v = post_deref(*(npy_type_src*)PyArray_GETPTR1(key_array, i)); \
1930+
table_pos = lookup_func(self, v, hash_func(v), kat); \
1931+
if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \
1932+
if (PyErr_Occurred()) { \
1933+
Py_DECREF(array); \
1934+
return NULL; \
1935+
} \
1936+
b[i] = -1; \
1937+
continue; \
1938+
} \
1939+
b[i] = (npy_int64)self->table[table_pos].keys_pos; \
1940+
} \
1941+
} \
1942+
1943+
# define GET_ALL_FILL_DT64(npy_type_src, npy_type_dst, kat, lookup_func, hash_func) \
1944+
{ \
1945+
npy_type_dst v; \
1946+
Py_ssize_t table_pos; \
1947+
for (; i < key_size; i++) { \
1948+
v = *(npy_type_src*)PyArray_GETPTR1(key_array, i); \
1949+
table_pos = lookup_func(self, v, hash_func(v), kat); \
1950+
if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \
1951+
if (PyErr_Occurred()) { \
1952+
Py_DECREF(array); \
1953+
return NULL; \
1954+
} \
1955+
b[i] = -1; \
1956+
continue; \
1957+
} \
1958+
b[i] = (npy_int64)self->table[table_pos].keys_pos; \
1959+
} \
1960+
} \
1961+
1962+
# define GET_ALL_FILL_FLEXIBLE(char_type, get_end_func, lookup_func, hash_func) \
1963+
{ \
1964+
char_type* v; \
1965+
Py_ssize_t dt_size = PyArray_ITEMSIZE(key_array) / sizeof(char_type); \
1966+
Py_ssize_t k_size; \
1967+
Py_ssize_t table_pos; \
1968+
for (; i < key_size; i++) { \
1969+
v = (char_type*)PyArray_GETPTR1(key_array, i); \
1970+
k_size = get_end_func(v, dt_size) - v; \
1971+
table_pos = lookup_func(self, v, k_size, hash_func(v, k_size)); \
1972+
if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \
1973+
if (PyErr_Occurred()) { \
1974+
Py_DECREF(array); \
1975+
return NULL; \
1976+
} \
1977+
b[i] = -1; \
1978+
continue; \
1979+
} \
1980+
b[i] = (npy_int64)self->table[table_pos].keys_pos; \
1981+
} \
1982+
} \
1983+
1984+
// Given a list or array of keys, return an input-aligned int64 array of looked-up
1985+
// positions, with -1 for any key not found (never raises KeyError). This is the
1986+
// "vectorized get with default -1" needed by left/outer joins. Immutable output.
1987+
static PyObject *
1988+
fam_get_all_fill(FAMObject *self, PyObject *key) {
1989+
Py_ssize_t key_size = 0;
1990+
Py_ssize_t keys_pos = -1;
1991+
PyObject* k = NULL;
1992+
PyObject *array = NULL;
1993+
Py_ssize_t i = 0;
1994+
1995+
int key_is_list;
1996+
if (PyList_CheckExact(key)) {
1997+
key_is_list = 1;
1998+
key_size = PyList_GET_SIZE(key);
1999+
}
2000+
else if (PyArray_Check(key)) {
2001+
key_is_list = 0;
2002+
key_size = PyArray_SIZE((PyArrayObject *)key);
2003+
}
2004+
else {
2005+
PyErr_SetString(PyExc_TypeError, "Must provide a list or array.");
2006+
return NULL;
2007+
}
2008+
2009+
npy_intp dims[] = {key_size};
2010+
array = PyArray_EMPTY(1, dims, NPY_INT64, 0);
2011+
if (array == NULL) {
2012+
return NULL;
2013+
}
2014+
npy_int64* b = (npy_int64*)PyArray_DATA((PyArrayObject*)array);
2015+
2016+
if (key_is_list) {
2017+
for (; i < key_size; i++) {
2018+
k = PyList_GET_ITEM(key, i); // borrow
2019+
keys_pos = lookup(self, k);
2020+
if (keys_pos < 0) {
2021+
if (PyErr_Occurred()) {
2022+
Py_DECREF(array);
2023+
return NULL;
2024+
}
2025+
b[i] = -1;
2026+
continue;
2027+
}
2028+
b[i] = (npy_int64)keys_pos;
2029+
}
2030+
}
2031+
else { // key is an array
2032+
PyArrayObject* key_array = (PyArrayObject *)key;
2033+
int key_array_t = PyArray_TYPE(key_array);
2034+
int use_typed = kat_is_kind(self->keys_array_type, PyArray_DESCR(key_array)->kind);
2035+
if (use_typed && key_array_t == NPY_DATETIME) {
2036+
NPY_DATETIMEUNIT key_unit = AK_dt_unit_from_array(key_array);
2037+
if (!kat_is_datetime_unit(self->keys_array_type, key_unit)) {
2038+
// mismatched units: use the coercing scalar path (below) rather than
2039+
// raising, so equal instants still resolve (and true misses give -1)
2040+
use_typed = 0;
2041+
}
2042+
}
2043+
if (use_typed) {
2044+
switch (key_array_t) {
2045+
case NPY_INT64:
2046+
GET_ALL_FILL_SCALARS(npy_int64, npy_int64, KAT_INT64, lookup_hash_int, int_to_hash,);
2047+
break;
2048+
case NPY_INT32:
2049+
GET_ALL_FILL_SCALARS(npy_int32, npy_int64, KAT_INT32, lookup_hash_int, int_to_hash,);
2050+
break;
2051+
case NPY_INT16:
2052+
GET_ALL_FILL_SCALARS(npy_int16, npy_int64, KAT_INT16, lookup_hash_int, int_to_hash,);
2053+
break;
2054+
case NPY_INT8:
2055+
GET_ALL_FILL_SCALARS(npy_int8, npy_int64, KAT_INT8, lookup_hash_int, int_to_hash,);
2056+
break;
2057+
case NPY_UINT64:
2058+
GET_ALL_FILL_SCALARS(npy_uint64, npy_uint64, KAT_UINT64, lookup_hash_uint, uint_to_hash,);
2059+
break;
2060+
case NPY_UINT32:
2061+
GET_ALL_FILL_SCALARS(npy_uint32, npy_uint64, KAT_UINT32, lookup_hash_uint, uint_to_hash,);
2062+
break;
2063+
case NPY_UINT16:
2064+
GET_ALL_FILL_SCALARS(npy_uint16, npy_uint64, KAT_UINT16, lookup_hash_uint, uint_to_hash,);
2065+
break;
2066+
case NPY_UINT8:
2067+
GET_ALL_FILL_SCALARS(npy_uint8, npy_uint64, KAT_UINT8, lookup_hash_uint, uint_to_hash,);
2068+
break;
2069+
case NPY_FLOAT64:
2070+
GET_ALL_FILL_SCALARS(npy_double, npy_double, KAT_FLOAT64, lookup_hash_double, double_to_hash,);
2071+
break;
2072+
case NPY_FLOAT32:
2073+
GET_ALL_FILL_SCALARS(npy_float, npy_double, KAT_FLOAT32, lookup_hash_double, double_to_hash,);
2074+
break;
2075+
case NPY_FLOAT16:
2076+
GET_ALL_FILL_SCALARS(npy_half, npy_double, KAT_FLOAT16, lookup_hash_double, double_to_hash, npy_half_to_double);
2077+
break;
2078+
case NPY_UNICODE:
2079+
GET_ALL_FILL_FLEXIBLE(Py_UCS4, ucs4_get_end_p, lookup_hash_unicode, unicode_to_hash);
2080+
break;
2081+
case NPY_STRING:
2082+
GET_ALL_FILL_FLEXIBLE(char, char_get_end_p, lookup_hash_string, string_to_hash);
2083+
break;
2084+
case NPY_DATETIME:
2085+
GET_ALL_FILL_DT64(npy_int64, npy_int64, KAT_INT64, lookup_hash_int, int_to_hash);
2086+
break;
2087+
default:
2088+
use_typed = 0; // unhandled kind: fall to the scalar path below
2089+
break;
2090+
}
2091+
}
2092+
if (!use_typed) {
2093+
for (; i < key_size; i++) {
2094+
k = PyArray_ToScalar(PyArray_GETPTR1(key_array, i), key_array);
2095+
if (k == NULL) {
2096+
Py_DECREF(array);
2097+
return NULL;
2098+
}
2099+
keys_pos = lookup(self, k);
2100+
if (keys_pos < 0) {
2101+
if (PyErr_Occurred()) {
2102+
Py_DECREF(k);
2103+
Py_DECREF(array);
2104+
return NULL;
2105+
}
2106+
Py_DECREF(k);
2107+
b[i] = -1;
2108+
continue;
2109+
}
2110+
Py_DECREF(k);
2111+
b[i] = (npy_int64)keys_pos;
2112+
}
2113+
}
2114+
}
2115+
2116+
PyArray_CLEARFLAGS((PyArrayObject *)array, NPY_ARRAY_WRITEABLE);
2117+
return array;
2118+
}
2119+
2120+
# undef GET_ALL_FILL_SCALARS
2121+
# undef GET_ALL_FILL_DT64
2122+
# undef GET_ALL_FILL_FLEXIBLE
2123+
2124+
19222125
static inline int
19232126
append_ssize_t(
19242127
PyObject* list,
@@ -2974,6 +3177,7 @@ static PyMethodDef fam_methods[] = {
29743177
{"keys", (PyCFunction) fam_keys, METH_NOARGS, NULL},
29753178
{"values", (PyCFunction) fam_values, METH_NOARGS, NULL},
29763179
{"get_all", (PyCFunction) fam_get_all, METH_O, NULL},
3180+
{"get_all_fill", (PyCFunction) fam_get_all_fill, METH_O, NULL},
29773181
{"get_any", (PyCFunction) fam_get_any, METH_O, NULL},
29783182
{NULL},
29793183
};

‎src/tri_map.c‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,47 @@ TriMap_register_one(TriMapObject *self, PyObject *args) {
269269
Py_RETURN_NONE;
270270
}
271271

272+
// Bulk one-to-one registration: given an int64 array `dst_pos` of length `src_len`,
273+
// register src position i to dst position dst_pos[i] (or -1 for an unmatched src) in a
274+
// single C loop -- equivalent to calling register_one(i, dst_pos[i]) for each i, but
275+
// without per-element Python overhead. Used by the hash-join fast path.
276+
PyObject *
277+
TriMap_register_many_from_one(TriMapObject *self, PyObject *arg) {
278+
if (self->finalized) {
279+
PyErr_SetString(PyExc_RuntimeError, "Cannot register post finalization");
280+
return NULL;
281+
}
282+
if (!PyArray_Check(arg)) {
283+
PyErr_SetString(PyExc_TypeError, "Must provide an array");
284+
return NULL;
285+
}
286+
PyArrayObject* a = (PyArrayObject*)arg;
287+
if (PyArray_TYPE(a) != NPY_INT64) {
288+
PyErr_SetString(PyExc_ValueError, "Array must be of type int64");
289+
return NULL;
290+
}
291+
if (PyArray_NDIM(a) != 1) {
292+
PyErr_SetString(PyExc_ValueError, "Array must be 1-dimensional");
293+
return NULL;
294+
}
295+
if (!PyArray_IS_C_CONTIGUOUS(a)) {
296+
PyErr_SetString(PyExc_ValueError, "Array must be contiguous");
297+
return NULL;
298+
}
299+
npy_intp n = PyArray_SIZE(a);
300+
if (n != self->src_len) {
301+
PyErr_SetString(PyExc_ValueError, "Array length must equal src_len");
302+
return NULL;
303+
}
304+
const npy_int64* d = (npy_int64*)PyArray_DATA(a);
305+
for (npy_intp i = 0; i < n; i++) {
306+
if (AK_TM_register_one(self, (Py_ssize_t)i, (Py_ssize_t)d[i])) {
307+
return NULL;
308+
}
309+
}
310+
Py_RETURN_NONE;
311+
}
312+
272313
PyObject *
273314
TriMap_register_unmatched_dst(TriMapObject *self) {
274315
if (self->finalized) {
@@ -1358,6 +1399,7 @@ TriMap_map_dst_fill(TriMapObject *self, PyObject *args) {
13581399

13591400
static PyMethodDef TriMap_methods[] = {
13601401
{"register_one", (PyCFunction)TriMap_register_one, METH_VARARGS, NULL},
1402+
{"register_many_from_one", (PyCFunction)TriMap_register_many_from_one, METH_O, NULL},
13611403
{"register_unmatched_dst", (PyCFunction)TriMap_register_unmatched_dst, METH_NOARGS, NULL},
13621404
{"register_many", (PyCFunction)TriMap_register_many, METH_VARARGS, NULL},
13631405
{"finalize", (PyCFunction)TriMap_finalize, METH_NOARGS, NULL},

0 commit comments

Comments
 (0)