在 pyspark 的范圍整數列表中查找整數列表的最佳方法是什么

Question

我有一個范圍列表，如下所示：

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

現在我需要做的是，對於“值列表”中的每個值，我需要獲取“范圍列表”中行的索引，以便該值介於“從”和“到”之間在“范圍列表”中的那一行。

假設：

• 兩個列表都是數據幀（我更願意這樣做是 pyspark）
• “范圍列表”中可能存在重疊，在這種情況下返回所有可能的范圍索引
• From 和 To 值具有包容性
• “范圍列表”可以有超過 1 億行
• “值列表”的長度將始終小於“范圍列表”。 （雖然我的圖片沒有顯示）

用例子解釋上面的，

“值列表”中的第一項是“17”，它落入的范圍將在“范圍列表”的索引 2 處，因為 17 介於 15 和 19 之間。

“值列表”中的第二項是“51”，它落入的范圍將在“范圍列表”的索引 9 處，因為 51 介於 50 和 54 之間。

我想知道是否有任何快速/有效的方法可以在 pyspark 的范圍內進行查找。

Answer 1

對於 Pyspark 1.6+：使用非對等連接來查找匹配范圍。 然后，由於您的行數將比最初的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|           []|
# +-----+-------------+

Answer 2

這應該適用於 Spark 2.1+：

import pyspark.sql.functions as F

df_ranges = df_ranges.withColumn("id", F.monotonically_increasing_id())
df = df_vals.crossJoin(df_ranges)
df = df.where((F.col('val') >= F.col('from')) & (F.col('value') <= F.col('to')))
df = df.groupby('val').agg(F.collect_set('id').alias('ids'))
df.show()

        +------+-----------+
        |   val|        ids|
        +------+-----------+
        |    17|        [2]|
        |    51|        [9]|
        +------+-----------+