ENH: missingdata: Finish implementation of scalar to array assignment

This still does not support a 'wheremask' which contains NA values,
but that can always be added later.

6 files changed, 552 insertions, 40 deletions

diff --git a/numpy/core/src/multiarray/array_assign.h b/numpy/core/src/multiarray/array_assign.h
index ddbcf6dae..75626000c 100644
--- a/numpy/core/src/multiarray/array_assign.h
+++ b/numpy/core/src/multiarray/array_assign.h
@@ -11,8 +11,13 @@
  * wheremask: If non-NULL, a boolean mask specifying where to copy.
  * casting: An exception is raised if the assignment violates this
  *          casting rule.
- * overwritena: If 1, overwrites everything in 'dst', if 0, it
- *              does not overwrite elements in 'dst' which are NA.
+ * preservena: If 0, overwrites everything in 'dst', if 1, it
+ *              preserves elements in 'dst' which are NA.
+ * preservewhichna: Must be NULL. When multi-NA support is implemented,
+ *                   this will be an array of flags for 'preservena=True',
+ *                   indicating which NA payload values to preserve.
+ *
+ * This function is implemented in array_assign_scalar.c.
  *
  * Returns 0 on success, -1 on failure.
  */
@@ -20,7 +25,8 @@ NPY_NO_EXPORT int
 array_assign_scalar(PyArrayObject *dst,
                     PyArray_Descr *src_dtype, char *src_data,
                     PyArrayObject *wheremask,
-                    NPY_CASTING casting, npy_bool overwritena);
+                    NPY_CASTING casting,
+                    npy_bool preservena, npy_bool *preservewhichna);
 
 /*
  * An array assignment function for copying arrays, broadcasting 'src' into
@@ -33,15 +39,19 @@ array_assign_scalar(PyArrayObject *dst,
  * wheremask: If non-NULL, a boolean mask specifying where to copy.
  * casting: An exception is raised if the copy violates this
  *          casting rule.
- * overwritena: If 1, overwrites everything in 'dst', if 0, it
- *              does not overwrite elements in 'dst' which are NA.
+ * preservena: If 0, overwrites everything in 'dst', if 1, it
+ *              preserves elements in 'dst' which are NA.
+ * preservewhichna: Must be NULL. When multi-NA support is implemented,
+ *                   this will be an array of flags for 'preservena=True',
+ *                   indicating which NA payload values to preserve.
  *
  * Returns 0 on success, -1 on failure.
  */
 NPY_NO_EXPORT int
-array_assign_broadcast(PyArrayObject *dst, PyArrayObject *src,
+array_assign_array(PyArrayObject *dst, PyArrayObject *src,
                     PyArrayObject *wheremask,
-                    NPY_CASTING casting, npy_bool overwritena);
+                    NPY_CASTING casting,
+                    npy_bool preservena, npy_bool *preservewhichna);
 
 /*
  * An array assignment function for copying arrays, treating the
@@ -57,8 +67,11 @@ array_assign_broadcast(PyArrayObject *dst, PyArrayObject *src,
  * wheremask: If non-NULL, a boolean mask specifying where to copy.
  * casting: An exception is raised if the copy violates this
  *          casting rule.
- * overwritena: If 1, overwrites everything in 'dst', if 0, it
- *              does not overwrite elements in 'dst' which are NA.
+ * preservena: If 0, overwrites everything in 'dst', if 1, it
+ *              preserves elements in 'dst' which are NA.
+ * preservewhichna: Must be NULL. When multi-NA support is implemented,
+ *                   this will be an array of flags for 'preservena=True',
+ *                   indicating which NA payload values to preserve.
  *
  * Returns 0 on success, -1 on failure.
  */
@@ -66,7 +79,8 @@ NPY_NO_EXPORT int
 array_assign_flat(PyArrayObject *dst, NPY_ORDER dst_order,
                   PyArrayObject *src, NPY_ORDER src_order,
                   PyArrayObject *wheremask,
-                  NPY_CASTING casting, npy_bool overwritena);
+                  NPY_CASTING casting,
+                  npy_bool preservena, npy_bool *preservewhichna);
 
 
 
diff --git a/numpy/core/src/multiarray/array_assign_scalar.c b/numpy/core/src/multiarray/array_assign_scalar.c
index 502d8c733..9ddb28f49 100644
--- a/numpy/core/src/multiarray/array_assign_scalar.c
+++ b/numpy/core/src/multiarray/array_assign_scalar.c
@@ -167,12 +167,12 @@ raw_array_wheremasked_assign_scalar(int ndim, npy_intp *shape,
 
 /*
  * Assigns the scalar value to every element of the destination raw array
- * except for those which are masked as NA.
+ * except for those which are masked as NA according to 'maskna'.
  *
  * Returns 0 on success, -1 on failure.
  */
 static int
-raw_array_assign_scalar_dont_overwritena(int ndim, npy_intp *shape,
+raw_array_assign_scalar_preservena(int ndim, npy_intp *shape,
         PyArray_Descr *dst_dtype, char *dst_data, npy_intp *dst_strides,
         PyArray_Descr *src_dtype, char *src_data,
         PyArray_Descr *maskna_dtype, char *maskna_data,
@@ -248,16 +248,28 @@ raw_array_assign_scalar_dont_overwritena(int ndim, npy_intp *shape,
 
     NPY_RAW_ITER_START(idim, ndim, coord, shape_it) {
         npy_intp buffered_count, count;
+        char *dst_d, *maskna_d;
         /* Process the innermost dimension a buffer size at a time */
         count = shape_it[0];
+        dst_d = dst_data;
+        maskna_d = maskna_data;
         do {
-            buffered_count = shape_it[0] > NPY_ARRAY_ASSIGN_BUFFERSIZE
+            buffered_count = count > NPY_ARRAY_ASSIGN_BUFFERSIZE
                                         ? count
                                         : NPY_ARRAY_ASSIGN_BUFFERSIZE;
-            /* TODO: invert the mask into the buffer here */
-            stransfer(dst_data, dst_strides_it[0], src_data, 0,
+
+            /* Invert the mask into the buffer */
+            maskinv_stransfer(maskna_buffer, maskna_itemsize,
+                        maskna_d, maskna_strides_it[0],
+                        buffered_count, maskna_itemsize, maskinv_transferdata);
+
+            /* Transfer the data based on the buffered mask */
+            stransfer(dst_d, dst_strides_it[0], src_data, 0,
                         (npy_mask *)maskna_buffer, maskna_itemsize,
                         buffered_count, src_itemsize, transferdata);
+
+            dst_d += buffered_count * dst_strides_it[0];
+            maskna_d += buffered_count * maskna_strides_it[0];
             count -= buffered_count;
         } while (count > 0);
     } NPY_RAW_ITER_TWO_NEXT(idim, ndim, coord, shape_it,
@@ -268,19 +280,157 @@ raw_array_assign_scalar_dont_overwritena(int ndim, npy_intp *shape,
         NPY_END_THREADS;
     }
 
+    PyArray_free(maskna_buffer);
     NPY_AUXDATA_FREE(transferdata);
     NPY_AUXDATA_FREE(maskinv_transferdata);
 
     return (needs_api && PyErr_Occurred()) ? -1 : 0;
 }
 
+/*
+ * Assigns the scalar value to every element of the destination raw array
+ * where the 'wheremask' is True, except for those which are masked as NA
+ * according to 'maskna'.
+ *
+ * Returns 0 on success, -1 on failure.
+ */
+static int
+raw_array_wheremasked_assign_scalar_preservena(int ndim, npy_intp *shape,
+        PyArray_Descr *dst_dtype, char *dst_data, npy_intp *dst_strides,
+        PyArray_Descr *src_dtype, char *src_data,
+        PyArray_Descr *maskna_dtype, char *maskna_data,
+        npy_intp *maskna_strides,
+        PyArray_Descr *wheremask_dtype, char *wheremask_data,
+        npy_intp *wheremask_strides)
+{
+    int idim;
+    npy_intp shape_it[NPY_MAXDIMS], dst_strides_it[NPY_MAXDIMS];
+    npy_intp maskna_strides_it[NPY_MAXDIMS];
+    npy_intp wheremask_strides_it[NPY_MAXDIMS];
+    npy_intp coord[NPY_MAXDIMS];
+    NPY_BEGIN_THREADS_DEF;
+
+    PyArray_MaskedStridedUnaryOp *stransfer = NULL;
+    NpyAuxData *transferdata = NULL;
+    int aligned, needs_api = 0;
+    npy_intp src_itemsize = src_dtype->elsize;
+
+    PyArray_StridedBinaryOp *maskand_stransfer = NULL;
+    NpyAuxData *maskand_transferdata = NULL;
+
+    char *maskna_buffer;
+    npy_intp maskna_itemsize;
+
+    /* Check alignment */
+    aligned = raw_array_is_aligned(ndim, dst_data, dst_strides,
+                                    dst_dtype->alignment);
+    if (((npy_intp)src_data & (src_dtype->alignment - 1)) != 0) {
+        aligned = 0;
+    }
+
+    /* Use raw iteration with no heap allocation */
+    if (PyArray_PrepareThreeRawArrayIter(
+                    ndim, shape,
+                    dst_data, dst_strides,
+                    maskna_data, maskna_strides,
+                    wheremask_data, wheremask_strides,
+                    &ndim, shape_it,
+                    &dst_data, dst_strides_it,
+                    &maskna_data, maskna_strides_it,
+                    &wheremask_data, wheremask_strides_it) < 0) {
+        return -1;
+    }
+
+    /* Allocate a buffer for inverting/anding the mask */
+    maskna_itemsize = maskna_dtype->elsize;
+    maskna_buffer = PyArray_malloc(NPY_ARRAY_ASSIGN_BUFFERSIZE *
+                                    maskna_itemsize);
+    if (maskna_buffer == NULL) {
+        PyErr_NoMemory();
+        return -1;
+    }
+
+    /* Get the function to do the casting */
+    if (PyArray_GetMaskedDTypeTransferFunction(aligned,
+                        0, dst_strides_it[0], maskna_itemsize,
+                        src_dtype, dst_dtype, maskna_dtype,
+                        0,
+                        &stransfer, &transferdata,
+                        &needs_api) != NPY_SUCCEED) {
+        PyArray_free(maskna_buffer);
+        return -1;
+    }
+
+    /*
+     * Get the function to invert the mask. The output
+     * of the binary operation is the dtype 'maskna_dtype'
+     */
+    if (PyArray_GetMaskAndFunction(
+                        maskna_strides_it[0], maskna_dtype, 1,
+                        wheremask_strides_it[0], wheremask_dtype, 0,
+                        &maskand_stransfer, &maskand_transferdata) < 0) {
+        PyArray_free(maskna_buffer);
+        NPY_AUXDATA_FREE(transferdata);
+        return -1;
+    }
+
+    if (!needs_api) {
+        NPY_BEGIN_THREADS;
+    }
+
+    NPY_RAW_ITER_START(idim, ndim, coord, shape_it) {
+        npy_intp buffered_count, count;
+        char *dst_d, *maskna_d, *wheremask_d;
+        /* Process the innermost dimension a buffer size at a time */
+        count = shape_it[0];
+        dst_d = dst_data;
+        maskna_d = maskna_data;
+        wheremask_d = wheremask_data;
+        do {
+            buffered_count = count > NPY_ARRAY_ASSIGN_BUFFERSIZE
+                                        ? count
+                                        : NPY_ARRAY_ASSIGN_BUFFERSIZE;
+
+            /* Prepare the mask into the buffer */
+            maskand_stransfer(maskna_buffer, maskna_itemsize,
+                        maskna_d, maskna_strides_it[0],
+                        wheremask_d, wheremask_strides_it[0],
+                        buffered_count, maskand_transferdata);
+
+            /* Transfer the data based on the buffered mask */
+            stransfer(dst_d, dst_strides_it[0], src_data, 0,
+                        (npy_mask *)maskna_buffer, maskna_itemsize,
+                        buffered_count, src_itemsize, transferdata);
+
+            dst_d += buffered_count * dst_strides_it[0];
+            maskna_d += buffered_count * maskna_strides_it[0];
+            wheremask_d += buffered_count * wheremask_strides_it[0];
+            count -= buffered_count;
+        } while (count > 0);
+    } NPY_RAW_ITER_THREE_NEXT(idim, ndim, coord, shape_it,
+                            dst_data, dst_strides_it,
+                            maskna_data, maskna_strides_it,
+                            wheremask_data, wheremask_strides_it);
+
+    if (!needs_api) {
+        NPY_END_THREADS;
+    }
+
+    PyArray_free(maskna_buffer);
+    NPY_AUXDATA_FREE(transferdata);
+    NPY_AUXDATA_FREE(maskand_transferdata);
+
+    return (needs_api && PyErr_Occurred()) ? -1 : 0;
+}
+
 
 /* See array_assign.h for documentation */
 NPY_NO_EXPORT int
 array_assign_scalar(PyArrayObject *dst,
                     PyArray_Descr *src_dtype, char *src_data,
                     PyArrayObject *wheremask,
-                    NPY_CASTING casting, npy_bool overwritena)
+                    NPY_CASTING casting,
+                    npy_bool preservena, npy_bool *preservewhichna)
 {
     int allocated_src_data = 0, dst_has_maskna = PyArray_HASMASKNA(dst);
     npy_longlong scalarbuffer[4];
@@ -303,6 +453,12 @@ array_assign_scalar(PyArrayObject *dst,
         return -1;
     }
 
+    if (preservewhichna != NULL) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "multi-NA support is not yet implemented");
+        return -1;
+    }
+
     /*
      * Make a copy of the src data if it's a different dtype than 'dst'
      * or isn't aligned, and the destination we're copying to has
@@ -336,7 +492,7 @@ array_assign_scalar(PyArrayObject *dst,
 
     if (wheremask == NULL) {
         /* A straightforward value assignment */
-        if (overwritena || !dst_has_maskna) {
+        if (!preservena || !dst_has_maskna) {
             /* If assigning to an array with an NA mask, set to all exposed */
             if (dst_has_maskna) {
                 if (PyArray_AssignMaskNA(dst, 1) < 0) {
@@ -353,7 +509,7 @@ array_assign_scalar(PyArrayObject *dst,
         }
         /* A value assignment without overwriting NA values */
         else {
-            if (raw_array_assign_scalar_dont_overwritena(
+            if (raw_array_assign_scalar_preservena(
                     PyArray_NDIM(dst), PyArray_DIMS(dst),
                     PyArray_DESCR(dst), PyArray_DATA(dst), PyArray_STRIDES(dst),
                     src_dtype, src_data,
@@ -364,6 +520,8 @@ array_assign_scalar(PyArrayObject *dst,
         }
     }
     else {
+        npy_intp wheremask_strides[NPY_MAXDIMS];
+
         if (PyArray_ContainsNA(wheremask)) {
             if (!dst_has_maskna) {
                 PyErr_SetString(PyExc_ValueError,
@@ -380,10 +538,16 @@ array_assign_scalar(PyArrayObject *dst,
             }
         }
 
-        /* A straightforward where-masked assignment */
-        if (overwritena || !dst_has_maskna) {
-            npy_intp wheremask_strides[NPY_MAXDIMS];
+        /* Broadcast the wheremask to 'dst' for raw iteration */
+        if (broadcast_strides(PyArray_NDIM(dst), PyArray_DIMS(dst),
+                    PyArray_NDIM(wheremask), PyArray_DIMS(wheremask),
+                    PyArray_STRIDES(wheremask), "where mask",
+                    wheremask_strides) < 0) {
+            goto fail;
+        }
 
+        /* A straightforward where-masked assignment */
+        if (!preservena || !dst_has_maskna) {
             /* If assigning to an array with an NA mask, set to all exposed */
             if (dst_has_maskna) {
                 /*
@@ -395,14 +559,6 @@ array_assign_scalar(PyArrayObject *dst,
                 }
             }
 
-            /* Broadcast the wheremask to 'dst' for raw iteration */
-            if (broadcast_strides(PyArray_NDIM(dst), PyArray_DIMS(dst),
-                        PyArray_NDIM(wheremask), PyArray_DIMS(wheremask),
-                        PyArray_STRIDES(wheremask), "where mask",
-                        wheremask_strides) < 0) {
-                goto fail;
-            }
-
             /* Do the masked assignment with raw array iteration */
             if (raw_array_wheremasked_assign_scalar(
                     PyArray_NDIM(dst), PyArray_DIMS(dst),
@@ -415,7 +571,16 @@ array_assign_scalar(PyArrayObject *dst,
         }
         /* A masked value assignment without overwriting NA values */
         else {
-            /* TODO: This function is next */
+            if (raw_array_wheremasked_assign_scalar_preservena(
+                    PyArray_NDIM(dst), PyArray_DIMS(dst),
+                    PyArray_DESCR(dst), PyArray_DATA(dst), PyArray_STRIDES(dst),
+                    src_dtype, src_data,
+                    PyArray_MASKNA_DTYPE(dst), PyArray_MASKNA_DATA(dst),
+                    PyArray_MASKNA_STRIDES(dst),
+                    PyArray_DESCR(wheremask), PyArray_DATA(wheremask),
+                    wheremask_strides) < 0) {
+                goto fail;
+            }
         }
     }
 
diff --git a/numpy/core/src/multiarray/dtype_transfer.c b/numpy/core/src/multiarray/dtype_transfer.c
index 650042902..a45f5dc5b 100644
--- a/numpy/core/src/multiarray/dtype_transfer.c
+++ b/numpy/core/src/multiarray/dtype_transfer.c
@@ -4093,7 +4093,146 @@ PyArray_PrepareTwoRawArrayIter(int ndim, npy_intp *shape,
     *out_dataB = dataB;
     *out_ndim = ndim;
     return 0;
+}
+
+/*
+ * The same as PyArray_PrepareOneRawArrayIter, but for three
+ * operands instead of one. Any broadcasting of the three operands
+ * should have already been done before calling this function,
+ * as the ndim and shape is only specified once for all operands.
+ *
+ * Only the strides of the first operand are used to reorder
+ * the dimensions, no attempt to consider all the strides together
+ * is made, as is done in the NpyIter object.
+ *
+ * You can use this together with NPY_RAW_ITER_START and
+ * NPY_RAW_ITER_THREE_NEXT to handle the looping boilerplate of everything
+ * but the innermost loop (which is for idim == 0).
+ *
+ * Returns 0 on success, -1 on failure.
+ */
+NPY_NO_EXPORT int
+PyArray_PrepareThreeRawArrayIter(int ndim, npy_intp *shape,
+                            char *dataA, npy_intp *stridesA,
+                            char *dataB, npy_intp *stridesB,
+                            char *dataC, npy_intp *stridesC,
+                            int *out_ndim, npy_intp *out_shape,
+                            char **out_dataA, npy_intp *out_stridesA,
+                            char **out_dataB, npy_intp *out_stridesB,
+                            char **out_dataC, npy_intp *out_stridesC)
+{
+    npy_stride_sort_item strideperm[NPY_MAXDIMS];
+    int i, j;
+
+    /* Special case 0 and 1 dimensions */
+    if (ndim == 0) {
+        *out_ndim = 1;
+        *out_dataA = dataA;
+        *out_dataB = dataB;
+        *out_dataC = dataC;
+        out_shape[0] = 1;
+        out_stridesA[0] = 0;
+        out_stridesB[0] = 0;
+        out_stridesC[0] = 0;
+        return 0;
+    }
+    else if (ndim == 1) {
+        npy_intp stride_entryA = stridesA[0];
+        npy_intp stride_entryB = stridesB[0];
+        npy_intp stride_entryC = stridesC[0];
+        npy_intp shape_entry = shape[0];
+        *out_ndim = 1;
+        out_shape[0] = shape[0];
+        /* Always make a positive stride for the first operand */
+        if (stride_entryA >= 0) {
+            *out_dataA = dataA;
+            *out_dataB = dataB;
+            *out_dataC = dataC;
+            out_stridesA[0] = stride_entryA;
+            out_stridesB[0] = stride_entryB;
+            out_stridesC[0] = stride_entryC;
+        }
+        else {
+            *out_dataA = dataA + stride_entryA * (shape_entry - 1);
+            *out_dataB = dataB + stride_entryB * (shape_entry - 1);
+            *out_dataC = dataC + stride_entryC * (shape_entry - 1);
+            out_stridesA[0] = -stride_entryA;
+            out_stridesB[0] = -stride_entryB;
+            out_stridesC[0] = -stride_entryC;
+        }
+        return 0;
+    }
+
+    /* Sort the axes based on the destination strides */
+    PyArray_CreateSortedStridePerm(ndim, stridesA, strideperm);
+    for (i = 0; i < ndim; ++i) {
+        int iperm = ndim - strideperm[i].perm - 1;
+        out_shape[i] = shape[iperm];
+        out_stridesA[i] = stridesA[iperm];
+        out_stridesB[i] = stridesB[iperm];
+        out_stridesC[i] = stridesC[iperm];
+    }
+
+    /* Reverse any negative strides of operand A */
+    for (i = 0; i < ndim; ++i) {
+        npy_intp stride_entryA = out_stridesA[i];
+        npy_intp stride_entryB = out_stridesB[i];
+        npy_intp stride_entryC = out_stridesC[i];
+        npy_intp shape_entry = out_shape[i];
+
+        if (stride_entryA < 0) {
+            dataA += stride_entryA * (shape_entry - 1);
+            dataB += stride_entryB * (shape_entry - 1);
+            dataC += stride_entryC * (shape_entry - 1);
+            out_stridesA[i] = -stride_entryA;
+            out_stridesB[i] = -stride_entryB;
+            out_stridesC[i] = -stride_entryC;
+        }
+        /* Detect 0-size arrays here */
+        if (shape_entry == 0) {
+            *out_ndim = 1;
+            *out_dataA = dataA;
+            *out_dataB = dataB;
+            *out_dataC = dataC;
+            out_shape[0] = 0;
+            out_stridesA[0] = 0;
+            out_stridesB[0] = 0;
+            out_stridesC[0] = 0;
+            return 0;
+        }
+    }
+
+    /* Coalesce any dimensions where possible */
+    i = 0;
+    for (j = 1; j < ndim; ++j) {
+        if (out_shape[i] == 1) {
+            /* Drop axis i */
+            out_shape[i] = out_shape[j];
+            out_stridesA[i] = out_stridesA[j];
+            out_stridesB[i] = out_stridesB[j];
+            out_stridesC[i] = out_stridesC[j];
+        }
+        else if (out_shape[j] == 1) {
+            /* Drop axis j */
+        }
+        else if (out_stridesA[i] * out_shape[i] == out_stridesA[j] &&
+                    out_stridesB[i] * out_shape[i] == out_stridesB[j] &&
+                    out_stridesC[i] * out_shape[i] == out_stridesC[j]) {
+            /* Coalesce axes i and j */
+            out_shape[i] *= out_shape[j];
+        }
+        else {
+            /* Can't coalesce, go to next i */
+            ++i;
+        }
+    }
+    ndim = i+1;
 
+    *out_dataA = dataA;
+    *out_dataB = dataB;
+    *out_dataC = dataC;
+    *out_ndim = ndim;
+    return 0;
 }
 
 /*
diff --git a/numpy/core/src/multiarray/na_mask.c b/numpy/core/src/multiarray/na_mask.c
index d6c8d1463..f43858c5b 100644
--- a/numpy/core/src/multiarray/na_mask.c
+++ b/numpy/core/src/multiarray/na_mask.c
@@ -582,7 +582,7 @@ static void
 _strided_bool_mask_inversion(char *dst, npy_intp dst_stride,
                             char *src, npy_intp src_stride,
                             npy_intp N, npy_intp NPY_UNUSED(src_itemsize),
-                            NpyAuxData *NPY_UNUSED(data))
+                            NpyAuxData *NPY_UNUSED(opdata))
 {
     while (N > 0) {
         *dst = !(*src);
@@ -593,12 +593,11 @@ _strided_bool_mask_inversion(char *dst, npy_intp dst_stride,
 }
 
 NPY_NO_EXPORT int
-PyArray_GetMaskInversionFunction(npy_intp mask_stride,
-                            PyArray_Descr *mask_dtype,
-                            PyArray_StridedUnaryOp **out_stransfer,
-                            NpyAuxData **out_transferdata)
+PyArray_GetMaskInversionFunction(
+        npy_intp mask_stride, PyArray_Descr *mask_dtype,
+        PyArray_StridedUnaryOp **out_unop, NpyAuxData **out_opdata)
 {
-    /* Will use the transferdata with the field version */
+    /* Will use the opdata with the field version */
     if (PyDataType_HASFIELDS(mask_dtype)) {
         PyErr_SetString(PyExc_RuntimeError,
                 "field-based masks are not supported yet");
@@ -613,7 +612,123 @@ PyArray_GetMaskInversionFunction(npy_intp mask_stride,
 
     /* TODO: Specialize for contiguous data */
 
-    *out_stransfer = &_strided_bool_mask_inversion;
-    *out_transferdata = NULL;
+    *out_unop = &_strided_bool_mask_inversion;
+    *out_opdata = NULL;
+    return 0;
+}
+
+static void
+_strided_bool_mask_noinv0_noinv1_and(char *dst, npy_intp dst_stride,
+                            char *src0, npy_intp src0_stride,
+                            char *src1, npy_intp src1_stride,
+                            npy_intp N, NpyAuxData *NPY_UNUSED(opdata))
+{
+    while (N > 0) {
+        *dst = (*src0) & (*src1);
+        dst += dst_stride;
+        src0 += src0_stride;
+        src1 += src1_stride;
+        --N;
+    }
+}
+
+static void
+_strided_bool_mask_inv0_noinv1_and(char *dst, npy_intp dst_stride,
+                            char *src0, npy_intp src0_stride,
+                            char *src1, npy_intp src1_stride,
+                            npy_intp N, NpyAuxData *NPY_UNUSED(opdata))
+{
+    while (N > 0) {
+        *dst = ((*src0) ^ 0x01) & (*src1);
+        dst += dst_stride;
+        src0 += src0_stride;
+        src1 += src1_stride;
+        --N;
+    }
+}
+
+static void
+_strided_bool_mask_noinv0_inv1_and(char *dst, npy_intp dst_stride,
+                            char *src0, npy_intp src0_stride,
+                            char *src1, npy_intp src1_stride,
+                            npy_intp N, NpyAuxData *NPY_UNUSED(opdata))
+{
+    while (N > 0) {
+        *dst = (*src0) & ((*src1) ^ 0x01);
+        dst += dst_stride;
+        src0 += src0_stride;
+        src1 += src1_stride;
+        --N;
+    }
+}
+
+static void
+_strided_bool_mask_inv0_inv1_and(char *dst, npy_intp dst_stride,
+                            char *src0, npy_intp src0_stride,
+                            char *src1, npy_intp src1_stride,
+                            npy_intp N, NpyAuxData *NPY_UNUSED(opdata))
+{
+    while (N > 0) {
+        *dst = ((*src0) | (*src1)) ^ 0x01;
+        dst += dst_stride;
+        src0 += src0_stride;
+        src1 += src1_stride;
+        --N;
+    }
+}
+
+/*
+ * Gets a function which ANDs together two masks, possibly inverting
+ * one or both of the masks as well.
+ *
+ * The dtype of the output must match 'mask0_dtype'.
+ */
+NPY_NO_EXPORT int
+PyArray_GetMaskAndFunction(
+        npy_intp mask0_stride, PyArray_Descr *mask0_dtype, int invert_mask0,
+        npy_intp mask1_stride, PyArray_Descr *mask1_dtype, int invert_mask1,
+        PyArray_StridedBinaryOp **out_binop, NpyAuxData **out_opdata)
+{
+    /* Will use the opdata with the field version */
+    if (PyDataType_HASFIELDS(mask0_dtype) ||
+                        PyDataType_HASFIELDS(mask1_dtype)) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "field-based masks are not supported yet");
+        return -1;
+    }
+
+    if (mask0_dtype->type_num == NPY_MASK ||
+                            mask1_dtype->type_num == NPY_MASK) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "multi-NA masks are not supported yet");
+        return -1;
+    }
+
+    if (mask0_dtype->type_num != NPY_BOOL ||
+                            mask1_dtype->type_num != NPY_BOOL) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "unsupported data type for mask");
+        return -1;
+    }
+
+    /* TODO: Specialize for contiguous data */
+
+    if (invert_mask0) {
+        if (invert_mask1) {
+            *out_binop = &_strided_bool_mask_inv0_inv1_and;
+        }
+        else {
+            *out_binop = &_strided_bool_mask_inv0_noinv1_and;
+        }
+    }
+    else {
+        if (invert_mask1) {
+            *out_binop = &_strided_bool_mask_noinv0_inv1_and;
+        }
+        else {
+            *out_binop = &_strided_bool_mask_noinv0_noinv1_and;
+        }
+    }
+    *out_opdata = NULL;
     return 0;
 }
diff --git a/numpy/core/src/multiarray/na_mask.h b/numpy/core/src/multiarray/na_mask.h
index d0538cda3..7e6b38384 100644
--- a/numpy/core/src/multiarray/na_mask.h
+++ b/numpy/core/src/multiarray/na_mask.h
@@ -34,7 +34,19 @@ PyArray_IsNA(PyObject *obj);
 NPY_NO_EXPORT int
 PyArray_GetMaskInversionFunction(npy_intp mask_stride,
                             PyArray_Descr *mask_dtype,
-                            PyArray_StridedUnaryOp **out_stransfer,
-                            NpyAuxData **out_transferdata);
+                            PyArray_StridedUnaryOp **out_unop,
+                            NpyAuxData **out_opdata);
+
+/*
+ * Gets a function which ANDs together two masks, possibly inverting
+ * one or both of the masks as well.
+ *
+ * The dtype of the output must match 'mask0_dtype'.
+ */
+NPY_NO_EXPORT int
+PyArray_GetMaskAndFunction(
+        npy_intp mask0_stride, PyArray_Descr *mask0_dtype, int invert_mask0,
+        npy_intp mask1_stride, PyArray_Descr *mask1_dtype, int invert_mask1,
+        PyArray_StridedBinaryOp **out_binop, NpyAuxData **out_opdata);
 
 #endif
diff --git a/numpy/core/src/private/lowlevel_strided_loops.h b/numpy/core/src/private/lowlevel_strided_loops.h
index f25019015..98004892b 100644
--- a/numpy/core/src/private/lowlevel_strided_loops.h
+++ b/numpy/core/src/private/lowlevel_strided_loops.h
@@ -42,6 +42,26 @@ typedef void (PyArray_MaskedStridedUnaryOp)(char *dst, npy_intp dst_stride,
                                     NpyAuxData *transferdata);
 
 /*
+ * This function pointer is for binary operations that input two
+ * arbitrarily strided one-dimensional array segments and output
+ * an arbitrarily strided array segment of the same size.
+ * It may be a fully general function, or a specialized function
+ * when the strides or item size have particular known values.
+ *
+ * Examples of binary operations are the basic arithmetic operations,
+ * logical operators AND, OR, and many others.
+ *
+ * The 'transferdata' parameter is slightly special, following a
+ * generic auxiliary data pattern defined in ndarraytypes.h
+ * Use NPY_AUXDATA_CLONE and NPY_AUXDATA_FREE to deal with this data.
+ *
+ */
+typedef void (PyArray_StridedBinaryOp)(char *dst, npy_intp dst_stride,
+                                    char *src0, npy_intp src0_stride,
+                                    char *src1, npy_intp src1_stride,
+                                    npy_intp N, NpyAuxData *transferdata);
+
+/*
  * Gives back a function pointer to a specialized function for copying
  * strided memory.  Returns NULL if there is a problem with the inputs.
  *
@@ -364,6 +384,32 @@ PyArray_PrepareTwoRawArrayIter(int ndim, npy_intp *shape,
                             char **out_dataA, npy_intp *out_stridesA,
                             char **out_dataB, npy_intp *out_stridesB);
 
+/*
+ * The same as PyArray_PrepareOneRawArrayIter, but for three
+ * operands instead of one. Any broadcasting of the three operands
+ * should have already been done before calling this function,
+ * as the ndim and shape is only specified once for all operands.
+ *
+ * Only the strides of the first operand are used to reorder
+ * the dimensions, no attempt to consider all the strides together
+ * is made, as is done in the NpyIter object.
+ *
+ * You can use this together with NPY_RAW_ITER_START and
+ * NPY_RAW_ITER_THREE_NEXT to handle the looping boilerplate of everything
+ * but the innermost loop (which is for idim == 0).
+ *
+ * Returns 0 on success, -1 on failure.
+ */
+NPY_NO_EXPORT int
+PyArray_PrepareThreeRawArrayIter(int ndim, npy_intp *shape,
+                            char *dataA, npy_intp *stridesA,
+                            char *dataB, npy_intp *stridesB,
+                            char *dataC, npy_intp *stridesC,
+                            int *out_ndim, npy_intp *out_shape,
+                            char **out_dataA, npy_intp *out_stridesA,
+                            char **out_dataB, npy_intp *out_stridesB,
+                            char **out_dataC, npy_intp *out_stridesC);
+
 /* Start raw iteration */
 #define NPY_RAW_ITER_START(idim, ndim, coord, shape) \
         memset((coord), 0, (ndim) * sizeof(coord[0])); \
@@ -400,6 +446,27 @@ PyArray_PrepareTwoRawArrayIter(int ndim, npy_intp *shape,
             } \
         } while ((idim) < (ndim))
 
+/* Increment to the next n-dimensional coordinate for three raw arrays */
+#define NPY_RAW_ITER_THREE_NEXT(idim, ndim, coord, shape, \
+                              dataA, stridesA, \
+                              dataB, stridesB, \
+                              dataC, stridesC) \
+            for ((idim) = 1; (idim) < (ndim); ++(idim)) { \
+                if (++(coord)[idim] == (shape)[idim]) { \
+                    (coord)[idim] = 0; \
+                    (dataA) -= ((shape)[idim] - 1) * (stridesA)[idim]; \
+                    (dataB) -= ((shape)[idim] - 1) * (stridesB)[idim]; \
+                    (dataC) -= ((shape)[idim] - 1) * (stridesC)[idim]; \
+                } \
+                else { \
+                    (dataA) += (stridesA)[idim]; \
+                    (dataB) += (stridesB)[idim]; \
+                    (dataC) += (stridesC)[idim]; \
+                    break; \
+                } \
+            } \
+        } while ((idim) < (ndim))
+
 
 /*
  *            TRIVIAL ITERATION