如何在Polars中执行计算并返回非DataFrame复杂对象

1. 全局模型训练（无分组） #

如果是对整个数据集训练模型，直接提取数据传给sklearn是最高效的方式——Polars的惰性模式核心优化列级延迟计算，而模型训练需要完整的样本矩阵，没必要强行嵌入惰性表达式链。原代码可优化数据转换逻辑，避免类型兼容问题：

import polars as pl
from sklearn.linear_model import ElasticNetCV
from sklearn import datasets

iris = datasets.load_iris()

dt = pl.DataFrame(iris["data"], columns=iris.feature_names)
dt = dt.with_columns(Species=pl.Series(iris.target))

X_cols = ["sepal length (cm)", "sepal width (cm)", "petal length (cm)"]
y_col = "petal width (cm)"

regr = ElasticNetCV(cv=5, random_state=0)
# 显式转换为numpy数组，适配sklearn输入要求
res = regr.fit(dt.select(X_cols).to_numpy(), dt.select(y_col).to_numpy().ravel())

若要统一Polars表达式风格（比如串联其他列操作），可通过pl.struct+map_elements包装训练逻辑，惰性模式下需调用collect()触发计算：

lazy_dt = dt.lazy()

model_expr = pl.struct(X_cols + [y_col]).map_elements(
    lambda s: ElasticNetCV(cv=5, random_state=0).fit(
        s[X_cols].to_numpy(),
        s[y_col].to_numpy().ravel()
    ),
    return_dtype=pl.Object
)

# 执行并获取训练好的模型
trained_model = lazy_dt.select(model_expr).collect().item()

2. 分组模型训练（对应data.table分组建模逻辑） #

如果需要按某列（如Species）分组训练模型，Polars的map_groups是等价于data.table分组操作的最佳方案，支持惰性模式延迟执行：

lazy_dt = dt.lazy()

grouped_models = (
    lazy_dt
    .group_by("Species")
    .agg(
        pl.struct(X_cols + [y_col]).map_elements(
            lambda s: ElasticNetCV(cv=5, random_state=0).fit(
                s[X_cols].to_numpy(),
                s[y_col].to_numpy().ravel()
            ),
            return_dtype=pl.Object
        ).alias("trained_model")
    )
    .collect()
)

# 查看分组结果
print(grouped_models)

关键注意点 #

• Polars惰性模式不适合强制包装全局模型训练逻辑，直接提取数据的效率远高于表达式嵌入。
• 分组场景下map_groups是标准解决方案，能在collect()阶段才执行分组内的模型训练。
• 存储复杂模型对象时，必须指定return_dtype=pl.Object，因为Polars无原生模型类型支持。

内容的提问来源于stack exchange，提问作者Klumpi