Spark DataFrame中基于w/h列范围计算color列区域求和

问题背景 #

现有一个代表灰度图像的Spark DataFrame，结构及前20行数据如下：

+---+---+-----+
|  w|  h|color|
+---+---+-----+
|  0|  0|255.0|
|  0|  1|255.0|
|  0|  2|255.0|
|  0|  3|255.0|
|  0|  4|255.0|
|  0|  5|255.0|
|  0|  6|255.0|
|  0|  7|255.0|
|  0|  8|255.0|
|  0|  9|255.0|
|  0| 10|255.0|
|  0| 11|255.0|
|  1|  0|255.0|
|  1|  1|255.0|
|  1|  2|255.0|
|  1|  3|255.0|
|  1|  4|255.0|
|  1|  5|255.0|
|  1|  6|255.0|
|  1|  7|255.0|
+---+---+-----+
top 20 rows

需求 #

对DataFrame中的每一行，计算满足以下条件的所有行的color列之和：

• w值在当前行w到w+num1范围内
• h值在当前行h到h+num2范围内

无效尝试 #

曾尝试用窗口函数实现，但写法有误，无法达到预期效果：

val windowW = Window.rangeBetween(Window.currentRow, Window.currentRow + num1)
val windowH = Window.rangeBetween(Window.currentRow, Window.currentRow + num2)

df.withColumn("color_sum", sum(col("color")).over(col("w").windowW and col("h").windowH))

预期输出示例 #

当num1和num2均为1时，第一行的计算结果如下：

+---+---+-----+----------+
|  w|  h|color|sum(color)|
+---+---+-----+----------+
|  0|  0|255.0|      1020|
+---+---+-----+----------+

该结果对应求和的行是(0, 0)、(0, 1)、(1, 0)、(1, 1)；行(1, 1)的求和范围则是(1, 1)、(1, 2)、(2, 1)、(2, 2)。

实现方法提示 #

1. 自连接方案（最通用） #

将原DataFrame与自身做左连接，连接条件直接匹配目标范围，之后分组求和：

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

val num1 = 1
val num2 = 1

val result = df.as("a")
  .join(df.as("b"),
    col("b.w").between(col("a.w"), col("a.w") + num1) &&
    col("b.h").between(col("a.h"), col("a.h") + num2),
    "left"
  )
  .groupBy("a.w", "a.h", "a.color")
  .agg(sum("b.color").alias("sum(color)"))
  .orderBy("a.w", "a.h")

这种方式逻辑直接，不需要依赖数据排序，适用于所有情况。

2. 窗口函数适配方案 #

原窗口函数写法错误，Spark窗口无法同时对两个字段设置独立的范围条件。若要使用窗口函数，需先确保数据按w和h排序，然后定义窗口并过滤，但局限性较强：

• 先按w分区、h排序，再在窗口内筛选h的范围，但无法处理w的跨分区范围，仅适用于num1=0的场景。
• 若num1不为0，窗口函数难以覆盖跨w值的行，因此自连接是更可靠的选择。

3. 性能优化 #

• 若数据量较大，可先按w字段对DataFrame进行分区，减少连接时的数据shuffle。
• 对w和h字段建立索引，提升连接条件的匹配效率。

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