Как обновить столбец в Spark Dataframe на основе группы значений столбцов?

Question

У меня есть Dataframe, значения которого должны быть из данной группы значений.

-----------------------
Id  Name    Department
-----------------------
1   John    Sales
2   Martin  Maintenance
3   Keith   Sales
4   Rob Unknown
5   Kevin   Unknown
6   Peter   Maintenance
------------------------

Допустимые значения для Department хранятся в String Array. [ 'Sales', 'техническое обслуживание', 'Training']

Если значение Департамента в DataFrame отличается от допустимых значений, его следует заменить на «Обучение». Таким образом, новый DataFrame будет -

-----------------------
Id  Name    Department
-----------------------
1   John    Sales
2   Martin  Maintenance
3   Keith   Sales
4   Rob     Training
5   Kevin   Training
6   Peter   Maintenance
------------------------

Что может быть возможным решением?

Answer 1

Вы можете выполнить свои требования, используя when/otherwise, concat и lit встроенные функции как

val validDepartments = Array("Sales","Maintenance","Training")

import org.apache.spark.sql.functions._
df.withColumn("Department", when(concat(validDepartments.map(x => lit(x)):_*).contains(col("Department")), col("Department")).otherwise("Training")).show(false)

, который должен дать вам

+---+------+---+-----------+
|Id |Name  |Age|Department |
+---+------+---+-----------+
|1  |John  |35 |Sales      |
|2  |Martin|34 |Maintenance|
|3  |Keith |33 |Sales      |
|4  |Rob   |34 |Training   |
|5  |Kevin |35 |Training   |
|6  |Peter |36 |Maintenance|
+---+------+---+-----------+

Простая udf функция должна также выполнить ваше требование как

val validDepartments = Array("Sales","Maintenance","Training")

import org.apache.spark.sql.functions._
def containsUdf = udf((department: String) => validDepartments.contains(department) match {case true => department; case false => "Training"} )

df.withColumn("Department", containsUdf(col("Department"))).show(false)

который должен дать вам тот же результат

Надеюсь, ответ полезен