如何在scikit-learn中直接使用Polars DataFrames进行机器学习训练？

1. 利用scikit-learn 1.3+的原生Polars支持 #

从scikit-learn 1.3版本开始，大部分估算器（estimator）已经原生支持Polars DataFrames作为输入，不需要手动转换为Pandas。你可以直接将Polars DataFrame传入fit()、predict()等方法，底层会自动处理数据转换（不会生成中间Pandas对象）。

示例代码：

import polars as pl
from sklearn.linear_model import LinearRegression

# 创建Polars DataFrame
df = pl.DataFrame({
    "feature1": [1, 2, 3, 4, 5],
    "feature2": [2, 4, 6, 8, 10],
    "target": [3, 6, 9, 12, 15]
})

# 初始化模型并直接使用Polars DataFrame训练
model = LinearRegression()
model.fit(df[["feature1", "feature2"]], df["target"])

# 直接用Polars DataFrame预测
predictions = model.predict(df[["feature1", "feature2"]])

2. 自定义Polars兼容的预处理转换器 #

如果需要在Pipeline中使用Polars做自定义预处理（比如特征选择、自定义特征工程），可以继承sklearn的BaseEstimator和TransformerMixin实现兼容Polars的转换器，全程不用转Pandas。

示例代码：

import polars as pl
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import Pipeline
from sklearn.ensemble import RandomForestRegressor

class PolarsFeatureSelector(BaseEstimator, TransformerMixin):
    def __init__(self, feature_cols):
        self.feature_cols = feature_cols
    
    def fit(self, X, y=None):
        # Polars DataFrame不需要额外fit操作
        return self
    
    def transform(self, X):
        # 返回Polars DataFrame，后续步骤sklearn会自动处理
        return X.select(self.feature_cols)

# 构建包含Polars转换器的Pipeline
pipe = Pipeline([
    ("select_features", PolarsFeatureSelector(feature_cols=["feature1", "feature2"])),
    ("model", RandomForestRegressor())
])

# 直接用Polars DataFrame训练Pipeline
pipe.fit(df, df["target"])

# 预测
pipe.predict(df)

3. 直接转换为Numpy数组（无需Pandas） #

如果你的scikit-learn版本较低（低于1.3），可以直接将Polars DataFrame的特征列转换为Numpy数组传递给模型，这比转Pandas更高效，也不需要依赖Pandas。

示例代码：

import polars as pl
from sklearn.ensemble import RandomForestClassifier

df = pl.DataFrame({
    "f1": [1, 2, 3, 4, 5],
    "f2": [0, 1, 0, 1, 0],
    "label": [0, 1, 0, 1, 0]
})

# 提取特征和目标，转换为Numpy数组
X = df.select(["f1", "f2"]).to_numpy()
y = df.select("label").to_numpy().ravel()  # ravel()将二维数组转为一维

# 训练模型
model = RandomForestClassifier()
model.fit(X, y)

注意事项 #

• 确保你的scikit-learn版本≥1.3，原生支持Polars的覆盖范围最广，包括大部分线性模型、树模型、聚类算法等。
• 对于极少数不支持Polars的sklearn组件，可以结合上述自定义转换器或Numpy转换的方式解决。

内容的提问来源于stack exchange，提问作者Regular Tech Guy