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:
TomAugspurger: I *think* this is what we want to do here. This may surprise some users, if their callable either expects or incidentally creates a pandas dataframe.

If a concrete dataframe really is needed, then `df.map_blocks(pd.DataFrame.assign, column=callable)` should be used instead I think.
mrocklin: Is anyone available to review this?

Patch:
diff --git a/dask/dataframe/categorical.py b/dask/dataframe/categorical.py
index 9900b54af..f372aa2ca 100644
--- a/dask/dataframe/categorical.py
+++ b/dask/dataframe/categorical.py
@@ -184,7 +184,7 @@ class CategoricalAccessor(Accessor):
             Keywords to pass on to the call to `compute`.
         """
         if self.known:
-            return self
+            return self._series
         categories = self._property_map('categories').unique().compute(**kwargs)
         return self.set_categories(categories.values)
 
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 807111334..855eed317 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2681,6 +2681,9 @@ class DataFrame(_Frame):
                     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_categorical.py b/dask/dataframe/tests/test_categorical.py
index b84011818..3b643e82e 100644
--- a/dask/dataframe/tests/test_categorical.py
+++ b/dask/dataframe/tests/test_categorical.py
@@ -271,6 +271,14 @@ def assert_array_index_eq(left, right):
     assert_eq(left, pd.Index(right) if isinstance(right, np.ndarray) else right)
 
 
+def test_return_type_known_categories():
+    df = pd.DataFrame({"A": ['a', 'b', 'c']})
+    df['A'] = df['A'].astype('category')
+    dask_df = dd.from_pandas(df, 2)
+    ret_type = dask_df.A.cat.as_known()
+    assert isinstance(ret_type, dd.core.Series)
+
+
 class TestCategoricalAccessor:
 
     @pytest.mark.parametrize('series', cat_series)
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 41a50385d..a9e61cc86 100644
--- 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)