优化Polars扁平DataFrame按表拆分的高效实现方案

需求背景 #

我们有一个扁平结构的Polars DataFrame，用于存储多数据库的结构化数据，包含以下字段：

• db_id：数据库唯一标识
• table_id：表唯一标识（不同数据库可共享相同表）
• column_id：列唯一标识（不同数据库可共享相同列）
• data：存储多值的列表字段

核心目标是按table_id拆分出独立的DataFrame：每个结果DataFrame的列由对应表的所有column_id加上db_id组成，data字段中的元素需要展开后对应到列的位置，完成转置与展开操作。

各版本实现与性能对比 #

初始实现 #

直接按table_id分组后进行透视+展开，Python循环处理每个分组，开销较大：

def initial_split(df: pl.DataFrame) -> dict[str, pl.DataFrame]:
    table_dfs = {}
    for table_id, group in df.group_by("table_id"):
        # 透视后展开所有列的列表值
        pivoted = group.pivot(index="db_id", columns="column_id", values="data")
        exploded = pivoted.explode(pl.exclude("db_id"))
        table_dfs[table_id] = exploded
    return table_dfs

性能：处理250万行数据耗时70分钟

第一次优化 #

先全局展开data字段，再分组透视，减少分组内的展开操作开销：

def optimized_split(df: pl.DataFrame) -> dict[str, pl.DataFrame]:
    # 先展开所有data列表，避免分组内重复操作
    exploded_df = df.explode("data")
    table_dfs = {}
    for table_id, group in exploded_df.group_by("table_id"):
        # 直接透视得到结构化结果
        pivoted = group.pivot(index="db_id", columns="column_id", values="data")
        table_dfs[table_id] = pivoted
    return table_dfs

性能：处理250万行数据耗时10分钟

最终高效方案 #

利用Polars的矢量化map_groups替代Python循环，完全在Polars内部处理分组逻辑，最大化性能：

def high_performance_split(df: pl.DataFrame) -> dict[str, pl.DataFrame]:
    # 全局展开data
    exploded_df = df.explode("data")
    # 分组后批量执行透视，返回字典格式
    grouped_result = exploded_df.group_by("table_id").map_groups(
        lambda group: group.pivot(index="db_id", columns="column_id", values="data")
    )
    # 转换为table_id到DataFrame的映射
    return {tid: df for tid, df in grouped_result.items()}

性能：处理250万行数据仅需50秒

性能测试结果 #

测试数据集：250万行扁平结构Polars DataFrame

• 初始实现：70分钟
• 第一次优化：10分钟
• 最终高效方案：50秒

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