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

时间:2017-11-13 09:31:21

标签: python apache-spark pyspark

我有一个用户定义的功能:

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)

我需要将参数param1传递给此udf。我该怎么办?

3 个答案:

答案 0 :(得分:5)

您可以使用littypedLit作为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)))

答案 1 :(得分:0)

如果你坚持,

def calculate(param1):
    return param1 * param1

sqlContext.udf.register("square", calculate)

答案 2 :(得分:-1)

我不确定这是否有用......但你可以试试吗?

calc = udf(calculate, FloatType())

param1 = "A"
#If not !A Give some dummy value other than A

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

def calculate(type, pos,param1):
   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)