Как уменьшить несколько строковых значений до предопределенных категорий в столбце

Question

Я хотел бы уменьшить значения определенного столбца в кадре данных на основе предварительно определенных категорий сопоставления с образцом.

Пример:

  val df = spark.createDataFrame(Seq(
  (1, "apple"),
  (2, "banana"),
  (3, "avocado"),
  (4, "potato"))).toDF("Id", "category")

Id  category
1   apple
2   banana
3   avocado
4   potato

Желаемый вывод:

  val df_reduced = spark.createDataFrame(Seq(
  (1, "fruit"),
  (2, "fruit"),
  (3, "vegetable"),
  (4, "vegetable"))).toDF("Id", "category")

Id  category
1   fruit
2   fruit
3   vegetable
4   vegetable

Вот решение, которое я придумал:

df.withColumn("category", when(col("category") === "apple", regexp_replace(col("category"), "apple", "fruit"))
              .otherwise(when(col("category") === "banana", regexp_replace(col("category"), "banana", "fruit"))
              .otherwise(when(col("category") === "avocado", regexp_replace(col("category"), "avocado", "vegetable"))
              .otherwise(when(col("category") === "potato", regexp_replace(col("category"), "potato", "vegetable"))
                         ))))
.show

Мне не очень нравится этот вложенный подход, когда есть необходимость, поэтому я хотел бы знать: есть ли лучшее, более идиоматическое решение дляэто задание?

Answer 1

Вы можете создать кадр данных для поиска как

val lookupDF = spark.createDataFrame(Seq(
  ("apple", "fruit"),
  ("banana", "fruit"),
  ("avocado", "vegetable"),
  ("potato", "vegetable"))).toDF("category", "category2")
//    +--------+---------+
//    |category|category2|
//    +--------+---------+
//    |apple   |fruit    |
//    |banana  |fruit    |
//    |avocado |vegetable|
//    |potato  |vegetable|
//    +--------+---------+

Поскольку фрейм данных lookup определенно будет маленьким, вы можете использовать функцию broadcast для join ing

import org.apache.spark.sql.functions._
df.join(broadcast(lookupDF), Seq("category"), "left")
    .select(col("Id"), col("category2").as("category"))
  .show(false)

что должно дать вам

+---+---------+
|Id |category |
+---+---------+
|1  |fruit    |
|2  |fruit    |
|3  |vegetable|
|4  |vegetable|
+---+---------+

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

Обновлено

Вы прокомментировали

как насчет пропущенных значений? если у меня есть категория в исходном df, которого нет в lookup df? Я получаю ноль, совет, как справиться с этим? Я бы предпочел сохранить исходное значение, если в таблице поиска не найдено совпадений, но я не могу сделать это с помощью объединений

Для решения такого случая вы можете использовать функцию when/otherwise как

import org.apache.spark.sql.functions._
df.join(broadcast(lookupDF), Seq("category"), "left")
  .select(col("Id"), when(col("category2").isNotNull, col("category2")).otherwise(col("category")).as("category"))
  .show(false)

Answer 2

Я думаю, вы должны воспользоваться помощью map и udf, как показано ниже

import org.apache.spark.sql.functions._

val map=Map("Apple"->"fruit","Mango"->"fruit","potato"->"vegetable","avocado"->"vegetable","Banana"->"fruit")

val replaceUDF=udf((name:String)=>map.getOrElse(name, name))
val outputdf=df.withColumn("new_category", replaceUDF(col("category"))

Пример вывода:

+---+--------+------------+
| Id|category|new_category|
+---+--------+------------+
|  1|   Apple|       fruit|
|  2|  Banana|       fruit|
|  3|  potato|   vegetable|
|  4| avocado|   vegetable|
|  5|   Mango|       fruit|
+---+--------+------------+