如何优化 Spark SQL 中的非等连接？ [复制]

Question

我有两个数据框需要使用具有两个连接谓词的非等连接（即不等连接）连接在一起。

一个数据框是一个直方图DataFrame[bin: bigint, lower_bound: double, upper_bound: double]
另一个数据框是观察集合DataFrame[id: bigint, observation: double]

我需要确定每个观察值属于直方图的哪个 bin，如下所示：

observations_df.join(histogram_df, 
    (
        (observations_df.observation >= histogram_df.lower_bound) &
        (observations_df.observation < histogram_df.upper_bound)
    )
   )

基本上它很慢，我正在寻找一些关于如何让它更快的建议。

下面是一些演示问题的示例代码。 observations_df 包含 100000 行，当 histogram_df 中的行数变得适当大（比如说number_of_bins = 500000）时，它变得非常非常慢，我确信这是因为我正在做一个非等值连接。如果您运行此代码，然后使用 number_of_rows 的值，从较低的值开始，然后增加，直到缓慢的性能很明显

from pyspark.sql.functions import lit, col, lead
from pyspark.sql.types import *
from pyspark.sql import SparkSession
from pyspark.sql.types import *
from pyspark.sql.functions import rand
from pyspark.sql import Window
spark = SparkSession \
    .builder \
    .getOrCreate()

number_of_bins = 500000

bin_width = 1.0 / number_of_bins
window = Window.orderBy('bin')
histogram_df = spark.range(0, number_of_bins)\
    .withColumnRenamed('id', 'bin')\
    .withColumn('lower_bound', 0 + lit(bin_width) * col('bin'))\
    .select('bin', 'lower_bound', lead('lower_bound', 1, 1.0).over(window).alias('upper_bound'))
observations_df = spark.range(0, 100000).withColumn('observation', rand())
observations_df.join(histogram_df, 
        (
            (observations_df.observation >= histogram_df.lower_bound) &
            (observations_df.observation < histogram_df.upper_bound)
        )
       ).groupBy('bin').count().head(15)

Answer 1

不建议将不等连接用于 spark join。通常，我会生成一个新列作为此类操作的连接键。 但是，对于您的情况，您不需要加入来确定每个观察值属于直方图的哪个 bin，因为可以预先计算每个 bin 的上限和下限，并且您可以使用观察值计算 bin。

您可以做的是编写一个 UDF，它会为您找到 bin 并将 bin 作为新列返回。 你可以参考pyspark: passing multiple dataframe fields to udf