在特定列上应用映射转换，同时保留 Scala 中的其他列

Question

我有一个 map 函数，我想将它应用于 dataframe 中的特定列。假设我有一个数据框 ABC，如下所示：

A	B
foo	1
bar	4
biz	3

我想将B列的所有元素乘以2，同时保留A列得到如下输出：

A	B
foo	2
bar	8
biz	6

我知道如何选择数据框中的 B 列并使用 map 转换元素，如下面的代码所示： ABC.select("B").columns.map(c => c*2)

但我的问题是我也无法获得 A 列。

我第一次尝试使用这样的东西： ABC.select("A", (ABC.select("B").columns.map(c => c*2):_*)

但是，这会引发“select”方法无法重载的错误——这是公平的，因为它不接受 Array[Column] 类型的参数。然后我尝试了这个：

val arrayB = Array("B")

ABC.select(ABC.columns.map(c => if (arrayB.contains(c)) c*2 else col(c)):_*)

这确实有效并将结果返回给我，但我想知道是否有更好的方法来做到这一点。提前致谢！

Answer 1

试试这个：

import org.apache.spark.sql.functions.{col, expr}

ABC.select(col("A"), expr("B * 2").as("B"))

您也可以使用selectExpr("A, B * 2 as B")（甚至可能是col("B") * 2？），但我更喜欢第一个，它更清洁且更易于维护。

更新（因为最新的 cmets）： 据我了解，您可能希望将特定函数应用于 df 中的所有列，如果我是正确的，那么我建议您这样做：

import org.apache.spark.{Column, DataFrame}
import org.apache.spark.sql.functions._

// This is the base function which will be curried //
def applyColumnTransformation(column: Column)(transformation: Column => Column): Column = transformation(column)

val myTransformation: Column => Column = inputColumn => {
    // your logic here, for instance:
    inputColumn * 2  // you can pass name alias here, like (inputColumn * 2).as(inputColumn.toString)
}

def applyMyTransformation(column: Column): Column = applyColumnTransformation(column)(myTransformation)

现在你已经拥有了你想要的一切，只需这样做：

val allColumns: Seq[Column] = myDf.columns.map(myDf.col)
val transformedColumns: Seq[Column] = allColumns.map(applymyTransformation)
myDf.select(transformedColumns: _*)

顺便说一下，一定要看看functions包，它实际上有很多有用的功能。