Dask DataFrame按索引选列报错Too many indexers求助

问题场景 #

我有一个Dask DataFrame ddf，其中ddf['X']是存储矩阵的Series列，还有一个索引列表indices，想要按这些索引选择矩阵的特征（列）。尝试了以下代码：

初始实现代码 #

def subselect_variables(df):
    subset = df.iloc[:, indices]
    return subset
ddf_X = (
        ddf['X']
        .map_partitions(subselect_variables, meta={col: 'f4'for col in range(len(indices))})
    )
ddf_X.to_parquet(
    my_path,
    engine='pyarrow',
    schema=my_schema,
    write_metadata_file=True,
    row_group_size=my_row_group_size
    )

也尝试过简化版本：

ddf_X = (
        ddf['X']
        .map_partitions(lambda df: df.iloc[:, indices], meta={col: 'f4'for col in range(len(indices))})
    )

替换:为slice(None)后仍然触发相同错误：pandas.errors.IndexingError: Too many indexers

问题根源 #

ddf['X']是单个Series列，每个元素是矩阵/数组对象，并非多列的DataFrame。直接对这个Series使用iloc[:, indices]会报错，因为iloc针对Series仅支持一维索引器，二维索引器会触发"Too many indexers"错误。

解决方案 #

方案1：对Series的每个元素单独处理（保留数组格式） #

通过apply对每个矩阵元素执行列选择，同时指定正确的元数据：

# 元数据格式说明：('列名', '数据类型', (输出数组长度,))
ddf_X = ddf['X'].apply(lambda mat: mat[:, indices], meta=('X', 'f4', (len(indices),)))

或者用map_partitions批量处理每个分区的Series：

def subselect_variables(series):
    return series.apply(lambda mat: mat[:, indices])

ddf_X = ddf['X'].map_partitions(subselect_variables, meta=('X', 'f4', (len(indices),)))

方案2：转换为Dask Array后选择列，再转回DataFrame（展开为单独列） #

如果希望将选择后的每个特征作为DataFrame的独立列，可以先转成Dask Array操作：

# 将Series转换为Dask Array（lengths=True处理变长数组，若矩阵长度固定可省略）
x_array = ddf['X'].to_dask_array(lengths=True)
# 按索引选择列
x_array_subset = x_array[:, indices]
# 转换为Dask DataFrame，指定列名
ddf_X = x_array_subset.to_dask_dataframe(columns=[f'feature_{i}' for i in range(len(indices))])

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