如何将spark数据帧列嵌入到地图列？

Question

我有一个包含许多列的spark数据帧。现在，我想将它们组合到一个地图中并构建一个新列。例如

col1:String col2:String col3:String... coln:String =>
col: Map(colname -> colval)

一种方法是：

df.withColumn("newcol", struct(df.columns.head, df.columns.tail: _*))

但是，我仍然需要将df转换为dataset。我不知道如何定义可以匹配struct类型的case类。

另一种选择是将列嵌入到Map类型中，但我不知道如何表达它。

Answer 1

出于性能原因，您可以使用现有的Spark函数来避免滚动自己的UDF：

org.apache.spark.sql.functions.map

这是一个完整的例子：

var mydata = Seq(("a", "b", "c"), ("d", "e", "f"), ("g", "h", "i"))
                 .toDF("f1", "f2", "f3")

var colnms_n_vals = mydata.columns.flatMap { c => Array(lit(c), col(c)) }

display(mydata.withColumn("myMap", map(colnms_n_vals:_*)))

结果如下：

f1  f2  f3  myMap
a   b   c   {"f1":"a","f2":"b","f3":"c"}
d   e   f   {"f1":"d","f2":"e","f3":"f"}
g   h   i   {"f1":"g","f2":"h","f3":"i"}

Answer 2

如果要从所有现有列构建新列，这是一个简单的解决方案。

import org.apache.spark.sql.functions._
val columnsName = ds.columns
val mkMap = udf((values: mutable.WrappedArray[Int]) => columnsName.zip(values).toMap)
ds.withColumn("new_col", mkMap(array(columnsName.head, columnsName.tail:_*)))