Spark Scala：基于DataFrame列表元素阈值添加标签的实现问题

需求说明 #

现有Spark DataFrame，其中grades列存储Grade对象列表，每个Grade对象包含name（字符串）和value（浮点数）字段。需要实现：当grades列表中存在name为HOME且value≥20.0的对象时，向tags列表添加PASS标签。

输入输出示例 #

输入：

+------+-----+----+-------+-------------------------------------------------------------+
| model| cnd | age| tags  |  grades                                                     |
+------+-----+----+-------+-------------------------------------------------------------+
|  foo1|   xx|  10|  []   |   [{name:"ATW", value: 10.0}, {name:"HOME", value: 20.0}]   | 
|  foo2|   xz|  12|  []   |   [{name:"ATW", value: 70.0}]                               | 
|  foo3|   xc|  13|  []   |   [{name:"ATW", value: 90.0}, {name:"HOME", value: 10.0}]   | 
+------+-----+----+-------+-------------------------------------------------------------+

输出：

+------+-----+----+-------+--------------------------------------------------------------+
| model| cnd | age| tags  |  grades                                                     |
+------+-----+----+-------+--------------------------------------------------------------+
|  foo1|   xx|  10| [PASS]|   [{name:"ATW", value: 10.0}, {name:"HOME", value: 20.0}]    | 
|  foo2|   xz|  12|  []   |   [{name:"ATW", value: 70.0}]                                | 
|  foo3|   xc|  13|  []   |   [{name:"ATW", value: 90.0}, {name:"HOME", value: 10.0}]    | 
+------+-----+----+-------+--------------------------------------------------------------+

错误原因 #

原代码抛出AnalysisException的核心原因：

• grades是嵌套对象数组，grades.value会被解析为Array<Double>类型，直接与20.0（单个Double值）比较，触发类型不匹配。
• 原逻辑仅判断是否存在name=HOME的元素，但无法关联对应元素的value是否满足≥20.0的条件，逻辑不严谨。

原错误代码：

dataFrame.withColumn("tags",
    when(
      array_contains(
        col("grades.name"),
        lit("HOME")
      ) && col("grades.value") >= lit(20.0),
      array_union(col("tags"), lit(Array("PASS")))
    ).otherwise(col("tags"))

正确实现 #

方案1：Spark 3.0+ 使用exists函数（推荐） #

exists函数可以直接遍历数组，检查是否存在符合条件的元素：

import org.apache.spark.sql.functions.{when, array_union, lit, exists, col}

val resultDF = dataFrame.withColumn("tags",
  when(
    // 遍历grades数组，检查是否存在符合条件的Grade对象
    exists(col("grades"), grade => 
      grade.getField("name") === lit("HOME") && grade.getField("value") >= lit(20.0)
    ),
    array_union(col("tags"), lit(Array("PASS")))
  ).otherwise(col("tags"))
)

方案2：低版本Spark 使用filter+size判断 #

如果Spark版本低于3.0，不支持exists，可以用filter过滤出符合条件的元素，再判断过滤后的数组长度是否大于0：

import org.apache.spark.sql.functions.{when, array_union, lit, filter, size, col}

val resultDF = dataFrame.withColumn("tags",
  when(
    // 过滤出name=HOME且value≥20的元素，判断是否有符合条件的结果
    size(filter(col("grades"), grade => 
      grade.getField("name") === lit("HOME") && grade.getField("value") >= lit(20.0)
    )) > 0,
    array_union(col("tags"), lit(Array("PASS")))
  ).otherwise(col("tags"))
)

说明 #

两种方案都能精准匹配“存在name为HOME且value≥20的Grade对象”的条件，避免了原代码的类型错误和逻辑漏洞，最终输出符合需求示例。

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