适配Polars与特征簇的Scikit-Learn置换重要性改写难题

我正在改写scikit-learn的置换重要性源码，目标是实现两个功能：

1. 兼容Polars数据框
2. 支持特征簇的置换重要性计算

测试代码 #

import polars as pl
import polars.selectors as cs
import numpy as np

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

X, y = make_classification(
    n_samples=1000,
    n_features=10,
    n_informative=3,
    n_redundant=0,
    n_repeated=0,
    n_classes=2,
    random_state=42,
    shuffle=False,
)
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=42)
feature_names = [f"feature_{i}" for i in range(X.shape[1])]

X_train_polars = pl.DataFrame(X_train, schema=feature_names)
X_test_polars = pl.DataFrame(X_test, schema=feature_names)
y_train_polars = pl.Series(y_train, schema=["target"])
y_test_polars = pl.Series(y_test, schema=["target"])

遇到的问题 #

计算特征簇重要性时，需要同时置换整个特征簇的所有列。置换后得到的shuffle_arr是(n_samples, 簇内特征数)的多维数组，无法直接用replace_column单步替换多列：

from sklearn.utils import check_random_state
random_state = check_random_state(42)
random_seed = random_state.randint(np.iinfo(np.int32).max + 1)

X_train_permuted = X_train_polars.clone()
shuffle_arr = np.array(X_train_permuted[:, ["feature_0", "feature_1"]])

random_state.shuffle(shuffle_arr)
# 这里无法直接用replace_column替换多列，会报错
X_train_permuted.replace_column(0, pl.Series(name="feature_0", values=shuffle_arr))

解决方案 #

方法1：用with_columns批量替换（推荐，非原地操作）

Polars推荐非原地操作，with_columns可以一次性传入多个Series完成批量替换，代码简洁且高效：

from sklearn.utils import check_random_state
random_state = check_random_state(42)

# 定义要置换的特征簇列名
cluster_cols = ["feature_0", "feature_1"]
# 提取簇内数据并置换
cluster_data = np.array(X_train_polars.select(cluster_cols))
random_state.shuffle(cluster_data)

# 将置换后的多维数组拆分为对应列的Series
permuted_series = [
    pl.Series(name=col, values=cluster_data[:, idx])
    for idx, col in enumerate(cluster_cols)
]

# 批量替换列，返回新的DataFrame
X_train_permuted = X_train_polars.with_columns(permuted_series)

方法2：循环调用replace_column实现原地替换

如果需要原地修改原DataFrame对象，可以循环遍历簇内列，逐个替换：

from sklearn.utils import check_random_state
random_state = check_random_state(42)

cluster_cols = ["feature_0", "feature_1"]
cluster_data = np.array(X_train_polars.select(cluster_cols))
random_state.shuffle(cluster_data)

X_train_permuted = X_train_polars.clone()
for idx, col in enumerate(cluster_cols):
    # 获取目标列在DataFrame中的索引位置
    col_position = X_train_permuted.columns.index(col)
    # 逐个替换列
    X_train_permuted.replace_column(col_position, pl.Series(name=col, values=cluster_data[:, idx]))

方法3：LazyFrame优化大数据场景

如果处理的是超大数据集，用Polars的懒加载模式（LazyFrame）可以更高效地管理内存：

from sklearn.utils import check_random_state
random_state = check_random_state(42)

cluster_cols = ["feature_0", "feature_1"]
cluster_data = np.array(X_train_polars.select(cluster_cols))
random_state.shuffle(cluster_data)

# 将置换后的簇数据转为临时DataFrame
permuted_cluster_df = pl.DataFrame(cluster_data, schema=cluster_cols)

# 用懒加载链式操作替换原列，保持原列顺序
X_train_permuted = (
    X_train_polars.lazy()
    .drop(cluster_cols)
    .with_columns(permuted_cluster_df.lazy())
    .select(X_train_polars.columns)  # 确保列顺序和原DataFrame一致
    .collect()
)

方法对比 #

• 方法1：符合Polars的函数式设计理念，无副作用，代码最简洁，优先推荐。
• 方法2：适合必须原地修改的场景，但簇列数较多时效率略低。
• 方法3：针对大数据集优化，减少内存占用，适合超大规模数据处理。

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