diff --git a/py-polars/polars/_utils/udfs.py b/py-polars/polars/_utils/udfs.py index 0ff968ed59ec..cfc9d06492b1 100644 --- a/py-polars/polars/_utils/udfs.py +++ b/py-polars/polars/_utils/udfs.py @@ -183,6 +183,7 @@ class OpNames: "endswith": "str.ends_with", "lower": "str.to_lowercase", "lstrip": "str.strip_chars_start", + "replace": "str.replace", "rstrip": "str.strip_chars_end", "startswith": "str.starts_with", "strip": "str.strip_chars", @@ -983,7 +984,7 @@ def _rewrite_methods( """Replace python method calls with synthetic POLARS_EXPRESSION op.""" LOAD_METHOD = OpNames.LOAD_ATTR if _MIN_PY312 else {"LOAD_METHOD"} if matching_instructions := ( - # method call with one basic arg, eg: "s.endswith('!')" + # method call with one arg, eg: "s.endswith('!')" self._matches( idx, opnames=[LOAD_METHOD, {"LOAD_CONST"}, OpNames.CALL], @@ -1016,6 +1017,47 @@ def _rewrite_methods( px = inst._replace(opname="POLARS_EXPRESSION", argval=expr, argrepr=expr) updated_instructions.append(px) + elif matching_instructions := ( + # method call with three args, eg: "s.replace('!','?',count=2)" + self._matches( + idx, + opnames=[ + LOAD_METHOD, + {"LOAD_CONST"}, + {"LOAD_CONST"}, + {"LOAD_CONST"}, + OpNames.CALL, + ], + argvals=[_PYTHON_METHODS_MAP], + ) + or + # method call with two args, eg: "s.replace('!','?')" + self._matches( + idx, + opnames=[LOAD_METHOD, {"LOAD_CONST"}, {"LOAD_CONST"}, OpNames.CALL], + argvals=[_PYTHON_METHODS_MAP], + ) + ): + inst = matching_instructions[0] + expr = _PYTHON_METHODS_MAP[inst.argval] + + param_values = [ + i.argval + for i in matching_instructions[1 : len(matching_instructions) - 1] + ] + if expr == "str.replace": + if len(param_values) == 3: + old, new, count = param_values + expr += f"({old!r},{new!r},n={count},literal=True)" + else: + old, new = param_values + expr = f"str.replace_all({old!r},{new!r},literal=True)" + else: + expr += f"({','.join(repr(v) for v in param_values)})" + + px = inst._replace(opname="POLARS_EXPRESSION", argval=expr, argrepr=expr) + updated_instructions.append(px) + return len(matching_instructions) @staticmethod diff --git a/py-polars/tests/unit/operations/map/test_inefficient_map_warning.py b/py-polars/tests/unit/operations/map/test_inefficient_map_warning.py index 74946f084d51..13eee2c731e3 100644 --- a/py-polars/tests/unit/operations/map/test_inefficient_map_warning.py +++ b/py-polars/tests/unit/operations/map/test_inefficient_map_warning.py @@ -176,6 +176,16 @@ """lambda x: x.lstrip().startswith(('!','#','?',"'"))""", """pl.col("b").str.strip_chars_start().str.contains(r"^(!|\\#|\\?|')")""", ), + ( + "b", + "lambda x: x.replace(':','')", + """pl.col("b").str.replace_all(':','',literal=True)""", + ), + ( + "b", + "lambda x: x.replace(':','',2)", + """pl.col("b").str.replace(':','',n=2,literal=True)""", + ), # --------------------------------------------- # json expr: load/extract # ---------------------------------------------