ENH: missingdata: Some tests/fixes for reduce method with skipna=True

2 files changed, 35 insertions, 7 deletions

diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index cdcc4ade6..d1b76d070 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -2980,6 +2980,16 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
                         "isn't implemented yet");
                 goto fail;
             }
+
+            /*
+             * If the result has a mask (i.e. from the out= parameter),
+             * Set it to all exposed.
+             */
+            if (PyArray_HASMASKNA(result)) {
+                if (PyArray_AssignMaskNA(result, 1) < 0) {
+                    goto fail;
+                }
+            }
         }
     }
 
@@ -3017,18 +3027,18 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
     /* Add mask-related flags */
     if (use_maskna) {
         if (skipna) {
-            /* Need the input's mask to determine what to skip */
-            op_flags[0] |= NPY_ITER_USE_MASKNA;
             /* The output's mask has been set to all exposed already */
-            op_flags[1] |= NPY_ITER_IGNORE_MASKNA;
+            op_flags[0] |= NPY_ITER_IGNORE_MASKNA;
+            /* Need the input's mask to determine what to skip */
+            op_flags[1] |= NPY_ITER_USE_MASKNA;
 
             /* TODO: allocate a temporary buffer for inverting the mask */
         }
         else {
-            /* The input's mask is already incorporated in the output's mask */
-            op_flags[0] |= NPY_ITER_IGNORE_MASKNA;
             /* Iterate over the output's mask */
-            op_flags[1] |= NPY_ITER_USE_MASKNA;
+            op_flags[0] |= NPY_ITER_USE_MASKNA;
+            /* The input's mask is already incorporated in the output's mask */
+            op_flags[1] |= NPY_ITER_IGNORE_MASKNA;
         }
     }
     else {
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index ec8fda00f..5ad6f0990 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -523,7 +523,8 @@ def test_ufunc_1D():
 def test_ufunc_reduce_1D():
     a = np.arange(3.0, maskna=True)
     b = np.array(0.5)
-    c = np.array(0.5, maskna=True)
+    c_orig = np.array(0.5)
+    c = c_orig.view(maskna=True)
 
     # Since 'a' has no NA values, this should work
     np.add.reduce(a, out=b)
@@ -534,7 +535,9 @@ def test_ufunc_reduce_1D():
     assert_raises(ValueError, np.add.reduce, a, out=b)
 
     # With an NA value, the output parameter can still be an NA-array
+    c_orig[...] = 0.5
     np.add.reduce(a, out=c)
+    assert_equal(c_orig, 0.5)
     assert_(np.isna(c))
 
     # Should not touch the out= element when assigning NA
@@ -548,6 +551,21 @@ def test_ufunc_reduce_1D():
     ret = np.add.reduce(a)
     assert_(np.isna(ret))
 
+    # With 'skipna=True'
+    ret = np.add.reduce(a, skipna=True)
+    assert_equal(ret, 2.0)
+
+    # With 'skipna=True', and out= parameter
+    b[...] = 0.5
+    np.add.reduce(a, skipna=True, out=b)
+    assert_equal(b, 2.0)
+
+    # With 'skipna=True', and out= parameter with a mask
+    c[...] = 0.5
+    c[...] = np.NA
+    np.add.reduce(a, skipna=True, out=c)
+    assert_(not np.isna(c))
+    assert_equal(c, 2.0)
 
 if __name__ == "__main__":
     run_module_suite()