Spark: Join within UDF or map function

Question

I have to write a complex UDF, in which I have to do a join with a different table, and return the number of matches. The actual use case is much more complex, but I've simplified the case here to minimum reproducible code. Here is the UDF code.

def predict_id(date,zip):
    filtered_ids = contest_savm.where((F.col('postal_code')==zip)  & (F.col('start_date')>=date))
    return filtered_ids.count()

When I define the UDF using the below code, I get a long list of console errors:

predict_id_udf = F.udf(predict_id,types.IntegerType())

The final line of the error is:

py4j.Py4JException: Method __getnewargs__([]) does not exist

I want to know what is the best way to go about it. I also tried map like this:

result_rdd = df.select("party_id").rdd\
  .map(lambda x: predict_id(x[0],x[1]))\
  .distinct()

It also resulted in a similar final error. I want to know, if there is anyway, I can do a join within UDF or map function, for each row of the original dataframe.

Answer 1

I have to write a complex UDF, in which I have to do a join with a different table, and return the number of matches.

It is not possible by design. I you want to achieve effect like this you have to use high level DF / RDD operators:

df.join(ontest_savm,
    (F.col('postal_code')==df["zip"])  & (F.col('start_date') >= df["date"])
).groupBy(*df.columns).count()