将Spark中的多个ArrayType列合并为一个ArrayType列

Question

我想合并火花中的多个数组类型[字符串类型]列以创建一个数组类型[字符串类型]。为了组合两列，我在这里找到了解决方案：

将Array[string]类型的两个spark sql列合并成一个新的Array[string]列

但是如果我不知道编译时的列数，我该如何进行组合呢？在运行时，我将知道所有要组合的列的名称。

一种选择是使用上面stackoverflow问题中定义的UDF，在循环中多次添加两列。但这涉及对整个数据框架的多次读取。有没有办法一次完成？

+------+------+---------+
| col1 | col2 | combined|
+------+------+---------+
| [a,b]| [i,j]|[a,b,i,j]|
| [c,d]| [k,l]|[c,d,k,l]|
| [e,f]| [m,n]|[e,f,m,n]|
| [g,h]| [o,p]|[g,h,o,p]|
+------+----+-----------+

Answer 1

> < li>

处理dataframe架构并获取< code > array type[string type]类型的所有列。

使用函数创建新的数据帧。前两列的array_union

遍历其余列并将它们中的每一个添加到组合列中

>>>from pyspark import Row
>>>from pyspark.sql.functions import array_union
>>>df = spark.createDataFrame([Row(col1=['aa1', 'bb1'], 
                                col2=['aa2', 'bb2'],
                                col3=['aa3', 'bb3'], 
                                col4= ['a', 'ee'], foo="bar"
                               )])
>>>df.show()
+----------+----------+----------+-------+---+
|      col1|      col2|      col3|   col4|foo|
+----------+----------+----------+-------+---+
|[aa1, bb1]|[aa2, bb2]|[aa3, bb3]|[a, ee]|bar|
+----------+----------+----------+-------+---+
>>>cols = [col_.name for col_ in df.schema 
...       if col_.dataType == ArrayType(StringType()) 
...        or col_.dataType == ArrayType(StringType(), False)
...       ]
>>>print(cols)
['col1', 'col2', 'col3', 'col4']
>>>
>>>final_df = df.withColumn("combined", array_union(cols[:2][0], cols[:2][1]))
>>>
>>>for col_ in cols[2:]:
...    final_df = final_df.withColumn("combined", array_union(col('combined'), col(col_)))
>>>
>>>final_df.select("combined").show(truncate=False)
+-------------------------------------+
|combined                             |
+-------------------------------------+
|[aa1, bb1, aa2, bb2, aa3, bb3, a, ee]|
+-------------------------------------+ 

Answer 2

val arrStr: Array[String] = Array("col1", "col2")

val arrCol: Array[Column] = arrString.map(c => df(c))

val assembleFunc = udf { r: Row => assemble(r.toSeq: _*)}

val outputDf = df.select(col("*"), assembleFunc(struct(arrCol: 
_*)).as("combined"))

def assemble(rowEntity: Any*): 
                    collection.mutable.WrappedArray[String] = {

 var outputArray = 
 rowEntity(0).asInstanceOf[collection.mutable.WrappedArray[String]]

  rowEntity.drop(1).foreach {
    case v: collection.mutable.WrappedArray[String] =>
      outputArray ++= v
    case null =>
      throw new SparkException("Values to assemble cannot be 
      null.")
    case o =>
      throw new SparkException(s"$o of type ${o.getClass.getName} 
      is not supported.")
 }

outputArray
}

outputDf.show(false)