如何合并Polars DataFrame指定列并重复其他列匹配行数

我有一个Polars DataFrame，需要将带后缀_0、_1的列合并成单列：比如a_0、a_1、a_2合并为a列，要求a_0的所有元素在前，接着是a_1的所有元素，以此类推；b_0、b_1、b_2同理合并为b列。同时，所有名称不以_加数字结尾的列（比如words、groups）需要重复足够次数，匹配合并后的列长度。

示例输入代码：

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,
    }
)
print(df)

输入DataFrame输出：

shape: (3, 8)
┌─────┬─────┬─────┬──────────┬──────────┬──────────┬─────────────────┬────────┐
│ a_0 ┆ a_1 ┆ a_2 ┆ b_0      ┆ b_1      ┆ b_2      ┆ words           ┆ groups │
│ --- ┆ --- ┆ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---             ┆ ---    │
│ f64 ┆ f64 ┆ f64 ┆ f64      ┆ f64      ┆ f64      ┆ str             ┆ str    │
╞═════╪═════╪═════╪══════════╪══════════╪══════════╪═════════════════╪════════╡
│ 0.0 ┆ 1.0 ┆ 2.0 ┆ 0.653892 ┆ 0.234362 ┆ 0.880558 ┆ OIww            ┆ W      │
│ 0.5 ┆ 1.5 ┆ 2.5 ┆ 0.408888 ┆ 0.213767 ┆ 1.833025 ┆ KkeB            ┆ Z      │
│ 1.0 ┆ 2.0 ┆ 3.0 ┆ 0.423949 ┆ 0.646378 ┆ 0.116173 ┆ NLOAgRxAtjWOHuQ ┆ O      │
└─────┴─────┴─────┴──────────┴──────────┴──────────┴─────────────────┴────────┘

期望输出：

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 │
└─────────────────┴────────┴─────┴──────────┘

通过重塑DataFrame结构即可实现需求，以下是两种实现方式：

针对固定前缀（如a、b）的实现 #

import polars as pl

# 提取静态列（不以_数字结尾的列）
static_cols = [col for col in df.columns if not col.endswith(("_0", "_1", "_2"))]

# 堆叠a开头的列，拆分索引并命名
a_stack = df.select(pl.col("^a_.*$")).stack().rename({"variable": "idx", "value": "a"})
# 堆叠b开头的列，拆分索引并命名
b_stack = df.select(pl.col("^b_.*$")).stack().rename({"variable": "idx", "value": "b"})

# 合并堆叠后的a、b列，按索引排序保证顺序
stacked = pl.concat([a_stack, b_stack.drop("idx")], how="horizontal").sort("idx")

# 重复静态列3次（对应每个前缀下的列数）
static_repeated = df.select(static_cols).repeat(3)

# 合并静态列与处理后的动态列
result = pl.concat([static_repeated, stacked.drop("idx")], how="horizontal")

print(result)

代码说明 #

1. 静态列筛选：直接通过后缀匹配筛选出不需要合并的列；如果列数多，也可以用正则表达式更灵活匹配。
2. 列堆叠：stack()方法将多列转为variable（原列名）和value（对应值）两列，重命名后方便后续排序。
3. 排序：按原列名排序，确保a_0的所有值在前，接着是a_1、a_2。
4. 重复静态列：repeat()方法让静态列重复对应次数，保证行数与堆叠后的动态列一致。
5. 合并结果：横向拼接静态列和动态列，得到最终格式的DataFrame。

自动识别前缀的通用实现 #

如果前缀不止a和b，可以自动识别所有前缀批量处理：

import polars as pl
import re

# 提取所有带数字后缀列的前缀
prefixes = list({re.match(r"(.*)_\d+$", col).group(1) for col in df.columns if re.match(r".*_\d+$", col)})
# 提取静态列
static_cols = [col for col in df.columns if not re.match(r".*_\d+$", col)]

# 批量处理每个前缀的列
stacked_dfs = []
for prefix in prefixes:
    stack_df = df.select(pl.col(f"^{prefix}_.*$")).stack()
    stack_df = stack_df.rename({"value": prefix})
    stacked_dfs.append(stack_df)

# 合并所有堆叠后的列，按原列名排序
combined_stacked = pl.concat(stacked_dfs, how="horizontal").sort("variable")

# 获取每个前缀下的列数，重复静态列
m = len([col for col in df.columns if col.startswith(f"{prefixes[0]}_")])
static_repeated = df.select(static_cols).repeat(m)

# 生成最终结果
result = pl.concat([static_repeated, combined_stacked.drop("variable")], how="horizontal")
print(result)

内容的提问来源于Stack Exchange，提问作者DeltaIV