ENH: missingdata: Implemented boolean assignment, working with NA masks

9 files changed, 472 insertions, 44 deletions

diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index 53eb21cdc..47b021e08 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -1586,6 +1586,34 @@ PyArray_GetArrayParamsFromObjectEx(PyObject *op,
         return 0;
     }
 
+    /* If op is a numpy.NA */
+    if (NpyNA_Check(op)) {
+        NpyNA_fields *fna = (NpyNA_fields *)op;
+
+        if (writeable) {
+            PyErr_SetString(PyExc_RuntimeError,
+                                "cannot write to numpy.NA");
+            return -1;
+        }
+        /* Use the NA's dtype if available */
+        if (fna->dtype != NULL) {
+            *out_dtype = fna->dtype;
+            Py_INCREF(*out_dtype);
+        }
+        /* Otherwise use the default NumPy dtype */
+        else {
+            *out_dtype = PyArray_DescrFromType(NPY_DEFAULT_TYPE);
+            if (*out_dtype == NULL) {
+                return -1;
+            }
+        }
+        *out_ndim = 0;
+        *out_arr = NULL;
+        *out_contains_na = 1;
+        return 0;
+
+    }
+
     /* If op supports the PEP 3118 buffer interface */
     if (!PyBytes_Check(op) && !PyUnicode_Check(op) &&
              _array_from_buffer_3118(op, (PyObject **)out_arr) == 0) {
@@ -1891,6 +1919,15 @@ PyArray_FromAny(PyObject *op, PyArray_Descr *newtype, int min_depth,
 
     /* If we got dimensions and dtype instead of an array */
     if (arr == NULL) {
+        /*
+         * If the input data is an NA object, and the ALLOWNA flag is
+         * enabled, produce an array with an NA mask.
+         */
+        if (contains_na && (flags & NPY_ARRAY_ALLOWNA) != 0 &&
+                                    NpyNA_Check(op)) {
+            flags |= NPY_ARRAY_MASKNA;
+        }
+
         if (flags & NPY_ARRAY_UPDATEIFCOPY) {
             Py_XDECREF(newtype);
             PyErr_SetString(PyExc_TypeError,
@@ -1962,6 +1999,10 @@ PyArray_FromAny(PyObject *op, PyArray_Descr *newtype, int min_depth,
                                          ndim, dims,
                                          NULL, NULL,
                                          flags&NPY_ARRAY_F_CONTIGUOUS, NULL);
+            if (ret == NULL) {
+                return NULL;
+            }
+
             /*
              * Add an NA mask if requested, or if allowed and the data
              * has NAs
@@ -1973,20 +2014,29 @@ PyArray_FromAny(PyObject *op, PyArray_Descr *newtype, int min_depth,
                     Py_DECREF(ret);
                     return NULL;
                 }
-            }
-            if (ret != NULL) {
-                if (ndim > 0) {
-                    if (PyArray_AssignFromSequence(ret, op) < 0) {
-                        Py_DECREF(ret);
-                        ret = NULL;
+
+                /* Special case assigning a single NA */
+                if (ndim == 0) {
+                    NpyNA *na = NpyNA_FromObject(op, 1);
+                    if (na != NULL) {
+                        PyArray_MASKNA_DATA(ret)[0] =
+                                        (char)NpyNA_AsMaskValue(na);
+                        return (PyObject *)ret;
                     }
                 }
-                else {
-                    if (PyArray_DESCR(ret)->f->setitem(op,
-                                                PyArray_DATA(ret), ret) < 0) {
-                        Py_DECREF(ret);
-                        ret = NULL;
-                    }
+            }
+
+            if (ndim > 0) {
+                if (PyArray_AssignFromSequence(ret, op) < 0) {
+                    Py_DECREF(ret);
+                    ret = NULL;
+                }
+            }
+            else {
+                if (PyArray_DESCR(ret)->f->setitem(op,
+                                            PyArray_DATA(ret), ret) < 0) {
+                    Py_DECREF(ret);
+                    ret = NULL;
                 }
             }
         }
@@ -2005,7 +2055,16 @@ PyArray_FromAny(PyObject *op, PyArray_Descr *newtype, int min_depth,
             ret = NULL;
         }
         else {
-            ret = (PyArrayObject *)PyArray_FromArray(arr, newtype, flags);
+            if (PyArray_HASMASKNA((PyArrayObject *)arr) &&
+                                (flags & NPY_ARRAY_ALLOWNA) == 0) {
+                PyErr_SetString(PyExc_ValueError,
+                        "this operation does not support "
+                        "arrays with NA masks");
+                ret = NULL;
+            }
+            else {
+                ret = (PyArrayObject *)PyArray_FromArray(arr, newtype, flags);
+            }
             Py_DECREF(arr);
         }
     }
diff --git a/numpy/core/src/multiarray/dtype_transfer.c b/numpy/core/src/multiarray/dtype_transfer.c
index 5ffd93886..023037abe 100644
--- a/numpy/core/src/multiarray/dtype_transfer.c
+++ b/numpy/core/src/multiarray/dtype_transfer.c
@@ -3897,6 +3897,30 @@ PyArray_PrepareOneRawArrayIter(int ndim, char *data,
     _npy_stride_sort_item strideperm[NPY_MAXDIMS];
     int i, j;
 
+    /* Special case 0 and 1 dimensions */
+    if (ndim == 0) {
+        *out_ndim = 1;
+        *out_data = data;
+        out_shape[0] = 1;
+        out_strides[0] = 0;
+        return 0;
+    }
+    else if (ndim == 1) {
+        npy_intp stride_entry = strides[0], shape_entry = shape[0];
+        *out_ndim = 1;
+        out_shape[0] = shape[0];
+        /* Always make a positive stride */
+        if (stride_entry >= 0) {
+            *out_data = data;
+            out_strides[0] = stride_entry;
+        }
+        else {
+            *out_data = data + stride_entry * (shape_entry - 1);
+            out_strides[0] = -stride_entry;
+        }
+        return 0;
+    }
+
     /* Sort the axes based on the destination strides */
     PyArray_CreateSortedStridePerm(ndim, strides, strideperm);
     for (i = 0; i < ndim; ++i) {
@@ -3913,6 +3937,14 @@ PyArray_PrepareOneRawArrayIter(int ndim, char *data,
             data += stride_entry * (shape_entry - 1);
             out_strides[i] = -stride_entry;
         }
+        /* Detect 0-size arrays here */
+        if (shape_entry == 0) {
+            *out_ndim = 1;
+            *out_data = data;
+            out_shape[0] = 0;
+            out_strides[0] = 0;
+            return 0;
+        }
     }
 
     /* Coalesce any dimensions where possible */
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
index 4ba7b7881..c0290174f 100644
--- a/numpy/core/src/multiarray/item_selection.c
+++ b/numpy/core/src/multiarray/item_selection.c
@@ -1712,6 +1712,11 @@ count_boolean_trues(int ndim, char *data, npy_intp *ashape, npy_intp *astrides)
         return -1;
     }
 
+    /* Handle zero-sized array */
+    if (shape[0] == 0) {
+        return 0;
+    }
+
     /* Do the iteration */
     memset(coord, 0, ndim * sizeof(npy_intp));
     /* Special case for contiguous inner loop */
diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index 10e3f82b2..cf294454d 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -686,13 +686,15 @@ array_subscript_simple(PyArrayObject *self, PyObject *op)
  * though this function allows different choices.
  */
 NPY_NO_EXPORT PyArrayObject *
-array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER order)
+array_boolean_subscript(PyArrayObject *self,
+                    PyArrayObject *bmask, NPY_ORDER order)
 {
     npy_intp size, itemsize;
     char *ret_data, *ret_maskna_data = NULL;
     PyArray_Descr *dtype;
     PyArrayObject *ret;
     int self_has_maskna = PyArray_HASMASKNA(self), needs_api = 0;
+    npy_intp bmask_size;
 
     if (PyArray_DESCR(bmask)->type_num != NPY_BOOL) {
         PyErr_SetString(PyExc_TypeError,
@@ -703,7 +705,8 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
     /* See the Boolean Indexing section of the missing data NEP */
     if (PyArray_ContainsNA(bmask)) {
         PyErr_SetString(PyExc_ValueError,
-                "The boolean mask array may not contain any NA values");
+                "The boolean mask indexing array "
+                "may not contain any NA values");
         return NULL;
     }
 
@@ -715,7 +718,10 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
     size = count_boolean_trues(PyArray_NDIM(bmask), PyArray_DATA(bmask),
                                 PyArray_DIMS(bmask), PyArray_STRIDES(bmask));
     /* Correction factor for broadcasting 'bmask' to 'self' */
-    size *= PyArray_SIZE(self) / PyArray_SIZE(bmask);
+    bmask_size = PyArray_SIZE(bmask);
+    if (bmask_size > 0) {
+        size *= PyArray_SIZE(self) / bmask_size;
+    }
 
     /* Allocate the output of the boolean indexing */
     dtype = PyArray_DESCR(self);
@@ -752,12 +758,15 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
         char **dataptrs;
 
         /* Set up the iterator */
-        flags = NPY_ITER_EXTERNAL_LOOP;
+        flags = NPY_ITER_EXTERNAL_LOOP | NPY_ITER_REFS_OK;
         if (self_has_maskna) {
-            op_flags[0] = NPY_ITER_READONLY | NPY_ITER_USE_MASKNA;
+            op_flags[0] = NPY_ITER_READONLY |
+                          NPY_ITER_NO_BROADCAST |
+                          NPY_ITER_USE_MASKNA;
         }
         else {
-            op_flags[0] = NPY_ITER_READONLY;
+            op_flags[0] = NPY_ITER_READONLY |
+                          NPY_ITER_NO_BROADCAST;
         }
         /*
          * Since we already checked PyArray_ContainsNA(bmask), can
@@ -781,6 +790,7 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
                         &stransfer, &transferdata,
                         &needs_api) != NPY_SUCCEED) {
             Py_DECREF(ret);
+            NpyIter_Deallocate(iter);
             return NULL;
         }
 
@@ -788,6 +798,7 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
         iternext = NpyIter_GetIterNext(iter, NULL);
         if (iternext == NULL) {
             Py_DECREF(ret);
+            NpyIter_Deallocate(iter);
             NPY_AUXDATA_FREE(transferdata);
             return NULL;
         }
@@ -823,6 +834,7 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
                                 subloopsize, itemsize, transferdata);
                     innersize -= subloopsize;
                     self_data += subloopsize * self_stride;
+                    ret_data += subloopsize * itemsize;
                 }
             } while (iternext(iter));
         }
@@ -848,6 +860,7 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
                     }
                     innersize -= subloopsize;
                     self_data += subloopsize * self_stride;
+                    maskna_data += subloopsize * maskna_stride;
                     /* Process unmasked values */
                     subloopsize = 0;
                     while (subloopsize < innersize && *bmask_data != 0) {
@@ -862,13 +875,15 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
                      */
                     stransfer(ret_data, itemsize, self_data, self_stride,
                                 subloopsize, itemsize, transferdata);
-                    /* Copy the mask */
+                    /* Copy the mask as well */
                     for (i = 0; i < subloopsize; ++i) {
-                        *ret_maskna_data++ = *maskna_data;
+                        *ret_maskna_data = *maskna_data;
+                        ++ret_maskna_data;
                         maskna_data += maskna_stride;
                     }
                     innersize -= subloopsize;
                     self_data += subloopsize * self_stride;
+                    ret_data += subloopsize * itemsize;
                 }
             } while (iternext(iter));
         }
@@ -880,6 +895,298 @@ array_boolean_subscript(PyArrayObject *self, PyArrayObject *bmask, NPY_ORDER ord
     return ret;
 }
 
+/*
+ * Implements boolean indexing assignment. This takes the one-dimensional
+ * array 'v' and assigns its values to all of the elements of 'self' for which
+ * the corresponding element of 'op' is True.
+ *
+ * This operation is somewhat unfortunate, because to match up with
+ * a one-dimensional output array, it has to choose a particular
+ * iteration order, in the case of NumPy that is always C order even
+ * though this function allows different choices.
+ *
+ * Returns 0 on success, -1 on failure.
+ */
+NPY_NO_EXPORT int
+array_ass_boolean_subscript(PyArrayObject *self,
+                    PyArrayObject *bmask, PyArrayObject *v, NPY_ORDER order)
+{
+    npy_intp size, src_itemsize, v_stride, v_maskna_stride = 0;
+    char *v_data, *v_maskna_data = NULL;
+    int self_has_maskna = PyArray_HASMASKNA(self);
+    int v_has_maskna = PyArray_HASMASKNA(v);
+    int needs_api = 0;
+    npy_intp bmask_size;
+    char constant_valid_mask = 1;
+
+    if (PyArray_DESCR(bmask)->type_num != NPY_BOOL) {
+        PyErr_SetString(PyExc_TypeError,
+                "NumPy boolean array indexing assignment "
+                "requires a boolean index");
+        return -1;
+    }
+
+    if (PyArray_NDIM(v) > 1) {
+        PyErr_SetString(PyExc_TypeError,
+                "NumPy boolean array indexing assignment "
+                "requires a 0 or 1-dimensional input");
+        return -1;
+    }
+
+    /* See the Boolean Indexing section of the missing data NEP */
+    if (PyArray_ContainsNA(bmask)) {
+        PyErr_SetString(PyExc_ValueError,
+                "The boolean mask assignment indexing array "
+                "may not contain any NA values");
+        return -1;
+    }
+
+    /* Can't assign an NA to an array which doesn't support it */
+    if (v_has_maskna && !self_has_maskna) {
+        if (PyArray_ContainsNA(v)) {
+            PyErr_SetString(PyExc_ValueError,
+                    "Cannot assign NA value to an array which "
+                    "does not support NAs");
+            return -1;
+        }
+        /* If there are no actual NAs, allow the assignment */
+        else {
+            v_has_maskna = 0;
+        }
+    }
+
+    /*
+     * Since we've checked that the mask contains no NAs, we
+     * can do a straightforward count of the boolean True values
+     * in the raw mask data array.
+     */
+    size = count_boolean_trues(PyArray_NDIM(bmask), PyArray_DATA(bmask),
+                                PyArray_DIMS(bmask), PyArray_STRIDES(bmask));
+    /* Correction factor for broadcasting 'bmask' to 'self' */
+    bmask_size = PyArray_SIZE(bmask);
+    if (bmask_size > 0) {
+        size *= PyArray_SIZE(self) / bmask_size;
+    }
+
+    /* Tweak the strides for 0-dim and broadcasting cases */
+    if (PyArray_NDIM(v) > 0 && PyArray_DIMS(v)[0] > 1) {
+        v_stride = PyArray_STRIDES(v)[0];
+        if (v_has_maskna) {
+            v_maskna_stride = PyArray_MASKNA_STRIDES(v)[0];
+        }
+        else {
+            v_maskna_stride = 0;
+        }
+
+        if (size != PyArray_DIMS(v)[0]) {
+            PyErr_Format(PyExc_TypeError,
+                    "NumPy boolean array indexing assignment "
+                    "requires an input shape matching the number of True "
+                    "values in the mask (%d) != (%d)",
+                    (int)PyArray_DIMS(v)[0], (int)size);
+            return -1;
+        }
+    }
+    else {
+        v_stride = 0;
+        v_maskna_stride = 0;
+    }
+
+    src_itemsize = PyArray_DESCR(v)->elsize;
+    v_data = PyArray_DATA(v);
+    if (v_has_maskna) {
+        v_maskna_data = PyArray_MASKNA_DATA(v);
+    }
+    /* If assigning unmasked to masked, use a 0-stride all valid mask */
+    else if (self_has_maskna) {
+        v_maskna_data = &constant_valid_mask;
+    }
+
+    /* Create an iterator for the data */
+    if (size > 0) {
+        NpyIter *iter;
+        PyArrayObject *op[2] = {self, bmask};
+        npy_uint32 flags, op_flags[2];
+        npy_intp fixed_strides[3];
+
+        NpyIter_IterNextFunc *iternext;
+        npy_intp innersize, *innerstrides;
+        char **dataptrs;
+
+        /* Set up the iterator */
+        flags = NPY_ITER_EXTERNAL_LOOP | NPY_ITER_REFS_OK;
+        if (self_has_maskna) {
+            op_flags[0] = NPY_ITER_WRITEONLY |
+                          NPY_ITER_NO_BROADCAST |
+                          NPY_ITER_USE_MASKNA;
+        }
+        else {
+            op_flags[0] = NPY_ITER_WRITEONLY |
+                          NPY_ITER_NO_BROADCAST;
+        }
+        /*
+         * Since we already checked PyArray_ContainsNA(bmask), can
+         * ignore any MASKNA of bmask.
+         */
+        op_flags[1] = NPY_ITER_READONLY | NPY_ITER_IGNORE_MASKNA;
+
+        iter = NpyIter_MultiNew(2, op, flags, order, NPY_NO_CASTING,
+                                op_flags, NULL);
+        if (iter == NULL) {
+            return -1;
+        }
+
+        /* Get the values needed for the inner loop */
+        iternext = NpyIter_GetIterNext(iter, NULL);
+        if (iternext == NULL) {
+            NpyIter_Deallocate(iter);
+            return -1;
+        }
+        innerstrides = NpyIter_GetInnerStrideArray(iter);
+        dataptrs = NpyIter_GetDataPtrArray(iter);
+
+        /* Regular inner loop */
+        if (!self_has_maskna) {
+            PyArray_StridedTransferFn *stransfer = NULL;
+            NpyAuxData *transferdata = NULL;
+            npy_intp self_stride = innerstrides[0];
+            npy_intp bmask_stride = innerstrides[1];
+            npy_intp subloopsize;
+
+            /* Get a dtype transfer function */
+            NpyIter_GetInnerFixedStrideArray(iter, fixed_strides);
+            if (PyArray_GetDTypeTransferFunction(
+                            PyArray_ISALIGNED(self) && PyArray_ISALIGNED(v),
+                            v_stride, fixed_strides[0],
+                            PyArray_DESCR(v), PyArray_DESCR(self),
+                            0,
+                            &stransfer, &transferdata,
+                            &needs_api) != NPY_SUCCEED) {
+                NpyIter_Deallocate(iter);
+                return -1;
+            }
+
+            do {
+                innersize = *NpyIter_GetInnerLoopSizePtr(iter);
+                char *self_data = dataptrs[0];
+                char *bmask_data = dataptrs[1];
+
+                while (innersize > 0) {
+                    /* Skip masked values */
+                    subloopsize = 0;
+                    while (subloopsize < innersize && *bmask_data == 0) {
+                        ++subloopsize;
+                        bmask_data += bmask_stride;
+                    }
+                    innersize -= subloopsize;
+                    self_data += subloopsize * self_stride;
+                    /* Process unmasked values */
+                    subloopsize = 0;
+                    while (subloopsize < innersize && *bmask_data != 0) {
+                        ++subloopsize;
+                        bmask_data += bmask_stride;
+                    }
+                    stransfer(self_data, self_stride, v_data, v_stride,
+                                subloopsize, src_itemsize, transferdata);
+                    innersize -= subloopsize;
+                    self_data += subloopsize * self_stride;
+                    v_data += subloopsize * v_stride;
+                }
+            } while (iternext(iter));
+
+            NPY_AUXDATA_FREE(transferdata);
+        }
+        /* NA masked inner loop */
+        else {
+            PyArray_MaskedStridedTransferFn *stransfer = NULL;
+            NpyAuxData *transferdata = NULL;
+            npy_intp i;
+            npy_intp self_stride = innerstrides[0];
+            npy_intp bmask_stride = innerstrides[1];
+            npy_intp self_maskna_stride = innerstrides[2];
+            npy_intp subloopsize;
+            PyArray_Descr *v_maskna_dtype;
+
+            if (PyArray_HASMASKNA(v)) {
+                v_maskna_dtype = PyArray_MASKNA_DTYPE(v);
+                Py_INCREF(v_maskna_dtype);
+            }
+            else {
+                v_maskna_dtype = PyArray_DescrFromType(NPY_BOOL);
+                if (v_maskna_dtype == NULL) {
+                    NpyIter_Deallocate(iter);
+                    return -1;
+                }
+            }
+
+            /* Get a dtype transfer function */
+            NpyIter_GetInnerFixedStrideArray(iter, fixed_strides);
+            if (PyArray_GetMaskedDTypeTransferFunction(
+                            PyArray_ISALIGNED(self) && PyArray_ISALIGNED(v),
+                            v_stride, fixed_strides[0], v_maskna_stride,
+                            PyArray_DESCR(v),
+                            PyArray_DESCR(self),
+                            v_maskna_dtype,
+                            0,
+                            &stransfer, &transferdata,
+                            &needs_api) != NPY_SUCCEED) {
+                Py_DECREF(v_maskna_dtype);
+                NpyIter_Deallocate(iter);
+                return -1;
+            }
+            Py_DECREF(v_maskna_dtype);
+
+            do {
+                innersize = *NpyIter_GetInnerLoopSizePtr(iter);
+                char *self_data = dataptrs[0];
+                char *bmask_data = dataptrs[1];
+                char *self_maskna_data = dataptrs[2];
+
+                while (innersize > 0) {
+                    /* Skip masked values */
+                    subloopsize = 0;
+                    while (subloopsize < innersize && *bmask_data == 0) {
+                        ++subloopsize;
+                        bmask_data += bmask_stride;
+                    }
+                    innersize -= subloopsize;
+                    self_data += subloopsize * self_stride;
+                    self_maskna_data += subloopsize * self_maskna_stride;
+                    /* Process unmasked values */
+                    subloopsize = 0;
+                    while (subloopsize < innersize && *bmask_data != 0) {
+                        ++subloopsize;
+                        bmask_data += bmask_stride;
+                    }
+                    /*
+                     * Because we're assigning to an existing array,
+                     * we have to be careful about not overwriting
+                     * NA masked values.
+                     */
+                    stransfer(self_data, self_stride, v_data, v_stride,
+                                (npy_mask *)v_maskna_data, v_maskna_stride,
+                                subloopsize, src_itemsize, transferdata);
+                    /* Copy the mask as well */
+                    for (i = 0; i < subloopsize; ++i) {
+                        *self_maskna_data = *v_maskna_data;
+                        self_maskna_data += self_maskna_stride;
+                        v_maskna_data += v_maskna_stride;
+                    }
+                    innersize -= subloopsize;
+                    self_data += subloopsize * self_stride;
+                    v_data += subloopsize * v_stride;
+                }
+            } while (iternext(iter));
+
+            NPY_AUXDATA_FREE(transferdata);
+        }
+
+        NpyIter_Deallocate(iter);
+    }
+
+    return 0;
+}
+
 NPY_NO_EXPORT PyObject *
 array_subscript(PyArrayObject *self, PyObject *op)
 {
@@ -991,6 +1298,12 @@ array_subscript(PyArrayObject *self, PyObject *op)
         return NULL;
     }
 
+    /* Boolean indexing */
+    if (PyArray_Check(op) && (PyArray_TYPE((PyArrayObject *)op) == NPY_BOOL)) {
+        return (PyObject *)array_boolean_subscript(self,
+                                        (PyArrayObject *)op, NPY_CORDER);
+    }
+
     fancy = fancy_indexing_check(op);
     if (fancy != SOBJ_NOTFANCY) {
         int oned;
@@ -1238,6 +1551,27 @@ array_ass_sub(PyArrayObject *self, PyObject *index, PyObject *op)
     }
     PyErr_Clear();
 
+    /* Boolean indexing */
+    if (PyArray_Check(index) &&
+                (PyArray_TYPE((PyArrayObject *)index) == NPY_BOOL)) {
+        PyArrayObject *op_arr;
+        PyArray_Descr *dtype = NULL;
+        /* If it's an NA with no dtype, specify it explicitly */
+        if (NpyNA_Check(op) && ((NpyNA_fields *)op)->dtype == NULL) {
+            dtype = PyArray_DESCR(self);
+            Py_INCREF(dtype);
+        }
+        op_arr = (PyArrayObject *)PyArray_FromAny(op, dtype, 0, 0,
+                                                      NPY_ARRAY_ALLOWNA, NULL);
+        if (op_arr == NULL) {
+            return -1;
+        }
+
+        return array_ass_boolean_subscript(self,
+                        (PyArrayObject *)index,
+                        (PyArrayObject *)op_arr, NPY_CORDER);
+    }
+
     fancy = fancy_indexing_check(index);
     if (fancy != SOBJ_NOTFANCY) {
         oned = ((PyArray_NDIM(self) == 1) &&
@@ -1408,7 +1742,7 @@ _nonzero_indices(PyObject *myBool, PyArrayIterObject **iters)
     PyArrayObject *ba = NULL, *new = NULL;
     int nd, j;
     npy_intp size, i, count;
-    Bool *ptr;
+    npy_bool *ptr;
     npy_intp coords[NPY_MAXDIMS], dims_m1[NPY_MAXDIMS];
     npy_intp *dptr[NPY_MAXDIMS];
 
@@ -1860,25 +2194,7 @@ PyArray_MapIterNew(PyObject *indexobj, int oned, int fancy)
      */
 
     /* convert all inputs to iterators */
-    if (PyArray_Check(indexobj) &&
-                    (PyArray_TYPE((PyArrayObject *)indexobj) == NPY_BOOL)) {
-        mit->numiter = _nonzero_indices(indexobj, mit->iters);
-        if (mit->numiter < 0) {
-            goto fail;
-        }
-        mit->nd = 1;
-        mit->dimensions[0] = mit->iters[0]->dims_m1[0]+1;
-        Py_DECREF(mit->indexobj);
-        mit->indexobj = PyTuple_New(mit->numiter);
-        if (mit->indexobj == NULL) {
-            goto fail;
-        }
-        for (i = 0; i < mit->numiter; i++) {
-            PyTuple_SET_ITEM(mit->indexobj, i, PyInt_FromLong(0));
-        }
-    }
-
-    else if (PyArray_Check(indexobj) || !PyTuple_Check(indexobj)) {
+    if (PyArray_Check(indexobj) || !PyTuple_Check(indexobj)) {
         mit->numiter = 1;
         indtype = PyArray_DescrFromType(NPY_INTP);
         arr = (PyArrayObject *)PyArray_FromAny(indexobj, indtype, 0, 0,
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index 5e4401c34..e0bc2ba91 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -463,7 +463,7 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,
      * reference arrays and flag it if so.
      */
     if (flags & NPY_ITER_REFS_OK) {
-        for (iop = 0; iop < nop; ++iop) {
+        for (iop = 0; iop < first_maskna_op; ++iop) {
             PyArray_Descr *rdt = op_dtype[iop];
             if ((rdt->flags & (NPY_ITEM_REFCOUNT |
                                      NPY_ITEM_IS_POINTER |
diff --git a/numpy/core/tests/test_na.py b/numpy/core/tests/test_na.py
index 98ad67f27..917803ff6 100644
--- a/numpy/core/tests/test_na.py
+++ b/numpy/core/tests/test_na.py
@@ -180,6 +180,11 @@ def test_array_maskna_construction():
     assert_equal(a.dtype, np.dtype('O'))
     assert_equal(type(a[2]), np.NAType)
 
+    # From np.NA as a straight scalar
+    a = np.array(np.NA, maskna=True)
+    assert_equal(type(a), np.ndarray)
+    assert_(np.isna(a))
+
 def test_isna():
     # Objects which are not np.NA or ndarray all return False
     assert_equal(np.isna(True), False)
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index 75ecd398a..db3a879ca 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -549,16 +549,20 @@ class TestFromiter(TestCase):
 class TestNonzero(TestCase):
     def test_nonzero_trivial(self):
         assert_equal(np.count_nonzero(array([])), 0)
+        assert_equal(np.count_nonzero(array([], dtype='?')), 0)
         assert_equal(np.nonzero(array([])), ([],))
 
         assert_equal(np.count_nonzero(array(0)), 0)
+        assert_equal(np.count_nonzero(array(0, dtype='?')), 0)
         assert_equal(np.nonzero(array(0)), ([],))
         assert_equal(np.count_nonzero(array(1)), 1)
+        assert_equal(np.count_nonzero(array(1, dtype='?')), 1)
         assert_equal(np.nonzero(array(1)), ([0],))
 
     def test_nonzero_onedim(self):
         x = array([1,0,2,-1,0,0,8])
         assert_equal(np.count_nonzero(x), 4)
+        assert_equal(np.count_nonzero(x), 4)
         assert_equal(np.nonzero(x), ([0, 2, 3, 6],))
 
         x = array([(1,2),(0,0),(1,1),(-1,3),(0,7)],
@@ -599,6 +603,13 @@ class TestIndex(TestCase):
         V[g1,g2] = -V[g1,g2]
         assert_((array([a[0][V>0],a[1][V>0],a[2][V>0]]) == a[:,V>0]).all())
 
+    def test_boolean_edgecase(self):
+        a = np.array([], dtype='int32')
+        b = np.array([], dtype='bool')
+        c = a[b]
+        assert_equal(c, [])
+        assert_equal(c.dtype, np.dtype('int32'))
+
 
 class TestBinaryRepr(TestCase):
     def test_zero(self):
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py
index 98597bc46..c4441d8a5 100644
--- a/numpy/lib/arraysetops.py
+++ b/numpy/lib/arraysetops.py
@@ -229,7 +229,7 @@ def intersect1d(ar1, ar2, assume_unique=False):
         ar2 = unique(ar2)
     aux = np.concatenate( (ar1, ar2) )
     aux.sort()
-    return aux[aux[1:] == aux[:-1]]
+    return aux[:-1][aux[1:] == aux[:-1]]
 
 def setxor1d(ar1, ar2, assume_unique=False):
     """
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index 27abcb2c1..bbc17d165 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -1090,7 +1090,7 @@ def intersect1d(ar1, ar2, assume_unique=False):
         # Might be faster than unique( intersect1d( ar1, ar2 ) )?
         aux = ma.concatenate((unique(ar1), unique(ar2)))
     aux.sort()
-    return aux[aux[1:] == aux[:-1]]
+    return aux[:-1][aux[1:] == aux[:-1]]
 
 
 def setxor1d(ar1, ar2, assume_unique=False):