Repository: dask/dask

Issue:
dask.dataframe __getitem__ does not work with subclasses of str in python 3.6
I am upgrading my code from python 2.7 to 3.6 and found that that the __getitem__ of a dask dataframe does not work as before: I use a subclass of `str` for the column names and then I want to access the columns with `df[column_name]`. This works fine with pandas and had worked with dask also in python 2.7. Here is a short example for comparing pandas and dask:

```
class MyString(str):
    pass

my_s = MyString('column_1')

import pandas as pd
df = df = pd.DataFrame({'column_1': [1, 1]})
print(df[my_s])

from dask import dataframe as dd
ddf = dd.from_pandas(df, npartitions=1)
print(ddf[my_s])
```

The last line results in following error:
```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-15-731527b3e767> in <module>()
      2 ddf = dd.from_pandas(df, npartitions=1)
----> 3 print(ddf[my_s])

/io/venv_debian_8_http/lib/python3.6/site-packages/dask/dataframe/core.py in __getitem__(self, key)
   2310             return new_dd_object(merge(self.dask, key.dask, dsk), name,
   2311                                  self, self.divisions)
-> 2312         raise NotImplementedError(key)
   2313 
   2314     def __setitem__(self, key, value):

NotImplementedError: column_1
```

I would expect pandas and dask dataframe to behave in the same way here. Software Versions are:
`
dask==0.17.1
pandas==0.22.0
Python is 3.6.4
`

Discussion:
mrocklin: This seems fine to me.  

It's odd seeing six in the codebase, which hasn't been used historically and isn't an explicit dependency.  I notice though that it has been used once in dataframe/io/sql.py for the last year though.  Thoughts on if we should add it as an explicit dependency of `dask[dataframe]` @jakirkham or work around with our own compatibility.py file?
jakirkham: Have no qualms with `six` other than it's style is a bit dated. Personally would encourage using a more modern Python 2/3 compat library like [`python-future`]( http://python-future.org/ ) or [`eight`]( https://eight.readthedocs.io/en/latest/ ). Using our own compat module would be an option.
jrbourbeau: Good to know about `string_types` in `compatibility.py`, I hadn't run across that before! I've got no preference re: `six` and can switch to using `dask.compatibility.string_types` here. 
mrocklin: It looks like we have string_types already in compatibility.py.

@jrbourbeau I recommend using this so that we can just avoid the question altogether.

If you felt like fixing the use of six in dask/dataframe/io/sql.py as well that would be welcome :)
mrocklin: Ah, you beat me to the comment 
jrbourbeau: Getting some `OSError: [Errno 12] Cannot allocate memory` errors for one of the builds on Travis. Should be unrelated to the changes made in this PR. 
mrocklin: Hrm, that's interesting.  I agree that it's likely unrelated.  Restarting.
We'll see if it recurs.

On Wed, May 2, 2018 at 5:43 PM, James Bourbeau <notifications@github.com>
wrote:

> Getting some OSError: [Errno 12] Cannot allocate memory errors for one of
> the builds on Travis. Should be unrelated to the changes made in this PR.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/pull/3461#issuecomment-386130958>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszEByXPa5YMFT1tdpalltqUdgOFc1ks5tuiiRgaJpZM4TwDKn>
> .
>

mrocklin: This looks good to me.  Thanks @jrbourbeau !

Patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 0ddbf62d4..59bd653ee 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -22,7 +22,7 @@ from .. import core
 
 from ..utils import partial_by_order
 from .. import threaded
-from ..compatibility import apply, operator_div, bind_method
+from ..compatibility import apply, operator_div, bind_method, string_types
 from ..context import globalmethod
 from ..utils import (random_state_data, pseudorandom, derived_from, funcname,
                      memory_repr, put_lines, M, key_split, OperatorMethodMixin,
@@ -2338,7 +2338,7 @@ class DataFrame(_Frame):
 
     def __getitem__(self, key):
         name = 'getitem-%s' % tokenize(self, key)
-        if np.isscalar(key) or isinstance(key, tuple):
+        if np.isscalar(key) or isinstance(key, (tuple, string_types)):
 
             if isinstance(self._meta.index, (pd.DatetimeIndex, pd.PeriodIndex)):
                 if key not in self._meta.columns:
diff --git a/dask/dataframe/io/sql.py b/dask/dataframe/io/sql.py
index cf3d4e19d..592b66403 100644
--- a/dask/dataframe/io/sql.py
+++ b/dask/dataframe/io/sql.py
@@ -1,8 +1,8 @@
 import numpy as np
 import pandas as pd
-import six
 
 from ... import delayed
+from ...compatibility import string_types
 from .io import from_delayed, from_pandas
 
 
@@ -83,27 +83,27 @@ def read_sql_table(table, uri, index_col, divisions=None, npartitions=None,
         raise ValueError("Must specify index column to partition on")
     engine = sa.create_engine(uri)
     m = sa.MetaData()
-    if isinstance(table, six.string_types):
+    if isinstance(table, string_types):
         table = sa.Table(table, m, autoload=True, autoload_with=engine,
                          schema=schema)
 
-    index = (table.columns[index_col] if isinstance(index_col, six.string_types)
+    index = (table.columns[index_col] if isinstance(index_col, string_types)
              else index_col)
-    if not isinstance(index_col, six.string_types + (elements.Label,)):
+    if not isinstance(index_col, string_types + (elements.Label,)):
         raise ValueError('Use label when passing an SQLAlchemy instance'
                          ' as the index (%s)' % index)
     if divisions and npartitions:
         raise TypeError('Must supply either divisions or npartitions, not both')
 
-    columns = ([(table.columns[c] if isinstance(c, six.string_types) else c)
+    columns = ([(table.columns[c] if isinstance(c, string_types) else c)
                 for c in columns]
                if columns else list(table.columns))
     if index_col not in columns:
         columns.append(table.columns[index_col]
-                       if isinstance(index_col, six.string_types)
+                       if isinstance(index_col, string_types)
                        else index_col)
 
-    if isinstance(index_col, six.string_types):
+    if isinstance(index_col, string_types):
         kwargs['index_col'] = index_col
     else:
         # function names get pandas auto-named
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 351cd9369..5a1d99131 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -13,7 +13,7 @@ Array
 Dataframe
 +++++++++
 
--
+- Add support for indexing Dask DataFrames with string subclasses (:pr:`3461`) `James Bourbeau`_
 
 Bag
 +++


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 35e290f7c..ec1802e2a 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -2600,6 +2600,17 @@ def test_getitem_multilevel():
     assert_eq(pdf[[('A', '0'), ('B', '1')]], ddf[[('A', '0'), ('B', '1')]])
 
 
+def test_getitem_string_subclass():
+    df = pd.DataFrame({'column_1': list(range(10))})
+    ddf = dd.from_pandas(df, npartitions=3)
+
+    class string_subclass(str):
+        pass
+    column_1 = string_subclass('column_1')
+
+    assert_eq(df[column_1], ddf[column_1])
+
+
 def test_diff():
     df = pd.DataFrame(np.random.randn(100, 5), columns=list('abcde'))
     ddf = dd.from_pandas(df, 5)