时间戳

首页 > 解决方案 > Scala——有条件地替换数据框的列值

scala - Scala——有条件地替换数据框的列值

问题描述

DataFrame 1 是我现在拥有的，我想编写一个 Scala 函数来使 DataFrame 1 看起来像 DataFrame 2。

转让是大类；电子转账和 IMT 是子类别。

逻辑是对于同一个ID（31898），如果Transfer和e-Transfer都标记了它，它应该只是e-Transfer；如果 Transfer 和 IMT 和 e-Transfer 都标记为同一个 ID（32614），则应为 e-Transfer + IMT；如果只是Transfer tagged to one ID (33987)，应该是Other；如果仅 e-Transfer 或 IMT 标记为 ID (34193)，则它应该只是 e-transfer pr IMT。

Scala 新手，不知道如何编写一个好的函数来做到这一点。请帮忙！！

DataFrame 1                        DataFrame 2
+---------+-------------+          +---------+------------------+
|   ID    | Category    |          |   ID    | Category         |
+---------+-------------+          +---------+------------------+  
|  31898  |   Transfer  |          |  31898  |  e-Transfer      |  
|  31898  |  e-Transfer |          |  32614  |  e-Transfer + IMT|
|  32614  |   Transfer  |  =====>  |  33987  |   Other          |
|  32614  |  e-Transfer |  =====>  |  34193  |  e-Transfer      |
|  32614  |     IMT     |          +---------+------------------+
|  33987  |   Transfer  |  
|  34193  |  e-Transfer |  
+---------+-------------+

标签： scalaapache-sparkdataframeuser-defined-functions

解决方案

您可以对 DataFrame 进行分组ID以聚合Category使用collect_set来组装类别数组，并使用以下命令根据类别数组中的内容创建一个新列array_contains：

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

val df = Seq(
  (31898, "Transfer"),
  (31898, "e-Transfer"),
  (32614, "Transfer"),
  (32614, "e-Transfer"),
  (32614, "IMT"),
  (33987, "Transfer"),
  (34193, "e-Transfer")
).toDF("ID", "Category")

df.groupBy("ID").agg(collect_set("Category").as("CategorySet")).
  withColumn( "Category",
    when(array_contains($"CategorySet", "e-Transfer") && array_contains($"CategorySet", "IMT"),
      "e-Transfer + IMT").otherwise(
    when(array_contains($"CategorySet", "e-Transfer") && array_contains($"CategorySet", "Transfer"),
      "e-Transfer").otherwise(
    when($"CategorySet" === Array("e-Transfer") || $"CategorySet" === Array("MIT"),
      $"CategorySet"(0)).otherwise(
    when($"CategorySet" === Array("Transfer"), "Other")
    )))
  ).
  show(false)
// +-----+---------------------------+----------------+
// |ID   |CategorySet                |Category        |
// +-----+---------------------------+----------------+
// |33987|[Transfer]                 |Other           |
// |32614|[Transfer, e-Transfer, IMT]|e-Transfer + IMT|
// |34193|[e-Transfer]               |e-Transfer      |
// |31898|[Transfer, e-Transfer]     |e-Transfer      |
// +-----+---------------------------+----------------+

您的样本数据可能未涵盖所有情况（例如[Transfer, MIT]）。现有示例代码将为null任何剩余案例生成类别值。如果发现其他情况，只需修改/扩展条件检查。

推荐阅读