Repository: mars-project/mars

Issue:
Support DataFrame.assign
Support DataFrame.assign. pandas docs see https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.assign.html .


Discussion:
Hi, 
Can I work on this for my first issue?
> Hi,
> Can I work on this for my first issue?

Sure, fork & send a PR :)
Is there any similar example in the current codebase that I can follow?
> Is there any similar example in the current codebase that I can follow?

Since Mars DataFrame already support `__setitem__`, so basically what you need to do is to add an `assign` function to 

https://github.com/mars-project/mars/blob/0c9114d1fffe56af7e3a3da01c7534e9597237a2/mars/dataframe/core.py#L1599

in which `__setitem__` can be used.

Resolving patch:
diff --git a/mars/dataframe/core.py b/mars/dataframe/core.py
--- a/mars/dataframe/core.py
+++ b/mars/dataframe/core.py
@@ -1838,6 +1838,84 @@ def itertuples(self, index=True, name='Pandas', batch_size=1000, session=None):
         return self._data.itertuples(batch_size=batch_size, session=session,
                                      index=index, name=name)
 
+    def assign(self, **kwargs):
+        """
+        Assign new columns to a DataFrame.
+        Returns a new object with all original columns in addition to new ones.
+        Existing columns that are re-assigned will be overwritten.
+
+        Parameters
+        ----------
+        **kwargs : dict of {str: callable or Series}
+            The column names are keywords. If the values are
+            callable, they are computed on the DataFrame and
+            assigned to the new columns. The callable must not
+            change input DataFrame (though pandas doesn't check it).
+            If the values are not callable, (e.g. a Series, scalar, or array),
+            they are simply assigned.
+
+        Returns
+        -------
+        DataFrame
+            A new DataFrame with the new columns in addition to
+            all the existing columns.
+
+        Notes
+        -----
+        Assigning multiple columns within the same ``assign`` is possible.
+        Later items in 'kwargs' may refer to newly created or modified
+        columns in 'df'; items are computed and assigned into 'df' in order.
+
+        Examples
+        --------
+        >>> import mars.dataframe as md
+        >>> df = md.DataFrame({'temp_c': [17.0, 25.0]},
+        ...                   index=['Portland', 'Berkeley'])
+        >>> df.execute()
+                  temp_c
+        Portland    17.0
+        Berkeley    25.0
+
+        Where the value is a callable, evaluated on `df`:
+
+        >>> df.assign(temp_f=lambda x: x.temp_c * 9 / 5 + 32)
+        >>> df.execute()
+                  temp_c  temp_f
+        Portland    17.0    62.6
+        Berkeley    25.0    77.0
+
+        Alternatively, the same behavior can be achieved by directly
+        referencing an existing Series or sequence:
+
+        >>> df.assign(temp_f=df['temp_c'] * 9 / 5 + 32)
+        >>> df.execute()
+                  temp_c  temp_f
+        Portland    17.0    62.6
+        Berkeley    25.0    77.0
+
+        You can create multiple columns within the same assign where one
+        of the columns depends on another one defined within the same assign:
+
+        >>> df.assign(temp_f=lambda x: x['temp_c'] * 9 / 5 + 32,
+        ...           temp_k=lambda x: (x['temp_f'] +  459.67) * 5 / 9)
+        >>> df.execute()
+                  temp_c  temp_f  temp_k
+        Portland    17.0    62.6  290.15
+        Berkeley    25.0    77.0  298.15
+        """
+
+        def apply_if_callable(maybe_callable, obj, **kwargs):
+            if callable(maybe_callable):
+                return maybe_callable(obj, **kwargs)
+
+            return maybe_callable
+
+        data = self.copy()
+
+        for k, v in kwargs.items():
+            data[k] = apply_if_callable(v, data)
+        return data
+
 
 class DataFrameGroupByChunkData(BaseDataFrameChunkData):
     type_name = 'DataFrameGroupBy'


Test patch:
diff --git a/mars/dataframe/tests/test_core.py b/mars/dataframe/tests/test_core.py
--- a/mars/dataframe/tests/test_core.py
+++ b/mars/dataframe/tests/test_core.py
@@ -201,6 +201,101 @@ def test_to_frame_or_series(setup):
     pd.testing.assert_series_equal(raw.to_series(name='new_name'), result)
 
 
+def test_assign(setup):
+    df = DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
+    original = df.copy()
+    result = df.assign(C=df.B / df.A)
+    result = result.execute().fetch()
+    expected = df.copy()
+    expected["C"] = [4, 2.5, 2]
+    expected_temp = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected_temp)
+
+    # lambda syntax
+    result = df.assign(C=lambda x: x.B / x.A)
+    result = result.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected_temp)
+
+    # original is unmodified
+    original_temp = original.copy().execute().fetch()
+    df_temp = df.copy().execute().fetch()
+    pd.testing.assert_frame_equal(df_temp, original_temp)
+
+    # Non-Series array-like
+    result = df.assign(C=[4, 2.5, 2])
+    result = result.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected_temp)
+
+    # original is unmodified
+    original_temp = original.copy().execute().fetch()
+    df_temp = df.copy().execute().fetch()
+    pd.testing.assert_frame_equal(original_temp, df_temp)
+
+    result = df.assign(B=df.B / df.A)
+    result = result.execute().fetch()
+    expected = expected.drop("B", axis=1).rename(columns={"C": "B"})
+    expected = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected)
+
+    # overwrite
+    result = df.assign(A=df.A + df.B)
+    result = result.execute().fetch()
+    expected = df.copy()
+    expected["A"] = [5, 7, 9]
+    expected["A"] = expected["A"].astype('int64')
+    expected_temp = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected_temp)
+
+    # lambda
+    result = df.assign(A=lambda x: x.A + x.B)
+    result = result.execute().fetch()
+    expected = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected)
+
+
+def test_assign_multiple(setup):
+    df = DataFrame([[1, 4], [2, 5], [3, 6]], columns=["A", "B"])
+    result = df.assign(C=[7, 8, 9], D=df.A, E=lambda x: x.B)
+    result["C"] = result["C"].astype('int64')
+    result = result.execute().fetch()
+    expected = DataFrame(
+        [[1, 4, 7, 1, 4], [2, 5, 8, 2, 5], [3, 6, 9, 3, 6]], columns=list("ABCDE")
+    )
+    expected = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected)
+
+
+def test_assign_order(setup):
+    df = DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
+    result = df.assign(D=df.A + df.B, C=df.A - df.B)
+    result = result.execute().fetch()
+    expected = DataFrame([[1, 2, 3, -1], [3, 4, 7, -1]], columns=list("ABDC"))
+    expected = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected)
+
+    result = df.assign(C=df.A - df.B, D=df.A + df.B)
+    result = result.execute().fetch()
+    expected = DataFrame([[1, 2, -1, 3], [3, 4, -1, 7]], columns=list("ABCD"))
+    expected = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected)
+
+
+def test_assign_dependent(setup):
+    df = DataFrame({"A": [1, 2], "B": [3, 4]})
+
+    result = df.assign(C=df.A, D=lambda x: x["A"] + x["C"])
+    result = result.execute().fetch()
+    expected = DataFrame([[1, 3, 1, 2], [2, 4, 2, 4]], columns=list("ABCD"))
+    expected = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected)
+
+    result = df.assign(C=lambda df: df.A, D=lambda df: df["A"] + df["C"])
+    result = result.execute().fetch()
+    expected = DataFrame([[1, 3, 1, 2], [2, 4, 2, 4]], columns=list("ABCD"))
+    expected = expected.execute().fetch()
+    pd.testing.assert_frame_equal(result, expected)
+
+
 def test_key_value(setup):
     raw = pd.DataFrame(np.random.rand(4, 3), columns=list('ABC'))
     df = DataFrame(raw)