In general: I'm trying to apply a function from an object that can't be pickled to an entire column of a pyspark data frame. How can I do this without a for loop? Since the object can't be pickled, I can't use withColumn and a udf.
Specifically: I'm trying to get the nearest neighbors for all rows of a Vector column in a pyspark DataFrame. I have a fitted BucketedRandomProjectionLSHModel and I want to use approxNearestNeighbors, but I can't pickle it.
Example:
from pyspark.ml.linalg import Vectors
from pyspark.sql.functions import col
from pyspark.ml.feature import BucketedRandomProjectionLSH
data = [(0, Vectors.dense([-1.0, -1.0 ]),),
(1, Vectors.dense([-1.0, 1.0 ]),),
(2, Vectors.dense([1.0, -1.0 ]),),
(3, Vectors.dense([1.0, 1.0]),)]
df = spark.createDataFrame(data, ["id", "features"])
# Create and fit brp
brp = BucketedRandomProjectionLSH(inputCol='features', outputCol="hashes")
brp_model = brp.fit(df)
# Get the nearest neighbor for each feature vector
df.withColumn("neighbors", brp_model.approxNearestNeighbors(df, col("features"), numNearestNeighbors=5))
This results in TypeError: can't pickle _thread.RLock objects -- and it seems that the brp_model is the locked object that can't be pickled.
How else can I use brp_model to get nearest neighbors for each row of the data frame without resorting to a for loop?
Note that using approxSpatialJoin instead is NOT a solution. (In order to get enough neighbors for every row in my data, I would need to set a threshold that generates thousands of neighbors for other rows.) Thanks!