diff --git a/dask_sql/datacontainer.py b/dask_sql/datacontainer.py
index cadc44f74..9f272a17f 100644
--- a/dask_sql/datacontainer.py
+++ b/dask_sql/datacontainer.py
@@ -171,14 +171,15 @@ def assign(self) -> dd.DataFrame:
         a dataframe which has the the columns specified in the
         stored ColumnContainer.
         """
-        df = self.df.assign(
-            **{
-                col_from: self.df[col_to]
-                for col_from, col_to in self.column_container.mapping()
-                if col_from in self.column_container.columns
-            }
-        )
-        return df[self.column_container.columns]
+        df = self.df[
+            [
+                self.column_container._frontend_backend_mapping[out_col]
+                for out_col in self.column_container.columns
+            ]
+        ]
+        df.columns = self.column_container.columns
+
+        return df
 
 
 class UDF: