Repository: dask/dask

Issue:
poor optimization of slicing operations on netCDF-backed xarray datasets
We have encountered a performance issue in xarray (pydata/xarray#1396) which appears to be a dask optimization issue at its core. In [correspondence on the xarray mailing list](https://groups.google.com/forum/#!topic/xarray/UbiQu7aZvkg), @mrocklin recommended I open an issue here.

The scenario is the following: I have a collection of large netCDF files on a disk. Each file can be though of as representing geospatial data: the first dimension is time, and the higher dimensions are spatial dimensions (e.g. x and y). I want to open the data as an xarray dataset, concatenating along the time dimension, select a single point in space, and lazily load the data for this point into memory, creating a timeseries. (This is an extremely common workflow in climate science and related fields.) xarray handles this by opening the files, wrapping the data in dask arrays, and then concatenating those arrays together, under the hood calling `dask.array.concatenate`. 

The problem is that the entire array from each file is being read eagerly before the slicing operation. This incurs a huge performance and memory penalty for what should naively be a very cheap operation. For very large datasets, it leads to memory overruns and total breakdown of the workflow.

Below is a reproducible example which requires xarray. A more detailed version of this example can be found in this [notebook](https://gist.github.com/rabernat/7175328ee04a3167fa4dede1821964c6)
```python
import numpy as np
import xarray as xr

### Create some test data
# This test data is typical of many climate datasets. Each file has 12 entires in the time
# dimension, representing e.g. each month of the year. There are 10 separate files,
# representing e.g. consecutive years. The code below writes about 9GB of data to disk.
# Change data_dir to a place where you can put 9GB of data

data_dir = '/data/scratch/rpa/xarray_test_data/'
nfiles = 10
nt = 12

for n in range(nfiles):
    data = np.random.rand(nt,1000,10000)
    time = (n*nt) + np.arange(nt)
    da = xr.DataArray(data, dims=['time', 'y', 'x'],
                      coords={'time': time})
    da.to_dataset(name='data').to_netcdf(data_dir + 'test_data.%03d.nc' % n)

### open the data via xarray
all_files = glob(data_dir + '*.nc')
ds = xr.open_mfdataset(all_files)

### Time loading the whole dataset into memory
# As a point of reference, let's see how long it takes to read all the data from disk and load it
# into memory. (We make a deep copy of the dataset first in order to preserve the original
# chunked dataset for later use.)

ds_copy = ds.copy(deep=True)
with ProgressBar():
    # on a high-performance workstation, this takes 15-20s
    %time _ = ds_copy.load()

### Time loading a single point
# This is the crux issue. We now want to just extract a timeseries of single point in space.
# If the loading is lazy, it should go much faster and use much less memory compared to
# the example above.

y, x = 200, 300
with ProgressBar():
    %time ts = ds.data[:, y, x].load()
```

The extraction of a single point takes about half the time of loading the whole dataset! Furthermore, by monitoring memory usage via top, I observe that resident memory increases strongly as the command executes. If I just loop through `all_files` manually and extract the point, it goes 500-1000 times faster and uses a tiny fraction of the memory. I explore this, and many other permutations (including bypassing xarray and doing `dask.array.concatenate` directly on the underling dask arrays of each dataset) in my [notebook](https://gist.github.com/rabernat/7175328ee04a3167fa4dede1821964c6). I believe this is a dask issue, not an xarray issue.

Here is the partial dot graph for `ds.data[:, y, x].dask`
![result_dot_graph](https://cloud.githubusercontent.com/assets/1197350/26169570/2c6ed1c2-3b0c-11e7-938e-a47a1833befb.png)

Here is the partial dot graph for `result._optimize(result.dask, result._keys())` where `result = ds.data[:, y, x].dask`
![result_dot_graph_optimized](https://cloud.githubusercontent.com/assets/1197350/26169591/4f5a8424-3b0c-11e7-9ca9-a09eb424d6c2.png)

The root cause of the issue, in the words of @mrocklin, is that "it looks like there is a getarray followed by a getitem.  We'll have to dive into why these two weren't fused into a single gatarray/getitem call."

```python
>>> print(xarray.__version__)
0.9.5-9-gd5c7e06
>>> print(dask.__version__)
0.14.3+19.gbcd0426
```


Discussion:
Looking into this now. The cause of the problem is that there's an alias in the graph (the `(concatenate, ...)` key points to the `(long_path_name, ...)` key), which is not something the optimizations are setup to deal with. Will fix.
Is this "alias" coming from xarray? Is it something we need to fix on the xarray side?
also worth keeping in mind that the dask prefixes of xarray data are about to change:
https://github.com/pydata/xarray/pull/1411

the node that is now called `/data/scratch/rpa/xarray_test_data/test_data.004.nc:/data-#01` will just be renamed `open_dataset-<token>`.

I don't know if this will impact the dask optimizations.
This has nothing to do with xarray, it's an issue with the internal optimizations of dask.

Resolving patch:
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -48,31 +48,37 @@ def hold_keys(dsk, dependencies):
     We don't want to fuse data present in the graph because it is easier to
     serialize as a raw value.
 
-    We don't want to fuse getitem/getarrays because we want to move around only
-    small pieces of data, rather than the underlying arrays.
+    We don't want to fuse chains after getitem/getarrays because we want to
+    move around only small pieces of data, rather than the underlying arrays.
     """
     dependents = reverse_dict(dependencies)
     data = {k for k, v in dsk.items() if type(v) not in (tuple, str)}
-    getters = (getitem, getarray, getarray_inline)
+    getters = (getarray, getarray_nofancy, getarray_inline, getitem)
 
     hold_keys = list(data)
     for dat in data:
         deps = dependents[dat]
         for dep in deps:
             task = dsk[dep]
-            try:
-                if task[0] in getters:
+            # If the task is a get* function, we walk up the chain, and stop
+            # when there's either more than one dependent, or the dependent is
+            # no longer a get* function or an alias. We then add the final
+            # key to the list of keys not to fuse.
+            if type(task) is tuple and task and task[0] in getters:
+                try:
                     while len(dependents[dep]) == 1:
                         new_dep = next(iter(dependents[dep]))
                         new_task = dsk[new_dep]
-                        if new_task[0] in (getitem, getarray):
+                        # If the task is a get* or an alias, continue up the
+                        # linear chain
+                        if new_task[0] in getters or new_task in dsk:
                             dep = new_dep
                             task = new_task
                         else:
                             break
-                    hold_keys.append(dep)
-            except (IndexError, TypeError):
-                pass
+                except (IndexError, TypeError):
+                    pass
+                hold_keys.append(dep)
     return hold_keys
 
 
@@ -91,9 +97,16 @@ def optimize_slices(dsk):
     for k, v in dsk.items():
         if type(v) is tuple and v[0] in getters and len(v) == 3:
             f, a, a_index = v
+            a_lock = None
             getter = f
-            while type(a) is tuple and a[0] in getters and len(a) == 3:
-                f2, b, b_index = a
+            while type(a) is tuple and a[0] in getters and len(a) in (3, 4):
+                if len(a) == 3:
+                    f2, b, b_index = a
+                    b_lock = None
+                else:
+                    f2, b, b_index, b_lock = a
+                if a_lock is not None and a_lock is not b_lock:
+                    break
                 if (type(a_index) is tuple) != (type(b_index) is tuple):
                     break
                 if type(a_index) is tuple:
@@ -118,22 +131,18 @@ def optimize_slices(dsk):
                         getter = getarray
                 except NotImplementedError:
                     break
-                a, a_index = b, c_index
-            if getter is not getitem:
-                dsk[k] = (getter, a, a_index)
-            elif (type(a_index) is slice and
-                    not a_index.start and
-                    a_index.stop is None and
-                    a_index.step is None):
-                dsk[k] = a
-            elif type(a_index) is tuple and all(type(s) is slice and
-                                                not s.start and
-                                                s.stop is None and
-                                                s.step is None
-                                                for s in a_index):
+                a, a_index, a_lock = b, c_index, b_lock
+            if (getter is getitem and
+                ((type(a_index) is slice and not a_index.start and
+                  a_index.stop is None and a_index.step is None) or
+                 (type(a_index) is tuple and
+                  all(type(s) is slice and not s.start and s.stop is None and
+                      s.step is None for s in a_index)))):
                 dsk[k] = a
+            elif a_lock is None:
+                dsk[k] = (getter, a, a_index)
             else:
-                dsk[k] = (getitem, a, a_index)
+                dsk[k] = (getter, a, a_index, a_lock)
     return dsk
 
 


Test patch:
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -5,12 +5,12 @@
 import dask
 import dask.array as da
 from dask.optimize import fuse
+from dask.utils import SerializableLock
+from dask.array.core import getarray, getarray_nofancy
 from dask.array.optimization import (getitem, optimize, optimize_slices,
                                      fuse_slice)
 from dask.array.utils import assert_eq
 
-from dask.array.core import getarray, getarray_nofancy
-
 
 def test_fuse_getitem():
     pairs = [((getarray, (getarray, 'x', slice(1000, 2000)), slice(15, 20)),
@@ -60,6 +60,29 @@ def test_fuse_getitem():
         assert result == {'y': expected}
 
 
+def test_fuse_getitem_lock():
+    lock1 = SerializableLock()
+    lock2 = SerializableLock()
+
+    pairs = [((getarray, (getarray, 'x', slice(1000, 2000), lock1), slice(15, 20)),
+              (getarray, 'x', slice(1015, 1020), lock1)),
+
+             ((getitem, (getarray, 'x', (slice(1000, 2000), slice(100, 200)), lock1),
+                        (slice(15, 20), slice(50, 60))),
+              (getarray, 'x', (slice(1015, 1020), slice(150, 160)), lock1)),
+
+             ((getitem, (getarray_nofancy, 'x', (slice(1000, 2000), slice(100, 200)), lock1),
+                        (slice(15, 20), slice(50, 60))),
+              (getarray_nofancy, 'x', (slice(1015, 1020), slice(150, 160)), lock1)),
+
+             ((getarray, (getarray, 'x', slice(1000, 2000), lock1), slice(15, 20), lock2),
+              (getarray, (getarray, 'x', slice(1000, 2000), lock1), slice(15, 20), lock2))]
+
+    for inp, expected in pairs:
+        result = optimize_slices({'y': inp})
+        assert result == {'y': expected}
+
+
 def test_optimize_with_getitem_fusion():
     dsk = {'a': 'some-array',
            'b': (getarray, 'a', (slice(10, 20), slice(100, 200))),
@@ -166,6 +189,18 @@ def test_minimize_data_transfer():
         assert dsk[dep][1] == big_key
 
 
+def test_fuse_slices_with_alias():
+    dsk = {'x': np.arange(16).reshape((4, 4)),
+           ('dx', 0, 0): (getarray, 'x', (slice(0, 4), slice(0, 4))),
+           ('alias', 0, 0): ('dx', 0, 0),
+           ('dx2', 0): (getitem, ('alias', 0, 0), (slice(None), 0))}
+    keys = [('dx2', 0)]
+    dsk2 = optimize(dsk, keys)
+    assert len(dsk2) == 3
+    fused_key = set(dsk2).difference(['x', ('dx2', 0)]).pop()
+    assert dsk2[fused_key] == (getarray, 'x', (slice(0, 4), 0))
+
+
 def test_dont_fuse_fancy_indexing_in_getarray_nofancy():
     dsk = {'a': (getitem, (getarray_nofancy, 'x', (slice(10, 20, None), slice(100, 200, None))),
                  ([1, 3], slice(50, 60, None)))}