Skip to content

Commit 45bc149

Browse files
jbrockmendelclaude
andcommitted
BUG: length-aware string hashtable keys (GH#34551)
pandas keys several khash tables on a bare `const char *`, which hashes to the first NUL and compares with strcmp, so any string containing an embedded NUL is truncated for both hashing and comparison. Distinct values that share a prefix up to their first NUL therefore collapse into one. Key those tables on an explicit (pointer, length) pair instead. The hash is X31 over exactly `len` bytes -- bit-identical to `__ac_X31_hash_string` for NUL-free input, so bucket distribution is unchanged -- and equality is a length check plus a byte comparison. `PyUnicode_AsUTF8AndSize` and `_token_len` supply the length, so no path pays an extra strlen. This fixes `unique`, `factorize`, `Series.unique`, `Index.unique` and `groupby` on object dtype, plus `read_csv`'s `dtype=object` interning and `dtype="category"` conversion. The na_values/true_values/false_values hashsets change key type but still derive their length with strlen, so NA handling is unchanged here; giving them an explicit length changes user-visible NA semantics and is left to a follow-up. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1 parent 644979a commit 45bc149

11 files changed

Lines changed: 257 additions & 65 deletions

File tree

doc/source/whatsnew/v3.1.0.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,7 @@ Strings
410410
- Bug in :meth:`Series.str.rsplit` and :meth:`Index.str.rsplit` silently accepting a compiled regex and returning incorrect results (:issue:`29633`)
411411
- Bug in :meth:`Series.str.split` with :class:`ArrowDtype` ``string`` not inferring regex for multi-character patterns when ``regex=None``, causing the pattern to be treated as a literal instead of a regular expression (:issue:`58321`)
412412
- Bug in :meth:`Series.unique`, :meth:`Index.unique` and :func:`factorize` on object dtype returning incorrect results when the values were not UTF-8 encodable, e.g. lone surrogates (:issue:`34550`)
413+
- Bug in :meth:`Series.unique`, :meth:`Index.unique`, :func:`factorize` and ``groupby`` on object dtype collapsing distinct strings that are identical up to an embedded NUL byte, e.g. ``""`` and ``"\x00"``, into a single value (:issue:`34551`)
413414

414415
Interval
415416
^^^^^^^^
@@ -496,6 +497,7 @@ I/O
496497
- Fixed bug in :func:`read_csv` with ``engine="pyarrow"`` where passing tuples in ``names`` produced flat columns instead of :class:`MultiIndex` columns as with the other engines (:issue:`65862`)
497498
- Fixed bug in :func:`read_csv` with the ``c`` engine where a quoted field containing an embedded NUL byte was silently truncated at the NUL under the default string dtype or ``dtype_backend="pyarrow"`` (:issue:`66415`)
498499
- Fixed bug in :func:`read_csv` with the ``c`` engine where an embedded ``\r`` followed by a space in an unquoted field could cause an infinite re-parsing loop, producing spurious rows or a buffer overflow (:issue:`51141`)
500+
- Fixed bug in :func:`read_csv` with the ``c`` engine where two fields differing only after an embedded NUL byte were read as the same value with ``dtype=object`` or ``dtype="category"`` (:issue:`19886`)
499501
- Fixed bug in :func:`read_excel` where usage of ``skiprows`` could lead to an infinite loop (:issue:`64027`)
500502
- Fixed bug in :func:`read_excel` with the ``openpyxl`` engine where reading a sheet set its ``max_row`` and ``max_column`` to ``None`` on the workbook exposed through ``ExcelFile.book`` (:issue:`63010`)
501503
- Fixed bug in :func:`read_sas` where ``encoding="infer"`` raised ``LookupError: unknown encoding: infer`` instead of falling back to latin-1, for SAS7BDAT files recording an encoding pandas does not recognize and for XPORT files, which record no encoding at all (:issue:`66470`)

pandas/_libs/hashtable.pxd

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ from pandas._libs.khash cimport (
2222
kh_int64_t,
2323
kh_pymap_t,
2424
kh_str_t,
25+
kh_strview_t,
2526
kh_uint8_t,
2627
kh_uint16_t,
2728
kh_uint32_t,

pandas/_libs/hashtable.pyx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ from pandas._libs.khash cimport (
2929
kh_needed_n_buckets,
3030
kh_python_hash_equal,
3131
kh_python_hash_func,
32+
kh_strview,
3233
khiter_t,
3334
)
3435
from pandas._libs.missing cimport (

pandas/_libs/hashtable_class_helper.pxi.in

Lines changed: 53 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ Template for each `dtype` helper function for hashtable
33

44
WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
55
"""
6-
from cpython.unicode cimport PyUnicode_AsUTF8
6+
from cpython.unicode cimport PyUnicode_AsUTF8AndSize
77

88
{{py:
99

@@ -964,6 +964,11 @@ cdef class {{name}}Factorizer(Factorizer):
964964

965965

966966
cdef class StringHashTable(HashTable):
967+
# Keys borrow each str's cached UTF-8 buffer via PyUnicode_AsUTF8AndSize,
968+
# which yields the length alongside the pointer, so keying on an explicit
969+
# length costs no extra pass. The buffer lives exactly as long as the str
970+
# does, so callers must keep the str alive while the key is in the table.
971+
967972
# these by-definition *must* be strings
968973
# or a sentinel np.nan / None missing value
969974
na_string_sentinel = '__nan__'
@@ -984,8 +989,8 @@ cdef class StringHashTable(HashTable):
984989
def sizeof(self, deep: bool = False) -> int:
985990
overhead = 4 * sizeof(uint32_t) + 3 * sizeof(uint32_t*)
986991
for_flags = max(1, self.table.n_buckets >> 5) * sizeof(uint32_t)
987-
for_pairs = self.table.n_buckets * (sizeof(char *) + # keys
988-
sizeof(Py_ssize_t)) # vals
992+
for_pairs = self.table.n_buckets * (sizeof(kh_strview_t) + # keys
993+
sizeof(Py_ssize_t)) # vals
989994
return overhead + for_flags + for_pairs
990995

991996
def get_state(self) -> dict[str, int]:
@@ -1000,8 +1005,11 @@ cdef class StringHashTable(HashTable):
10001005
cpdef get_item(self, str val):
10011006
cdef:
10021007
khiter_t k
1003-
const char *v
1004-
v = PyUnicode_AsUTF8(val)
1008+
kh_strview_t v
1009+
Py_ssize_t length
1010+
const char *ptr
1011+
ptr = PyUnicode_AsUTF8AndSize(val, &length)
1012+
v = kh_strview(ptr, length)
10051013

10061014
k = kh_get_str(self.table, v)
10071015
if k != self.table.n_buckets:
@@ -1013,9 +1021,12 @@ cdef class StringHashTable(HashTable):
10131021
cdef:
10141022
khiter_t k
10151023
int ret = 0
1016-
const char *v
1024+
kh_strview_t v
1025+
Py_ssize_t length
1026+
const char *ptr
10171027

1018-
v = PyUnicode_AsUTF8(key)
1028+
ptr = PyUnicode_AsUTF8AndSize(key, &length)
1029+
v = kh_strview(ptr, length)
10191030

10201031
k = kh_put_str(self.table, v, &ret)
10211032
if kh_exist_str(self.table, k):
@@ -1033,16 +1044,17 @@ cdef class StringHashTable(HashTable):
10331044
intp_t *resbuf = <intp_t*>labels.data
10341045
khiter_t k
10351046
kh_str_t *table = self.table
1036-
const char *v
1037-
const char **vecs
1047+
kh_strview_t *vecs
1048+
Py_ssize_t length
1049+
const char *ptr
10381050

1039-
vecs = <const char **>malloc(n * sizeof(char *))
1051+
vecs = <kh_strview_t *>malloc(n * sizeof(kh_strview_t))
10401052
if vecs is NULL:
10411053
raise MemoryError()
10421054
for i in range(n):
10431055
val = values[i]
1044-
v = PyUnicode_AsUTF8(val)
1045-
vecs[i] = v
1056+
ptr = PyUnicode_AsUTF8AndSize(val, &length)
1057+
vecs[i] = kh_strview(ptr, length)
10461058

10471059
with nogil:
10481060
for i in range(n):
@@ -1064,24 +1076,27 @@ cdef class StringHashTable(HashTable):
10641076
Py_ssize_t i, n = len(values)
10651077
int ret = 0
10661078
object val
1067-
const char *v
1079+
kh_strview_t v
1080+
kh_strview_t *vecs
1081+
Py_ssize_t length
1082+
const char *ptr
10681083
khiter_t k
10691084
intp_t[::1] locs = np.empty(n, dtype=np.intp)
10701085

10711086
# these by-definition *must* be strings
1072-
vecs = <const char **>malloc(n * sizeof(char *))
1087+
vecs = <kh_strview_t *>malloc(n * sizeof(kh_strview_t))
10731088
if vecs is NULL:
10741089
raise MemoryError()
10751090
for i in range(n):
10761091
val = values[i]
10771092

10781093
if isinstance(val, str):
1079-
# GH#31499 if we have an np.str_ PyUnicode_AsUTF8 won't recognize
1080-
# it as a str, even though isinstance does.
1081-
v = PyUnicode_AsUTF8(<str>val)
1094+
# GH#31499 if we have an np.str_ PyUnicode_AsUTF8AndSize won't
1095+
# recognize it as a str, even though isinstance does.
1096+
ptr = PyUnicode_AsUTF8AndSize(<str>val, &length)
10821097
else:
1083-
v = PyUnicode_AsUTF8(self.na_string_sentinel)
1084-
vecs[i] = v
1098+
ptr = PyUnicode_AsUTF8AndSize(self.na_string_sentinel, &length)
1099+
vecs[i] = kh_strview(ptr, length)
10851100

10861101
with nogil:
10871102
for i in range(n):
@@ -1103,24 +1118,26 @@ cdef class StringHashTable(HashTable):
11031118
Py_ssize_t i, n = len(values)
11041119
int ret = 0
11051120
object val
1106-
const char *v
1107-
const char **vecs
1121+
kh_strview_t v
1122+
kh_strview_t *vecs
1123+
Py_ssize_t length
1124+
const char *ptr
11081125
khiter_t k
11091126

11101127
# these by-definition *must* be strings
1111-
vecs = <const char **>malloc(n * sizeof(char *))
1128+
vecs = <kh_strview_t *>malloc(n * sizeof(kh_strview_t))
11121129
if vecs is NULL:
11131130
raise MemoryError()
11141131
for i in range(n):
11151132
val = values[i]
11161133

11171134
if isinstance(val, str):
1118-
# GH#31499 if we have an np.str_ PyUnicode_AsUTF8 won't recognize
1119-
# it as a str, even though isinstance does.
1120-
v = PyUnicode_AsUTF8(<str>val)
1135+
# GH#31499 if we have an np.str_ PyUnicode_AsUTF8AndSize won't
1136+
# recognize it as a str, even though isinstance does.
1137+
ptr = PyUnicode_AsUTF8AndSize(<str>val, &length)
11211138
else:
1122-
v = PyUnicode_AsUTF8(self.na_string_sentinel)
1123-
vecs[i] = v
1139+
ptr = PyUnicode_AsUTF8AndSize(self.na_string_sentinel, &length)
1140+
vecs[i] = kh_strview(ptr, length)
11241141

11251142
with nogil:
11261143
for i in range(n):
@@ -1174,8 +1191,10 @@ cdef class StringHashTable(HashTable):
11741191
int64_t[::1] uindexer
11751192
int ret = 0
11761193
object val
1177-
const char *v
1178-
const char **vecs
1194+
kh_strview_t v
1195+
kh_strview_t *vecs
1196+
Py_ssize_t length
1197+
const char *ptr
11791198
khiter_t k
11801199
bint use_na_value
11811200
bint non_null_na_value
@@ -1188,7 +1207,7 @@ cdef class StringHashTable(HashTable):
11881207
non_null_na_value = not checknull(na_value)
11891208

11901209
# assign pointers and pre-filter out missing (if ignore_na)
1191-
vecs = <const char **>malloc(n * sizeof(char *))
1210+
vecs = <kh_strview_t *>malloc(n * sizeof(kh_strview_t))
11921211
if vecs is NULL:
11931212
raise MemoryError()
11941213
for i in range(n):
@@ -1209,15 +1228,15 @@ cdef class StringHashTable(HashTable):
12091228
else:
12101229
# if ignore_na is False, we also stringify NaN/None/etc.
12111230
try:
1212-
v = PyUnicode_AsUTF8(<str>val)
1231+
ptr = PyUnicode_AsUTF8AndSize(<str>val, &length)
12131232
except UnicodeEncodeError:
12141233
# GH#34550 val is not utf-8 encodable (e.g. a lone
12151234
# surrogate). Retain the repr in `refs` so the buffer
1216-
# `v` points to outlives the nogil loop below.
1235+
# `ptr` points to outlives the nogil loop below.
12171236
val = repr(val)
12181237
refs.append(val)
1219-
v = PyUnicode_AsUTF8(<str>val)
1220-
vecs[i] = v
1238+
ptr = PyUnicode_AsUTF8AndSize(<str>val, &length)
1239+
vecs[i] = kh_strview(ptr, length)
12211240

12221241
# compute
12231242
with nogil:

pandas/_libs/include/pandas/vendored/klib/khash.h

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -749,7 +749,10 @@ typedef const char *kh_cstr_t;
749749
#define kh_exist_set_int8(h, k) (kh_exist(h, k))
750750
#define kh_exist_set_uint8(h, k) (kh_exist(h, k))
751751

752-
KHASH_MAP_INIT_STR(str, size_t)
752+
// kh_str_t is deliberately *not* instantiated here. pandas keys its string
753+
// tables on an explicit (pointer, length) pair so that values containing
754+
// embedded NUL bytes are not truncated; see KHASH_MAP_INIT_STRVIEW in
755+
// khash_python.h.
753756
KHASH_MAP_INIT_INT(int32, size_t)
754757
KHASH_MAP_INIT_UINT(uint32, size_t)
755758
KHASH_MAP_INIT_INT64(int64, size_t)

pandas/_libs/include/pandas/vendored/klib/khash_python.h

Lines changed: 83 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -380,11 +380,71 @@ KHASH_SET_INIT_PYOBJECT(pyset)
380380
#define kh_exist_pymap(h, k) (kh_exist(h, k))
381381
#define kh_exist_pyset(h, k) (kh_exist(h, k))
382382

383-
KHASH_MAP_INIT_STR(strbox, kh_pyobject_t)
383+
// Length-aware string key.
384+
//
385+
// A bare `const char *` key (klib's kh_cstr_t) hashes and compares only up to
386+
// the first NUL byte, so distinct values that share a NUL-terminated prefix
387+
// collapse into a single entry. Python strings, CSV fields and Arrow string
388+
// buffers can all contain embedded NULs, which made that silently return wrong
389+
// results (GH#34551, GH#19886). Keying on an explicit (pointer, length) pair
390+
// fixes it, and additionally lets a key point at a slice of a larger buffer
391+
// without copying or NUL-terminating it.
392+
//
393+
// The pointed-to bytes are borrowed, exactly as for kh_cstr_t: the caller must
394+
// keep them alive for as long as the entry is in the table.
395+
typedef struct {
396+
const char *ptr;
397+
size_t len;
398+
} kh_strview_t;
399+
400+
static inline kh_strview_t kh_strview(const char *ptr, size_t len) {
401+
kh_strview_t result;
402+
result.ptr = ptr;
403+
result.len = len;
404+
return result;
405+
}
406+
407+
// X31 over exactly `len` bytes. Bit-identical to __ac_X31_hash_string for
408+
// NUL-free input, so bucket distribution -- and hence performance -- is
409+
// unchanged for the data that used to work.
410+
static inline khuint_t kh_strview_hash_func(kh_strview_t key) {
411+
khuint_t h = 0;
412+
if (key.len) {
413+
h = (khuint_t)key.ptr[0];
414+
for (size_t i = 1; i < key.len; ++i)
415+
h = (h << 5) - h + (khuint_t)key.ptr[i];
416+
}
417+
return h;
418+
}
419+
420+
// The length check alone settles most probes, and the byte loop is deliberate:
421+
// memcmp with a runtime length does not inline, and one libc call per token
422+
// costs ~9% on a bool column, where a hashset lookup hits on nearly every
423+
// token and there is little other per-token work to absorb it.
424+
static inline int kh_strview_hash_equal(kh_strview_t a, kh_strview_t b) {
425+
if (a.len != b.len)
426+
return 0;
427+
for (size_t i = 0; i < a.len; ++i)
428+
if (a.ptr[i] != b.ptr[i])
429+
return 0;
430+
return 1;
431+
}
432+
433+
#define KHASH_SET_INIT_STRVIEW(name) \
434+
KHASH_INIT(name, kh_strview_t, char, 0, kh_strview_hash_func, \
435+
kh_strview_hash_equal)
436+
437+
#define KHASH_MAP_INIT_STRVIEW(name, khval_t) \
438+
KHASH_INIT(name, kh_strview_t, khval_t, 1, kh_strview_hash_func, \
439+
kh_strview_hash_equal)
440+
441+
KHASH_MAP_INIT_STRVIEW(str, size_t)
442+
KHASH_MAP_INIT_STRVIEW(strbox, kh_pyobject_t)
384443

385444
typedef struct {
386445
kh_str_t *table;
387446
int starts[256];
447+
int has_empty;
388448
} kh_str_starts_t;
389449

390450
typedef kh_str_starts_t *p_kh_str_starts_t;
@@ -396,22 +456,35 @@ static inline p_kh_str_starts_t kh_init_str_starts(void) {
396456
return result;
397457
}
398458

399-
static inline khuint_t kh_put_str_starts_item(kh_str_starts_t *table, char *key,
400-
int *ret) {
401-
khuint_t result = kh_put_str(table->table, key, ret);
459+
// NOTE: these two derive the length with strlen, so an na_value or a field
460+
// containing an embedded NUL is still compared only up to that NUL -- exactly
461+
// the pre-existing behavior. Giving them an explicit length is what fixes
462+
// GH#19886, and is deliberately left to a follow-up because it changes
463+
// user-visible NA semantics (and hence dtype inference).
464+
static inline khuint_t kh_put_str_starts_item(kh_str_starts_t *table,
465+
const char *key, int *ret) {
466+
size_t len = strlen(key);
467+
khuint_t result = kh_put_str(table->table, kh_strview(key, len), ret);
402468
if (*ret != 0) {
403-
table->starts[(unsigned char)key[0]] = 1;
469+
// The empty key gets its own flag rather than a slot in starts[]. Sharing
470+
// starts['\0'] with it would make every token that merely *begins* with an
471+
// embedded NUL look like a candidate.
472+
if (len == 0)
473+
table->has_empty = 1;
474+
else
475+
table->starts[(unsigned char)key[0]] = 1;
404476
}
405477
return result;
406478
}
407479

408480
static inline khuint_t kh_get_str_starts_item(const kh_str_starts_t *table,
409481
const char *key) {
410-
unsigned char ch = *key;
411-
if (table->starts[ch]) {
412-
if (ch == '\0' || kh_get_str(table->table, key) != table->table->n_buckets)
413-
return 1;
414-
}
482+
size_t len = strlen(key);
483+
if (len == 0)
484+
return (khuint_t)table->has_empty;
485+
if (table->starts[(unsigned char)key[0]] &&
486+
kh_get_str(table->table, kh_strview(key, len)) != table->table->n_buckets)
487+
return 1;
415488
return 0;
416489
}
417490

0 commit comments

Comments
 (0)