Merge pull request #7466 from jaimefrio/inplace_reduceat_segfault

BUG: segfault inplace object reduceat, fixes #7465

2 files changed, 37 insertions, 2 deletions

diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 6eb0aae55..52299c1be 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -3700,10 +3700,18 @@ PyUFunc_Reduceat(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *ind,
 
                 /* Copy the first element to start the reduction */
                 if (otype == NPY_OBJECT) {
+                    /*
+                     * Input (dataptr[0]) and output (dataptr[1]) may
+                     * point to the same memory, e.g.
+                     * np.add.reduceat(a, np.arange(len(a)), out=a).
+                     * In that case need to 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];
-                    Py_XINCREF(*(PyObject **)dataptr_copy[0]);
                 }
                 else {
                     memcpy(dataptr_copy[0], dataptr_copy[1], itemsize);
@@ -3758,10 +3766,18 @@ PyUFunc_Reduceat(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *ind,
 
             /* Copy the first element to start the reduction */
             if (otype == NPY_OBJECT) {
+                /*
+                 * Input (dataptr[0]) and output (dataptr[1]) may point
+                 * to the same memory, e.g.
+                 * np.add.reduceat(a, np.arange(len(a)), out=a).
+                 * In that case need to 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];
-                Py_XINCREF(*(PyObject **)dataptr_copy[0]);
             }
             else {
                 memcpy(dataptr_copy[0], dataptr_copy[1], itemsize);
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index ab8cecff0..62fe3c04b 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -665,6 +665,25 @@ class TestUfunc(TestCase):
         np.add.accumulate(arr, out=arr, axis=-1)
         assert_array_equal(arr[0, :], np.array([[2]*i for i in [1, 3, 6, 10]]))
 
+    def test_object_array_reduceat_inplace(self):
+        # Checks that in-place reduceats work, see also gh-7465
+        arr = np.empty(4, dtype=object)
+        arr[:] = [[1] for i in range(4)]
+        out = np.empty(4, dtype=object)
+        out[:] = [[1] for i in range(4)]
+        np.add.reduceat(arr, np.arange(4), out=arr)
+        np.add.reduceat(arr, np.arange(4), out=arr)
+        assert_array_equal(arr, out)
+
+        # And the same if the axis argument is used
+        arr = np.ones((2, 4), dtype=object)
+        arr[0, :] = [[2] for i in range(4)]
+        out = np.ones((2, 4), dtype=object)
+        out[0, :] = [[2] for i in range(4)]
+        np.add.reduceat(arr, np.arange(4), out=arr, axis=-1)
+        np.add.reduceat(arr, np.arange(4), out=arr, axis=-1)
+        assert_array_equal(arr, out)
+
     def test_object_scalar_multiply(self):
         # Tickets #2469 and #4482
         arr = np.matrix([1, 2], dtype=object)