Spark实现行列转换pivot和unpivot

背景

做过数据清洗ETL工作的都知道，行列转换是一个常见的数据整理需求。在不同的编程语言中有不同的实现方法，比如SQL中使用case+group，或者Power BI的M语言中用拖放组件实现。今天正好需要在pyspark中处理一个数据行列转换，就把这个方法记录下来。

首先明确一下啥叫行列转换，因为这个叫法也不是很统一，有的地方叫转置，有的地方叫透视，不一而足。我们就以下图为例，定义如下：

• 从左边这种变成右边这种，叫透视（pivot）
• 反之叫逆透视（unpivot）

 

 

 

Spark实现

构造样本数据

首先我们构造一个以行格式保存数据的数据集

from pyspark.sql import SparkSession
spark = SparkSession.builder.appName('JupyterPySpark').enableHiveSupport().getOrCreate()

import pyspark.sql.functions as F

# 原始数据 
df = spark.createDataFrame([('2018-01','项目1',100), ('2018-01','项目2',200), ('2018-01','项目3',300),
                            ('2018-02','项目1',1000), ('2018-02','项目2',2000), ('2018-03','项目x',999)
                           ], ['年月','项目','收入'])
复制代码

样本数据如下，我们可以看到，每一个项目在指定月份都只有一行记录，并且项目是稀疏的。即，不是每个项目都会出现在每一个月份中，如项目2仅出现在2018-01当中。

+-------+---+----+
|  年月| 项目|  收入|
+-------+---+----+
|2018-01|项目1| 100|
|2018-01|项目2| 200|
|2018-01|项目3| 300|
|2018-02|项目1|1000|
|2018-02|项目2|2000|
|2018-03|项目x| 999|
+-------+---+----+
复制代码

透视Pivot

透视操作简单直接，逻辑如下

• 按照不需要转换的字段分组，本例中是年月；
• 使用pivot函数进行透视，透视过程中可以提供第二个参数来明确指定使用哪些数据项；
• 汇总数字字段，本例中是收入；

代码如下

df_pivot = df.groupBy('年月')\
                .pivot('项目', ['项目1','项目2','项目3','项目x'])\
                .agg(F.sum('收入'))\
                .fillna(0)
复制代码

结果如下

+-------+----+----+---+---+
| 年月| 项目1| 项目2|项目3|项目x|
+-------+----+----+---+---+
|2018-03|   0|   0|  0|999|
|2018-02|1000|2000|  0|  0|
|2018-01| 100| 200|300|  0|
+-------+----+----+---+---+
复制代码

逆透视Unpivot

Spark没有提供内置函数来实现unpivot操作，不过我们可以使用Spark SQL提供的stack函数来间接实现需求。有几点需要特别注意：

• 使用selectExpr在Spark中执行SQL片段；
• 如果字段名称有中文，要使用反引号**`** 把字段包起来；

代码如下

df_pivot.selectExpr("`年月`", 
                    "stack(4, '项目1', `项目1`,'项目2', `项目2`, '项目3', `项目3`, '项目x', `项目x`) as (`项目`,`收入`)")\
            .filter("`收入` > 0 ")\
            .orderBy(["`年月`", "`项目`"])\
            .show()
复制代码

结果如下

+-------+---+----+
|     年月| 项目|  收入|
+-------+---+----+
|2018-01|项目1| 100|
|2018-01|项目2| 200|
|2018-01|项目3| 300|
|2018-02|项目1|1000|
|2018-02|项目2|2000|
|2018-03|项目x| 999|
+-------+---+----+

作者：wait4friend
链接：https://juejin.im/post/5b1e343f518825137c1c6a27
来源：掘金
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

SparkSQL 实现

 

spark从1.6开始引入，到现在2.4版本，pivot算子有了进一步增强，这使得后续无论是交给pandas继续做处理，还是交给R继续分析，都简化了不少。大家无论在使用pandas、numpy或是R的时候，首先会做的就是处理数据，尤其是将列表，转成成合适的形状。

列表

在说透视表之前，我们先看看，什么是列表，在传统观念上，列表的每一行代表一条记录，而每一列代表一个属性。

+-------+-------+-----+

|   date|project|value|

+-------+-------+-----+

|2018-01|     p1|  100|

|2018-01|     p2|  200|

|2018-01|     p3|  300|

|2018-02|     p1| 1000|

|2018-02|     p2| 2000|

|2018-03|     px|  999|

+-------+-------+-----+

举个简单的例子，如上表，一条记录可能代表某个项目，在某个年月创造的价值。而在这个表里面，某一列，就代表一个属性，比如date代表日期，project代表项目名称。而这里每一行，代表一条独立，完整的记录，一条与另外一条记录，没有直接的关系。

这种结构，也是一般关系型数据库的数据结构。

透视表

透视表没有一个明确的定义，一般是观念上是指，为了方便进行数据分析，而对数据进行一定的重排，方便后续分析，计算等操作。透视表每一个元素及其对应的“坐标”一起形成一条完整的记录。

+-------+------+------+-----+-----+

|   date|    p1|    p2|   p3|   px|

+-------+------+------+-----+-----+

|2018-01| 100.0| 200.0|300.0|  0.0|

|2018-02|1000.0|2000.0|  0.0|  0.0|

|2018-03|   0.0|   0.0|  0.0|999.0|

+-------+------+------+-----+-----+

上面的表，是将列表进行重排后的透视表，其第一行和第一列可以理解成索引，而在表中根据索引可以确定一条唯一的值，他们一起组成一条相当于列表里的数据。

通过一般的定义，我们能看出，透视表主要用于分析，所以，一般的场景我们都会先对数据进行聚合，以后再对数据分析，这样也更有意义。就好像，将话费清单，做成透视表，尽管逻辑上没有任何问题，但是结果是可能比现在的清单列表更难查阅。

PS：一些可以借鉴的名词，目前维基百科并没有收录，也只能权且理解一下吧

建模拟数据

先来模拟个数据吧，按照前面的例子，建个csv，这里多加了一列s2，是为了做多透视列的，

date,project,value,s2
2018-01,p1,100,12
2018-01,p2,200,33
2018-01,p3,300,44
2018-02,p1,1000,22
2018-02,p2,2000,41
2018-03,px,999,22

spark API

我们先来看下DEMO程序

SparkConf sparkConf = new SparkConf().setAppName("JavaWordCount").setMaster("local");
SparkContext sc = SparkContext.getOrCreate(sparkConf);
SparkSession ss = new SparkSession(sc);
Dataset<Row> ds = ss.read()
         //csv分隔符 
        .option("sep", ",")
         //是否包含header
        .option("header", "true")
        //加载csv路径
        .csv("E:\\devlop\\workspace\\sparkdemo\\src\\main\\java\\com\\dafei1288\\spark\\data1.csv");
Dataset<Row>  r = 
        //设置分组
        ds.groupBy(col("date"))
        //设置pivot
        .pivot("project")
        //设置聚合
        .agg(sum("value"));
r.show();

在加载csv的时候，我们设置了分隔符，以及读取表头。

对加载后的dataset只需要进行3步设置

groupBy 设置分组列

pivot 设置pivot列

agg 设置聚合方式，可以是求和、平均等聚合函数

我们得到的输出结果如下：

+-------+------+------+-----+-----+

|   date|    p1|    p2|   p3|   px|

+-------+------+------+-----+-----+

|2018-03|  null|  null| null|999.0|

|2018-02|1000.0|2000.0| null| null|

|2018-01| 100.0| 200.0|300.0| null|

+-------+------+------+-----+-----+

请注意，这里和sql有些区别，就是groupBy的时候，不需要将project列写入了，如果写入成了

groupBy(col("date"),col("project"))
那么结果就是这样了

+-------+-------+------+------+-----+-----+

|   date|project|    p1|    p2|   p3|   px|

+-------+-------+------+------+-----+-----+

|2018-01|     p3|  null|  null|300.0| null|

|2018-01|     p2|  null| 200.0| null| null|

|2018-01|     p1| 100.0|  null| null| null|

|2018-03|     px|  null|  null| null|999.0|

|2018-02|     p1|1000.0|  null| null| null|

|2018-02|     p2|  null|2000.0| null| null|

+-------+-------+------+------+-----+-----+

sparkSQL 

SparkConf sparkConf = new SparkConf().setAppName("JavaWordCount").setMaster("local");
SparkContext sc = SparkContext.getOrCreate(sparkConf);
SparkSession ss = new SparkSession(sc);
Dataset<Row> ds = ss.read() .option("sep", ",")
        .option("header", "true").csv("E:\\devlop\\workspace\\sparkdemo\\src\\main\\java\\com\\dafei1288\\spark\\data1.csv");
ds.registerTempTable("f");
Dataset<Row>  r = ds.sqlContext().sql(
"select * from (
    select date,project as p,sum(value) as ss from f group by date,project
   )
  pivot (  
      sum(ss) 
      for p in ( 'p1','p2','p3','px' )  
   ) 
   order by date");
r.na().fill(0).show();

可以看到，这里我们将读取的csv注册成了表f，使用spark sql语句，这里和oracle的透视语句类似

pivot语法： pivot( 聚合列  for  待转换列 in (列值) )     

其语法还是比较简单的。

为了展示数据好看一点，我特意使用语句

r.na().fill(0)
将空值`null`替换成了0。

+-------+------+------+-----+-----+

|   date|    p1|    p2|   p3|   px|

+-------+------+------+-----+-----+

|2018-01| 100.0| 200.0|300.0|  0.0|

|2018-02|1000.0|2000.0|  0.0|  0.0|

|2018-03|   0.0|   0.0|  0.0|999.0|

+-------+------+------+-----+-----+

多聚合列

上文提到了，多做了一列，就是为了这个DEMO准备的，使用如下SparkSQL语句，设置多聚合列透视表

select * from (
    select date,project as p,sum(value) as ss,sum(s2) as ss2 from f group by date,project
)
pivot (  
      sum(ss),sum(ss2)  
     for p in ( 'p1','p2','p3','px' ) 
) 
order by date

这里为例方便看，我就截图了

为了防止OOM的情况，spark对pivot的数据量进行了限制，其可以通过spark.sql.pivotMaxValues 来进行修改，默认值为10000，这里是指piovt后的列数。