Spark:在UDF或映射函数中连接

我必须编写一个复杂的UDF,其中我必须使用不同的表进行连接,并返回匹配数。实际的用例要复杂得多,但我在这里将用例简化为最小的可重复代码。这是UDF代码。

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

当我使用下面的代码定义UDF时,会得到一长串控制台错误:

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

错误的最后一行是:

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

我想知道最好的方法是什么。我也试过 map 这样地:

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