From df0ad0e6aa33cd206129d863553a7c08acf71bfd Mon Sep 17 00:00:00 2001
From: Maximilian Roos <5635139+max-sixty@users.noreply.github.com>
Date: Thu, 17 Sep 2020 17:25:52 -0700
Subject: [PATCH] Revert "Fix optimize for chunked DataArray (#4432)"

This reverts commit 9a8a62ba551e737dc87e39aded2f7cc788ff118d.
---
 doc/whats-new.rst         | 2 --
 xarray/core/variable.py   | 3 +++
 xarray/tests/test_dask.py | 7 -------
 3 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 60a53512622..d8b1fc2fba9 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -84,13 +84,11 @@ Bug fixes
 - Fix `KeyError` when doing linear interpolation to an nd `DataArray`
   that contains NaNs (:pull:`4233`).
   By `Jens Svensmark <https://github.com/jenssss>`_
-- Fix ``dask.optimize`` on ``DataArray`` producing an invalid Dask task graph (:issue:`3698`)
 - Fix incorrect legend labels for :py:meth:`Dataset.plot.scatter` (:issue:`4126`).
   By `Peter Hausamann <https://github.com/phausamann>`_.
 - Fix indexing with datetime64 scalars with pandas 1.1 (:issue:`4283`).
   By `Stephan Hoyer <https://github.com/shoyer>`_ and
   `Justus Magin <https://github.com/keewis>`_.
-  
 
 Documentation
 ~~~~~~~~~~~~~
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index c55e61cb816..6de00ee882a 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -501,6 +501,9 @@ def __dask_postpersist__(self):
 
     @staticmethod
     def _dask_finalize(results, array_func, array_args, dims, attrs, encoding):
+        if isinstance(results, dict):  # persist case
+            name = array_args[0]
+            results = {k: v for k, v in results.items() if k[0] == name}
         data = array_func(results, *array_args)
         return Variable(dims, data, attrs=attrs, encoding=encoding)
 
diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py
index 489bf09fa3c..46685a29a47 100644
--- a/xarray/tests/test_dask.py
+++ b/xarray/tests/test_dask.py
@@ -1607,10 +1607,3 @@ def test_more_transforms_pass_lazy_array_equiv(map_da, map_ds):
         assert_equal(map_da._from_temp_dataset(map_da._to_temp_dataset()), map_da)
         assert_equal(map_da.astype(map_da.dtype), map_da)
         assert_equal(map_da.transpose("y", "x", transpose_coords=False).cxy, map_da.cxy)
-
-
-def test_optimize():
-    a = dask.array.ones((10, 5), chunks=(1, 3))
-    arr = xr.DataArray(a).chunk(5)
-    (arr2,) = dask.optimize(arr)
-    arr2.compute()