Spark/Scala：移除DataFrame数组类型列中的部分元素

嘿，我来帮你搞定Spark DataFrame里移除数组列元素的问题！先把你给出的初始DataFrame代码再贴一遍，方便后续演示：

val df = Seq(
  (1, "CS", 0, (0.1, 0.2, 0.4, 0.5)),
  (4, "Ed", 0, (0.4, 0.8, 0.3, 0.6)),
  (7, "CS", 0, (0.2, 0.5, 0.4, 0.7)),
  (101, "CS", 1, (0.5, 0.7, 0.3, 0.8)),
  (5, "CS", 1, (0.4, 0.2, 0.6, 0.9))
).toDF("id", "dept", "test", "array")

你提到的“移除部分元素”通常有几种常见场景，我分别给出实现方案：

1. 按索引移除指定位置的元素 #

如果你的需求是移除数组中特定索引位置的元素（比如第2个、第4个，注意Spark数组索引从0开始），推荐用Spark 3.0+支持的transform函数结合过滤来实现，性能比自定义UDF好：

import org.apache.spark.sql.functions.{transform, when, lit}

// 定义要移除的索引集合（比如移除索引1和3的元素）
val removeIndices = Set(1, 3)

val dfTrimmedByIndex = df.withColumn(
  "array_trimmed",
  // 遍历数组元素和索引，保留不在移除列表中的元素
  transform(df("array"), (elem, idx) => when(!removeIndices.contains(idx), elem))
    // 过滤掉被标记为null的元素
    .filter(_.isNotNull)
)

dfTrimmedByIndex.show(false)

执行后，新的array_trimmed列就会去掉原数组中索引1和3的元素，比如第一行的结果会变成[0.1, 0.4]。

2. 按元素值移除特定值 #

如果要移除数组中所有等于某个值的元素（比如所有0.4），直接用filter内置函数最方便：

import org.apache.spark.sql.functions.filter

val dfFilteredByValue = df.withColumn(
  "array_trimmed",
  // 保留所有不等于0.4的元素
  filter(df("array"), elem => elem =!= lit(0.4))
)

dfFilteredByValue.show(false)

3. 移除满足条件的元素 #

如果需要移除数组中符合某个条件的元素（比如大于0.6的元素），同样用filter函数：

val dfFilteredByCondition = df.withColumn(
  "array_trimmed",
  // 保留所有小于等于0.6的元素
  filter(df("array"), elem => elem <= lit(0.6))
)

dfFilteredByCondition.show(false)

兼容低版本Spark（低于3.0）的方案 #

如果你的Spark版本低于3.0，不支持transform和filter内置函数，可以用自定义UDF来实现：

import org.apache.spark.sql.functions.udf

// UDF：移除指定索引的元素
val removeIndicesUDF = udf((arr: Seq[Double], indicesToRemove: Set[Int]) => {
  arr.zipWithIndex
    // 过滤掉要移除的索引对应的元素
    .filterNot { case (_, idx) => indicesToRemove.contains(idx) }
    // 只保留元素值
    .map(_._1)
})

val dfWithUDF = df.withColumn(
  "array_trimmed",
  removeIndicesUDF(df("array"), lit(Set(1, 3)))
)

dfWithUDF.show(false)

小提示 #

• 优先使用Spark内置函数（transform、filter），它们的性能比自定义UDF好，因为Spark可以对内置函数进行优化。
• 如果需要更复杂的移除逻辑，比如根据其他列的值动态决定移除规则，也可以基于上述思路扩展实现。

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