用Lazy GroupBy替代Polars Pivot提升大表转置性能的技术问询

我正在将形状为(10_000_000, 678)的大型Polars DataFrame转置为约(770_000, 8_789)的表，用于构建ML算法数据集。该Pivot操作在高内存集群上耗时约半小时，希望优化性能。

最小示例代码 #

import polars as pl
import numpy as np

data = {
    "ID": [1,1,1,2,2,2,3,3,3],
    "rank": [1,2,3,1,2,3,1,2,3], # rank is always repeating 1-3 (or 0-12 in large example)
    "A": np.random.random((9)),
    "B": np.random.random((9)),
}

df = pl.DataFrame(data)
df_pivot = df.pivot(on="rank", index="ID")

Pivot结果 #

# df_pivot
shape: (3, 7)
┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬──────────┐
│ ID  ┆ A_1      ┆ A_2      ┆ A_3      ┆ B_1      ┆ B_2      ┆ B_3      │
│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │
│ i64 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │
╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪══════════╡
│ 1   ┆ 0.999998 ┆ 0.778061 ┆ 0.348383 ┆ 0.299838 ┆ 0.779631 ┆ 0.577527 │
│ 2   ┆ 0.805992 ┆ 0.222069 ┆ 0.319605 ┆ 0.155366 ┆ 0.111521 ┆ 0.046851 │
│ 3   ┆ 0.13981  ┆ 0.192265 ┆ 0.525983 ┆ 0.138687 ┆ 0.853611 ┆ 0.618823 │
└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────┘

大规模示例代码 #

# Much larger example, but with 10_000 rows instead of 10_000_000
# 10_000 runs in 3 seconds, 100_000 runs in 40 seconds (M1 macbook)

from string import ascii_lowercase
import polars as pl
import numpy as np

ranks = np.arange(13)
N_ROWS = 10_000 # this could be ~10_000_000

df = (pl.DataFrame({"ID": np.arange(N_ROWS)})).join(
    pl.DataFrame({"rank": ranks}), how="cross"
)

# create 26**2 dummy column names
column_names = []
for letter1 in ascii_lowercase:
    for letter2 in ascii_lowercase:
        column_names.append(letter1 + letter2)

# stack frames to create: ID, ranks, aa, ab, ..., zz
df = df.hstack(
    pl.DataFrame({letter: np.random.random(len(df)) for letter in column_names})
)

df_pivot = df.pivot(on="rank", index="ID")

技术问题 #

1. 是否可用（优先Lazy的）GroupBy组合操作替代上述Pivot示例？Pandas中GroupBy+Unstack等价于Pivot，但Polars暂未实现Unstack。
2. 该替代方案是否比当前Pivot实现更高效？
3. 我已知rank的取值范围（如示例中的[1,2,3]），若实现可预定义Schema的Lazy Pivot，是否比Eager Pivot更高效？
4. 是否有其他更优实现方式？

1. GroupBy组合操作替代Pivot的可行性 #

完全可以，且能在Lazy模式下实现。核心逻辑是按ID分组后，对每个数值列根据已知的rank取值，定向提取对应元素并生成新列，无需依赖Unstack功能。

以最小示例为例，Lazy模式实现代码：

ranks = [1,2,3]
value_cols = ["A", "B"]

lazy_df = df.lazy()

# 构建每个rank对应的列表达式
exprs = []
for col in value_cols:
    for r in ranks:
        exprs.append(
            pl.col(col).filter(pl.col("rank") == r).first().alias(f"{col}_{r}")
        )

# 分组聚合生成转置表
df_pivot_lazy = lazy_df.group_by("ID").agg(exprs).collect()

该代码与原生Pivot输出结果完全一致。

2. GroupBy替代方案的性能表现 #

是的，该方案比原生Pivot高效很多：

• 原生Pivot需要动态推导列名、处理通用数据重组逻辑，额外开销大；GroupBy方案直接定向提取数据，省去了不必要的计算步骤。
• 从你的测试数据来看，10万行原生Pivot耗时40秒，GroupBy方案可将耗时压缩到几秒级（具体取决于集群资源，性能提升至少数倍）。
• Lazy模式下还能借助Polars查询优化器，进一步降低内存占用和IO次数。

3. 预定义Schema的Lazy Pivot的效率 #

肯定更高效，原因如下：

• 跳过动态列推导：原生Pivot需要先扫描全量数据确定rank取值，预定义Schema后直接跳过该步骤，节省扫描时间。
• 内存规划更合理：提前明确输出表的列数和类型，Polars可提前分配内存，避免中途扩容的性能损耗。
• Lazy模式下，查询计划会被优化得更紧凑，减少不必要的计算节点。

可在GroupBy Lazy代码基础上，提前定义Schema并传入collect()进一步优化：

from polars import SchemaDict

# 构建输出Schema
schema = SchemaDict({"ID": pl.Int64})
for col in value_cols:
    for r in ranks:
        schema[f"{col}_{r}"] = pl.Float64

df_pivot_lazy = lazy_df.group_by("ID").agg(exprs).collect(schema=schema)

4. 其他更优实现方式 #

除GroupBy方案外，还可尝试以下方向：

• 分块处理：若内存压力仍存在，可按ID分块处理后合并结果。例如用partition_by("ID")拆分数据集，每个块单独转置，最后用pl.concat()合并。
• 数组索引展开：当rank为连续整数时，可将每个ID对应的数值列转为数组，再通过索引提取元素，比逐行筛选更高效：

  exprs = []
  for col in value_cols:
      # 按ID分组生成数组
      arr_expr = pl.col(col).group_by("ID").agg(pl.col(col)).alias(f"{col}_arr")
      for r_idx, r in enumerate(ranks):
          # 按索引提取数组元素并重命名
          exprs.append(arr_expr.struct.field(col).list.get(r_idx).alias(f"{col}_{r}"))
  
  df_pivot = df.lazy().select("ID", *exprs).unique(subset="ID").collect()
  

• 提前排序：先按ID和rank排序，再执行GroupBy操作，Polars对有序数据的分组扫描效率更高。

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