Repository: dask/dask

Issue:
Subset dask.dataframe columns with an Index
I think this should be doable:


```python
In [1]: import pandas as pd; import dask.dataframe as dd

In [2]: df = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]})

In [3]: df[pd.Index(['C', 'A', 'B'])]
Out[3]:
   C  A  B
0  5  1  3
1  6  2  4

In [4]: ddf = dd.from_pandas(df, 2)

In [5]: ddf[pd.Index(['C', 'A', 'B'])]
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-5-3ae043d06f34> in <module>()
----> 1 ddf[pd.Index(['C', 'A', 'B'])]

~/Envs/dask-dev/lib/python3.6/site-packages/dask/dask/dataframe/core.py in __getitem__(self, key)
   2157             return new_dd_object(merge(self.dask, key.dask, dsk), name,
   2158                                  self, self.divisions)
-> 2159         raise NotImplementedError(key)
   2160
   2161     def __setitem__(self, key, value):

NotImplementedError: Index(['C', 'A', 'B'], dtype='object')
```

We check for `isinstance(key, list)`. That could probably be relaxed a bit.

Discussion:

Patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 02e13575f..804732df2 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -4,7 +4,7 @@ from ..utils import ignoring
 from .core import (Array, block, concatenate, stack, from_array, store,
                    map_blocks, atop, to_hdf5, to_npy_stack, from_npy_stack,
                    from_delayed, asarray, asanyarray,
-                   broadcast_arrays, broadcast_to)
+                   broadcast_arrays, broadcast_to, from_zarr, to_zarr)
 from .routines import (take, choose, argwhere, where, coarsen, insert,
                        ravel, roll, unique, squeeze, ptp, diff, ediff1d,
                        gradient, bincount, digitize, histogram, cov, array,
diff --git a/dask/array/core.py b/dask/array/core.py
index f53f63627..3ecc6d9e6 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -47,6 +47,7 @@ from .. import threaded, core
 from .. import sharedict
 from ..sharedict import ShareDict
 from .numpy_compat import _Recurser
+from ..bytes.core import get_mapper, get_fs_token_paths
 
 
 config.update_defaults({'array': {
@@ -1887,6 +1888,15 @@ class Array(DaskMethodsMixin):
         from .routines import nonzero
         return nonzero(self)
 
+    def to_zarr(self, *args, **kwargs):
+        """Save array to the zarr storage format
+
+        See https://zarr.readthedocs.io for details about the format.
+
+        See function ``to_zarr()`` for parameters.
+        """
+        return to_zarr(self, *args, **kwargs)
+
 
 def ensure_int(f):
     i = int(f)
@@ -2189,6 +2199,128 @@ def from_array(x, chunks, name=None, lock=False, asarray=True, fancy=True,
     return Array(dsk, name, chunks, dtype=x.dtype)
 
 
+def from_zarr(url, component=None, storage_options=None, chunks=None, **kwargs):
+    """Load array from the zarr storage format
+
+    See https://zarr.readthedocs.io for details about the format.
+
+    Parameters
+    ----------
+    url: str or MutableMapping
+        Location of the data. A URL can include a protocol specifier like s3://
+        for remote data. Can also be any MutableMapping instance, which should
+        be serializable if used in multiple processes.
+    component: str or None
+        If the location is a zarr group rather than an array, this is the
+        subcomponent that should be loaded, something like ``'foo/bar'``.
+    storage_options: dict
+        Any additional parameters for the storage backend (ignored for local
+        paths)
+    chunks: tuple of ints or tuples of ints
+        Passed to ``da.from_array``, allows setting the chunks on
+        initialisation, if the chunking scheme in the on-disc dataset is not
+        optimal for the calculations to follow.
+    kwargs: passed to ``zarr.Array``.
+    """
+    import zarr
+    storage_options = storage_options or {}
+    if isinstance(url, str):
+        fs, fs_token, path = get_fs_token_paths(
+            url, 'rb', storage_options=storage_options)
+        assert len(path) == 1
+        mapper = get_mapper(fs, path[0])
+    else:
+        mapper = url
+    z = zarr.Array(mapper, read_only=True, path=component, **kwargs)
+    chunks = chunks if chunks is not None else z.chunks
+    return from_array(z, chunks, name='zarr-%s' % url)
+
+
+def to_zarr(arr, url, component=None, storage_options=None,
+            overwrite=False, compute=True, return_stored=False, **kwargs):
+    """Save array to the zarr storage format
+
+    See https://zarr.readthedocs.io for details about the format.
+
+    Parameters
+    ----------
+    arr: dask.array
+        Data to store
+    url: str or MutableMapping
+        Location of the data. A URL can include a protocol specifier like s3://
+        for remote data. Can also be any MutableMapping instance, which should
+        be serializable if used in multiple processes.
+    component: str or None
+        If the location is a zarr group rather than an array, this is the
+        subcomponent that should be created/over-written.
+    storage_options: dict
+        Any additional parameters for the storage backend (ignored for local
+        paths)
+    overwrite: bool
+        If given array already exists, overwrite=False will cause an error,
+        where overwrite=True will replace the existing data.
+    compute, return_stored: see ``store()``
+    kwargs: passed to the ``zarr.create()`` function, e.g., compression options
+    """
+    import zarr
+    if not _check_regular_chunks(arr.chunks):
+        raise ValueError('Attempt to save array to zarr with irregular '
+                         'chunking, please call `arr.rechunk(...)` first.')
+    storage_options = storage_options or {}
+    if isinstance(url, str):
+        fs, fs_token, path = get_fs_token_paths(
+            url, 'rb', storage_options=storage_options)
+        assert len(path) == 1
+        mapper = get_mapper(fs, path[0])
+    else:
+        # assume the object passed is already a mapper
+        mapper = url
+    chunks = [c[0] for c in arr.chunks]
+    z = zarr.create(shape=arr.shape, chunks=chunks, dtype=arr.dtype,
+                    store=mapper, path=component, overwrite=overwrite, **kwargs)
+    return store(arr, z, compute=compute, return_stored=return_stored)
+
+
+def _check_regular_chunks(chunkset):
+    """Check if the chunks are regular
+
+    "Regular" in this context means that along every axis, the chunks all
+    have the same size, except the last one, which may be smaller
+
+    Parameters
+    ----------
+    chunkset: tuple of tuples of ints
+        From the ``.chunks`` attribute of an ``Array``
+
+    Returns
+    -------
+    True if chunkset passes, else False
+
+    Examples
+    --------
+    >>> import dask.array as da
+    >>> arr = da.zeros(10, chunks=(5, ))
+    >>> _check_regular_chunks(arr.chunks)
+    True
+
+    >>> arr = da.zeros(10, chunks=((3, 3, 3, 1), ))
+    >>> _check_regular_chunks(arr.chunks)
+    True
+
+    >>> arr = da.zeros(10, chunks=((3, 1, 3, 3), ))
+    >>> _check_regular_chunks(arr.chunks)
+    False
+    """
+    for chunks in chunkset:
+        if len(chunks) == 1:
+            continue
+        if len(set(chunks[:-1])) > 1:
+            return False
+        if chunks[-1] > chunks[0]:
+            return False
+    return True
+
+
 def from_delayed(value, shape, dtype, name=None):
     """ Create a dask array from a dask delayed value
 
diff --git a/dask/array/fft.py b/dask/array/fft.py
index abcd2ffaa..b1e624d62 100644
--- a/dask/array/fft.py
+++ b/dask/array/fft.py
@@ -168,8 +168,8 @@ def fft_wrap(fft_func, kind=None, dtype=None):
 
         _dtype = dtype
         if _dtype is None:
-            _dtype = fft_func(np.ones(len(axes) * (8,),
-                                      dtype=a.dtype)).dtype
+            sample = np.ones(a.ndim * (8,), dtype=a.dtype)
+            _dtype = fft_func(sample).dtype
 
         for each_axis in axes:
             if len(a.chunks[each_axis]) != 1:
@@ -210,20 +210,20 @@ def fft_wrap(fft_func, kind=None, dtype=None):
     return func
 
 
-fft = fft_wrap(np.fft.fft, dtype=np.complex_)
-fft2 = fft_wrap(np.fft.fft2, dtype=np.complex_)
-fftn = fft_wrap(np.fft.fftn, dtype=np.complex_)
-ifft = fft_wrap(np.fft.ifft, dtype=np.complex_)
-ifft2 = fft_wrap(np.fft.ifft2, dtype=np.complex_)
-ifftn = fft_wrap(np.fft.ifftn, dtype=np.complex_)
-rfft = fft_wrap(np.fft.rfft, dtype=np.complex_)
-rfft2 = fft_wrap(np.fft.rfft2, dtype=np.complex_)
-rfftn = fft_wrap(np.fft.rfftn, dtype=np.complex_)
-irfft = fft_wrap(np.fft.irfft, dtype=np.float_)
-irfft2 = fft_wrap(np.fft.irfft2, dtype=np.float_)
-irfftn = fft_wrap(np.fft.irfftn, dtype=np.float_)
-hfft = fft_wrap(np.fft.hfft, dtype=np.float_)
-ihfft = fft_wrap(np.fft.ihfft, dtype=np.complex_)
+fft = fft_wrap(np.fft.fft)
+fft2 = fft_wrap(np.fft.fft2)
+fftn = fft_wrap(np.fft.fftn)
+ifft = fft_wrap(np.fft.ifft)
+ifft2 = fft_wrap(np.fft.ifft2)
+ifftn = fft_wrap(np.fft.ifftn)
+rfft = fft_wrap(np.fft.rfft)
+rfft2 = fft_wrap(np.fft.rfft2)
+rfftn = fft_wrap(np.fft.rfftn)
+irfft = fft_wrap(np.fft.irfft)
+irfft2 = fft_wrap(np.fft.irfft2)
+irfftn = fft_wrap(np.fft.irfftn)
+hfft = fft_wrap(np.fft.hfft)
+ihfft = fft_wrap(np.fft.ihfft)
 
 
 def _fftfreq_block(i, n, d):
diff --git a/dask/bytes/core.py b/dask/bytes/core.py
index c2f7499d9..73b0eae33 100644
--- a/dask/bytes/core.py
+++ b/dask/bytes/core.py
@@ -313,10 +313,30 @@ def get_fs_token_paths(urlpath, mode='rb', num=1, name_function=None,
 
     else:
         raise TypeError('url type not understood: %s' % urlpath)
+    fs.protocol = protocol
 
     return fs, fs_token, paths
 
 
+def get_mapper(fs, path):
+    # This is not the right way to do this.
+    # At the very least, we should have the correct failed import messages
+    if fs.protocol == 'file':
+        from zarr.storage import DirectoryStore
+        return DirectoryStore(path)
+    elif fs.protocol == 's3':
+        from s3fs.mapping import S3Map
+        return S3Map(path, fs)
+    elif fs.protocol in ['gcs', 'gs']:
+        from gcsfs.mapping import GCSMap
+        return GCSMap(path, fs)
+    elif fs.protocol == 'hdfs':
+        from hdfs3.mapping import HDFSMap
+        return HDFSMap(fs, path)
+    else:
+        raise ValueError('No mapper for protocol "%s"' % fs.protocol)
+
+
 def open_text_files(urlpath, compression=None, mode='rt', encoding='utf8',
                     errors='strict', **kwargs):
     """ Given path return dask.delayed file-like objects in text mode
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 54f256c24..86bf3277f 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2359,7 +2359,7 @@ class DataFrame(_Frame):
         elif isinstance(key, slice):
             return self.loc[key]
 
-        if isinstance(key, list):
+        if isinstance(key, (pd.Series, np.ndarray, pd.Index, list)):
             # error is raised from pandas
             meta = self._meta[_extract_meta(key)]
 
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
index 773061acd..6c004289d 100644
--- a/docs/source/array-api.rst
+++ b/docs/source/array-api.rst
@@ -321,8 +321,10 @@ Create and Store Arrays
    from_array
    from_delayed
    from_npy_stack
+   from_zarr
    store
    to_hdf5
+   to_zarr
    to_npy_stack
 
 Generalized Ufuncs
diff --git a/docs/source/array-creation.rst b/docs/source/array-creation.rst
index 51eb859da..fb3692683 100644
--- a/docs/source/array-creation.rst
+++ b/docs/source/array-creation.rst
@@ -11,6 +11,7 @@ supports Numpy-style slicing.
    from_array
    from_delayed
    from_npy_stack
+   from_zarr
    stack
    concatenate
 
@@ -177,6 +178,7 @@ Store Dask Arrays
    store
    to_hdf5
    to_npy_stack
+   to_zarr
    compute
 
 In Memory
@@ -243,6 +245,36 @@ Store several arrays in one computation with the function
    >>> da.to_hdf5('myfile.hdf5', {'/x': x, '/y': y})  # doctest: +SKIP
 
 
+Zarr
+----
+
+The `zarr`_ format is a chunk-wise binary array storage file format, with a good selection
+of encoding and compression options. Due to each chunk being stored in a separate file, it
+is ideal for parallel access in both reading and writing (for the latter, if the dask array
+chunks are alligned with the target). Furthermore, storage in
+:doc:`remote data services<remote_data_services>`_ such as
+S3 and GCS is supported.
+
+.. _zarr: https://zarr.readthedocs.io
+
+For example, to save data to a local zarr dataset:
+
+.. code-block:: Python
+
+   >>> arr.to_zarr('output.zarr')
+
+or to save to a particular bucket on S3:
+
+.. code-block:: Python
+
+   >>> arr.to_zarr('s3://mybucket/output.zarr', storage_option={'key': 'mykey',
+                   'secret': 'mysecret'})
+
+To retrieve those data, you would do ``da.read_zarr`` with exactly the same arguments. The
+chunking of the resultant dask.Array is defined by how the files were saved, unless
+otherwise specified.
+
+
 Plugins
 =======
 
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 6bc853540..cfb948819 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -8,6 +8,9 @@ Changelog
 Array
 +++++
 
+- Fix ``rechunk`` with chunksize of -1 in a dict (:pr:`3469`) `Stephan Hoyer`_
+- ``einsum`` now accepts the ``split_every`` parameter (:pr:`3396`) `Guido Imperiale`_
+- Add to/read_zarr for Zarr-format datasets and arrays (:pr:`3460`) `Martin Durant`_
 - Experimental addition of generalized ufunc support, ``apply_gufunc``, ``gufunc``, and ``as_gufunc`` (:pr:`#3109`) `Markus Gonser`_
 
 Dataframe
@@ -15,6 +18,7 @@ Dataframe
 
 - Add to/read_json (:pr:`3494`) `Martin Durant`_
 - Adds ``index`` to unsupported arguments for ``DataFrame.rename`` method (:pr:`3522`) `James Bourbeau`_
+- Adds support to subset Dask DataFrame columns using ``numpy.ndarray``, ``pandas.Series``, and ``pandas.Index`` objects (:pr:`3536`) `James Bourbeau`_
 
 Bag
 +++


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index b70c506d3..973aada17 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -3328,3 +3328,81 @@ def test_normalize_chunks_nan():
     with pytest.raises(ValueError) as info:
         normalize_chunks(((np.nan, np.nan), 'auto'), (10, 10), limit=10, dtype=np.uint8)
     assert "auto" in str(info.value)
+
+
+def test_zarr_roundtrip():
+    pytest.importorskip('zarr')
+    with tmpdir() as d:
+        a = da.zeros((3, 3), chunks=(1, 1))
+        a.to_zarr(d)
+        a2 = da.from_zarr(d)
+        assert_eq(a, a2)
+        assert a2.chunks == a.chunks
+
+
+def test_read_zarr_chunks():
+    pytest.importorskip('zarr')
+    a = da.zeros((9, ), chunks=(3, ))
+    with tmpdir() as d:
+        a.to_zarr(d)
+        arr = da.from_zarr(d, chunks=(5, ))
+        assert arr.chunks == ((5, 4), )
+
+
+def test_zarr_pass_mapper():
+    pytest.importorskip('zarr')
+    import zarr.storage
+    with tmpdir() as d:
+        mapper = zarr.storage.DirectoryStore(d)
+        a = da.zeros((3, 3), chunks=(1, 1))
+        a.to_zarr(mapper)
+        a2 = da.from_zarr(mapper)
+        assert_eq(a, a2)
+        assert a2.chunks == a.chunks
+
+
+def test_zarr_group():
+    zarr = pytest.importorskip('zarr')
+    with tmpdir() as d:
+        a = da.zeros((3, 3), chunks=(1, 1))
+        a.to_zarr(d, component='test')
+        with pytest.raises((OSError, ValueError)):
+            a.to_zarr(d, component='test', overwrite=False)
+        a.to_zarr(d, component='test', overwrite=True)
+
+        # second time is fine, group exists
+        a.to_zarr(d, component='test2', overwrite=False)
+        a.to_zarr(d, component='nested/test', overwrite=False)
+        group = zarr.open_group(d, mode='r')
+        assert list(group) == ['nested', 'test', 'test2']
+        assert 'test' in group['nested']
+
+        a2 = da.from_zarr(d, component='test')
+        assert_eq(a, a2)
+        assert a2.chunks == a.chunks
+
+
+@pytest.mark.parametrize('data', [[( ), True],
+                                  [((1, ),), True],
+                                  [((1, 1, 1),), True],
+                                  [((1, ), (1, )), True],
+                                  [((2, 2, 1), ), True],
+                                  [((2, 2, 3), ), False],
+                                  [((1, 1, 1), (2, 2, 3)), False],
+                                  [((1, 2, 1), ), False]
+                                  ])
+def test_regular_chunks(data):
+    chunkset, expected = data
+    assert da.core._check_regular_chunks(chunkset) == expected
+
+
+def test_zarr_nocompute():
+    pytest.importorskip('zarr')
+    with tmpdir() as d:
+        a = da.zeros((3, 3), chunks=(1, 1))
+        out = a.to_zarr(d, compute=False)
+        assert isinstance(out, Delayed)
+        dask.compute(out)
+        a2 = da.from_zarr(d)
+        assert_eq(a, a2)
+        assert a2.chunks == a.chunks
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 75caaa623..b155de6c9 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -2621,6 +2621,15 @@ def test_getitem_string_subclass():
     assert_eq(df[column_1], ddf[column_1])
 
 
+@pytest.mark.parametrize('col_type', [list, np.array, pd.Series, pd.Index])
+def test_getitem_column_types(col_type):
+    df = pd.DataFrame({'A': [1, 2], 'B': [3, 4], 'C': [5, 6]})
+    ddf = dd.from_pandas(df, 2)
+    cols = col_type(['C', 'A', 'B'])
+
+    assert_eq(df[cols], ddf[cols])
+
+
 def test_diff():
     df = pd.DataFrame(np.random.randn(100, 5), columns=list('abcde'))
     ddf = dd.from_pandas(df, 5)