如何将参数传递给用户定义的函数？

Question

我有一个用户自定义函数：

calc = udf(calculate, FloatType())

param1 = "A"

result = df.withColumn('col1', calc(col('type'), col('pos'))).groupBy('pk').sum('events')

def calculate(type, pos):
   if param1=="A":
       a, b = [ 0.05, -0.06 ]
   else:
       a, b = [ 0.15, -0.16 ]
   return a * math.pow(type, b) * max(pos, 1)

我需要向这个 udf 传递一个参数 param1。我该怎么做？

- Dinosaurius

---

你可以重载UDF。 - Maulik Doshi

@MaulikDoshi：你能展示一下如何做吗？ - Dinosaurius

你能展示一下UDF的语法吗？ - Maulik Doshi

@MaulikDoshi：这是 def calculate(type, pos):。你的意思是需要看完整的函数吗？ - Dinosaurius

请注意我的更新。 - Dinosaurius

Answer 1

你可以像这样将 lit 或 typedLit 用作你的 udf 的参数:

在Python中:

from pyspark.sql.functions import udf, col, lit
mult = udf(lambda value, multiplier: value * multiplier)
df = spark.sparkContext.parallelize([(1,),(2,),(3,)]).toDF()
df.select(mult(col("_1"), lit(3)))

在Scala中:

import org.apache.spark.sql.functions.{udf, col, lit}
val mult = udf((value: Double, multiplier: Double) => value * multiplier)
val df = sparkContext.parallelize((1 to 10)).toDF
df.select(mult(col("value"), lit(3)))

- Paul V

---

那么，在我的情况下我可以使用 lit(param1) 吗？ - Dinosaurius

是的！lit将只作为另一列具有唯一值的列。 - Paul V

抱歉我用Scala而不是Python给你答案，但是思路是一样的！ - Paul V

那么，lit会在df中创建一个额外的列吗？ - Dinosaurius

不确定我是否理解了您的关注点，但这不会改变您的原始数据框。它将使用lit的值并像是一个单独的列一样运作。 - Paul V

让我来测试一下。我使用的是 PySpark，而不是 Scala。 - Dinosaurius