Repository: dask/dask

Issue:
DataFrame.assign doesn't support "dependent" columns
**What happened**:

pandas' assign supports "dependent" columns, where later arguments to `.assign` can refer to earlier ones: https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.assign.html?highlight=assign#pandas.DataFrame.assign

```python
import pandas as pd

x = pd.DataFrame({"A": [1, 2]})

x.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)
```

With dask, this raises

```pytb
import pandas as pd
import dask.dataframe as dd

x = dd.from_pandas(pd.DataFrame({"A": [1, 2]}), npartitions=2)

x.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_351/661796857.py in <module>
      4 x = dd.from_pandas(pd.DataFrame({"A": [1, 2]}), npartitions=2)
      5 
----> 6 x.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)

/srv/conda/envs/notebook/lib/python3.8/site-packages/dask/dataframe/core.py in assign(self, **kwargs)
   4307                 )
   4308             if callable(v):
-> 4309                 kwargs[k] = v(self)
   4310 
   4311             if isinstance(v, Array):

/tmp/ipykernel_351/661796857.py in <lambda>(df)
      4 x = dd.from_pandas(pd.DataFrame({"A": [1, 2]}), npartitions=2)
      5 
----> 6 x.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)

/srv/conda/envs/notebook/lib/python3.8/site-packages/dask/dataframe/core.py in __getattr__(self, key)
   4042             return self[key]
   4043         else:
-> 4044             raise AttributeError("'DataFrame' object has no attribute %r" % key)
   4045 
   4046     def __dir__(self):

AttributeError: 'DataFrame' object has no attribute 'B'
```


**What you expected to happen**:

Same as pandas.

**Minimal Complete Verifiable Example**:

```python
import pandas as pd
import dask.dataframe as dd

x = dd.from_pandas(pd.DataFrame({"A": [1, 2]}), npartitions=2)

x.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)
```

In pandas, we do that by assigning to the result as we iterate through the kwargs: https://github.com/pandas-dev/pandas/blob/8814de9fbc7d53c6db201eab717fd3d6b607654b/pandas/core/frame.py#L4480-L4484

```python
        data = self.copy()

        for k, v in kwargs.items():
            data[k] = com.apply_if_callable(v, data)
        return data
```

I think something similar would work for dask, but I haven't checked.


Discussion:
@TomAugspurger, I'm new here, I would like to solve this issue.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -4400,6 +4400,7 @@ def categorize(self, columns=None, index=None, split_every=None, **kwargs):
 
     @derived_from(pd.DataFrame)
     def assign(self, **kwargs):
+        data = self.copy()
         for k, v in kwargs.items():
             if not (
                 isinstance(v, Scalar)
@@ -4414,26 +4415,29 @@ def assign(self, **kwargs):
                     "{0}".format(typename(type(v)))
                 )
             if callable(v):
-                kwargs[k] = v(self)
-
+                kwargs[k] = v(data)
             if isinstance(v, Array):
                 from .io import from_dask_array
 
                 if len(v.shape) > 1:
                     raise ValueError("Array assignment only supports 1-D arrays")
-                if v.npartitions != self.npartitions:
+                if v.npartitions != data.npartitions:
                     raise ValueError(
                         "Number of partitions do not match ({0} != {1})".format(
-                            v.npartitions, self.npartitions
+                            v.npartitions, data.npartitions
                         )
                     )
-                kwargs[k] = from_dask_array(v, index=self.index, meta=self._meta)
+                kwargs[k] = from_dask_array(v, index=data.index, meta=data._meta)
+
+            pairs = [k, kwargs[k]]
 
-        pairs = list(sum(kwargs.items(), ()))
+            # Figure out columns of the output
+            df2 = data._meta_nonempty.assign(
+                **_extract_meta({k: kwargs[k]}, nonempty=True)
+            )
+            data = elemwise(methods.assign, data, *pairs, meta=df2)
 
-        # Figure out columns of the output
-        df2 = self._meta_nonempty.assign(**_extract_meta(kwargs, nonempty=True))
-        return elemwise(methods.assign, self, *pairs, meta=df2)
+        return data
 
     @derived_from(pd.DataFrame, ua_args=["index"])
     def rename(self, index=None, columns=None):


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
@@ -1484,7 +1484,7 @@ def test_assign():
         d="string",
         e=ddf.a.sum(),
         f=ddf.a + ddf.b,
-        g=lambda x: x.a + x.b,
+        g=lambda x: x.a + x.c,
         dt=pd.Timestamp(2018, 2, 13),
     )
     res_unknown = ddf_unknown.assign(
@@ -1492,7 +1492,7 @@ def test_assign():
         d="string",
         e=ddf_unknown.a.sum(),
         f=ddf_unknown.a + ddf_unknown.b,
-        g=lambda x: x.a + x.b,
+        g=lambda x: x.a + x.c,
         dt=pd.Timestamp(2018, 2, 13),
     )
     sol = df.assign(
@@ -1500,7 +1500,7 @@ def test_assign():
         d="string",
         e=df.a.sum(),
         f=df.a + df.b,
-        g=lambda x: x.a + x.b,
+        g=lambda x: x.a + x.c,
         dt=pd.Timestamp(2018, 2, 13),
     )
     assert_eq(res, sol)
@@ -1527,6 +1527,14 @@ def test_assign():
     with pytest.raises(ValueError):
         ddf.assign(foo=ddf_unknown.a)
 
+    df = pd.DataFrame({"A": [1, 2]})
+    df.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)
+
+    ddf = dd.from_pandas(pd.DataFrame({"A": [1, 2]}), npartitions=2)
+    ddf.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)
+
+    assert_eq(df, ddf)
+
 
 def test_assign_callable():
     df = dd.from_pandas(pd.DataFrame({"A": range(10)}), npartitions=2)