The most performant and proper way to get similarities between large embeddings/vectors - pyspark

Question

I have 2 data frames with 69 & ~230.000 rows and (KEY, Embedding) columns.
The embedding columns are array type and each has a length of 768 which I obtained from a fine-tuned transformer model (i.e., all-distilroberta-v1). What I wanna achieve is to calculate cosine similarities of those embedding pairs from different DFs.
To this end, I calculated the dot products of the embeddings for each raw using UDFs as:
spark = SparkSession.builder \
.appName('app_similarity') \
.master('local[*]') \
.config('spark.sql.execution.arrow.pyspark.enabled', True) \
.config('spark.sql.session.timeZone', 'UTC') \
.config('spark.driver.memory','6G') \
.config('spark.ui.showConsoleProgress', True) \
.config('spark.sql.repl.eagerEval.enabled', True) \
.config('spark.executor.cores', "5") \
.config('spark.executor.memory', "4g") \
.config('spark.num.executors', "5") \
.getOrCreate()
dfSim = df_main.crossJoin(df_other).select("KEY1","KEY2", "Emb","EMBe")
dfSim = dfSim.withColumn("Emb", F.regexp_replace(F.col("Emb"), "\[", "")).withColumn("Emb", F.regexp_replace(F.col("Emb"), "]",""))
dfSim = dfSim.withColumn("EMBe", F.regexp_replace(F.col("EMBe"), "\[", "")).withColumn("EMBe", F.regexp_replace(F.col("EMBe"), "]",""))
dfSim = dfSim.withColumn("Emb", F.split(F.col("Emb"), ",").cast(ArrayType(DoubleType())))
dfSim = dfSim.withColumn("EMBe", F.split(F.col("EMBe"), ",").cast(ArrayType(DoubleType())))
sim_cos = F.udf(lambda x,y : float(np.array(x).dot(np.array(y))), FloatType())
norm = F.udf(lambda x : float(np.linalg.norm(x)), FloatType())
dfSim = dfSim.withColumn("dot", sim_cos(F.col("Emb"),F.col("EMBe")))
dfSim = dfSim.withColumn("Sim", F.round(F.col("dot"),5))
As seen, I use crossjoin which results in all the pairs of keys, i.e., 15 million rows.
With this approach I get some errors like:
when I apply a similarity filter e.g., dfSim.filter(F.col("Sim") > 0.3).show() it works but somehow when I increase the sim score to 0.4 or 0.5 it throws socket.timeout: timed out error
or I cannot upload to sql db etc.
I was wondering if it would be possible to apply an LSH or minHash approach to directly use embeddings and get their neighbors and similarities instead of creating 15 million pairs?
EDIT:
The code below worked. In the beginning, LSH approach did not work due to starting the spark session with "local[*]" causing all the CPUs to be used.
When restricting it by 5 it worked. However, bear in mind, this is still
memory intense process. Playing around 'spark.executor.memory' & 'spark.driver.memory' might be necessary.
from pyspark.ml.feature import BucketedRandomProjectionLSH
from pyspark.ml.linalg import Vectors
from pyspark.sql.functions import col
spark = SparkSession.builder \
.appName('app_similarity') \
.master('local[5]') \
.config('spark.sql.execution.arrow.pyspark.enabled', True) \
.config('spark.sql.session.timeZone', 'UTC') \
.config('spark.driver.memory','8G') \
.config('spark.ui.showConsoleProgress', True) \
.config('spark.sql.repl.eagerEval.enabled', True) \
.config('spark.executor.cores', "3") \
.config('spark.executor.memory', "2g") \
.config('spark.num.executors', "3") \
.getOrCreate()
ut_rdd = df_main.drop("_c0").rdd.map(lambda row:row[0])
emb_rdd = df_main.drop("_c0").rdd.map(lambda row:row[1])
dfM = ut_rdd.zip(emb_rdd.map(lambda x:Vectors.dense(x))).toDF(schema=['ut','emb'])
uto_rdd = df_other.drop("_c0").rdd.map(lambda row:row[0])
embo_rdd = df_other.drop("_c0").rdd.map(lambda row:row[1])
dfO = uto_rdd.zip(embo_rdd.map(lambda x:Vectors.dense(x))).toDF(schema=['ut','emb'])
brp = BucketedRandomProjectionLSH(inputCol="emb", outputCol="hashes", bucketLength=2.0, numHashTables=3)
model = brp.fit(dfM)
dfSimLSH = model.approxSimilarityJoin(dfM, dfO, 1.2, distCol="EuclideanDistance").select(col("datasetA.ut").alias("utm"),
col("datasetB.ut").alias("uto"),
col("EuclideanDistance"))