Repository: dask/dask

Issue:
`.assign` leads to different results when using lambda vs not using lambda

`.assign` leads to different results when using lambda vs not using lambda. Please see the example below.

The difference seems to happen only in the divisions limits. I guess this is a bug? Any idea on why this is happening?


### Imports


```python
import pandas as pd
import numpy as np
import dask.dataframe as dd
```

### Create Dataframe


```python
idx = pd.date_range(start='2013-01-01 00:00:00', end='2013-01-01 06:00:00', freq='30t')
df_pd = pd.DataFrame({'col_1':[1,2,3,4,5,6,7,8,9,10,11,12,13,]}, index=idx)
df = dd.from_pandas(df_pd, npartitions=2)
df.compute()
```




<div>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th></th>
      <th>col_1</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>2013-01-01 00:00:00</th>
      <td>1</td>
    </tr>
    <tr>
      <th>2013-01-01 00:30:00</th>
      <td>2</td>
    </tr>
    <tr>
      <th>2013-01-01 01:00:00</th>
      <td>3</td>
    </tr>
    <tr>
      <th>2013-01-01 01:30:00</th>
      <td>4</td>
    </tr>
    <tr>
      <th>2013-01-01 02:00:00</th>
      <td>5</td>
    </tr>
    <tr>
      <th>2013-01-01 02:30:00</th>
      <td>6</td>
    </tr>
    <tr>
      <th>2013-01-01 03:00:00</th>
      <td>7</td>
    </tr>
    <tr>
      <th>2013-01-01 03:30:00</th>
      <td>8</td>
    </tr>
    <tr>
      <th>2013-01-01 04:00:00</th>
      <td>9</td>
    </tr>
    <tr>
      <th>2013-01-01 04:30:00</th>
      <td>10</td>
    </tr>
    <tr>
      <th>2013-01-01 05:00:00</th>
      <td>11</td>
    </tr>
    <tr>
      <th>2013-01-01 05:30:00</th>
      <td>12</td>
    </tr>
    <tr>
      <th>2013-01-01 06:00:00</th>
      <td>13</td>
    </tr>
  </tbody>
</table>
</div>




```python
df.divisions
```




    (Timestamp('2013-01-01 00:00:00', freq='30T'),
     Timestamp('2013-01-01 03:30:00', freq='30T'),
     Timestamp('2013-01-01 06:00:00', freq='30T'))



### This is what the shifted column looks like:


```python
df.col_1.shift(freq='1h').compute().to_frame()
```




<div>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th></th>
      <th>col_1</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>2013-01-01 01:00:00</th>
      <td>1</td>
    </tr>
    <tr>
      <th>2013-01-01 01:30:00</th>
      <td>2</td>
    </tr>
    <tr>
      <th>2013-01-01 02:00:00</th>
      <td>3</td>
    </tr>
    <tr>
      <th>2013-01-01 02:30:00</th>
      <td>4</td>
    </tr>
    <tr>
      <th>2013-01-01 03:00:00</th>
      <td>5</td>
    </tr>
    <tr>
      <th>2013-01-01 03:30:00</th>
      <td>6</td>
    </tr>
    <tr>
      <th>2013-01-01 04:00:00</th>
      <td>7</td>
    </tr>
    <tr>
      <th>2013-01-01 04:30:00</th>
      <td>8</td>
    </tr>
    <tr>
      <th>2013-01-01 05:00:00</th>
      <td>9</td>
    </tr>
    <tr>
      <th>2013-01-01 05:30:00</th>
      <td>10</td>
    </tr>
    <tr>
      <th>2013-01-01 06:00:00</th>
      <td>11</td>
    </tr>
    <tr>
      <th>2013-01-01 06:30:00</th>
      <td>12</td>
    </tr>
    <tr>
      <th>2013-01-01 07:00:00</th>
      <td>13</td>
    </tr>
  </tbody>
</table>
</div>



### This is what  assigning the shifted column without using lambda, looks like:


```python
df.assign(shifted=df.col_1.shift(freq='1h')).compute()
```




<div>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th></th>
      <th>col_1</th>
      <th>shifted</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>2013-01-01 00:00:00</th>
      <td>1</td>
      <td>NaN</td>
    </tr>
    <tr>
      <th>2013-01-01 00:30:00</th>
      <td>2</td>
      <td>NaN</td>
    </tr>
    <tr>
      <th>2013-01-01 01:00:00</th>
      <td>3</td>
      <td>1.0</td>
    </tr>
    <tr>
      <th>2013-01-01 01:30:00</th>
      <td>4</td>
      <td>2.0</td>
    </tr>
    <tr>
      <th>2013-01-01 02:00:00</th>
      <td>5</td>
      <td>3.0</td>
    </tr>
    <tr>
      <th>2013-01-01 02:30:00</th>
      <td>6</td>
      <td>4.0</td>
    </tr>
    <tr>
      <th>2013-01-01 03:00:00</th>
      <td>7</td>
      <td>5.0</td>
    </tr>
    <tr>
      <th>2013-01-01 03:30:00</th>
      <td>8</td>
      <td>6.0</td>
    </tr>
    <tr>
      <th>2013-01-01 04:00:00</th>
      <td>9</td>
      <td>7.0</td>
    </tr>
    <tr>
      <th>2013-01-01 04:30:00</th>
      <td>10</td>
      <td>8.0</td>
    </tr>
    <tr>
      <th>2013-01-01 05:00:00</th>
      <td>11</td>
      <td>9.0</td>
    </tr>
    <tr>
      <th>2013-01-01 05:30:00</th>
      <td>12</td>
      <td>10.0</td>
    </tr>
    <tr>
      <th>2013-01-01 06:00:00</th>
      <td>13</td>
      <td>11.0</td>
    </tr>
  </tbody>
</table>
</div>



### This is what  assigning the shifted column using lambda, looks like:


```python
df.assign(shifted=lambda x: x.col_1.shift(freq='1h')).compute()
```




<div>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th></th>
      <th>col_1</th>
      <th>shifted</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>2013-01-01 00:00:00</th>
      <td>1</td>
      <td>NaN</td>
    </tr>
    <tr>
      <th>2013-01-01 00:30:00</th>
      <td>2</td>
      <td>NaN</td>
    </tr>
    <tr>
      <th>2013-01-01 01:00:00</th>
      <td>3</td>
      <td>1.0</td>
    </tr>
    <tr>
      <th>2013-01-01 01:30:00</th>
      <td>4</td>
      <td>2.0</td>
    </tr>
    <tr>
      <th>2013-01-01 02:00:00</th>
      <td>5</td>
      <td>3.0</td>
    </tr>
    <tr>
      <th>2013-01-01 02:30:00</th>
      <td>6</td>
      <td>4.0</td>
    </tr>
    <tr>
      <th>2013-01-01 03:00:00</th>
      <td>7</td>
      <td>5.0</td>
    </tr>
    <tr>
      <th>2013-01-01 03:30:00</th>
      <td>8</td>
      <td>NaN</td>
    </tr>
    <tr>
      <th>2013-01-01 04:00:00</th>
      <td>9</td>
      <td>NaN</td>
    </tr>
    <tr>
      <th>2013-01-01 04:30:00</th>
      <td>10</td>
      <td>8.0</td>
    </tr>
    <tr>
      <th>2013-01-01 05:00:00</th>
      <td>11</td>
      <td>9.0</td>
    </tr>
    <tr>
      <th>2013-01-01 05:30:00</th>
      <td>12</td>
      <td>10.0</td>
    </tr>
    <tr>
      <th>2013-01-01 06:00:00</th>
      <td>13</td>
      <td>11.0</td>
    </tr>
  </tbody>
</table>
</div>




Discussion:
> The difference seems to happen only in the divisions limits. I guess this is a bug? Any idea on why this is happening?

Yes, this is a bug.

Callables and DataFrames are handled differently in elemwise: https://github.com/dask/dask/blob/83fdc8dc98418fbcb2393401644036dda40f5d37/dask/dataframe/core.py

With `shifted=df.col_1.shift(freq='1h')` the new column is "known" and aligned.

With a callable, the new column is not "known", so we need to call it. I suspect we call it on each partition, rather than on the entire dataframe. I'll do some digging later today.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2678,6 +2678,9 @@ def assign(self, **kwargs):
                     callable(v) or pd.api.types.is_scalar(v)):
                 raise TypeError("Column assignment doesn't support type "
                                 "{0}".format(type(v).__name__))
+            if callable(v):
+                kwargs[k] = v(self)
+
         pairs = list(sum(kwargs.items(), ()))
 
         # Figure out columns of the output


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
@@ -925,6 +925,13 @@ def test_assign():
         d.assign(foo=d_unknown.a)
 
 
+def test_assign_callable():
+    df = dd.from_pandas(pd.DataFrame({"A": range(10)}), npartitions=2)
+    a = df.assign(B=df.A.shift())
+    b = df.assign(B=lambda x: x.A.shift())
+    assert_eq(a, b)
+
+
 def test_map():
     assert_eq(d.a.map(lambda x: x + 1), full.a.map(lambda x: x + 1))
     lk = dict((v, v + 1) for v in full.a.values)