How to use generic polars expressions and register new ones.

polars_expressions.py

import polars as pl

df1 = pl.DataFrame({"A":[1,2 ,3 ,4 ,5 ,6], "B":[1,5,9,1,2,3]})
df2 = pl.DataFrame({"A":[5,10,15,20,25,30], "B":[1,5,9,1,2,3]})

dbl_expr = pl.col("A") * 2
sq_expr = pl.col("A") ** 2

df1.with_columns(dbl_A = dbl_expr,
                 sq_A = sq_expr)



# Register your own custom functrionality in a reusable name
@pl.api.register_expr_namespace("double_and_compute_n") # This is the name of the expression group that is run on a selection
class DoubleCompN:
    def __init__(self, expr: pl.Expr) -> None:
        self._expr = expr # This represents the expression where it is called e.g. `pl.col("A")` in  `pl.col("A").double_and_compute_n.add(5)`

    def add(self, n: int)-> pl.Expr:
        return ((self._expr * 2) + n)
    
    def subtract(self, n: int)-> pl.Expr:
        return ((self._expr * 2) - n)
      
      
      
df1.with_columns(
    pl.col("A"),
    (pl.col("A") * 2).alias('2n'),
    (pl.col("A").double_and_compute_n.add(3)).alias("da3"),
    (pl.col("A").double_and_compute_n.subtract(10).alias("ds10"))
    )

df1.with_columns(dbl_A = dbl_expr, sq_A = sq_expr)

	A	B	dbl_A	sq_A
0	1	1	2	1
1	2	5	4	4
2	3	9	6	9
3	4	1	8	16
4	5	2	10	25
5	6	3	12	36

df1.with_columns(
    pl.col("A"),
    (pl.col("A") * 2).alias('2n'),
    (pl.col("A").double_and_compute_n.add(3)).alias("da3"),
    (pl.col("A").double_and_compute_n.subtract(10).alias("ds10"))
    )

	A	B	2n	da3	ds10
0	1	1	2	5	-8
1	2	5	4	7	-6
2	3	9	6	9	-4
3	4	1	8	11	-2
4	5	2	10	13	0
5	6	3	12	15	2