Merge pull request numpy#26261 from ngoldbaum/compatible-stringdtype-ufuncs

ngoldbaum · web-flow · commit 3c09f1666022 · 2024-04-16T11:14:08.000-06:00
ENH: introduce a notion of "compatible" stringdtype instances
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
@@ -185,6 +185,30 @@ _eq_comparison(int scoerce, int ocoerce, PyObject *sna, PyObject *ona)
     return na_eq_cmp(sna, ona);
 }
 
+// Currently this can only return 0 or -1, the latter indicating that the
+// error indicator is set. Pass in out_na if you want to figure out which
+// na is valid.
+NPY_NO_EXPORT int
+stringdtype_compatible_na(PyObject *na1, PyObject *na2, PyObject **out_na) {
+    if ((na1 != NULL) && (na2 != NULL)) {
+        int na_eq = na_eq_cmp(na1, na2);
+
+        if (na_eq < 0) {
+            return -1;
+        }
+        else if (na_eq == 0) {
+            PyErr_Format(PyExc_TypeError,
+                         "Cannot find a compatible null string value for "
+                         "null strings '%R' and '%R'", na1, na2);
+            return -1;
+        }
+    }
+    if (out_na != NULL) {
+        *out_na = na1 ? na1 : na2;
+    }
+    return 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
@@ -193,18 +217,18 @@ _eq_comparison(int scoerce, int ocoerce, PyObject *sna, PyObject *ona)
 static PyArray_StringDTypeObject *
 common_instance(PyArray_StringDTypeObject *dtype1, PyArray_StringDTypeObject *dtype2)
 {
-    int eq = _eq_comparison(dtype1->coerce, dtype2->coerce, dtype1->na_object,
-                            dtype2->na_object);
+    PyObject *out_na_object = NULL;
 
-    if (eq <= 0) {
-        PyErr_SetString(
-                PyExc_ValueError,
-                "Cannot find common instance for unequal dtype instances");
+    if (stringdtype_compatible_na(
+                dtype1->na_object, dtype2->na_object, &out_na_object) == -1) {
+        PyErr_Format(PyExc_TypeError,
+                     "Cannot find common instance for incompatible dtypes "
+                     "'%R' and '%R'", (PyObject *)dtype1, (PyObject *)dtype2);
         return NULL;
     }
 
     return (PyArray_StringDTypeObject *)new_stringdtype_instance(
-            dtype1->na_object, dtype1->coerce);
+            out_na_object, dtype1->coerce && dtype1->coerce);
 }
 
 /*
@@ -280,30 +304,22 @@ stringdtype_setitem(PyArray_StringDTypeObject *descr, PyObject *obj, char **data
 {
     npy_packed_static_string *sdata = (npy_packed_static_string *)dataptr;
 
-    int is_cmp = 0;
-
     // borrow reference
     PyObject *na_object = descr->na_object;
 
-    // Note there are two different na_object != NULL checks here.
-    //
-    // Do not refactor this!
-    //
     // We need the result of the comparison after acquiring the allocator, but
     // cannot use functions requiring the GIL when the allocator is acquired,
     // 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) {
-            return -1;
-        }
+    int na_cmp = na_eq_cmp(obj, na_object);
+    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 int
+stringdtype_compatible_na(PyObject *na1, PyObject *na2, PyObject **out_na);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/numpy/_core/src/umath/stringdtype_ufuncs.cpp b/numpy/_core/src/umath/stringdtype_ufuncs.cpp
@@ -246,20 +246,11 @@ 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];
+    int out_coerce = descr1->coerce && descr1->coerce;
+    PyObject *out_na_object = NULL;
 
-    // _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 (stringdtype_compatible_na(
+                descr1->na_object, descr2->na_object, &out_na_object) == -1) {
         return (NPY_CASTING)-1;
     }
 
@@ -272,8 +263,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);
+                out_na_object, out_coerce);
 
         if (out_descr == NULL) {
             return (NPY_CASTING)-1;
@@ -562,6 +552,13 @@ 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];
+
+    if (stringdtype_compatible_na(descr1->na_object, descr2->na_object, NULL) == -1) {
+        return (NPY_CASTING)-1;
+    }
+
     Py_INCREF(given_descrs[0]);
     loop_descrs[0] = given_descrs[0];
     Py_INCREF(given_descrs[1]);
@@ -789,19 +786,7 @@ string_findlike_resolve_descriptors(
     PyArray_StringDTypeObject *descr1 = (PyArray_StringDTypeObject *)given_descrs[0];
     PyArray_StringDTypeObject *descr2 = (PyArray_StringDTypeObject *)given_descrs[1];
 
-    // _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 (stringdtype_compatible_na(descr1->na_object, descr2->na_object, NULL) == -1) {
         return (NPY_CASTING)-1;
     }
 
@@ -850,19 +835,7 @@ string_startswith_endswith_resolve_descriptors(
     PyArray_StringDTypeObject *descr1 = (PyArray_StringDTypeObject *)given_descrs[0];
     PyArray_StringDTypeObject *descr2 = (PyArray_StringDTypeObject *)given_descrs[1];
 
-    // _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 (stringdtype_compatible_na(descr1->na_object, descr2->na_object, NULL) == -1) {
         return (NPY_CASTING)-1;
     }
 
@@ -1061,46 +1034,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[],
@@ -1308,22 +1241,16 @@ replace_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 *descr3 = (PyArray_StringDTypeObject *)given_descrs[2];
+    int out_coerce = descr1->coerce && descr2->coerce && descr3->coerce;
+    PyObject *out_na_object = NULL;
 
-    // _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) {
+    if (stringdtype_compatible_na(
+                descr1->na_object, descr2->na_object, &out_na_object) == -1) {
         return (NPY_CASTING)-1;
     }
 
-    if (eq_res != 1) {
-        PyErr_SetString(PyExc_TypeError,
-                        "String replace is only supported with equal StringDType "
-                        "instances.");
+    if (stringdtype_compatible_na(
+                out_na_object, descr3->na_object, &out_na_object) == -1) {
         return (NPY_CASTING)-1;
     }
 
@@ -1340,8 +1267,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);
+                out_na_object, out_coerce);
 
         if (out_descr == NULL) {
             return (NPY_CASTING)-1;
@@ -1588,18 +1514,11 @@ center_ljust_rjust_resolve_descriptors(
 {
     PyArray_StringDTypeObject *input_descr = (PyArray_StringDTypeObject *)given_descrs[0];
     PyArray_StringDTypeObject *fill_descr = (PyArray_StringDTypeObject *)given_descrs[2];
+    int out_coerce = input_descr->coerce && fill_descr->coerce;
+    PyObject *out_na_object = NULL;
 
-    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 (stringdtype_compatible_na(
+                input_descr->na_object, fill_descr->na_object, &out_na_object) == -1) {
         return (NPY_CASTING)-1;
     }
 
@@ -1614,8 +1533,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);
+                out_na_object, out_coerce);
 
         if (out_descr == NULL) {
             return (NPY_CASTING)-1;
@@ -1888,6 +1806,7 @@ zfill_strided_loop(PyArrayMethod_Context *context,
     return -1;
 }
 
+
 static NPY_CASTING
 string_partition_resolve_descriptors(
         PyArrayMethodObject *self,
@@ -1901,14 +1820,25 @@ 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];
+    int out_coerce = descr1->coerce && descr2->coerce;
+    PyObject *out_na_object = NULL;
+
+    if (stringdtype_compatible_na(
+                descr1->na_object, descr2->na_object, &out_na_object) == -1) {
+        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);
+                out_na_object, out_coerce);
         if (loop_descrs[i] == NULL) {
             return (NPY_CASTING)-1;
         }
@@ -2655,7 +2585,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/_natype.py b/numpy/_core/tests/_natype.py
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 int`
	`53`	`+stringdtype_compatible_na(PyObject na1, PyObject na2, PyObject **out_na);`
	`54`	`+`
`52`	`55`	`#ifdef __cplusplus`
`53`	`56`	`}`
`54`	`57`	`#endif`