Scala实现：依据第二个DataFrame为第一个DataFrame的tags列表添加MAIN标签

需求说明 #

我有两个DataFrame：

• 第一个包含model、cnd、age、tags（字符串列表/数组，支持重复）、min、max列
• 第二个仅包含main_model列

需要实现逻辑：若第一个DataFrame的model值与第二个DataFrame中任意main_model匹配，则为该条数据的tags列表添加MAIN标签。

输入输出示例 #

输入 #

第一个DataFrame：

+------+-----+----+-------+------+-----+
| model| cnd | age| tags  |  min | max |
+------+-----+----+-------+------+-----+
|  foo1|   xx|  10|  []   |   1  |  2  |
|  foo2|   yy|  20|  []   |   2  |  3  | 
|  foo3|   zz|  30|  []   |   3  |  4  | 
+------+-----+----+-------+------+-----+

第二个DataFrame（主模型列表）：

+-----------+
| main_model|
+-----------+
|  foo1     |
|  foo3     | 
|  foo5     | 
+-----------+

输出 #

+------+-----+----+-------+------+-----+
| model| cnd | age| tags  |  min | max |
+------+-----+----+-------+------+-----+
|  foo1|   xx|  10|[MAIN] |   1  |  2  |
|  foo2|   yy|  20|  []   |   2  |  3  | 
|  foo3|   zz|  30|[MAIN] |   3  |  4  | 
+------+-----+----+-------+------+-----+

当前尝试的代码 #

dataFrame1.join(dataFrame2, dataFrame1("model") === dataFrame2("main_model"), "left_outer") .select(
dataFrame1("model"),
dataFrame1("cnd"),
dataFrame1("age"),
when(dataFrame2("main_model").isNotNull, concat(dataFrame1("tags"), lit(", MAIN"))).otherwise(dataFrame1("tags")).alias("tags"),
dataFrame1("min"),
dataFrame1("max")
)

问题：concat是字符串拼接函数，但tags是数组类型，直接拼接会导致类型不兼容，无法得到正确的数组结果。

正确解决方案 #

代码实现 #

首先导入Spark数组操作相关函数：

import org.apache.spark.sql.functions.{array, lit, when, array_union, col, array_append, array_contains}

基础版本（允许重复添加MAIN）

如果不担心tags中已存在MAIN，可以用array_append直接添加元素：

val mainModelsDF = dataFrame2.select("main_model")

val resultDF = dataFrame1.join(mainModelsDF, dataFrame1("model") === mainModelsDF("main_model"), "left_outer")
  .select(
    dataFrame1("model"),
    dataFrame1("cnd"),
    dataFrame1("age"),
    when(col("main_model").isNotNull, array_append(dataFrame1("tags"), lit("MAIN")))
      .otherwise(dataFrame1("tags"))
      .alias("tags"),
    dataFrame1("min"),
    dataFrame1("max")
  )

去重版本（避免重复添加MAIN）

如果需要确保tags中最多只有一个MAIN，可以先判断是否已包含该标签：

val resultDF = dataFrame1.join(mainModelsDF, dataFrame1("model") === mainModelsDF("main_model"), "left_outer")
  .select(
    dataFrame1("model"),
    dataFrame1("cnd"),
    dataFrame1("age"),
    when(
      col("main_model").isNotNull && !array_contains(dataFrame1("tags"), "MAIN"),
      array_union(dataFrame1("tags"), array(lit("MAIN")))
    ).otherwise(dataFrame1("tags")).alias("tags"),
    dataFrame1("min"),
    dataFrame1("max")
  )

关键说明 #

• array_append：将元素添加到数组末尾，允许重复元素，适合不需要去重的场景
• array_union：合并两个数组并自动去重，适合需要避免重复标签的场景
• 左连接后通过main_model是否为null判断当前model是否在主模型列表中

内容的提问来源于stack exchange，提问作者xard4sTR