BUG: fancy indexing copy (numpy#26558)

tylerjereddy · seberg · DimitriPapadopoulos · commit c2024f7d8c21 · 2024-07-04T22:43:00.000+02:00
* BUG: fancy indexing copy * Fixes numpygh-26542 * For these very simple advanced indexing cases, if the `result` and `self` arrays share the same data pointers, use a copy of `result` to assign values to `self`, otherwise the outcome of the fancy indexing suffers from mutation of `result` before the assignments are complete. * MAINT, BUG: PR 26558 revisions * Avoid leaking the old `tmp_arr`, based on reviewer feedback. * Add a test for a similar 2D case that fails, then hoist `solve_may_share_memory()` check farther up in the control flow such that the test passes. * Add a reviewer-requested test for index overlap. * BUG: PR 26558 revisions * The usage of `solve_may_share_memory` in the above PR wasn't quite right since it ignored the case of failing to solve the overlap problem. This has been revised according to reviewer feedback. Co-authored-by: Sebastian Berg <sebastianb@nvidia.com> --------- Co-authored-by: Sebastian Berg <sebastianb@nvidia.com>
diff --git a/numpy/_core/src/multiarray/mapping.c b/numpy/_core/src/multiarray/mapping.c
@@ -1960,6 +1960,10 @@ array_assign_subscript(PyArrayObject *self, PyObject *ind, PyObject *op)
         tmp_arr = (PyArrayObject *)op;
     }
 
+    if (tmp_arr && solve_may_share_memory(self, tmp_arr, 1) != 0) {
+        Py_SETREF(tmp_arr, (PyArrayObject *)PyArray_NewCopy(tmp_arr, NPY_ANYORDER));
+    }
+
     /*
      * Special case for very simple 1-d fancy indexing, which however
      * is quite common. This saves not only a lot of setup time in the
diff --git a/numpy/_core/tests/test_indexing.py b/numpy/_core/tests/test_indexing.py
@@ -133,6 +133,28 @@ def test_empty_fancy_index(self):
         b = np.array([])
         assert_raises(IndexError, a.__getitem__, b)
 
+    def test_gh_26542(self):
+        a = np.array([0, 1, 2])
+        idx = np.array([2, 1, 0])
+        a[idx] = a
+        expected = np.array([2, 1, 0])
+        assert_equal(a, expected)
+
+    def test_gh_26542_2d(self):
+        a = np.array([[0, 1, 2]])
+        idx_row = np.zeros(3, dtype=int)
+        idx_col = np.array([2, 1, 0])
+        a[idx_row, idx_col] = a
+        expected = np.array([[2, 1, 0]])
+        assert_equal(a, expected)
+
+    def test_gh_26542_index_overlap(self):
+        arr = np.arange(100)
+        expected_vals = np.copy(arr[:-10])
+        arr[10:] = arr[:-10]
+        actual_vals = arr[10:]
+        assert_equal(actual_vals, expected_vals)
+
     def test_ellipsis_index(self):
         a = np.array([[1, 2, 3],
                       [4, 5, 6],
