如何用Polars实现类似R中pivot_longer的宽表转长表操作

原R代码及需求 #

需要复刻的R代码功能为将宽格式数据转换为长格式，提取年份与生育率数据：

wide_data <- read_csv('https://raw.githubusercontent.com/rafalab/dslabs/master/inst/extdata/life-expectancy-and-fertility-two-countries-example.csv')

new_tidy_data <- pivot_longer(wide_data, `1960`:`2015`, names_to = "year", values_to = "fertility")

数据集结构说明：

• 共113列，第一列为country
• 后续列以{年份}_{指标}命名，如1960_fertility、1960_life_expectancy，覆盖1960至2015年
• 仅包含德国、韩国两行数据

期望输出 #

转换后的数据格式示例（取前6行）：

head(new_tidy_data)
#> # A tibble: 6 × 3
#>   country year  fertility
#>   <chr>   <chr>     <dbl>
#> 1 Germany 1960       2.41
#> 2 Germany 1961       2.44
#> 3 Germany 1962       2.47
#> 4 Germany 1963       2.49
#> 5 Germany 1964       2.49
#> # ℹ 1 more row

现有Polars代码 #

当前未完成功能的Polars代码：

import polars as pl
import polars.selectors as cs

df = pl.read_csv('https://raw.githubusercontent.com/rafalab/dslabs/master/inst/extdata/life-expectancy-and-fertility-two-countries-example.csv')
df.pivot()

正确的Polars实现 #

Polars中用melt函数对应R的pivot_longer，结合字符串拆分提取年份与指标，最终筛选出生率数据：

import polars as pl

# 读取数据
df = pl.read_csv('https://raw.githubusercontent.com/rafalab/dslabs/master/inst/extdata/life-expectancy-and-fertility-two-countries-example.csv')

# 转换为长格式，保留country作为标识列
melted_df = df.melt(id_vars="country", variable_name="year_indicator")

# 拆分年份与指标列
result_df = melted_df.with_columns(
    pl.col("year_indicator").str.split("_").map_elements(lambda x: x[0]).alias("year"),
    pl.col("year_indicator").str.split("_").map_elements(lambda x: x[1]).alias("indicator")
)

# 筛选生育率数据，整理列顺序并排序
final_df = result_df.filter(pl.col("indicator") == "fertility") \
                  .select("country", "year", pl.col("value").alias("fertility")) \
                  .sort("country", "year")

# 查看前6行结果
print(final_df.head(6))

执行后输出结果与期望一致：

shape: (6, 3)
┌─────────┬──────┬───────────┐
│ country ┆ year ┆ fertility │
│ ---     ┆ ---  ┆ ---       │
│ str     ┆ str  ┆ f64       │
╞═════════╪══════╪═══════════╡
│ Germany ┆ 1960 ┆ 2.41      │
│ Germany ┆ 1961 ┆ 2.44      │
│ Germany ┆ 1962 ┆ 2.47      │
│ Germany ┆ 1963 ┆ 2.49      │
│ Germany ┆ 1964 ┆ 2.49      │
│ Germany ┆ 1965 ┆ 2.46      │
└─────────┴──────┴───────────┘

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