【Pyspark基础】行转列和列转行（超多列时）

一、问题

现在pyspark中有字段user_id和k个item_id列，目标是实现类似sql中经典任务的行转列和列转行，即一项项的user_id和item_id。可以通过df.printSchema()查看当前df的字段：

root
 |-- user_id: double (nullable = true)
 |-- beat_id[0]: double (nullable = true)
 |-- beat_id[1]: double (nullable = true)
 |-- beat_id[2]: double (nullable = true)
 |-- beat_id[3]: double (nullable = true)
 .......

二、方法一

先从一个栗子开始，可能会疑惑的地方在selectExpr里面的stack，可以理解成将对应的原来的一个个字段进行“堆叠”，然后再一一送入到后面as重命名的项目字段中：

# test_example
from pyspark.sql import SparkSession
spark = SparkSession.builder.appName('JupyterPySpark').enableHiveSupport().getOrCreate()
import pyspark.sql.functions as F
# 原始数据 
test = spark.createDataFrame([('2018-01','项目1',100), ('2018-01','项目2',200), ('2018-01','项目3',300),
                            ('2018-02','项目1',1000), ('2018-02','项目2',2000), ('2018-03','项目4',999),
                            ('2018-05','项目1',6000), ('2018-05','项目2',4000), ('2018-05','项目4',1999)
                           ], ['月份','项目','收入'])
# test.show()

# 一、行转列
test_pivot = test.groupBy('月份') \
        .pivot('项目', ['项目1', '项目2', '项目3', '项目4']) \
        .agg(F.sum('收入')) \
        .fillna(0)
test_pivot.show()

# 二、列转行
 # 逆透视Unpivot
unpivot_test =test_pivot.selectExpr("`月份`",
                        "stack(4, '项目1', `项目1`,'项目2', `项目2`, '项目3', `项目3`, '项目4', `项目4`) as (`项目`,`收入`)") \
        .filter("`收入` > 0 ") \
        .orderBy(["`月份`", "`项目`"]) \

unpivot_test.show()

+-------+-----+-----+-----+-----+
|   月份|项目1|项目2|项目3|项目4|
+-------+-----+-----+-----+-----+
|2018-03|    0|    0|    0|  999|
|2018-02| 1000| 2000|    0|    0|
|2018-05| 6000| 4000|    0| 1999|
|2018-01|  100|  200|  300|    0|
+-------+-----+-----+-----+-----+

+-------+-----+----+
|   月份| 项目|收入|
+-------+-----+----+
|2018-01|项目1| 100|
|2018-01|项目2| 200|
|2018-01|项目3| 300|
|2018-02|项目1|1000|
|2018-02|项目2|2000|
|2018-03|项目4| 999|
|2018-05|项目1|6000|
|2018-05|项目2|4000|
|2018-05|项目4|1999|
+-------+-----+----+

三、解决方案

• 和上面的思路一样，但是如果字段是中文时，需要将在stack中对中文加上``符号，但是不利于后期的处理，所以最好用正则表达式将其去掉。
• 如果需要转的列很多，就更需要用下面的实现了，定义如下的unpivot函数。

from pyspark.sql.functions import regexp_replace

def unpivot(df, keys,feature,value):
    '''df:待转换的数据框
       keys：待转换表中需要保留的主键key，以list[]类型传入
       feature, value：转换后的列名，可自定义
    '''
    # 转换类型是为了避免字段类不匹配，统一将数据转换为double类型(string也行)，如果保证数据类型完全一致，可以省略该句
    df = df.select(*[col(x).astype("double") for x in df.columns])
    cols = [x for x in df.columns if x not in keys]
    stack_str = ','.join(map(lambda x: "'`%s`', `%s`" % (x, x), cols))#这里join是为了用连接符‘,’将各个('`x`',`x`)连接起来
    
    df = (df.selectExpr(*keys, "stack(%s, %s) as (%s, %s)" % (len(cols), stack_str,feature,value))
          .withColumn(feature,regexp_replace(feature,'\`',''))
         )
    return df


keys = ['user_id']
feature,value = 'features','beat_id'
# df_test.new = unpivot(df_test, keys,feature,value)
df_result3 = unpivot(df_result2, keys,feature,value)
df_result3.show() 

+-----------+-----------+---------+
|    user_id|   features|  beat_id|
+-----------+-----------+---------+
|1.9079423E7| beat_id[0]|1018216.0|
|1.9079423E7| beat_id[1]| 886351.0|
|1.9079423E7| beat_id[2]|1051107.0|
|1.9079423E7| beat_id[3]|1018226.0|
+-----------+-----------+---------+

Reference

[1] https://zhuanlan.zhihu.com/p/337437504