diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index 0ca63e43da40..27791afb5e7d 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -22,6 +22,7 @@ #include "ctors.h" #include "lowlevel_strided_loops.h" #include "array_assign.h" +#include "refcount.h" #include "npy_sort.h" #include "npy_partition.h" @@ -1070,6 +1071,9 @@ _new_sortlike(PyArrayObject *op, int axis, PyArray_SortFunc *sort, ret = -1; goto fail; } + if (PyDataType_FLAGCHK(PyArray_DESCR(op), NPY_NEEDS_INIT)) { + memset(buffer, 0, N * elsize); + } } NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(op)); @@ -1114,16 +1118,7 @@ _new_sortlike(PyArrayObject *op, int axis, PyArray_SortFunc *sort, } if (needcopy) { - if (hasrefs) { - if (swap) { - copyswapn(buffer, elsize, NULL, 0, N, swap, op); - } - _unaligned_strided_byte_copy(it->dataptr, astride, - buffer, elsize, N, elsize); - } - else { - copyswapn(it->dataptr, astride, buffer, elsize, N, swap, op); - } + copyswapn(it->dataptr, astride, buffer, elsize, N, swap, op); } PyArray_ITER_NEXT(it); @@ -1132,7 +1127,10 @@ _new_sortlike(PyArrayObject *op, int axis, PyArray_SortFunc *sort, fail: NPY_END_THREADS_DESCR(PyArray_DESCR(op)); /* cleanup internal buffer */ - PyDataMem_UserFREE(buffer, N * elsize, mem_handler); + if (needcopy) { + PyArray_ClearBuffer(PyArray_DESCR(op), buffer, elsize, N, 1); + PyDataMem_UserFREE(buffer, N * elsize, mem_handler); + } if (ret < 0 && !PyErr_Occurred()) { /* Out of memory during sorting or buffer creation */ PyErr_NoMemory(); @@ -1206,6 +1204,9 @@ _new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort, ret = -1; goto fail; } + if (PyDataType_FLAGCHK(PyArray_DESCR(op), NPY_NEEDS_INIT)) { + memset(valbuffer, 0, N * elsize); + } } if (needidxbuffer) { @@ -1281,7 +1282,10 @@ _new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort, fail: NPY_END_THREADS_DESCR(PyArray_DESCR(op)); /* cleanup internal buffers */ - PyDataMem_UserFREE(valbuffer, N * elsize, mem_handler); + if (needcopy) { + PyArray_ClearBuffer(PyArray_DESCR(op), valbuffer, elsize, N, 1); + PyDataMem_UserFREE(valbuffer, N * elsize, mem_handler); + } PyDataMem_UserFREE(idxbuffer, N * sizeof(npy_intp), mem_handler); if (ret < 0) { if (!PyErr_Occurred()) { diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py index 6a803a68da14..0799c0a5acf9 100644 --- a/numpy/core/tests/test_multiarray.py +++ b/numpy/core/tests/test_multiarray.py @@ -2116,19 +2116,26 @@ def test_sort_object(self): c.sort(kind=kind) assert_equal(c, a, msg) - def test_sort_structured(self): + @pytest.mark.parametrize("dt", [ + np.dtype([('f', float), ('i', int)]), + np.dtype([('f', float), ('i', object)])]) + @pytest.mark.parametrize("step", [1, 2]) + def test_sort_structured(self, dt, step): # test record array sorts. - dt = np.dtype([('f', float), ('i', int)]) - a = np.array([(i, i) for i in range(101)], dtype=dt) + a = np.array([(i, i) for i in range(101*step)], dtype=dt) b = a[::-1] for kind in ['q', 'h', 'm']: msg = "kind=%s" % kind - c = a.copy() + c = a.copy()[::step] + indx = c.argsort(kind=kind) c.sort(kind=kind) - assert_equal(c, a, msg) - c = b.copy() + assert_equal(c, a[::step], msg) + assert_equal(a[::step][indx], a[::step], msg) + c = b.copy()[::step] + indx = c.argsort(kind=kind) c.sort(kind=kind) - assert_equal(c, a, msg) + assert_equal(c, a[step-1::step], msg) + assert_equal(b[::step][indx], a[step-1::step], msg) @pytest.mark.parametrize('dtype', ['datetime64[D]', 'timedelta64[D]']) def test_sort_time(self, dtype):