Repository: dask/dask

Issue:
DataFrame.assign with lambda
Example:

```
import pandas as pd
import dask.dataframe as dd

In [18]: concrete_df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})

In [19]: concrete_df.assign(new_col = lambda x: x.a + x.b)
Out[19]: 
   a  b  new_col
0  1  4        5
1  2  5        7
2  3  6        9

In [20]: dd.from_pandas(concrete_df, npartitions=1).assign(new_col = lambda x: x.a + x.b)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-20-54c795f1117b> in <module>()
----> 1 dd.from_pandas(concrete_df, npartitions=1).assign(new_col = lambda x: x.a + x.b)

~\AppData\Local\Continuum\Anaconda3\lib\site-packages\dask\dataframe\core.py in assign(self, **kwargs)
   2295                     np.isscalar(v)):
   2296                 raise TypeError("Column assignment doesn't support type "
-> 2297                                 "{0}".format(type(v).__name__))
   2298         pairs = list(sum(kwargs.items(), ()))
   2299 

TypeError: Column assignment doesn't support type function
```

I suspect it wouldn't be too hard to implement this, but in the meantime could also give a more meaningful error message (`NotImplementedError`)

pandas `0.20.1`
dask `0.15.0`


Discussion:
The following does work - does it meet your requirements?

    d = dd.from_pandas(concrete_df, npartitions=1)
    d.assign(new_col=(d.a + d.b))
Sorry, my minimum example was a little too minimal, where I actually ran into this was a `groupby` chain where I can't refer back to original frame.

```python
In [67]: (concrete_df
    ...:  .groupby('a')
    ...:  [['b']].sum()
    ...:  .assign(j=lambda x: x.b + 1))
Out[67]: 
   b  j
a      
1  4  5
2  5  6
3  6  7

In [68]: ddf = dd.from_pandas(concrete_df, npartitions=1)


In [71]: (ddf
    ...:  .groupby('a')
    ...:  [['b']].sum()
    ...:  .assign(j=ddf.b + 1))
ValueError: Not all divisions are known, can't align partitions. 
Please use `set_index` to set the index.
```

I was able to workaround fine by computing the intermediate step - was opening this mostly as an error reporting issue as the original error message had me stumped for a minute.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2307,7 +2307,7 @@ def categorize(self, columns=None, index=None, split_every=None, **kwargs):
     def assign(self, **kwargs):
         for k, v in kwargs.items():
             if not (isinstance(v, (Series, Scalar, pd.Series)) or
-                    np.isscalar(v)):
+                    callable(v) or np.isscalar(v)):
                 raise TypeError("Column assignment doesn't support type "
                                 "{0}".format(type(v).__name__))
         pairs = list(sum(kwargs.items(), ()))


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
@@ -745,15 +745,18 @@ def test_assign():
     res = d.assign(c=1,
                    d='string',
                    e=d.a.sum(),
-                   f=d.a + d.b)
+                   f=d.a + d.b,
+                   g=lambda x: x.a + x.b)
     res_unknown = d_unknown.assign(c=1,
                                    d='string',
                                    e=d_unknown.a.sum(),
-                                   f=d_unknown.a + d_unknown.b)
+                                   f=d_unknown.a + d_unknown.b,
+                                   g=lambda x: x.a + x.b)
     sol = full.assign(c=1,
                       d='string',
                       e=full.a.sum(),
-                      f=full.a + full.b)
+                      f=full.a + full.b,
+                      g=lambda x: x.a + x.b)
     assert_eq(res, sol)
     assert_eq(res_unknown, sol)