在spark scala 中将行合并为单个struct 列存在效率问题，我们如何做得更好？

Question

我正在尝试加快和限制获取多列及其值并将它们插入到同一行的地图中的成本。 这是一项要求，因为我们有一个正在读取此作业的遗留系统，但尚未准备好进行重构。 还有另外一张地图，有一些数据需要与此结合。

目前我们有一些解决方案，所有这些解决方案似乎都会在同一集群上产生大约相同的运行时间，并且在 Parquet 中存储了大约 1TB 的数据：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.json4s._
import org.json4s.jackson.JsonMethods._
import spark.implicits._

def jsonToMap(s: String, map: Map[String, String]): Map[String, String] = { 
  implicit val formats = org.json4s.DefaultFormats
    val jsonMap = if(!s.isEmpty){
      parse(s).extract[Map[String, String]]
    } else {
      Map[String, String]()
    }
    if(map != null){
      map ++ jsonMap
    } else {
      jsonMap
    }
  }
val udfJsonToMap = udf(jsonToMap _)

def addMap(key:String, value:String, map: Map[String,String]): Map[String,String] = {
  if(map == null) {
    Map(key -> value)
  } else {
    map + (key -> value)
  }
}

val addMapUdf = udf(addMap _)

val output = raw.columns.foldLeft(raw.withColumn("allMap", typedLit(Map.empty[String, String]))) { (memoDF, colName) =>
    if(colName.startsWith("columnPrefix/")){
        memoDF.withColumn("allMap", when(col(colName).isNotNull, addMapUdf(substring_index(lit(colName), "/", -1), col(colName), col("allTagsMap")) ))
    } else if(colName.equals("originalMap")){
        memoDF.withColumn("allMap", when(col(colName).isNotNull, udfJsonToMap(col(colName), col("allMap"))))
    } else {
      memoDF
    }
}

在 9 m5.xlarge 上大约需要 1 小时

val resourceTagColumnNames = raw.columns.filter(colName => colName.startsWith("columnPrefix/"))
def structToMap: Row => Map[String,String] = { row =>
  row.getValuesMap[String](resourceTagColumnNames)
}
val structToMapUdf = udf(structToMap)

val experiment = raw
  .withColumn("allStruct", struct(resourceTagColumnNames.head, resourceTagColumnNames.tail:_*))
  .select("allStruct")
  .withColumn("allMap", structToMapUdf(col("allStruct")))
  .select("allMap")

也在同一个集群上运行大约 1 小时

这段代码一切正常，但速度不够快，它比我们现在拥有的所有其他变换长约 10 倍，这对我们来说是一个瓶颈。

有没有另一种更有效的方法来获得这个结果？

编辑：我也尝试通过键限制数据，但是因为尽管键保持不变，我正在合并的列中的值可能会发生变化，我无法在不冒数据丢失风险的情况下限制数据大小。

Answer 1

Tl;DR：仅使用 spark sql 内置函数可以显着加快计算速度

如本答案所述，spark sql 本机函数比用户定义函数的性能更高。 因此，我们可以尝试仅使用 spark sql 本机函数来实现您的问题的解决方案。

我展示了两个主要的实现版本。 一个使用在我写这个答案时可用的最新版本 Spark 中存在的所有 sql 函数，即 Spark 3.0。 当提出问题时，另一个仅使用 spark 版本中存在的 sql 函数，因此 Spark 2.3 中存在的函数。 这个版本中用到的所有函数在Spark 2.2中也可以使用

Spark 3.0 使用 sql 函数实现

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{MapType, StringType}

val mapFromPrefixedColumns = map_filter(
  map(raw.columns.filter(_.startsWith("columnPrefix/")).flatMap(c => Seq(lit(c.dropWhile(_ != '/').tail), col(c))): _*),
  (_, v) => v.isNotNull
)

val mapFromOriginalMap = when(col("originalMap").isNotNull && col("originalMap").notEqual(""),
  from_json(col("originalMap"), MapType(StringType, StringType))
).otherwise(
  map()
)

val comprehensiveMapExpr = map_concat(mapFromPrefixedColumns, mapFromOriginalMap)

raw.withColumn("allMap", comprehensiveMapExpr)

Spark 2.2 用 sql 函数实现

在 spark 2.2 中，我们没有函数map_concat （在 spark 2.4 中可用）和map_filter （在 spark 3.0 中可用）。 我用用户定义的函数替换它们：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{MapType, StringType}

def filterNull(map: Map[String, String]): Map[String, String] = map.toSeq.filter(_._2 != null).toMap
val filter_null_udf = udf(filterNull _)

def mapConcat(map1: Map[String, String], map2: Map[String, String]): Map[String, String] = map1 ++ map2
val map_concat_udf = udf(mapConcat _)

val mapFromPrefixedColumns = filter_null_udf(
  map(raw.columns.filter(_.startsWith("columnPrefix/")).flatMap(c => Seq(lit(c.dropWhile(_ != '/').tail), col(c))): _*)
)

val mapFromOriginalMap = when(col("originalMap").isNotNull && col("originalMap").notEqual(""),
  from_json(col("originalMap"), MapType(StringType, StringType))
).otherwise(
  map()
)

val comprehensiveMapExpr = map_concat_udf(mapFromPrefixedColumns, mapFromOriginalMap)

raw.withColumn("allMap", comprehensiveMapExpr)

用sql函数实现，不用json映射

问题的最后一部分包含一个简化的代码，没有映射 json 列，也没有过滤结果映射中的空值。 我为这个特定案例创建了以下实现。 由于我不使用在 spark 2.2 和 spark 3.0 之间添加的函数，因此我不需要此实现的两个版本：

import org.apache.spark.sql.functions._

val mapFromPrefixedColumns = map(raw.columns.filter(_.startsWith("columnPrefix/")).flatMap(c => Seq(lit(c), col(c))): _*)
raw.withColumn("allMap", mapFromPrefixedColumns)

跑

对于以下数据帧作为输入：

+--------------------+--------------------+--------------------+----------------+
|columnPrefix/column1|columnPrefix/column2|columnPrefix/column3|originalMap     |
+--------------------+--------------------+--------------------+----------------+
|a                   |1                   |x                   |{"column4": "k"}|
|b                   |null                |null                |null            |
|c                   |null                |null                |{}              |
|null                |null                |null                |null            |
|d                   |2                   |null                |                |
+--------------------+--------------------+--------------------+----------------+

我获得以下allMap列：

+--------------------------------------------------------+
|allMap                                                  |
+--------------------------------------------------------+
|[column1 -> a, column2 -> 1, column3 -> x, column4 -> k]|
|[column1 -> b]                                          |
|[column1 -> c]                                          |
|[]                                                      |
|[column1 -> d, column2 -> 2]                            |
+--------------------------------------------------------+

对于没有 json 列的映射：

+---------------------------------------------------------------------------------+
|allMap                                                                           |
+---------------------------------------------------------------------------------+
|[columnPrefix/column1 -> a, columnPrefix/column2 -> 1, columnPrefix/column3 -> x]|
|[columnPrefix/column1 -> b, columnPrefix/column2 ->, columnPrefix/column3 ->]    |
|[columnPrefix/column1 -> c, columnPrefix/column2 ->, columnPrefix/column3 ->]    |
|[columnPrefix/column1 ->, columnPrefix/column2 ->, columnPrefix/column3 ->]      |
|[columnPrefix/column1 -> d, columnPrefix/column2 -> 2, columnPrefix/column3 ->]  |
+---------------------------------------------------------------------------------+

基准

我生成了一个 1000 万行的 csv 文件，未压缩（大约 800 Mo），其中包含一列没有列前缀的列，九列带有列前缀的列，以及一个包含 json 作为字符串的冒号：

+---+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+-------------------+
|id |columnPrefix/column1|columnPrefix/column2|columnPrefix/column3|columnPrefix/column4|columnPrefix/column5|columnPrefix/column6|columnPrefix/column7|columnPrefix/column8|columnPrefix/column9|originalMap        |
+---+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+-------------------+
|1  |iwajedhor           |ijoefzi             |der                 |ob                  |galsu               |ril                 |le                  |zaahuz              |fuzi                |{"column10":"true"}|
|2  |ofo                 |davfiwir            |lebfim              |roapej              |lus                 |roum                |te                  |javes               |karutare            |{"column10":"true"}|
|3  |jais                |epciel              |uv                  |piubnak             |saajo               |doke                |ber                 |pi                  |igzici              |{"column10":"true"}|
|4  |agami               |zuhepuk             |er                  |pizfe               |lafudbo             |zan                 |hoho                |terbauv             |ma                  |{"column10":"true"}|
...

基准是读取此 csv 文件，创建列allMap ，并将此列写入 parquet。 我在我的本地机器上运行了这个，我得到了以下结果

+--------------------------+--------------------+-------------------------+-------------------------+
|     implementations      | current (with udf) | sql functions spark 3.0 | sql functions spark 2.2 |
+--------------------------+--------------------+-------------------------+-------------------------+
| execution time           | 138 seconds        | 48 seconds              | 82 seconds              |
| improvement from current | 0 % faster         | 64 % faster             | 40 % faster             |
+--------------------------+--------------------+-------------------------+-------------------------+

我还遇到了问题中的第二个实现，即删除 json 列的映射和 map 中空值的过滤。

+--------------------------+-----------------------+------------------------------------+
| implementations          | current (with struct) | sql functions without json mapping |
+--------------------------+-----------------------+------------------------------------+
| execution time           | 46 seconds            | 35 seconds                         |
| improvement from current | 0 %                   | 23 % faster                        |
+--------------------------+-----------------------+------------------------------------+

当然，基准测试相当基础，但与使用用户定义函数的实现相比，我们可以看到改进

结论

当您遇到性能问题并且使用用户定义的函数时，尝试用 spark sql 函数替换这些用户定义的函数可能是个好主意