Repository: dask/dask

Issue:
assign(lambda df: df.str.split(x, expand=True)[1]) gives KeyError
pandas `0.20.3`, dask `0.15.1+33.gaafeedc`

In pandas, I can extract the domain name of many URLs like so:

```python
import pandas as pd
df = pd.DataFrame(["http://example.com/"], columns=["url"])
df = df.assign(domain=df.url.str.split("/", expand=True)[2])
print(df)
```

```
                   url       domain
0  http://example.com/  example.com
```

But in dask this fails.

My mental model of dask is that pandas should operate as it always would inside the closure passed to `assign()`, so is this a bug?

```python
from dask import dataframe as dd
import pandas as pd

df = pd.DataFrame(["http://example.com/"], columns=["url"])

(
    dd
    .from_pandas(df, npartitions=1)
    .assign(domain=lambda df: df.url.str.split(" ", expand=True)[2])
).compute()
```

```
~/.local/lib/python3.5/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   2442                 return self._engine.get_loc(key)
   2443             except KeyError:
-> 2444                 return self._engine.get_loc(self._maybe_cast_indexer(key))
   2445 
   2446         indexer = self.get_indexer([key], method=method, tolerance=tolerance)

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc (pandas/_libs/index.c:5280)()

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc (pandas/_libs/index.c:5126)()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item (pandas/_libs/hashtable.c:20523)()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item (pandas/_libs/hashtable.c:20477)()

KeyError: 2
```

# Full traceback collapsed below

<details>
<summary>Click to expand python traceback `KeyError`</summary>

```ipytb
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/.local/lib/python3.5/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   2441             try:
-> 2442                 return self._engine.get_loc(key)
   2443             except KeyError:

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc (pandas/_libs/index.c:5280)()

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc (pandas/_libs/index.c:5126)()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item (pandas/_libs/hashtable.c:20523)()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item (pandas/_libs/hashtable.c:20477)()

KeyError: 2

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
<ipython-input-7-ab744d790f75> in <module>()
      3     dd
      4     .from_pandas(df, npartitions=1)
----> 5     .assign(domain=lambda df: df.url.str.split(" ", expand=True)[0])
      6 ).compute()

~/.local/src/github.com/dask/dask/dask/dataframe/core.py in assign(self, **kwargs)
   2317 
   2318         # Figure out columns of the output
-> 2319         df2 = self._meta.assign(**_extract_meta(kwargs))
   2320         return elemwise(methods.assign, self, *pairs, meta=df2)
   2321 

~/.local/lib/python3.5/site-packages/pandas/core/frame.py in assign(self, **kwargs)
   2495         results = {}
   2496         for k, v in kwargs.items():
-> 2497             results[k] = com._apply_if_callable(v, data)
   2498 
   2499         # ... and then assign

~/.local/lib/python3.5/site-packages/pandas/core/common.py in _apply_if_callable(maybe_callable, obj, **kwargs)
    439     """
    440     if callable(maybe_callable):
--> 441         return maybe_callable(obj, **kwargs)
    442     return maybe_callable
    443 

<ipython-input-7-ab744d790f75> in <lambda>(df)
      3     dd
      4     .from_pandas(df, npartitions=1)
----> 5     .assign(domain=lambda df: df.url.str.split(" ", expand=True)[0])
      6 ).compute()

~/.local/lib/python3.5/site-packages/pandas/core/frame.py in __getitem__(self, key)
   1962             return self._getitem_multilevel(key)
   1963         else:
-> 1964             return self._getitem_column(key)
   1965 
   1966     def _getitem_column(self, key):

~/.local/lib/python3.5/site-packages/pandas/core/frame.py in _getitem_column(self, key)
   1969         # get column
   1970         if self.columns.is_unique:
-> 1971             return self._get_item_cache(key)
   1972 
   1973         # duplicate columns & possible reduce dimensionality

~/.local/lib/python3.5/site-packages/pandas/core/generic.py in _get_item_cache(self, item)
   1643         res = cache.get(item)
   1644         if res is None:
-> 1645             values = self._data.get(item)
   1646             res = self._box_item_values(item, values)
   1647             cache[item] = res

~/.local/lib/python3.5/site-packages/pandas/core/internals.py in get(self, item, fastpath)
   3588 
   3589             if not isnull(item):
-> 3590                 loc = self.items.get_loc(item)
   3591             else:
   3592                 indexer = np.arange(len(self.items))[isnull(self.items)]

~/.local/lib/python3.5/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   2442                 return self._engine.get_loc(key)
   2443             except KeyError:
-> 2444                 return self._engine.get_loc(self._maybe_cast_indexer(key))
   2445 
   2446         indexer = self.get_indexer([key], method=method, tolerance=tolerance)

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc (pandas/_libs/index.c:5280)()

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc (pandas/_libs/index.c:5126)()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item (pandas/_libs/hashtable.c:20523)()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item (pandas/_libs/hashtable.c:20477)()

KeyError: 2
```

</details>


Discussion:

Resolving patch:
diff --git a/dask/dataframe/accessor.py b/dask/dataframe/accessor.py
--- a/dask/dataframe/accessor.py
+++ b/dask/dataframe/accessor.py
@@ -4,6 +4,8 @@
 import pandas as pd
 from toolz import partial
 
+from ..utils import derived_from
+
 
 def maybe_wrap_pandas(obj, x):
     if isinstance(x, np.ndarray):
@@ -24,6 +26,8 @@ class Accessor(object):
     * _accessor
     * _accessor_name
     """
+    _not_implemented = set()
+
     def __init__(self, series):
         from .core import Series
         if not isinstance(series, Series):
@@ -60,12 +64,18 @@ def _function_map(self, attr, *args, **kwargs):
                                            self._accessor_name, attr, args,
                                            kwargs, meta=meta, token=token)
 
+    @property
+    def _delegates(self):
+        return set(dir(self._accessor)).difference(self._not_implemented)
+
     def __dir__(self):
-        return sorted(set(dir(type(self)) + list(self.__dict__) +
-                      dir(self._accessor)))
+        o = self._delegates
+        o.update(self.__dict__)
+        o.update(dir(type(self)))
+        return list(o)
 
     def __getattr__(self, key):
-        if key in dir(self._accessor):
+        if key in self._delegates:
             if isinstance(getattr(self._accessor, key), property):
                 return self._property_map(key)
             else:
@@ -96,7 +106,12 @@ class StringAccessor(Accessor):
     """
     _accessor = pd.Series.str
     _accessor_name = 'str'
+    _not_implemented = {'get_dummies'}
 
     def _validate(self, series):
         if not series.dtype == 'object':
             raise AttributeError("Can only use .str accessor with object dtype")
+
+    @derived_from(pd.core.strings.StringMethods)
+    def split(self, pat=None, n=-1):
+        return self._function_map('split', pat=pat, n=n)


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
@@ -1348,6 +1348,10 @@ def test_str_accessor():
     assert 'str' not in dir(a.y)
     assert not hasattr(a.y, 'str')
 
+    # not implemented methods don't show up
+    assert 'get_dummies' not in dir(a.x.str)
+    assert not hasattr(a.x.str, 'get_dummies')
+
     assert 'upper' in dir(a.x.str)
     assert_eq(a.x.str.upper(), df.x.str.upper())
     assert set(a.x.str.upper().dask) == set(a.x.str.upper().dask)