implement fully checked kh_pymap interface

Author: matiaslindgren

pandas/_libs/hashtable_class_helper.pxi.in

@@ -4,8 +4,6 @@ Template for each `dtype` helper function for hashtable
 WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
 """
 from cpython.unicode cimport PyUnicode_AsUTF8
-from cpython.exc cimport PyErr_Occurred, PyErr_Fetch
-from cpython.ref cimport Py_XDECREF
 
 {{py:
 
@@ -75,7 +73,7 @@ cimported_types = ['complex64',
                    'int16',
                    'int32',
                    'int64',
-                   'pymap',
+                   'pymap_checked',
                    'str',
                    'strbox',
                    'uint8',
@@ -1311,32 +1309,16 @@ cdef class StringHashTable(HashTable):
         return labels
 
 
-cdef raise_if_errors():
-    cdef:
-        object exc
-        PyObject *type
-        PyObject *value
-        PyObject *traceback
-
-    PyErr_Fetch(&type, &value, &traceback)
-    if value != NULL:
-        exc = <object>value
-        Py_XDECREF(value)
-        Py_XDECREF(type)
-        Py_XDECREF(traceback)
-        raise exc
-
-
 cdef class PyObjectHashTable(HashTable):
 
     def __init__(self, int64_t size_hint=1):
-        self.table = kh_init_pymap()
+        self.table = kh_init_pymap_checked()
         size_hint = min(kh_needed_n_buckets(size_hint), SIZE_HINT_LIMIT)
-        kh_resize_pymap(self.table, size_hint)
+        kh_resize_pymap_checked(self.table, size_hint)
 
     def __dealloc__(self):
         if self.table is not NULL:
-            kh_destroy_pymap(self.table)
+            kh_destroy_pymap_checked(self.table)
             self.table = NULL
 
     def __len__(self) -> int:
@@ -1347,7 +1329,7 @@ cdef class PyObjectHashTable(HashTable):
             khiter_t k
         hash(key)
 
-        k = kh_get_pymap(self.table, <PyObject*>key)
+        k = kh_get_pymap_checked(self.table, <PyObject*>key)
         return k != self.table.n_buckets
 
     def sizeof(self, deep: bool = False) -> int:
@@ -1374,8 +1356,7 @@ cdef class PyObjectHashTable(HashTable):
         cdef:
             khiter_t k
 
-        k = kh_get_pymap(self.table, <PyObject*>val)
-        raise_if_errors()
+        k = kh_get_pymap_checked(self.table, <PyObject*>val)
         if k != self.table.n_buckets:
             return self.table.vals[k]
         else:
@@ -1389,9 +1370,8 @@ cdef class PyObjectHashTable(HashTable):
 
         hash(key)
 
-        k = kh_put_pymap(self.table, <PyObject*>key, &ret)
-        raise_if_errors()
-        if kh_exist_pymap(self.table, k):
+        k = kh_put_pymap_checked(self.table, <PyObject*>key, &ret)
+        if kh_exist_pymap_checked(self.table, k):
             self.table.vals[k] = val
         else:
             raise KeyError(key)
@@ -1408,8 +1388,7 @@ cdef class PyObjectHashTable(HashTable):
             val = values[i]
             hash(val)
 
-            k = kh_put_pymap(self.table, <PyObject*>val, &ret)
-            raise_if_errors()
+            k = kh_put_pymap_checked(self.table, <PyObject*>val, &ret)
             self.table.vals[k] = i
 
     def lookup(self, ndarray[object] values, object mask = None) -> ndarray:
@@ -1426,7 +1405,7 @@ cdef class PyObjectHashTable(HashTable):
             val = values[i]
             hash(val)
 
-            k = kh_get_pymap(self.table, <PyObject*>val)
+            k = kh_get_pymap_checked(self.table, <PyObject*>val)
             if k != self.table.n_buckets:
                 locs[i] = self.table.vals[k]
             else:
@@ -1504,10 +1483,10 @@ cdef class PyObjectHashTable(HashTable):
                 labels[i] = na_sentinel
                 continue
 
-            k = kh_get_pymap(self.table, <PyObject*>val)
+            k = kh_get_pymap_checked(self.table, <PyObject*>val)
             if k == self.table.n_buckets:
                 # k hasn't been seen yet
-                k = kh_put_pymap(self.table, <PyObject*>val, &ret)
+                k = kh_put_pymap_checked(self.table, <PyObject*>val, &ret)
                 uniques.append(val)
                 if return_inverse:
                     self.table.vals[k] = count

pandas/_libs/hashtable_func_helper.pxi.in

@@ -6,26 +6,24 @@ WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
 
 {{py:
 
-# name, dtype, ttype, c_type, to_c_type
-dtypes = [('Complex128', 'complex128', 'complex128',
-                         'khcomplex128_t', 'to_khcomplex128_t'),
-          ('Complex64', 'complex64', 'complex64',
-                        'khcomplex64_t', 'to_khcomplex64_t'),
-          ('Float64', 'float64', 'float64', 'float64_t', ''),
-          ('Float32', 'float32', 'float32', 'float32_t', ''),
-          ('UInt64', 'uint64', 'uint64', 'uint64_t', ''),
-          ('UInt32', 'uint32', 'uint32', 'uint32_t', ''),
-          ('UInt16', 'uint16', 'uint16', 'uint16_t', ''),
-          ('UInt8', 'uint8', 'uint8', 'uint8_t', ''),
-          ('Object', 'object', 'pymap', 'object', '<PyObject*>'),
-          ('Int64', 'int64', 'int64', 'int64_t', ''),
-          ('Int32', 'int32', 'int32', 'int32_t', ''),
-          ('Int16', 'int16', 'int16', 'int16_t', ''),
-          ('Int8', 'int8', 'int8', 'int8_t', '')]
+# name, dtype, ttype, tfunc_type, c_type, to_c_type
+dtypes = [('Complex128', 'complex128', 'complex128', 'complex128', 'khcomplex128_t', 'to_khcomplex128_t'),
+          ('Complex64', 'complex64', 'complex64', 'complex64', 'khcomplex64_t', 'to_khcomplex64_t'),
+          ('Float64', 'float64', 'float64', 'float64', 'float64_t', ''),
+          ('Float32', 'float32', 'float32', 'float32', 'float32_t', ''),
+          ('UInt64', 'uint64', 'uint64', 'uint64', 'uint64_t', ''),
+          ('UInt32', 'uint32', 'uint32', 'uint32', 'uint32_t', ''),
+          ('UInt16', 'uint16', 'uint16', 'uint16', 'uint16_t', ''),
+          ('UInt8', 'uint8', 'uint8', 'uint8', 'uint8_t', ''),
+          ('Object', 'object', 'pymap', 'pymap_checked', 'object', '<PyObject*>'),
+          ('Int64', 'int64', 'int64', 'int64', 'int64_t', ''),
+          ('Int32', 'int32', 'int32', 'int32', 'int32_t', ''),
+          ('Int16', 'int16', 'int16', 'int16', 'int16_t', ''),
+          ('Int8', 'int8', 'int8', 'int8', 'int8_t', '')]
 
 }}
 
-{{for name, dtype, ttype, c_type, to_c_type in dtypes}}
+{{for name, dtype, ttype, tfunc_type, c_type, to_c_type in dtypes}}
 
 
 @cython.wraparound(False)
@@ -55,26 +53,26 @@ cdef value_count_{{dtype}}(const {{dtype}}_t[:] values, bint dropna, const uint8
     #    result_keys remembers the original order of keys
 
     result_keys = {{name}}Vector()
-    table = kh_init_{{ttype}}()
+    table = kh_init_{{tfunc_type}}()
 
     {{if dtype == 'object'}}
     if uses_mask:
         raise NotImplementedError("uses_mask not implemented with object dtype")
 
-    kh_resize_{{ttype}}(table, n // 10)
+    kh_resize_{{tfunc_type}}(table, n // 10)
 
     for i in range(n):
         val = values[i]
         if not dropna or not checknull(val):
-            k = kh_get_{{ttype}}(table, {{to_c_type}}val)
+            k = kh_get_{{tfunc_type}}(table, {{to_c_type}}val)
             if k != table.n_buckets:
                 table.vals[k] += 1
             else:
-                k = kh_put_{{ttype}}(table, {{to_c_type}}val, &ret)
+                k = kh_put_{{tfunc_type}}(table, {{to_c_type}}val, &ret)
                 table.vals[k] = 1
                 result_keys.append(val)
     {{else}}
-    kh_resize_{{ttype}}(table, n)
+    kh_resize_{{tfunc_type}}(table, n)
 
     for i in range(n):
         val = {{to_c_type}}(values[i])
@@ -90,11 +88,11 @@ cdef value_count_{{dtype}}(const {{dtype}}_t[:] values, bint dropna, const uint8
             if uses_mask and isna_entry:
                 na_counter += 1
             else:
-                k = kh_get_{{ttype}}(table, val)
+                k = kh_get_{{tfunc_type}}(table, val)
                 if k != table.n_buckets:
                     table.vals[k] += 1
                 else:
-                    k = kh_put_{{ttype}}(table, val, &ret)
+                    k = kh_put_{{tfunc_type}}(table, val, &ret)
                     table.vals[k] = 1
                     result_keys.append(val)
     {{endif}}
@@ -107,17 +105,17 @@ cdef value_count_{{dtype}}(const {{dtype}}_t[:] values, bint dropna, const uint8
 
     for i in range(table.size):
         {{if dtype == 'object'}}
-        k = kh_get_{{ttype}}(table, result_keys.data[i])
+        k = kh_get_{{tfunc_type}}(table, result_keys.data[i])
         {{else}}
-        k = kh_get_{{ttype}}(table, result_keys.data.data[i])
+        k = kh_get_{{tfunc_type}}(table, result_keys.data.data[i])
         {{endif}}
         result_counts[i] = table.vals[k]
 
     if na_counter > 0:
         result_counts[table.size] = na_counter
         result_keys.append(val)
 
-    kh_destroy_{{ttype}}(table)
+    kh_destroy_{{tfunc_type}}(table)
 
     return result_keys.to_array(), result_counts.base, na_counter
 
@@ -138,12 +136,12 @@ cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first', cons
         {{endif}}
         Py_ssize_t i, n = len(values), first_na = -1
         khiter_t k
-        kh_{{ttype}}_t *table = kh_init_{{ttype}}()
+        kh_{{ttype}}_t *table = kh_init_{{tfunc_type}}()
         ndarray[uint8_t, ndim=1, cast=True] out = np.empty(n, dtype='bool')
         bint seen_na = False, uses_mask = mask is not None
         bint seen_multiple_na = False
 
-    kh_resize_{{ttype}}(table, min(kh_needed_n_buckets(n), SIZE_HINT_LIMIT))
+    kh_resize_{{tfunc_type}}(table, min(kh_needed_n_buckets(n), SIZE_HINT_LIMIT))
 
     if keep not in ('last', 'first', False):
         raise ValueError('keep must be either "first", "last" or False')
@@ -168,7 +166,7 @@ cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first', cons
                         seen_na = True
                 else:
                     value = {{to_c_type}}(values[i])
-                    kh_put_{{ttype}}(table, value, &ret)
+                    kh_put_{{tfunc_type}}(table, value, &ret)
                     out[i] = ret == 0
     {{endfor}}
 
@@ -193,16 +191,16 @@ cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first', cons
 
                 else:
                     value = {{to_c_type}}(values[i])
-                    k = kh_get_{{ttype}}(table, value)
+                    k = kh_get_{{tfunc_type}}(table, value)
                     if k != table.n_buckets:
                         out[table.vals[k]] = 1
                         out[i] = 1
                     else:
-                        k = kh_put_{{ttype}}(table, value, &ret)
+                        k = kh_put_{{tfunc_type}}(table, value, &ret)
                         table.vals[k] = i
                         out[i] = 0
 
-    kh_destroy_{{ttype}}(table)
+    kh_destroy_{{tfunc_type}}(table)
     return out
 
 
@@ -243,11 +241,11 @@ cdef ismember_{{dtype}}(const {{dtype}}_t[:] arr, const {{dtype}}_t[:] values):
         {{c_type}} val
         {{endif}}
 
-        kh_{{ttype}}_t *table = kh_init_{{ttype}}()
+        kh_{{ttype}}_t *table = kh_init_{{tfunc_type}}()
 
     # construct the table
     n = len(values)
-    kh_resize_{{ttype}}(table, n)
+    kh_resize_{{tfunc_type}}(table, n)
 
     {{if dtype == 'object'}}
     if True:
@@ -256,7 +254,7 @@ cdef ismember_{{dtype}}(const {{dtype}}_t[:] arr, const {{dtype}}_t[:] values):
     {{endif}}
         for i in range(n):
             val = {{to_c_type}}(values[i])
-            kh_put_{{ttype}}(table, val, &ret)
+            kh_put_{{tfunc_type}}(table, val, &ret)
 
     # test membership
     n = len(arr)
@@ -269,10 +267,10 @@ cdef ismember_{{dtype}}(const {{dtype}}_t[:] arr, const {{dtype}}_t[:] values):
     {{endif}}
         for i in range(n):
             val = {{to_c_type}}(arr[i])
-            k = kh_get_{{ttype}}(table, val)
+            k = kh_get_{{tfunc_type}}(table, val)
             result[i] = (k != table.n_buckets)
 
-    kh_destroy_{{ttype}}(table)
+    kh_destroy_{{tfunc_type}}(table)
     return result.view(np.bool_)
 
 # ----------------------------------------------------------------------

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

@@ -218,6 +218,9 @@ static inline int pyobject_cmp(PyObject *a, PyObject *b) {
 
   int result = PyObject_RichCompareBool(a, b, Py_EQ);
   if (result < 0) {
+    if (PyErr_Occurred() != NULL) {
+      return 0;
+    }
     return 0;
   }
   return result;
@@ -320,6 +323,9 @@ static inline khuint32_t kh_python_hash_func(PyObject *key) {
   }
 
   if (hash == -1) {
+    if (PyErr_Occurred() != NULL) {
+      return 0;
+    }
     return 0;
   }
 #if SIZEOF_PY_HASH_T == 4

pandas/_libs/khash.pxd

@@ -125,5 +125,13 @@ cdef extern from "pandas/vendored/klib/khash_python.h":
 
     khuint_t kh_needed_n_buckets(khuint_t element_n) nogil
 
+cdef kh_pymap_t* kh_init_pymap_checked()
+cdef void kh_destroy_pymap_checked(kh_pymap_t*)
+cdef void kh_clear_pymap_checked(kh_pymap_t*)
+cdef khuint_t kh_get_pymap_checked(kh_pymap_t*, PyObject*)
+cdef void kh_resize_pymap_checked(kh_pymap_t*, khuint_t)
+cdef khuint_t kh_put_pymap_checked(kh_pymap_t*, PyObject*, int*)
+cdef void kh_del_pymap_checked(kh_pymap_t*, khuint_t)
+cdef bint kh_exist_pymap_checked(kh_pymap_t*, khiter_t)
 
 include "khash_for_primitive_helper.pxi"