У меня есть df в искре, которая как следующая структура:

amount gender status
1000   male   married
1313   female single
1000   male   married

В основном я хочу создать новый столбец, где пол это число

amount gender status  gender_num
1000   male   married 1
1313   female single  2
1000   male   married 1

Я устал от следующего:

  val gender = df.gender

  val gender_num = gender match {
case male => 1
case female => 2
}

Я получаю следующую ошибку:

<console>:125: error: value pa_gender_category is not a member of org.apache.spark.sql.DataFrame
val gender = data.pa_gender_category

Я знаю, что есть функция stringtoindex, но я хотел бы сделать это вручную

4
Anubhav Dikshit 24 Апр 2017 в 15:15

2 ответа

Лучший ответ

Используйте withColumn

val input = // load input DataFrame
val withGender = input.withColumn("gender_num", when($"gender" === "female", 2).otherwise(1))

Вы можете связать больше вариантов:

val withGender = input.withColumn("gender_num", when($"gender" === "female", 2).when($"gender" == "other", 3).otherwise(1))

Вы также можете использовать UDF, как в ответе Акаша. Имейте в виду, что иногда пользовательские функции нельзя оптимизировать так же сильно, как встроенные функции, но они могут быть более читабельными

10
T. Gawęda 24 Апр 2017 в 12:28

Вы можете использовать UDF Spark

import org.apache.spark.sql.functions.udf
def genderToNumber: UserDefinedFunction = {
    udf((gender: String) => {
                             gender match {
                                           case "male" => 1
                                           case "female" => 2
                                          }
                          }               })

Вы можете применить UDF этим

   val newDF = df.withColumn("gender_num", genderToNumber(df("gender")))
2
Akash Sethi 24 Апр 2017 в 15:40