pass multiple columns to UDF at once through regex about polars HOT 6 OPEN

df = pl.DataFrame({
    "B_1": [1, 2, 3, 4],
    "B_2": [5, 6, 7, 8],
    "B_3": [9, 10, 11, 12]
}).with_row_index()

df.with_columns(
    pl.map_batches(
        pl.struct("^B_.*$"),
        lambda x: x[0].struct[0] + x[0].struct[1] + x[0].struct[2]
    ).name.prefix("C_")
)

# shape: (4, 5)
# ┌───────┬─────┬─────┬─────┬───────┐
# │ index ┆ B_1 ┆ B_2 ┆ B_3 ┆ C_B_1 │
# │ ---   ┆ --- ┆ --- ┆ --- ┆ ---   │
# │ u32   ┆ i64 ┆ i64 ┆ i64 ┆ i64   │
# ╞═══════╪═════╪═════╪═════╪═══════╡
# │ 0     ┆ 1   ┆ 5   ┆ 9   ┆ 15    │
# │ 1     ┆ 2   ┆ 6   ┆ 10  ┆ 18    │
# │ 2     ┆ 3   ┆ 7   ┆ 11  ┆ 21    │
# │ 3     ┆ 4   ┆ 8   ┆ 12  ┆ 24    │
# └───────┴─────┴─────┴─────┴───────┘

df.with_columns(
    pl.map_batches(
        list(pl.Series(df.columns).str.extract('^(B_.*)$').drop_nulls()),
        lambda x: x[0] + x[1] + x[2]
    )
    .name.prefix("C_")
)

import polars.selectors as cs

df.with_columns(
    pl.map_batches(
        cs.expand_selector(df, cs.matches("^B_.*")),
        lambda x: x[0] + x[1] + x[2]
    )
    .name.prefix("C_")
)

import polars as pl
import polars.selectors as cs

df = pl.DataFrame({
    "A": [9, 10, 11, 12],
    "B": [1, 2, 3, 4],
}).with_row_index()

df = df.with_columns(df.to_dummies('B'))


def func(yx):
    return yx[0] + yx[1] + yx[2]


def lstsq_1(y, *x):
    return pl.map_batches([y, *x], lambda yx: func(yx))


def lstsq_2(y, *x):
    # it use the `df`, not good
    z = cs.expand_selector(df, x[0])
    return pl.map_batches([y, *z], lambda yx: func(yx))


out = df.with_columns([
    # polars.exceptions.ComputeError: the name: 'resid' passed to `LazyFrame.with_columns` is duplicate
    # lstsq_1(pl.col('A'), cs.matches('^B_.*$')).alias('resid'),
    # not good
    lstsq_2(pl.col('A'), cs.matches("^B_.*$")).alias('resid'),
])
print(out)
"""
shape: (4, 8)
┌───────┬─────┬─────┬─────┬─────┬─────┬─────┬───────┐
│ index ┆ A   ┆ B   ┆ B_1 ┆ B_2 ┆ B_3 ┆ B_4 ┆ resid │
│ ---   ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ ---   │
│ u32   ┆ i64 ┆ i64 ┆ u8  ┆ u8  ┆ u8  ┆ u8  ┆ i64   │
╞═══════╪═════╪═════╪═════╪═════╪═════╪═════╪═══════╡
│ 0     ┆ 9   ┆ 1   ┆ 1   ┆ 0   ┆ 0   ┆ 0   ┆ 10    │
│ 1     ┆ 10  ┆ 2   ┆ 0   ┆ 1   ┆ 0   ┆ 0   ┆ 11    │
│ 2     ┆ 11  ┆ 3   ┆ 0   ┆ 0   ┆ 1   ┆ 0   ┆ 11    │
│ 3     ┆ 12  ┆ 4   ┆ 0   ┆ 0   ┆ 0   ┆ 1   ┆ 12    │
└───────┴─────┴─────┴─────┴─────┴─────┴─────┴───────┘
"""

from functools import reduce

import polars as pl
import polars.selectors as cs

df = pl.DataFrame({
    "A": [9, 10, 11, 12],
    "B": [1, 2, 3, 4],
}).with_row_index()

df = df.with_columns(df.to_dummies('B'))


def func_2(yx):
    y = yx[0]
    x = yx[1:]
    return y - reduce(lambda a, b: a + b, x)


def func_1(yx):
    y = yx[0].struct[0]
    x = list(yx[0].struct)[1:]
    return y - reduce(lambda a, b: a + b, x)


def lstsq_1(y, *x):
    return pl.map_batches(pl.struct([y, *x]), lambda yx: func_1(yx))


def lstsq_2(y, *x):
    # it use the `df`, not good
    z = cs.expand_selector(df, x[0])
    return pl.map_batches([y, *z], lambda yx: func_2(yx))


out = df.with_columns([
    lstsq_1(pl.col('A'), pl.col('^B_.*$')).name.prefix('C_'),
    lstsq_2(pl.col('A'), cs.matches("^B_.*$")).name.prefix('D_'),
])
print(out)
"""
shape: (4, 9)
┌───────┬─────┬─────┬─────┬───┬─────┬─────┬─────┬─────┐
│ index ┆ A   ┆ B   ┆ B_1 ┆ … ┆ B_3 ┆ B_4 ┆ C_A ┆ D_A │
│ ---   ┆ --- ┆ --- ┆ --- ┆   ┆ --- ┆ --- ┆ --- ┆ --- │
│ u32   ┆ i64 ┆ i64 ┆ u8  ┆   ┆ u8  ┆ u8  ┆ i64 ┆ i64 │
╞═══════╪═════╪═════╪═════╪═══╪═════╪═════╪═════╪═════╡
│ 0     ┆ 9   ┆ 1   ┆ 1   ┆ … ┆ 0   ┆ 0   ┆ 8   ┆ 8   │
│ 1     ┆ 10  ┆ 2   ┆ 0   ┆ … ┆ 0   ┆ 0   ┆ 9   ┆ 9   │
│ 2     ┆ 11  ┆ 3   ┆ 0   ┆ … ┆ 1   ┆ 0   ┆ 10  ┆ 10  │
│ 3     ┆ 12  ┆ 4   ┆ 0   ┆ … ┆ 0   ┆ 1   ┆ 11  ┆ 11  │
└───────┴─────┴─────┴─────┴───┴─────┴─────┴─────┴─────┘
"""

df.with_columns(
   pl.col("^B_.*$").map_batches(...)
)

df.with_columns(
   pl.col("B_1").map_batches(...),
   pl.col("B_2").map_batches(...),
   pl.col("B_3").map_batches(...)
)

pl.map_batches([pl.col("B_1"), pl.col("B_2"), pl.col("B_3")], ...