Merge pull request #2974 from seberg/issue-427

Fix some large array sorts and memory error handling

2 files changed, 36 insertions, 14 deletions

diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
index 31916cf40..7c8041cf1 100644
--- a/numpy/core/src/multiarray/item_selection.c
+++ b/numpy/core/src/multiarray/item_selection.c
@@ -750,6 +750,9 @@ _new_sort(PyArrayObject *op, int axis, NPY_SORTKIND which)
                 (astride != (npy_intp) elsize) || swap;
     if (needcopy) {
         char *buffer = PyDataMem_NEW(N*elsize);
+        if (buffer == NULL) {
+            goto fail;
+        }
 
         while (size--) {
             _unaligned_strided_byte_copy(buffer, (npy_intp) elsize, it->dataptr,
@@ -783,9 +786,11 @@ _new_sort(PyArrayObject *op, int axis, NPY_SORTKIND which)
     return 0;
 
  fail:
+    /* Out of memory during sorting or buffer creation */
     NPY_END_THREADS;
+    PyErr_NoMemory();
     Py_DECREF(it);
-    return 0;
+    return -1;
 }
 
 static PyObject*
@@ -795,10 +800,10 @@ _new_argsort(PyArrayObject *op, int axis, NPY_SORTKIND which)
     PyArrayIterObject *it = NULL;
     PyArrayIterObject *rit = NULL;
     PyArrayObject *ret;
-    int needcopy = 0, i;
-    npy_intp N, size;
-    int elsize, swap;
+    npy_intp N, size, i;
     npy_intp astride, rstride, *iptr;
+    int elsize;
+    int needcopy = 0, swap;
     PyArray_ArgSortFunc *argsort;
     NPY_BEGIN_THREADS_DEF;
 
@@ -832,7 +837,14 @@ _new_argsort(PyArrayObject *op, int axis, NPY_SORTKIND which)
         char *valbuffer, *indbuffer;
 
         valbuffer = PyDataMem_NEW(N*elsize);
+        if (valbuffer == NULL) {
+            goto fail;
+        }
         indbuffer = PyDataMem_NEW(N*sizeof(npy_intp));
+        if (indbuffer == NULL) {
+            PyDataMem_FREE(valbuffer);
+            goto fail;
+        }
         while (size--) {
             _unaligned_strided_byte_copy(valbuffer, (npy_intp) elsize, it->dataptr,
                                          astride, N, elsize);
@@ -878,6 +890,10 @@ _new_argsort(PyArrayObject *op, int axis, NPY_SORTKIND which)
 
  fail:
     NPY_END_THREADS;
+    if (!PyErr_Occurred()) {
+        /* Out of memory during sorting or buffer creation */
+        PyErr_NoMemory();
+    }
     Py_DECREF(ret);
     Py_XDECREF(it);
     Py_XDECREF(rit);
@@ -952,7 +968,8 @@ PyArray_Sort(PyArrayObject *op, int axis, NPY_SORTKIND which)
 {
     PyArrayObject *ap = NULL, *store_arr = NULL;
     char *ip;
-    int i, n, m, elsize, orign;
+    npy_intp i, n, m;
+    int elsize, orign;
     int res = 0;
     int axis_orig = axis;
     int (*sort)(void *, size_t, size_t, npy_comparator);
@@ -1208,13 +1225,12 @@ PyArray_LexSort(PyObject *sort_keys, int axis)
     PyArrayIterObject **its;
     PyArrayObject *ret = NULL;
     PyArrayIterObject *rit = NULL;
-    int n;
+    npy_intp n, N, size, i, j;
+    npy_intp astride, rstride, *iptr;
     int nd;
-    int needcopy = 0, i,j;
-    npy_intp N, size;
+    int needcopy = 0;
     int elsize;
     int maxelsize;
-    npy_intp astride, rstride, *iptr;
     int object = 0;
     PyArray_ArgSortFunc *argsort;
     NPY_BEGIN_THREADS_DEF;
@@ -1331,7 +1347,14 @@ PyArray_LexSort(PyObject *sort_keys, int axis)
         int *swaps;
 
         valbuffer = PyDataMem_NEW(N*maxelsize);
+        if (valbuffer == NULL) {
+            goto fail;
+        }
         indbuffer = PyDataMem_NEW(N*sizeof(npy_intp));
+        if (indbuffer == NULL) {
+            PyDataMem_FREE(indbuffer);
+            goto fail;
+        }
         swaps = malloc(n*sizeof(int));
         for (j = 0; j < n; j++) {
             swaps[j] = PyArray_ISBYTESWAPPED(mps[j]);
@@ -1400,6 +1423,10 @@ PyArray_LexSort(PyObject *sort_keys, int axis)
 
  fail:
     NPY_END_THREADS;
+    if (!PyErr_Occurred()) {
+        /* Out of memory during sorting or buffer creation */
+        PyErr_NoMemory();
+    }
     Py_XDECREF(rit);
     Py_XDECREF(ret);
     for (i = 0; i < n; i++) {
diff --git a/numpy/core/src/npysort/mergesort.c.src b/numpy/core/src/npysort/mergesort.c.src
index 1003b95d1..55dcab4fb 100644
--- a/numpy/core/src/npysort/mergesort.c.src
+++ b/numpy/core/src/npysort/mergesort.c.src
@@ -112,7 +112,6 @@ mergesort_@suff@(@type@ *start, npy_intp num, void *NOT_USED)
     pr = pl + num;
     pw = (@type@ *) malloc((num/2) * sizeof(@type@));
     if (pw == NULL) {
-        PyErr_NoMemory();
         return -NPY_ENOMEM;
     }
     mergesort0_@suff@(pl, pr, pw);
@@ -176,7 +175,6 @@ amergesort_@suff@(@type@ *v, npy_intp *tosort, npy_intp num, void *NOT_USED)
     pr = pl + num;
     pw = (npy_intp *) malloc((num/2) * sizeof(npy_intp));
     if (pw == NULL) {
-        PyErr_NoMemory();
         return -NPY_ENOMEM;
     }
     amergesort0_@suff@(pl, pr, v, pw);
@@ -259,13 +257,11 @@ mergesort_@suff@(@type@ *start, npy_intp num, PyArrayObject *arr)
     pr = pl + num*len;
     pw = (@type@ *) malloc((num/2) * elsize);
     if (pw == NULL) {
-        PyErr_NoMemory();
         err = -NPY_ENOMEM;
         goto fail_0;
     }
     vp = (@type@ *) malloc(elsize);
     if (vp == NULL) {
-        PyErr_NoMemory();
         err = -NPY_ENOMEM;
         goto fail_1;
     }
@@ -335,7 +331,6 @@ amergesort_@suff@(@type@ *v, npy_intp *tosort, npy_intp num, PyArrayObject *arr)
     pr = pl + num;
     pw = (npy_intp *) malloc((num/2) * sizeof(npy_intp));
     if (pw == NULL) {
-        PyErr_NoMemory();
         return -NPY_ENOMEM;
     }
     amergesort0_@suff@(pl, pr, v, pw, len);