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:
Makes sense to me

On Fri, Oct 27, 2017 at 9:34 AM, Tom Augspurger <notifications@github.com>
wrote:

> I think this should be doable:
>
> 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.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/2832>, or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszG1GfEAVzZLaKCrcqix6iER71WqWks5swdvggaJpZM4QJEPY>
> .
>


Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2359,7 +2359,7 @@ def __getitem__(self, key):
         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)]
 


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -2621,6 +2621,15 @@ class string_subclass(str):
     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)