Spark DataFrame列的合并与拆分

版本说明：Spark-2.3.0

使用Spark SQL在对数据进行处理的过程中，可能会遇到对一列数据拆分为多列，或者把多列数据合并为一列。这里记录一下目前想到的对DataFrame列数据进行合并和拆分的几种方法。

1 DataFrame列数据的合并
例如：我们有如下数据，想要将三列数据合并为一列，并以“,”分割

1+----+---+-----------+
2|name|age|      phone|
3+----+---+-----------+
4|Ming| 20|15552211521|
5|hong| 19|13287994007|
6| zhi| 21|15552211523|
7+----+---+-----------+

1.1 使用map方法重写

使用map方法重写就是将DataFrame使用map取值之后，然后使用toSeq方法转成Seq格式，最后使用Seq的foldLeft方法拼接数据，并返回，如下所示：

1//方法1：利用map重写
2    val separator = ","
3    df.map(_.toSeq.foldLeft("")(_ + separator + _).substring(1)).show()
4
5    /**
6      * +-------------------+
7      * |              value|
8      * +-------------------+
9      * |Ming,20,15552211521|
10      * |hong,19,13287994007|
11      * | zhi,21,15552211523|
12      * +-------------------+
13      */

1.2 使用内置函数concat_ws

合并多列数据也可以使用SparkSQL的内置函数concat_ws()

1//方法2： 使用内置函数 concat_ws
2    import org.apache.spark.sql.functions._
3    df.select(concat_ws(separator, $"name", $"age", $"phone").cast(StringType).as("value")).show()
4
5    /**
6      * +-------------------+
7      * |              value|
8      * +-------------------+
9      * |Ming,20,15552211521|
10      * |hong,19,13287994007|
11      * | zhi,21,15552211523|
12      * +-------------------+
13      */

1.3 使用自定义UDF函数

自己编写UDF函数，实现多列合并

1 //方法3：使用自定义UDF函数
2
3    // 编写udf函数
4    def mergeCols(row: Row): String = {
5      row.toSeq.foldLeft("")(_ + separator + _).substring(1)
6    }
7
8    val mergeColsUDF = udf(mergeCols _)
9    df.select(mergeColsUDF(struct($"name", $"age", $"phone")).as("value")).show()

完整代码：

1import org.apache.spark.sql.{Row, SparkSession}
2import org.apache.spark.sql.types.StringType
3
4/**
5  * Created by shirukai on 2018/9/12
6  * DataFrame 合并列
7  */
8object MergeColsTest {
9  def main(args: Array[String]): Unit = {
10    val spark = SparkSession
11      .builder()
12      .appName(this.getClass.getSimpleName)
13      .master("local")
14      .getOrCreate()
15
16    //从内存创建一组DataFrame数据
17    import spark.implicits._
18    val df = Seq(("Ming", 20, 15552211521L), ("hong", 19, 13287994007L), ("zhi", 21, 15552211523L))
19      .toDF("name", "age", "phone")
20    df.show()
21    /**
22      * +----+---+-----------+
23      * |name|age|      phone|
24      * +----+---+-----------+
25      * |Ming| 20|15552211521|
26      * |hong| 19|13287994007|
27      * | zhi| 21|15552211523|
28      * +----+---+-----------+
29      */
30    //方法1：利用map重写
31    val separator = ","
32    df.map(_.toSeq.foldLeft("")(_ + separator + _).substring(1)).show()
33
34    /**
35      * +-------------------+
36      * |              value|
37      * +-------------------+
38      * |Ming,20,15552211521|
39      * |hong,19,13287994007|
40      * | zhi,21,15552211523|
41      * +-------------------+
42      */
43    //方法2： 使用内置函数 concat_ws
44    import org.apache.spark.sql.functions._
45    df.select(concat_ws(separator, $"name", $"age", $"phone").cast(StringType).as("value")).show()
46
47    /**
48      * +-------------------+
49      * |              value|
50      * +-------------------+
51      * |Ming,20,15552211521|
52      * |hong,19,13287994007|
53      * | zhi,21,15552211523|
54      * +-------------------+
55      */
56    //方法3：使用自定义UDF函数
57
58    // 编写udf函数
59    def mergeCols(row: Row): String = {
60      row.toSeq.foldLeft("")(_ + separator + _).substring(1)
61    }
62
63    val mergeColsUDF = udf(mergeCols _)
64    df.select(mergeColsUDF(struct($"name", $"age", $"phone")).as("value")).show()
65
66    /**
67      * /**
68      * * +-------------------+
69      * * |              value|
70      * * +-------------------+
71      * * |Ming,20,15552211521|
72      * * |hong,19,13287994007|
73      * * | zhi,21,15552211523|
74      * * +-------------------+
75      **/
76      */
77  }
78}

2 DataFrame列数据的拆分

上面我们将DataFrame的多列数据合并为一列如下所示，有时候我们也需要将单列数据，以某种拆分规则，拆分为多列。下面提供几种将一列拆分为多列的方法。

1+-------------------+
2|              value|
3+-------------------+
4|Ming,20,15552211521|
5|hong,19,13287994007|
6| zhi,21,15552211523|
7+-------------------+

2.1 使用内置函数split，然后遍历添加列

该方法，先利用内置函数split将单列的数据拆分，然后遍历使用getItem(角标)方法获取拆分后的数据，依次使用withColumn方法添加新列，代码如下所示：

1  //方法1： 使用内置函数split，然后遍历添加列
2    val separator = ","
3    lazy val first = df.first()
4
5    val numAttrs = first.toString().split(separator).length
6    val attrs = Array.tabulate(numAttrs)(n => "col_" + n)
7    //按指定分隔符拆分value列，生成splitCols列
8    var newDF = df.withColumn("splitCols", split($"value", separator))
9    attrs.zipWithIndex.foreach(x => {
10      newDF = newDF.withColumn(x._1, $"splitCols".getItem(x._2))
11    })
12    newDF.show()
13  /**
14      * +-------------------+--------------------+-----+-----+-----------+
15      * |              value|           splitCols|col_0|col_1|      col_2|
16      * +-------------------+--------------------+-----+-----+-----------+
17      * |Ming,20,15552211521|[Ming, 20, 155522...| Ming|   20|15552211521|
18      * |hong,19,13287994007|[hong, 19, 132879...| hong|   19|13287994007|
19      * | zhi,21,15552211523|[zhi, 21, 1555221...|  zhi|   21|15552211523|
20      * +-------------------+--------------------+-----+-----+-----------+

2.2 使用UDF函数创建多列数据，然后合并
该方法是使用udf函数，生成多个列，然后合并到原来的数据。该方法参考了VectorDisassembler（与spark ml官网提供的VectorAssembler相反），这是一个第三方的spark ml向量拆分算法，该方法github地址：https://github.com/jamesbconner/VectorDisassembler。代码如下所示：

1//方法2：使用udf函数创建多列，然后合并
2    val attributes: Array[Attribute] = {
3      val numAttrs = first.toString().split(separator).length
4      //生成attributes
5      Array.tabulate(numAttrs)(i => NumericAttribute.defaultAttr.withName("value" + "_" + i))
6    }
7    //创建多列数据
8    val fieldCols = attributes.zipWithIndex.map(x => {
9      val assembleFunc = udf {
10        str: String =>
11          str.split(separator)(x._2)
12      }
13      assembleFunc(df("value").cast(StringType)).as(x._1.name.get, x._1.toMetadata())
14    })
15    //合并数据
16    df.select(col("*") +: fieldCols: _*).show()
17
18    /**
19      * +-------------------+-------+-------+-----------+
20      * |              value|value_0|value_1|    value_2|
21      * +-------------------+-------+-------+-----------+
22      * |Ming,20,15552211521|   Ming|     20|15552211521|
23      * |hong,19,13287994007|   hong|     19|13287994007|
24      * | zhi,21,15552211523|    zhi|     21|15552211523|
25      * +-------------------+-------+-------+-----------+
26      */

完整代码：

1import org.apache.spark.ml.attribute.{Attribute, NumericAttribute}
2import org.apache.spark.sql.SparkSession
3import org.apache.spark.sql.types.StringType
4
5/**
6  * Created by shirukai on 2018/9/12
7  * 拆分列
8  */
9object SplitColTest {
10  def main(args: Array[String]): Unit = {
11    val spark = SparkSession
12      .builder()
13      .appName(this.getClass.getSimpleName)
14      .master("local")
15      .getOrCreate()
16
17    //从内存中创建DataFrame
18    import spark.implicits._
19    val df = Seq("Ming,20,15552211521", "hong,19,13287994007", "zhi,21,15552211523")
20      .toDF("value")
21    df.show()
22
23    /**
24      * +-------------------+
25      * |              value|
26      * +-------------------+
27      * |Ming,20,15552211521|
28      * |hong,19,13287994007|
29      * | zhi,21,15552211523|
30      * +-------------------+
31      */
32
33    import org.apache.spark.sql.functions._
34    //方法1： 使用内置函数split，然后遍历添加列
35    val separator = ","
36    lazy val first = df.first()
37
38    val numAttrs = first.toString().split(separator).length
39    val attrs = Array.tabulate(numAttrs)(n => "col_" + n)
40    //按指定分隔符拆分value列，生成splitCols列
41    var newDF = df.withColumn("splitCols", split($"value", separator))
42    attrs.zipWithIndex.foreach(x => {
43      newDF = newDF.withColumn(x._1, $"splitCols".getItem(x._2))
44    })
45    newDF.show()
46
47    /**
48      * +-------------------+--------------------+-----+-----+-----------+
49      * |              value|           splitCols|col_0|col_1|      col_2|
50      * +-------------------+--------------------+-----+-----+-----------+
51      * |Ming,20,15552211521|[Ming, 20, 155522...| Ming|   20|15552211521|
52      * |hong,19,13287994007|[hong, 19, 132879...| hong|   19|13287994007|
53      * | zhi,21,15552211523|[zhi, 21, 1555221...|  zhi|   21|15552211523|
54      * +-------------------+--------------------+-----+-----+-----------+
55      */
56
57    //方法2：使用udf函数创建多列，然后合并
58    val attributes: Array[Attribute] = {
59      val numAttrs = first.toString().split(separator).length
60      //生成attributes
61      Array.tabulate(numAttrs)(i => NumericAttribute.defaultAttr.withName("value" + "_" + i))
62    }
63    //创建多列数据
64    val fieldCols = attributes.zipWithIndex.map(x => {
65      val assembleFunc = udf {
66        str: String =>
67          str.split(separator)(x._2)
68      }
69      assembleFunc(df("value").cast(StringType)).as(x._1.name.get, x._1.toMetadata())
70    })
71    //合并数据
72    df.select(col("*") +: fieldCols: _*).show()
73
74    /**
75      * +-------------------+-------+-------+-----------+
76      * |              value|value_0|value_1|    value_2|
77      * +-------------------+-------+-------+-----------+
78      * |Ming,20,15552211521|   Ming|     20|15552211521|
79      * |hong,19,13287994007|   hong|     19|13287994007|
80      * | zhi,21,15552211523|    zhi|     21|15552211523|
81      * +-------------------+-------+-------+-----------+
82      */
83  }
84}