ENH: introduce 'compatible' stringdtype instances

ngoldbaum · ngoldbaum · commit 9817861a6724 · 2024-04-12T14:30:05.000-06:00
diff --git a/doc/neps/nep-0055-string_dtype.rst b/doc/neps/nep-0055-string_dtype.rst
@@ -534,11 +534,33 @@ future NumPy or a downstream library may add locale-aware sorting, case folding,
 and normalization for NumPy unicode strings arrays, but we are not proposing
 adding these features at this time.
 
-Two ``StringDType`` instances are considered identical if they are created with
-the same ``na_object`` and ``coerce`` parameter. We propose checking for unequal
-``StringDType`` instances in the ``resolve_descriptors`` function of binary
-ufuncs that take two string arrays and raising an error if an operation is
-performed with unequal ``StringDType`` instances.
+Two ``StringDType`` instances are considered equal if they are created with the
+same ``na_object`` and ``coerce`` parameter. For ufuncs that accept more than
+one string argument we also introduce the concept of "compatible"
+``StringDType`` instances. We allow distinct DType instances to be used in ufunc
+operations together if have the same ``na_object`` or if only one
+or the other DType has an ``na_object`` explicitly set. We do not consider
+string coercion for determining whether instances are compatible, although if
+the result of the operation is a string, the result will inherit the stricter
+string coercion setting of the original operands.
+
+This notion of "compatible" instances will be enforced in the
+``resolve_descriptors`` function of binary ufuncs. This choice makes it easier
+to work with non-default ``StringDType`` instances, because python strings are
+coerced to the default ``StringDType`` instance, so the following idiomatic
+expression is allowed::
+
+  >>> arr = np.array(["hello", "world"], dtype=StringDType(na_object=None))
+  >>> arr + "!"
+  array(['hello!', 'world!'], dtype=StringDType(na_object=None))
+
+If we only considered equality of ``StringDType`` instances, this would
+be an error, making for an awkward user experience. If the operands have
+distinct ``na_object`` settings, NumPy will raise an error because the choice
+for the result DType is ambiguous::
+
+  >>> arr + np.array("!", dtype=StringDType(na_object=""))
+  TypeError: Cannot find common instance for incompatible dtype instances
 
 ``np.strings`` namespace
 ************************
diff --git a/numpy/_core/src/multiarray/stringdtype/dtype.c b/numpy/_core/src/multiarray/stringdtype/dtype.c
@@ -144,14 +144,21 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
     return NULL;
 }
 
-static int
-na_eq_cmp(PyObject *a, PyObject *b) {
+NPY_NO_EXPORT int
+na_eq_cmp(PyObject *a, PyObject *b, int coerce_nulls) {
     if (a == b) {
         // catches None and other singletons like Pandas.NA
         return 1;
     }
     if (a == NULL || b == NULL) {
-        return 0;
+        if (coerce_nulls) {
+            // an object with an explictly set NA object is considered
+            // compatible for binary operations to one with no explicitly set NA
+            return 1;
+        }
+        else {
+            return 0;
+        }
     }
     if (PyFloat_Check(a) && PyFloat_Check(b)) {
         // nan check catches np.nan and float('nan')
@@ -182,29 +189,29 @@ _eq_comparison(int scoerce, int ocoerce, PyObject *sna, PyObject *ona)
     if (scoerce != ocoerce) {
         return 0;
     }
-    return na_eq_cmp(sna, ona);
+    return na_eq_cmp(sna, ona, 0);
 }
 
 /*
  * This is used to determine the correct dtype to return when dealing
  * with a mix of different dtypes (for example when creating an array
  * from a list of scalars).
  */
-static PyArray_StringDTypeObject *
+NPY_NO_EXPORT PyArray_StringDTypeObject *
 common_instance(PyArray_StringDTypeObject *dtype1, PyArray_StringDTypeObject *dtype2)
 {
-    int eq = _eq_comparison(dtype1->coerce, dtype2->coerce, dtype1->na_object,
-                            dtype2->na_object);
+    int eq = na_eq_cmp(dtype1->na_object, dtype2->na_object, 1);
 
     if (eq <= 0) {
         PyErr_SetString(
-                PyExc_ValueError,
-                "Cannot find common instance for unequal dtype instances");
+                PyExc_TypeError,
+                "Cannot find common instance for incompatible dtype instances");
         return NULL;
     }
 
     return (PyArray_StringDTypeObject *)new_stringdtype_instance(
-            dtype1->na_object, dtype1->coerce);
+            dtype1->na_object != NULL ? dtype1->na_object : dtype2->na_object,
+            !((dtype1->coerce == 0) || (dtype2->coerce == 0)));
 }
 
 /*
@@ -280,7 +287,7 @@ stringdtype_setitem(PyArray_StringDTypeObject *descr, PyObject *obj, char **data
 {
     npy_packed_static_string *sdata = (npy_packed_static_string *)dataptr;
 
-    int is_cmp = 0;
+    int na_cmp = 0;
 
     // borrow reference
     PyObject *na_object = descr->na_object;
@@ -294,16 +301,16 @@ stringdtype_setitem(PyArray_StringDTypeObject *descr, PyObject *obj, char **data
     // so we do the comparison before acquiring the allocator.
 
     if (na_object != NULL) {
-        is_cmp = na_eq_cmp(obj, na_object);
-        if (is_cmp == -1) {
+        na_cmp = na_eq_cmp(obj, na_object, 1);
+        if (na_cmp == -1) {
             return -1;
         }
     }
 
     npy_string_allocator *allocator = NpyString_acquire_allocator(descr);
 
     if (na_object != NULL) {
-        if (is_cmp) {
+        if (na_cmp) {
             if (NpyString_pack_null(allocator, sdata) < 0) {
                 PyErr_SetString(PyExc_MemoryError,
                                 "Failed to pack null string during StringDType "
diff --git a/numpy/_core/src/multiarray/stringdtype/dtype.h b/numpy/_core/src/multiarray/stringdtype/dtype.h
@@ -49,6 +49,9 @@ stringdtype_finalize_descr(PyArray_Descr *dtype);
 NPY_NO_EXPORT int
 _eq_comparison(int scoerce, int ocoerce, PyObject *sna, PyObject *ona);
 
+NPY_NO_EXPORT PyArray_StringDTypeObject *
+common_instance(PyArray_StringDTypeObject *dtype1, PyArray_StringDTypeObject *dtype2);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/numpy/_core/src/umath/stringdtype_ufuncs.cpp b/numpy/_core/src/umath/stringdtype_ufuncs.cpp
@@ -246,20 +246,9 @@ binary_resolve_descriptors(struct PyArrayMethodObject_tag *NPY_UNUSED(method),
 {
     PyArray_StringDTypeObject *descr1 = (PyArray_StringDTypeObject *)given_descrs[0];
     PyArray_StringDTypeObject *descr2 = (PyArray_StringDTypeObject *)given_descrs[1];
+    PyArray_StringDTypeObject *common_descr = common_instance(descr1, descr2);
 
-    // _eq_comparison has a short-circuit pointer comparison fast path,
-    // so no need to check here
-    int eq_res = _eq_comparison(descr1->coerce, descr2->coerce,
-                                descr1->na_object, descr2->na_object);
-
-    if (eq_res < 0) {
-        return (NPY_CASTING)-1;
-    }
-
-    if (eq_res != 1) {
-        PyErr_SetString(PyExc_TypeError,
-                        "Can only do binary operations with equal StringDType "
-                        "instances.");
+    if (common_descr == NULL) {
         return (NPY_CASTING)-1;
     }
 
@@ -272,8 +261,7 @@ binary_resolve_descriptors(struct PyArrayMethodObject_tag *NPY_UNUSED(method),
 
     if (given_descrs[2] == NULL) {
         out_descr = (PyArray_Descr *)new_stringdtype_instance(
-                ((PyArray_StringDTypeObject *)given_descrs[1])->na_object,
-                ((PyArray_StringDTypeObject *)given_descrs[1])->coerce);
+                common_descr->na_object, common_descr->coerce);
 
         if (out_descr == NULL) {
             return (NPY_CASTING)-1;
@@ -562,6 +550,14 @@ string_comparison_resolve_descriptors(
         PyArray_Descr *const given_descrs[],
         PyArray_Descr *loop_descrs[], npy_intp *NPY_UNUSED(view_offset))
 {
+    PyArray_StringDTypeObject *descr1 = (PyArray_StringDTypeObject *)given_descrs[0];
+    PyArray_StringDTypeObject *descr2 = (PyArray_StringDTypeObject *)given_descrs[1];
+    PyArray_StringDTypeObject *common_descr = common_instance(descr1, descr2);
+
+    if (common_descr == NULL) {
+        return (NPY_CASTING)-1;
+    }
+
     Py_INCREF(given_descrs[0]);
     loop_descrs[0] = given_descrs[0];
     Py_INCREF(given_descrs[1]);
@@ -788,20 +784,9 @@ string_findlike_resolve_descriptors(
 {
     PyArray_StringDTypeObject *descr1 = (PyArray_StringDTypeObject *)given_descrs[0];
     PyArray_StringDTypeObject *descr2 = (PyArray_StringDTypeObject *)given_descrs[1];
+    PyArray_StringDTypeObject *common_descr = common_instance(descr1, descr2);
 
-    // _eq_comparison has a short-circuit pointer comparison fast path,
-    // so no need to check here
-    int eq_res = _eq_comparison(descr1->coerce, descr2->coerce,
-                                descr1->na_object, descr2->na_object);
-
-    if (eq_res < 0) {
-        return (NPY_CASTING)-1;
-    }
-
-    if (eq_res != 1) {
-        PyErr_SetString(PyExc_TypeError,
-                        "Can only do binary operations with equal StringDType "
-                        "instances.");
+    if (common_descr == NULL) {
         return (NPY_CASTING)-1;
     }
 
@@ -849,20 +834,9 @@ string_startswith_endswith_resolve_descriptors(
 {
     PyArray_StringDTypeObject *descr1 = (PyArray_StringDTypeObject *)given_descrs[0];
     PyArray_StringDTypeObject *descr2 = (PyArray_StringDTypeObject *)given_descrs[1];
+    PyArray_StringDTypeObject *common_descr = common_instance(descr1, descr2);
 
-    // _eq_comparison has a short-circuit pointer comparison fast path, so
-    // no need to do it here
-    int eq_res = _eq_comparison(descr1->coerce, descr2->coerce,
-                                descr1->na_object, descr2->na_object);
-
-    if (eq_res < 0) {
-        return (NPY_CASTING)-1;
-    }
-
-    if (eq_res != 1) {
-        PyErr_SetString(PyExc_TypeError,
-                        "Can only do binary operations with equal StringDType "
-                        "instances.");
+    if (common_descr == NULL) {
         return (NPY_CASTING)-1;
     }
 
@@ -1061,46 +1035,6 @@ all_strings_promoter(PyObject *NPY_UNUSED(ufunc),
     return 0;
 }
 
-static NPY_CASTING
-strip_chars_resolve_descriptors(
-        struct PyArrayMethodObject_tag *NPY_UNUSED(method),
-        PyArray_DTypeMeta *const NPY_UNUSED(dtypes[]),
-        PyArray_Descr *const given_descrs[],
-        PyArray_Descr *loop_descrs[],
-        npy_intp *NPY_UNUSED(view_offset))
-{
-    Py_INCREF(given_descrs[0]);
-    loop_descrs[0] = given_descrs[0];
-
-    // we don't actually care about the null behavior of the second argument,
-    // so no need to check if the first two descrs are equal like in
-    // binary_resolve_descriptors
-
-    Py_INCREF(given_descrs[1]);
-    loop_descrs[1] = given_descrs[1];
-
-    PyArray_Descr *out_descr = NULL;
-
-    if (given_descrs[2] == NULL) {
-        out_descr = (PyArray_Descr *)new_stringdtype_instance(
-                ((PyArray_StringDTypeObject *)given_descrs[0])->na_object,
-                ((PyArray_StringDTypeObject *)given_descrs[0])->coerce);
-
-        if (out_descr == NULL) {
-            return (NPY_CASTING)-1;
-        }
-    }
-    else {
-        Py_INCREF(given_descrs[2]);
-        out_descr = given_descrs[2];
-    }
-
-    loop_descrs[2] = out_descr;
-
-    return NPY_NO_CASTING;
-}
-
-
 NPY_NO_EXPORT int
 string_lrstrip_chars_strided_loop(
         PyArrayMethod_Context *context, char *const data[],
@@ -1309,21 +1243,10 @@ replace_resolve_descriptors(struct PyArrayMethodObject_tag *NPY_UNUSED(method),
     PyArray_StringDTypeObject *descr2 = (PyArray_StringDTypeObject *)given_descrs[1];
     PyArray_StringDTypeObject *descr3 = (PyArray_StringDTypeObject *)given_descrs[2];
 
-    // _eq_comparison has a short-circuit pointer comparison fast path, so
-    // no need to do it here
-    int eq_res = (_eq_comparison(descr1->coerce, descr2->coerce,
-                                 descr1->na_object, descr2->na_object) &&
-                  _eq_comparison(descr1->coerce, descr3->coerce,
-                                 descr1->na_object, descr3->na_object));
-
-    if (eq_res < 0) {
-        return (NPY_CASTING)-1;
-    }
+    PyArray_StringDTypeObject *common_descr = common_instance(
+            common_instance(descr1, descr2), descr3);
 
-    if (eq_res != 1) {
-        PyErr_SetString(PyExc_TypeError,
-                        "String replace is only supported with equal StringDType "
-                        "instances.");
+    if (common_descr == NULL) {
         return (NPY_CASTING)-1;
     }
 
@@ -1340,8 +1263,7 @@ replace_resolve_descriptors(struct PyArrayMethodObject_tag *NPY_UNUSED(method),
 
     if (given_descrs[4] == NULL) {
         out_descr = (PyArray_Descr *)new_stringdtype_instance(
-                ((PyArray_StringDTypeObject *)given_descrs[0])->na_object,
-                ((PyArray_StringDTypeObject *)given_descrs[0])->coerce);
+                common_descr->na_object, common_descr->coerce);
 
         if (out_descr == NULL) {
             return (NPY_CASTING)-1;
@@ -1588,18 +1510,9 @@ center_ljust_rjust_resolve_descriptors(
 {
     PyArray_StringDTypeObject *input_descr = (PyArray_StringDTypeObject *)given_descrs[0];
     PyArray_StringDTypeObject *fill_descr = (PyArray_StringDTypeObject *)given_descrs[2];
+    PyArray_StringDTypeObject *common_descr = common_instance(input_descr, fill_descr);
 
-    int eq_res = _eq_comparison(input_descr->coerce, fill_descr->coerce,
-                                input_descr->na_object, fill_descr->na_object);
-
-    if (eq_res < 0) {
-        return (NPY_CASTING)-1;
-    }
-
-    if (eq_res != 1) {
-        PyErr_SetString(PyExc_TypeError,
-                        "Can only do text justification operations with equal"
-                        "StringDType instances.");
+    if (common_descr == NULL) {
         return (NPY_CASTING)-1;
     }
 
@@ -1614,8 +1527,7 @@ center_ljust_rjust_resolve_descriptors(
 
     if (given_descrs[3] == NULL) {
         out_descr = (PyArray_Descr *)new_stringdtype_instance(
-                ((PyArray_StringDTypeObject *)given_descrs[1])->na_object,
-                ((PyArray_StringDTypeObject *)given_descrs[1])->coerce);
+                common_descr->na_object, common_descr->coerce);
 
         if (out_descr == NULL) {
             return (NPY_CASTING)-1;
@@ -1888,6 +1800,7 @@ zfill_strided_loop(PyArrayMethod_Context *context,
     return -1;
 }
 
+
 static NPY_CASTING
 string_partition_resolve_descriptors(
         PyArrayMethodObject *self,
@@ -1901,14 +1814,23 @@ string_partition_resolve_descriptors(
                      "currently support the 'out' keyword", self->name);
         return (NPY_CASTING)-1;
     }
-    for (int i=0; i<2; i++) {
-        Py_INCREF(given_descrs[i]);
-        loop_descrs[i] = given_descrs[i];
+
+    PyArray_StringDTypeObject *descr1 = (PyArray_StringDTypeObject *)given_descrs[0];
+    PyArray_StringDTypeObject *descr2 = (PyArray_StringDTypeObject *)given_descrs[1];
+    PyArray_StringDTypeObject *common_descr = common_instance(descr1, descr2);
+
+    if (common_descr == NULL) {
+        return (NPY_CASTING)-1;
     }
-    PyArray_StringDTypeObject *adescr = (PyArray_StringDTypeObject *)given_descrs[0];
+
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[0] = given_descrs[0];
+    Py_INCREF(given_descrs[1]);
+    loop_descrs[1] = given_descrs[1];
+
     for (int i=2; i<5; i++) {
         loop_descrs[i] = (PyArray_Descr *)new_stringdtype_instance(
-                adescr->na_object, adescr->coerce);
+                common_descr->na_object, common_descr->coerce);
         if (loop_descrs[i] == NULL) {
             return (NPY_CASTING)-1;
         }
@@ -2655,7 +2577,7 @@ init_stringdtype_ufuncs(PyObject *umath)
 
     for (int i=0; i<3; i++) {
         if (init_ufunc(umath, strip_chars_names[i], strip_chars_dtypes,
-                       &strip_chars_resolve_descriptors,
+                       &binary_resolve_descriptors,
                        &string_lrstrip_chars_strided_loop,
                        2, 1, NPY_NO_CASTING, (NPY_ARRAYMETHOD_FLAGS) 0,
                        &strip_types[i]) < 0) {
diff --git a/numpy/_core/tests/test_stringdtype.py b/numpy/_core/tests/test_stringdtype.py

Original file line number	Diff line number	Diff line change
`@@ -49,6 +49,9 @@ stringdtype_finalize_descr(PyArray_Descr *dtype);`
`49`	`49`	`NPY_NO_EXPORT int`
`50`	`50`	`_eq_comparison(int scoerce, int ocoerce, PyObject sna, PyObject ona);`
`51`	`51`
	`52`	`+NPY_NO_EXPORT PyArray_StringDTypeObject *`
	`53`	`+common_instance(PyArray_StringDTypeObject dtype1, PyArray_StringDTypeObject dtype2);`
	`54`	`+`
`52`	`55`	`#ifdef __cplusplus`
`53`	`56`	`}`
`54`	`57`	`#endif`