From bdc1444b893b09099eaa3e1c1d9f5cabd543515a Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 20 Sep 2023 18:40:15 +0100 Subject: [PATCH] Support callables in DataFrame.assign --- python/cudf/cudf/core/dataframe.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/cudf/cudf/core/dataframe.py b/python/cudf/cudf/core/dataframe.py index 84c16b71997..7eb7ede9c70 100644 --- a/python/cudf/cudf/core/dataframe.py +++ b/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