Optimized way to join pyspark dataframes on closest distance - apache-spark

Question

I'm trying to join the GHCN weather dataset and another dataset:
Weather Dataset: (called "weather" in the code)
station_id
date
PRCP
SNOW
SNWD
TMAX
TMIN
Latitude
Longitude
Elevation
State
date_final
CA001010235
19730707
null
null
0.0
0.0
null
48.4
-123.4833
17.0
BC
1973-07-07
CA001010235
19780337
14
8
0.0
0.0
null
48.4
-123.4833
17.0
BC
1978-03-30
CA001010595
19690607
null
null
0.0
0.0
null
48.5833
-123.5167
17.0
BC
1969-06-07
Species Dataset: (called "ebird" in the code where "ebird_id" is unique for each row in the dataset)
speciesCode
comName
sciName
locId
locName
ObsDt
howMany
lat
lng
obsValid
obsReview
locationPrivate
subId
ebird_id
nswowl
Northern Saw-whet Owl
Aegolius acadicus
L787133
Kumdis Slough
2017-03-20 23:15
1
53.7392187
-132.1612358
TRUE
FALSE
TRUE
S35611913
eff-178121-fff
wilsni1
Wilson's Snipe
Gallinago delicata
L1166559
Hornby Island--Ford Cove
2017-03-20 21:44
1
49.4973435
-124.6768427
TRUE
FALSE
FALSE
S35323282
abc-1920192-fff
cacgoo1
Cackling Goose
Branta hutchinsii
L833055
Central Saanich--ȾIKEL (Maber Flats)
2017-03-20 19:24
5
48.5724686
-123.4305167
TRUE
FALSE
FALSE
S35322116
yhj-9102910-fff
Result Expected: I need to join these tables by finding the closest weather station for each row in the species dataset for the same date. So in this example, ebird_id "eff-178121-fff" is closest to the weather station "CA001010235" and the distance is around 20 kms.
speciesCode
comName
sciName
locId
locName
ObsDt
howMany
lat
lng
obsValid
obsReview
locationPrivate
subId
ebird_id
station_id
date
PRCP
SNOW
SNWD
TMAX
TMIN
Latitude
Longitude
Elevation
State
date_final
distance(kms)
nswowl
Northern Saw-whet Owl
Aegolius acadicus
L787133
Kumdis Slough
2017-03-20 23:15
1
53.7392187
-132.1612358
TRUE
FALSE
TRUE
S35611913
eff-178121-fff
CA001010235
20170320
null
null
0.0
0.0
null
48.4
-123.4833
17.0
BC
2017-03-20
20
wilsni1
Wilson's Snipe
Gallinago delicata
L1166559
Hornby Island--Ford Cove
2017-03-20 21:44
1
49.4973435
-124.6768427
TRUE
FALSE
FALSE
S35323282
abc-1920192-fff
CA001010595
20170320
null
null
0.0
0.0
null
48.5833
-123.5167
17.0
BC
2017-03-20
What I have tried so far: I referred to this link and it works for a sample of the datasets but when I tried to run the below code for the entire weather and entire species dataset, the cross join worked but the partitionBy and windows function line was taking too long. I also tried replacing the partionBy and windows function with pyspark SQL queries in case that would be faster but it's still taking time. Is there any optimized way to do this?
join_df = ebird.crossJoin(weather).withColumn("dist_longit", radians(weather["Longitude"]) - radians(ebird["lng"])).withColumn("dist_latit", radians(weather["Latitude"]) - radians(ebird["lat"]))
join_df = join_df.withColumn("haversine_distance_kms", asin(sqrt(
sin(join_df["dist_latit"] / 2) ** 2 + cos(radians(join_df["lat"]))
* cos(radians(join_df["Latitude"])) * sin(join_df["dist_longit"] / 2) ** 2
)
) * 2 * 6371).drop("dist_longit","dist_latit")
W = W.partitionBy("ebird_id")
result = join_df.withColumn("min_dist", min(join_df['haversine_distance_kms']).over(W))\
.filter(col("min_dist") == col('haversine_distance_kms'))
print(result.show(1))
Edit:
Size of the datasets:
print(weather.count()) #output: 8211812
print(ebird.count()) #output: 1564574