BUG: Fix repeat, accumulate for strings and accumulate API logic (numpy#27773)

seberg · ngoldbaum · web-flow · commit fb91abc0e888 · 2024-11-18T17:53:08.000+01:00
* BUG: Fix repeat, accumulate for strings and accumulate API logic This fixes three (actually four) things: 1. `needs_api` was wrongly overwritten, this was just a line not deleted earlier. 2. Accumulate didn't allow string dtypes because it only allowed objects with references, this fixes it and adds a test. 3. `repeat` was just broken with string dtype... I guess there was no actual test. 4. `repeat` internals passed on `cast_info` not by reference. I guess that isn't a bug, but it's weird. Tests cover things relatively well, although things like GIL release being right is of course harder to test. Closes numpygh-27709 * DOC: Add small code comment about helping compiler * Update numpy/_core/src/umath/ufunc_object.c Co-authored-by: Nathan Goldbaum <nathan.goldbaum@gmail.com> --------- Co-authored-by: Nathan Goldbaum <nathan.goldbaum@gmail.com>
diff --git a/numpy/_core/src/multiarray/item_selection.c b/numpy/_core/src/multiarray/item_selection.c
@@ -785,21 +785,21 @@ static NPY_GCC_OPT_3 inline int
 npy_fastrepeat_impl(
     npy_intp n_outer, npy_intp n, npy_intp nel, npy_intp chunk,
     npy_bool broadcast, npy_intp* counts, char* new_data, char* old_data,
-    npy_intp elsize, NPY_cast_info cast_info, int needs_refcounting)
+    npy_intp elsize, NPY_cast_info *cast_info, int needs_custom_copy)
 {
     npy_intp i, j, k;
     for (i = 0; i < n_outer; i++) {
         for (j = 0; j < n; j++) {
             npy_intp tmp = broadcast ? counts[0] : counts[j];
             for (k = 0; k < tmp; k++) {
-                if (!needs_refcounting) {
+                if (!needs_custom_copy) {
                     memcpy(new_data, old_data, chunk);
                 }
                 else {
                     char *data[2] = {old_data, new_data};
                     npy_intp strides[2] = {elsize, elsize};
-                    if (cast_info.func(&cast_info.context, data, &nel,
-                                       strides, cast_info.auxdata) < 0) {
+                    if (cast_info->func(&cast_info->context, data, &nel,
+                                       strides, cast_info->auxdata) < 0) {
                         return -1;
                     }
                 }
@@ -811,48 +811,53 @@ npy_fastrepeat_impl(
     return 0;
 }
 
+
+/*
+ * Helper to allow the compiler to specialize for all direct element copy
+ * cases (e.g. all numerical dtypes).
+ */
 static NPY_GCC_OPT_3 int
 npy_fastrepeat(
     npy_intp n_outer, npy_intp n, npy_intp nel, npy_intp chunk,
     npy_bool broadcast, npy_intp* counts, char* new_data, char* old_data,
-    npy_intp elsize, NPY_cast_info cast_info, int needs_refcounting)
+    npy_intp elsize, NPY_cast_info *cast_info, int needs_custom_copy)
 {
-    if (!needs_refcounting) {
+    if (!needs_custom_copy) {
         if (chunk == 1) {
             return npy_fastrepeat_impl(
                 n_outer, n, nel, chunk, broadcast, counts, new_data, old_data,
-                elsize, cast_info, needs_refcounting);
+                elsize, cast_info, needs_custom_copy);
         }
         if (chunk == 2) {
             return npy_fastrepeat_impl(
                 n_outer, n, nel, chunk, broadcast, counts, new_data, old_data,
-                elsize, cast_info, needs_refcounting);
+                elsize, cast_info, needs_custom_copy);
         }
         if (chunk == 4) {
             return npy_fastrepeat_impl(
                 n_outer, n, nel, chunk, broadcast, counts, new_data, old_data,
-                elsize, cast_info, needs_refcounting);
+                elsize, cast_info, needs_custom_copy);
         }
         if (chunk == 8) {
             return npy_fastrepeat_impl(
                 n_outer, n, nel, chunk, broadcast, counts, new_data, old_data,
-                elsize, cast_info, needs_refcounting);
+                elsize, cast_info, needs_custom_copy);
         }
         if (chunk == 16) {
             return npy_fastrepeat_impl(
                 n_outer, n, nel, chunk, broadcast, counts, new_data, old_data,
-                elsize, cast_info, needs_refcounting);
+                elsize, cast_info, needs_custom_copy);
         }
         if (chunk == 32) {
             return npy_fastrepeat_impl(
                 n_outer, n, nel, chunk, broadcast, counts, new_data, old_data,
-                elsize, cast_info, needs_refcounting);
+                elsize, cast_info, needs_custom_copy);
         }
     }
 
     return npy_fastrepeat_impl(
         n_outer, n, nel, chunk, broadcast, counts, new_data, old_data, elsize,
-        cast_info, needs_refcounting);
+        cast_info, needs_custom_copy);
 }
 
 
@@ -872,7 +877,6 @@ PyArray_Repeat(PyArrayObject *aop, PyObject *op, int axis)
     char *new_data, *old_data;
     NPY_cast_info cast_info;
     NPY_ARRAYMETHOD_FLAGS flags;
-    int needs_refcounting;
 
     repeats = (PyArrayObject *)PyArray_ContiguousFromAny(op, NPY_INTP, 0, 1);
     if (repeats == NULL) {
@@ -897,7 +901,6 @@ PyArray_Repeat(PyArrayObject *aop, PyObject *op, int axis)
     aop = (PyArrayObject *)ap;
     n = PyArray_DIM(aop, axis);
     NPY_cast_info_init(&cast_info);
-    needs_refcounting = PyDataType_REFCHK(PyArray_DESCR(aop));
 
     if (!broadcast && PyArray_SIZE(repeats) != n) {
         PyErr_Format(PyExc_ValueError,
@@ -947,16 +950,18 @@ PyArray_Repeat(PyArrayObject *aop, PyObject *op, int axis)
         n_outer *= PyArray_DIMS(aop)[i];
     }
 
-    if (needs_refcounting) {
+    int needs_custom_copy = 0;
+    if (PyDataType_REFCHK(PyArray_DESCR(ret))) {
+        needs_custom_copy = 1;
         if (PyArray_GetDTypeTransferFunction(
-                1, elsize, elsize, PyArray_DESCR(aop), PyArray_DESCR(aop), 0,
+                1, elsize, elsize, PyArray_DESCR(aop), PyArray_DESCR(ret), 0,
                 &cast_info, &flags) < 0) {
             goto fail;
         }
     }
 
     if (npy_fastrepeat(n_outer, n, nel, chunk, broadcast, counts, new_data,
-                       old_data, elsize, cast_info, needs_refcounting) < 0) {
+                       old_data, elsize, &cast_info, needs_custom_copy) < 0) {
         goto fail;
     }
 
diff --git a/numpy/_core/src/umath/ufunc_object.c b/numpy/_core/src/umath/ufunc_object.c
@@ -2593,6 +2593,10 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
     int idim, ndim;
     int needs_api, need_outer_iterator;
     int res = 0;
+
+    NPY_cast_info copy_info;
+    NPY_cast_info_init(&copy_info);
+
 #if NPY_UF_DBG_TRACING
     const char *ufunc_name = ufunc_get_name_cstr(ufunc);
 #endif
@@ -2637,14 +2641,6 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
     assert(PyArray_EquivTypes(descrs[0], descrs[1])
            && PyArray_EquivTypes(descrs[0], descrs[2]));
 
-    if (PyDataType_REFCHK(descrs[2]) && descrs[2]->type_num != NPY_OBJECT) {
-        /* This can be removed, but the initial element copy needs fixing */
-        PyErr_SetString(PyExc_TypeError,
-                "accumulation currently only supports `object` dtype with "
-                "references");
-        goto fail;
-    }
-
     PyArrayMethod_Context context = {
         .caller = (PyObject *)ufunc,
         .method = ufuncimpl,
@@ -2740,10 +2736,10 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
         else {
             PyArray_Descr *dtype = descrs[0];
             Py_INCREF(dtype);
-            op[0] = out = (PyArrayObject *)PyArray_NewFromDescr(
+            op[0] = out = (PyArrayObject *)PyArray_NewFromDescr_int(
                                     &PyArray_Type, dtype,
                                     ndim, PyArray_DIMS(op[1]), NULL, NULL,
-                                    0, NULL);
+                                    0, NULL, NULL, _NPY_ARRAY_ENSURE_DTYPE_IDENTITY);
             if (out == NULL) {
                 goto fail;
             }
@@ -2766,6 +2762,18 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
             1, 0, fixed_strides, &strided_loop, &auxdata, &flags) < 0) {
         goto fail;
     }
+    /* Set up function to copy the first element if it has references */
+    if (PyDataType_REFCHK(descrs[2])) {
+        NPY_ARRAYMETHOD_FLAGS copy_flags;
+        /* Setup guarantees aligned here. */
+        if (PyArray_GetDTypeTransferFunction(
+                1, 0, 0, descrs[1], descrs[2], 0, &copy_info,
+                &copy_flags) == NPY_FAIL) {
+            goto fail;
+        }
+        flags = PyArrayMethod_COMBINED_FLAGS(flags, copy_flags);
+    }
+
     needs_api = (flags & NPY_METH_REQUIRES_PYAPI) != 0;
     if (!(flags & NPY_METH_NO_FLOATINGPOINT_ERRORS)) {
         /* Start with the floating-point exception flags cleared */
@@ -2829,18 +2837,17 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
              * Output (dataptr[0]) and input (dataptr[1]) may point to
              * the same memory, e.g. np.add.accumulate(a, out=a).
              */
-            if (descrs[2]->type_num == NPY_OBJECT) {
-                /*
-                 * Incref before decref to avoid the possibility of the
-                 * reference count being zero temporarily.
-                 */
-                Py_XINCREF(*(PyObject **)dataptr_copy[1]);
-                Py_XDECREF(*(PyObject **)dataptr_copy[0]);
-                *(PyObject **)dataptr_copy[0] =
-                                    *(PyObject **)dataptr_copy[1];
+            if (copy_info.func) {
+                const npy_intp one = 1;
+                if (copy_info.func(
+                        &copy_info.context, &dataptr_copy[1], &one,
+                        &stride_copy[1], copy_info.auxdata) < 0) {
+                    NPY_END_THREADS;
+                    goto fail;
+                }
             }
             else {
-                memmove(dataptr_copy[0], dataptr_copy[1], itemsize);
+                memmove(dataptr_copy[2], dataptr_copy[1], itemsize);
             }
 
             if (count_m1 > 0) {
@@ -2889,18 +2896,17 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
          * Output (dataptr[0]) and input (dataptr[1]) may point to the
          * same memory, e.g. np.add.accumulate(a, out=a).
          */
-        if (descrs[2]->type_num == NPY_OBJECT) {
-            /*
-             * Incref before decref to avoid the possibility of the
-             * reference count being zero temporarily.
-             */
-            Py_XINCREF(*(PyObject **)dataptr_copy[1]);
-            Py_XDECREF(*(PyObject **)dataptr_copy[0]);
-            *(PyObject **)dataptr_copy[0] =
-                                *(PyObject **)dataptr_copy[1];
+        if (copy_info.func) {
+            const npy_intp one = 1;
+            const npy_intp strides[2] = {itemsize, itemsize};
+            if (copy_info.func(
+                    &copy_info.context, &dataptr_copy[1], &one,
+                    strides, copy_info.auxdata) < 0) {
+                goto fail;
+            }
         }
         else {
-            memmove(dataptr_copy[0], dataptr_copy[1], itemsize);
+            memmove(dataptr_copy[2], dataptr_copy[1], itemsize);
         }
 
         if (count > 1) {
@@ -2910,8 +2916,6 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
 
             NPY_UF_DBG_PRINT1("iterator loop count %d\n", (int)count);
 
-            needs_api = PyDataType_REFCHK(descrs[0]);
-
             if (!needs_api) {
                 NPY_BEGIN_THREADS_THRESHOLDED(count);
             }
@@ -2925,6 +2929,7 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
 
 finish:
     NPY_AUXDATA_FREE(auxdata);
+    NPY_cast_info_xfree(&copy_info);
     Py_DECREF(descrs[0]);
     Py_DECREF(descrs[1]);
     Py_DECREF(descrs[2]);
@@ -2949,6 +2954,8 @@ PyUFunc_Accumulate(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
     Py_XDECREF(out);
 
     NPY_AUXDATA_FREE(auxdata);
+    NPY_cast_info_xfree(&copy_info);
+
     Py_XDECREF(descrs[0]);
     Py_XDECREF(descrs[1]);
     Py_XDECREF(descrs[2]);
diff --git a/numpy/_core/tests/test_stringdtype.py b/numpy/_core/tests/test_stringdtype.py
@@ -1574,6 +1574,40 @@ def test_unset_na_coercion():
             arr == op
 
 
+def test_repeat(string_array):
+    res = string_array.repeat(1000)
+    # Create an empty array with expanded dimension, and fill it.  Then,
+    # reshape it to the expected result.
+    expected = np.empty_like(string_array, shape=string_array.shape + (1000,))
+    expected[...] = string_array[:, np.newaxis]
+    expected = expected.reshape(-1)
+
+    assert_array_equal(res, expected, strict=True)
+
+
+@pytest.mark.parametrize("tile", [1, 6, (2, 5)])
+def test_accumulation(string_array, tile):
+    """Accumulation is odd for StringDType but tests dtypes with references.
+    """
+    # Fill with mostly empty strings to not create absurdly big strings
+    arr = np.zeros_like(string_array, shape=(100,))
+    arr[:len(string_array)] = string_array
+    arr[-len(string_array):] = string_array
+
+    # Bloat size a bit (get above thresholds and test >1 ndim).
+    arr = np.tile(string_array, tile)
+
+    res = np.add.accumulate(arr, axis=0)
+    res_obj = np.add.accumulate(arr.astype(object), axis=0)
+    assert_array_equal(res, res_obj.astype(arr.dtype), strict=True)
+
+    if arr.ndim > 1:
+        res = np.add.accumulate(arr, axis=-1)
+        res_obj = np.add.accumulate(arr.astype(object), axis=-1)
+
+        assert_array_equal(res, res_obj.astype(arr.dtype), strict=True)
+
+
 class TestImplementation:
     """Check that strings are stored in the arena when possible.
 
