python - 在pyspark的范围整数列表中查找整数列表的最佳方法是什么

Question

我有一个范围列表，如下所示：

我有一个值列表，如下所示：

现在我需要做的是，对于“值列表”中的每个值，我需要获取“范围列表”中行的索引，以便该值介于“从”和“到”之间在“范围列表”中的该行。

假设：

• 两个列表都是数据框（我更愿意这样做是 pyspark）
• “范围列表”中可能有重叠，在这种情况下返回所有可能的范围索引
• From 和 To 值包含在内
• “范围列表”可以有超过 1 亿行
• “值列表”的长度将始终小于“范围列表”。（虽然我的图片没有显示）

用例子解释上面的内容，

“值列表”中的第一项是“17”，它所属的范围将在“范围列表”的索引 2 处，因为 17 介于 15 和 19 之间。

“值列表”中的第二项是“51”，它所属的范围将在“范围列表”的索引 9 处，因为 51 介于 50 和 54 之间。

我想知道是否有任何快速/有效的方法来查找 pyspark 中的范围。

标签： pythonpysparkdatabricks

Answer 1

对于 Pyspark 1.6+：使用非 equi 连接来查找匹配范围。values然后，由于范围可能重叠，因此您将拥有比最初更多的行，按值分组并调用collect_list。

我添加了一个不在范围内的值的示例。如果您不关心它们，请更改数据框的连接方式left。

ranges = spark.createDataFrame(((1, 0, 5), (2, 4, 7), (3, 8, 10)),
                               schema=("index", "from", "to"))
values = spark.createDataFrame(((-1,), (3,), (5,), (8,), (100,)),
                               schema=("value",))
df2 = ranges.join(values,
                  values.value.between(ranges["from"], ranges["to"]),
                  how="right")  # change to left outliers should be ignored
df2.groupBy("value").agg(collect_list("index").alias("range_indices")).show()
# +-----+-------------+
# |value|range_indices|
# +-----+-------------+
# |    5|       [1, 2]|
# |  100|           []|
# |    3|          [1]|
# |    8|          [3]|
# |   -1|           []|
# +-----+-------------+