pyspark数据帧转置问题

Posted 2023-04-15

【中文标题】pyspark数据帧转置问题

【英文标题】：pyspark dataframe transpose issue

【发布时间】：2020-04-11 12:35:03

【问题描述】：

我有一个带有相应列的以下输入数据框：

dim1,dim2,tran_clmn,input1,input2
101,201,Y1,1,2
102,202,Y2,2,3
103,203,Y3,3,4
104,204,Y4,4,5
105,205,Y5,5,6

我需要根据列 tran_clmn

将输入数据转置为下面的转置输出

dim1,dim2,new_trn_clm,Y1,Y2,Y3,Y4,Y5
101,201,input1,1,,,,
101,201,input2,2,,,,
102,202,input1,,2,,,
102,202,input2,,3,,,
103,203,input1,,,3,,
103,203,input2,,,4,,
104,204,input1,,,,4,
104,204,input2,,,,5,
105,205,input1,,,,,5
105,205,input2,,,,,6

如何实现这种场景？没有聚合选项。是否可以使用 groupBy 和 pivot 方法完成并获得结果？

【参考方案1】：

两年前，我在这个post 中发现了一个执行这种转换（从宽到长）的函数。这对我很有帮助：也许对你也有帮助。

import pyspark.sql.functions as psf

def melt(df: DataFrame,
         id_vars: Iterable[str], value_vars: Iterable[str],
         var_name: str = "variable", value_name: str = "value"):
    """
    Convert :class:`DataFrame` from wide to long format.
    """
    # Create array<struct<variable: str, value: ...>>
    _vars_and_vals = psf.array(*(
        psf.struct(psf.lit(c).alias(var_name), psf.col(c).alias(value_name))
        for c in value_vars))
    # Add to the DataFrame and explode
    _tmp = df.withColumn("_vars_and_vals", psf.explode(_vars_and_vals))
    cols = id_vars + [
        psf.col("_vars_and_vals")[x].alias(x) for x in [var_name, value_name]]
    return _tmp.select(*cols)

在你的情况下，这可能是这样的（未经测试）：

melt(df, id_vars=['dim1','dim2'],
                       value_vars=[input1','input2'],
                       var_name="tran_clmn", value_name="tran_clmn")