Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Polars select multiple element-wise products

Suppose I have the following dataframe:

the_df = pl.DataFrame({'x1': [1,1,1], 'x2': [2,2,2], 'y1': [1,1,1], 'y2': [2,2,2]})
┌─────┬─────┬─────┬─────┐
│ x1  ┆ x2  ┆ y1  ┆ y2  │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╡
│ 1   ┆ 2   ┆ 1   ┆  2  │
│ 1   ┆ 2   ┆ 1   ┆  2  │
│ 1   ┆ 2   ┆ 1   ┆  2  │
└─────┴─────┴─────┴─────┘

And and two lists, xs = ['x1', 'x2'], ys = ['y1', 'y2'].

Is there a good way to add the products between x1/y1 and x2/y2 using .select()? So the result should look like the following. Specifically, I want to use the lists rather than writing out z1=x1*y1, z2=x2*y2 (the real data has more terms I want to multiply).

┌─────┬─────┬─────┬─────┬─────┬─────┐
│ x1  ┆ x2  ┆ y1  ┆ y2  ┆ z1  ┆ z2  │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╪═════╪═════╡
│ 1   ┆ 2   ┆ 1   ┆ 2   ┆ 1   ┆ 4   │
│ 1   ┆ 2   ┆ 1   ┆ 2   ┆ 1   ┆ 4   │
│ 1   ┆ 2   ┆ 1   ┆ 2   ┆ 1   ┆ 4   │
└─────┴─────┴─────┴─────┴─────┴─────┘
like image 856
dfried Avatar asked Oct 12 '25 01:10

dfried


1 Answers

you can do something like this:

zs = ['z1','z2']

df.with_columns(
    (pl.col(xc) * pl.col(yc)).alias(zc) for xc, yc, zc in zip(xs, ys, zs)
)

┌─────┬─────┬─────┬─────┬─────┬─────┐
│ x1  ┆ x2  ┆ y1  ┆ y2  ┆ z1  ┆ z2  │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╪═════╪═════╡
│ 1   ┆ 2   ┆ 1   ┆ 2   ┆ 1   ┆ 4   │
│ 1   ┆ 2   ┆ 1   ┆ 2   ┆ 1   ┆ 4   │
│ 1   ┆ 2   ┆ 1   ┆ 2   ┆ 1   ┆ 4   │
└─────┴─────┴─────┴─────┴─────┴─────┘
like image 116
Roman Pekar Avatar answered Oct 14 '25 16:10

Roman Pekar