pyspark 自定义函数

一：自定义函数的一般流程

# 1.创建普通的python函数

def to_upper(s):
        if s is not None:
            return s.upper()

# 2.注册自定义函数
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# 根据python的返回值类型定义好spark对应的数据类型
# python函数中返回的是string，对应的pyspark是StringType
toDateUDF=udf(to_upper, StringType())  

# 使用自定义函数
df1.withColumn('color',toDateUDF('color')).show()
 

二：简单的自定义函数
最简单的就是通过lambda函数，不需要定义返回值类型，可以直接使用

slen = udf(lambda s: len(s), IntegerType())#这里用匿名函数

三：自定义函数，需要定义返回值类型

slen = udf(lambda s: len(s), IntegerType())#这里用匿名函数
def to_upper(s):
    if s is not None:
        return s.upper()
def add_one(x):
    if x is not None:
        return x + 1
df = spark.createDataFrame([(1, "John Doe", 21)], ("id", "name", "age"))
add_udfs = udf(add_one,IntegerType())#定义返回值类型
df.select(slen('name').alias("slen(name)"),to_upper('name'),add_udfs('age')).show() 

补充：

自定义函数的重点在于定义返回值类型的数据格式，其数据类型基本都是从from pyspark.sql.types import * 导入，常用的包括：

StructType()：结构体
StructField()：结构体中的元素
LongType()：长整型
StringType()：字符串
IntegerType()：一般整型
FloatType()：浮点型

参考：Spark UDF使用详解及代码示例

 

 

 

后来还看到一种方法，试过也可以用：编写udf

from pyspark.sql.types import *
def pow1(m,n):
return float(m)**float(n)
udf = spark.udf
udf.register('pow1',pow1,returnType=DoubleType())
df = spark.range(0,10,2,3)
df.createOrReplaceTempView('A')
spark.sql('select pow1(id,2) from A').show()
spark.sql('select pow1(id,3) from A').show()