如何在spark中获取两个最接近的值 scala DataFrame

Question

大家好我是 Spark scala 的新手。我想使用 spark scala 通过分区找到最接近的值。我的输入是这样的：

例如第一行：值 1 在 value2 列中介于 2 和 7 之间

+--------+----------+----------+
|id      |value1    |value2    |
+--------+----------+----------+
|1       |3         |1         |
|1       |3         |2         |
|1       |3         |7         |

|2       |4         |2         |
|2       |4         |3         |
|2       |4         |8         |

|3       |5         |3         |
|3       |5         |6         |
|3       |5         |7         |
|3       |5         |8         |

我的 output 应该是这样的：

+--------+----------+----------+
|id      |value1    |value2    |
+--------+----------+----------+
|1       |3         |2         |
|1       |3         |7         |

|2       |4         |3         |
|2       |4         |8         |

|3       |5         |3         |
|3       |5         |6         |

有人可以指导我如何解决这个问题。

Answer 1

我没有提供您似乎想要学习的代码答案，而是为您提供了伪代码和参考，以便您自己找到答案。

1. 对元素进行 分组(select id, value1)（使用collect_list在 value2 上聚合），以便您可以将所有 value2 收集到一个数组中。
2. select (id, and (add( concat ) value1 to the collect_list array)) 排序数组。
3. 在数组中找到( array_position ) value1。
4. splice数组。 检索 ( array_position ) 结果之前的值和之后的值
5. 如果数组少于 3 个元素做错误处理
6. 现在数组中的最后一个值和数组中的第一个值是你的“最接近的数字”。

Answer 2

为此，您将需要window 个函数。

val window = Window
  .partitionBy("id", "value1")
  .orderBy(asc("value2"))

val result = df
  .withColumn("prev", lag("value2").over(window))
  .withColumn("next", lead("value2").over(window))
  .withColumn("dist_prev", col("value2").minus(col("prev")))
  .withColumn("dist_next", col("next").minus(col("value2")))
  .withColumn("min", min(col("dist_prev")).over(window))
  .filter(col("dist_prev") === col("min") || col("dist_next") === col("min"))
  .drop("prev", "next", "dist_prev", "dist_next", "min")

我还没有测试过它，所以将它更多地看作是概念的说明，而不是现成可用的示例。

这是这里发生的事情：

• 首先，创建一个window来描述您的分组规则：我们希望行按前两列分组，并按每组中的第三列排序。
• 接下来，将prev和next列添加到 dataframe，它们分别包含组中上一行和下一行的value2列的值。 （组中第一行的prev将是 null，最后一行的next将是 null - 没关系）。
• 添加dist_prev和dist_next分别包含value2与prev一个值和next值之间的距离。 （请注意，每一行的dist_prev将与前一行的dist_next具有相同的值）。
• 在每个组中找到dist_prev的最小值，并将其添加为min列（注意， dist_next的最小值在构造上是相同的，因此我们这里只需要一列）。
• 筛选行，选择dist_next或dist_prev中具有最小值的行。 除非有多行彼此之间的距离相同，否则这会找到最紧密的对 - 这种情况未在您的问题中考虑，因此我们不知道在这种情况下您想要什么样的行为。 此实现将简单地返回所有这些行。
• 最后，删除所有添加到 dataframe 的额外列，使其恢复到原来的形状。