ENH: missingdata: Make the nditer USE_MASKNA mode work with buffering

3 files changed, 94 insertions, 9 deletions

diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c
index ddd227bf6..ae330e294 100644
--- a/numpy/core/src/multiarray/nditer_api.c
+++ b/numpy/core/src/multiarray/nditer_api.c
@@ -1817,6 +1817,7 @@ npyiter_copy_from_buffers(NpyIter *iter)
     int ndim = NIT_NDIM(iter);
     int iop, nop = NIT_NOP(iter);
     int maskop = NIT_MASKOP(iter);
+    int first_maskna_op = NIT_FIRST_MASKNA_OP(iter);
 
     char *op_itflags = NIT_OPITFLAGS(iter);
     NpyIter_BufferData *bufferdata = NIT_BUFFERDATA(iter);
@@ -1832,6 +1833,7 @@ npyiter_copy_from_buffers(NpyIter *iter)
     char **ptrs = NBF_PTRS(bufferdata), **ad_ptrs = NAD_PTRS(axisdata);
     char **buffers = NBF_BUFFERS(bufferdata);
     char *buffer;
+    npy_int8 *maskna_indices = NIT_MASKNA_INDICES(iter);
 
     npy_intp reduce_outerdim = 0;
     npy_intp *reduce_outerstrides = NULL;
@@ -1934,7 +1936,37 @@ npyiter_copy_from_buffers(NpyIter *iter)
                                     "operand %d (%d items)\n",
                                     (int)iop, (int)op_transfersize);
 
-                if (op_itflags[iop] & NPY_OP_ITFLAG_WRITEMASKED) {
+                /* USE_MASKNA operand */
+                if (iop < first_maskna_op && maskna_indices[iop] >= 0) {
+                    int iop_maskna = maskna_indices[iop];
+                    npy_mask *maskptr;
+                    /* TODO: support WRITEMASKED + USE_MASKNA together */
+
+                    /*
+                     * The mask pointer may be in the buffer or in
+                     * the array, detect which one.
+                     */
+                    delta = (ptrs[iop_maskna] - buffers[iop_maskna]);
+                    if (0 <= delta &&
+                            delta <= buffersize*dtypes[iop_maskna]->elsize) {
+                        maskptr = (npy_mask *)buffers[iop_maskna];
+                    }
+                    else {
+                        maskptr = (npy_mask *)ad_ptrs[iop_maskna];
+                    }
+
+                    PyArray_TransferMaskedStridedToNDim(ndim_transfer,
+                            ad_ptrs[iop], dst_strides, axisdata_incr,
+                            buffer, src_stride,
+                            maskptr, strides[iop_maskna],
+                            dst_coords, axisdata_incr,
+                            dst_shape, axisdata_incr,
+                            op_transfersize, dtypes[iop]->elsize,
+                            (PyArray_MaskedStridedTransferFn *)stransfer,
+                            transferdata);
+                }
+                /* WRITEMASKED operand */
+                else if (op_itflags[iop] & NPY_OP_ITFLAG_WRITEMASKED) {
                     npy_mask *maskptr;
 
                     /*
@@ -1960,6 +1992,7 @@ npyiter_copy_from_buffers(NpyIter *iter)
                             (PyArray_MaskedStridedTransferFn *)stransfer,
                             transferdata);
                 }
+                /* Regular operand */
                 else {
                     PyArray_TransferStridedToNDim(ndim_transfer,
                             ad_ptrs[iop], dst_strides, axisdata_incr,
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index c530ec844..52a69797f 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -968,6 +968,18 @@ npyiter_check_per_op_flags(npy_uint32 op_flags, char *op_itflags)
                 "be used together with ARRAYMASK");
             return 0;
         }
+        /*
+         * When WRITEMASKED and USE_MASKNA are supported together,
+         * it will probably require that buffering is enabled as well,
+         * because that will need yet another temporary mask buffer to combine
+         * the two masks before doing the masked copies.
+         */
+        if ((op_flags & NPY_ITER_USE_MASKNA) != 0) {
+            PyErr_SetString(PyExc_ValueError,
+                "The combination of iterator flags WRITEMASKED "
+                "and USE_MASKNA is not yet supported");
+            return 0;
+        }
         *op_itflags |= NPY_OP_ITFLAG_WRITEMASKED;
     }
 
@@ -3265,6 +3277,7 @@ npyiter_allocate_transfer_functions(NpyIter *iter)
                         **writetransferfn = NBF_WRITETRANSFERFN(bufferdata);
     NpyAuxData **readtransferdata = NBF_READTRANSFERDATA(bufferdata),
                **writetransferdata = NBF_WRITETRANSFERDATA(bufferdata);
+    npy_int8 *maskna_indices = NIT_MASKNA_INDICES(iter);
 
     PyArray_StridedTransferFn *stransfer = NULL;
     NpyAuxData *transferdata = NULL;
@@ -3313,9 +3326,39 @@ npyiter_allocate_transfer_functions(NpyIter *iter)
                 int move_references = 1;
 
                 /*
-                 * If the operand is WRITEMASKED, use a masked transfer fn.
+                 * If the operand has USE_MASKNA, use a masked transfer fn.
+                 * The masks for the maskna operands can be copied straight
+                 * unless the operand is also WRITEMASKED.
                  */
-                if (flags & NPY_OP_ITFLAG_WRITEMASKED) {
+                if (iop < first_maskna_op && maskna_indices[iop] >= 0) {
+                    /* TODO: support USE_MASKNA + WRITEMASKED together */
+                    PyArray_Descr *mask_dtype =
+                                    PyArray_MASKNA_DTYPE(op[iop]);
+                    int iop_maskna = maskna_indices[iop];
+
+                    /*
+                     * If the mask's stride is contiguous, use it, otherwise
+                     * the mask may or may not be buffered, so the stride
+                     * could be inconsistent.
+                     */
+                    if (PyArray_GetMaskedDTypeTransferFunction(
+                                (flags & NPY_OP_ITFLAG_ALIGNED) != 0,
+                                op_dtype[iop]->elsize,
+                                op_stride,
+                            (flags & NPY_OP_ITFLAG_REDUCE) ? NPY_MAX_INTP :
+                                                   strides[iop_maskna],
+                                op_dtype[iop],
+                                op_orig_dtype,
+                                mask_dtype,
+                                move_references,
+                                (PyArray_MaskedStridedTransferFn **)&stransfer,
+                                &transferdata,
+                                &needs_api) != NPY_SUCCEED) {
+                        goto fail;
+                    }
+                }
+                /* If the operand is WRITEMASKED, use a masked transfer fn */
+                else if (flags & NPY_OP_ITFLAG_WRITEMASKED) {
                     int maskop = NIT_MASKOP(iter);
                     PyArray_Descr *mask_dtype = PyArray_DESCR(op[maskop]);
 
diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
index fda2f9b4c..8892c643d 100644
--- a/numpy/core/tests/test_nditer.py
+++ b/numpy/core/tests/test_nditer.py
@@ -2434,20 +2434,29 @@ def test_iter_maskna():
     assert_equal(a[1], 0)
     assert_equal(np.isna(a), [1,0,1])
 
-    # Copying values with buffering
-    a_orig[...] = [1,2,3]
-    b_orig[...] = [4,5,6]
-    a[...] = [np.NA, np.NA, 5]
-    b[...] = [np.NA, 3, np.NA]
+    # Copying NA values with buffering
+    a_orig[...] = [1.5,2.5,3.5]
+    b_orig[...] = [4.5,5.5,6.5]
+    a[...] = [np.NA, np.NA, 5.5]
+    b[...] = [np.NA, 3.5, np.NA]
     it = np.nditer([a,b], ['buffered'], [['writeonly','use_maskna'],
                                ['readonly','use_maskna']],
                     op_dtypes=['i4','i4'],
                     casting='unsafe')
     for x, y in it:
         x[...] = y
+    # The 3.5 in b gets truncated to 3, because the iterator is processing
+    # elements as int32 values.
     assert_equal(a[1], 3)
     assert_equal(np.isna(a), [1,0,1])
-    assert_equal(a_orig, [1,3,3])
+    assert_equal(a_orig, [1.5,3,5.5])
+
+    # WRITEMASKED and MASKNA aren't supported together yet
+    mask = np.array([1,1,0], dtype='?')
+    assert_raises(ValueError, np.nditer, [a,b,mask], ['buffered'],
+                                [['writeonly','use_maskna','writemasked'],
+                                 ['readonly','use_maskna'],
+                                 ['readonly','arraymask']])
 
 if __name__ == "__main__":
     run_module_suite()