Repository: pydata/xarray

Issue:
scatter plot is slow 
### What happened?

scatter plot is slow when the dataset has large (length) coordinates even though those coordinates are not involved in the scatter plot.

### What did you expect to happen?

scatter plot speed does not depend on coordinates that are not involved in the scatter plot, which was the case at some point in the past


### Minimal Complete Verifiable Example

```Python
import numpy as np
import xarray as xr
from matplotlib import pyplot as plt
%config InlineBackend.figure_format = 'retina'
%matplotlib inline

# Define coordinates
month = np.arange(1, 13, dtype=np.int64)
L = np.arange(1, 13, dtype=np.int64)

# Create random values for the variables SP and SE
np.random.seed(0)  # For reproducibility
SP_values = np.random.rand(len(L), len(month))
SE_values = SP_values + np.random.rand(len(L), len(month))

# Create the dataset
ds = xr.Dataset(
    {
        "SP": (["L", "month"], SP_values),
        "SE": (["L", "month"], SE_values)
    },
    coords={
        "L": L,
        "month": month,
        "S": np.arange(250),
        "model": np.arange(7),
        "M": np.arange(30)
    }
)
# slow
ds.plot.scatter(x='SP', y='SE')

ds = xr.Dataset(
    {
        "SP": (["L", "month"], SP_values),
        "SE": (["L", "month"], SE_values)
    },
    coords={
        "L": L,
        "month": month
    }
)
# fast
ds.plot.scatter(x='SP', y='SE')
```


### MVCE confirmation

- [X] Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
- [X] Complete example — the example is self-contained, including all data and the text of any traceback.
- [X] Verifiable example — the example copy & pastes into an IPython prompt or [Binder notebook](https://mybinder.org/v2/gh/pydata/xarray/main?urlpath=lab/tree/doc/examples/blank_template.ipynb), returning the result.
- [X] New issue — a search of GitHub Issues suggests this is not a duplicate.
- [X] Recent environment — the issue occurs with the latest version of xarray and its dependencies.

### Relevant log output

_No response_

### Anything else we need to know?

For me, slow = 25 seconds and fast = instantaneous 

### Environment

<details>

INSTALLED VERSIONS
------------------
commit: None
python: 3.11.9 | packaged by conda-forge | (main, Apr 19 2024, 18:45:13) [Clang 16.0.6 ]
python-bits: 64
OS: Darwin
OS-release: 23.5.0
machine: x86_64
processor: i386
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.14.3
libnetcdf: 4.9.2

xarray: 2024.6.0
pandas: 2.2.2
numpy: 1.26.4
scipy: 1.13.1
netCDF4: 1.6.5
pydap: installed
h5netcdf: 1.3.0
h5py: 3.11.0
zarr: 2.18.2
cftime: 1.6.4
nc_time_axis: 1.4.1
iris: None
bottleneck: 1.3.8
dask: 2024.6.0
distributed: 2024.6.0
matplotlib: 3.8.4
cartopy: 0.23.0
seaborn: 0.13.2
numbagg: 0.8.1
fsspec: 2024.6.0
cupy: None
pint: 0.24
sparse: 0.15.4
flox: 0.9.8
numpy_groupies: 0.11.1
setuptools: 70.0.0
pip: 24.0
conda: None
pytest: 8.2.2
mypy: None
IPython: 8.17.2
sphinx: None</details>


Discussion:

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 753586c3..0c401c23 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -37,6 +37,8 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+- Fix scatter plot broadcasting unneccesarily. (:issue:`9129`, :pull:`9206`)
+  By `Jimmy Westling <https://github.com/illviljan>`_.
 - Don't convert custom indexes to ``pandas`` indexes when computing a diff (:pull:`9157`)
   By `Justus Magin <https://github.com/keewis>`_.
 - Make :py:func:`testing.assert_allclose` work with numpy 2.0 (:issue:`9165`, :pull:`9166`).
@@ -49,6 +51,9 @@ Bug fixes
   By `Michael Niklas <https://github.com/headtr1ck>`_.
 - Dark themes are now properly detected for ``html[data-theme=dark]``-tags (:pull:`9200`).
   By `Dieter Werthmüller <https://github.com/prisae>`_.
+- Reductions no longer fail for ``np.complex_`` dtype arrays when numbagg is
+  installed.
+  By `Maximilian Roos <https://github.com/max-sixty>`_
 
 Documentation
 ~~~~~~~~~~~~~
diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py
index 38f8f8cd..65ff8667 100644
--- a/xarray/core/datatree.py
+++ b/xarray/core/datatree.py
@@ -61,7 +61,7 @@ if TYPE_CHECKING:
     import pandas as pd
 
     from xarray.core.datatree_io import T_DataTreeNetcdfEngine, T_DataTreeNetcdfTypes
-    from xarray.core.merge import CoercibleValue
+    from xarray.core.merge import CoercibleMapping, CoercibleValue
     from xarray.core.types import ErrorOptions, NetcdfWriteModes, ZarrWriteModes
 
 # """
@@ -954,23 +954,29 @@ class DataTree(
 
         Just like `dict.update` this is an in-place operation.
         """
-        # TODO separate by type
         new_children: dict[str, DataTree] = {}
-        new_variables = {}
-        for k, v in other.items():
-            if isinstance(v, DataTree):
-                # avoid named node being stored under inconsistent key
-                new_child: DataTree = v.copy()
-                # Datatree's name is always a string until we fix that (#8836)
-                new_child.name = str(k)
-                new_children[str(k)] = new_child
-            elif isinstance(v, (DataArray, Variable)):
-                # TODO this should also accommodate other types that can be coerced into Variables
-                new_variables[k] = v
-            else:
-                raise TypeError(f"Type {type(v)} cannot be assigned to a DataTree")
-
-        vars_merge_result = dataset_update_method(self.to_dataset(), new_variables)
+        new_variables: CoercibleMapping
+
+        if isinstance(other, Dataset):
+            new_variables = other
+        else:
+            new_variables = {}
+            for k, v in other.items():
+                if isinstance(v, DataTree):
+                    # avoid named node being stored under inconsistent key
+                    new_child: DataTree = v.copy()
+                    # Datatree's name is always a string until we fix that (#8836)
+                    new_child.name = str(k)
+                    new_children[str(k)] = new_child
+                elif isinstance(v, (DataArray, Variable)):
+                    # TODO this should also accommodate other types that can be coerced into Variables
+                    new_variables[k] = v
+                else:
+                    raise TypeError(f"Type {type(v)} cannot be assigned to a DataTree")
+
+        vars_merge_result = dataset_update_method(
+            self.to_dataset(inherited=False), new_variables
+        )
         data = Dataset._construct_direct(**vars_merge_result._asdict())
 
         # TODO are there any subtleties with preserving order of children like this?
diff --git a/xarray/core/nputils.py b/xarray/core/nputils.py
index 6970d374..1d30fe9d 100644
--- a/xarray/core/nputils.py
+++ b/xarray/core/nputils.py
@@ -191,7 +191,7 @@ def _create_method(name, npmodule=np) -> Callable:
                 or kwargs.get("ddof", 0) == 1
             )
             # TODO: bool?
-            and values.dtype.kind in "uifc"
+            and values.dtype.kind in "uif"
             # and values.dtype.isnative
             and (dtype is None or np.dtype(dtype) == values.dtype)
             # numbagg.nanquantile only available after 0.8.0 and with linear method
diff --git a/xarray/plot/dataset_plot.py b/xarray/plot/dataset_plot.py
index edc2bf43..96b59f61 100644
--- a/xarray/plot/dataset_plot.py
+++ b/xarray/plot/dataset_plot.py
@@ -721,8 +721,8 @@ def _temp_dataarray(ds: Dataset, y: Hashable, locals_: dict[str, Any]) -> DataAr
     """Create a temporary datarray with extra coords."""
     from xarray.core.dataarray import DataArray
 
-    # Base coords:
-    coords = dict(ds.coords)
+    coords = dict(ds[y].coords)
+    dims = set(ds[y].dims)
 
     # Add extra coords to the DataArray from valid kwargs, if using all
     # kwargs there is a risk that we add unnecessary dataarrays as
@@ -732,12 +732,17 @@ def _temp_dataarray(ds: Dataset, y: Hashable, locals_: dict[str, Any]) -> DataAr
     coord_kwargs = locals_.keys() & valid_coord_kwargs
     for k in coord_kwargs:
         key = locals_[k]
-        if ds.data_vars.get(key) is not None:
-            coords[key] = ds[key]
+        darray = ds.get(key)
+        if darray is not None:
+            coords[key] = darray
+            dims.update(darray.dims)
+
+    # Trim dataset from unneccessary dims:
+    ds_trimmed = ds.drop_dims(ds.sizes.keys() - dims)  # TODO: Use ds.dims in the future
 
     # The dataarray has to include all the dims. Broadcast to that shape
     # and add the additional coords:
-    _y = ds[y].broadcast_like(ds)
+    _y = ds[y].broadcast_like(ds_trimmed)
 
     return DataArray(_y, coords=coords)
 


Test patch:
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index 080447ce..b000de31 100644
--- a/xarray/tests/test_computation.py
+++ b/xarray/tests/test_computation.py
@@ -2598,3 +2598,11 @@ def test_cross(a, b, ae, be, dim: str, axis: int, use_dask: bool) -> None:
 
     actual = xr.cross(a, b, dim=dim)
     xr.testing.assert_duckarray_allclose(expected, actual)
+
+
+@pytest.mark.parametrize("compute_backend", ["numbagg"], indirect=True)
+def test_complex_number_reduce(compute_backend):
+    da = xr.DataArray(np.ones((2,), dtype=np.complex64), dims=["x"])
+    # Check that xarray doesn't call into numbagg, which doesn't compile for complex
+    # numbers at the moment (but will when numba supports dynamic compilation)
+    da.min()
diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py
index f2b58fa2..f7cff17b 100644
--- a/xarray/tests/test_datatree.py
+++ b/xarray/tests/test_datatree.py
@@ -244,11 +244,6 @@ class TestUpdate:
         dt: DataTree = DataTree()
         dt.update({"foo": xr.DataArray(0), "a": DataTree()})
         expected = DataTree.from_dict({"/": xr.Dataset({"foo": 0}), "a": None})
-        print(dt)
-        print(dt.children)
-        print(dt._children)
-        print(dt["a"])
-        print(expected)
         assert_equal(dt, expected)
 
     def test_update_new_named_dataarray(self):
@@ -268,14 +263,38 @@ class TestUpdate:
     def test_update_overwrite(self):
         actual = DataTree.from_dict({"a": DataTree(xr.Dataset({"x": 1}))})
         actual.update({"a": DataTree(xr.Dataset({"x": 2}))})
-
         expected = DataTree.from_dict({"a": DataTree(xr.Dataset({"x": 2}))})
+        assert_equal(actual, expected)
 
-        print(actual)
-        print(expected)
-
+    def test_update_coordinates(self):
+        expected = DataTree.from_dict({"/": xr.Dataset(coords={"a": 1})})
+        actual = DataTree.from_dict({"/": xr.Dataset()})
+        actual.update(xr.Dataset(coords={"a": 1}))
         assert_equal(actual, expected)
 
+    def test_update_inherited_coords(self):
+        expected = DataTree.from_dict(
+            {
+                "/": xr.Dataset(coords={"a": 1}),
+                "/b": xr.Dataset(coords={"c": 1}),
+            }
+        )
+        actual = DataTree.from_dict(
+            {
+                "/": xr.Dataset(coords={"a": 1}),
+                "/b": xr.Dataset(),
+            }
+        )
+        actual["/b"].update(xr.Dataset(coords={"c": 1}))
+        assert_identical(actual, expected)
+
+        # DataTree.identical() currently does not require that non-inherited
+        # coordinates are defined identically, so we need to check this
+        # explicitly
+        actual_node = actual.children["b"].to_dataset(inherited=False)
+        expected_node = expected.children["b"].to_dataset(inherited=False)
+        assert_identical(actual_node, expected_node)
+
 
 class TestCopy:
     def test_copy(self, create_test_datatree):
diff --git a/xarray/tests/test_missing.py b/xarray/tests/test_missing.py
index da9513a7..bd75f633 100644
--- a/xarray/tests/test_missing.py
+++ b/xarray/tests/test_missing.py
@@ -421,46 +421,37 @@ def test_ffill():
     assert_equal(actual, expected)
 
 
-def test_ffill_use_bottleneck_numbagg():
+@pytest.mark.parametrize("compute_backend", [None], indirect=True)
+@pytest.mark.parametrize("method", ["ffill", "bfill"])
+def test_b_ffill_use_bottleneck_numbagg(method, compute_backend):
+    """
+    bfill & ffill fail if both bottleneck and numba are disabled
+    """
     da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
-    with xr.set_options(use_bottleneck=False, use_numbagg=False):
-        with pytest.raises(RuntimeError):
-            da.ffill("x")
+    with pytest.raises(RuntimeError):
+        getattr(da, method)("x")
 
 
 @requires_dask
-def test_ffill_use_bottleneck_dask():
+@pytest.mark.parametrize("compute_backend", [None], indirect=True)
+@pytest.mark.parametrize("method", ["ffill", "bfill"])
+def test_b_ffill_use_bottleneck_dask(method, compute_backend):
+    """
+    ffill fails if both bottleneck and numba are disabled, on dask arrays
+    """
     da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
-    da = da.chunk({"x": 1})
-    with xr.set_options(use_bottleneck=False, use_numbagg=False):
-        with pytest.raises(RuntimeError):
-            da.ffill("x")
+    with pytest.raises(RuntimeError):
+        getattr(da, method)("x")
 
 
 @requires_numbagg
 @requires_dask
-def test_ffill_use_numbagg_dask():
-    with xr.set_options(use_bottleneck=False):
-        da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
-        da = da.chunk(x=-1)
-        # Succeeds with a single chunk:
-        _ = da.ffill("x").compute()
-
-
-def test_bfill_use_bottleneck():
-    da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
-    with xr.set_options(use_bottleneck=False, use_numbagg=False):
-        with pytest.raises(RuntimeError):
-            da.bfill("x")
-
-
-@requires_dask
-def test_bfill_use_bottleneck_dask():
+@pytest.mark.parametrize("compute_backend", ["numbagg"], indirect=True)
+def test_ffill_use_numbagg_dask(compute_backend):
     da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
-    da = da.chunk({"x": 1})
-    with xr.set_options(use_bottleneck=False, use_numbagg=False):
-        with pytest.raises(RuntimeError):
-            da.bfill("x")
+    da = da.chunk(x=-1)
+    # Succeeds with a single chunk:
+    _ = da.ffill("x").compute()
 
 
 @requires_bottleneck
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
index b302ad3a..fa08e997 100644
--- a/xarray/tests/test_plot.py
+++ b/xarray/tests/test_plot.py
@@ -3416,3 +3416,43 @@ def test_9155() -> None:
         data = xr.DataArray([1, 2, 3], dims=["x"])
         fig, ax = plt.subplots(ncols=1, nrows=1)
         data.plot(ax=ax)
+
+
+@requires_matplotlib
+def test_temp_dataarray() -> None:
+    from xarray.plot.dataset_plot import _temp_dataarray
+
+    x = np.arange(1, 4)
+    y = np.arange(4, 6)
+    var1 = np.arange(x.size * y.size).reshape((x.size, y.size))
+    var2 = np.arange(x.size * y.size).reshape((x.size, y.size))
+    ds = xr.Dataset(
+        {
+            "var1": (["x", "y"], var1),
+            "var2": (["x", "y"], 2 * var2),
+            "var3": (["x"], 3 * x),
+        },
+        coords={
+            "x": x,
+            "y": y,
+            "model": np.arange(7),
+        },
+    )
+
+    # No broadcasting:
+    y_ = "var1"
+    locals_ = {"x": "var2"}
+    da = _temp_dataarray(ds, y_, locals_)
+    assert da.shape == (3, 2)
+
+    # Broadcast from 1 to 2dim:
+    y_ = "var3"
+    locals_ = {"x": "var1"}
+    da = _temp_dataarray(ds, y_, locals_)
+    assert da.shape == (3, 2)
+
+    # Ignore non-valid coord kwargs:
+    y_ = "var3"
+    locals_ = dict(x="x", extend="var2")
+    da = _temp_dataarray(ds, y_, locals_)
+    assert da.shape == (3,)