Scala:在保留空值的情况下在 Window 上收集列表()

Question

我有一个如下所示的数据框:

+----+----+----+
|colA|colB|colC|
+----+----+----+
|1   |1   |23  |
|1   |2   |63  |
|1   |3   |null|
|1   |4   |32  |
|2   |2   |56  |
+----+----+----+

我应用以下说明在 C 列中创建一系列值:

import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions._
df.withColumn("colD", 
collect_list("colC").over(Window.partitionBy("colA").orderBy("colB")))

结果是这样的，D 列被创建并包含 C 列的值作为一个序列，同时它删除了 null 值:

+----+----+----+------------+
|colA|colB|colC|colD        |
+----+----+----+------------+
|1   |1   |23  |[23]        |
|1   |2   |63  |[23, 63]    |
|1   |3   |null|[23, 63]    |
|1   |4   |32  |[23,63,32]  |
|2   |2   |56  |[56]        |
+----+----+----+------------+

但是，我想在新列中保留空值并得到以下结果:

+----+----+----+-----------------+
|colA|colB|colC|colD             |
+----+----+----+-----------------+
|1   |1   |23  |[23]             |
|1   |2   |63  |[23, 63]         |
|1   |3   |null|[23, 63, null]   |
|1   |4   |32  |[23,63,null, 32] |
|2   |2   |56  |[56]             |
+----+----+----+-----------------+

如您所见，结果中仍然有 null 值。你知道我该怎么做吗？

Answer 1

正如 LeoC 提到的，collect_list 将丢弃空值。此行为似乎有解决方法。通过将每个标量包装到数组中，然后是 collect_list 将导致 [[23], [63], [], [32]] 然后当您执行 flatten 你会得到 [23, 63,, 32]。数组中的那些缺失值是空值。

collect_list 和 flatten 我相信是在 Spark 2.4 中引入的内置 sql 函数。我没有研究实现来验证这是预期的行为，所以我不知道这个解决方案有多可靠。

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

val df = Seq(
  (Some(1), Some(1), Some(23)),
  (Some(1), Some(2), Some(63)),
  (Some(1), Some(3), None),
  (Some(1), Some(4), Some(32)),
  (Some(2), Some(2), Some(56))
).toDF("colA", "colB", "colC")

val newDf = df.withColumn("colD", flatten(collect_list(array("colC"))
    .over(Window.partitionBy("colA").orderBy("colB"))))


+----+----+----+-------------+
|colA|colB|colC|         colD|
+----+----+----+-------------+
|   1|   1|  23|         [23]|
|   1|   2|  63|     [23, 63]|
|   1|   3|null|    [23, 63,]|
|   1|   4|  32|[23, 63,, 32]|
|   2|   2|  56|         [56]|
+----+----+----+-------------+