Skip to content

Commit

Permalink
perf(python): Lower overhead for BytecodeParser on introspection of…
Browse files Browse the repository at this point in the history
… incompatible UDFs (#20280)
  • Loading branch information
alexander-beedie authored Dec 13, 2024
1 parent df8672d commit 84d317a
Showing 1 changed file with 15 additions and 7 deletions.
22 changes: 15 additions & 7 deletions py-polars/polars/_utils/udfs.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ class OpNames:
| set(SYNTHETIC)
| LOAD_VALUES
)
MATCHABLE_OPS = PARSEABLE_OPS | set(BINARY) | LOAD_ATTR | CALL
UNARY_VALUES = frozenset(UNARY.values())


Expand Down Expand Up @@ -753,11 +754,18 @@ def __init__(
self._function = function
self._caller_variables = caller_variables
self._original_instructions = list(instructions)
self._rewritten_instructions = self._rewrite(
self._upgrade_instruction(inst)
for inst in self._unpack_superinstructions(self._original_instructions)
if inst.opname not in self._ignored_ops
)

normalised_instructions = []

for inst in self._unpack_superinstructions(self._original_instructions):
if inst.opname not in self._ignored_ops:
if inst.opname not in OpNames.MATCHABLE_OPS:
self._rewritten_instructions = []
return
upgraded_inst = self._upgrade_instruction(inst)
normalised_instructions.append(upgraded_inst)

self._rewritten_instructions = self._rewrite(normalised_instructions)

def __len__(self) -> int:
return len(self._rewritten_instructions)
Expand Down Expand Up @@ -810,15 +818,15 @@ def _matches(
return instructions
return []

def _rewrite(self, instructions: Iterator[Instruction]) -> list[Instruction]:
def _rewrite(self, instructions: list[Instruction]) -> list[Instruction]:
"""
Apply rewrite rules, potentially injecting synthetic operations.
Rules operate on the instruction stream and can examine/modify
it as needed, pushing updates into "updated_instructions" and
returning True/False to indicate if any changes were made.
"""
self._instructions = list(instructions)
self._instructions = instructions
updated_instructions: list[Instruction] = []
idx = 0
while idx < len(self._instructions):
Expand Down

0 comments on commit 84d317a

Please sign in to comment.