introducing complex hash tables

realead · realead · commit 6fe568ed446c · 2020-12-29T21:38:47.000+01:00
diff --git a/pandas/_libs/hashtable.pxd b/pandas/_libs/hashtable.pxd
@@ -1,12 +1,16 @@
 from numpy cimport intp_t, ndarray
 
 from pandas._libs.khash cimport (
+    complex64_t,
+    complex128_t,
     float32_t,
     float64_t,
     int8_t,
     int16_t,
     int32_t,
     int64_t,
+    kh_complex64_t,
+    kh_complex128_t,
     kh_float32_t,
     kh_float64_t,
     kh_int8_t,
@@ -19,6 +23,8 @@ from pandas._libs.khash cimport (
     kh_uint16_t,
     kh_uint32_t,
     kh_uint64_t,
+    khcomplex64_t,
+    khcomplex128_t,
     uint8_t,
     uint16_t,
     uint32_t,
@@ -90,6 +96,18 @@ cdef class Float32HashTable(HashTable):
     cpdef get_item(self, float32_t val)
     cpdef set_item(self, float32_t key, Py_ssize_t val)
 
+cdef class Complex64HashTable(HashTable):
+    cdef kh_complex64_t *table
+
+    cpdef get_item(self, complex64_t val)
+    cpdef set_item(self, complex64_t key, Py_ssize_t val)
+
+cdef class Complex128HashTable(HashTable):
+    cdef kh_complex128_t *table
+
+    cpdef get_item(self, complex128_t val)
+    cpdef set_item(self, complex128_t key, Py_ssize_t val)
+
 cdef class PyObjectHashTable(HashTable):
     cdef kh_pymap_t *table
 
diff --git a/pandas/_libs/hashtable.pyx b/pandas/_libs/hashtable.pyx
@@ -13,7 +13,15 @@ cnp.import_array()
 
 
 from pandas._libs cimport util
-from pandas._libs.khash cimport KHASH_TRACE_DOMAIN, kh_str_t, khiter_t
+from pandas._libs.khash cimport (
+    KHASH_TRACE_DOMAIN,
+    are_equal_khcomplex64_t,
+    are_equal_khcomplex128_t,
+    kh_str_t,
+    khcomplex64_t,
+    khcomplex128_t,
+    khiter_t,
+)
 from pandas._libs.missing cimport checknull
 
 
diff --git a/pandas/_libs/hashtable_class_helper.pxi.in b/pandas/_libs/hashtable_class_helper.pxi.in
@@ -8,7 +8,34 @@ WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
 {{py:
 
 # name
-cimported_types = ['float32',
+complex_types = ['complex64',
+                 'complex128']
+}}
+
+{{for name in complex_types}}
+cdef kh{{name}}_t to_kh{{name}}_t({{name}}_t val) nogil:
+    cdef kh{{name}}_t res
+    res.real = val.real
+    res.imag = val.imag
+    return res
+
+cdef {{name}}_t to_{{name}}(kh{{name}}_t val) nogil:
+    cdef {{name}}_t res
+    res.real = val.real
+    res.imag = val.imag
+    return res
+
+cdef bint is_nan_kh{{name}}_t(kh{{name}}_t val) nogil:
+    return val.real != val.real or val.imag != val.imag
+{{endfor}}
+
+
+{{py:
+
+# name
+cimported_types = ['complex64',
+                   'complex128',
+                   'float32',
                    'float64',
                    'int8',
                    'int16',
@@ -48,7 +75,9 @@ from pandas._libs.missing cimport C_NA
 # but is included for completeness (rather ObjectVector is used
 # for uniques in hashtables)
 
-dtypes = [('Float64', 'float64', 'float64_t'),
+dtypes = [('Complex128', 'complex128', 'khcomplex128_t'),
+          ('Complex64', 'complex64', 'khcomplex64_t'),
+          ('Float64', 'float64', 'float64_t'),
           ('Float32', 'float32', 'float32_t'),
           ('Int64', 'int64', 'int64_t'),
           ('Int32', 'int32', 'int32_t'),
@@ -94,6 +123,8 @@ ctypedef fused vector_data:
     UInt8VectorData
     Float64VectorData
     Float32VectorData
+    Complex128VectorData
+    Complex64VectorData
     StringVectorData
 
 cdef inline bint needs_resize(vector_data *data) nogil:
@@ -106,7 +137,9 @@ cdef inline bint needs_resize(vector_data *data) nogil:
 {{py:
 
 # name, dtype, c_type
-dtypes = [('Float64', 'float64', 'float64_t'),
+dtypes = [('Complex128', 'complex128', 'khcomplex128_t'),
+          ('Complex64', 'complex64', 'khcomplex64_t'),
+          ('Float64', 'float64', 'float64_t'),
           ('UInt64', 'uint64', 'uint64_t'),
           ('Int64', 'int64', 'int64_t'),
           ('Float32', 'float32', 'float32_t'),
@@ -303,22 +336,24 @@ cdef class HashTable:
 
 {{py:
 
-# name, dtype, c_type, float_group
-dtypes = [('Float64', 'float64', 'float64_t', True),
-          ('UInt64', 'uint64', 'uint64_t', False),
-          ('Int64', 'int64', 'int64_t', False),
-          ('Float32', 'float32', 'float32_t', True),
-          ('UInt32', 'uint32', 'uint32_t', False),
-          ('Int32', 'int32', 'int32_t', False),
-          ('UInt16', 'uint16', 'uint16_t', False),
-          ('Int16', 'int16', 'int16_t', False),
-          ('UInt8', 'uint8', 'uint8_t', False),
-          ('Int8', 'int8', 'int8_t', False)]
+# name, dtype, c_type, float_group, complex_group
+dtypes = [('Complex128', 'complex128', 'khcomplex128_t', True, True),
+          ('Float64', 'float64', 'float64_t', True, False),
+          ('UInt64', 'uint64', 'uint64_t', False, False),
+          ('Int64', 'int64', 'int64_t', False, False),
+          ('Complex64', 'complex64', 'khcomplex64_t', True, True),
+          ('Float32', 'float32', 'float32_t', True, False),
+          ('UInt32', 'uint32', 'uint32_t', False, False),
+          ('Int32', 'int32', 'int32_t', False, False),
+          ('UInt16', 'uint16', 'uint16_t', False, False),
+          ('Int16', 'int16', 'int16_t', False, False),
+          ('UInt8', 'uint8', 'uint8_t', False, False),
+          ('Int8', 'int8', 'int8_t', False, False)]
 
 }}
 
 
-{{for name, dtype, c_type, float_group in dtypes}}
+{{for name, dtype, c_type, float_group, complex_group in dtypes}}
 
 cdef class {{name}}HashTable(HashTable):
 
@@ -339,7 +374,13 @@ cdef class {{name}}HashTable(HashTable):
     def __contains__(self, object key):
         cdef:
             khiter_t k
-        k = kh_get_{{dtype}}(self.table, key)
+            {{c_type}} ckey
+        {{if complex_group}}
+        ckey = to_{{c_type}}(key)
+        {{else}}
+        ckey = key
+        {{endif}}
+        k = kh_get_{{dtype}}(self.table, ckey)
         return k != self.table.n_buckets
 
     def sizeof(self, deep=False):
@@ -353,7 +394,13 @@ cdef class {{name}}HashTable(HashTable):
     cpdef get_item(self, {{dtype}}_t val):
         cdef:
             khiter_t k
-        k = kh_get_{{dtype}}(self.table, val)
+            {{c_type}} cval
+        {{if complex_group}}
+        cval = to_{{c_type}}(val)
+        {{else}}
+        cval = val
+        {{endif}}
+        k = kh_get_{{dtype}}(self.table, cval)
         if k != self.table.n_buckets:
             return self.table.vals[k]
         else:
@@ -363,8 +410,13 @@ cdef class {{name}}HashTable(HashTable):
         cdef:
             khiter_t k
             int ret = 0
-
-        k = kh_put_{{dtype}}(self.table, key, &ret)
+            {{c_type}} ckey
+        {{if complex_group}}
+        ckey = to_{{c_type}}(key)
+        {{else}}
+        ckey = key
+        {{endif}}
+        k = kh_put_{{dtype}}(self.table, ckey, &ret)
         if kh_exist_{{dtype}}(self.table, k):
             self.table.vals[k] = val
         else:
@@ -486,9 +538,17 @@ cdef class {{name}}HashTable(HashTable):
             # We use None, to make it optional, which requires `object` type
             # for the parameter. To please the compiler, we use na_value2,
             # which is only used if it's *specified*.
-            na_value2 = <{{dtype}}_t>na_value
+            {{if complex_group}}
+            na_value2 = to_{{c_type}}(na_value)
+            {{else}}
+            na_value2 = na_value
+            {{endif}}       
         else:
+            {{if complex_group}}
+            na_value2 = to_{{c_type}}(0)
+            {{else}}
             na_value2 = 0
+            {{endif}}   
 
         with nogil:
             for i in range(n):
@@ -499,10 +559,14 @@ cdef class {{name}}HashTable(HashTable):
                         labels[i] = na_sentinel
                         continue
                 elif ignore_na and (
-                {{if not name.lower().startswith(("uint", "int"))}}
-                val != val or
-                {{endif}}
+                {{if complex_group}}
+                not is_nan_{{c_type}}(val) or 
+                (use_na_value and are_equal_{{c_type}}(val,na_value2))
+                {{elif float_group}}
+                val != val or (use_na_value and val == na_value2)
+                {{else}}
                 (use_na_value and val == na_value2)
+                {{endif}}             
                 ):
                     # if missing values do not count as unique values (i.e. if
                     # ignore_na is True), skip the hashtable entry for them,
@@ -625,7 +689,12 @@ cdef class {{name}}HashTable(HashTable):
                 val = values[i]
 
                 # specific for groupby
-                {{if dtype != 'uint64'}}
+                {{if dtype == 'complex64' or dtype== 'complex128'}}
+                # TODO: what should be done here?
+                if val.real < 0:
+                    labels[i] = -1
+                    continue
+                {{elif dtype != 'uint64'}}
                 if val < 0:
                     labels[i] = -1
                     continue
diff --git a/pandas/_libs/hashtable_func_helper.pxi.in b/pandas/_libs/hashtable_func_helper.pxi.in
@@ -6,22 +6,24 @@ WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
 
 {{py:
 
-# dtype, ttype, c_type
-dtypes = [('float64', 'float64', 'float64_t'),
-          ('float32', 'float32', 'float32_t'),
-          ('uint64', 'uint64', 'uint64_t'),
-          ('uint32', 'uint32', 'uint32_t'),
-          ('uint16', 'uint16', 'uint16_t'),
-          ('uint8', 'uint8', 'uint8_t'),
-          ('object', 'pymap', 'object'),
-          ('int64', 'int64', 'int64_t'),
-          ('int32', 'int32', 'int32_t'),
-          ('int16', 'int16', 'int16_t'),
-          ('int8', 'int8', 'int8_t')]
+# dtype, ttype, c_type, complex_group
+dtypes = [('complex128', 'complex128', 'khcomplex128_t', True),
+          ('complex64', 'complex64', 'khcomplex64_t', True),
+          ('float64', 'float64', 'float64_t', False),
+          ('float32', 'float32', 'float32_t', False),
+          ('uint64', 'uint64', 'uint64_t', False),
+          ('uint32', 'uint32', 'uint32_t', False),
+          ('uint16', 'uint16', 'uint16_t', False),
+          ('uint8', 'uint8', 'uint8_t', False),
+          ('object', 'pymap', 'object', False),
+          ('int64', 'int64', 'int64_t', False),
+          ('int32', 'int32', 'int32_t', False),
+          ('int16', 'int16', 'int16_t', False),
+          ('int8', 'int8', 'int8_t', False)]
 
 }}
 
-{{for dtype, ttype, c_type in dtypes}}
+{{for dtype, ttype, c_type, complex_group in dtypes}}
 
 
 @cython.wraparound(False)
@@ -63,6 +65,8 @@ cdef build_count_table_{{dtype}}({{c_type}}[:] values,
 
             {{if dtype == 'float64' or dtype == 'float32'}}
             if val == val or not dropna:
+            {{elif complex_group}}
+            if not is_nan_{{c_type}}(val) or not dropna:
             {{else}}
             if True:
             {{endif}}
@@ -114,7 +118,11 @@ cpdef value_count_{{dtype}}({{c_type}}[:] values, bint dropna):
     with nogil:
         for k in range(table.n_buckets):
             if kh_exist_{{ttype}}(table, k):
+                {{if complex_group}}
+                result_keys[i] = to_{{dtype}}(table.keys[k])
+                {{else}}
                 result_keys[i] = table.keys[k]
+                {{endif}}         
                 result_counts[i] = table.vals[k]
                 i += 1
     {{endif}}
@@ -279,7 +287,9 @@ def ismember_{{dtype}}(const {{c_type}}[:] arr, const {{c_type}}[:] values):
 {{py:
 
 # dtype, ctype, table_type, npy_dtype
-dtypes = [('float64', 'float64_t', 'float64', 'float64'),
+dtypes = [('complex128', 'khcomplex128_t', 'complex128', 'complex128'),
+          ('complex64', 'khcomplex64_t', 'complex64', 'complex64'),
+          ('float64', 'float64_t', 'float64', 'float64'),
           ('float32', 'float32_t', 'float32', 'float32'),
           ('int64', 'int64_t', 'int64', 'int64'),
           ('int32', 'int32_t', 'int32', 'int32'),
diff --git a/pandas/_libs/khash.pxd b/pandas/_libs/khash.pxd
@@ -1,5 +1,7 @@
 from cpython.object cimport PyObject
 from numpy cimport (
+    complex64_t,
+    complex128_t,
     float32_t,
     float64_t,
     int8_t,
@@ -19,6 +21,18 @@ cdef extern from "khash_python.h":
     ctypedef uint32_t khint_t
     ctypedef khint_t khiter_t
 
+    ctypedef struct khcomplex128_t:
+        double real
+        double imag
+
+    bint are_equal_khcomplex128_t "kh_complex_hash_equal" (khcomplex128_t a, khcomplex128_t b) nogil
+
+    ctypedef struct khcomplex64_t:
+        float real
+        float imag
+
+    bint are_equal_khcomplex64_t "kh_complex_hash_equal" (khcomplex64_t a, khcomplex64_t b) nogil
+
     ctypedef struct kh_pymap_t:
         khint_t n_buckets, size, n_occupied, upper_bound
         uint32_t *flags
diff --git a/pandas/_libs/khash_for_primitive_helper.pxi.in b/pandas/_libs/khash_for_primitive_helper.pxi.in
@@ -17,6 +17,8 @@ primitive_types = [('int64', 'int64_t'),
                    ('uint16', 'uint16_t'),
                    ('int8', 'int8_t'),
                    ('uint8', 'uint8_t'),
+                   ('complex64', 'khcomplex64_t'),
+                   ('complex128', 'khcomplex128_t'),
                   ]
 }}
 
diff --git a/pandas/_libs/src/klib/khash_python.h b/pandas/_libs/src/klib/khash_python.h

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,8 @@ primitive_types = [('int64', 'int64_t'),`
`17`	`17`	`('uint16', 'uint16_t'),`
`18`	`18`	`('int8', 'int8_t'),`
`19`	`19`	`('uint8', 'uint8_t'),`
	`20`	`+ ('complex64', 'khcomplex64_t'),`
	`21`	`+ ('complex128', 'khcomplex128_t'),`
`20`	`22`	`]`
`21`	`23`	`}}`
`22`	`24`