WIP: adding promoter for equals and not_equals ufuncs

ngoldbaum · ngoldbaum · commit 1ce658bb2d16 · 2023-06-01T13:56:24.000-06:00
diff --git a/stringdtype/stringdtype/src/umath.c b/stringdtype/stringdtype/src/umath.c
@@ -164,8 +164,47 @@ string_equal_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
     return 0;
 }
 
+static int
+string_not_equal_strided_loop(PyArrayMethod_Context *NPY_UNUSED(context),
+                              char *const data[], npy_intp const dimensions[],
+                              npy_intp const strides[],
+                              NpyAuxData *NPY_UNUSED(auxdata))
+{
+    npy_intp N = dimensions[0];
+    char *in1 = data[0];
+    char *in2 = data[1];
+    npy_bool *out = (npy_bool *)data[2];
+    npy_intp in1_stride = strides[0];
+    npy_intp in2_stride = strides[1];
+    npy_intp out_stride = strides[2];
+
+    ss *s1 = NULL, *s2 = NULL;
+
+    while (N--) {
+        s1 = (ss *)in1;
+        s2 = (ss *)in2;
+        if (ss_isnull(s1) || ss_isnull(s2)) {
+            // s1 or s2 is NA
+            *out = (npy_bool)0;
+        }
+        else if (s1->len == s2->len &&
+                 strncmp(s1->buf, s2->buf, s1->len) == 0) {
+            *out = (npy_bool)0;
+        }
+        else {
+            *out = (npy_bool)1;
+        }
+
+        in1 += in1_stride;
+        in2 += in2_stride;
+        out += out_stride;
+    }
+
+    return 0;
+}
+
 static NPY_CASTING
-string_equal_resolve_descriptors(
+string_comparison_resolve_descriptors(
         struct PyArrayMethodObject_tag *NPY_UNUSED(method),
         PyArray_DTypeMeta *NPY_UNUSED(dtypes[]), PyArray_Descr *given_descrs[],
         PyArray_Descr *loop_descrs[], npy_intp *NPY_UNUSED(view_offset))
@@ -227,9 +266,10 @@ string_isnan_resolve_descriptors(
  * Copied from NumPy, because NumPy doesn't always use it :)
  */
 static int
-default_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
-                       PyArray_DTypeMeta *signature[],
-                       PyArray_DTypeMeta *new_op_dtypes[])
+ufunc_promoter_internal(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
+                        PyArray_DTypeMeta *signature[],
+                        PyArray_DTypeMeta *new_op_dtypes[],
+                        PyArray_DTypeMeta *final_dtype)
 {
     /* If nin < 2 promotion is a no-op, so it should not be registered */
     assert(ufunc->nin > 1);
@@ -261,19 +301,11 @@ default_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
             }
         }
     }
-    /* Otherwise, use the common DType of all input operands */
-    if (common == NULL) {
-        common = PyArray_PromoteDTypeSequence(ufunc->nin, op_dtypes);
-        if (common == NULL) {
-            if (PyErr_ExceptionMatches(PyExc_TypeError)) {
-                PyErr_Clear(); /* Do not propagate normal promotion errors */
-            }
-            return -1;
-        }
-    }
+    Py_XDECREF(common);
 
+    /* Otherwise, set all input operands to StringDType */
     for (int i = 0; i < ufunc->nargs; i++) {
-        PyArray_DTypeMeta *tmp = common;
+        PyArray_DTypeMeta *tmp = final_dtype;
         if (signature[i]) {
             tmp = signature[i]; /* never replace a fixed one. */
         }
@@ -285,10 +317,27 @@ default_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
         new_op_dtypes[i] = op_dtypes[i];
     }
 
-    Py_DECREF(common);
     return 0;
 }
 
+static int
+string_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
+                      PyArray_DTypeMeta *signature[],
+                      PyArray_DTypeMeta *new_op_dtypes[])
+{
+    return ufunc_promoter_internal(ufunc, op_dtypes, signature, new_op_dtypes,
+                                   (PyArray_DTypeMeta *)&StringDType);
+}
+
+static int
+pandas_string_ufunc_promoter(PyUFuncObject *ufunc, PyArray_DTypeMeta *op_dtypes[],
+                             PyArray_DTypeMeta *signature[],
+                             PyArray_DTypeMeta *new_op_dtypes[])
+{
+    return ufunc_promoter_internal(ufunc, op_dtypes, signature, new_op_dtypes,
+                                   (PyArray_DTypeMeta *)&PandasStringDType);
+}
+
 // Register a ufunc.
 //
 // Pass NULL for resolve_func to use the default_resolve_descriptors.
@@ -334,23 +383,33 @@ init_ufunc(PyObject *numpy, const char *ufunc_name, PyArray_DTypeMeta **dtypes,
 
 int
 add_promoter(PyObject *numpy, const char *ufunc_name,
-             PyArray_DTypeMeta **dtypes)
+             PyArray_DTypeMeta *ldtype, PyArray_DTypeMeta *rdtype,
+             PyArray_DTypeMeta *edtype, int is_pandas)
 {
     PyObject *ufunc = PyObject_GetAttrString(numpy, ufunc_name);
 
     if (ufunc == NULL) {
         return -1;
     }
 
-    PyObject *DType_tuple = PyTuple_Pack(3, dtypes[0], dtypes[1], dtypes[2]);
+    PyObject *DType_tuple = PyTuple_Pack(3, ldtype, rdtype, edtype);
 
     if (DType_tuple == NULL) {
         Py_DECREF(ufunc);
         return -1;
     }
 
-    PyObject *promoter_capsule = PyCapsule_New((void *)&default_ufunc_promoter,
-                                               "numpy._ufunc_promoter", NULL);
+    PyObject *promoter_capsule = NULL;
+    
+    if (is_pandas == 0) {
+        promoter_capsule = PyCapsule_New((void *)&string_ufunc_promoter,
+                                          "numpy._ufunc_promoter", NULL);
+    }
+    else {
+        promoter_capsule = PyCapsule_New((void *)&pandas_string_ufunc_promoter,
+                                         "numpy._ufunc_promoter", NULL);
+    }
+            
 
     if (promoter_capsule == NULL) {
         Py_DECREF(ufunc);
@@ -380,30 +439,46 @@ init_ufuncs(void)
         return -1;
     }
 
-    PyArray_DTypeMeta *eq_dtypes[] = {(PyArray_DTypeMeta *)&StringDType,
-                                      (PyArray_DTypeMeta *)&StringDType,
-                                      &PyArray_BoolDType};
+    PyArray_DTypeMeta *comparison_dtypes[] = {(PyArray_DTypeMeta *)&StringDType,
+                                              (PyArray_DTypeMeta *)&StringDType,
+                                              &PyArray_BoolDType};
 
-    if (init_ufunc(numpy, "equal", eq_dtypes,
-                   &string_equal_resolve_descriptors,
+    if (init_ufunc(numpy, "equal", comparison_dtypes,
+                   &string_comparison_resolve_descriptors,
                    &string_equal_strided_loop, "string_equal", 2, 1,
                    NPY_NO_CASTING, 0) < 0) {
         goto error;
     }
 
-    PyArray_DTypeMeta *promoter_dtypes[2][3] = {
-            {(PyArray_DTypeMeta *)&StringDType, &PyArray_UnicodeDType,
-             &PyArray_BoolDType},
-            {&PyArray_UnicodeDType, (PyArray_DTypeMeta *)&StringDType,
-             &PyArray_BoolDType},
-    };
-
-    if (add_promoter(numpy, "equal", promoter_dtypes[0]) < 0) {
+    if (init_ufunc(numpy, "not_equal", comparison_dtypes,
+                   &string_comparison_resolve_descriptors,
+                   &string_not_equal_strided_loop, "string_not_equal", 2, 1,
+                   NPY_NO_CASTING, 0) < 0) {
         goto error;
     }
 
-    if (add_promoter(numpy, "equal", promoter_dtypes[1]) < 0) {
-        goto error;
+    char *ufunc_names[2] = {"equal", "not_equal"};
+
+    for (int i = 0; i < 2; i++) {
+        if (add_promoter(numpy, ufunc_names[i], (PyArray_DTypeMeta *)&StringDType,
+                         &PyArray_UnicodeDType, &PyArray_BoolDType, 0) < 0) {
+            goto error;
+        }
+
+        if (add_promoter(numpy, ufunc_names[i], &PyArray_UnicodeDType,
+                         (PyArray_DTypeMeta *)&StringDType, &PyArray_BoolDType, 0) < 0) {
+            goto error;
+        }
+
+        if (add_promoter(numpy, ufunc_names[i], &PyArray_ObjectDType,
+                         (PyArray_DTypeMeta *)&StringDType, &PyArray_BoolDType, 0) < 0) {
+            goto error;
+        }
+
+        if (add_promoter(numpy, ufunc_names[i], (PyArray_DTypeMeta *)&StringDType,
+                         &PyArray_ObjectDType, &PyArray_BoolDType, 0) < 0) {
+            goto error;
+        }
     }
 
     PyArray_DTypeMeta *isnan_dtypes[] = {(PyArray_DTypeMeta *)&StringDType,
@@ -448,30 +523,45 @@ init_ufuncs(void)
         goto finish;
     }
 
-    PyArray_DTypeMeta *peq_dtypes[] = {(PyArray_DTypeMeta *)&PandasStringDType,
-                                       (PyArray_DTypeMeta *)&PandasStringDType,
-                                       &PyArray_BoolDType};
+    PyArray_DTypeMeta *p_comparison_dtypes[] =
+            {(PyArray_DTypeMeta *)&PandasStringDType,
+             (PyArray_DTypeMeta *)&PandasStringDType,
+             &PyArray_BoolDType};
 
-    if (init_ufunc(numpy, "equal", peq_dtypes,
-                   &string_equal_resolve_descriptors,
+    if (init_ufunc(numpy, "equal", p_comparison_dtypes,
+                   &string_comparison_resolve_descriptors,
                    &string_equal_strided_loop, "string_equal", 2, 1,
                    NPY_NO_CASTING, 0) < 0) {
         goto error;
     }
 
-    PyArray_DTypeMeta *p_promoter_dtypes[2][3] = {
-            {(PyArray_DTypeMeta *)&PandasStringDType, &PyArray_UnicodeDType,
-             &PyArray_BoolDType},
-            {&PyArray_UnicodeDType, (PyArray_DTypeMeta *)&PandasStringDType,
-             &PyArray_BoolDType},
-    };
-
-    if (add_promoter(numpy, "equal", p_promoter_dtypes[0]) < 0) {
+    if (init_ufunc(numpy, "not_equal", p_comparison_dtypes,
+                   &string_comparison_resolve_descriptors,
+                   &string_not_equal_strided_loop, "string_not_equal", 2, 1,
+                   NPY_NO_CASTING, 0) < 0) {
         goto error;
     }
 
-    if (add_promoter(numpy, "equal", p_promoter_dtypes[1]) < 0) {
-        goto error;
+    for (int i = 0; i < 2; i++) {
+        if (add_promoter(numpy, ufunc_names[i], (PyArray_DTypeMeta *)&PandasStringDType,
+                         &PyArray_UnicodeDType, &PyArray_BoolDType, 1) < 0) {
+            goto error;
+        }
+
+        if (add_promoter(numpy, ufunc_names[i], &PyArray_UnicodeDType,
+                         (PyArray_DTypeMeta *)&PandasStringDType, &PyArray_BoolDType, 1) < 0) {
+            goto error;
+        }
+
+        if (add_promoter(numpy, ufunc_names[i], &PyArray_ObjectDType,
+                         (PyArray_DTypeMeta *)&PandasStringDType, &PyArray_BoolDType, 1) < 0) {
+            goto error;
+        }
+
+        if (add_promoter(numpy, ufunc_names[i], (PyArray_DTypeMeta *)&PandasStringDType,
+                         &PyArray_ObjectDType, &PyArray_BoolDType, 1) < 0) {
+            goto error;
+        }
     }
 
     PyArray_DTypeMeta *p_isnan_dtypes[] = {
diff --git a/stringdtype/tests/test_stringdtype.py b/stringdtype/tests/test_stringdtype.py
@@ -142,12 +142,15 @@ def test_insert_scalar(dtype, scalar, string_list):
         )
 
 
-def test_equality_promotion(dtype, string_list):
+@pytest.mark.parametrize("o_dtype", [np.str_, object])
+def test_equality_promotion(string_list, dtype, o_dtype):
     sarr = np.array(string_list, dtype=dtype)
-    uarr = np.array(string_list, dtype=np.str_)
+    oarr = np.array(string_list, dtype=o_dtype)
 
-    np.testing.assert_array_equal(sarr, uarr)
-    np.testing.assert_array_equal(uarr, sarr)
+    np.testing.assert_array_equal(sarr, oarr)
+    np.testing.assert_array_equal(oarr, sarr)
+    assert not np.any(sarr != oarr)
+    assert not np.any(oarr != sarr)
 
 
 def test_isnan(dtype, string_list):
