使用相同的条件替换数据帧中的两个不同的列值，在 scala 中具有最小的复杂性

Question

函数withColumn可用于一次替换一列。我需要在相同条件下替换两个不同的列。

我用了两次withColumn。但我需要优化它。

val newDf: DataFrame = df.withColumn("a", when(col("b")===1, 0))
.withColumn("c", when(col("b")===1, "Y"))

第一

| a | b | c |
| 5 | 1 | N |

预期：

| a | b | c |
| 0 | 1 | Y|

Answer 1

val cols = df.columns:+"b":+"c"
val newDf: DataFrame = df.select(
"*",
when(col("b")===1, 0),
when(col("b")===1, "Y")
).toDF(cols:_*)

或 SQL 查询

df.createTempView("df")
spark.sql("
select df.*, 
case when df.b = 1 then 0 else NULL end as a,
case when df.b = 1 then 0 else NULL end as c, 
from df df
")

Answer 2

尝试匹配大小写。 给出条件并定义案例。从该案例中调用两个 withColumn。 但它仍然两次调用列

Answer 3

Spark 中的函数when返回类型为Column的对象。如果您不想重复您的条件，您可以将它们提取到单独的函数中：

def whenBIsOneThen[B](value: B): Column = when(col("b")===1, value)

然后像这样使用它：

df
.withColumn("a", whenBIsOneThen(0))
.withColumn("c", whenBIsOneThen("Y"))

Answer 4

我的代码与问题中的代码不完全一样，但对我来说，这样的东西可以使用地图和列列表。

val tobechanged = List("a","c")
val changeRelation = Map("a"->0,"c"->"Y")
var colList=df.columns.toSeq
var updatedDF = df
for(colName <- colList){
updatedDF = updatedDF.withColumn(colName, when(tobechanged.contains(colName), lit(changeRelation.get(colName)))
.otherwise(col(colName)))
}