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，并在列null的值被删除时将列C的值作为序列包含进来：

+----+----+----+------------+
|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

由于collect_list自动删除了所有null，因此一种方法是在应用该方法之前，将null临时替换为指定的数字，例如Int.MinValue，然后使用一个UDF，以便随后将这些数字恢复回null：

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")

def replaceWithNull(n: Int) = udf( (arr: Seq[Int]) =>
  arr.map( i => if (i != n) Some(i) else None )
)

df.withColumn( "colD", replaceWithNull(Int.MinValue)(
    collect_list(when($"colC".isNull, Int.MinValue).otherwise($"colC")).
      over(Window.partitionBy("colA").orderBy("colB"))
  )
).show
// +----+----+----+------------------+
// |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]|
// +----+----+----+------------------+

Answer 2

正如LeoC提到的collect_list将丢弃空值。似乎有解决此问题的方法。通过将每个标量包装到数组中，紧跟collect_list将得到[[23], [63], [], [32]]，然后对它进行flatten时，将得到[23, 63,, 32]。数组中缺少的这些值为null。

我相信

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

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]|
+----+----+----+-------------+

Scala：Window上的collect_list（），保留空值

2 个答案: