ENH: int64 tinkering with khash

Author: wesm

pandas/src/khash.pxd

@@ -58,7 +58,7 @@ cdef extern from "khash.h":
     ctypedef struct kh_int64_t:
         khint_t n_buckets, size, n_occupied, upper_bound
         uint32_t *flags
-        kh_cstr_t *keys
+        int64_t *keys
         Py_ssize_t *vals
 
     inline kh_int64_t* kh_init_int64()

pandas/src/sandbox.pyx

@@ -321,6 +321,113 @@ cdef class StringHashTable:
         # return None
         return reverse, labels, counts[:count].copy()
 
+cdef class Int64HashTable:
+
+    cdef:
+        kh_int64_t *table
+
+    def __init__(self, size_hint=1):
+        if size_hint is not None:
+            kh_resize_int64(self.table, size_hint)
+
+    def __cinit__(self):
+        self.table = kh_init_int64()
+
+    def __dealloc__(self):
+        kh_destroy_int64(self.table)
+
+    cdef inline int check_type(self, object val):
+        return PyString_Check(val)
+
+    cpdef get_item(self, int64_t val):
+        cdef khiter_t k
+        k = kh_get_int64(self.table, val)
+        if k != self.table.n_buckets:
+            return self.table.vals[k]
+        else:
+            raise KeyError(val)
+
+    def get_iter_test(self, int64_t key, Py_ssize_t iterations):
+        cdef Py_ssize_t i, val
+        for i in range(iterations):
+            k = kh_get_int64(self.table, val)
+            if k != self.table.n_buckets:
+                val = self.table.vals[k]
+
+    cpdef set_item(self, int64_t key, Py_ssize_t val):
+        cdef:
+            khiter_t k
+            int ret
+
+        k = kh_put_int64(self.table, key, &ret)
+        self.table.keys[k] = key
+        if kh_exist_int64(self.table, k):
+            self.table.vals[k] = val
+        else:
+            raise KeyError(key)
+
+    def map_locations(self, ndarray[int64_t] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            int ret
+            int64_t val
+            khiter_t k
+
+        for i in range(n):
+            val = values[i]
+            k = kh_put_int64(self.table, val, &ret)
+            # print 'putting %s, %s' % (val, count)
+            self.table.vals[k] = i
+
+    def lookup_locations(self, ndarray[int64_t] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            int ret
+            int64_t val
+            khiter_t k
+            ndarray[int32_t] locs = np.empty(n, dtype='i4')
+
+        for i in range(n):
+            val = values[i]
+            k = kh_get_int64(self.table, val)
+            if k != self.table.n_buckets:
+                locs[i] = self.table.vals[k]
+            else:
+                locs[i] = -1
+
+        return locs
+
+    def factorize(self, ndarray[object] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            ndarray[int32_t] labels = np.empty(n, dtype=np.int32)
+            ndarray[int32_t] counts = np.empty(n, dtype=np.int32)
+            dict reverse = {}
+            Py_ssize_t idx, count = 0
+            int ret
+            int64_t val
+            khiter_t k
+
+        for i in range(n):
+            val = values[i]
+            k = kh_get_int64(self.table, val)
+            if k != self.table.n_buckets:
+                idx = self.table.vals[k]
+                labels[i] = idx
+                counts[idx] = counts[idx] + 1
+            else:
+                k = kh_put_int64(self.table, val, &ret)
+                if not ret:
+                    kh_del_int64(self.table, k)
+                self.table.vals[k] = count
+                reverse[count] = val
+                labels[i] = count
+                counts[count] = 1
+                count += 1
+
+        # return None
+        return reverse, labels, counts[:count].copy()
+
 from libc.stdlib cimport free
 
 cdef class PyObjectHashTable: