Polars多值列pivot+unpivot实现：求更优替代Pandas stack/unstack方案

需求与Pandas参考实现 #

需要完成的操作：对包含2个值列的表执行pivot，对pivot后的DataFrame应用自定义函数（UDF），再执行unpivot（对应Pandas的stack操作）。

以下是Pandas中的实现示例：

import pandas as pd
import polars as pl
from polars import col, lit, selectors as cs

df = pl.DataFrame(
    {
        "foo": [1, 1, 2, 2, 3, 3],
        "bar": ["y", "x", "y", "x", "y", "x"],
        "baz1": [1, 2, None, 4, 5, None],
        "baz2": [1, None, 3, 4, None, 6]
    }
)
# df输出如下：
'''
┌─────┬─────┬──────┬──────┐
│ foo ┆ bar ┆ baz1 ┆ baz2 │
│ --- ┆ --- ┆ ---  ┆ ---  │
│ i64 ┆ str ┆ i64  ┆ i64  │
╞═════╪═════╪══════╪══════╡
│ 1   ┆ y   ┆ 1    ┆ 1    │
│ 1   ┆ x   ┆ 2    ┆ null │
│ 2   ┆ y   ┆ null ┆ 3    │
│ 2   ┆ x   ┆ 4    ┆ 4    │
│ 3   ┆ y   ┆ 5    ┆ null │
│ 3   ┆ x   ┆ null ┆ 6    │
└─────┴─────┴──────┴──────┘
'''

pd_df = df.to_pandas()

index_col = ['foo']
columns_col = ['bar']
values_col = ['baz1', 'baz2']

def pd_udf(df): # 示例函数，实际场景会更复杂
    return (
        df.ffill() * 3
    )

pd_res = (
    pd_df.groupby(index_col + columns_col).first() # 无重复数据时等价于set_index
    .unstack()
    .pipe(pd_udf)
    .stack()
    .reset_index()
    .sort_values(index_col + columns_col)
    .pipe(pl.from_pandas)
)
# pd_res输出如下：
'''
┌─────┬─────┬──────┬──────┐
│ foo ┆ bar ┆ baz1 ┆ baz2 │
│ --- ┆ --- ┆ ---  ┆ ---  │
│ i64 ┆ str ┆ f64  ┆ f64  │
╞═════╪═════╪══════╪══════╡
│ 1   ┆ x   ┆ 6.0  ┆ null │
│ 1   ┆ y   ┆ 3.0  ┆ 3.0  │
│ 2   ┆ x   ┆ 12.0 ┆ 12.0 │
│ 2   ┆ y   ┆ 3.0  ┆ 9.0  │
│ 3   ┆ x   ┆ 12.0 ┆ 18.0 │
│ 3   ┆ y   ┆ 15.0 ┆ 9.0  │
└─────┴─────┴──────┴──────┘
'''

已有的两种Polars实现方案 #

1. Pivot→Unpivot→二次Pivot方案 #

该方案需要处理pivot后的列命名（格式为{value_name}_{column_name}_{column_value}），实现stack逻辑较繁琐：

def pl_udf(df):
    return (
        df.with_columns(
            pl.exclude(index_col).forward_fill() * lit(3)
        )
    )

lazy_df_1 = (
    df
    .pivot(on = columns_col,
           index = index_col,
           values = values_col)
    .lazy()
    .pipe(pl_udf)
    
    # 对应Pandas的stack操作
    .unpivot(index = index_col)
    .select(
        col(index_col+['value']),
        col('variable').str.split('_').list.get(0).alias('temp'),
        col('variable').str.split('_').list.get(1).alias(columns_col[0]),
    )
    .collect()
    .pivot(on = 'temp',
          index = index_col+columns_col,
          values = 'value')
    .lazy()
    .sort(index_col+columns_col)
)
lazy_df_1.collect()

2. GroupBy+Explode+Series级UDF方案 #

该方案中，Series级UDF需要先转为DataFrame才能使用DataFrame方法：

def pl_udf_series(s):
    '''直接应用于Series的UDF'''
    return (
        s.to_frame() # 转为DataFrame以使用DataFrame方法（实际UDF会用到这些方法）
        .select(col(s.name).forward_fill() * lit(3))
        .to_series().to_list()
    )

lazy_df_2 = (
    df.lazy()
    .group_by(columns_col)
    .agg(
        col(index_col),
        col(values_col).map_elements(pl_udf_series)  
    )
    .explode(columns=index_col + values_col)
    .sort(index_col+columns_col)
    .select(col(index_col+columns_col+values_col)) # 调整列顺序
)
lazy_df_2.collect()

两种方案均能得到与Pandas一致的结果：

pd_res.equals(lazy_df_1.collect())
# True
pd_res.equals(lazy_df_2.collect())
# True

性能对比 #

• Pandas方案：约1000微秒
• Pivot-Unpivot-Pivot方案：约380微秒
• GroupBy-Explode方案：约450微秒（部分场景下性能更优）

更优实现方案 #

可以利用Polars的melt+pivot组合，结合over窗口函数简化流程，避免多次pivot/unpivot的列拆分操作，同时保持高效：

def pl_optimized_udf(df):
    return df.with_columns(
        col('value').forward_fill().over(index_col) * lit(3)
    )

optimized_df = (
    df.lazy()
    # 先 melt 把值列转为长格式
    .melt(id_vars=index_col + columns_col, value_vars=values_col, variable_name='temp')
    # 应用UDF，利用窗口函数处理ffill
    .pipe(pl_optimized_udf)
    # 再 pivot 转回宽格式
    .pivot(
        index=index_col + columns_col,
        on='temp',
        values='value'
    )
    .sort(index_col + columns_col)
    .collect()
)

方案优势 #

1. 代码更简洁：无需处理pivot后的复杂列名拆分，逻辑更直观
2. 性能更稳定：减少了中间的列转换和多次pivot操作，在大数据量下优势更明显
3. 可扩展性强：UDF中可以直接使用Polars的窗口函数、表达式等，无需在Series和DataFrame之间来回转换

验证结果一致性：

pd_res.equals(optimized_df)
# True

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