From 00c1995b4aa67f94bfd88ca85a509c058a219ae6 Mon Sep 17 00:00:00 2001
From: Bidek56 <darek.bidek@gmail.com>
Date: Thu, 13 Jun 2024 06:58:21 -0400
Subject: [PATCH] Adding string to replace param list

---
 __tests__/expr.test.ts    | 43 ++++++++++++++++++++++++++++++++++-----
 polars/lazy/expr/index.ts | 12 +++++------
 2 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/__tests__/expr.test.ts b/__tests__/expr.test.ts
index 2c72b04b4..4a324034f 100644
--- a/__tests__/expr.test.ts
+++ b/__tests__/expr.test.ts
@@ -1178,17 +1178,43 @@ describe("expr.str", () => {
     expect(seriesActual).toFrameEqual(expected);
   });
   test("expr.replace", () => {
-    const df = pl.DataFrame({ a: [1, 2, 2, 3] });
+    const df = pl.DataFrame({ a: [1, 2, 2, 3], b: ["a", "b", "c", "d"] });
     let actual = df.withColumns(pl.col("a").replace(2, 100).alias("replaced"));
     let expected = pl.DataFrame({
       a: [1, 2, 2, 3],
+      b: ["a", "b", "c", "d"],
       replaced: [1, 100, 100, 3],
     });
     expect(actual).toFrameEqual(expected);
     actual = df.withColumns(
       pl.col("a").replace([2, 3], [100, 200], -1, pl.Float64).alias("replaced"),
     );
-    expected = pl.DataFrame({ a: [1, 2, 2, 3], replaced: [-1, 100, 100, 200] });
+    expected = pl.DataFrame({
+      a: [1, 2, 2, 3],
+      b: ["a", "b", "c", "d"],
+      replaced: [-1, 100, 100, 200],
+    });
+    expect(actual).toFrameEqual(expected);
+    actual = df.withColumns(
+      pl.col("b").replace("a", "c", "e", pl.Utf8).alias("replaced"),
+    );
+    expected = pl.DataFrame({
+      a: [1, 2, 2, 3],
+      b: ["a", "b", "c", "d"],
+      replaced: ["c", "e", "e", "e"],
+    });
+    expect(actual).toFrameEqual(expected);
+    actual = df.withColumns(
+      pl
+        .col("b")
+        .replace(["a", "b"], ["c", "d"], "e", pl.Utf8)
+        .alias("replaced"),
+    );
+    expected = pl.DataFrame({
+      a: [1, 2, 2, 3],
+      b: ["a", "b", "c", "d"],
+      replaced: ["c", "d", "e", "e"],
+    });
     expect(actual).toFrameEqual(expected);
     const mapping = { 2: 100, 3: 200 };
     actual = df.withColumns(
@@ -1197,13 +1223,20 @@ describe("expr.str", () => {
         .replace({ old: mapping, default_: -1, returnDtype: pl.Int64 })
         .alias("replaced"),
     );
-    expected = pl.DataFrame({ a: [1, 2, 2, 3], replaced: [-1, 100, 100, 200] });
+    expected = pl.DataFrame({
+      a: [1, 2, 2, 3],
+      b: ["a", "b", "c", "d"],
+      replaced: [-1, 100, 100, 200],
+    });
     expect(actual).toFrameEqual(expected);
-
     actual = df.withColumns(
       pl.col("a").replace({ old: mapping }).alias("replaced"),
     );
-    expected = pl.DataFrame({ a: [1, 2, 2, 3], replaced: [1, 100, 100, 200] });
+    expected = pl.DataFrame({
+      a: [1, 2, 2, 3],
+      b: ["a", "b", "c", "d"],
+      replaced: [1, 100, 100, 200],
+    });
     expect(actual).toFrameEqual(expected);
   });
   test("slice", () => {
diff --git a/polars/lazy/expr/index.ts b/polars/lazy/expr/index.ts
index 21aacba7d..f8a6b7f48 100644
--- a/polars/lazy/expr/index.ts
+++ b/polars/lazy/expr/index.ts
@@ -871,9 +871,9 @@ export interface Expr
    * ``` 
    */
   replace(
-    old: Expr | number | number[],
-    new_: Expr | number | number[],
-    default_?: Expr | number | number[],
+    old: Expr | string | number | (number | string)[],
+    new_: Expr | string | number | (number | string)[],
+    default_?: Expr | string | number | (number | string)[],
     returnDtype?: DataType,
   ): Expr;
   replace({
@@ -882,9 +882,9 @@ export interface Expr
     default_,
     returnDtype,
   }: {
-    old: unknown | Expr | number | number[];
-    new_?: Expr | number | number[];
-    default_?: Expr | number | number[];
+    old: unknown | Expr | string | number | (number | string)[];
+    new_?: Expr | string | number | (number | string)[];
+    default_?: Expr | string | number | (number | string)[];
     returnDtype?: DataType;
   }): Expr;
   /** Reverse the arrays in the list */