Polars 中按前缀合并列并重复非模式列的完整教程

2026-01-30 Python

本文介绍如何在 polars 中将具有相同前缀（如 `a_0`, `a_1`, `a_2`）的多列纵向堆叠为单列（如 `a`），同时自动复制其他非模式列（如 `words`, `groups`）以匹配扩展后的行数。

在数据预处理中，常需将“宽格式”的分组列（如 a_0, a_1, a_2 和 b_0, b_1, b_2）转换为“长格式”——即把每个前缀对应的所有列值垂直拼接（[a_0; a_1; a_2]），而非横向聚合。与此同时，其余标识性列（如 words、groups）需按需重复，确保每条新记录仍能正确关联原始语义。这本质上是 unpivot + 分组索引 + pivot 的组合操作，关键在于构造稳定的行序关系。

以下为完整实现方案（基于 Polars 0.20+）：

import polars as pl
import numpy as np
import string

# 构造示例数据（同问题）
rng = np.random.default_rng(42)
nr = 3
letters = list(string.ascii_letters)
uppercase = list(string.ascii_uppercase)
words, groups = [], []
for i in range(nr):
    word = ''.join([rng.choice(letters) for _ in range(rng.integers(3, 20))])
    words.append(word)
    group = rng.choice(uppercase)
    groups.append(group)
df = pl.DataFrame({
    "a_0": np.linspace(0, 1, nr),
    "a_1": np.linspace(1, 2, nr),
    "a_2": np.linspace(2, 3, nr),
    "b_0": np.random.rand(nr),
    "b_1": 2 * np.random.rand(nr),
    "b_2": 3 * np.random.rand(nr),
    "words": words,
    "groups": groups,
})

# ✅ 核心转换：三步法
result = (
    df
    .unpivot(
        index=["words", "groups"],          # 指定需保留并重复的列（不参与展开）
        on=[col for col in df.columns if "_" in col and col.split("_")[-1].isdigit()]  # 显式指定要展开的列（更安全）
    )
    .with_columns(
        pl.col("variable").str.replace(r"_\d+$", "")  # 提取前缀：a_0 → "a"，b_1 → "b"
    )
    .with_columns(
        index=pl.int_range(0, pl.len()).over("variable")  # 每个前缀组内独立编号：a组0,1,2；b组0,1,2
    )
    .pivot(
        on="variable",
        index=["index", "words", "groups"],
        values="value",
        aggregate_function=None  # 禁用聚合，确保一一映射
    )
    .drop("index")  # 移除辅助索引列
)

print(result)

输出结果与预期一致：

Grokipedia

xAI推出的AI在线百科全书

shape: (9, 4)
┌─────────────────┬────────┬─────┬──────────┐
│ words           ┆ groups ┆ a   ┆ b        │
│ ---             ┆ ---    ┆ --- ┆ ---      │
│ str             ┆ str    ┆ f64 ┆ f64      │
╞═════════════════╪════════╪═════╪══════════╡
│ OIww            ┆ W      ┆ 0.0 ┆ 0.653892 │
│ KkeB            ┆ Z      ┆ 0.5 ┆ 0.408888 │
│ NLOAgRxAtjWOHuQ ┆ O      ┆ 1.0 ┆ 0.423949 │
│ OIww            ┆ W      ┆ 1.0 ┆ 0.234362 │
│ KkeB            ┆ Z      ┆ 1.5 ┆ 0.213767 │
│ NLOAgRxAtjWOHuQ ┆ O      ┆ 2.0 ┆ 0.646378 │
│ OIww            ┆ W      ┆ 2.0 ┆ 0.880558 │
│ KkeB            ┆ Z      ┆ 2.5 ┆ 1.833025 │
│ NLOAgRxAtjWOHuQ ┆ O      ┆ 3.0 ┆ 0.116173 │
└─────────────────┴────────┴─────┴──────────┘

关键要点说明：

unpivot(index=...) 是基石：它将所有非 index 列转为两列（variable, value），同时自动复制 index 列内容到每一行，为后续对齐奠定基础。
正则提取前缀：使用 str.replace(r"_\d+$", "") 安全剥离末尾 _数字，避免误伤含下划线的其他列名（如 user_id_2024）。
over("variable") 构造组内序号：确保 a_0, a_1, a_2 的值在堆叠后仍保持原始行序，并与 b_0, b_1, b_2 严格对齐。
pivot(..., index=[...])：以 index（组内序号）和标识列（words, groups）共同作为行键，保证 pivot 后每行唯一且语义完整。
显式指定 on= 列（推荐）：比依赖字符串匹配更鲁棒，尤其当数据中存在非模式列（如 c_extra）时可避免误卷积。