ENH: umath: Fix reduce with NAs for ufuncs that have no unit

11 files changed, 70 insertions, 79 deletions

diff --git a/numpy/core/code_generators/genapi.py b/numpy/core/code_generators/genapi.py
index b083f6e77..b3b56fa47 100644
--- a/numpy/core/code_generators/genapi.py
+++ b/numpy/core/code_generators/genapi.py
@@ -23,6 +23,8 @@ __docformat__ = 'restructuredtext'
 # The files under src/ that are scanned for API functions
 API_FILES = [join('multiarray', 'methods.c'),
              join('multiarray', 'arrayobject.c'),
+             join('multiarray', 'array_assign_scalar.c'),
+             join('multiarray', 'array_assign_array.c'),
              join('multiarray', 'flagsobject.c'),
              join('multiarray', 'descriptor.c'),
              join('multiarray', 'iterators.c'),
diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index 22f0a349a..dc73df1b3 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -330,6 +330,8 @@ multiarray_funcs_api = {
     'PyArray_AssignOne':                    291,
     'PyArray_AssignNA':                     292,
     'PyArray_AssignMaskNA':                 293,
+    'PyArray_AssignRawScalar':              294,
+    'PyArray_AssignArray':                  295,
 }
 
 ufunc_types_api = {
diff --git a/numpy/core/src/multiarray/array_assign.h b/numpy/core/src/multiarray/array_assign.h
index 0f9f613b0..c58a63d22 100644
--- a/numpy/core/src/multiarray/array_assign.h
+++ b/numpy/core/src/multiarray/array_assign.h
@@ -2,58 +2,6 @@
 #define _NPY_PRIVATE__ARRAY_ASSIGN_H_
 
 /*
- * Assigns a scalar value specified by 'src_dtype' and 'src_data'
- * to elements of 'dst'.
- *
- * dst: The destination array.
- * src_dtype: The data type of the source scalar.
- * src_data: The memory element of the source scalar.
- * wheremask: If non-NULL, a boolean mask specifying where to copy.
- * casting: An exception is raised if the assignment violates this
- *          casting rule.
- * 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.
- */
-NPY_NO_EXPORT int
-array_assign_scalar(PyArrayObject *dst,
-                    PyArray_Descr *src_dtype, char *src_data,
-                    PyArrayObject *wheremask,
-                    NPY_CASTING casting,
-                    npy_bool preservena, npy_bool *preservewhichna);
-
-/*
- * An array assignment function for copying arrays, broadcasting 'src' into
- * 'dst'. This function makes a temporary copy of 'src' if 'src' and
- * 'dst' overlap, to be able to handle views of the same data with
- * different strides.
- *
- * dst: The destination array.
- * src: The source array.
- * wheremask: If non-NULL, a boolean mask specifying where to copy.
- * casting: An exception is raised if the copy violates this
- *          casting rule.
- * 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_array(PyArrayObject *dst, PyArrayObject *src,
-                    PyArrayObject *wheremask,
-                    NPY_CASTING casting,
-                    npy_bool preservena, npy_bool *preservewhichna);
-
-/*
  * An array assignment function for copying arrays, treating the
  * arrays as flat according to their respective ordering rules.
  * This function makes a temporary copy of 'src' if 'src' and
@@ -64,7 +12,6 @@ array_assign_array(PyArrayObject *dst, PyArrayObject *src,
  * dst_order: The rule for how 'dst' is to be made flat.
  * src: The source array.
  * src_order: The rule for how 'src' is to be made flat.
- * wheremask: If non-NULL, a boolean mask specifying where to copy.
  * casting: An exception is raised if the copy violates this
  *          casting rule.
  * preservena: If 0, overwrites everything in 'dst', if 1, it
@@ -75,12 +22,13 @@ array_assign_array(PyArrayObject *dst, PyArrayObject *src,
  *
  * Returns 0 on success, -1 on failure.
  */
+/* Not yet implemented
 NPY_NO_EXPORT int
-array_assign_flat(PyArrayObject *dst, NPY_ORDER dst_order,
+PyArray_AssignArrayAsFlat(PyArrayObject *dst, NPY_ORDER dst_order,
                   PyArrayObject *src, NPY_ORDER src_order,
-                  PyArrayObject *wheremask,
                   NPY_CASTING casting,
                   npy_bool preservena, npy_bool *preservewhichna);
+*/
 
 
 
diff --git a/numpy/core/src/multiarray/array_assign_array.c b/numpy/core/src/multiarray/array_assign_array.c
index 96bfd1f45..7e5524714 100644
--- a/numpy/core/src/multiarray/array_assign_array.c
+++ b/numpy/core/src/multiarray/array_assign_array.c
@@ -361,10 +361,28 @@ raw_array_wheremasked_assign_array_preservena(int ndim, npy_intp *shape,
     return (needs_api && PyErr_Occurred()) ? -1 : 0;
 }
 
-
-/* See array_assign.h for documentation */
+/*NUMPY_API
+ *
+ * An array assignment function for copying arrays, broadcasting 'src' into
+ * 'dst'. This function makes a temporary copy of 'src' if 'src' and
+ * 'dst' overlap, to be able to handle views of the same data with
+ * different strides.
+ *
+ * dst: The destination array.
+ * src: The source array.
+ * wheremask: If non-NULL, a boolean mask specifying where to copy.
+ * casting: An exception is raised if the copy violates this
+ *          casting rule.
+ * 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_array(PyArrayObject *dst, PyArrayObject *src,
+PyArray_AssignArray(PyArrayObject *dst, PyArrayObject *src,
                     PyArrayObject *wheremask,
                     NPY_CASTING casting,
                     npy_bool preservena, npy_bool *preservewhichna)
@@ -389,7 +407,8 @@ array_assign_array(PyArrayObject *dst, PyArrayObject *src,
             }
         }
 
-        return array_assign_scalar(dst, PyArray_DESCR(src), PyArray_DATA(src),
+        return PyArray_AssignRawScalar(
+                            dst, PyArray_DESCR(src), PyArray_DATA(src),
                             wheremask, casting, preservena, preservewhichna);
     }
 
@@ -499,7 +518,7 @@ array_assign_array(PyArrayObject *dst, PyArrayObject *src,
             }
         }
 
-        if (array_assign_array(tmp, src,
+        if (PyArray_AssignArray(tmp, src,
                                 NULL, NPY_UNSAFE_CASTING, 0, NULL) < 0) {
             Py_DECREF(tmp);
             goto fail;
diff --git a/numpy/core/src/multiarray/array_assign_scalar.c b/numpy/core/src/multiarray/array_assign_scalar.c
index 5f2675176..66d7555d7 100644
--- a/numpy/core/src/multiarray/array_assign_scalar.c
+++ b/numpy/core/src/multiarray/array_assign_scalar.c
@@ -302,9 +302,29 @@ raw_array_wheremasked_assign_scalar_preservena(int ndim, npy_intp *shape,
 }
 
 
-/* See array_assign.h for documentation */
+/*NUMPY_API
+ *
+ * Assigns a scalar value specified by 'src_dtype' and 'src_data'
+ * to elements of 'dst'.
+ *
+ * dst: The destination array.
+ * src_dtype: The data type of the source scalar.
+ * src_data: The memory element of the source scalar.
+ * wheremask: If non-NULL, a boolean mask specifying where to copy.
+ * casting: An exception is raised if the assignment violates this
+ *          casting rule.
+ * 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.
+ */
 NPY_NO_EXPORT int
-array_assign_scalar(PyArrayObject *dst,
+PyArray_AssignRawScalar(PyArrayObject *dst,
                     PyArray_Descr *src_dtype, char *src_data,
                     PyArrayObject *wheremask,
                     NPY_CASTING casting,
diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
index 7e4aadca7..9189839b5 100644
--- a/numpy/core/src/multiarray/arrayobject.c
+++ b/numpy/core/src/multiarray/arrayobject.c
@@ -52,7 +52,6 @@ maintainer email:  oliphant.travis@ieee.org
 #include "buffer.h"
 #include "na_singleton.h"
 #include "na_mask.h"
-#include "array_assign.h"
 
 /*NUMPY_API
   Compute the size of an array (in number of items)
@@ -203,7 +202,7 @@ PyArray_CopyObject(PyArrayObject *dest, PyObject *src_object)
                 }
 
                 /* TODO: switch to SAME_KIND casting */
-                retcode = array_assign_scalar(dest, dtype, value,
+                retcode = PyArray_AssignRawScalar(dest, dtype, value,
                                         NULL, NPY_UNSAFE_CASTING, 0, NULL);
                 Py_DECREF(dtype);
                 Py_DECREF(src_object);
diff --git a/numpy/core/src/multiarray/convert.c b/numpy/core/src/multiarray/convert.c
index 1dbb387b2..5640f1978 100644
--- a/numpy/core/src/multiarray/convert.c
+++ b/numpy/core/src/multiarray/convert.c
@@ -15,7 +15,6 @@
 #include "arrayobject.h"
 #include "mapping.h"
 #include "lowlevel_strided_loops.h"
-#include "array_assign.h"
 #include "scalartypes.h"
 
 #include "convert.h"
@@ -330,7 +329,7 @@ PyArray_FillWithScalar(PyArrayObject *arr, PyObject *obj)
             return -1;
         }
     }
-    /* Use array_assign_scalar if 'obj' is a numpy scalar object */
+    /* NumPy scalar */
     else if (PyArray_IsScalar(obj, Generic)) {
         dtype = PyArray_DescrFromScalar(obj);
         if (dtype == NULL) {
@@ -405,7 +404,7 @@ PyArray_FillWithScalar(PyArrayObject *arr, PyObject *obj)
     /* Use the value pointer we got if possible */
     if (value != NULL) {
         /* TODO: switch to SAME_KIND casting */
-        retcode = array_assign_scalar(arr, dtype, value,
+        retcode = PyArray_AssignRawScalar(arr, dtype, value,
                                 NULL, NPY_UNSAFE_CASTING, 0, NULL);
         Py_DECREF(dtype);
         return retcode;
@@ -464,7 +463,7 @@ PyArray_AssignZero(PyArrayObject *dst,
     }
     value = 0;
 
-    retcode = array_assign_scalar(dst, bool_dtype, (char *)&value,
+    retcode = PyArray_AssignRawScalar(dst, bool_dtype, (char *)&value,
                                 wheremask, NPY_SAFE_CASTING,
                                 preservena, preservewhichna);
 
@@ -502,7 +501,7 @@ PyArray_AssignOne(PyArrayObject *dst,
     }
     value = 1;
 
-    retcode = array_assign_scalar(dst, bool_dtype, (char *)&value,
+    retcode = PyArray_AssignRawScalar(dst, bool_dtype, (char *)&value,
                                 wheremask, NPY_SAFE_CASTING,
                                 preservena, preservewhichna);
 
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index b60d6ba2d..48d81d75b 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -25,7 +25,6 @@
 #include "_datetime.h"
 #include "datetime_strings.h"
 #include "na_singleton.h"
-#include "array_assign.h"
 
 /*
  * Reading from a file or a string.
@@ -2922,7 +2921,7 @@ PyArray_CopyAnyInto(PyArrayObject *dst, PyArrayObject *src)
 NPY_NO_EXPORT int
 PyArray_CopyInto(PyArrayObject *dst, PyArrayObject *src)
 {
-    return array_assign_array(dst, src, NULL, NPY_UNSAFE_CASTING, 0, NULL);
+    return PyArray_AssignArray(dst, src, NULL, NPY_UNSAFE_CASTING, 0, NULL);
 }
 
 /*NUMPY_API
@@ -2933,7 +2932,7 @@ PyArray_CopyInto(PyArrayObject *dst, PyArrayObject *src)
 NPY_NO_EXPORT int
 PyArray_MoveInto(PyArrayObject *dst, PyArrayObject *src)
 {
-    return array_assign_array(dst, src, NULL, NPY_UNSAFE_CASTING, 0, NULL);
+    return PyArray_AssignArray(dst, src, NULL, NPY_UNSAFE_CASTING, 0, NULL);
 }
 
 /*NUMPY_API
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 97248ca38..b4c212871 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -50,7 +50,6 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0;
 #include "datetime_busdaycal.h"
 #include "na_singleton.h"
 #include "na_mask.h"
-#include "array_assign.h"
 
 /* Only here for API compatibility */
 NPY_NO_EXPORT PyTypeObject PyBigArray_Type;
@@ -1730,7 +1729,7 @@ array_copyto(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds)
         }
     }
 
-    if (array_assign_array(dst, src,
+    if (PyArray_AssignArray(dst, src,
                     wheremask, casting, preservena, NULL) < 0) {
         goto fail;
     }
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 3c20ac5c6..8ab632d5b 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -2848,8 +2848,13 @@ initialize_reduce_result(int identity, PyArrayObject *result,
         }
     }
 
-    /* Copy the elements into the result to start */
-    if (PyArray_CopyInto(result, arr_view) < 0) {
+    /*
+     * Copy the elements into the result to start, with
+     * 'preservena' set to True so that we don't overwrite
+     * what we already calculated in ReduceNAMask.
+     */
+    if (PyArray_AssignArray(result, arr_view, NULL, NPY_UNSAFE_CASTING,
+                            1, NULL) < 0) {
         Py_DECREF(arr_view);
         return NULL;
     }
diff --git a/numpy/core/tests/test_api.py b/numpy/core/tests/test_api.py
index f98c8c8c1..2c98810fa 100644
--- a/numpy/core/tests/test_api.py
+++ b/numpy/core/tests/test_api.py
@@ -171,7 +171,7 @@ def test_copyto_maskna():
     a_orig = np.zeros((2,3), dtype='f8')
     a = a_orig.view(maskna=True)
 
-    # Simple copy to from non-masked to NA-masked
+    # Simple copy from non-masked to NA-masked
     a[...] = np.NA
     np.copyto(a, np.arange(6).reshape(2,3))
     assert_equal(a, [[0,1,2],[3,4,5]])
@@ -179,7 +179,7 @@ def test_copyto_maskna():
     np.copyto(a.T, np.arange(6).reshape(3,2) + 1)
     assert_equal(a, [[1,3,5],[2,4,6]])
 
-    # Simple copy to from NA-masked to NA-masked
+    # Simple copy from NA-masked to NA-masked
     a[...] = np.NA
     a[1,2] = 12
     tmp = np.arange(6, maskna=True).reshape(2,3)
@@ -240,6 +240,5 @@ def test_copyto_maskna():
     assert_equal(a_orig, [[0,3,3],[3,4,12]])
     assert_equal(np.isna(a), [[0,1,0],[1,1,1]])
 
-
 if __name__ == "__main__":
     run_module_suite()