ENH: missingdata: Make the ndarray.take function work with NA masks

5 files changed, 286 insertions, 48 deletions

diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index 16e912c95..ac3a6648b 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -236,7 +236,7 @@ _update_descr_and_dimensions(PyArray_Descr **des, npy_intp *newdims,
 
 
     newnd = oldnd + numnew;
-    if (newnd > MAX_DIMS) {
+    if (newnd > NPY_MAXDIMS) {
         goto finish;
     }
     if (tuple) {
@@ -2453,7 +2453,7 @@ PyArray_FromInterface(PyObject *input)
     char *data;
     Py_ssize_t buffer_len;
     int res, i, n;
-    intp dims[MAX_DIMS], strides[MAX_DIMS];
+    intp dims[NPY_MAXDIMS], strides[NPY_MAXDIMS];
     int dataflags = NPY_ARRAY_BEHAVED;
 
     /* Get the memory from __array_data__ and __array_offset__ */
@@ -2732,7 +2732,7 @@ PyArray_FromDimsAndDataAndDescr(int nd, int *d,
 {
     PyObject *ret;
     int i;
-    npy_intp newd[MAX_DIMS];
+    npy_intp newd[NPY_MAXDIMS];
     char msg[] = "PyArray_FromDimsAndDataAndDescr: use PyArray_NewFromDescr.";
 
     if (DEPRECATE(msg) < 0) {
@@ -3822,14 +3822,14 @@ PyArray_CheckAxis(PyArrayObject *arr, int *axis, int flags)
     PyObject *temp1, *temp2;
     int n = PyArray_NDIM(arr);
 
-    if (*axis == MAX_DIMS || n == 0) {
+    if (*axis == NPY_MAXDIMS || n == 0) {
         if (n != 1) {
             temp1 = PyArray_Ravel(arr,0);
             if (temp1 == NULL) {
                 *axis = 0;
                 return NULL;
             }
-            if (*axis == MAX_DIMS) {
+            if (*axis == NPY_MAXDIMS) {
                 *axis = PyArray_NDIM((PyArrayObject *)temp1)-1;
             }
         }
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
index 5be85aaef..8e861d460 100644
--- a/numpy/core/src/multiarray/item_selection.c
+++ b/numpy/core/src/multiarray/item_selection.c
@@ -32,10 +32,11 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
     intp nd, i, j, n, m, max_item, tmp, chunk, nelem;
     intp shape[MAX_DIMS];
     char *src, *dest;
-    int err;
+    int err, use_maskna = 0;
 
     indices = NULL;
-    self = (PyArrayObject *)PyArray_CheckAxis(self0, &axis, NPY_ARRAY_CARRAY);
+    self = (PyArrayObject *)PyArray_CheckAxis(self0, &axis,
+                                    NPY_ARRAY_CARRAY | NPY_ARRAY_ALLOWNA);
     if (self == NULL) {
         return NULL;
     }
@@ -45,6 +46,9 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
     if (indices == NULL) {
         goto fail;
     }
+
+
+
     n = m = chunk = 1;
     nd = PyArray_NDIM(self) + PyArray_NDIM(indices) - 1;
     for (i = 0; i < nd; i++) {
@@ -75,14 +79,24 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
         if (obj == NULL) {
             goto fail;
         }
+
+        /* Allocate an NA mask if necessary */
+        if (PyArray_HASMASKNA(self)) {
+            if (PyArray_AllocateMaskNA(obj, 1, 0, 1) < 0) {
+                goto fail;
+            }
+            use_maskna = 1;
+        }
     }
     else {
-        int flags = NPY_ARRAY_CARRAY | NPY_ARRAY_UPDATEIFCOPY;
+        int flags = NPY_ARRAY_CARRAY |
+                    NPY_ARRAY_UPDATEIFCOPY |
+                    NPY_ARRAY_ALLOWNA;
 
         if ((PyArray_NDIM(out) != nd) ||
             !PyArray_CompareLists(PyArray_DIMS(out), shape, nd)) {
             PyErr_SetString(PyExc_ValueError,
-                            "bad shape in output array");
+                        "output array does not match result of ndarray.take");
             goto fail;
         }
 
@@ -100,6 +114,18 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
         if (obj == NULL) {
             goto fail;
         }
+
+        if (PyArray_HASMASKNA(self)) {
+            if (PyArray_HASMASKNA(obj)) {
+                use_maskna = 1;
+            }
+            else if (PyArray_ContainsNA(self)) {
+                PyErr_SetString(PyExc_ValueError,
+                        "Cannot assign NA value to an array which "
+                        "does not support NAs");
+                goto fail;
+            }
+        }
     }
 
     max_item = PyArray_DIMS(self)[axis];
@@ -109,7 +135,121 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
     dest = PyArray_DATA(obj);
 
     func = PyArray_DESCR(self)->f->fasttake;
-    if (func == NULL) {
+    if (use_maskna) {
+        char *dst_maskna = NULL, *src_maskna = NULL;
+        npy_intp itemsize = PyArray_DESCR(obj)->elsize;
+        PyArray_MaskedStridedTransferFn *maskedstransfer = NULL;
+        NpyAuxData *transferdata = NULL;
+        int needs_api = 0;
+
+        if (PyArray_GetMaskedDTypeTransferFunction(
+                        1,
+                        itemsize,
+                        itemsize,
+                        1,
+                        PyArray_DESCR(obj),
+                        PyArray_DESCR(obj),
+                        PyArray_MASKNA_DTYPE(obj),
+                        0,
+                        &maskedstransfer, &transferdata,
+                        &needs_api) != NPY_SUCCEED) {
+            goto fail;
+        }
+
+
+        src_maskna = PyArray_MASKNA_DATA(self);
+        dst_maskna = PyArray_MASKNA_DATA(obj);
+        if (PyDataType_REFCHK(PyArray_DESCR(self))) {
+            /*
+             * TODO: Should use PyArray_GetDTypeTransferFunction
+             *       instead of raw memmove to remedy this.
+             */
+            PyErr_SetString(PyExc_RuntimeError,
+                    "ndarray.take doesn't support object arrays with "
+                    "masks yet");
+            NPY_AUXDATA_FREE(transferdata);
+            goto fail;
+        }
+
+        switch(clipmode) {
+        case NPY_RAISE:
+            for (i = 0; i < n; i++) {
+                for (j = 0; j < m; j++) {
+                    tmp = ((intp *)(PyArray_DATA(indices)))[j];
+                    if (tmp < 0) {
+                        tmp = tmp + max_item;
+                    }
+                    if ((tmp < 0) || (tmp >= max_item)) {
+                        PyErr_SetString(PyExc_IndexError,
+                                "index out of range "\
+                                "for array");
+                        NPY_AUXDATA_FREE(transferdata);
+                        goto fail;
+                    }
+                    maskedstransfer(dest, itemsize,
+                                    src + tmp*chunk, itemsize,
+                                    (npy_mask *)src_maskna, 1,
+                                    nelem, itemsize, transferdata);
+                    dest += chunk;
+                    memmove(dst_maskna, src_maskna + tmp*nelem, nelem);
+                    dst_maskna += nelem;
+                }
+                src += chunk*max_item;
+                src_maskna += nelem*max_item;
+            }
+            break;
+        case NPY_WRAP:
+            for (i = 0; i < n; i++) {
+                for (j = 0; j < m; j++) {
+                    tmp = ((intp *)(PyArray_DATA(indices)))[j];
+                    if (tmp < 0) {
+                        while (tmp < 0) {
+                            tmp += max_item;
+                        }
+                    }
+                    else if (tmp >= max_item) {
+                        while (tmp >= max_item) {
+                            tmp -= max_item;
+                        }
+                    }
+                    maskedstransfer(dest, itemsize,
+                                    src + tmp*chunk, itemsize,
+                                    (npy_mask *)src_maskna, 1,
+                                    nelem, itemsize, transferdata);
+                    dest += chunk;
+                    memmove(dst_maskna, src_maskna + tmp*nelem, nelem);
+                    dst_maskna += nelem;
+                }
+                src += chunk*max_item;
+                src_maskna += nelem*max_item;
+            }
+            break;
+        case NPY_CLIP:
+            for (i = 0; i < n; i++) {
+                for (j = 0; j < m; j++) {
+                    tmp = ((intp *)(PyArray_DATA(indices)))[j];
+                    if (tmp < 0) {
+                        tmp = 0;
+                    }
+                    else if (tmp >= max_item) {
+                        tmp = max_item - 1;
+                    }
+                    maskedstransfer(dest, itemsize,
+                                    src + tmp*chunk, itemsize,
+                                    (npy_mask *)src_maskna, 1,
+                                    nelem, itemsize, transferdata);
+                    dest += chunk;
+                    memmove(dst_maskna, src_maskna + tmp*nelem, nelem);
+                    dst_maskna += nelem;
+                }
+                src += chunk*max_item;
+                src_maskna += nelem*max_item;
+            }
+            break;
+        }
+        NPY_AUXDATA_FREE(transferdata);
+    }
+    else if (func == NULL) {
         switch(clipmode) {
         case NPY_RAISE:
             for (i = 0; i < n; i++) {
@@ -1667,14 +1807,26 @@ PyArray_Compress(PyArrayObject *self, PyObject *condition, int axis,
     PyArrayObject *cond;
     PyObject *res, *ret;
 
-    cond = (PyArrayObject *)PyArray_FROM_O(condition);
-    if (cond == NULL) {
-        return NULL;
+    if (PyArray_Check(condition)) {
+        cond = (PyArrayObject *)condition;
+        Py_INCREF(cond);
     }
+    else {
+        PyArray_Descr *dtype = PyArray_DescrFromType(NPY_BOOL);
+        if (dtype == NULL) {
+            return NULL;
+        }
+        cond = (PyArrayObject *)PyArray_FromAny(condition, dtype,
+                                    0, 0, NPY_ARRAY_ALLOWNA, NULL);
+        if (cond == NULL) {
+            return NULL;
+        }
+    }
+
     if (PyArray_NDIM(cond) != 1) {
         Py_DECREF(cond);
         PyErr_SetString(PyExc_ValueError,
-                        "condition must be 1-d array");
+                        "condition must be a 1-d array");
         return NULL;
     }
 
@@ -1742,7 +1894,8 @@ count_boolean_trues(int ndim, char *data, npy_intp *ashape, npy_intp *astrides)
 }
 
 /*NUMPY_API
- * Counts the number of non-zero elements in the array
+ * Counts the number of non-zero elements in the array. Raises
+ * an error if the array contains an NA.
  *
  * Returns -1 on error.
  */
@@ -1759,6 +1912,16 @@ PyArray_CountNonzero(PyArrayObject *self)
     char **dataptr;
     npy_intp *strideptr, *innersizeptr;
 
+    /* If 'self' has an NA mask, make sure it has no NA values */
+    if (PyArray_HASMASKNA(self)) {
+        if (PyArray_ContainsNA(self)) {
+            PyErr_SetString(PyExc_ValueError,
+                    "Cannot count the number of nonzeros in an array "
+                    "which contains an NA");
+            return -1;
+        }
+    }
+
     /* Special low-overhead version specific to the boolean type */
     if (PyArray_DESCR(self)->type_num == NPY_BOOL) {
         return count_boolean_trues(PyArray_NDIM(self), PyArray_DATA(self),
@@ -1789,9 +1952,14 @@ PyArray_CountNonzero(PyArrayObject *self)
         return 0;
     }
 
-    /* Otherwise create and use an iterator to count the nonzeros */
-    iter = NpyIter_New(self, NPY_ITER_READONLY|
-                             NPY_ITER_EXTERNAL_LOOP|
+    /*
+     * Otherwise create and use an iterator to count the nonzeros.
+     * Can ignore any NAs because we already checked PyArray_ContainsNA
+     * earlier.
+     */
+    iter = NpyIter_New(self, NPY_ITER_READONLY |
+                             NPY_ITER_IGNORE_MASKNA |
+                             NPY_ITER_EXTERNAL_LOOP |
                              NPY_ITER_REFS_OK,
                         NPY_KEEPORDER, NPY_NO_CASTING,
                         NULL);
@@ -1844,7 +2012,7 @@ PyArray_Nonzero(PyArrayObject *self)
     PyArray_NonzeroFunc *nonzero = PyArray_DESCR(self)->f->nonzero;
     char *data;
     npy_intp stride, count;
-    npy_intp nonzero_count = PyArray_CountNonzero(self);
+    npy_intp nonzero_count;
     npy_intp *multi_index;
 
     NpyIter *iter;
@@ -1852,6 +2020,16 @@ PyArray_Nonzero(PyArrayObject *self)
     NpyIter_GetMultiIndexFunc *get_multi_index;
     char **dataptr;
 
+    /*
+     * First count the number of non-zeros in 'self'. If 'self' contains
+     * an NA value, this will raise an error, so after this call
+     * we can assume 'self' contains no NAs.
+     */
+    nonzero_count = PyArray_CountNonzero(self);
+    if (nonzero_count < 0) {
+        return NULL;
+    }
+
     /* Allocate the result as a 2D array */
     ret_dims[0] = nonzero_count;
     ret_dims[1] = (ndim == 0) ? 1 : ndim;
@@ -1881,10 +2059,15 @@ PyArray_Nonzero(PyArrayObject *self)
         goto finish;
     }
 
-    /* Build an iterator tracking a multi-index, in C order */
-    iter = NpyIter_New(self, NPY_ITER_READONLY|
-                             NPY_ITER_MULTI_INDEX|
-                             NPY_ITER_ZEROSIZE_OK|
+    /*
+     * Build an iterator tracking a multi-index, in C order. We
+     * can ignore NAs because the PyArray_CountNonzero call checked
+     * that there were no NAs already.
+     */
+    iter = NpyIter_New(self, NPY_ITER_READONLY |
+                             NPY_ITER_IGNORE_MASKNA |
+                             NPY_ITER_MULTI_INDEX |
+                             NPY_ITER_ZEROSIZE_OK |
                              NPY_ITER_REFS_OK,
                         NPY_CORDER, NPY_NO_CASTING,
                         NULL);
diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
index 7f850f4dd..c728f17e0 100644
--- a/numpy/core/src/multiarray/methods.c
+++ b/numpy/core/src/multiarray/methods.c
@@ -46,9 +46,6 @@ NpyArg_ParseKeywords(PyObject *keys, const char *format, char **kwlist, ...)
     return ret;
 }
 
-/* Should only be used if x is known to be an nd-array */
-#define _ARET(x) PyArray_Return((PyArrayObject *)(x))
-
 static PyObject *
 array_take(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
@@ -65,7 +62,8 @@ array_take(PyArrayObject *self, PyObject *args, PyObject *kwds)
                                      PyArray_ClipmodeConverter, &mode))
         return NULL;
 
-    return _ARET(PyArray_TakeFrom(self, indices, dimension, out, mode));
+    return PyArray_Return((PyArrayObject *)
+                PyArray_TakeFrom(self, indices, dimension, out, mode));
 }
 
 static PyObject *
@@ -103,7 +101,7 @@ array_reshape(PyArrayObject *self, PyObject *args, PyObject *kwds)
     static char *keywords[] = {"order", NULL};
     PyArray_Dims newshape;
     PyObject *ret;
-    PyArray_ORDER order = PyArray_CORDER;
+    PyArray_ORDER order = NPY_CORDER;
     Py_ssize_t n = PyTuple_Size(args);
 
     if (!NpyArg_ParseKeywords(kwds, "|O&", keywords,
@@ -222,7 +220,7 @@ array_argmax(PyArrayObject *self, PyObject *args, PyObject *kwds)
                                      PyArray_OutputConverter, &out))
         return NULL;
 
-    return _ARET(PyArray_ArgMax(self, axis, out));
+    return PyArray_Return((PyArrayObject *)PyArray_ArgMax(self, axis, out));
 }
 
 static PyObject *
@@ -237,7 +235,7 @@ array_argmin(PyArrayObject *self, PyObject *args, PyObject *kwds)
                                      PyArray_OutputConverter, &out))
         return NULL;
 
-    return _ARET(PyArray_ArgMin(self, axis, out));
+    return PyArray_Return((PyArrayObject *)PyArray_ArgMin(self, axis, out));
 }
 
 static PyObject *
@@ -1042,7 +1040,7 @@ array_getarray(PyArrayObject *self, PyObject *args)
 static PyObject *
 array_copy(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    PyArray_ORDER order = PyArray_CORDER;
+    PyArray_ORDER order = NPY_CORDER;
     static char *kwlist[] = {"order", NULL};
 
     if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O&", kwlist,
@@ -1087,7 +1085,7 @@ array_resize(PyArrayObject *self, PyObject *args, PyObject *kwds)
         return NULL;
     }
 
-    ret = PyArray_Resize(self, &newshape, refcheck, PyArray_CORDER);
+    ret = PyArray_Resize(self, &newshape, refcheck, NPY_CORDER);
     PyDimMem_FREE(newshape.ptr);
     if (ret == NULL) {
         return NULL;
@@ -1108,7 +1106,7 @@ array_repeat(PyArrayObject *self, PyObject *args, PyObject *kwds) {
                                      PyArray_AxisConverter, &axis)) {
         return NULL;
     }
-    return _ARET(PyArray_Repeat(self, repeats, axis));
+    return PyArray_Return((PyArrayObject *)PyArray_Repeat(self, repeats, axis));
 }
 
 static PyObject *
@@ -1135,7 +1133,7 @@ array_choose(PyArrayObject *self, PyObject *args, PyObject *kwds)
         return NULL;
     }
 
-    return _ARET(PyArray_Choose(self, choices, out, clipmode));
+    return PyArray_Return((PyArrayObject *)PyArray_Choose(self, choices, out, clipmode));
 }
 
 static PyObject *
@@ -1241,7 +1239,7 @@ array_argsort(PyArrayObject *self, PyObject *args, PyObject *kwds)
         Py_XDECREF(PyArray_DESCR(self));
         ((PyArrayObject_fieldaccess *)self)->descr = saved;
     }
-    return _ARET(res);
+    return PyArray_Return((PyArrayObject *)res);
 }
 
 static PyObject *
@@ -1256,7 +1254,7 @@ array_searchsorted(PyArrayObject *self, PyObject *args, PyObject *kwds)
                                      PyArray_SearchsideConverter, &side)) {
         return NULL;
     }
-    return _ARET(PyArray_SearchSorted(self, keys, side));
+    return PyArray_Return((PyArrayObject *)PyArray_SearchSorted(self, keys, side));
 }
 
 static void
@@ -2047,7 +2045,8 @@ array_compress(PyArrayObject *self, PyObject *args, PyObject *kwds)
                                      PyArray_OutputConverter, &out)) {
         return NULL;
     }
-    return _ARET(PyArray_Compress(self, condition, axis, out));
+    return PyArray_Return(
+                (PyArrayObject *)PyArray_Compress(self, condition, axis, out));
 }
 
 
@@ -2082,7 +2081,7 @@ array_trace(PyArrayObject *self, PyObject *args, PyObject *kwds)
 
     rtype = _CHKTYPENUM(dtype);
     Py_XDECREF(dtype);
-    return _ARET(PyArray_Trace(self, offset, axis1, axis2, rtype, out));
+    return PyArray_Return((PyArrayObject *)PyArray_Trace(self, offset, axis1, axis2, rtype, out));
 }
 
 #undef _CHKTYPENUM
@@ -2105,7 +2104,7 @@ array_clip(PyArrayObject *self, PyObject *args, PyObject *kwds)
         PyErr_SetString(PyExc_ValueError, "One of max or min must be given.");
         return NULL;
     }
-    return _ARET(PyArray_Clip(self, min, max, out));
+    return PyArray_Return((PyArrayObject *)PyArray_Clip(self, min, max, out));
 }
 
 
@@ -2135,14 +2134,14 @@ array_diagonal(PyArrayObject *self, PyObject *args, PyObject *kwds)
                                      &axis2)) {
         return NULL;
     }
-    return _ARET(PyArray_Diagonal(self, offset, axis1, axis2));
+    return PyArray_Return((PyArrayObject *)PyArray_Diagonal(self, offset, axis1, axis2));
 }
 
 
 static PyObject *
 array_flatten(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    PyArray_ORDER order = PyArray_CORDER;
+    PyArray_ORDER order = NPY_CORDER;
     static char *kwlist[] = {"order", NULL};
 
     if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O&", kwlist,
@@ -2156,7 +2155,7 @@ array_flatten(PyArrayObject *self, PyObject *args, PyObject *kwds)
 static PyObject *
 array_ravel(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    PyArray_ORDER order = PyArray_CORDER;
+    PyArray_ORDER order = NPY_CORDER;
     static char *kwlist[] = {"order", NULL};
 
     if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O&", kwlist,
@@ -2179,7 +2178,7 @@ array_round(PyArrayObject *self, PyObject *args, PyObject *kwds)
                                      PyArray_OutputConverter, &out)) {
         return NULL;
     }
-    return _ARET(PyArray_Round(self, decimals, out));
+    return PyArray_Return((PyArrayObject *)PyArray_Round(self, decimals, out));
 }
 
 
@@ -2467,5 +2466,3 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = {
         METH_VARARGS | METH_KEYWORDS, NULL},
     {NULL, NULL, 0, NULL}           /* sentinel */
 };
-
-#undef _ARET
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index c7e0cfa9d..9f73139b9 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -1995,7 +1995,8 @@ array_count_nonzero(PyObject *NPY_UNUSED(self), PyObject *args)
         return NULL;
     }
 
-    array = (PyArrayObject *)PyArray_FromAny(array_in, NULL, 0, 0, 0, NULL);
+    array = (PyArrayObject *)PyArray_FromAny(array_in, NULL,
+                                        0, 0, NPY_ARRAY_ALLOWNA, NULL);
     if (array == NULL) {
         return NULL;
     }
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index 393edf74d..a7127868a 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -517,7 +517,64 @@ def test_array_maskna_view_array_assignment_1D():
 
     # TODO: fancy indexing is next...
 
-def test_ufunc_1D():
+def test_maskna_nonzero_1D():
+    a = np.zeros((5,), maskna=True)
+
+    # The nonzeros without any NAs
+    assert_equal(np.count_nonzero(a), 0)
+    assert_equal(np.nonzero(a)[0], [])
+    a[2] = 3
+    assert_equal(np.count_nonzero(a), 1)
+    assert_equal(np.nonzero(a)[0], [2])
+    a[3:] = 2
+    assert_equal(np.count_nonzero(a), 3)
+    assert_equal(np.nonzero(a)[0], [2,3,4])
+
+    # The nonzeros with an NA
+    a[2] = np.NA
+    assert_raises(ValueError, np.count_nonzero, a)
+    assert_raises(ValueError, np.nonzero, a)
+
+def test_maskna_take_1D():
+    a = np.arange(5, maskna=True)
+    b = np.arange(3)
+    c = b.view(maskna=True)
+
+    # Take without any NAs
+    assert_equal(a.take([0,2,4]), [0,2,4])
+
+    # Take without any NAs, into non-NA output parameter
+    a.take([0,2,4], out=b)
+    assert_equal(b, [0,2,4])
+
+    # Take without any NAs, into NA output parameter
+    b[...] = 1
+    c[...] = np.NA
+    a.take([0,2,4], out=c)
+    assert_equal(c, [0,2,4])
+
+    # Take with some NAs
+    a[2] = np.NA
+    a[3] = np.NA
+    ret = a.take([0,2,4])
+    assert_equal([ret[0], ret[2]], [0,4])
+    assert_equal(np.isna(ret), [0,1,0])
+
+    # Take with some NAs, into NA output parameter
+    b[...] = 1
+    c[...] = np.NA
+    a.take([0,2,4], out=c)
+    assert_equal(b, [0,1,4])
+    assert_equal([c[0], c[2]], [0,4])
+    assert_equal(np.isna(c), [0,1,0])
+
+    c[...] = 1
+    a.take([0,2,4], out=c)
+    assert_equal(b, [0,1,4])
+    assert_equal([c[0], c[2]], [0,4])
+    assert_equal(np.isna(c), [0,1,0])
+
+def test_maskna_ufunc_1D():
     a = np.arange(3, maskna=True)
     b = np.arange(3)
 
@@ -526,7 +583,7 @@ def test_ufunc_1D():
     assert_(c.flags.maskna)
     #assert_equal(c, [0,2,4])
 
-def test_ufunc_add_reduce_1D():
+def test_maskna_ufunc_add_reduce_1D():
     a = np.arange(3.0, maskna=True)
     b = np.array(0.5)
     c_orig = np.array(0.5)