python - Pyspark Dataframe - 如何根据列数组作为输入连接列

Question

我有 10 列的数据框，想要基于作为输入的列数组进行函数串联:

arr = ["col1", "col2", "col3"]

这是迄今为止的最新情况:

newDF = rawDF.select(concat(col("col1"), col("col2"), col("col3") )).exceptAll(updateDF.select( concat(col("col1"), col("col2"), col("col3") ) ) )

另外:

df3 = df2.join(df1, concat( df2.col1, df2.col2, df2.col3, df2.col3 ) == df1.col5 ) 

但我想创建一个循环或函数来根据输入数组执行此操作(而不是像现在那样对其进行硬编码)。

最好的方法是什么？

Answer 1

您可以使用 (*) 解压列。在 pyspark.sql 文档中，如果任何函数具有 (*cols)，则意味着您可以解压 cols。对于连接:

pyspark.sql.functions.concat(*cols)

from pyspark.sql import functions as F
arr = ["col1", "col2", "col3"]
newDF = rawDF.select(F.concat(*(F.col(col) for col in arr))).exceptAll(updateDF.select(F.concat(*(F.col(col) for col in arr))))

对于连接:

arr=['col1','col2','col3']
df3 = df2.join(df1, F.concat(*(F.col(col) for col in arr)) == df1.col5 )