Column op

Author: xinrong-meng

databricks/koalas/series.py

@@ -2023,13 +2023,13 @@ def factorize(self, sort: bool = True, na_sentinel: Optional[int] = -1):
         kvs = list(
             chain(*([(F.lit(unique), F.lit(code)) for unique, code in unique_to_code.items()]))
         )
-        sdf = sdf.withColumn(new_scol_name, F.create_map(kvs))
 
-        null_scol = F.when(scol_for(sdf, scol_name).isNull(), F.lit(na_sentinel_code))
-        mapped_scol = scol_for(sdf, new_scol_name).getItem(scol_for(sdf, scol_name))
-        sdf = sdf.withColumn(scol_name, null_scol.otherwise(mapped_scol))
+        map_scol = F.create_map(kvs)
+        null_scol = F.when(self.spark.column.isNull(), F.lit(na_sentinel_code))
+        mapped_scol = map_scol.getItem(self.spark.column)
+        internal = self._internal.with_new_columns([null_scol.otherwise(mapped_scol)])
 
-        codes = first_series(DataFrame(self._internal.with_new_sdf(sdf))).astype(np.int64)
+        codes = first_series(DataFrame(internal)).astype(np.int64)
 
         if na_sentinel is not None:
             uniques = ks.Index([x for x in uniques_list if x is not None])