Repository: pydata/xarray

Issue:
# Reading data after saving data from masked arrays results in different numbers

I hit an issue that can go silent for many users.

I am masking my data using [`dask.array.ma.masked_invalid`](https://docs.dask.org/en/stable/generated/dask.array.ma.masked_invalid.html) (my dataset is much larger than my memory). When saving the results and loading them again, the data is changed. In particular, `0.` is assigned to those elements that where masked instead of `np.nan` or `fill_value`, which is `1e+20`.

Below an example that illustrates the issue:
```python
In [1]: import dask.array as da
   ...: import numpy as np
   ...: import xarray as xr

In [2]: ds = xr.DataArray(
   ...:     da.stack(
   ...:         [da.from_array(np.array([[np.nan, np.nan], [np.nan, 2]])) for _ in range(
   ...: 10)],
   ...:         axis=0
   ...:     ).astype('float32'),
   ...:     dims=('time', 'lat', 'lon')
   ...: ).to_dataset(name='mydata')

In [3]: # Mask my data

In [4]: ds = xr.apply_ufunc(da.ma.masked_invalid, ds, dask='allowed') 

In [5]: ds.mean('time').compute()
Out[5]: 
<xarray.Dataset> Size: 16B
Dimensions:  (lat: 2, lon: 2)
Dimensions without coordinates: lat, lon
Data variables:
    mydata   (lat, lon) float32 16B nan nan nan 2.0

In [6]: # Write to file

In [7]: ds.mean('time').to_netcdf('foo.nc')

In [8]: # Read foo.nc

In [9]: foo = xr.open_dataset('foo.nc')

In [10]: foo.compute()
Out[10]: 
<xarray.Dataset> Size: 16B
Dimensions:  (lat: 2, lon: 2)
Dimensions without coordinates: lat, lon
Data variables:
    mydata   (lat, lon) float32 16B 0.0 0.0 0.0 2.0
```

I expected `mydata` to be either `[np.nan, np.nan, np.nan, 2.0]`, `numpy.MaskedArray`, or `[1e+20, 1e+20, 1e+20, 2.0]`, since:
```python
In [11]: ds.mean('time')['mydata'].data.compute()
Out[11]: 
masked_array(
  data=[[--, --],
        [--, 2.0]],
  mask=[[ True,  True],
        [ True, False]],
  fill_value=1e+20,
  dtype=float32)
```
instead of `[0.0, 0.0, 0.0, 2.0]`. No warning is raised, and I do not get why the fill values are replaced by `0.`

A way to address this for my data (but might not be true for all cases) is to use [`xarray.where`](https://docs.xarray.dev/en/stable/generated/xarray.where.html) before writing to file as
```python
In [12]: xr.where(ds.mean('time'), ds.mean('time'), np.nan).to_netcdf('foo.nc')

In [13]: foo = xr.open_dataset('foo.nc')

In [14]: foo.compute()
Out[14]: 
<xarray.Dataset> Size: 16B
Dimensions:  (lat: 2, lon: 2)
Dimensions without coordinates: lat, lon
Data variables:
    mydata   (lat, lon) float32 16B nan nan nan 2.0
```
or mask the data in some other way, e.g. using xarray.where in the beginning instead of xarray.apply_ufunc.

<details>
  <summary>Output of <code>xr.show_versions()</code></summary>
<pre><code>
INSTALLED VERSIONS
------------------
commit: None
python: 3.12.5 | packaged by conda-forge | (main, Aug  8 2024, 18:36:51) [GCC 12.4.0]
python-bits: 64
OS: Linux
OS-release: 5.15.153.1-microsoft-standard-WSL2
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: C.UTF-8
LOCALE: ('C', 'UTF-8')
libhdf5: 1.14.3
libnetcdf: 4.9.2

xarray: 2024.9.0
pandas: 2.2.2
numpy: 2.1.1
scipy: None
netCDF4: 1.7.1
pydap: None
h5netcdf: None
h5py: None
zarr: None
cftime: 1.6.4
nc_time_axis: None
iris: None
bottleneck: None
dask: 2024.9.0
distributed: 2024.9.0
matplotlib: None
cartopy: None
seaborn: None
numbagg: None
fsspec: 2024.9.0
cupy: None
pint: None
sparse: None
flox: None
numpy_groupies: None
setuptools: 73.0.1
pip: 24.2
conda: None
pytest: None
mypy: None
IPython: 8.27.0
sphinx: None
</pre></code>
</details>

Implementation patch:
diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py
index 85ef75f352c..e10ab5a3558 100644
--- a/xarray/core/duck_array_ops.py
+++ b/xarray/core/duck_array_ops.py
@@ -127,6 +127,10 @@ def fail_on_dask_array_input(values, msg=None, func_name=None):
     "masked_invalid", eager_module=np.ma, dask_module="dask.array.ma"
 )
 
+getmaskarray = _dask_or_eager_func(
+    "getmaskarray", eager_module=np.ma, dask_module="dask.array.ma"
+)
+
 
 def sliding_window_view(array, window_shape, axis=None, **kwargs):
     # TODO: some libraries (e.g. jax) don't have this, implement an alternative?
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index 7478d48bb23..ab8b4be775b 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -50,6 +50,7 @@
 from xarray.namedarray.core import NamedArray, _raise_if_any_duplicate_dimensions
 from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.pycompat import (
+    array_type,
     async_to_duck_array,
     integer_types,
     is_0d_dask_array,
@@ -292,13 +293,12 @@ def convert_non_numpy_type(data):
         else:
             data = pandas_data
 
-    if isinstance(data, np.ma.MaskedArray):
-        mask = np.ma.getmaskarray(data)
-        if mask.any():
-            _dtype, fill_value = dtypes.maybe_promote(data.dtype)
-            data = duck_array_ops.where_method(data, ~mask, fill_value)
-        else:
-            data = np.asarray(data)
+    if isinstance(data, np.ma.MaskedArray) or (
+        isinstance(data, array_type("dask"))
+        and isinstance(getattr(data, "_meta", None), np.ma.MaskedArray)
+    ):
+        mask = duck_array_ops.getmaskarray(data)
+        data = duck_array_ops.where_method(data, ~mask)
 
     if isinstance(data, np.matrix):
         data = np.asarray(data)


Test patch:
diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py
index f7d2b516a05..ff5470094ec 100644
--- a/xarray/tests/test_backends.py
+++ b/xarray/tests/test_backends.py
@@ -85,6 +85,7 @@
     mock,
     network,
     parametrize_zarr_format,
+    raise_if_dask_computes,
     requires_cftime,
     requires_dask,
     requires_fsspec,
@@ -2231,6 +2232,28 @@ def test_create_default_indexes(self, tmp_path, create_default_indexes) -> None:
             else:
                 assert len(loaded_ds.xindexes) == 0
 
+    @requires_dask
+    def test_encoding_masked_arrays(self, tmp_path) -> None:
+        store_path = tmp_path / "tmp.nc"
+
+        with raise_if_dask_computes():
+            ds = xr.DataArray(
+                dask.array.from_array(
+                    np.ma.masked_array(
+                        np.array([[np.nan, np.nan], [np.nan, 2]]),
+                        np.array([[True, True], [True, False]]),
+                    )
+                ).astype("float32"),
+                dims=("x", "y"),
+            ).to_dataset(name="mydata")
+
+        expected = ds.mean("x")
+        expected.to_netcdf(
+            store_path, encoding=dict(mydata=dict(_FillValue=np.float32(1e20)))
+        )
+        with open_dataset(store_path, engine=self.engine) as actual:
+            assert_identical(expected.compute(), actual.compute())
+
 
 @requires_netCDF4
 class TestNetCDF4Data(NetCDF4Base):
diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py
index 44a044275e2..1ef6adc9caf 100644
--- a/xarray/tests/test_variable.py
+++ b/xarray/tests/test_variable.py
@@ -2757,7 +2757,7 @@ def test_masked_array(self):
         expected = np.arange(5)
         actual: Any = as_compatible_data(original)
         assert_array_equal(expected, actual)
-        assert np.dtype(int) == actual.dtype
+        assert np.dtype(float) == actual.dtype
 
         original1: Any = np.ma.MaskedArray(np.arange(5), mask=4 * [False] + [True])
         expected1: Any = np.arange(5.0)