rewrite hash fix, add tests

Author: matiaslindgren

pandas/_libs/hashtable_class_helper.pxi.in

@@ -4,6 +4,8 @@ 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:
 
@@ -1309,6 +1311,22 @@ 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):
@@ -1356,15 +1374,8 @@ cdef class PyObjectHashTable(HashTable):
         cdef:
             khiter_t k
 
-        # GH 57052
-        # in khash_python.h, kh_python_hash_equal and kh_python_hash_func will be called repeatedly by khash in a loop.
-        # if object implements custom __hash__ and __eq__ methods that can raise exceptions,
-        # kh_python_hash_{equal,func} will suppress the exceptions without warnings.
-        # as a workaround: try triggering exceptions here, before starting the khash loop
-        hash(val)
-        val == val
-
         k = kh_get_pymap(self.table, <PyObject*>val)
+        raise_if_errors()
         if k != self.table.n_buckets:
             return self.table.vals[k]
         else:
@@ -1377,10 +1388,9 @@ cdef class PyObjectHashTable(HashTable):
             char* buf
 
         hash(key)
-        # GH 57052
-        key == key
 
         k = kh_put_pymap(self.table, <PyObject*>key, &ret)
+        raise_if_errors()
         if kh_exist_pymap(self.table, k):
             self.table.vals[k] = val
         else:

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

@@ -193,6 +193,9 @@ static inline int tupleobject_cmp(PyTupleObject *a, PyTupleObject *b) {
 }
 
 static inline int pyobject_cmp(PyObject *a, PyObject *b) {
+  if (PyErr_Occurred() != NULL) {
+    return 0;
+  }
   if (a == b) {
     return 1;
   }
@@ -215,7 +218,6 @@ static inline int pyobject_cmp(PyObject *a, PyObject *b) {
 
   int result = PyObject_RichCompareBool(a, b, Py_EQ);
   if (result < 0) {
-    PyErr_Clear();
     return 0;
   }
   return result;
@@ -292,6 +294,9 @@ static inline Py_hash_t tupleobject_hash(PyTupleObject *key) {
 }
 
 static inline khuint32_t kh_python_hash_func(PyObject *key) {
+  if (PyErr_Occurred() != NULL) {
+    return 0;
+  }
   Py_hash_t hash;
   // For PyObject_Hash holds:
   //    hash(0.0) == 0 == hash(-0.0)
@@ -315,7 +320,6 @@ static inline khuint32_t kh_python_hash_func(PyObject *key) {
   }
 
   if (hash == -1) {
-    PyErr_Clear();
     return 0;
   }
 #if SIZEOF_PY_HASH_T == 4

pandas/tests/frame/indexing/test_indexing.py

@@ -1943,3 +1943,27 @@ def test_setitem_validation_scalar_int(self, invalid, any_int_numpy_dtype, index
     def test_setitem_validation_scalar_float(self, invalid, float_numpy_dtype, indexer):
         df = DataFrame({"a": [1, 2, None]}, dtype=float_numpy_dtype)
         self._check_setitem_invalid(df, invalid, indexer)
+
+
+def test_error_raised_from_custom_hash_method():
+    # GH 57052
+    class testkey:
+        def __init__(self, value):
+            self.value = value
+
+        def __hash__(self):
+            raise RuntimeError(f"exception in {self!r}.__hash__")
+
+        def __eq__(self, other):
+            return self.value == other.value
+
+        def __repr__(self):
+            return f"testkey({self.value})"
+
+    df = DataFrame({"i": map(testkey, range(10))}).set_index("i")
+    for i in range(len(df.index)):
+        key = testkey(i)
+        with pytest.raises(
+            RuntimeError, match=re.escape(f"exception in {key!r}.__hash__")
+        ):
+            df.loc[key]

pandas/tests/libs/test_hashtable.py

@@ -783,33 +783,57 @@ def test_float_complex_int_are_equal_as_objects():
     tm.assert_numpy_array_equal(result, expected)
 
 
-@pytest.mark.parametrize(
-    "throw1hash, throw2hash, throw1eq, throw2eq",
-    product([True, False], repeat=4),
-)
-def test_exceptions_thrown_from_custom_hash_and_eq_methods(
-    throw1hash, throw2hash, throw1eq, throw2eq
-):
+class testkey:
     # GH 57052
-    class testkey:
-        def __init__(self, value, throw_hash=False, throw_eq=False):
-            self.value = value
-            self.throw_hash = throw_hash
-            self.throw_eq = throw_eq
+    def __init__(self, value, throw_hash=False, throw_eq=False):
+        self.value = value
+        self.throw_hash = throw_hash
+        self.throw_eq = throw_eq
+
+    def __hash__(self):
+        if self.throw_hash:
+            raise RuntimeError(f"exception in {self!r}.__hash__")
+        return hash(self.value)
+
+    def __eq__(self, other):
+        if self.throw_eq:
+            raise RuntimeError(f"exception in {self!r}.__eq__")
+        return self.value == other.value
 
-        def __hash__(self):
-            if self.throw_hash:
-                raise RuntimeError(f"exception in {self!r}.__hash__")
-            return hash(self.value)
+    def __repr__(self):
+        return f"testkey({self.value}, {self.throw_hash}, {self.throw_eq})"
+
+
+@pytest.mark.parametrize("throw1, throw2", product([True, False], repeat=2))
+def test_error_raised_from_hash_method_in_set_item(throw1, throw2):
+    # GH 57052
+    table = ht.PyObjectHashTable()
+
+    key1 = testkey(value="hello1", throw_hash=throw1)
+    key2 = testkey(value="hello2", throw_hash=throw2)
+
+    if throw1:
+        with pytest.raises(
+            RuntimeError, match=re.escape(f"exception in {key1!r}.__hash__")
+        ):
+            table.set_item(key1, 123)
+    else:
+        table.set_item(key1, 123)
+        assert table.get_item(key1) == 123
 
-        def __eq__(self, other):
-            if self.throw_eq:
-                raise RuntimeError(f"exception in {self!r}.__eq__")
-            return self.value == other.value
+    if throw2:
+        with pytest.raises(
+            RuntimeError, match=re.escape(f"exception in {key2!r}.__hash__")
+        ):
+            table.set_item(key2, 456)
+    else:
+        table.set_item(key2, 456)
+        assert table.get_item(key2) == 456
 
-        def __repr__(self):
-            return f"testkey({self.value}, {self.throw_hash}, {self.throw_eq})"
 
+@pytest.mark.parametrize("throw1, throw2", product([True, False], repeat=2))
+def test_error_raised_from_hash_method_in_get_item(throw1, throw2):
+    # GH 57052
     table = ht.PyObjectHashTable()
 
     key1 = testkey(value="hello1")
@@ -818,43 +842,63 @@ def __repr__(self):
     table.set_item(key1, 123)
     table.set_item(key2, 456)
 
-    key1.throw_hash = throw1hash
-    key2.throw_hash = throw2hash
-    key1.throw_eq = throw1eq
-    key2.throw_eq = throw2eq
+    key1.throw_hash = throw1
+    key2.throw_hash = throw2
 
-    if throw1hash and throw1eq:
-        with pytest.raises(
-            RuntimeError, match=re.escape(f"exception in {key1!r}.") + "__(hash|eq)__"
-        ):
-            table.get_item(key1)
-    elif throw1hash:
+    if throw1:
         with pytest.raises(
             RuntimeError, match=re.escape(f"exception in {key1!r}.__hash__")
         ):
             table.get_item(key1)
-    elif throw1eq:
-        with pytest.raises(
-            RuntimeError, match=re.escape(f"exception in {key1!r}.__eq__")
-        ):
-            table.get_item(key1)
     else:
         assert table.get_item(key1) == 123
 
-    if throw2hash and throw2eq:
+    if throw2:
         with pytest.raises(
-            RuntimeError, match=re.escape(f"exception in {key2!r}.") + "__(hash|eq)__"
+            RuntimeError, match=re.escape(f"exception in {key2!r}.__hash__")
         ):
             table.get_item(key2)
-    elif throw2hash:
+    else:
+        assert table.get_item(key2) == 456
+
+
+@pytest.mark.parametrize("throw", [True, False])
+def test_error_raised_from_eq_method_in_set_item(throw):
+    # GH 57052
+    table = ht.PyObjectHashTable()
+
+    key1 = testkey(value="hello", throw_eq=throw)
+    key2 = testkey(value=key1.value)
+
+    if throw:
+        table.set_item(key1, 123)
         with pytest.raises(
-            RuntimeError, match=re.escape(f"exception in {key2!r}.__hash__")
+            RuntimeError, match=re.escape(f"exception in {key1!r}.__eq__")
         ):
-            table.get_item(key2)
-    elif throw2eq:
+            table.set_item(key2, 456)
+    else:
+        table.set_item(key2, 456)
+        assert table.get_item(key2) == 456
+
+
+@pytest.mark.parametrize("throw", [True, False])
+def test_error_raised_from_eq_method_in_get_item(throw):
+    # GH 57052
+    table = ht.PyObjectHashTable()
+
+    key1 = testkey(value="hello")
+    key2 = testkey(value=key1.value)
+
+    table.set_item(key1, 123)
+    table.set_item(key2, 456)
+
+    if throw:
+        key1.throw_eq = True
         with pytest.raises(
-            RuntimeError, match=re.escape(f"exception in {key2!r}.__eq__")
+            RuntimeError, match=re.escape(f"exception in {key1!r}.__eq__")
         ):
             table.get_item(key2)
     else:
+        # this looks odd but it is because key1.value == key2.value
+        assert table.get_item(key1) == 456
         assert table.get_item(key2) == 456