Add hashtable for C buffer type (ptr + len) avoiding strlen overhead

Author: immerrr

pandas/hashtable.pyx

@@ -1,4 +1,5 @@
-from cpython cimport PyObject, Py_INCREF, PyList_Check, PyTuple_Check
+from cpython cimport (PyObject, Py_INCREF, PyList_Check, PyTuple_Check,
+                      PyString_AsStringAndSize, PyDict_Copy)
 
 from khash cimport *
 from numpy cimport *
@@ -843,6 +844,127 @@ cdef class PyObjectHashTable(HashTable):
         return labels
 
 
+cdef inline cbuf_t to_cbuf(object s):
+    cdef cbuf_t output
+    PyString_AsStringAndSize(s, <char**>&output.buf, &output.len)
+    return output
+
+
+cdef class CBufHashTable(HashTable):
+    cdef kh_cbuf_map_t *table
+
+    def __cinit__(self, int size_hint=1):
+        self.table = kh_init_cbuf_map()
+        if size_hint is not None:
+            kh_resize_cbuf_map(self.table, size_hint)
+
+    def __dealloc__(self):
+        kh_destroy_cbuf_map(self.table)
+
+    cdef inline int check_type(self, object val):
+        return util.is_string_object(val)
+
+    cpdef get_item(self, object val):
+        cdef khiter_t it
+        it = kh_get_cbuf_map(self.table, to_cbuf(val))
+        if it != self.table.n_buckets:
+            return self.table.vals[it]
+        else:
+            raise KeyError(val)
+
+    def get_iter_test(self, object key, Py_ssize_t iterations):
+        cdef khiter_t it
+        cdef Py_ssize_t i, val
+        for i in range(iterations):
+            it = kh_get_cbuf_map(self.table, to_cbuf(key))
+            if it != self.table.n_buckets:
+                val = self.table.vals[it]
+
+    cpdef set_item(self, object key, Py_ssize_t val):
+        cdef:
+            khiter_t it
+            int ret = 0
+            cbuf_t buf
+
+        buf = to_cbuf(key)
+
+        it = kh_put_cbuf_map(self.table, buf, &ret)
+        self.table.keys[it] = buf
+        if kh_exist_cbuf_map(self.table, it):
+            self.table.vals[it] = val
+        else:
+            raise KeyError(key)
+
+    def get_indexer(self, ndarray[object] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            ndarray[int64_t] labels = np.empty(n, dtype=np.int64)
+            cbuf_t buf
+            int64_t[::1] out = labels
+            khiter_t it
+            kh_cbuf_map_t *table = self.table
+
+        for i in range(n):
+            buf = to_cbuf(values[i])
+            it = kh_get_cbuf_map(table, buf)
+            if it != table.n_buckets:
+                out[i] = table.vals[it]
+            else:
+                out[i] = -1
+        return labels
+
+    def unique(self, ndarray[object] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            Py_ssize_t idx, count = 0
+            int ret = 0
+            object val
+            cbuf_t buf
+            khiter_t it
+            ObjectVector uniques = ObjectVector()
+
+        for i in range(n):
+            val = values[i]
+            buf = to_cbuf(val)
+            it = kh_get_cbuf_map(self.table, buf)
+            if it == self.table.n_buckets:
+                it = kh_put_cbuf_map(self.table, buf, &ret)
+                count += 1
+                uniques.append(val)
+
+        return uniques.to_array()
+
+    def factorize(self, ndarray[object] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            ndarray[int64_t] labels = np.empty(n, dtype=np.int64)
+            list reverse = []
+            Py_ssize_t idx, count = 0
+            int ret = 0
+            object val
+            cbuf_t buf
+            khiter_t it
+
+        for i in range(n):
+            val = values[i]
+            buf = to_cbuf(val)
+            it = kh_get_cbuf_map(self.table, buf)
+            if it != self.table.n_buckets:
+                idx = self.table.vals[it]
+                labels[i] = idx
+            else:
+                it = kh_put_cbuf_map(self.table, buf, &ret)
+
+                self.table.vals[it] = count
+                reverse.append(val)
+                labels[i] = count
+                count += 1
+
+        return PyDict_Copy(enumerate(reverse)), labels
+
+
+
+
 cdef class Factorizer:
     cdef public PyObjectHashTable table
     cdef public ObjectVector uniques

pandas/src/khash.pxd

@@ -115,3 +115,21 @@ cdef extern from "khash_python.h":
 
     bint kh_exist_strbox(kh_strbox_t*, khiter_t)
 
+    ctypedef struct cbuf_t:
+        kh_cstr_t buf
+        Py_ssize_t len
+
+    ctypedef struct kh_cbuf_map_t:
+        khint_t n_buckets, size, n_occupied, upper_bound
+        uint32_t *flags
+        cbuf_t *keys
+        size_t *vals
+
+    inline kh_cbuf_map_t* kh_init_cbuf_map()
+    inline void kh_destroy_cbuf_map(kh_cbuf_map_t*)
+    inline void kh_clear_cbuf_map(kh_cbuf_map_t*)
+    inline khint_t kh_get_cbuf_map(kh_cbuf_map_t*, cbuf_t)
+    inline void kh_resize_cbuf_map(kh_cbuf_map_t*, khint_t)
+    inline khint_t kh_put_cbuf_map(kh_cbuf_map_t*, cbuf_t, int*)
+
+    bint kh_exist_cbuf_map(kh_cbuf_map_t*, khiter_t)

pandas/src/klib/khash_python.h

@@ -168,4 +168,41 @@ KHASH_SET_INIT_PYOBJECT(pyset)
 KHASH_INIT(strbox, kh_cstr_t, kh_pyobject_t, 1,
            str_xxhash_hash_func, kh_str_hash_equal)
 
+/* Plain old C buffer structure */
+typedef struct {
+    kh_cstr_t buf;
+    Py_ssize_t len;
+} cbuf_t;
+
+static khint_t PANDAS_INLINE cbuf_xxhash(cbuf_t val) {
+    switch (val.len) {
+    case 0:
+        return XXH32_EMPTY_HASH;
+    case 1:
+        return XXH32_ONECHAR_HASH[(uint8_t)val.buf[0]];
+    default:
+        return XXH32(val.buf, val.len, XXH32_SEED);
+    }
+}
+
+static int PANDAS_INLINE cbuf_equal(cbuf_t a, cbuf_t b) {
+    int i;
+    if (a.len != b.len) {
+        return 0;
+    }
+    if (a.buf == b.buf) {
+        return 1;
+    }
+    for (i = 0; i < a.len; ++i) {
+        if (a.buf[i] != b.buf[i]) {
+            return 0;
+        }
+    }
+    return 1;
+}
+
+/* [cbuf_t -> size_t] hash map */
+KHASH_INIT(cbuf_map, cbuf_t, size_t, 1, cbuf_xxhash, cbuf_equal)
+#define kh_exist_cbuf_map(h, k) (kh_exist(h, k))
+
 #endif /* _KLIB_KHASH_PYTHON_H_ */