如何在PySpark DataFrame中实现基于向下最近匹配值的关联查询

Absolutely! You can replicate that SQL subquery logic in PySpark using DataFrame operations. Below are two practical approaches, tailored to different data size scenarios, that will give you the exact same result as your SQL query.

First: Set Up Test Data #

Let's start by creating the sample DataFrames to match your tables, so you can run and test the code immediately:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# Initialize SparkSession
spark = SparkSession.builder.appName("NearestLowerMatch").getOrCreate()

# Create tblActuals DataFrame
actuals_data = [
    ("09/02/2020", 14.1),
    ("10/02/2020", 15.3),
    ("11/02/2020", 12.2),
    ("12/02/2020", 12.4),
    ("13/02/2020", 12.5),
    ("14/02/2020", 11),
    ("15/02/2020", 14.6)
]
tblActuals = spark.createDataFrame(actuals_data, ["Date", "Temperature"])

# Create tblCoefficients DataFrame
coeffs_data = [
    (10.5, 0.997825593),
    (11, 0.997825593),
    (11.5, 0.997663198),
    (12, 0.997307614),
    (12.5, 0.996848773),
    (13, 0.996468537),
    (13.5, 0.99638519),
    (14, 0.996726301),
    (14.5, 0.997435894),
    (15, 0.998311153),
    (15.5, 0.999135509)
]
tblCoefficients = spark.createDataFrame(coeffs_data, ["Metric", "Coefficient"])

Approach 1: Cross Join + Window Function (Good for Small Datasets) #

This method mirrors your SQL logic closely: we join each temperature with all qualifying coefficients (where Metric <= Temperature), then use a window function to pick the largest matching Metric for each row.

# Step 1: Cross join and filter to keep only valid (Metric <= Temperature) pairs
cross_joined = tblActuals.crossJoin(tblCoefficients).filter(F.col("Metric") <= F.col("Temperature"))

# Step 2: Define a window to group by Date/Temperature, sorted by Metric descending
window_spec = Window.partitionBy("Date", "Temperature").orderBy(F.desc("Metric"))

# Step 3: Add row numbers and keep only the first row (largest Metric)
result = cross_joined.withColumn("row_num", F.row_number().over(window_spec)) \
                     .filter(F.col("row_num") == 1) \
                     .drop("row_num", "Metric") \
                     .orderBy("Date")

# View the result
result.show()

Approach 2: Broadcast Join + Aggregation (More Efficient for Small Dimension Tables) #

If your tblCoefficients is a small dimension table (like in your example), using a broadcast join will reduce data shuffling. We'll collect and sort valid coefficient pairs, then pick the largest matching one.

# Step 1: Broadcast the small coefficients table to optimize the join
broadcast_coeffs = F.broadcast(tblCoefficients)

# Step 2: Join and collect all valid (Metric, Coefficient) pairs per temperature
grouped = tblActuals.join(broadcast_coeffs, F.col("Metric") <= F.col("Temperature"), "left") \
                    .groupBy("Date", "Temperature") \
                    .agg(F.array_sort(F.collect_list(F.struct(F.col("Metric"), F.col("Coefficient")))).alias("sorted_coeffs"))

# Step 3: Extract the last element (largest Metric) from the sorted list
result = grouped.withColumn("max_metric_pair", F.element_at(F.col("sorted_coeffs"), -1)) \
                .select("Date", "Temperature", F.col("max_metric_pair.Coefficient").alias("Coefficient")) \
                .orderBy("Date")

# View the result
result.show()

What the Output Looks Like #

Both methods will produce this result, matching your SQL query exactly:

+-----------+-----------+-----------+
|       Date|Temperature|Coefficient|
+-----------+-----------+-----------+
|09/02/2020|       14.1|0.996726301|
|10/02/2020|       15.3|0.998311153|
|11/02/2020|       12.2|0.997307614|
|12/02/2020|       12.4|0.997307614|
|13/02/2020|       12.5|0.996848773|
|14/02/2020|       11.0|0.997825593|
|15/02/2020|       14.6|0.997435894|
+-----------+-----------+-----------+

Notes on Performance #

• Use Approach 1 only if both tables are small—cross joins can get expensive with large datasets.
• Use Approach 2 when tblCoefficients is small (a dimension table), as broadcasting avoids expensive data shuffling across nodes.

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