Я пытаюсь сделать в Polars LazyFrame то, что я умею делать с Pandas.
Здесь я могу применить пользовательскую функцию к нескольким столбцам, и эта функция использует другой столбец в качестве второстепенный аргумент.
import statsmodels.stats.proportion as ssp
import pandas as pd
df = pd.DataFrame(
{
"a": [0, 1, 0, 1],
"b": [1, 2, 3, 4],
"nobs": [10, 20, 30, 40],
}
)
proportion_confint_func = lambda c: ssp.proportion_confint(c, df["nobs"], method="agresti_coull")[0]
df[["a", "b"]].apply(proportion_confint_func)
a b
0 0.0 0.000000
1 0.0 0.015656
2 0.0 0.026639
3 0.0 0.033880
Как мне это сделать на Polars LazyFrame? Моя попытка использовать map_batches
from functools import partial
from scipy import stats
import polars as pl
def proportion_confint_pl(count_a, nobs_a, alpha:float=0.05, method="agresti_coull", bound="lower"):
crit = stats.norm.isf(alpha / 2.0)
nobs_c = nobs_a + crit**2
q_c = (count_a + crit**2 / 2.0) / nobs_c
std_c = ((q_c * (1.0 - q_c) / nobs_c)).sqrt()
dist = crit * std_c
if bound == "lower":
ci = (q_c - dist)
elif bound == "upper":
ci = (q_c + dist)
return ci.clip(0, 1)
ldf = pl.LazyFrame(
{
"a": [0, 1, 0, 1],
"b": [1, 2, 3, 4],
"nobs": [10, 20, 30, 40],
}
)
proportion_confint_pl_func = partial(proportion_confint_pl, nobs_a=pl.col("nobs"))
ldf.select(pl.col(["a", "b"]).map_batches(proportion_confint_pl_func)).collect()
Но он дает мне столбцы объекта.
shape: (1, 2)
┌─────────────────────────────────┬─────────────────────────────────┐
│ a ┆ b │
│ --- ┆ --- │
│ object ┆ object │
╞═════════════════════════════════╪═════════════════════════════════╡
│ [([(Series[a]) / ([(col("nobs"… ┆ [([(Series) / ([(col("nobs"… │
└─────────────────────────────────┴─────────────────────────────────┘
Подробнее здесь: https://stackoverflow.com/questions/790 ... an-argumen