ENH: missingdata: In progress exposing USE_MASKNA to Python numpy.nditer

4 files changed, 246 insertions, 38 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index ec8388311..49d38c32f 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -324,6 +324,8 @@ multiarray_funcs_api = {
     'PyArray_HasNASupport':                 284,
     'PyArray_ContainsNA':                   285,
     'PyArray_AllocateMaskNA':               286,
+    'NpyIter_GetFirstMaskNAOp':             288,
+    'NpyIter_GetMaskNAIndices':             289,
 }
 
 ufunc_types_api = {
diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c
index 96715acd8..ddd227bf6 100644
--- a/numpy/core/src/multiarray/nditer_api.c
+++ b/numpy/core/src/multiarray/nditer_api.c
@@ -272,7 +272,9 @@ NpyIter_Reset(NpyIter *iter, char **errmsg)
 }
 
 /*NUMPY_API
- * Resets the iterator to its initial state, with new base data pointers
+ * Resets the iterator to its initial state, with new base data pointers.
+ * This function requires great caution, even more so if any
+ * NPY_ITER_USE_MASKNA operands were specified.
  *
  * If errmsg is non-NULL, it should point to a variable which will
  * receive the error message, and no Python exception will be set.
@@ -748,6 +750,35 @@ NpyIter_GetNOp(NpyIter *iter)
 }
 
 /*NUMPY_API
+ * Gets the index of the first operand which is the
+ * mask for an NPY_ITER_USE_MASKNA operand.
+ */
+NPY_NO_EXPORT int
+NpyIter_GetFirstMaskNAOp(NpyIter *iter)
+{
+    return NIT_FIRST_MASKNA_OP(iter);
+}
+
+/*NUMPY_API
+ * Gets the correspondences between the operands with
+ * NPY_ITER_USEMASKNA set and their corresponding masks.
+ *
+ * If i < NpyIter_GetFirstMaskNAOp(iter), then
+ * NpyIter_GetMaskNAIndices(iter)[i] is either -1 or
+ * an index >= NpyIter_GetFirstMaskNAOp(iter) of the corresponding
+ * mask.
+ *
+ * If i >= NpyIter_GetFirstMaskNAOp(iter), then
+ * NpyIter_GetMaskNAIndices(iter)[i] is the index
+ * of the corresponding maskna operand for the mask.
+ */
+NPY_NO_EXPORT npy_int8 *
+NpyIter_GetMaskNAIndices(NpyIter *iter)
+{
+    return NIT_MASKNA_INDICES(iter);
+}
+
+/*NUMPY_API
  * Gets the number of elements being iterated
  */
 NPY_NO_EXPORT npy_intp
@@ -1004,6 +1035,7 @@ NpyIter_GetIterView(NpyIter *iter, npy_intp i)
     npy_uint32 itflags = NIT_ITFLAGS(iter);
     int idim, ndim = NIT_NDIM(iter);
     int nop = NIT_NOP(iter);
+    int first_maskna_op = NIT_FIRST_MASKNA_OP(iter);
 
     npy_intp shape[NPY_MAXDIMS], strides[NPY_MAXDIMS];
     PyArrayObject *obj, *view;
@@ -1012,8 +1044,9 @@ NpyIter_GetIterView(NpyIter *iter, npy_intp i)
     NpyIter_AxisData *axisdata;
     npy_intp sizeof_axisdata;
     int writeable;
+    npy_int8 *maskna_indices = NIT_MASKNA_INDICES(iter);
 
-    if (i < 0 || i >= nop) {
+    if (i < 0 || i >= first_maskna_op) {
         PyErr_SetString(PyExc_IndexError,
                 "index provided for an iterator view was out of bounds");
         return NULL;
@@ -1034,9 +1067,11 @@ NpyIter_GetIterView(NpyIter *iter, npy_intp i)
     sizeof_axisdata = NIT_AXISDATA_SIZEOF(itflags, ndim, nop);
 
     /* Retrieve the shape and strides from the axisdata */
-    for (idim = 0; idim < ndim; ++idim, NIT_ADVANCE_AXISDATA(axisdata, 1)) {
+    for (idim = 0; idim < ndim; ++idim) {
         shape[ndim-idim-1] = NAD_SHAPE(axisdata);
         strides[ndim-idim-1] = NAD_STRIDES(axisdata)[i];
+
+        NIT_ADVANCE_AXISDATA(axisdata, 1);
     }
 
     Py_INCREF(dtype);
@@ -1055,6 +1090,29 @@ NpyIter_GetIterView(NpyIter *iter, npy_intp i)
     }
     /* Make sure all the flags are good */
     PyArray_UpdateFlags(view, NPY_ARRAY_UPDATE_ALL);
+    /*
+     * Add the mask to the view if the operand was NPY_ITER_USE_MASKNA.
+     */
+    if (maskna_indices[i] >= 0) {
+        PyArrayObject_fieldaccess *fview = (PyArrayObject_fieldaccess *)view;
+        int i_maskna = maskna_indices[i];
+        npy_intp *maskna_strides = fview->maskna_strides;
+
+        fview->maskna_dtype = PyArray_MASKNA_DTYPE(obj);
+        Py_INCREF(fview->maskna_dtype);
+        fview->maskna_data = NIT_RESETDATAPTR(iter)[i_maskna];
+
+        axisdata = NIT_AXISDATA(iter);
+        for (idim = 0; idim < ndim; ++idim) {
+            maskna_strides[ndim-idim-1] = NAD_STRIDES(axisdata)[i_maskna];
+
+            NIT_ADVANCE_AXISDATA(axisdata, 1);
+        }
+
+        /* This view doesn't own the mask */
+        fview->flags |= NPY_ARRAY_MASKNA;
+        fview->flags &= ~NPY_ARRAY_OWNMASKNA;
+    }
 
     return view;
 }
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index 2054a4fdb..5e4401c34 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -85,6 +85,8 @@ npyiter_allocate_arrays(NpyIter *iter,
                         PyArray_Descr **op_dtype, PyTypeObject *subtype,
                         npy_uint32 *op_flags, char *op_itflags,
                         int **op_axes, int output_scalars);
+static int
+npyiter_fill_maskna_axisdata(NpyIter *iter, int **op_axes);
 static void
 npyiter_get_priority_subtype(int first_maskna_op, PyArrayObject **op,
                             char *op_itflags,
@@ -406,6 +408,18 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,
         return NULL;
     }
 
+    /*
+     * If there were any NA masks added to the iteration, fill in
+     * the strides and other data they need. This is being done
+     * after all the 'allocate' arrays are finished.
+     */
+    if (maskna_nop > 0) {
+        if (!npyiter_fill_maskna_axisdata(iter, op_axes)) {
+            NpyIter_Deallocate(iter);
+            return NULL;
+        }
+    }
+
     NPY_IT_TIME_POINT(c_allocate_arrays);
 
     /*
@@ -1149,27 +1163,6 @@ npyiter_prepare_one_operand(PyArrayObject **op,
 }
 
 /*
- * Prepares the maskna virtual operand for one of the constructor
- * operands.  Assumes that 'op' has already been prepared by
- * npyiter_prepare_one_operand. Fills in 'op_maskna_dataptr',
- * 'op_maskna_dtype', and 'op_maskna_itflags'.
- *
- * This needs to be called after any 'allocate' operands have
- * been allocated.
- *
- * Returns 1 on success, 0 on failure.
- */
-static int
-npyiter_prepare_one_maskna_operand(PyArrayObject *op,
-                        char **op_maskna_dataptr,
-                        PyArray_Descr **op_maskna_dtype,
-                        npy_uint32 flags,
-                        npy_uint32 op_flags, char *op_maskna_itflags)
-{
-    return 1;
-}
-
-/*
  * Process all the operands, copying new references so further processing
  * can replace the arrays if copying is necessary.
  */
@@ -3128,6 +3121,84 @@ npyiter_allocate_arrays(NpyIter *iter,
 }
 
 /*
+ * Prepares the maskna virtual operands for the constructor
+ * operands, and fills in the axisdata. Fills in 'op_maskna_dataptr',
+ * 'op_maskna_dtype', and may modify 'op_maskna_itflags'.
+ *
+ * This needs to be called after any 'allocate' operands have
+ * been allocated. There is no validation of the shape/strides done,
+ * because the shape of a mask exactly matches the shape of the
+ * operand to which it attached.
+ *
+ * Returns 1 on success, 0 on failure.
+ */
+static int
+npyiter_fill_maskna_axisdata(NpyIter *iter, int **op_axes)
+{
+    npy_uint32 itflags = NIT_ITFLAGS(iter);
+    int idim, ndim = NIT_NDIM(iter);
+    int iop, iop_maskna, nop = NIT_NOP(iter);
+    int first_maskna_op = NIT_FIRST_MASKNA_OP(iter);
+
+    npy_int8 *maskna_indices = NIT_MASKNA_INDICES(iter);
+    NpyIter_AxisData *axisdata;
+    npy_intp sizeof_axisdata;
+    PyArrayObject **op = NIT_OPERANDS(iter), *op_cur;
+    char **op_dataptr = NIT_RESETDATAPTR(iter);
+
+    axisdata = NIT_AXISDATA(iter);
+    sizeof_axisdata = NIT_AXISDATA_SIZEOF(itflags, ndim, nop);
+
+    /* Fill in the reset dataptr array with the mask pointers */
+    for (iop = first_maskna_op; iop < nop; ++iop) {
+        op_dataptr[iop] = PyArray_MASKNA_DATA(op[maskna_indices[iop]]);
+    }
+
+    /* Process the maskna operands, filling in the axisdata */
+    for (idim = 0; idim < ndim; ++idim) {
+        npy_intp *strides = NAD_STRIDES(axisdata);
+
+        for (iop = first_maskna_op; iop < nop; ++iop) {
+            /*
+             * iop_maskna is the index of the USE_MASKNA input,
+             * iop is the index of the corresponding mask.
+             */
+            iop_maskna = maskna_indices[iop];
+            op_cur = op[iop_maskna];
+
+            /*
+             * The strides of the mask will be zero exactly
+             * where they're zero for the main data
+             */
+            if (strides[iop_maskna] == 0) {
+                strides[iop] = 0;
+            }
+            else {
+                int i;
+
+                if (op_axes == NULL || op_axes[iop] == NULL) {
+                    i = PyArray_NDIM(op_cur) - idim - 1;
+                }
+                else {
+                    i = op_axes[iop][ndim-idim-1];
+                }
+
+                strides[iop] = PyArray_MASKNA_STRIDES(op_cur)[i];
+            }
+        }
+
+        /* Initialize the mask data pointers */
+        memcpy(NAD_PTRS(axisdata) + first_maskna_op,
+                op_dataptr + first_maskna_op,
+                NPY_SIZEOF_INTP*(nop - first_maskna_op));
+
+        NIT_ADVANCE_AXISDATA(axisdata, 1);
+    }
+
+    return 1;
+}
+
+/*
  * The __array_priority__ attribute of the inputs determines
  * the subtype of any output arrays.  This function finds the
  * subtype of the input array with highest priority.
diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c
index 5ece85eca..0d4da1436 100644
--- a/numpy/core/src/multiarray/nditer_pywrap.c
+++ b/numpy/core/src/multiarray/nditer_pywrap.c
@@ -428,8 +428,17 @@ NpyIter_OpFlagsConverter(PyObject *op_flags_in,
                 }
                 break;
             case 'u':
-                if (strcmp(str, "updateifcopy") == 0) {
-                    flag = NPY_ITER_UPDATEIFCOPY;
+                switch (str[1]) {
+                    case 'p':
+                        if (strcmp(str, "updateifcopy") == 0) {
+                            flag = NPY_ITER_UPDATEIFCOPY;
+                        }
+                        break;
+                    case 's':
+                        if (strcmp(str, "use_maskna") == 0) {
+                            flag = NPY_ITER_USE_MASKNA;
+                        }
+                        break;
                 }
                 break;
             case 'v':
@@ -659,9 +668,9 @@ npyiter_convert_op_axes(PyObject *op_axes_in, npy_intp nop,
 }
 
 /*
- * Converts the operand array and op_flags array into the form NpyIter_AdvancedNew
- * needs.  Sets nop, and on success, each op[i] owns a reference
- * to an array object.
+ * Converts the operand array and op_flags array into the form
+ * NpyIter_AdvancedNew needs.  Sets nop, and on success, each
+ * op[i] owns a reference to an array object.
  */
 static int
 npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
@@ -1977,7 +1986,11 @@ static PyObject *npyiter_nop_get(NewNpyArrayIterObject *self)
         return NULL;
     }
 
-    return PyInt_FromLong(NpyIter_GetNOp(self->iter));
+    /*
+     * We only expose the provided operands, which is everything
+     * before the first MASKNA operand.
+     */
+    return PyInt_FromLong(NpyIter_GetFirstMaskNAOp(self->iter));
 }
 
 static PyObject *npyiter_itersize_get(NewNpyArrayIterObject *self)
@@ -2008,7 +2021,11 @@ npyiter_seq_length(NewNpyArrayIterObject *self)
         return 0;
     }
     else {
-        return NpyIter_GetNOp(self->iter);
+        /*
+         * We only expose the provided operands, which is everything
+         * before the first MASKNA operand.
+         */
+        return NpyIter_GetFirstMaskNAOp(self->iter);
     }
 }
 
@@ -2017,10 +2034,12 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
 {
     PyArrayObject *ret;
 
+    npy_int8 *maskna_indices;
     npy_intp ret_ndim;
     npy_intp nop, innerloopsize, innerstride;
     char *dataptr;
     PyArray_Descr *dtype;
+    int has_external_loop;
 
     if (self->iter == NULL || self->finished) {
         PyErr_SetString(PyExc_ValueError,
@@ -2035,7 +2054,11 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
         return NULL;
     }
 
-    nop = NpyIter_GetNOp(self->iter);
+    /*
+     * We only expose the provided operands, which is everything
+     * before the first MASKNA operand.
+     */
+    nop = NpyIter_GetFirstMaskNAOp(self->iter);
     if (i < 0 || i >= nop) {
         PyErr_Format(PyExc_IndexError,
                 "Iterator operand index %d is out of bounds", (int)i);
@@ -2059,8 +2082,10 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
 
     dataptr = self->dataptrs[i];
     dtype = self->dtypes[i];
+    has_external_loop = NpyIter_HasExternalLoop(self->iter);
+    maskna_indices = NpyIter_GetMaskNAIndices(self->iter);
 
-    if (NpyIter_HasExternalLoop(self->iter)) {
+    if (has_external_loop) {
         innerloopsize = *self->innerloopsizeptr;
         innerstride = self->innerstrides[i];
         ret_ndim = 1;
@@ -2085,6 +2110,23 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
 
     PyArray_UpdateFlags(ret, NPY_ARRAY_UPDATE_ALL);
 
+    /* If this is a USE_MASKNA operand, include the mask */
+    if (maskna_indices[i] >= 0) {
+        PyArrayObject *obj = NpyIter_GetOperandArray(self->iter)[i];
+        PyArrayObject_fieldaccess *fret = (PyArrayObject_fieldaccess *)ret;
+        int i_maskna = maskna_indices[i];
+
+        fret->maskna_dtype = PyArray_MASKNA_DTYPE(obj);
+        Py_INCREF(fret->maskna_dtype);
+        fret->maskna_data = self->dataptrs[i_maskna];
+        if (has_external_loop) {
+            fret->maskna_strides[0] = self->innerstrides[i_maskna];
+        }
+
+        fret->flags |= NPY_ARRAY_MASKNA;
+        fret->flags &= ~NPY_ARRAY_OWNMASKNA;
+    }
+
     return (PyObject *)ret;
 }
 
@@ -2109,7 +2151,11 @@ npyiter_seq_slice(NewNpyArrayIterObject *self,
         return NULL;
     }
 
-    nop = NpyIter_GetNOp(self->iter);
+    /*
+     * We only expose the provided operands, which is everything
+     * before the first MASKNA operand.
+     */
+    nop = NpyIter_GetFirstMaskNAOp(self->iter);
     if (ilow < 0) {
         ilow = 0;
     }
@@ -2143,10 +2189,11 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
 {
 
     npy_intp nop, innerloopsize, innerstride;
+    npy_int8 *maskna_indices;
     char *dataptr;
     PyArray_Descr *dtype;
     PyArrayObject *tmp;
-    int ret;
+    int ret, has_external_loop;
 
     if (v == NULL) {
         PyErr_SetString(PyExc_ValueError,
@@ -2167,7 +2214,11 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
         return -1;
     }
 
-    nop = NpyIter_GetNOp(self->iter);
+    /*
+     * We only expose the provided operands, which is everything
+     * before the first MASKNA operand.
+     */
+    nop = NpyIter_GetFirstMaskNAOp(self->iter);
     if (i < 0 || i >= nop) {
         PyErr_Format(PyExc_IndexError,
                 "Iterator operand index %d is out of bounds", (int)i);
@@ -2181,8 +2232,9 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
 
     dataptr = self->dataptrs[i];
     dtype = self->dtypes[i];
+    has_external_loop = NpyIter_HasExternalLoop(self->iter);
 
-    if (NpyIter_HasExternalLoop(self->iter)) {
+    if (has_external_loop) {
         innerloopsize = *self->innerloopsizeptr;
         innerstride = self->innerstrides[i];
     }
@@ -2191,6 +2243,8 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
         innerstride = 0;
     }
 
+    maskna_indices = NpyIter_GetMaskNAIndices(self->iter);
+
     /* TODO - there should be a better way than this... */
     Py_INCREF(dtype);
     tmp = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype,
@@ -2200,6 +2254,25 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
     if (tmp == NULL) {
         return -1;
     }
+    /* If this is a USE_MASKNA operand, include the mask */
+    if (maskna_indices[i] >= 0) {
+        PyArrayObject *obj = NpyIter_GetOperandArray(self->iter)[i];
+        PyArrayObject_fieldaccess *ftmp = (PyArrayObject_fieldaccess *)tmp;
+        int i_maskna = maskna_indices[i];
+
+        ftmp->maskna_dtype = PyArray_MASKNA_DTYPE(obj);
+        Py_INCREF(ftmp->maskna_dtype);
+        ftmp->maskna_data = self->dataptrs[i_maskna];
+        if (has_external_loop) {
+            ftmp->maskna_strides[0] = self->innerstrides[i_maskna];
+        }
+        else {
+            ftmp->maskna_strides[0] = 0;
+        }
+
+        ftmp->flags |= NPY_ARRAY_MASKNA;
+        ftmp->flags &= ~NPY_ARRAY_OWNMASKNA;
+    }
     PyArray_UpdateFlags(tmp, NPY_ARRAY_UPDATE_ALL);
     ret = PyArray_CopyObject(tmp, v);
     Py_DECREF(tmp);
@@ -2232,7 +2305,11 @@ npyiter_seq_ass_slice(NewNpyArrayIterObject *self, Py_ssize_t ilow,
         return -1;
     }
 
-    nop = NpyIter_GetNOp(self->iter);
+    /*
+     * We only expose the provided operands, which is everything
+     * before the first MASKNA operand.
+     */
+    nop = NpyIter_GetFirstMaskNAOp(self->iter);
     if (ilow < 0) {
         ilow = 0;
     }