How do you "pipe" an expression in Polars?
Consider this code:
def transformation(col:pl.Series)->pl.Series: return col.tanh().suffix('_tanh') It'd be nice to be able to do this:
df.with_columns([ pl.col('colA').pipe(transformation), pl.col('colB').pipe(transformation), pl.col('colC').pipe(transformation), pl.col('colD').pipe(transformation), ]) But I don't think Polars supports .pipe for Series / expressions.
The alternative is
df.with_columns([ transformation(pl.col('colA')), transformation(pl.col('colB')), transformation(pl.col('colC')), transformation(pl.col('colD')), ]) But this gets messy (IMO) when you have arguments to the transformation function
Edit:
I implemented this and it "works" for me
def _pipe(self, func, *args, **kwargs): return func(self, *args, **kwargs) pl.Expr.pipe = _pipe 2 Answers
Typically (like pandas) you'd apply pipe at the DataFrame level.
Especially in conjunction with lazy-eval, this would be equivalent to chaining expressions; your function will receive the underlying eager/lazy frame, along with any optional *args and **kwargs, and by making it lazy() you ensure that your chain of operations can still take advantage of the query optimiser and parallelisation.
For example:
import polars as pl # define some UDFs def extend_with_tan( df ): return df.with_columns( pl.all().tanh().suffix("_tanh") ) def mul_in_place( df, n ): return df.select( (pl.all() * n).suffix(f"_x{n}") ) # init lazyframe df = pl.DataFrame({ "colA": [-4], "colB": [-2], "colC": [10], }).lazy() # pipe/result dfx = df.pipe( extend_with_tan ).pipe( mul_in_place,n=3 ) dfx.collect() # ┌─────────┬─────────┬─────────┬──────────────┬──────────────┬──────────────┐ # │ colA_x3 ┆ colB_x3 ┆ colC_x3 ┆ colA_tanh_x3 ┆ colB_tanh_x3 ┆ colC_tanh_x3 │ # │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │ # │ i64 ┆ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 │ # ╞═════════╪═════════╪═════════╪══════════════╪══════════════╪══════════════╡ # │ -12 ┆ -6 ┆ 30 ┆ -2.997988 ┆ -2.892083 ┆ 3.0 │ # └─────────┴─────────┴─────────┴──────────────┴──────────────┴──────────────┘ - API Docs: polars "pipe" method
As you've probably realized, adding custom methods in order to be able to do method chaining is unfortunately not a first-class citizen in python.
In polars, a canonical way that hopefully satisfies you is to instead write a function that returns an expression. You do this already (although the type hint is incorrectly set to pl.Series), but can save some space by giving a string argument to our transformation function:
import polars as pl df = pl.DataFrame({"colA": [-4], "colB": [-2], "colC": [0], "colD": [2]}) def transformation(name: str | list[str]) -> pl.Expr: return pl.col(name).tanh().suffix("_tanh") df1 = df.with_columns( [ transformation("colA"), transformation("colB"), transformation("colC"), transformation("colD"), ] ) I realise this doesn't quite do what you wanted, but perhaps the following will cheer you up a bit. Since pl.col() can take a list of column names, we can do the following:
df2 = df.with_column(transformation(["colA", "colB", "colC", "colD"])) assert df1.frame_equal(df2) # True And we can even target all of them using a regular expression:
# ^col\w+$ is a regular expression matching `col<anything>` df3 = df.with_column(transformation("^col\w+$")) assert df1.frame_equal(df3) # True 
