Замените два разных значения столбца в кадре данных, используя одно и то же условие с минимальной сложностью в 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 из этого случая. Но он все еще вызывает withColumn два раза

Answer 3

Мой код был не совсем таким, как в вопросе, но для меня что-то вроде этого работало с использованием карт и списка столбцов.

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)))

    }

Answer 4

Функция when в Spark возвращает объект типа Column. Если вы не хотите повторять свои условия, вы можете извлечь их в отдельную функцию:

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

, а затем использовать ее следующим образом:

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