如何在Polars中实现Pandas的wide_to_long功能？

Polars 可以通过 unpivot 结合字符串提取、透视操作，实现和 Pandas wide_to_long 完全一致的效果，不管是单个存根名还是多个存根名的场景都能支持。

1. 单个存根名场景（对应 Pandas 示例） #

直接熔解目标列，通过正则提取后缀作为新列：

import polars as pl

df = pl.DataFrame({
    'famid': [1, 1, 1, 2, 2, 2, 3, 3, 3],
    'birth': [1, 2, 3, 1, 2, 3, 1, 2, 3],
    'ht_one': [2.8, 2.9, 2.2, 2, 1.8, 1.9, 2.2, 2.3, 2.1],
    'ht_two': [3.4, 3.8, 2.9, 3.2, 2.8, 2.4, 3.3, 3.4, 2.9]
})

pl_long_df = (
    df
    # 熔解ht开头的列，保留索引列
    .unpivot(index=['famid', 'birth'], variable_name='col', value_name='ht')
    # 提取列名中的后缀作为age列
    .with_columns(age=pl.col('col').str.extract(r'ht_(\w+)', 1))
    .drop('col')
    # 调整列顺序匹配Pandas输出
    .select(['famid', 'birth', 'age', 'ht'])
)

print(pl_long_df)

2. 多个存根名场景（你的调查数据示例） #

针对Familiarity和Affinity这类多存根，先拆分存根名和循环编号，再透视得到独立列：

import polars as pl
import numpy as np

# 生成测试数据
age_groups = np.random.choice(['0-18', '19-35', '36-50', '51-65', '65+'], size=10)
genders = np.random.choice(['Male', 'Female', 'Other'], size=10)
fam_aff = np.random.rand(10, 4)
cols = ['Age_group', 'Gender', 'Familiarity_loop1', 'Familiarity_loop2', 'Affinity_loop1', 'Affinity_loop2']
data = np.column_stack([age_groups, genders, fam_aff])
df = pl.DataFrame(data=data, columns=cols)
df = df.with_columns(unique_records=pl.arange(0, len(df)))

pl_result = (
    df
    # 熔解所有带_loop的列，保留非目标列作为索引
    .unpivot(index=['Age_group', 'Gender', 'unique_records'], variable_name='col', value_name='value')
    # 从列名中提取存根名和循环编号
    .with_columns(
        stubname=pl.col('col').str.extract(r'(\w+)_loop\d+', 1),
        loop=pl.col('col').str.extract(r'\w+_(loop\d+)', 1)
    )
    .drop('col')
    # 透视将存根名转为独立列
    .pivot(
        index=['Age_group', 'Gender', 'unique_records', 'loop'],
        columns='stubname',
        values='value',
        aggregate_function=pl.element()
    )
    # 调整列顺序匹配预期输出
    .select(['Age_group', 'Gender', 'unique_records', 'loop', 'Familiarity', 'Affinity'])
)

print(pl_result)

核心要点 #

• 用 str.extract 精准拆分存根名和后缀，对应 Pandas 的 sep 和 suffix 参数
• 多存根场景下，先熔解所有目标列，提取存根和标识后再透视，即可得到和 wide_to_long 传入存根列表一致的结果
• pl.element() 作为聚合函数，确保单个分组值直接保留，避免不必要的聚合

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