feat(python): Identify inefficient use of Python string removeprefix, removesuffix, and zfill in map_elements

[alexander-beedie]

Following on from #19668; a few other string functions that we can easily detect and replace 1:1 with an expression.

Example

import polars as pl

df = pl.DataFrame({"px1": ["$.25"]})
df.with_columns(
    px2 = pl.col("px1").map_elements(lambda s: s.removeprefix("$").zfill(5))
)

# PolarsInefficientMapWarning:
#   Replace this expression...
#     - pl.col("px1").map_elements(lambda s: ...)
#   with this one instead:
#     + pl.col("px1").str.strip_prefix('$').str.zfill(5)

shape: (1, 2)
┌──────┬───────┐
│ px1  ┆ px2   │
│ ---  ┆ ---   │
│ str  ┆ str   │
╞══════╪═══════╡
│ $.25 ┆ 00.25 │
└──────┴───────┘

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 79.74%. Comparing base (4b03406) to head (6fde05a).
Report is 9 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #19672   +/-   ##
=======================================
  Coverage   79.73%   79.74%           
=======================================
  Files        1541     1541           
  Lines      212227   212227           
  Branches     2449     2449           
=======================================
+ Hits       169226   169240   +14     
+ Misses      42446    42432   -14     
  Partials      555      555           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ritchie46]

Am I seeing it correctly that we do that recursively? :o

[alexander-beedie]

Am I seeing it correctly that we do that recursively? :o

More or less ;) We walk all of the disassembled bytecode and can construct a compound expression if all of the individually recognised fragments are compatible. The fiddliest part was reconstructing parenthetical nesting from a stack of RPN with jump offsets...

Example

from polars._utils.udfs import BytecodeParser

bp = BytecodeParser(
    function = lambda x: (x > 1 and x != 2) or ((x % 2 == 0) and (x < 3)), 
    map_target = "expr",
)

We actually read the Instruction objects, but the dis output shows the outline:

bp.dis()
# 2           0 RESUME                   0
#             2 LOAD_FAST                0 (x)
#             4 LOAD_CONST               1 (1)
#             6 COMPARE_OP              68 (>)
#            10 COPY                     1
#            12 POP_JUMP_IF_FALSE        5 (to 24)
#            14 POP_TOP
#            16 LOAD_FAST                0 (x)
#            18 LOAD_CONST               2 (2)
#            20 COMPARE_OP              55 (!=)
#       >>   24 COPY                     1
#            26 POP_JUMP_IF_TRUE        15 (to 58)
#            28 POP_TOP
#            30 LOAD_FAST                0 (x)
#            32 LOAD_CONST               2 (2)
#            34 BINARY_OP                6 (%)
#            38 LOAD_CONST               3 (0)
#            40 COMPARE_OP              40 (==)
#            44 COPY                     1
#            46 POP_JUMP_IF_FALSE        5 (to 58)
#            48 POP_TOP
#            50 LOAD_FAST                0 (x)
#            52 LOAD_CONST               4 (3)
#            54 COMPARE_OP               2 (<)
#       >>   58 RETURN_VALUE

bp.can_attempt_rewrite()
True

print(
    bp.to_expression("colx")
)
# ((pl.col("colx") > 1) & (pl.col("colx") != 2)) 
#   | (((pl.col("colx") % 2) == 0) & (pl.col("colx") < 3))