Spark/Scala：从具有 Array 类型列的 DataFrame 中删除一些组件答案

【问题标题】：Spark/Scala: Remove some component from a DataFrame with Array typed columnSpark/Scala：从具有 Array 类型列的 DataFrame 中删除一些组件
【发布时间】：2018-09-27 19:40:50
【问题描述】：

标题可能不是很清楚。让我用一个例子来解释我想要实现的目标。以 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")
+---+----+----+--------------------+
| id|dept|test|               array|
+---+----+----+--------------------+
|  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]|
+---+----+----+--------------------+

我想根据 id、dept 和 test 列中的信息删除/删除数组列的一些元素。具体来说，每个数组的4个元素对应CS部门的4个id，编号是按照id升序生成的（意思是1、5、7、101）。现在我想删除每个数组中对应于测试列为 1 的 id 的元素。在此示例中，将删除第 2 和第 4 个元素，最终结果将如下所示：

+---+----+----+----------+
| id|dept|test|     array|
+---+----+----+----------+
|  1|  CS|   0|[0.1, 0.4]|
|  4|  Ed|   0|[0.4, 0.3]|
|  7|  CS|   0|[0.2, 0.4]|
|101|  CS|   1|[0.5, 0.3]|
|  5|  CS|   1|[0.4, 0.6]|
+---+----+----+----------+

为了避免收集所有结果并在 Scala 中进行操作。如果可能，我想将操作保留在 Spark DataFrame 中。我解决这个问题的想法包括两个步骤：

找出需要删除的数组元素的索引
应用删除/删除操作

到目前为止，我想我已经弄清楚了第 1 步：

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

    val w = Window.partitionBy("dept").orderBy("id")
    val studentIdIdx = df.select("id", "dept")
      .withColumn("Index", row_number().over(w))
      .where("dept = 'CS'").drop("dept")
    studentIdIdx.show()
    +---+-----+
    | id|Index|
    +---+-----+
    |  1|    1|
    |  5|    2|
    |  7|    3|
    |101|    4|
    +---+-----+
    val testIds = df.where("test = 1")
      .select($"id".as("test_id"))
    val testMask = studentIdIdx
      .join(testIds, studentIdIdx("id") === testIds("test_id"))
      .drop("id","test_id")
    testMask.show()
    +-----+
    |Index|
    +-----+
    |    2|
    |    4|
    +-----+

所以我的两个相关问题是：

如何将删除/删除功能应用到每行中带有索引的每个数组？（我也愿意提出更好的方法来计算索引）
我想要的真正最终 DataFrame 应该在上述结果之上删除一些元素。具体来说，对于test=0 & dept=CS，应该去掉id的Index对应的数组元素。在这个例子中，id=1 的行中的第 1 个元素和 id=7 的行中的第 3 个元素（删除前的原始索引）应该被删除，真正的最终结果是：
```
+---+----+----+----------+    
| id|dept|test|     array|
+---+----+----+----------+
|  1|  CS|   0|[0.4]     |
|  4|  Ed|   0|[0.4, 0.3]|
|  7|  CS|   0|[0.2]     |
|101|  CS|   1|[0.5, 0.3]|
|  5|  CS|   1|[0.4, 0.6]|
+---+----+----+----------+
```

我提到第二点，以防万一可以应用更有效的方法来同时实现两个删除操作。如果没有，我想一旦我知道如何使用索引信息进行删除操作，我应该能够弄清楚如何进行第二次删除。谢谢！

【问题讨论】：

在您的示例中，“array”列不是数组类型，而是结构类型
你说the 3rd element (original index before any removal) in the row with id=7 should be removed 那么| 7| CS| 0|[0.4] 不应该是| 7| CS| 0|[0.2] 吗？如果那是真的，那么我在下面的回答符合您的要求
是的，你是对的。我将编辑帖子。

标签： arrays scala apache-spark

【解决方案1】：

假设你有 initial dataframe 作为

+---+----+----+--------------------+
| id|dept|test|               array|
+---+----+----+--------------------+
|  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]|
+---+----+----+--------------------+

root
 |-- id: integer (nullable = false)
 |-- dept: string (nullable = true)
 |-- test: integer (nullable = false)
 |-- array: array (nullable = true)
 |    |-- element: double (containsNull = false)

你可以应用窗口函数来获取行号为

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

val w = Window.partitionBy("dept").orderBy("id")

val tempdf = df.withColumn("Index", row_number().over(w))

这会给你

+---+----+----+--------------------+-----+
|id |dept|test|array               |Index|
+---+----+----+--------------------+-----+
|1  |CS  |0   |[0.1, 0.2, 0.4, 0.5]|1    |
|5  |CS  |1   |[0.4, 0.2, 0.6, 0.9]|2    |
|7  |CS  |0   |[0.2, 0.5, 0.4, 0.7]|3    |
|101|CS  |1   |[0.5, 0.7, 0.3, 0.8]|4    |
|4  |Ed  |0   |[0.4, 0.8, 0.3, 0.6]|1    |
+---+----+----+--------------------+-----+

下一步是选择 dept = CS 和 test = 1 的行并获取索引列表

val csStudentIdIdxToRemove = tempdf.filter("dept = 'CS' and test = '1'").select(collect_list(tempdf("Index"))).collect()(0).getAs[Seq[Int]](0)
//WrappedArray(2, 4)

然后你定义一个 udf 函数来使用你的所有逻辑从数组列中删除元素

def removeUdf = udf((array: Seq[Double], additionalIndex: Int) => additionalIndex match{
  case 0 => array.zipWithIndex.filterNot(x => csStudentIdIdxToRemove.contains(x._2 + 1)).map(_._1)
  case _ => {
    val withAdditionalIndex = csStudentIdIdxToRemove ++ Seq(additionalIndex)
    array.zipWithIndex.filterNot(x => withAdditionalIndex.contains(x._2 + 1)).map(_._1)
  }
})

然后调用udf函数并删除索引列

tempdf.withColumn("array", removeUdf(tempdf("array"), when(tempdf("dept") === "CS" && tempdf("test") === 0, tempdf("Index")).otherwise(lit(0))))
    .drop("Index")

最后你应该得到你想要的结果

+---+----+----+----------+
|1  |CS  |0   |[0.4]     |
|5  |CS  |1   |[0.4, 0.6]|
|7  |CS  |0   |[0.2]     |
|101|CS  |1   |[0.5, 0.3]|
|4  |Ed  |0   |[0.4, 0.3]|
+---+----+----+----------+

希望答案简洁实用

【讨论】：

感谢 Ramesh，我在最后一行代码中遇到了不可序列化的错误：名称：org.apache.spark.SparkException 消息：任务不可序列化.....原因：java.io。 NotSerializableException：org.apache.spark.sql.expressions.WindowSpec。该错误似乎与 Window 函数有关。我在独立模式下使用 Spark 2.3。
添加@transient ，如stackoverflow.com/a/37334595/5880706中所述
谢谢。在 Window 函数之前添加 @transient 后它会运行。我对此不是很熟悉。添加这个对性能/资源要求有什么影响吗？任何我可以阅读更多关于此的参考资料都会非常有帮助。
@transient 注解仅表示该对象不用于序列化。你可以找到更多信息artima.com/pins1ed/annotations.html和help.semmle.com/wiki/display/SCALA/…
@GuanghuaShu 如果答案帮助您解决了问题，那么您应该考虑支持和接受:)

【解决方案2】：

这里是我的解决方案。要删除索引，我会使用 UDF：

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


val dropElements = udf(
  (array: Seq[Double], indices: Seq[Int]) =>
    array.zipWithIndex.filterNot { case (x, i) => indices.contains(i + 1) }.map(_._1)
)

df
  .withColumn("index_to_drop", row_number().over(Window.partitionBy($"dept").orderBy($"id")))
  .withColumn("index_to_drop", when($"test" === 1, $"index_to_drop"))
  .withColumn("indices_to_drop", collect_list($"index_to_drop").over(Window.partitionBy($"dept")))
  .withColumn("array", dropElements($"array", $"indices_to_drop"))
  .select($"id", $"dept", $"test", $"array")
  .show()


+---+----+----+--------------------+
| id|dept|test|               array|
+---+----+----+--------------------+
|  1|  CS|   0|          [0.1, 0.4]|
|  5|  CS|   1|          [0.4, 0.6]|
|  7|  CS|   0|          [0.2, 0.4]|
|101|  CS|   1|          [0.5, 0.3]|
|  4|  Ed|   0|[0.4, 0.8, 0.3, 0.6]|
+---+----+----+--------------------+

【讨论】：

感谢您为 udf 提供了一个很好的解决方案。我对您的解决方案有两个跟进：1）是否也可以删除 dept=Ed 的相应元素？ 2) 正如我在上一部分中提到的真实最终结果，由于一些元素已被删除，我现在无法知道元素的原始索引，实际上第二次删除并不直接。您是否看到将两个删除操作组合在一起的有效方法？

【解决方案3】：

这是没有 UDF 的不同解决方案。我建议您尽可能避免使用 UDF。从 spark 2.1.0 可以使用poseexplode 函数。也没有添加cmets 抱歉。

import org.apache.spark.sql.functions.posexplode
import org.apache.spark.sql.expressions.Window

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

scala> df.show()
+---+----+----+--------------------+
| id|dept|test|      arraytoprocess|
+---+----+----+--------------------+
|  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]|
+---+----+----+--------------------+

val columnIndicestoDrop = df.withColumn("zipRank",row_number().over(Window.partitionBy("dept")
    .orderBy("id")))
    .withColumn("pos",when($"test" === 1, $"zipRank"-1))
    .filter('pos.isNotNull)
    .select('pos)
    .distinct()

scala> columnIndicestoDrop.show()
+---+
|pos|
+---+
|  1|
|  3|
+---+

val dfwitharrayIndices = df.select('id,
    'dept,
    'test,
    'arraytoprocess,
    posexplode($"arraytoprocess") as Seq("pos", "val"))

scala> dfwitharrayIndices.show()
+---+----+----+--------------------+---+---+
| id|dept|test|      arraytoprocess|pos|val|
+---+----+----+--------------------+---+---+
|  1|  CS|   0|[0.1, 0.2, 0.4, 0.5]|  0|0.1|
|  1|  CS|   0|[0.1, 0.2, 0.4, 0.5]|  1|0.2|
|  1|  CS|   0|[0.1, 0.2, 0.4, 0.5]|  2|0.4|
|  1|  CS|   0|[0.1, 0.2, 0.4, 0.5]|  3|0.5|
|  4|  Ed|   0|[0.4, 0.8, 0.3, 0.6]|  0|0.4|
|  4|  Ed|   0|[0.4, 0.8, 0.3, 0.6]|  1|0.8|
|  4|  Ed|   0|[0.4, 0.8, 0.3, 0.6]|  2|0.3|
|  4|  Ed|   0|[0.4, 0.8, 0.3, 0.6]|  3|0.6|
|  7|  CS|   0|[0.2, 0.5, 0.4, 0.7]|  0|0.2|
|  7|  CS|   0|[0.2, 0.5, 0.4, 0.7]|  1|0.5|
|  7|  CS|   0|[0.2, 0.5, 0.4, 0.7]|  2|0.4|
|  7|  CS|   0|[0.2, 0.5, 0.4, 0.7]|  3|0.7|
|101|  CS|   1|[0.5, 0.7, 0.3, 0.8]|  0|0.5|
|101|  CS|   1|[0.5, 0.7, 0.3, 0.8]|  1|0.7|
|101|  CS|   1|[0.5, 0.7, 0.3, 0.8]|  2|0.3|
|101|  CS|   1|[0.5, 0.7, 0.3, 0.8]|  3|0.8|
|  5|  CS|   1|[0.4, 0.2, 0.6, 0.9]|  0|0.4|
|  5|  CS|   1|[0.4, 0.2, 0.6, 0.9]|  1|0.2|
|  5|  CS|   1|[0.4, 0.2, 0.6, 0.9]|  2|0.6|
|  5|  CS|   1|[0.4, 0.2, 0.6, 0.9]|  3|0.9|
+---+----+----+--------------------+---+---+

val finaldataFrame = dfwitharrayIndices
    .join(broadcast(columnIndicestoDrop),Seq("pos"),"leftanti")
    .select('id,'dept,'test,'val)
    .groupBy('id,'dept,'test)
    .agg(collect_list('val).as("finalarray"))

scala> finaldataFrame.show()
+---+----+----+----------+
| id|dept|test|finalarray|
+---+----+----+----------+
|  5|  CS|   1|[0.4, 0.6]|
|  4|  Ed|   0|[0.4, 0.3]|
|  1|  CS|   0|[0.1, 0.4]|
|  7|  CS|   0|[0.2, 0.4]|
|101|  CS|   1|[0.5, 0.3]|
+---+----+----+----------+

【讨论】：

感谢 Vijay，您能否提供一些关于您为什么更喜欢posexplode 而不是udf 的见解？当数据帧很大时，poseexplode 似乎很昂贵。
我不是 100% 确定这一点。但我觉得 udf 正在进行字符串比较，这会触发 spark 将所有字符串转换为 java.lang.string，这是一个昂贵的过程，因为 java 字符串占用了大量空间。但是使用posexplode，您并没有进行任何字符串操作，而是我们只是加入了两个数据帧/数据集，这是 spark 非常擅长的。但是让我知道上面的代码是如何扩展的。避免 udf 总是好的，因为 spark 函数总是有效的。