Spark将函数应用于某些列

Question

我有一个包含N列的数据框，我有一个hashmap，其中键是列名称，值是元组（字符串，布尔值）。

对于hasmap上的每一列并满足某个条件 - 与元组的一些比较 - 我必须将一个函数应用于该列的所有元素。

最后，我必须根据列名修改满足某些条件的数据帧的某些列。

这样做的最佳方式是什么？

Answer 1

您可以尝试在hashmap上使用foldLeft。这是Scala的一个例子：

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

val colMap = Map(
  "col1" -> ("a", true),
  "col2" -> ("a", true),
  "col3" -> ("a", false)
)

val df = Seq((1, 1, 1), (2, 2, 2)).toDF("col1", "col2", "col3")

def verify(strValue: String, boolValue: Boolean): Boolean = {
  strValue == "a" && boolValue
}
def processColumn(c: Column): Column = c * 2

colMap.foldLeft(df) {
  case (curDF, (colName, (strValue, boolValue))) => 
    if(verify(strValue, boolValue)) curDF.withColumn(colName, processColumn(col(colName)))
    else curDF
}