Support callables in DataFrame.assign

python/cudf/cudf/core/dataframe.py

@@ -1407,13 +1407,15 @@ def assign(self, **kwargs):
         """
         new_df = cudf.DataFrame(index=self.index.copy())
         for name, col in self._data.items():
-            if name in kwargs:
-                new_df[name] = kwargs.pop(name)
-            else:
+            # Copy non-overwritten kwargs
+            if name not in kwargs:
                 new_df._data[name] = col.copy()
 
         for k, v in kwargs.items():
-            new_df[k] = v
+            if callable(v):
+                new_df[k] = v(new_df)
+            else:
+                new_df[k] = v
         return new_df
 
     @classmethod