Filter null values before plotting (pydata#8535)

Illviljan · web-flow · commit 2971994ef1dd · 2023-12-13T08:44:54.000-07:00
* Remove nulls when plotting.

* Update test_plot.py

* Update test_plot.py

* Update whats-new.rst

* Update test_plot.py

* only filter on scatter plots as that is safe.
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -46,6 +46,8 @@ Documentation
 Internal Changes
 ~~~~~~~~~~~~~~~~
 
+- Remove null values before plotting. (:pull:`8535`).
+  By `Jimmy Westling <https://github.com/illviljan>`_.
 
 .. _whats-new.2023.12.0:
 
diff --git a/xarray/plot/dataarray_plot.py b/xarray/plot/dataarray_plot.py
@@ -944,6 +944,12 @@ def newplotfunc(
         if plotfunc.__name__ == "scatter":
             size_ = kwargs.pop("_size", markersize)
             size_r = _MARKERSIZE_RANGE
+
+            # Remove any nulls, .where(m, drop=True) doesn't work when m is
+            # a dask array, so load the array to memory.
+            # It will have to be loaded to memory at some point anyway:
+            darray = darray.load()
+            darray = darray.where(darray.notnull(), drop=True)
         else:
             size_ = kwargs.pop("_size", linewidth)
             size_r = _LINEWIDTH_RANGE
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
@@ -3372,3 +3372,16 @@ def test_plot1d_default_rcparams() -> None:
         np.testing.assert_allclose(
             ax.collections[0].get_edgecolor(), mpl.colors.to_rgba_array("k")
         )
+
+
+@requires_matplotlib
+def test_plot1d_filtered_nulls() -> None:
+    ds = xr.tutorial.scatter_example_dataset(seed=42)
+    y = ds.y.where(ds.y > 0.2)
+    expected = y.notnull().sum().item()
+
+    with figure_context():
+        pc = y.plot.scatter()
+        actual = pc.get_offsets().shape[0]
+
+        assert expected == actual
