如何在 Spark 数据集中计算组的行数并添加计数为零的组？答案

【问题标题】：How to count rows of a group and add groups of count zero in Spark Dataset?如何在 Spark 数据集中计算组的行数并添加计数为零的组？
【发布时间】：2021-07-29 12:48:17
【问题描述】：

我在 Scala 中使用 Spark，Datasets，我面临以下情况。（编辑：我已将示例重写为更具体。这是为了放入一个方法中并与不同的数据集重用。）

假设我们有以下数据集，表示在特定时间戳区域 a 到区域 b 之间的移动：

case class Move(id: Int, ra: Int, rb: Int, time_a: Timestamp, time_b: Timestamp)

val ds = Seq(
  (1, 123, 125, "2021-07-25 13:05:20", "2021-07-25 15:05:20"),
  (2, 470, 125, "2021-07-25 00:05:20", "2021-07-25 02:05:20"),
  (1, 470, 125, "2021-07-25 02:05:20", "2021-07-25 04:05:20"),
  (3, 123, 125, "2021-07-26 00:45:20", "2021-07-26 02:45:20"),
  (3, 125, 123, "2021-07-28 16:05:20", "2021-07-28 18:05:20"),
  (1, 125, 123, "2021-07-29 20:05:20", "2021-07-30 01:05:20")
).toDF("id", "ra", "rb", "time_a", "time_b")
  .withColumn("time_a", to_timestamp($"time_a"))
  .withColumn("time_b", to_timestamp($"time_b"))
  .as[Move]

我想要的是计算每天的移动次数，这样我就可以得到一个带有一些时间划分的起点-终点矩阵。此外，我希望能够沿这些时间划分应用一些标签/类别。（如何进行此分类的代码无关紧要，我们可以假设它在这里有效，并且每个类别都在val categs: Seq[TimeCategory] 中进行了描述）。可以这样做：

case class Flow(ra: Int, rb: Int, time: Long, categ: Long, inflow: Long, outflow: Long)

val Row(t0: Timestamp, tf: Timestamp) = ds.select(min($"time_a"), max($"time_b")).head
val ldt0 = t0.toLocalDateTime()
val ldtf = tf.toLocalDateTime()
val unit = ChronoUnit.DAYS
val fmds = ds.map({
  case Move(id, ra, rb, ta, tb) => {
    val ldta = ta.toLocalDateTime()
    val ldtb = tb.toLocalDateTime()
    (
      id, ra, rb,
      unit.between(ldt0, ldta),
      unit.between(ldt0, ldtb),
      TimeCategory.encode(categs, ldta),
      TimeCategory.encode(categs, ldtb)
    )
  }
})
val outfds = fmds
  .groupBy($"ra", $"rb", col("time_a").as("time"), col("categ_a").as("categ"))
  .agg(count("id").as("outflow"))
  .withColumn("inflow", lit(0))
val infds = fmds
  .groupBy($"ra", $"rb", col("time_b").as("time"), col("categ_b").as("categ"))
  .agg(count("id").as("inflow"))
  .withColumn("outflow", lit(0))
val fds = outfds
  .unionByName(infds)
  .groupBy("ra", "rb", "time", "categ")
  .agg(sum("outflow").as("outflow"), sum("inflow").as("inflow"))
  .orderBy("ra", "rb", "time", "categ")
  .as[Flow]

这将产生以下结果：

+---+---+----+-----+-------+------+
| ra| rb|time|categ|outflow|inflow|
+---+---+----+-----+-------+------+
|123|125|   0|   19|      1|     1|
|123|125|   1|   13|      1|     0|
|123|125|   1|   14|      0|     1|
|125|123|   3|   11|      1|     0|
|125|123|   3|   12|      0|     1|
|125|123|   4|   12|      1|     0|
|125|123|   5|   14|      0|     1|
|470|125|   0|   21|      1|     0|
|470|125|   0|   22|      1|     2|
+---+---+----+-----+-------+------+

问题是，如果我想获得每对区域每天的平均流入量，就会缺少流入量 = 0 的许多天。例如，如果完成以下 agg：

fds.groupBy("ra", "rb", "categ").agg(avg("inflow"))
// output:
+---+---+-----+-----------+
| ra| rb|categ|avg(inflow)|
+---+---+-----+-----------+
|123|125|   14|        1.0|
|123|125|   13|        0.0|
|470|125|   21|        0.0|
|125|123|   12|        0.5|
|125|123|   11|        0.0|
|123|125|   19|        1.0|
|125|123|   14|        1.0|
|470|125|   22|        2.0|
+---+---+-----+-----------+

对于125 -> 123 和类别12，平均值为0.5，但考虑到整个数据集的开始和结束时间戳，类别12 应该有5 天，而不仅仅是2。平均值应该是1 / 5 = 0.2。这第二种情况是我要计算的。考虑到我也想计算其他 agg 函数（如 stddev），我想最灵活的选择是“填充”值应该为零的行。就性能/可扩展性而言，最好的方法是什么？

（编辑：我也想到了一个更好的方法。） 到目前为止，我想到的一个可能的解决方案是创建另一个带有“时隙”的DataFrame（在在这种情况下，每个“槽”将是一个包含所有合适类别的日期索引）并再做一个联合，如下所示：

// TimeCategory.encodeAll basically just returns
// every category that suits that timestamp
val timeindex: Seq[(Long, Long)] = (0L to unit.between(ldt0, ldtf))
  .flatMap(i => {
    val t = ldt0.plus(i, unit)
    val categCodes = TimeCategory.encodeAll(categs, unit, t)
    categCodes.map((i, _))
  })
val zds = infds
  .select($"ra", $"rb", explode(typedLit(timeindex)).as("time"))
  .select($"ra", $"rb", col("time")("_1").as("time"), col("time")("_2").as("categ"),
    lit(0).as("outflow"), lit(0).as("inflow"))
val fds = outfds
  .unionByName(infds)
  .unionByName(zds)
  .groupBy("ra", "rb", "time", "categ")
  .agg(sum("outflow").as("outflow"), sum("inflow").as("inflow"))
  .orderBy("ra", "rb", "time", "categ")
  .as[Flow]
  .show()
fds.groupBy("ra", "rb", "categ")
  .agg(avg("inflow"))
  .where($"avg(inflow)" > 0.0)
  .show()

结果：

+---+---+----+-----+-------+------+
| ra| rb|time|categ|outflow|inflow|
+---+---+----+-----+-------+------+
|123|125|   0|   17|      0|     0|
|123|125|   0|   18|      0|     0|
|123|125|   0|   19|      1|     1|
|123|125|   0|   20|      0|     0|
|123|125|   0|   21|      0|     0|
|123|125|   0|   22|      0|     0|
|123|125|   1|    9|      0|     0|
|123|125|   1|   10|      0|     0|
|123|125|   1|   11|      0|     0|
|123|125|   1|   12|      0|     0|
|123|125|   1|   13|      1|     0|
|123|125|   1|   14|      0|     1|
|123|125|   2|    9|      0|     0|
|123|125|   2|   10|      0|     0|
|123|125|   2|   11|      0|     0|
|123|125|   2|   12|      0|     0|
|123|125|   2|   13|      0|     0|
|123|125|   2|   14|      0|     0|
|123|125|   3|    9|      0|     0|
|123|125|   3|   10|      0|     0|
+---+---+----+-----+-------+------+
only showing top 20 rows

+---+---+-----+-----------+
| ra| rb|categ|avg(inflow)|
+---+---+-----+-----------+
|123|125|   14|        0.2|
|125|123|   12|        0.2| // <- the correct avg
|123|125|   19|        1.0|
|125|123|   14|        0.2|
|470|125|   22|        2.0|
+---+---+-----+-----------+

是否有可能改善这一点？

【问题讨论】：

没有任何表join，在执行groupBy/agg 之后如何生成数据集involving other columns besides a time "index"？特别是，对于由“间隔日期”组成的行，这些“其他列”会是什么样子？也许一些示例输出数据有助于澄清。
我建议探索窗口函数：window function
只是为了确定：agg(count("id")) 完成后您的数据有多大？因为如果不是数十亿行，您可以直接在带有 scala 集合的驱动程序上收集它并计算平均值，填充孔。我这么说是因为很多人忘记了一旦数据减少到几MB，使用像spark这样的复杂分布式框架就有点过头了。之后您始终可以重新创建 spark 数据集，例如使用 spark 在某处显示或写入。
@Juh_ 它旨在成为库中的一种方法，因此它也应该适用于更大的数据集。但是，在agg(count(id)) 之后，看看常见的用例，它真的不应该达到数十亿行。
@LeoC 我将使用更有意义的示例和我想到的更新方法来编辑问题

标签： scala apache-spark apache-spark-sql

【解决方案1】：

我只能对您的初始解决方案提出一点改进，将 join 替换为 unionByBame。应该比join效率高，但是我没有测试性能：

val ds = Seq(
  1L -> java.sql.Timestamp.valueOf("2021-07-25 13:05:20"),
  2L -> java.sql.Timestamp.valueOf("2021-07-25 00:05:20"),
  1L -> java.sql.Timestamp.valueOf("2021-07-25 02:05:20"),
  3L -> java.sql.Timestamp.valueOf("2021-07-26 00:45:20"),
  3L -> java.sql.Timestamp.valueOf("2021-07-28 16:05:20"),
  1L -> java.sql.Timestamp.valueOf("2021-07-29 20:05:20")
).toDF("id", "timestamp")

val ddf = ds
  .withColumn("day", dayofmonth($"timestamp"))
  .drop("timestamp")

val r = ds.select(min($"timestamp"), max($"timestamp")).head()
val t0 = r.getTimestamp(0)
val tf = r.getTimestamp(1)

val idf = (t0.getDate() to tf.getDate()).toDF("day")
  .withColumn("id", lit(null).cast(LongType))

ddf.unionByName(idf)
  .groupBy($"day")
  .agg(sum(when($"id".isNotNull, 1).otherwise(0)).as("count"))
  .orderBy("day")
  .show

【讨论】：

我采纳了您在我的解决方案中使用unionByName 的建议；）谢谢