I am not able to run the code with the examples provided but three things immediately stick out for me that might be the cause of the long runtime.

1: LazyFrame:
Lazy evaluation is great in many circumstances where operations are reasonably linear. In this case, there are some linear and some dependent calculations. This may be creating a large backlog of lazy queries, especially in side of the for loop.

2: For Loop:
If the same actions are being taken for each combination of shift and code, it would probably be much faster to vectorize these nested loops. Without know what data you are passing into this function it is hard for me to know how to exactly write a vectorized version that gets rid of the for loops.

3: DataFrame vs Array:
This is more of a general remark, but its often the case that numpy arrays are significantly more performant for math operations than DataFrames. If the long runtime of the module is a concern, I would recommend bouncing that DataFrame into an array when it enters the function, doing all the required math on it, then converting it back to a DataFrame on the way out. Doing this back and forth conversion is normally much faster when there is a larger volume of math that needs to happen in between.