使用polars.pivot()转换数据框（类似R中的pivot_longer）

Question

我之前在用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列：第一列是国家名，接下来是1960年的生育率、1960年的预期寿命、1961年的生育率、1961年的预期寿命……一直到2015年的生育率和预期寿命。

而且数据中有两行，分别是德国和韩国。

我期望的结果是：

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

到目前为止，我的代码是这样的：

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()

谢谢！！

Answer 1

在polars中，你可以使用pivot这个功能把数据变得更宽，而melt（想象一下冰柱）则是把数据变得更长。

不过，melt不会把单个列中的内容用分隔符拆分成两个列，这个你得自己来做。

这看起来像这样……

(
    df
    .melt('country', value_name='fertility')
    .with_columns(
        pl.col('variable').str.splitn('_',2).struct.rename_fields(['year','var'])
    )
    .unnest('variable')
    .filter(pl.col('var')=='fertility')
    .drop('var')
    .sort('country')
)

polars的表达式工作方式是，对于每一个输入，只有一个输出，但有一种叫做struct的类型，可以嵌套多个列。这样我们就可以用splitn把变量列拆分成年份部分和变量部分。我们可以用unnest把这个结构转换成两个普通的列，这个操作是在数据框（df）层面进行的，而不是作为一个表达式。

再想想，既然你想过滤出只有生育率的数据，你可以在melt之前先过滤和重命名列，像这样：

(
    df
    .select(
        pl.col("^(country|.+fertility)$")
        .name.map(lambda x: x.replace("_fertility", "")))
    .melt('country', variable_name='year', value_name='fertility')
    .sort('country')
)
shape: (112, 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      │
│ …           ┆ …    ┆ …         │
│ South Korea ┆ 2011 ┆ 1.29      │
│ South Korea ┆ 2012 ┆ 1.3       │
│ South Korea ┆ 2013 ┆ 1.32      │
│ South Korea ┆ 2014 ┆ 1.34      │
│ South Korea ┆ 2015 ┆ 1.36      │
└─────────────┴──────┴───────────┘

最后

如果你想要一个包含生育率和预期寿命的列，你需要把前面的方法和最后的pivot结合起来，像这样：

(
    df
    .melt('country')
    .with_columns(
        pl.col('variable').str.splitn('_',2).struct.rename_fields(['year','var'])
    )
    .unnest('variable')
    .pivot(
        values='value', index=['country','year'], 
        columns='var', aggregate_function='first')
    .sort('country')
shape: (112, 4)
┌─────────────┬──────┬───────────┬─────────────────┐
│ country     ┆ year ┆ fertility ┆ life_expectancy │
│ ---         ┆ ---  ┆ ---       ┆ ---             │
│ str         ┆ str  ┆ f64       ┆ f64             │
╞═════════════╪══════╪═══════════╪═════════════════╡
│ Germany     ┆ 1960 ┆ 2.41      ┆ 69.26           │
│ Germany     ┆ 1961 ┆ 2.44      ┆ 69.85           │
│ Germany     ┆ 1962 ┆ 2.47      ┆ 70.01           │
│ Germany     ┆ 1963 ┆ 2.49      ┆ 70.1            │
│ Germany     ┆ 1964 ┆ 2.49      ┆ 70.66           │
│ …           ┆ …    ┆ …         ┆ …               │
│ South Korea ┆ 2011 ┆ 1.29      ┆ 80.6            │
│ South Korea ┆ 2012 ┆ 1.3       ┆ 80.7            │
│ South Korea ┆ 2013 ┆ 1.32      ┆ 80.9            │
│ South Korea ┆ 2014 ┆ 1.34      ┆ 80.9            │
│ South Korea ┆ 2015 ┆ 1.36      ┆ 81.0            │
└─────────────┴──────┴───────────┴─────────────────┘
)