Pyspark，如何使用udf计算泊松分布？

我有一个如下所示的数据帧：

df_schema = StructType([StructField("date", StringType(), True),\
                              StructField("col1", FloatType(), True),\
                             StructField("col2", FloatType(), True)])
df_data = [('2020-08-01',0.09,0.8),\
                 ('2020-08-02',0.0483,0.8)]
rdd = sc.parallelize(df_data)
df = sqlContext.createDataFrame(df_data, df_schema)
df = df.withColumn("date",to_date("date", 'yyyy-MM-dd'))
df.show() 

+----------+------+----+
|      date|  col1|col2|
+----------+------+----+
|2020-08-01|  0.09| 0.8|
|2020-08-02|0.0483| 0.8|
+----------+------+----+

我想用col1和col2计算泊松CDF

我们可以很容易地从scipy.stats导入熊猫数据框中的泊松，但我不知道如何处理pyspark

prob=poisson.cdf（x，mu），其中在我们的例子中x=col1，mu=col2

尝试1：

from scipy.stats import poisson
from pyspark.sql.functions import udf,col
def poisson_calc(a,b):
    return poisson.cdf(a,b,axis=1)

poisson_calc = udf(poisson_calc, FloatType())

df_new = df.select(
  poisson_calc(col('col1'),col('col2')).alias("want") )

df_new.show()

---

给我一个错误：TypeError:_parse_args（）得到一个意外的关键字参数“axis”

我发现您的尝试有一些问题

• 您将

  udf

  命名为与基础函数相同的名称。令人惊讶的是，这本身并不是一个问题，但我会避免它

• scipy.stats.poisson.cdf

• 您必须显式地将输出转换为

  float

  ，否则会遇到

解决所有问题时，以下各项应起作用：

from scipy.stats import poisson
from pyspark.sql.functions import udf,col

def poisson_calc(a,b):
    return float(poisson.cdf(a,b))

poisson_calc_udf = udf(poisson_calc, FloatType())

df_new = df.select(
  poisson_calc_udf(col('col1'),col('col2')).alias("want") 
)

df_new.show()
#+----------+
#|      want|
#+----------+
#|0.44932896|
#|0.44932896|
#+----------+

---

嗨，pault，我正在使用你的udf方法进行泊松计算。而且我认为这需要很长的时间来处理，因为我有超过数百万的数据。有什么方法可以优化它吗？@dakjdlajsl您可以尝试使用本机spark重写函数，以避免使用

udf

，但我不知道这有多困难（如果可能的话）。如果您对近似解决方案没有意见，另一种方法是加载泊松分布表并进行连接，而不是精确的CDF计算。