ENH: missingdata: Move ReduceMaskNAArray out of the public API

Keep this function private, to keep flexibility open for longer.
Also added what will likely be the necessary parameters for future
where mask and multi-NA expansion.

4 files changed, 59 insertions, 17 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index e21cc96b4..527307055 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -325,17 +325,16 @@ multiarray_funcs_api = {
     'PyArray_HasNASupport':                 286,
     'PyArray_ContainsNA':                   287,
     'PyArray_AllocateMaskNA':               288,
-    'PyArray_ReduceMaskNAArray':            289,
-    'PyArray_CreateSortedStridePerm':       290,
-    'PyArray_AssignZero':                   291,
-    'PyArray_AssignOne':                    292,
-    'PyArray_AssignNA':                     293,
-    'PyArray_AssignMaskNA':                 294,
-    'PyArray_AssignRawScalar':              295,
-    'PyArray_AssignArray':                  296,
-    'PyArray_ReduceWrapper':                297,
-    'PyArray_RemoveAxesInPlace':            298,
-    'PyArray_DebugPrint':                   299,
+    'PyArray_CreateSortedStridePerm':       289,
+    'PyArray_AssignZero':                   290,
+    'PyArray_AssignOne':                    291,
+    'PyArray_AssignNA':                     292,
+    'PyArray_AssignMaskNA':                 293,
+    'PyArray_AssignRawScalar':              294,
+    'PyArray_AssignArray':                  295,
+    'PyArray_ReduceWrapper':                296,
+    'PyArray_RemoveAxesInPlace':            297,
+    'PyArray_DebugPrint':                   298,
 }
 
 ufunc_types_api = {
diff --git a/numpy/core/src/multiarray/na_mask.c b/numpy/core/src/multiarray/na_mask.c
index f6267beac..2684c472c 100644
--- a/numpy/core/src/multiarray/na_mask.c
+++ b/numpy/core/src/multiarray/na_mask.c
@@ -699,22 +699,28 @@ raw_reduce_maskna_array(int ndim, npy_intp *shape,
     return 0;
 }
 
-/*NUMPY_API
- *
+/*
  * This function performs a reduction on the masks for an array.
  *
  * This is for use with a reduction where 'skipna=False'.
  *
+ * operand: The operand for which the reduction is being done. This array
+ *          must have an NA mask.
  * result: The result array, which should have the same 'ndim' as
  *         'operand' but with dimensions of size one for every reduction
  *         axis. This array must have an NA mask.
- * operand: The operand for which the reduction is being done. This array
- *          must have an NA mask.
+ * wheremask: NOT SUPPORTED YET, but is for a boolean mask which can
+ *            broadcast to 'result', indicating where to do the reduction.
+ *            Should pass in NULL.
+ * skipwhichna: NOT SUPPORTED YET, but for future expansion to multi-NA,
+ *              where reductions can be done on NAs with a subset of
+ *              the possible payloads. Should pass in NULL.
  *
  * Returns 0 on success, -1 on failure.
  */
 NPY_NO_EXPORT int
-PyArray_ReduceMaskNAArray(PyArrayObject *result, PyArrayObject *operand)
+PyArray_ReduceMaskNAArray(PyArrayObject *operand, PyArrayObject *result,
+                            PyArrayObject *wheremask, npy_bool *skipwhichna)
 {
     int idim, ndim;
     npy_intp result_strides[NPY_MAXDIMS];
@@ -735,6 +741,20 @@ PyArray_ReduceMaskNAArray(PyArrayObject *result, PyArrayObject *operand)
         return -1;
     }
 
+    /* Validate that the parameters for future expansion are NULL */
+    if (wheremask != NULL) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "the NA mask reduction operation in NumPy does not "
+                "yet support a where mask");
+        return -1;
+    }
+    if (skipwhichna != NULL) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "multi-NA support is not yet implemented in "
+                "the NA mask reduction operation");
+        return -1;
+    }
+
     /* Need to make sure the appropriate strides are 0 in 'result' */
     result_shape = PyArray_SHAPE(result);
     operand_shape = PyArray_SHAPE(operand);
diff --git a/numpy/core/src/multiarray/na_mask.h b/numpy/core/src/multiarray/na_mask.h
index 920b32dfa..547389e77 100644
--- a/numpy/core/src/multiarray/na_mask.h
+++ b/numpy/core/src/multiarray/na_mask.h
@@ -43,4 +43,27 @@ PyArray_GetMaskAndFunction(
         npy_intp mask1_stride, PyArray_Descr *mask1_dtype, int invert_mask1,
         PyArray_StridedBinaryOp **out_binop, NpyAuxData **out_opdata);
 
+/*
+ * This function performs a reduction on the masks for an array.
+ *
+ * This is for use with a reduction where 'skipna=False'.
+ *
+ * operand: The operand for which the reduction is being done. This array
+ *          must have an NA mask.
+ * result: The result array, which should have the same 'ndim' as
+ *         'operand' but with dimensions of size one for every reduction
+ *         axis. This array must have an NA mask.
+ * wheremask: NOT SUPPORTED YET, but is for a boolean mask which can
+ *            broadcast to 'result', indicating where to do the reduction.
+ *            Should pass in NULL.
+ * skipwhichna: NOT SUPPORTED YET, but for future expansion to multi-NA,
+ *              where reductions can be done on NAs with a subset of
+ *              the possible payloads. Should pass in NULL.
+ *
+ * Returns 0 on success, -1 on failure.
+ */
+NPY_NO_EXPORT int
+PyArray_ReduceMaskNAArray(PyArrayObject *operand, PyArrayObject *result,
+                            PyArrayObject *wheremask, npy_bool *skipwhichna);
+
 #endif
diff --git a/numpy/core/src/multiarray/reduction.c b/numpy/core/src/multiarray/reduction.c
index aafe370ab..ecf0935ae 100644
--- a/numpy/core/src/multiarray/reduction.c
+++ b/numpy/core/src/multiarray/reduction.c
@@ -613,7 +613,7 @@ PyArray_ReduceWrapper(PyArrayObject *operand, PyArrayObject *out,
      * the required NA masking semantics.
      */
     if (use_maskna && !skipna) {
-        if (PyArray_ReduceMaskNAArray(result, operand) < 0) {
+        if (PyArray_ReduceMaskNAArray(operand, result, NULL, NULL) < 0) {
             goto fail;
         }