BUG: float64 hash table for handling NAs in Series.unique, close #714

Author: wesm

pandas/core/series.py

@@ -814,11 +814,16 @@ def unique(self):
         uniques : ndarray
         """
         values = self.values
-        if not values.dtype == np.object_:
-            values = values.astype('O')
-        table = lib.PyObjectHashTable(len(values))
-        uniques = lib.list_to_object_array(table.unique(values))
-        return lib.maybe_convert_objects(uniques)
+        if issubclass(values.dtype.type, np.floating):
+            table = lib.Float64HashTable(len(values))
+            uniques = np.array(table.unique(values), dtype='f8')
+        else:
+            if not values.dtype == np.object_:
+                values = values.astype('O')
+            table = lib.PyObjectHashTable(len(values))
+            uniques = lib.list_to_object_array(table.unique(values))
+            uniques = lib.maybe_convert_objects(uniques)
+        return uniques
 
     def nunique(self):
         """

pandas/src/hashtable.pyx

@@ -579,6 +579,55 @@ cdef class Int64HashTable:
 
         return uniques
 
+ONAN = np.nan
+
+cdef class Float64HashTable:
+
+    cdef:
+        kh_float64_t *table
+
+    def __init__(self, size_hint=1):
+        if size_hint is not None:
+            kh_resize_float64(self.table, size_hint)
+
+    def __cinit__(self):
+        self.table = kh_init_float64()
+
+    def __dealloc__(self):
+        kh_destroy_float64(self.table)
+
+    def factorize(self, ndarray[object] values):
+        reverse = {}
+        labels, counts = self.get_labels(values, reverse, 0)
+        return reverse, labels, counts
+
+    def unique(self, ndarray[float64_t] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            Py_ssize_t idx, count = 0
+            int ret
+            float64_t val
+            khiter_t k
+            list uniques = []
+            bint seen_na = 0
+
+        # TODO: kvec
+
+        for i in range(n):
+            val = values[i]
+
+            if val == val:
+                k = kh_get_float64(self.table, val)
+                if k == self.table.n_buckets:
+                    k = kh_put_float64(self.table, val, &ret)
+                    uniques.append(val)
+                    count += 1
+            elif not seen_na:
+                seen_na = 1
+                uniques.append(ONAN)
+
+        return uniques
+
 cdef class PyObjectHashTable:
 
     cdef:

pandas/src/khash.h

@@ -112,6 +112,7 @@ int main() {
 #include <stdlib.h>
 #include <string.h>
 #include <limits.h>
+#include <Python.h>
 
 /* compipler specific configuration */
 
@@ -129,6 +130,8 @@ typedef unsigned long long khuint64_t;
 typedef signed long long khint64_t;
 #endif
 
+typedef double khfloat64_t;
+
 #ifndef PANDAS_INLINE
   #if defined(__GNUC__)
     #define PANDAS_INLINE __inline__
@@ -346,6 +349,12 @@ static const double __ac_HASH_UPPER = 0.77;
   @abstract     64-bit integer comparison function
  */
 #define kh_int64_hash_equal(a, b) ((a) == (b))
+
+// kludge
+
+#define kh_float64_hash_func _Py_HashDouble
+#define kh_float64_hash_equal kh_int64_hash_equal
+
 /*! @function
   @abstract     const char* hash function
   @param  s     Pointer to a null terminated string
@@ -544,6 +553,9 @@ static PANDAS_INLINE khint_t __ac_Wang_hash(khint_t key)
 #define KHASH_MAP_INIT_INT64(name, khval_t)								\
 	KHASH_INIT(name, khint64_t, khval_t, 1, kh_int64_hash_func, kh_int64_hash_equal)
 
+#define KHASH_MAP_INIT_FLOAT64(name, khval_t)								\
+	KHASH_INIT(name, khfloat64_t, khval_t, 1, kh_float64_hash_func, kh_float64_hash_equal)
+
 typedef const char *kh_cstr_t;
 /*! @function
   @abstract     Instantiate a hash map containing const char* keys
@@ -584,12 +596,14 @@ KHASH_SET_INIT_PYOBJECT(pyset)
 #define kh_exist_pymap(h, k) (kh_exist(h, k))
 #define kh_exist_pyset(h, k) (kh_exist(h, k))
 #define kh_exist_str(h, k) (kh_exist(h, k))
+#define kh_exist_float64(h, k) (kh_exist(h, k))
 #define kh_exist_int64(h, k) (kh_exist(h, k))
 #define kh_exist_int32(h, k) (kh_exist(h, k))
 
 KHASH_MAP_INIT_STR(str, Py_ssize_t)
 
 KHASH_MAP_INIT_INT(int32, Py_ssize_t)
 KHASH_MAP_INIT_INT64(int64, Py_ssize_t)
+KHASH_MAP_INIT_FLOAT64(float64, Py_ssize_t)
 
 #endif /* __AC_KHASH_H */

pandas/src/khash.pxd

@@ -1,5 +1,5 @@
 from cpython cimport PyObject
-from numpy cimport int64_t, int32_t, uint32_t
+from numpy cimport int64_t, int32_t, uint32_t, float64_t
 
 cdef extern from "khash.h":
     ctypedef uint32_t khint_t
@@ -71,6 +71,22 @@ cdef extern from "khash.h":
 
     bint kh_exist_int64(kh_int64_t*, khiter_t)
 
+    ctypedef struct kh_float64_t:
+        khint_t n_buckets, size, n_occupied, upper_bound
+        uint32_t *flags
+        float64_t *keys
+        Py_ssize_t *vals
+
+    inline kh_float64_t* kh_init_float64()
+    inline void kh_destroy_float64(kh_float64_t*)
+    inline void kh_clear_float64(kh_float64_t*)
+    inline khint_t kh_get_float64(kh_float64_t*, float64_t)
+    inline void kh_resize_float64(kh_float64_t*, khint_t)
+    inline khint_t kh_put_float64(kh_float64_t*, float64_t, int*)
+    inline void kh_del_float64(kh_float64_t*, khint_t)
+
+    bint kh_exist_float64(kh_float64_t*, khiter_t)
+
     ctypedef struct kh_int32_t:
         khint_t n_buckets, size, n_occupied, upper_bound
         uint32_t *flags