pythonapplypython-polars

How to apply custom functions with multiple parameters in Polars?


Now I have a dataframe:

df = pd.DataFrame({
    "a":[1,2,3,4,5],
    "b":[2,3,4,5,6],
    "c":[3,4,5,6,7]
})

The function:

def fun(a,b,shift_len): 
     return a+b*shift_len,b-shift_len

Using Pandas, I can get the result by:

df[["d","e"]] = df.apply(lambda row:fun(row["a"],row["b"],3),axis=1,result_type="expand")

I want to know how can I use polars to get the same result?


Solution

  • Passing arguments with args

    import pandas as pd
    df1 = pd.DataFrame({"a":[1,2,3,4,5],"b":[2,3,4,5,6],"c":[3,4,5,6,7]})
    
    
    def t(df, row1, row2, shift_len):
        return df[row1] + df[row2] * shift_len, df[row2] - shift_len
    
    
    df1[["d", "e"]] = df1.apply(t, args=("a", "b", 3), axis=1, result_type="expand")
    print(df1)
    
    OUTPUT:
    
       a  b  c   d  e
    0  1  2  3   7 -1
    1  2  3  4  11  0
    2  3  4  5  15  1
    3  4  5  6  19  2
    4  5  6  7  23  3