fix complex case

Author: realead

pandas/_libs/src/klib/khash_python.h

@@ -163,17 +163,50 @@ KHASH_MAP_INIT_COMPLEX128(complex128, size_t)
 #define kh_exist_complex128(h, k) (kh_exist(h, k))
 
 
+int PANDAS_INLINE floatobject_cmp(PyObject* a, PyObject* b){
+    return Py_IS_NAN(PyFloat_AS_DOUBLE(a)) &&
+           Py_IS_NAN(PyFloat_AS_DOUBLE(b));
+}
+
+
+int PANDAS_INLINE complexobject_cmp(PyObject* a, PyObject* b){
+    return (
+                Py_IS_NAN(PyComplex_RealAsDouble(a)) &&
+                Py_IS_NAN(PyComplex_RealAsDouble(b)) &&
+                Py_IS_NAN(PyComplex_ImagAsDouble(a)) &&
+                Py_IS_NAN(PyComplex_ImagAsDouble(b))
+           )
+           ||
+           (
+                Py_IS_NAN(PyComplex_RealAsDouble(a)) &&
+                Py_IS_NAN(PyComplex_RealAsDouble(b)) &&
+                PyComplex_ImagAsDouble(a) == PyComplex_ImagAsDouble(b)
+           )
+           ||
+           (
+                PyComplex_RealAsDouble(a) == PyComplex_RealAsDouble(b) &&
+                Py_IS_NAN(PyComplex_ImagAsDouble(a)) &&
+                Py_IS_NAN(PyComplex_ImagAsDouble(b))
+           );
+}
+
+
 int PANDAS_INLINE pyobject_cmp(PyObject* a, PyObject* b) {
 	int result = PyObject_RichCompareBool(a, b, Py_EQ);
 	if (result < 0) {
 		PyErr_Clear();
 		return 0;
 	}
-    if (result == 0) {  // still could be two NaNs
-        return PyFloat_CheckExact(a) &&
-               PyFloat_CheckExact(b) &&
-               Py_IS_NAN(PyFloat_AS_DOUBLE(a)) &&
-               Py_IS_NAN(PyFloat_AS_DOUBLE(b));
+    if (result == 0) {  // still could be built-ins with NaNs
+        if (Py_TYPE(a) != Py_TYPE(b)) {
+            return 0;
+        }
+        if (PyFloat_CheckExact(a)) {
+            return floatobject_cmp(a, b);
+        }
+        if (PyComplex_CheckExact(a)) {
+            return complexobject_cmp(a, b);
+        }
     }
 	return result;
 }

pandas/tests/libs/test_hashtable.py

@@ -187,14 +187,38 @@ def test_nan_float(self):
         table.set_item(nan1, 42)
         assert table.get_item(nan2) == 42
 
-    def test_nan_complex(self):
-        nan1 = 1j * float("nan")
-        nan2 = 1j * float("nan")
+    def test_nan_complex_both(self):
+        nan1 = complex(float("nan"), float("nan"))
+        nan2 = complex(float("nan"), float("nan"))
         assert nan1 is not nan2
         table = ht.PyObjectHashTable()
         table.set_item(nan1, 42)
         assert table.get_item(nan2) == 42
 
+    def test_nan_complex_real(self):
+        nan1 = complex(float("nan"), 1)
+        nan2 = complex(float("nan"), 1)
+        other = complex(float("nan"), 2)
+        assert nan1 is not nan2
+        table = ht.PyObjectHashTable()
+        table.set_item(nan1, 42)
+        assert table.get_item(nan2) == 42
+        with pytest.raises(KeyError, match=None) as error:
+            table.get_item(other)
+        assert str(error.value) == str(other)
+
+    def test_nan_complex_imag(self):
+        nan1 = complex(1, float("nan"))
+        nan2 = complex(1, float("nan"))
+        other = complex(2, float("nan"))
+        assert nan1 is not nan2
+        table = ht.PyObjectHashTable()
+        table.set_item(nan1, 42)
+        assert table.get_item(nan2) == 42
+        with pytest.raises(KeyError, match=None) as error:
+            table.get_item(other)
+        assert str(error.value) == str(other)
+
 
 def test_get_labels_groupby_for_Int64(writable):
     table = ht.Int64HashTable()