From 33aa81527c98c440e5a7c873b6d58c5b4eb22d87 Mon Sep 17 00:00:00 2001 From: David Fulford Date: Sat, 28 Mar 2026 18:00:48 -0500 Subject: [PATCH 01/10] =?UTF-8?q?perf:=20=5Faccel.c=20v2=20=E2=80=94=20pow?= =?UTF-8?q?er-of-2=20hashing,=20CSR=20joins,=20single-pass=20output?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace modulo hash probing with bitmask (& mask) on power-of-2 tables - Replace linked-list join hash table with CSR (flat contiguous row buffer) - Eliminate double-pass join probing with growable IndexBuf output - Unroll composite_key inner loop for 1-3 column fast paths - Use memcpy for first_seen output, PyArray_EMPTY where safe Co-Authored-By: Claude Opus 4.6 (1M context) --- tafra/_accel.c | 424 +++++++++++++++++++++++++++++++------------------ 1 file changed, 269 insertions(+), 155 deletions(-) diff --git a/tafra/_accel.c b/tafra/_accel.c index 8971962..673b14e 100644 --- a/tafra/_accel.c +++ b/tafra/_accel.c @@ -1,8 +1,12 @@ /* - * tafra/_accel.c — Minimal C extension for hot-path acceleration. + * tafra/_accel.c — C extension for hot-path acceleration (v2). * - * Provides single-pass grouped aggregation and hash-based join index - * construction, eliminating multiple numpy passes and temporary arrays. + * Performance improvements over v1: + * - Power-of-2 hash tables with bitmask probing (no division) + * - CSR (flat contiguous) layout for join hash tables (no linked lists) + * - Single-pass join index construction (no counting pre-pass) + * - Unrolled composite_key for 1-3 column fast paths + * - Reduced redundant initialization * * All functions accept and return numpy arrays. Falls back gracefully * if this module fails to compile — pure-Python paths remain available. @@ -14,6 +18,7 @@ #define PY_SSIZE_T_CLEAN #include #include +#include #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION #include @@ -21,9 +26,21 @@ #include /* Golden ratio constant for multiplicative hashing (Knuth). - * phi = (sqrt(5)-1)/2 * 2^64 ≈ 0x9E3779B97F4A7C15 */ + * phi = (sqrt(5)-1)/2 * 2^64 ~ 0x9E3779B97F4A7C15 */ #define GOLDEN_HASH 0x9E3779B97F4A7C15ULL -#define HASH_KEY(k, sz) ((npy_intp)(((npy_uint64)(k) * GOLDEN_HASH) >> 1) % (sz)) + +/* Power-of-2 hash: multiply then mask. table_size must be power of 2. */ +#define HASH_KEY(k, mask) ((npy_intp)(((npy_uint64)(k) * GOLDEN_HASH) >> 1) & (mask)) + + +/* Round up to next power of 2 (minimum 16). */ +static npy_intp +next_pow2(npy_intp n) +{ + npy_intp p = 16; + while (p < n) p <<= 1; + return p; +} /* Helper: coerce to contiguous array of given type. Caller must Py_DECREF. */ @@ -305,7 +322,7 @@ accel_groupby_max(PyObject *self, PyObject *args) * cardinalities: tuple of ints (max value + 1 for each column) * * Computes: key[i] = col0[i]*card1*card2*... + col1[i]*card2*... + ... - * Single pass, no temporaries. + * Unrolled fast paths for 1-3 columns (common case). * ================================================================ */ static PyObject * @@ -375,9 +392,9 @@ accel_composite_key(PyObject *self, PyObject *args) for (Py_ssize_t c = n_cols - 2; c >= 0; c--) mult[c] = mult[c + 1] * cards[c + 1]; - /* Build output in single pass */ + /* Build output */ npy_intp dims[1] = {n_rows}; - PyArrayObject *out = (PyArrayObject *)PyArray_ZEROS(1, dims, NPY_INT64, 0); + PyArrayObject *out = (PyArrayObject *)PyArray_EMPTY(1, dims, NPY_INT64, 0); if (!out) { for (Py_ssize_t c = 0; c < n_cols; c++) Py_DECREF(col_arrs[c]); free(col_arrs); free(cards); free(mult); @@ -385,27 +402,46 @@ accel_composite_key(PyObject *self, PyObject *args) } npy_int64 *result = (npy_int64 *)PyArray_DATA(out); - /* Get data pointers */ - const npy_int64 **col_data = (const npy_int64 **)malloc(n_cols * sizeof(npy_int64 *)); - if (!col_data) { - for (Py_ssize_t c = 0; c < n_cols; c++) Py_DECREF(col_arrs[c]); - free(col_arrs); free(cards); free(mult); - Py_DECREF(out); - return PyErr_NoMemory(); - } - for (Py_ssize_t c = 0; c < n_cols; c++) - col_data[c] = (const npy_int64 *)PyArray_DATA(col_arrs[c]); - - /* Single pass: result[i] = sum(col_data[c][i] * mult[c]) */ - for (npy_intp i = 0; i < n_rows; i++) { - npy_int64 key = 0; + /* Unrolled fast paths for common column counts */ + if (n_cols == 1) { + const npy_int64 *d0 = (const npy_int64 *)PyArray_DATA(col_arrs[0]); + memcpy(result, d0, n_rows * sizeof(npy_int64)); + } else if (n_cols == 2) { + const npy_int64 *d0 = (const npy_int64 *)PyArray_DATA(col_arrs[0]); + const npy_int64 *d1 = (const npy_int64 *)PyArray_DATA(col_arrs[1]); + npy_int64 m0 = mult[0]; + for (npy_intp i = 0; i < n_rows; i++) + result[i] = d0[i] * m0 + d1[i]; + } else if (n_cols == 3) { + const npy_int64 *d0 = (const npy_int64 *)PyArray_DATA(col_arrs[0]); + const npy_int64 *d1 = (const npy_int64 *)PyArray_DATA(col_arrs[1]); + const npy_int64 *d2 = (const npy_int64 *)PyArray_DATA(col_arrs[2]); + npy_int64 m0 = mult[0], m1 = mult[1]; + for (npy_intp i = 0; i < n_rows; i++) + result[i] = d0[i] * m0 + d1[i] * m1 + d2[i]; + } else { + /* General case */ + const npy_int64 **col_data = (const npy_int64 **)malloc(n_cols * sizeof(npy_int64 *)); + if (!col_data) { + for (Py_ssize_t c = 0; c < n_cols; c++) Py_DECREF(col_arrs[c]); + free(col_arrs); free(cards); free(mult); + Py_DECREF(out); + return PyErr_NoMemory(); + } for (Py_ssize_t c = 0; c < n_cols; c++) - key += col_data[c][i] * mult[c]; - result[i] = key; + col_data[c] = (const npy_int64 *)PyArray_DATA(col_arrs[c]); + + for (npy_intp i = 0; i < n_rows; i++) { + npy_int64 key = 0; + for (Py_ssize_t c = 0; c < n_cols; c++) + key += col_data[c][i] * mult[c]; + result[i] = key; + } + free((void *)col_data); } for (Py_ssize_t c = 0; c < n_cols; c++) Py_DECREF(col_arrs[c]); - free(col_arrs); free(cards); free(mult); free((void *)col_data); + free(col_arrs); free(cards); free(mult); return (PyObject *)out; } @@ -420,7 +456,7 @@ accel_composite_key(PyObject *self, PyObject *args) * codes: int64 array of integer codes (0..n_unique-1), first-seen order * n_unique: int * - * Replaces np.unique(return_inverse=True) which is O(n log n). + * Power-of-2 table with bitmask probing. * ================================================================ */ typedef struct { @@ -444,9 +480,9 @@ accel_encode_strings(PyObject *self, PyObject *args) npy_intp n = PyArray_SIZE(arr); PyObject **data = (PyObject **)PyArray_DATA(arr); - /* Hash table */ - npy_intp table_size = n * 2; - if (table_size < 16) table_size = 16; + /* Power-of-2 hash table at ~50% load factor */ + npy_intp table_size = next_pow2(n * 2); + npy_intp mask = table_size - 1; StrHashEntry *table = (StrHashEntry *)calloc(table_size, sizeof(StrHashEntry)); if (!table) { Py_DECREF(arr); return PyErr_NoMemory(); } @@ -466,24 +502,23 @@ accel_encode_strings(PyObject *self, PyObject *args) return NULL; } - npy_intp slot = HASH_KEY(h, table_size); + npy_intp slot = (npy_intp)(((npy_uint64)h * GOLDEN_HASH) >> 1) & mask; while (table[slot].occupied) { - /* Check if same object or equal value */ if (table[slot].hash == h) { int eq = PyObject_RichCompareBool(table[slot].obj, obj, Py_EQ); if (eq < 0) { free(table); Py_DECREF(arr); Py_DECREF(codes_out); return NULL; } - if (eq) break; /* found existing entry */ + if (eq) break; } - slot = (slot + 1) % table_size; + slot = (slot + 1) & mask; } if (!table[slot].occupied) { table[slot].hash = h; - table[slot].obj = obj; /* borrowed ref, arr keeps it alive */ + table[slot].obj = obj; table[slot].code = n_unique; table[slot].occupied = 1; n_unique++; @@ -506,17 +541,12 @@ accel_encode_strings(PyObject *self, PyObject *args) * group_indices(key_array) -> (first_seen_indices, group_row_lists, n_groups) * key_array: int64 array (composite or single-column key) * - * Returns: - * first_seen_indices: int64 array of length n_groups - * group_row_lists: Python list of int64 arrays (row indices per group) - * n_groups: int - * - * Replaces np.unique + argsort + split with a single O(n) pass. + * Power-of-2 table with bitmask probing. * ================================================================ */ typedef struct { npy_int64 key; - npy_intp label; /* assigned group label */ + npy_intp label; int occupied; } GroupHashEntry; @@ -533,9 +563,9 @@ accel_group_indices(PyObject *self, PyObject *args) npy_intp n = PyArray_SIZE(key_arr); const npy_int64 *keys = (const npy_int64 *)PyArray_DATA(key_arr); - /* Allocate hash table */ - npy_intp table_size = n * 2; - if (table_size < 16) table_size = 16; + /* Power-of-2 hash table */ + npy_intp table_size = next_pow2(n * 2); + npy_intp mask = table_size - 1; GroupHashEntry *table = (GroupHashEntry *)calloc(table_size, sizeof(GroupHashEntry)); if (!table) { Py_DECREF(key_arr); return PyErr_NoMemory(); } @@ -553,9 +583,9 @@ accel_group_indices(PyObject *self, PyObject *args) for (npy_intp i = 0; i < n; i++) { npy_int64 k = keys[i]; - npy_intp h = HASH_KEY(k, table_size); + npy_intp h = HASH_KEY(k, mask); while (table[h].occupied && table[h].key != k) - h = (h + 1) % table_size; + h = (h + 1) & mask; if (!table[h].occupied) { table[h].key = k; @@ -612,7 +642,7 @@ accel_group_indices(PyObject *self, PyObject *args) return NULL; } group_data[g] = (npy_intp *)PyArray_DATA(arr); - PyList_SET_ITEM(group_list, g, (PyObject *)arr); /* steals ref */ + PyList_SET_ITEM(group_list, g, (PyObject *)arr); } free(counts); @@ -629,9 +659,7 @@ accel_group_indices(PyObject *self, PyObject *args) npy_intp fs_dims[1] = {n_groups}; PyArrayObject *fs_out = (PyArrayObject *)PyArray_EMPTY(1, fs_dims, NPY_INTP, 0); if (!fs_out) { free(first_seen); Py_DECREF(group_list); return NULL; } - npy_intp *fs_data = (npy_intp *)PyArray_DATA(fs_out); - for (npy_intp g = 0; g < n_groups; g++) - fs_data[g] = first_seen[g]; + memcpy(PyArray_DATA(fs_out), first_seen, n_groups * sizeof(npy_intp)); free(first_seen); PyObject *result = Py_BuildValue("(OOn)", fs_out, group_list, (Py_ssize_t)n_groups); @@ -642,66 +670,137 @@ accel_group_indices(PyObject *self, PyObject *args) /* ================================================================ - * Hash join: build (left_indices, right_indices) for equi-join + * Hash join v2: CSR layout + single-pass index construction * - * Uses a simple open-addressing hash table on the right key, - * then probes with each left key. + * Build phase: + * 1. Hash right keys into power-of-2 table, count per bucket. + * 2. Prefix sum -> offsets into flat contiguous row buffer. + * 3. Scatter right row indices into flat buffer. * - * Input arrays are coerced to contiguous int64. + * Probe phase reads sequential memory per bucket (cache-friendly). + * + * Output construction uses a growable buffer (realloc doubling) + * to avoid a separate counting pre-pass over left keys. * ================================================================ */ typedef struct { npy_int64 key; - npy_intp first; /* index into chain array */ - npy_intp count; -} HashEntry; - -typedef struct { - npy_intp row; - npy_intp next; /* -1 = end */ -} ChainNode; + npy_intp offset; /* start index into flat row buffer */ + npy_intp count; /* number of right rows with this key */ +} CSRHashEntry; -/* Build hash table on right-side keys. Returns 0 on success, -1 on error. */ +/* Build CSR hash table on right-side keys. + * Returns 0 on success, -1 on error. + * Caller must free *out_table and *out_rows. */ static int -build_hash_table(const npy_int64 *right, npy_intp right_n, - HashEntry **out_table, npy_intp *out_table_size, - ChainNode **out_chain) +build_csr_hash_table(const npy_int64 *right, npy_intp right_n, + CSRHashEntry **out_table, npy_intp *out_mask, + npy_intp **out_rows) { - npy_intp table_size = right_n * 2; - if (table_size < 16) table_size = 16; + npy_intp table_size = next_pow2(right_n * 2); + npy_intp mask = table_size - 1; - HashEntry *table = (HashEntry *)calloc(table_size, sizeof(HashEntry)); - ChainNode *chain = (ChainNode *)malloc(right_n * sizeof(ChainNode)); - if (!table || !chain) { - free(table); free(chain); - PyErr_NoMemory(); - return -1; + /* calloc zeros count and offset to 0; count==0 means empty slot */ + CSRHashEntry *table = (CSRHashEntry *)calloc(table_size, sizeof(CSRHashEntry)); + if (!table) { PyErr_NoMemory(); return -1; } + + /* Pass 1: count occurrences per key */ + for (npy_intp i = 0; i < right_n; i++) { + npy_int64 k = right[i]; + npy_intp h = HASH_KEY(k, mask); + while (table[h].count > 0 && table[h].key != k) + h = (h + 1) & mask; + table[h].key = k; + table[h].count++; } - for (npy_intp i = 0; i < table_size; i++) { - table[i].count = 0; - table[i].first = -1; + /* Pass 2: prefix sum to compute offsets */ + npy_intp running = 0; + for (npy_intp h = 0; h < table_size; h++) { + if (table[h].count > 0) { + table[h].offset = running; + running += table[h].count; + } } + /* Allocate flat row buffer */ + npy_intp *rows = (npy_intp *)malloc(right_n * sizeof(npy_intp)); + if (!rows) { free(table); PyErr_NoMemory(); return -1; } + + /* Temporary write positions (reuse a stack of counts) */ + npy_intp *write_pos = (npy_intp *)calloc(table_size, sizeof(npy_intp)); + if (!write_pos) { free(table); free(rows); PyErr_NoMemory(); return -1; } + + /* Pass 3: scatter right row indices into flat buffer */ for (npy_intp i = 0; i < right_n; i++) { npy_int64 k = right[i]; - npy_intp h = HASH_KEY(k, table_size); + npy_intp h = HASH_KEY(k, mask); while (table[h].count > 0 && table[h].key != k) - h = (h + 1) % table_size; - chain[i].row = i; - chain[i].next = table[h].first; - table[h].key = k; - table[h].first = i; - table[h].count++; + h = (h + 1) & mask; + rows[table[h].offset + write_pos[h]] = i; + write_pos[h]++; } + free(write_pos); *out_table = table; - *out_table_size = table_size; - *out_chain = chain; + *out_mask = mask; + *out_rows = rows; + return 0; +} + + +/* Growable index buffer for single-pass join output */ +typedef struct { + npy_intp *data; + npy_intp len; + npy_intp cap; +} IndexBuf; + +static int +indexbuf_init(IndexBuf *buf, npy_intp initial_cap) +{ + buf->data = (npy_intp *)malloc(initial_cap * sizeof(npy_intp)); + if (!buf->data) { PyErr_NoMemory(); return -1; } + buf->len = 0; + buf->cap = initial_cap; + return 0; +} + +static int +indexbuf_ensure(IndexBuf *buf, npy_intp additional) +{ + if (buf->len + additional <= buf->cap) return 0; + npy_intp new_cap = buf->cap; + while (new_cap < buf->len + additional) new_cap <<= 1; + npy_intp *tmp = (npy_intp *)realloc(buf->data, new_cap * sizeof(npy_intp)); + if (!tmp) { PyErr_NoMemory(); return -1; } + buf->data = tmp; + buf->cap = new_cap; return 0; } +static void +indexbuf_free(IndexBuf *buf) +{ + free(buf->data); + buf->data = NULL; + buf->len = buf->cap = 0; +} + +/* Transfer ownership of buffer contents into a numpy array. + * On success, buf->data is set to NULL (numpy owns it via its own copy). + * Returns new reference or NULL on error. */ +static PyArrayObject * +indexbuf_to_array(IndexBuf *buf) +{ + npy_intp dims[1] = {buf->len}; + PyArrayObject *arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + if (!arr) return NULL; + memcpy(PyArray_DATA(arr), buf->data, buf->len * sizeof(npy_intp)); + return arr; +} + static PyObject * accel_inner_join(PyObject *self, PyObject *args) @@ -722,54 +821,58 @@ accel_inner_join(PyObject *self, PyObject *args) const npy_int64 *left = (const npy_int64 *)PyArray_DATA(left_arr); const npy_int64 *right = (const npy_int64 *)PyArray_DATA(right_arr); - HashEntry *table; npy_intp table_size; ChainNode *chain; - if (build_hash_table(right, right_n, &table, &table_size, &chain) < 0) { + CSRHashEntry *table; npy_intp mask; npy_intp *rows; + if (build_csr_hash_table(right, right_n, &table, &mask, &rows) < 0) { Py_DECREF(left_arr); Py_DECREF(right_arr); return NULL; } - /* Count total output rows */ - npy_intp total = 0; - for (npy_intp i = 0; i < left_n; i++) { - npy_intp h = HASH_KEY(left[i], table_size); - while (table[h].count > 0 && table[h].key != left[i]) - h = (h + 1) % table_size; - if (table[h].count > 0 && table[h].key == left[i]) - total += table[h].count; - } - - npy_intp dims[1] = {total}; - PyArrayObject *li_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); - PyArrayObject *ri_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); - if (!li_arr || !ri_arr) { - Py_XDECREF(li_arr); Py_XDECREF(ri_arr); - free(table); free(chain); + /* Single-pass probe with growable output */ + npy_intp init_cap = left_n; /* reasonable starting guess */ + if (init_cap < 64) init_cap = 64; + IndexBuf li_buf, ri_buf; + if (indexbuf_init(&li_buf, init_cap) < 0 || + indexbuf_init(&ri_buf, init_cap) < 0) { + indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + free(table); free(rows); Py_DECREF(left_arr); Py_DECREF(right_arr); return NULL; } - npy_intp *li = (npy_intp *)PyArray_DATA(li_arr); - npy_intp *ri = (npy_intp *)PyArray_DATA(ri_arr); - npy_intp pos = 0; - for (npy_intp i = 0; i < left_n; i++) { - npy_intp h = HASH_KEY(left[i], table_size); + npy_intp h = HASH_KEY(left[i], mask); while (table[h].count > 0 && table[h].key != left[i]) - h = (h + 1) % table_size; + h = (h + 1) & mask; if (table[h].count > 0 && table[h].key == left[i]) { - npy_intp ci = table[h].first; - while (ci >= 0) { - li[pos] = i; - ri[pos] = chain[ci].row; - pos++; - ci = chain[ci].next; + npy_intp cnt = table[h].count; + npy_intp off = table[h].offset; + if (indexbuf_ensure(&li_buf, cnt) < 0 || + indexbuf_ensure(&ri_buf, cnt) < 0) { + indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + free(table); free(rows); + Py_DECREF(left_arr); Py_DECREF(right_arr); + return NULL; + } + for (npy_intp j = 0; j < cnt; j++) { + li_buf.data[li_buf.len] = i; + ri_buf.data[ri_buf.len] = rows[off + j]; + li_buf.len++; + ri_buf.len++; } } } - free(table); free(chain); + free(table); free(rows); Py_DECREF(left_arr); Py_DECREF(right_arr); + PyArrayObject *li_arr = indexbuf_to_array(&li_buf); + PyArrayObject *ri_arr = indexbuf_to_array(&ri_buf); + indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + if (!li_arr || !ri_arr) { + Py_XDECREF(li_arr); Py_XDECREF(ri_arr); + return NULL; + } + PyObject *result = PyTuple_Pack(2, (PyObject *)li_arr, (PyObject *)ri_arr); Py_DECREF(li_arr); Py_DECREF(ri_arr); return result; @@ -795,62 +898,73 @@ accel_left_join(PyObject *self, PyObject *args) const npy_int64 *left = (const npy_int64 *)PyArray_DATA(left_arr); const npy_int64 *right = (const npy_int64 *)PyArray_DATA(right_arr); - HashEntry *table; npy_intp table_size; ChainNode *chain; - if (build_hash_table(right, right_n, &table, &table_size, &chain) < 0) { + CSRHashEntry *table; npy_intp mask; npy_intp *rows; + if (build_csr_hash_table(right, right_n, &table, &mask, &rows) < 0) { Py_DECREF(left_arr); Py_DECREF(right_arr); return NULL; } - npy_intp total = 0; - int has_null = 0; - for (npy_intp i = 0; i < left_n; i++) { - npy_intp h = HASH_KEY(left[i], table_size); - while (table[h].count > 0 && table[h].key != left[i]) - h = (h + 1) % table_size; - if (table[h].count > 0 && table[h].key == left[i]) - total += table[h].count; - else { - total++; - has_null = 1; - } - } - - npy_intp dims[1] = {total}; - PyArrayObject *li_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); - PyArrayObject *ri_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); - if (!li_arr || !ri_arr) { - Py_XDECREF(li_arr); Py_XDECREF(ri_arr); - free(table); free(chain); + /* Single-pass probe with growable output */ + npy_intp init_cap = left_n; + if (init_cap < 64) init_cap = 64; + IndexBuf li_buf, ri_buf; + if (indexbuf_init(&li_buf, init_cap) < 0 || + indexbuf_init(&ri_buf, init_cap) < 0) { + indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + free(table); free(rows); Py_DECREF(left_arr); Py_DECREF(right_arr); return NULL; } - npy_intp *li = (npy_intp *)PyArray_DATA(li_arr); - npy_intp *ri = (npy_intp *)PyArray_DATA(ri_arr); - npy_intp pos = 0; + int has_null = 0; for (npy_intp i = 0; i < left_n; i++) { - npy_intp h = HASH_KEY(left[i], table_size); + npy_intp h = HASH_KEY(left[i], mask); while (table[h].count > 0 && table[h].key != left[i]) - h = (h + 1) % table_size; + h = (h + 1) & mask; if (table[h].count > 0 && table[h].key == left[i]) { - npy_intp ci = table[h].first; - while (ci >= 0) { - li[pos] = i; - ri[pos] = chain[ci].row; - pos++; - ci = chain[ci].next; + npy_intp cnt = table[h].count; + npy_intp off = table[h].offset; + if (indexbuf_ensure(&li_buf, cnt) < 0 || + indexbuf_ensure(&ri_buf, cnt) < 0) { + indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + free(table); free(rows); + Py_DECREF(left_arr); Py_DECREF(right_arr); + return NULL; + } + for (npy_intp j = 0; j < cnt; j++) { + li_buf.data[li_buf.len] = i; + ri_buf.data[ri_buf.len] = rows[off + j]; + li_buf.len++; + ri_buf.len++; } } else { - li[pos] = i; - ri[pos] = -1; - pos++; + if (indexbuf_ensure(&li_buf, 1) < 0 || + indexbuf_ensure(&ri_buf, 1) < 0) { + indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + free(table); free(rows); + Py_DECREF(left_arr); Py_DECREF(right_arr); + return NULL; + } + li_buf.data[li_buf.len] = i; + ri_buf.data[ri_buf.len] = -1; + li_buf.len++; + ri_buf.len++; + has_null = 1; } } - free(table); free(chain); + free(table); free(rows); Py_DECREF(left_arr); Py_DECREF(right_arr); + PyArrayObject *li_arr = indexbuf_to_array(&li_buf); + PyArrayObject *ri_arr = indexbuf_to_array(&ri_buf); + indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + if (!li_arr || !ri_arr) { + Py_XDECREF(li_arr); Py_XDECREF(ri_arr); + return NULL; + } + PyObject *result = PyTuple_Pack(3, (PyObject *)li_arr, (PyObject *)ri_arr, has_null ? Py_True : Py_False); @@ -892,7 +1006,7 @@ static PyMethodDef AccelMethods[] = { static struct PyModuleDef accelmodule = { PyModuleDef_HEAD_INIT, "_accel", - "C acceleration for tafra hot paths", + "C acceleration for tafra hot paths (v2: CSR joins, bitmask hashing)", -1, AccelMethods }; From f1e074b81e25e116399b7ea3bfe5099dbc7335ab Mon Sep 17 00:00:00 2001 From: David Fulford Date: Sat, 28 Mar 2026 19:32:46 -0500 Subject: [PATCH 02/10] =?UTF-8?q?perf:=20=5Faccel.c=20v3=20=E2=80=94=20com?= =?UTF-8?q?bine=20bitmask=20hashing=20with=20slot-caching=20joins?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merges best of v2 (power-of-2 bitmask probing, unrolled composite_key, CSR flat row buffer) with Derrick's slot-caching two-pass join approach (no re-probing, exact pre-allocation). Co-Authored-By: Claude Opus 4.6 (1M context) --- tafra/_accel.c | 217 +++++++++++++++++++------------------------------ 1 file changed, 85 insertions(+), 132 deletions(-) diff --git a/tafra/_accel.c b/tafra/_accel.c index 673b14e..2265c77 100644 --- a/tafra/_accel.c +++ b/tafra/_accel.c @@ -1,10 +1,10 @@ /* - * tafra/_accel.c — C extension for hot-path acceleration (v2). + * tafra/_accel.c — C extension for hot-path acceleration (v3). * * Performance improvements over v1: * - Power-of-2 hash tables with bitmask probing (no division) * - CSR (flat contiguous) layout for join hash tables (no linked lists) - * - Single-pass join index construction (no counting pre-pass) + * - Slot-caching two-pass joins (no re-probing, exact pre-allocation) * - Unrolled composite_key for 1-3 column fast paths * - Reduced redundant initialization * @@ -670,17 +670,18 @@ accel_group_indices(PyObject *self, PyObject *args) /* ================================================================ - * Hash join v2: CSR layout + single-pass index construction + * Hash join v3: CSR layout + slot-caching two-pass * - * Build phase: + * Build phase (right side): * 1. Hash right keys into power-of-2 table, count per bucket. * 2. Prefix sum -> offsets into flat contiguous row buffer. * 3. Scatter right row indices into flat buffer. * - * Probe phase reads sequential memory per bucket (cache-friendly). - * - * Output construction uses a growable buffer (realloc doubling) - * to avoid a separate counting pre-pass over left keys. + * Probe phase (left side): + * Pass 1: probe each left key, cache the slot index, sum counts + * for exact output pre-allocation. + * Pass 2: emit pairs using cached slots — no re-probing, + * sequential reads from rows[offset..offset+count]. * ================================================================ */ typedef struct { @@ -728,7 +729,7 @@ build_csr_hash_table(const npy_int64 *right, npy_intp right_n, npy_intp *rows = (npy_intp *)malloc(right_n * sizeof(npy_intp)); if (!rows) { free(table); PyErr_NoMemory(); return -1; } - /* Temporary write positions (reuse a stack of counts) */ + /* Temporary write positions */ npy_intp *write_pos = (npy_intp *)calloc(table_size, sizeof(npy_intp)); if (!write_pos) { free(table); free(rows); PyErr_NoMemory(); return -1; } @@ -750,58 +751,6 @@ build_csr_hash_table(const npy_int64 *right, npy_intp right_n, } -/* Growable index buffer for single-pass join output */ -typedef struct { - npy_intp *data; - npy_intp len; - npy_intp cap; -} IndexBuf; - -static int -indexbuf_init(IndexBuf *buf, npy_intp initial_cap) -{ - buf->data = (npy_intp *)malloc(initial_cap * sizeof(npy_intp)); - if (!buf->data) { PyErr_NoMemory(); return -1; } - buf->len = 0; - buf->cap = initial_cap; - return 0; -} - -static int -indexbuf_ensure(IndexBuf *buf, npy_intp additional) -{ - if (buf->len + additional <= buf->cap) return 0; - npy_intp new_cap = buf->cap; - while (new_cap < buf->len + additional) new_cap <<= 1; - npy_intp *tmp = (npy_intp *)realloc(buf->data, new_cap * sizeof(npy_intp)); - if (!tmp) { PyErr_NoMemory(); return -1; } - buf->data = tmp; - buf->cap = new_cap; - return 0; -} - -static void -indexbuf_free(IndexBuf *buf) -{ - free(buf->data); - buf->data = NULL; - buf->len = buf->cap = 0; -} - -/* Transfer ownership of buffer contents into a numpy array. - * On success, buf->data is set to NULL (numpy owns it via its own copy). - * Returns new reference or NULL on error. */ -static PyArrayObject * -indexbuf_to_array(IndexBuf *buf) -{ - npy_intp dims[1] = {buf->len}; - PyArrayObject *arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); - if (!arr) return NULL; - memcpy(PyArray_DATA(arr), buf->data, buf->len * sizeof(npy_intp)); - return arr; -} - - static PyObject * accel_inner_join(PyObject *self, PyObject *args) { @@ -827,52 +776,58 @@ accel_inner_join(PyObject *self, PyObject *args) return NULL; } - /* Single-pass probe with growable output */ - npy_intp init_cap = left_n; /* reasonable starting guess */ - if (init_cap < 64) init_cap = 64; - IndexBuf li_buf, ri_buf; - if (indexbuf_init(&li_buf, init_cap) < 0 || - indexbuf_init(&ri_buf, init_cap) < 0) { - indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + /* Probe pass 1: cache slot per left key, sum total output rows */ + npy_intp *slots = (npy_intp *)malloc(left_n * sizeof(npy_intp)); + if (!slots) { free(table); free(rows); Py_DECREF(left_arr); Py_DECREF(right_arr); - return NULL; + return PyErr_NoMemory(); } + npy_intp total = 0; for (npy_intp i = 0; i < left_n; i++) { npy_intp h = HASH_KEY(left[i], mask); while (table[h].count > 0 && table[h].key != left[i]) h = (h + 1) & mask; if (table[h].count > 0 && table[h].key == left[i]) { - npy_intp cnt = table[h].count; - npy_intp off = table[h].offset; - if (indexbuf_ensure(&li_buf, cnt) < 0 || - indexbuf_ensure(&ri_buf, cnt) < 0) { - indexbuf_free(&li_buf); indexbuf_free(&ri_buf); - free(table); free(rows); - Py_DECREF(left_arr); Py_DECREF(right_arr); - return NULL; - } - for (npy_intp j = 0; j < cnt; j++) { - li_buf.data[li_buf.len] = i; - ri_buf.data[ri_buf.len] = rows[off + j]; - li_buf.len++; - ri_buf.len++; - } + slots[i] = h; + total += table[h].count; + } else { + slots[i] = -1; } } - free(table); free(rows); - Py_DECREF(left_arr); Py_DECREF(right_arr); - - PyArrayObject *li_arr = indexbuf_to_array(&li_buf); - PyArrayObject *ri_arr = indexbuf_to_array(&ri_buf); - indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + npy_intp dims[1] = {total}; + PyArrayObject *li_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + PyArrayObject *ri_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); if (!li_arr || !ri_arr) { Py_XDECREF(li_arr); Py_XDECREF(ri_arr); + free(slots); free(table); free(rows); + Py_DECREF(left_arr); Py_DECREF(right_arr); return NULL; } + /* Emit pass 2: sequential reads from cached slots */ + npy_intp *li = (npy_intp *)PyArray_DATA(li_arr); + npy_intp *ri = (npy_intp *)PyArray_DATA(ri_arr); + npy_intp pos = 0; + + for (npy_intp i = 0; i < left_n; i++) { + npy_intp h = slots[i]; + if (h >= 0) { + npy_intp off = table[h].offset; + npy_intp cnt = table[h].count; + for (npy_intp j = 0; j < cnt; j++) { + li[pos] = i; + ri[pos] = rows[off + j]; + pos++; + } + } + } + + free(slots); free(table); free(rows); + Py_DECREF(left_arr); Py_DECREF(right_arr); + PyObject *result = PyTuple_Pack(2, (PyObject *)li_arr, (PyObject *)ri_arr); Py_DECREF(li_arr); Py_DECREF(ri_arr); return result; @@ -904,67 +859,65 @@ accel_left_join(PyObject *self, PyObject *args) return NULL; } - /* Single-pass probe with growable output */ - npy_intp init_cap = left_n; - if (init_cap < 64) init_cap = 64; - IndexBuf li_buf, ri_buf; - if (indexbuf_init(&li_buf, init_cap) < 0 || - indexbuf_init(&ri_buf, init_cap) < 0) { - indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + /* Probe pass 1: cache slot per left key, sum total output rows */ + npy_intp *slots = (npy_intp *)malloc(left_n * sizeof(npy_intp)); + if (!slots) { free(table); free(rows); Py_DECREF(left_arr); Py_DECREF(right_arr); - return NULL; + return PyErr_NoMemory(); } + npy_intp total = 0; int has_null = 0; - for (npy_intp i = 0; i < left_n; i++) { npy_intp h = HASH_KEY(left[i], mask); while (table[h].count > 0 && table[h].key != left[i]) h = (h + 1) & mask; if (table[h].count > 0 && table[h].key == left[i]) { - npy_intp cnt = table[h].count; - npy_intp off = table[h].offset; - if (indexbuf_ensure(&li_buf, cnt) < 0 || - indexbuf_ensure(&ri_buf, cnt) < 0) { - indexbuf_free(&li_buf); indexbuf_free(&ri_buf); - free(table); free(rows); - Py_DECREF(left_arr); Py_DECREF(right_arr); - return NULL; - } - for (npy_intp j = 0; j < cnt; j++) { - li_buf.data[li_buf.len] = i; - ri_buf.data[ri_buf.len] = rows[off + j]; - li_buf.len++; - ri_buf.len++; - } + slots[i] = h; + total += table[h].count; } else { - if (indexbuf_ensure(&li_buf, 1) < 0 || - indexbuf_ensure(&ri_buf, 1) < 0) { - indexbuf_free(&li_buf); indexbuf_free(&ri_buf); - free(table); free(rows); - Py_DECREF(left_arr); Py_DECREF(right_arr); - return NULL; - } - li_buf.data[li_buf.len] = i; - ri_buf.data[ri_buf.len] = -1; - li_buf.len++; - ri_buf.len++; + slots[i] = -1; + total++; has_null = 1; } } - free(table); free(rows); - Py_DECREF(left_arr); Py_DECREF(right_arr); - - PyArrayObject *li_arr = indexbuf_to_array(&li_buf); - PyArrayObject *ri_arr = indexbuf_to_array(&ri_buf); - indexbuf_free(&li_buf); indexbuf_free(&ri_buf); + npy_intp dims[1] = {total}; + PyArrayObject *li_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + PyArrayObject *ri_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); if (!li_arr || !ri_arr) { Py_XDECREF(li_arr); Py_XDECREF(ri_arr); + free(slots); free(table); free(rows); + Py_DECREF(left_arr); Py_DECREF(right_arr); return NULL; } + /* Emit pass 2: sequential reads from cached slots */ + npy_intp *li = (npy_intp *)PyArray_DATA(li_arr); + npy_intp *ri = (npy_intp *)PyArray_DATA(ri_arr); + npy_intp pos = 0; + + for (npy_intp i = 0; i < left_n; i++) { + npy_intp h = slots[i]; + if (h >= 0) { + npy_intp off = table[h].offset; + npy_intp cnt = table[h].count; + for (npy_intp j = 0; j < cnt; j++) { + li[pos] = i; + ri[pos] = rows[off + j]; + pos++; + } + } else { + li[pos] = i; + ri[pos] = -1; + pos++; + } + } + + free(slots); free(table); free(rows); + Py_DECREF(left_arr); Py_DECREF(right_arr); + PyObject *result = PyTuple_Pack(3, (PyObject *)li_arr, (PyObject *)ri_arr, has_null ? Py_True : Py_False); @@ -1006,7 +959,7 @@ static PyMethodDef AccelMethods[] = { static struct PyModuleDef accelmodule = { PyModuleDef_HEAD_INIT, "_accel", - "C acceleration for tafra hot paths (v2: CSR joins, bitmask hashing)", + "C acceleration for tafra hot paths (v3: bitmask hashing, CSR joins, slot caching)", -1, AccelMethods }; From 82bb428b06a2e31608e211e9893df338d8721d40 Mon Sep 17 00:00:00 2001 From: David Fulford Date: Sat, 28 Mar 2026 19:32:46 -0500 Subject: [PATCH 03/10] =?UTF-8?q?perf:=20=5Faccel.c=20v3=20=E2=80=94=20com?= =?UTF-8?q?bine=20bitmask=20hashing=20with=20slot-caching=20joins?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merges best of v2 (power-of-2 bitmask probing, unrolled composite_key, CSR flat row buffer) with Derrick's slot-caching two-pass join approach (no re-probing, exact pre-allocation). Co-Authored-By: Claude Opus 4.6 (1M context) --- tafra/_accel.c | 364 ++++++++++++++++++++++--------------------------- 1 file changed, 165 insertions(+), 199 deletions(-) diff --git a/tafra/_accel.c b/tafra/_accel.c index b8c40d7..2265c77 100644 --- a/tafra/_accel.c +++ b/tafra/_accel.c @@ -1,8 +1,12 @@ /* - * tafra/_accel.c — Minimal C extension for hot-path acceleration. + * tafra/_accel.c — C extension for hot-path acceleration (v3). * - * Provides single-pass grouped aggregation and hash-based join index - * construction, eliminating multiple numpy passes and temporary arrays. + * Performance improvements over v1: + * - Power-of-2 hash tables with bitmask probing (no division) + * - CSR (flat contiguous) layout for join hash tables (no linked lists) + * - Slot-caching two-pass joins (no re-probing, exact pre-allocation) + * - Unrolled composite_key for 1-3 column fast paths + * - Reduced redundant initialization * * All functions accept and return numpy arrays. Falls back gracefully * if this module fails to compile — pure-Python paths remain available. @@ -14,6 +18,7 @@ #define PY_SSIZE_T_CLEAN #include #include +#include #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION #include @@ -21,9 +26,21 @@ #include /* Golden ratio constant for multiplicative hashing (Knuth). - * phi = (sqrt(5)-1)/2 * 2^64 ≈ 0x9E3779B97F4A7C15 */ + * phi = (sqrt(5)-1)/2 * 2^64 ~ 0x9E3779B97F4A7C15 */ #define GOLDEN_HASH 0x9E3779B97F4A7C15ULL -#define HASH_KEY(k, sz) ((npy_intp)(((npy_uint64)(k) * GOLDEN_HASH) >> 1) % (sz)) + +/* Power-of-2 hash: multiply then mask. table_size must be power of 2. */ +#define HASH_KEY(k, mask) ((npy_intp)(((npy_uint64)(k) * GOLDEN_HASH) >> 1) & (mask)) + + +/* Round up to next power of 2 (minimum 16). */ +static npy_intp +next_pow2(npy_intp n) +{ + npy_intp p = 16; + while (p < n) p <<= 1; + return p; +} /* Helper: coerce to contiguous array of given type. Caller must Py_DECREF. */ @@ -35,55 +52,6 @@ as_contig(PyObject *obj, int typenum) } -/* ================================================================ - * Generic hash table probing - * - * Callbacks return: 1 = yes, 0 = no, -1 = error (for fallible comparisons). - * The probe function uses void* so a single implementation covers all - * entry types; per-type static callbacks cast back to the concrete struct. - * ================================================================ */ - -typedef int (*ht_occupied_fn)(const void *entry); -typedef int (*ht_equal_fn)(const void *entry, npy_int64 hash_key, void *ctx); - -/* - * Linear-probe to find either the slot matching hash_key or the first - * empty slot. On return: - * *found = 1 if an occupied matching slot was found, 0 if empty slot - * *err = 1 if keys_equal returned -1 (comparison error), else 0 - */ -static inline npy_intp -ht_probe(const void *table, npy_intp entry_size, npy_intp table_size, - npy_int64 hash_key, void *ctx, int *found, int *err, - ht_occupied_fn is_occupied, ht_equal_fn keys_equal) -{ - npy_intp slot = HASH_KEY(hash_key, table_size); - const char *base = (const char *)table; - *err = 0; - for (;;) { - const void *entry = base + slot * entry_size; - if (!is_occupied(entry)) { - *found = 0; - return slot; - } - int eq = keys_equal(entry, hash_key, ctx); - if (eq < 0) { *found = 0; *err = 1; return slot; } - if (eq) { *found = 1; return slot; } - slot = (slot + 1) % table_size; - } -} - -/* Convenience: alloc a zero-initialized hash table with 2x headroom. */ -static inline void * -ht_alloc(npy_intp n_items, npy_intp entry_size, npy_intp *out_table_size) -{ - npy_intp sz = n_items * 2; - if (sz < 16) sz = 16; - *out_table_size = sz; - return calloc(sz, entry_size); -} - - /* ================================================================ * GroupBy aggregations: single pass over (labels, data) * @@ -354,7 +322,7 @@ accel_groupby_max(PyObject *self, PyObject *args) * cardinalities: tuple of ints (max value + 1 for each column) * * Computes: key[i] = col0[i]*card1*card2*... + col1[i]*card2*... + ... - * Single pass, no temporaries. + * Unrolled fast paths for 1-3 columns (common case). * ================================================================ */ static PyObject * @@ -424,9 +392,9 @@ accel_composite_key(PyObject *self, PyObject *args) for (Py_ssize_t c = n_cols - 2; c >= 0; c--) mult[c] = mult[c + 1] * cards[c + 1]; - /* Build output in single pass */ + /* Build output */ npy_intp dims[1] = {n_rows}; - PyArrayObject *out = (PyArrayObject *)PyArray_ZEROS(1, dims, NPY_INT64, 0); + PyArrayObject *out = (PyArrayObject *)PyArray_EMPTY(1, dims, NPY_INT64, 0); if (!out) { for (Py_ssize_t c = 0; c < n_cols; c++) Py_DECREF(col_arrs[c]); free(col_arrs); free(cards); free(mult); @@ -434,27 +402,46 @@ accel_composite_key(PyObject *self, PyObject *args) } npy_int64 *result = (npy_int64 *)PyArray_DATA(out); - /* Get data pointers */ - const npy_int64 **col_data = (const npy_int64 **)malloc(n_cols * sizeof(npy_int64 *)); - if (!col_data) { - for (Py_ssize_t c = 0; c < n_cols; c++) Py_DECREF(col_arrs[c]); - free(col_arrs); free(cards); free(mult); - Py_DECREF(out); - return PyErr_NoMemory(); - } - for (Py_ssize_t c = 0; c < n_cols; c++) - col_data[c] = (const npy_int64 *)PyArray_DATA(col_arrs[c]); - - /* Single pass: result[i] = sum(col_data[c][i] * mult[c]) */ - for (npy_intp i = 0; i < n_rows; i++) { - npy_int64 key = 0; + /* Unrolled fast paths for common column counts */ + if (n_cols == 1) { + const npy_int64 *d0 = (const npy_int64 *)PyArray_DATA(col_arrs[0]); + memcpy(result, d0, n_rows * sizeof(npy_int64)); + } else if (n_cols == 2) { + const npy_int64 *d0 = (const npy_int64 *)PyArray_DATA(col_arrs[0]); + const npy_int64 *d1 = (const npy_int64 *)PyArray_DATA(col_arrs[1]); + npy_int64 m0 = mult[0]; + for (npy_intp i = 0; i < n_rows; i++) + result[i] = d0[i] * m0 + d1[i]; + } else if (n_cols == 3) { + const npy_int64 *d0 = (const npy_int64 *)PyArray_DATA(col_arrs[0]); + const npy_int64 *d1 = (const npy_int64 *)PyArray_DATA(col_arrs[1]); + const npy_int64 *d2 = (const npy_int64 *)PyArray_DATA(col_arrs[2]); + npy_int64 m0 = mult[0], m1 = mult[1]; + for (npy_intp i = 0; i < n_rows; i++) + result[i] = d0[i] * m0 + d1[i] * m1 + d2[i]; + } else { + /* General case */ + const npy_int64 **col_data = (const npy_int64 **)malloc(n_cols * sizeof(npy_int64 *)); + if (!col_data) { + for (Py_ssize_t c = 0; c < n_cols; c++) Py_DECREF(col_arrs[c]); + free(col_arrs); free(cards); free(mult); + Py_DECREF(out); + return PyErr_NoMemory(); + } for (Py_ssize_t c = 0; c < n_cols; c++) - key += col_data[c][i] * mult[c]; - result[i] = key; + col_data[c] = (const npy_int64 *)PyArray_DATA(col_arrs[c]); + + for (npy_intp i = 0; i < n_rows; i++) { + npy_int64 key = 0; + for (Py_ssize_t c = 0; c < n_cols; c++) + key += col_data[c][i] * mult[c]; + result[i] = key; + } + free((void *)col_data); } for (Py_ssize_t c = 0; c < n_cols; c++) Py_DECREF(col_arrs[c]); - free(col_arrs); free(cards); free(mult); free((void *)col_data); + free(col_arrs); free(cards); free(mult); return (PyObject *)out; } @@ -469,27 +456,16 @@ accel_composite_key(PyObject *self, PyObject *args) * codes: int64 array of integer codes (0..n_unique-1), first-seen order * n_unique: int * - * Replaces np.unique(return_inverse=True) which is O(n log n). + * Power-of-2 table with bitmask probing. * ================================================================ */ typedef struct { Py_hash_t hash; - PyObject *obj; /* borrowed ref; NULL = unoccupied */ + PyObject *obj; /* borrowed reference to the original object */ npy_intp code; + int occupied; } StrHashEntry; -static int str_occupied(const void *entry) - { return ((const StrHashEntry *)entry)->obj != NULL; } - -/* ctx carries the PyObject* being looked up. Returns -1/0/1 - * matching PyObject_RichCompareBool's convention. */ -static int str_equal(const void *entry, npy_int64 hash_key, void *ctx) -{ - const StrHashEntry *e = (const StrHashEntry *)entry; - if (e->hash != (Py_hash_t)hash_key) return 0; - return PyObject_RichCompareBool(e->obj, (PyObject *)ctx, Py_EQ); -} - static PyObject * accel_encode_strings(PyObject *self, PyObject *args) { @@ -504,10 +480,10 @@ accel_encode_strings(PyObject *self, PyObject *args) npy_intp n = PyArray_SIZE(arr); PyObject **data = (PyObject **)PyArray_DATA(arr); - /* Hash table */ - npy_intp table_size; - StrHashEntry *table = (StrHashEntry *)ht_alloc(n, sizeof(StrHashEntry), - &table_size); + /* Power-of-2 hash table at ~50% load factor */ + npy_intp table_size = next_pow2(n * 2); + npy_intp mask = table_size - 1; + StrHashEntry *table = (StrHashEntry *)calloc(table_size, sizeof(StrHashEntry)); if (!table) { Py_DECREF(arr); return PyErr_NoMemory(); } /* Output codes */ @@ -526,19 +502,25 @@ accel_encode_strings(PyObject *self, PyObject *args) return NULL; } - int found, err; - npy_intp slot = ht_probe(table, sizeof(StrHashEntry), table_size, - (npy_int64)h, obj, &found, &err, - str_occupied, str_equal); - if (err) { - free(table); Py_DECREF(arr); Py_DECREF(codes_out); - return NULL; + npy_intp slot = (npy_intp)(((npy_uint64)h * GOLDEN_HASH) >> 1) & mask; + + while (table[slot].occupied) { + if (table[slot].hash == h) { + int eq = PyObject_RichCompareBool(table[slot].obj, obj, Py_EQ); + if (eq < 0) { + free(table); Py_DECREF(arr); Py_DECREF(codes_out); + return NULL; + } + if (eq) break; + } + slot = (slot + 1) & mask; } - if (!found) { + if (!table[slot].occupied) { table[slot].hash = h; - table[slot].obj = obj; /* borrowed ref, arr keeps it alive */ + table[slot].obj = obj; table[slot].code = n_unique; + table[slot].occupied = 1; n_unique++; } codes[i] = table[slot].code; @@ -559,25 +541,15 @@ accel_encode_strings(PyObject *self, PyObject *args) * group_indices(key_array) -> (first_seen_indices, group_row_lists, n_groups) * key_array: int64 array (composite or single-column key) * - * Returns: - * first_seen_indices: int64 array of length n_groups - * group_row_lists: Python list of int64 arrays (row indices per group) - * n_groups: int - * - * Replaces np.unique + argsort + split with a single O(n) pass. + * Power-of-2 table with bitmask probing. * ================================================================ */ typedef struct { npy_int64 key; - npy_intp label; /* assigned group label */ + npy_intp label; int occupied; } GroupHashEntry; -static int group_occupied(const void *entry) - { return ((const GroupHashEntry *)entry)->occupied; } -static int group_equal(const void *entry, npy_int64 hash_key, void *ctx) - { (void)ctx; return ((const GroupHashEntry *)entry)->key == hash_key; } - static PyObject * accel_group_indices(PyObject *self, PyObject *args) { @@ -591,10 +563,10 @@ accel_group_indices(PyObject *self, PyObject *args) npy_intp n = PyArray_SIZE(key_arr); const npy_int64 *keys = (const npy_int64 *)PyArray_DATA(key_arr); - /* Allocate hash table */ - npy_intp table_size; - GroupHashEntry *table = (GroupHashEntry *)ht_alloc( - n, sizeof(GroupHashEntry), &table_size); + /* Power-of-2 hash table */ + npy_intp table_size = next_pow2(n * 2); + npy_intp mask = table_size - 1; + GroupHashEntry *table = (GroupHashEntry *)calloc(table_size, sizeof(GroupHashEntry)); if (!table) { Py_DECREF(key_arr); return PyErr_NoMemory(); } /* Pass 1: assign labels, count per group */ @@ -609,14 +581,13 @@ accel_group_indices(PyObject *self, PyObject *args) } npy_intp n_groups = 0; - int found, err; /* err always 0 — group callbacks are infallible */ for (npy_intp i = 0; i < n; i++) { npy_int64 k = keys[i]; - npy_intp h = ht_probe(table, sizeof(GroupHashEntry), table_size, - k, NULL, &found, &err, - group_occupied, group_equal); + npy_intp h = HASH_KEY(k, mask); + while (table[h].occupied && table[h].key != k) + h = (h + 1) & mask; - if (!found) { + if (!table[h].occupied) { table[h].key = k; table[h].occupied = 1; table[h].label = n_groups; @@ -671,7 +642,7 @@ accel_group_indices(PyObject *self, PyObject *args) return NULL; } group_data[g] = (npy_intp *)PyArray_DATA(arr); - PyList_SET_ITEM(group_list, g, (PyObject *)arr); /* steals ref */ + PyList_SET_ITEM(group_list, g, (PyObject *)arr); } free(counts); @@ -688,9 +659,7 @@ accel_group_indices(PyObject *self, PyObject *args) npy_intp fs_dims[1] = {n_groups}; PyArrayObject *fs_out = (PyArrayObject *)PyArray_EMPTY(1, fs_dims, NPY_INTP, 0); if (!fs_out) { free(first_seen); Py_DECREF(group_list); return NULL; } - npy_intp *fs_data = (npy_intp *)PyArray_DATA(fs_out); - for (npy_intp g = 0; g < n_groups; g++) - fs_data[g] = first_seen[g]; + memcpy(PyArray_DATA(fs_out), first_seen, n_groups * sizeof(npy_intp)); free(first_seen); PyObject *result = Py_BuildValue("(OOn)", fs_out, group_list, (Py_ssize_t)n_groups); @@ -701,83 +670,82 @@ accel_group_indices(PyObject *self, PyObject *args) /* ================================================================ - * Hash join: build (left_indices, right_indices) for equi-join + * Hash join v3: CSR layout + slot-caching two-pass * - * CSR layout: two-pass build produces contiguous per-key storage. - * Pass 1 counts right-side rows per key, prefix sum computes offsets, - * Pass 2 scatters row indices into a flat rows[] array. - * Emit phase does sequential reads — no pointer chasing. + * Build phase (right side): + * 1. Hash right keys into power-of-2 table, count per bucket. + * 2. Prefix sum -> offsets into flat contiguous row buffer. + * 3. Scatter right row indices into flat buffer. * - * Input arrays are coerced to contiguous int64. + * Probe phase (left side): + * Pass 1: probe each left key, cache the slot index, sum counts + * for exact output pre-allocation. + * Pass 2: emit pairs using cached slots — no re-probing, + * sequential reads from rows[offset..offset+count]. * ================================================================ */ typedef struct { npy_int64 key; - npy_intp count; /* number of right-side rows for this key */ - npy_intp offset; /* index into flat rows[] array */ -} JoinHashEntry; - -static int join_occupied(const void *entry) - { return ((const JoinHashEntry *)entry)->count > 0; } -static int join_equal(const void *entry, npy_int64 hash_key, void *ctx) - { (void)ctx; return ((const JoinHashEntry *)entry)->key == hash_key; } + npy_intp offset; /* start index into flat row buffer */ + npy_intp count; /* number of right rows with this key */ +} CSRHashEntry; -/* Build CSR hash table on right-side keys. Returns 0 on success, -1 on error. - * Caller must free(*out_table) and free(*out_rows). */ +/* Build CSR hash table on right-side keys. + * Returns 0 on success, -1 on error. + * Caller must free *out_table and *out_rows. */ static int -build_join_table(const npy_int64 *right, npy_intp right_n, - JoinHashEntry **out_table, npy_intp *out_table_size, - npy_intp **out_rows) +build_csr_hash_table(const npy_int64 *right, npy_intp right_n, + CSRHashEntry **out_table, npy_intp *out_mask, + npy_intp **out_rows) { - npy_intp table_size; - JoinHashEntry *table = (JoinHashEntry *)ht_alloc( - right_n, sizeof(JoinHashEntry), &table_size); + npy_intp table_size = next_pow2(right_n * 2); + npy_intp mask = table_size - 1; + + /* calloc zeros count and offset to 0; count==0 means empty slot */ + CSRHashEntry *table = (CSRHashEntry *)calloc(table_size, sizeof(CSRHashEntry)); if (!table) { PyErr_NoMemory(); return -1; } - /* Pass 1: probe + count */ - int found, err; + /* Pass 1: count occurrences per key */ for (npy_intp i = 0; i < right_n; i++) { - npy_intp h = ht_probe(table, sizeof(JoinHashEntry), table_size, - right[i], NULL, &found, &err, - join_occupied, join_equal); - if (!found) - table[h].key = right[i]; + npy_int64 k = right[i]; + npy_intp h = HASH_KEY(k, mask); + while (table[h].count > 0 && table[h].key != k) + h = (h + 1) & mask; + table[h].key = k; table[h].count++; } - /* Exclusive prefix sum over occupied slots to fill offset */ + /* Pass 2: prefix sum to compute offsets */ npy_intp running = 0; - for (npy_intp i = 0; i < table_size; i++) { - if (table[i].count > 0) { - table[i].offset = running; - running += table[i].count; + for (npy_intp h = 0; h < table_size; h++) { + if (table[h].count > 0) { + table[h].offset = running; + running += table[h].count; } } - /* Allocate flat rows array */ + /* Allocate flat row buffer */ npy_intp *rows = (npy_intp *)malloc(right_n * sizeof(npy_intp)); - /* Temporary write positions (copy of offsets, bumped during scatter) */ - npy_intp *wpos = (npy_intp *)malloc(table_size * sizeof(npy_intp)); - if (!rows || !wpos) { - free(table); free(rows); free(wpos); - PyErr_NoMemory(); - return -1; - } - for (npy_intp i = 0; i < table_size; i++) - wpos[i] = table[i].offset; + if (!rows) { free(table); PyErr_NoMemory(); return -1; } + + /* Temporary write positions */ + npy_intp *write_pos = (npy_intp *)calloc(table_size, sizeof(npy_intp)); + if (!write_pos) { free(table); free(rows); PyErr_NoMemory(); return -1; } - /* Pass 2: probe + scatter row indices */ + /* Pass 3: scatter right row indices into flat buffer */ for (npy_intp i = 0; i < right_n; i++) { - npy_intp h = ht_probe(table, sizeof(JoinHashEntry), table_size, - right[i], NULL, &found, &err, - join_occupied, join_equal); - rows[wpos[h]++] = i; + npy_int64 k = right[i]; + npy_intp h = HASH_KEY(k, mask); + while (table[h].count > 0 && table[h].key != k) + h = (h + 1) & mask; + rows[table[h].offset + write_pos[h]] = i; + write_pos[h]++; } - free(wpos); + free(write_pos); *out_table = table; - *out_table_size = table_size; + *out_mask = mask; *out_rows = rows; return 0; } @@ -800,15 +768,15 @@ accel_inner_join(PyObject *self, PyObject *args) npy_intp left_n = PyArray_SIZE(left_arr); npy_intp right_n = PyArray_SIZE(right_arr); const npy_int64 *left = (const npy_int64 *)PyArray_DATA(left_arr); - const npy_int64 *right_data = (const npy_int64 *)PyArray_DATA(right_arr); + const npy_int64 *right = (const npy_int64 *)PyArray_DATA(right_arr); - JoinHashEntry *table; npy_intp table_size; npy_intp *rows; - if (build_join_table(right_data, right_n, &table, &table_size, &rows) < 0) { + CSRHashEntry *table; npy_intp mask; npy_intp *rows; + if (build_csr_hash_table(right, right_n, &table, &mask, &rows) < 0) { Py_DECREF(left_arr); Py_DECREF(right_arr); return NULL; } - /* Count pass: probe each left key, cache slot, sum counts */ + /* Probe pass 1: cache slot per left key, sum total output rows */ npy_intp *slots = (npy_intp *)malloc(left_n * sizeof(npy_intp)); if (!slots) { free(table); free(rows); @@ -816,13 +784,12 @@ accel_inner_join(PyObject *self, PyObject *args) return PyErr_NoMemory(); } - int found, err; npy_intp total = 0; for (npy_intp i = 0; i < left_n; i++) { - npy_intp h = ht_probe(table, sizeof(JoinHashEntry), table_size, - left[i], NULL, &found, &err, - join_occupied, join_equal); - if (found) { + npy_intp h = HASH_KEY(left[i], mask); + while (table[h].count > 0 && table[h].key != left[i]) + h = (h + 1) & mask; + if (table[h].count > 0 && table[h].key == left[i]) { slots[i] = h; total += table[h].count; } else { @@ -840,7 +807,7 @@ accel_inner_join(PyObject *self, PyObject *args) return NULL; } - /* Emit pass: sequential scan of rows[offset..offset+count] */ + /* Emit pass 2: sequential reads from cached slots */ npy_intp *li = (npy_intp *)PyArray_DATA(li_arr); npy_intp *ri = (npy_intp *)PyArray_DATA(ri_arr); npy_intp pos = 0; @@ -884,15 +851,15 @@ accel_left_join(PyObject *self, PyObject *args) npy_intp left_n = PyArray_SIZE(left_arr); npy_intp right_n = PyArray_SIZE(right_arr); const npy_int64 *left = (const npy_int64 *)PyArray_DATA(left_arr); - const npy_int64 *right_data = (const npy_int64 *)PyArray_DATA(right_arr); + const npy_int64 *right = (const npy_int64 *)PyArray_DATA(right_arr); - JoinHashEntry *table; npy_intp table_size; npy_intp *rows; - if (build_join_table(right_data, right_n, &table, &table_size, &rows) < 0) { + CSRHashEntry *table; npy_intp mask; npy_intp *rows; + if (build_csr_hash_table(right, right_n, &table, &mask, &rows) < 0) { Py_DECREF(left_arr); Py_DECREF(right_arr); return NULL; } - /* Count pass: probe each left key, cache slot, sum counts */ + /* Probe pass 1: cache slot per left key, sum total output rows */ npy_intp *slots = (npy_intp *)malloc(left_n * sizeof(npy_intp)); if (!slots) { free(table); free(rows); @@ -900,14 +867,13 @@ accel_left_join(PyObject *self, PyObject *args) return PyErr_NoMemory(); } - int found, err; npy_intp total = 0; int has_null = 0; for (npy_intp i = 0; i < left_n; i++) { - npy_intp h = ht_probe(table, sizeof(JoinHashEntry), table_size, - left[i], NULL, &found, &err, - join_occupied, join_equal); - if (found) { + npy_intp h = HASH_KEY(left[i], mask); + while (table[h].count > 0 && table[h].key != left[i]) + h = (h + 1) & mask; + if (table[h].count > 0 && table[h].key == left[i]) { slots[i] = h; total += table[h].count; } else { @@ -927,7 +893,7 @@ accel_left_join(PyObject *self, PyObject *args) return NULL; } - /* Emit pass: sequential scan of rows[offset..offset+count] */ + /* Emit pass 2: sequential reads from cached slots */ npy_intp *li = (npy_intp *)PyArray_DATA(li_arr); npy_intp *ri = (npy_intp *)PyArray_DATA(ri_arr); npy_intp pos = 0; @@ -987,13 +953,13 @@ static PyMethodDef AccelMethods[] = { "Hash inner join: inner_join(left_key, right_key) -> (li, ri)"}, {"left_join", accel_left_join, METH_VARARGS, "Hash left join: left_join(left_key, right_key) -> (li, ri, has_null)"}, -{NULL, NULL, 0, NULL} + {NULL, NULL, 0, NULL} }; static struct PyModuleDef accelmodule = { PyModuleDef_HEAD_INIT, "_accel", - "C acceleration for tafra hot paths", + "C acceleration for tafra hot paths (v3: bitmask hashing, CSR joins, slot caching)", -1, AccelMethods }; From e03f94f9f13a6be829856651ef222fd8a518eccd Mon Sep 17 00:00:00 2001 From: David Fulford Date: Thu, 2 Apr 2026 17:16:58 -0500 Subject: [PATCH 04/10] fix: composite key collision in multi-column joins with asymmetric cardinalities _build_composite_key computed per-column cardinalities independently on each call. When left and right sides had different unique value counts, the positional encoding multipliers diverged, mapping distinct key tuples to the same composite integer. Co-Authored-By: Claude Opus 4.6 (1M context) --- tafra/group.py | 44 +++++++++++++++++++++++++++++++++++++------- 1 file changed, 37 insertions(+), 7 deletions(-) diff --git a/tafra/group.py b/tafra/group.py index 207f7c5..2efb72a 100644 --- a/tafra/group.py +++ b/tafra/group.py @@ -540,11 +540,19 @@ def _encode_columns_paired( return left_enc, right_enc @staticmethod - def _build_composite_key(encoded: list[np.ndarray[Any, Any]]) -> np.ndarray[Any, Any]: + def _build_composite_key( + encoded: list[np.ndarray[Any, Any]], + cards: list[int] | None = None, + ) -> np.ndarray[Any, Any]: """ Combine multiple integer-coded columns into a single flat key. Uses positional encoding: key = c0 * N1*N2*... + c1 * N2*... + c2 * ... Raises ValueError if the product of cardinalities would overflow int64. + + If *cards* is provided it is used directly; otherwise cardinalities are + inferred from the per-column max. For joins, callers MUST pass shared + cardinalities (max over both sides) so left and right keys use the same + positional encoding. """ if len(encoded) == 1: return encoded[0] @@ -557,7 +565,8 @@ def _build_composite_key(encoded: list[np.ndarray[Any, Any]]) -> np.ndarray[Any, # Callers _encode_columns (GroupBy/Transform) and _encode_columns_paired # (joins) ensure this by encoding strings via np.unique and shifting # negative integers to non-negative range. - cards = [int(c.max()) + 1 for c in encoded] + if cards is None: + cards = [int(c.max()) + 1 for c in encoded] # check for overflow: product of all cardinalities must fit int64 product = 1 @@ -1066,7 +1075,10 @@ def _validate_ops(ops: Iterable[str]) -> None: raise TypeError(f"The operator {op} is not valid.") @staticmethod - def _build_composite_key(cols: list[np.ndarray[Any, Any]]) -> np.ndarray[Any, Any]: + def _build_composite_key( + cols: list[np.ndarray[Any, Any]], + cards: list[int] | None = None, + ) -> np.ndarray[Any, Any]: """Build a single sortable key array from multiple columns.""" if len(cols) == 1: return cols[0] @@ -1253,8 +1265,17 @@ def apply(self, left_t: "Tafra", right_t: "Tafra") -> "Tafra": right_cols_filt = right_cols_data l_enc, r_enc = GroupSet._encode_columns_paired(left_cols_filt, right_cols_filt) - left_key = GroupSet._build_composite_key(l_enc) - right_key = GroupSet._build_composite_key(r_enc) + # Shared cardinalities: max over both sides so positional encoding + # is consistent between left and right keys. + if len(l_enc) > 1: + cards: list[int] | None = [ + max(int(lc.max()), int(rc.max())) + 1 + for lc, rc in zip(l_enc, r_enc) + ] + else: + cards = None + left_key = GroupSet._build_composite_key(l_enc, cards) + right_key = GroupSet._build_composite_key(r_enc, cards) if _HAS_ACCEL: li, ri = _c_inner_join( @@ -1526,8 +1547,17 @@ def apply(self, left_t: "Tafra", right_t: "Tafra") -> "Tafra": left_null_idx = np.array([], dtype=np.intp) l_enc, r_enc = GroupSet._encode_columns_paired(left_cols_filt, right_cols_filt) - left_key = GroupSet._build_composite_key(l_enc) - right_key = GroupSet._build_composite_key(r_enc) + # Shared cardinalities: max over both sides so positional encoding + # is consistent between left and right keys. + if len(l_enc) > 1: + cards: list[int] | None = [ + max(int(lc.max()), int(rc.max())) + 1 + for lc, rc in zip(l_enc, r_enc) + ] + else: + cards = None + left_key = GroupSet._build_composite_key(l_enc, cards) + right_key = GroupSet._build_composite_key(r_enc, cards) if _HAS_ACCEL: li, ri, has_null = _c_left_join( From ddd42cce9bb15443e4015f8ca6f76eef0f0dc9fb Mon Sep 17 00:00:00 2001 From: David Fulford Date: Thu, 2 Apr 2026 17:22:45 -0500 Subject: [PATCH 05/10] test: add regression tests for asymmetric-cardinality composite key joins Also guard cards computation against empty encoded arrays (e.g. when all rows on one side are null-filtered). Co-Authored-By: Claude Opus 4.6 (1M context) --- tafra/group.py | 4 ++-- test/test_tafra.py | 57 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tafra/group.py b/tafra/group.py index 2efb72a..cc938bf 100644 --- a/tafra/group.py +++ b/tafra/group.py @@ -1267,7 +1267,7 @@ def apply(self, left_t: "Tafra", right_t: "Tafra") -> "Tafra": l_enc, r_enc = GroupSet._encode_columns_paired(left_cols_filt, right_cols_filt) # Shared cardinalities: max over both sides so positional encoding # is consistent between left and right keys. - if len(l_enc) > 1: + if len(l_enc) > 1 and len(l_enc[0]) > 0 and len(r_enc[0]) > 0: cards: list[int] | None = [ max(int(lc.max()), int(rc.max())) + 1 for lc, rc in zip(l_enc, r_enc) @@ -1549,7 +1549,7 @@ def apply(self, left_t: "Tafra", right_t: "Tafra") -> "Tafra": l_enc, r_enc = GroupSet._encode_columns_paired(left_cols_filt, right_cols_filt) # Shared cardinalities: max over both sides so positional encoding # is consistent between left and right keys. - if len(l_enc) > 1: + if len(l_enc) > 1 and len(l_enc[0]) > 0 and len(r_enc[0]) > 0: cards: list[int] | None = [ max(int(lc.max()), int(rc.max())) + 1 for lc, rc in zip(l_enc, r_enc) diff --git a/test/test_tafra.py b/test/test_tafra.py index c186ee8..de0ab39 100644 --- a/test/test_tafra.py +++ b/test/test_tafra.py @@ -2851,6 +2851,35 @@ def test_inner_join_c_vs_python_same_results(self) -> None: assert _tafra_to_row_set(t_c) == _tafra_to_row_set(t_py) + def test_inner_join_asymmetric_cardinality(self) -> None: + """Multi-column join where left has fewer unique values than right. + + Regression test: _build_composite_key must use shared cardinalities + so that positional encoding is consistent across both sides. + """ + right = Tafra( + { + "a": np.array(["X", "X", "Y", "Y"], dtype=object), + "b": np.array(["P", "Q", "P", "Q"], dtype=object), + "val": np.array([10, 20, 30, 40]), + } + ) + # Left has only one unique value per column — different max codes + left = Tafra( + { + "a": np.array(["Y"], dtype=object), + "b": np.array(["Q"], dtype=object), + "id": np.array([0]), + } + ) + t = left.inner_join( + right, + [("a", "a", "=="), ("b", "b", "==")], + select=["a", "b", "id", "val"], + ) + assert len(t) == 1 + assert t["val"][0] == 40 + class TestLeftJoin: def test_left_join_equi(self) -> None: @@ -3456,6 +3485,34 @@ def test_left_join_c_vs_python_same_results(self) -> None: py_rows = sorted([(t_py["k"][i], t_py["lv"][i]) for i in range(len(t_py))]) assert c_rows == py_rows + def test_left_join_asymmetric_cardinality(self) -> None: + """Multi-column left join where left has fewer unique values than right. + + Regression test: _build_composite_key must use shared cardinalities + so that positional encoding is consistent across both sides. + """ + right = Tafra( + { + "a": np.array(["X", "X", "Y", "Y"], dtype=object), + "b": np.array(["P", "Q", "P", "Q"], dtype=object), + "val": np.array([10, 20, 30, 40]), + } + ) + left = Tafra( + { + "a": np.array(["Y"], dtype=object), + "b": np.array(["Q"], dtype=object), + "id": np.array([0]), + } + ) + t = left.left_join( + right, + [("a", "a", "=="), ("b", "b", "==")], + select=["a", "b", "id", "val"], + ) + assert len(t) == 1 + assert t["val"][0] == 40 + class TestCrossJoin: def test_cross_join(self) -> None: From 2a2ea8a8aac032be694482792437da041f120c6a Mon Sep 17 00:00:00 2001 From: David Fulford Date: Thu, 2 Apr 2026 17:26:17 -0500 Subject: [PATCH 06/10] chore: bump version to 2.2.2 Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/changelog.md | 8 ++++++++ pyproject.toml | 2 +- recipe/meta.yaml | 2 +- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index cc4973c..872de3c 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -1,5 +1,13 @@ # Version History +## 2.2.2 + +* **Fix**: Multi-column join composite key collision when left and right sides + have different numbers of unique values. `_build_composite_key` now uses + shared cardinalities (max over both sides) so positional encoding is + consistent between left and right keys. +* **Perf**: `_accel.c` v3 — combine bitmask hashing with slot-caching joins. + ## 2.2.1 * **Fix**: Join and Union dtype validation now compares `_dtypes` metadata diff --git a/pyproject.toml b/pyproject.toml index ef7241d..418b939 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "tafra" -version = "2.2.1" +version = "2.2.2" description = "Tafra: essence of a dataframe" readme = "README.md" license = "MIT" diff --git a/recipe/meta.yaml b/recipe/meta.yaml index 1a31bff..6181b5c 100644 --- a/recipe/meta.yaml +++ b/recipe/meta.yaml @@ -1,5 +1,5 @@ {% set name = "tafra" %} -{% set version = "2.2.1" %} +{% set version = "2.2.2" %} package: From 4d157a26b09393e2e0980d6a5df333a6c6879ec8 Mon Sep 17 00:00:00 2001 From: David Fulford Date: Thu, 2 Apr 2026 17:32:12 -0500 Subject: [PATCH 07/10] fix: preserve object dtype for string arrays in constructor ObjectFormatter.parse_dtype unconditionally converted object arrays of strings to StringDType, breaking union() when one side came from an external source that preserves object dtype. String conversion is now opt-in via convert_strings parameter; parse_object_dtypes_inplace() opts in explicitly. Co-Authored-By: Claude Opus 4.6 (1M context) --- CLAUDE.md | 2 +- docs/changelog.md | 4 ++++ tafra/base.py | 2 +- tafra/formatter.py | 14 +++++++++++--- test/test_tafra.py | 10 +++++++--- 5 files changed, 24 insertions(+), 8 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index ec73ccf..507af7a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -67,7 +67,7 @@ The library has one core abstraction and a set of aggregation/partitioning opera **Supporting modules:** - `protocol.py` — Typing protocols for duck-typing compatibility (Series, DataFrame, Cursor) -- `formatter.py` — `ObjectFormatter` for custom dtype parsing (e.g., Decimal → float); auto-converts object arrays of strings to `StringDType(na_object=None)` for nullable string support +- `formatter.py` — `ObjectFormatter` for custom dtype parsing (e.g., Decimal → float); string conversion from object to `StringDType(na_object=None)` is opt-in via `parse_object_dtypes_inplace()` - `csvreader.py` — CSV reader with type inference; string columns produce `StringDType(na_object=None)` for nullable string support ## Testing diff --git a/docs/changelog.md b/docs/changelog.md index 872de3c..bc171e0 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -6,6 +6,10 @@ have different numbers of unique values. `_build_composite_key` now uses shared cardinalities (max over both sides) so positional encoding is consistent between left and right keys. +* **Fix**: Constructor no longer auto-converts object arrays of strings to + `StringDType`. This preserves the caller's dtype and prevents `union()` + mismatches when one Tafra comes from an external source with object dtype. + Use `parse_object_dtypes_inplace()` for explicit conversion. * **Perf**: `_accel.c` v3 — combine bitmask hashing with slot-caching joins. ## 2.2.1 diff --git a/tafra/base.py b/tafra/base.py index e1a2af5..9dbab82 100644 --- a/tafra/base.py +++ b/tafra/base.py @@ -836,7 +836,7 @@ def parse_object_dtypes_inplace(self) -> None: Parse the object dtypes using the `ObjectFormatter` instance. """ for column, value in self._data.items(): - parsed_value = object_formatter.parse_dtype(value) + parsed_value = object_formatter.parse_dtype(value, convert_strings=True) if parsed_value is not None: self._data[column] = parsed_value self._dtypes[column] = self._format_dtype(parsed_value.dtype) diff --git a/tafra/formatter.py b/tafra/formatter.py index 89c8047..0fd3776 100644 --- a/tafra/formatter.py +++ b/tafra/formatter.py @@ -90,7 +90,11 @@ def __len__(self) -> int: def copy(self) -> dict[str, Any]: return {k: dict.__getitem__(self, k) for k in self} - def parse_dtype(self, value: np.ndarray[Any, Any]) -> np.ndarray[Any, Any] | None: + def parse_dtype( + self, + value: np.ndarray[Any, Any], + convert_strings: bool = False, + ) -> np.ndarray[Any, Any] | None: """ Parse an object dtype. @@ -98,6 +102,10 @@ def parse_dtype(self, value: np.ndarray[Any, Any]) -> np.ndarray[Any, Any] | Non ---------- value: np.ndarray The `np.ndarray` to be parsed. + convert_strings: bool + If True, convert object arrays of strings to + ``StringDType(na_object=None)``. Default False — callers must + opt in (e.g. ``parse_object_dtypes_inplace``). Returns ------- @@ -112,8 +120,8 @@ def parse_dtype(self, value: np.ndarray[Any, Any]) -> np.ndarray[Any, Any] | Non value = self[type_name](value) return value - # convert object arrays of strings to StringDType - if type_name == "str": + # convert object arrays of strings to StringDType (opt-in only) + if convert_strings and type_name == "str": return value.astype(np.dtypes.StringDType(na_object=None)) # type: ignore[call-arg] return None diff --git a/test/test_tafra.py b/test/test_tafra.py index de0ab39..8615084 100644 --- a/test/test_tafra.py +++ b/test/test_tafra.py @@ -2122,13 +2122,17 @@ def test_string_column_uses_stringdtype(self) -> None: t2["label"] = "constant" assert t2["label"].dtype.kind == "T" - def test_object_string_array_converted_to_stringdtype( + def test_object_string_array_preserves_dtype( self, ) -> None: - """Object arrays of strings should be auto-converted to - StringDType.""" + """Object arrays of strings should preserve object dtype. + Use parse_object_dtypes_inplace() for explicit StringDType conversion.""" obj_arr = np.array(["a", "b", "c"], dtype=object) t = Tafra({"s": obj_arr}) + assert t["s"].dtype.kind == "O" + + # Explicit conversion should still work + t.parse_object_dtypes_inplace() assert t["s"].dtype.kind == "T" def test_drop_duplicates_string(self) -> None: From 0a6379bc988b753902c1209c73030e87632f9c7d Mon Sep 17 00:00:00 2001 From: David Fulford Date: Thu, 2 Apr 2026 17:38:45 -0500 Subject: [PATCH 08/10] fix: guard parse_dtype against empty object arrays, fix stale docstring Co-Authored-By: Claude Opus 4.6 (1M context) --- tafra/formatter.py | 6 +++--- test/test_tafra.py | 8 ++++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tafra/formatter.py b/tafra/formatter.py index 0fd3776..99748d1 100644 --- a/tafra/formatter.py +++ b/tafra/formatter.py @@ -109,10 +109,10 @@ def parse_dtype( Returns ------- - value, modified: Tuple(np.ndarray, bool) - The `np.ndarray` and whether it was modified or not. + np.ndarray or None + The converted array, or None if no conversion was applied. """ - if value.dtype.kind != "O": + if value.dtype.kind != "O" or len(value) == 0: return None type_name = type(value[0]).__name__ diff --git a/test/test_tafra.py b/test/test_tafra.py index 8615084..e68150e 100644 --- a/test/test_tafra.py +++ b/test/test_tafra.py @@ -2135,6 +2135,14 @@ def test_object_string_array_preserves_dtype( t.parse_object_dtypes_inplace() assert t["s"].dtype.kind == "T" + def test_empty_object_array_no_crash(self) -> None: + """Empty object arrays must not crash parse_dtype (value[0] guard).""" + t = Tafra({"s": np.array([], dtype=object), "v": np.array([], dtype=float)}) + assert t["s"].dtype.kind == "O" + # parse_object_dtypes_inplace should also be safe + t.parse_object_dtypes_inplace() + assert t["s"].dtype.kind == "O" + def test_drop_duplicates_string(self) -> None: """drop_duplicates works with StringDType columns.""" t = Tafra( From 941bcec56200d27210e53b58eacf50d112911e13 Mon Sep 17 00:00:00 2001 From: David Fulford Date: Thu, 2 Apr 2026 17:59:41 -0500 Subject: [PATCH 09/10] fix: guard _accel.c join functions against empty input arrays malloc(0) can return NULL per C standard, causing spurious MemoryError. Add early returns for zero-length inputs in inner_join and left_join. Co-Authored-By: Claude Opus 4.6 (1M context) --- tafra/_accel.c | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/tafra/_accel.c b/tafra/_accel.c index 2265c77..901c776 100644 --- a/tafra/_accel.c +++ b/tafra/_accel.c @@ -767,6 +767,22 @@ accel_inner_join(PyObject *self, PyObject *args) npy_intp left_n = PyArray_SIZE(left_arr); npy_intp right_n = PyArray_SIZE(right_arr); + + /* Empty input on either side → empty inner join result */ + if (left_n == 0 || right_n == 0) { + Py_DECREF(left_arr); Py_DECREF(right_arr); + npy_intp dims[1] = {0}; + PyArrayObject *li_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + PyArrayObject *ri_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + if (!li_arr || !ri_arr) { + Py_XDECREF(li_arr); Py_XDECREF(ri_arr); + return NULL; + } + PyObject *result = PyTuple_Pack(2, (PyObject *)li_arr, (PyObject *)ri_arr); + Py_DECREF(li_arr); Py_DECREF(ri_arr); + return result; + } + const npy_int64 *left = (const npy_int64 *)PyArray_DATA(left_arr); const npy_int64 *right = (const npy_int64 *)PyArray_DATA(right_arr); @@ -850,6 +866,43 @@ accel_left_join(PyObject *self, PyObject *args) npy_intp left_n = PyArray_SIZE(left_arr); npy_intp right_n = PyArray_SIZE(right_arr); + + /* Empty left → empty result; empty right → all left rows unmatched */ + if (left_n == 0) { + Py_DECREF(left_arr); Py_DECREF(right_arr); + npy_intp dims[1] = {0}; + PyArrayObject *li_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + PyArrayObject *ri_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + if (!li_arr || !ri_arr) { + Py_XDECREF(li_arr); Py_XDECREF(ri_arr); + return NULL; + } + PyObject *result = PyTuple_Pack(3, (PyObject *)li_arr, (PyObject *)ri_arr, Py_False); + Py_DECREF(li_arr); Py_DECREF(ri_arr); + return result; + } + if (right_n == 0) { + npy_intp dims[1] = {left_n}; + PyArrayObject *li_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + PyArrayObject *ri_arr = (PyArrayObject *)PyArray_SimpleNew(1, dims, NPY_INTP); + if (!li_arr || !ri_arr) { + Py_XDECREF(li_arr); Py_XDECREF(ri_arr); + Py_DECREF(left_arr); Py_DECREF(right_arr); + return NULL; + } + npy_intp *li = (npy_intp *)PyArray_DATA(li_arr); + npy_intp *ri = (npy_intp *)PyArray_DATA(ri_arr); + for (npy_intp i = 0; i < left_n; i++) { + li[i] = i; + ri[i] = -1; + } + Py_DECREF(left_arr); Py_DECREF(right_arr); + PyObject *result = PyTuple_Pack(3, (PyObject *)li_arr, (PyObject *)ri_arr, + left_n > 0 ? Py_True : Py_False); + Py_DECREF(li_arr); Py_DECREF(ri_arr); + return result; + } + const npy_int64 *left = (const npy_int64 *)PyArray_DATA(left_arr); const npy_int64 *right = (const npy_int64 *)PyArray_DATA(right_arr); From 7f328920a6aa2896f3daa729e8fb59ad21e45722 Mon Sep 17 00:00:00 2001 From: David Fulford Date: Thu, 2 Apr 2026 18:16:43 -0500 Subject: [PATCH 10/10] fix: add assertion validating cards length in _build_composite_key Co-Authored-By: Claude Opus 4.6 (1M context) --- tafra/group.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tafra/group.py b/tafra/group.py index cc938bf..897c72a 100644 --- a/tafra/group.py +++ b/tafra/group.py @@ -567,6 +567,7 @@ def _build_composite_key( # negative integers to non-negative range. if cards is None: cards = [int(c.max()) + 1 for c in encoded] + assert len(cards) == len(encoded), "cards length must match encoded columns" # check for overflow: product of all cardinalities must fit int64 product = 1