Fix bug that string_embed function Only replace one columns when repl…

…ace=True. (#5497)

* add for loop to replace all Columns.

* add replace test for string_embds function

python/friesian/src/bigdl/friesian/feature/table.py

@@ -2545,7 +2545,8 @@ def string_embed(
                     .withColumn(c + "_embds", tolist(c + "_embds"))
 
         if replace:
-            df = df.drop(c).withColumnRenamed(c + "_embds", c)
+            for c in cols:
+                df = df.drop(c).withColumnRenamed(c + "_embds", c)
         return FeatureTable(df)
 
 

python/friesian/test/bigdl/friesian/feature/test_table.py

@@ -1512,6 +1512,10 @@ def test_string_embed(self):
         invalidInputError(text_embeds.select("text1_embds").size() == 3, "size error")
         invalidInputError(text_embeds.select("text2_embds").size() == 3, "size error")
 
+        text_embeds_replaced = tbl.string_embed(["text1", "text2"], reduce_dim=5, replace=True)
+        invalidInputError('text1_embds' not in text_embeds_replaced.columns, 'replace failed')
+        invalidInputError('text2_embds' not in text_embeds_replaced.columns, 'replace failed')
+
         with self.assertRaises(Exception) as context:
             text_embeds = tbl.string_embed(["text1"], reduce_dim=1000)
         self.assertTrue("must be no less than k=1000" in str(context.exception))