使用 Scala 将多列转换为 Spark Dataframe 上的一列地图

Posted 2023-03-31

【中文标题】使用 Scala 将多列转换为 Spark Dataframe 上的一列地图

【英文标题】：Convert multiple columns into a column of map on Spark Dataframe using Scala

【发布时间】：2019-02-21 00:20:26

【问题描述】：

我有一个具有可变列数的数据框，例如 Col1、Col2、Col3。 我需要使用下面的代码将 Col1 和 Col2 组合成一列数据类型映射。

val df_converted = df.withColumn("ConvertedCols", map(lit("Col1"), col("Col1"), lit("Col2"), col("Col2")))

但是当我不知道列的数量和名称时，如何对所有列执行此操作？

【参考方案1】：

一种方法是通过flatMap 将DataFrame 的列列表扩展为Seq(lit(c1), col(c1), lit(c2), col(c2), ...) 并应用Spark 的map，如下所示：

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

val df = Seq(
  ("a", "b", "c", "d"),
  ("e", "f", "g", "h")
).toDF("c1", "c2", "c3", "c4")

val kvCols = df.columns.flatMap(c => Seq(lit(c), col(c)))

df.withColumn("ConvertedCols", map(kvCols: _*)).show(false)
// +---+---+---+---+---------------------------------------+
// |c1 |c2 |c3 |c4 |ConvertedCols                          |
// +---+---+---+---+---------------------------------------+
// |a  |b  |c  |d  |Map(c1 -> a, c2 -> b, c3 -> c, c4 -> d)|
// |e  |f  |g  |h  |Map(c1 -> e, c2 -> f, c3 -> g, c4 -> h)|
// +---+---+---+---+---------------------------------------+

【讨论】：

谢谢！我们通过 flatmap() 得到一个新的 DF。如果我们在这里处理大数据，会有性能问题吗？

@Chi CHEN，像这样生成一个简单的 MapType 列并不比许多其他一般转换在计算上更昂贵，尽管在空间方面它确实需要相当大的内存/存储来容纳生成的数据，其中包括所有列值以及重复的列名（作为键）。

【参考方案2】：

另一种方法是使用from_json和to_json获取地图类型列：

val df2 = df.withColumn(
    "ConvertedCols", 
    from_json(to_json(struct("*")), lit("map<string,string>"))
)

df2.show(false)
+---+---+---+---+------------------------------------+
|c1 |c2 |c3 |c4 |ConvertedCols                       |
+---+---+---+---+------------------------------------+
|a  |b  |c  |d  |[c1 -> a, c2 -> b, c3 -> c, c4 -> d]|
|e  |f  |g  |h  |[c1 -> e, c2 -> f, c3 -> g, c4 -> h]|
+---+---+---+---+------------------------------------+