[SNOW-1879403] Fix replace when passing Column object to avoid incorrect usage of lit(...) from within. (#2858)

<!---
Please answer these questions before creating your pull request. Thanks!
--->

1. Which Jira issue is this PR addressing? Make sure that there is an
accompanying issue to your PR.

   Fixes SNOW-1879403

2. Fill out the following pre-review checklist:

- [x] I am adding a new automated test(s) to verify correctness of my
new code
- [ ] If this test skips Local Testing mode, I'm requesting review from
@snowflakedb/local-testing
   - [ ] I am adding new logging messages
   - [ ] I am adding a new telemetry message
   - [ ] I am adding new credentials
   - [ ] I am adding a new dependency
- [ ] If this is a new feature/behavior, I'm adding the Local Testing
parity changes.
- [x] I acknowledge that I have ensured my changes to be thread-safe.
Follow the link for more information: [Thread-safe Developer
Guidelines](https://github.com/snowflakedb/snowpark-python/blob/main/CONTRIBUTING.md#thread-safe-development)

3. Please describe how your code solves the related issue.

`replace` supports both `Column` and literal objects. The code within
the function directly called `lit(...)`, however `lit(...)` can't be
called with a `Column` object. Adding if to fix this.

Author: sfc-gh-lspiegelberg

CHANGELOG.md

@@ -51,8 +51,9 @@
 #### Bug Fixes
 
 - Fixed a bug in local testing mode that caused a column to contain None when it should contain 0
-- Fixed a bug in `StructField.from_json` that prevented TimestampTypes with tzinfo from being parsed correctly.
+- Fixed a bug in `StructField.from_json` that prevented TimestampTypes with `tzinfo` from being parsed correctly.
 - Fixed a bug in function `date_format` that caused an error when the input column was date type or timestamp type.
+- Fixed a bug in `replace` when passing `Column` expression objects.
 
 ### Snowpark pandas API Updates
 

src/snowflake/snowpark/functions.py

@@ -3384,8 +3384,12 @@ def replace(
         if _emit_ast
         else None
     )
-    pat = lit(pattern, _emit_ast=_emit_ast)
-    rep = lit(replacement, _emit_ast=_emit_ast)
+    pat = pattern if isinstance(pattern, Column) else lit(pattern, _emit_ast=_emit_ast)
+    rep = (
+        replacement
+        if isinstance(replacement, Column)
+        else lit(replacement, _emit_ast=_emit_ast)
+    )
     return builtin(sql_func_name, _ast=ast, _emit_ast=False)(sub, pat, rep)
 
 

tests/integ/test_dataframe.py

@@ -4459,3 +4459,23 @@ def test_with_column_keep_column_order(session):
     assert df3.columns == ["A", "B", "C"]
     df3 = df.with_columns(["C", "A"], [lit(0), lit(0)], keep_column_order=True)
     assert df3.columns == ["A", "B", "C"]
+
+
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="replace function is not supported in Local Testing",
+)
+def test_SNOW_1879403_replace_with_lit(session):
+
+    # TODO SNOW-1880749: support for local testing mode.
+
+    from snowflake.snowpark.functions import replace
+
+    df = session.create_dataframe(
+        [["apple"], ["apple pie"], ["apple juice"]], schema=["a"]
+    )
+    ans = df.select(
+        replace(col("a"), lit("apple"), lit("orange")).alias("result")
+    ).collect()
+
+    Utils.check_answer(ans, [Row("orange"), Row("orange pie"), Row("orange juice")])