Repository: modin-project/modin

Issue:
BUG: modin.numpy.mean does not propagate NaN values correctly
### Modin version checks

- [X] I have checked that this issue has not already been reported.

- [X] I have confirmed this bug exists on the latest released version of Modin.

- [X] I have confirmed this bug exists on the main branch of Modin. (In order to do this you can follow [this guide](https://modin.readthedocs.io/en/stable/getting_started/installation.html#installing-from-the-github-master-branch).)


### Reproducible Example

```python
In [1]: import modin.numpy as np

In [2]: arr = np.array([[1, 2], [3, 4], [5, np.nan]])

In [3]: arr.mean()
Out[3]: 3.0

In [4]: arr._to_numpy().mean()
Out[4]: nan
```


### Issue Description

modin.numpy.mean does not propagate NA values present in the input. We could do so by using `skipna=False` in the `mean` argument, but we use the `skipna=True` behavior to support the `where` argument, by replacing masked values with `nan`. We'd need to replace this behavior and swap `skipna` to `False`.

### Expected Behavior

propagate nans

### Error Logs

<details>

```python-traceback

Replace this line with the error backtrace (if applicable).

```

</details>


### Installed Versions

<details>

Replace this line with the output of pd.show_versions()

</details>



Discussion:

Resolving patch:
diff --git a/modin/numpy/arr.py b/modin/numpy/arr.py
--- a/modin/numpy/arr.py
+++ b/modin/numpy/arr.py
@@ -717,6 +717,24 @@ def __ne__(self, x2):
         result = caller._query_compiler.ne(callee._query_compiler, **kwargs)
         return array(_query_compiler=result, _ndim=new_ndim)
 
+    def _compute_masked_mean(self, mask, output_dtype, axis):
+        # By default, pandas ignores NaN values when doing computations.
+        # NumPy; however, propagates the value by default. We use pandas
+        # default behaviour in order to mask values (by replacing them)
+        # with NaN when initially computing the mean, but we need to propagate
+        # NaN values that were not masked to the final output, so we do a
+        # sum along the same axis (where masked values are 0) to see where
+        # NumPy would propagate NaN, and swap out those values in our result
+        # with NaN.
+        target = mask.where(self, numpy.nan)._query_compiler
+        target = target.astype(
+            {col_name: output_dtype for col_name in target.columns}
+        ).mean(axis=axis)
+        na_propagation_mask = mask.where(self, 0)._query_compiler
+        na_propagation_mask = na_propagation_mask.sum(axis=axis, skipna=False)
+        target = target.where(na_propagation_mask.notna(), numpy.nan)
+        return target
+
     def mean(self, axis=None, dtype=None, out=None, keepdims=None, *, where=True):
         out_dtype = (
             dtype
@@ -730,10 +748,12 @@ def mean(self, axis=None, dtype=None, out=None, keepdims=None, *, where=True):
         if self._ndim == 1:
             if axis == 1:
                 raise numpy.AxisError(1, 1)
-            target = where.where(self, numpy.nan) if isinstance(where, array) else self
-            result = target._query_compiler.astype(
-                {col_name: out_dtype for col_name in target._query_compiler.columns}
-            ).mean(axis=0)
+            if isinstance(where, array):
+                result = self._compute_masked_mean(where, out_dtype, 0)
+            else:
+                result = self._query_compiler.astype(
+                    {col_name: out_dtype for col_name in self._query_compiler.columns}
+                ).mean(axis=0, skipna=False)
             if keepdims:
                 if out is not None and out.shape != (1,):
                     raise ValueError(
@@ -755,9 +775,20 @@ def mean(self, axis=None, dtype=None, out=None, keepdims=None, *, where=True):
             # we return `numpy.nan`, since that is what NumPy would do.
             return result.to_numpy()[0, 0] if where else numpy.nan
         if axis is None:
-            result = self
-            if isinstance(where, array):
-                result = where.where(self, numpy.nan)
+            # If any of the (non-masked) elements of our array are `NaN`, we know that the
+            # result of `mean` must be `NaN`. This is a fastpath to see if any unmasked elements
+            # are `NaN`.
+            contains_na_check = (
+                where.where(self, 0) if isinstance(where, array) else self
+            )
+            if (
+                contains_na_check._query_compiler.isna()
+                .any(axis=1)
+                .any(axis=0)
+                .to_numpy()[0, 0]
+            ):
+                return numpy.nan
+            result = where.where(self, numpy.nan) if isinstance(where, array) else self
             # Since our current QueryCompiler does not have a mean that reduces 2D objects to
             # a single value, we need to calculate the mean ourselves. First though, we need
             # to figure out how many objects that we are taking the mean over (since any
@@ -796,10 +827,12 @@ def mean(self, axis=None, dtype=None, out=None, keepdims=None, *, where=True):
                 else:
                     return array([[numpy.nan]], dtype=out_dtype)
             return result if where is not False else numpy.nan
-        target = where.where(self, numpy.nan) if isinstance(where, array) else self
-        result = target._query_compiler.astype(
-            {col_name: out_dtype for col_name in self._query_compiler.columns}
-        ).mean(axis=axis)
+        if isinstance(where, array):
+            result = self._compute_masked_mean(where, out_dtype, axis)
+        else:
+            result = self._query_compiler.astype(
+                {col_name: out_dtype for col_name in self._query_compiler.columns}
+            ).mean(axis=axis, skipna=False)
         new_ndim = self._ndim - 1 if not keepdims else self._ndim
         if new_ndim == 0:
             return result.to_numpy()[0, 0] if where is not False else numpy.nan
@@ -1147,7 +1180,7 @@ def prod(
             return result if where is not False else initial
         target = where.where(self, 1) if isinstance(where, array) else self
         result = target._query_compiler.astype(
-            {col_name: out_dtype for col_name in self._query_compiler.columns}
+            {col_name: out_dtype for col_name in target._query_compiler.columns}
         ).prod(axis=axis, skipna=False)
         result = result.mul(initial)
         new_ndim = self._ndim - 1 if not keepdims else self._ndim
@@ -1447,7 +1480,7 @@ def sum(
             return result if where is not False else initial
         target = where.where(self, 0) if isinstance(where, array) else self
         result = target._query_compiler.astype(
-            {col_name: out_dtype for col_name in self._query_compiler.columns}
+            {col_name: out_dtype for col_name in target._query_compiler.columns}
         ).sum(axis=axis, skipna=False)
         result = result.add(initial)
         new_ndim = self._ndim - 1 if not keepdims else self._ndim


Test patch:
diff --git a/modin/numpy/test/test_array_axis_functions.py b/modin/numpy/test/test_array_axis_functions.py
--- a/modin/numpy/test/test_array_axis_functions.py
+++ b/modin/numpy/test/test_array_axis_functions.py
@@ -420,6 +420,19 @@ def test_mean():
     modin_result = modin_arr.mean(where=modin_where)
     numpy_result = numpy_arr.mean(where=numpy_where)
     assert modin_result == numpy_result
+    # Test NA propagation
+    numpy_arr = numpy.array([[1, 2], [3, 4], [5, numpy.nan]])
+    modin_arr = np.array([[1, 2], [3, 4], [5, np.nan]])
+    assert numpy.isnan(modin_arr.mean())
+    numpy.testing.assert_array_equal(
+        numpy_arr.mean(axis=1), modin_arr.mean(axis=1)._to_numpy()
+    )
+    numpy.testing.assert_array_equal(
+        numpy_arr.mean(axis=0), modin_arr.mean(axis=0)._to_numpy()
+    )
+    numpy_where = numpy.array([[True, True], [True, True], [True, False]])
+    modin_where = np.array(numpy_where)
+    assert modin_arr.mean(where=modin_where) == numpy_arr.mean(where=numpy_where)
 
 
 def test_prod():