Implement Ball tree on pyspark for nearest position between two dataframes

I need to retrieve, to a dataframe presenting individual positions (latitude/longitude) of 6 million rows, the information about the nearest store (within 3km).

So I thought about the BallTree of scikit learn, building a UDF... Except that the information is presented in two different dataframes (and that the broadcasting() doesn't work in my working environment).

Do you have any tips to perform this join?

Find hereunder a replicable exemple and what I tried:

from pyspark.sql import functions as f
from sklearn.neighbors import BallTree
from pyspark.sql.functions import udf, pandas_udf
from pyspark.sql.types import StructType, StructField, FloatType, StringType, IntegerType
import numpy as np


individual_positions = spark.createDataFrame([
    (1, 52.03078, 0.29924), 
    (2, 55.27098, -4.35238), 
    (3, 56.40228, -4.62219), 
    (4, 55.32196, -1.90366), 
    (5, 52.73526, -0.25418)], 
    ['ind','lat','lon'])


df_store = spark.createDataFrame([
    (1, 52.03180, 0.30000), 
    (2, 57.22800, -3.22241), 
    (3, 55.27098, -4.32380), 
    (4, 58.41750, -3.67292), 
    (5, 57.44247, -5.18581)], 
    ['Store_ID','lat_store','lon_store'])

# deg radians
df_store = df_store.withColumn('lat_store_rad', f.radians(f.col('lat_store')))
df_store = df_store.withColumn('lon_store_rad', f.radians(f.col('lon_store')))

Define the function:

def get_nearest(src_points, candidates = df_store, k_neighbors=1):
        """Find nearest neighbors for all source points from a set of candidate points"""
        candidates_a = np.array(candidates.select('lat_store_rad', 'lon_store_rad').collect())
        src_points = np.array(src_points)
        tree = BallTree(candidates_a, leaf_size=15, metric='haversine')
        distances, indices = tree.query(src_points, k=k_neighbors)
        distances = distances.transpose()
        indices = indices.transpose()
        closest = indices[0]
        closest_dist = distances[0]
        closest_store = candidates.filter(f.col('index')==int(closest[0]))
        earth_radius = 6371
        closest_store = closest_store.withColumn('distance', f.lit(closest_dist[0]) * f.lit(earth_radius) )
        # remove the closest if distance > 3km
        closest_store = closest_store.withColumn('Store_ID', f.when(f.col('distance')>3,'999999999').otherwise(f.col('Store_ID')))
        return (closest_store.select('Store_ID').collect())

# register udf
get_nearest_udf = f.udf(get_nearest, IntegerType())

Apply it:

individual_positions = individual_positions.withColumn('lat_rad', f.radians(f.col('lat')))
individual_positions = individual_positions.withColumn('lon_rad', f.radians(f.col('lon')))
# apply udf
individual_positions = individual_positions.withColumn('Store_ID', get_nearest_udf(src_points = individual_positions.select('lat_rad', 'lon_rad').collect()))

The function get_nearest() is working well (without spark / udf registration), but as we have two datasets it's not working in the spark/udf logic get_nearest_udf().

The expected output would be the individual_positions added with the Store_ID information.

Many thanks for your advices