大数据组件
菜单
本页目录

Spark-SQL-UDF

好的，让我们详细举例讲解Spark中UDF的两种注册方式，分别是spark.udf.register()和pyspark.sql.functions.udf()，并给出适合它们的使用场景和具体代码。

一、spark.udf.register() 注册UDF

SparkSession.udf.register()用于将自定义函数注册为SQL函数。注册之后，我们可以在SQL查询语句中直接使用这个函数。这个方法通常在你想要通过SQL语句查询数据并且希望在SQL中调用UDF时非常有用。

1. 示例：将字符串转换为大写，并在SQL查询中使用

以下是详细步骤和代码示例：

  • 我们会定义一个Python函数，将名字转换为大写。
  • 使用spark.udf.register()注册这个函数，使它可以在SQL查询中使用。
from pyspark.sql import SparkSession
from pyspark.sql.types import StringType

# 初始化SparkSession
spark = SparkSession.builder.appName("Register UDF Example").getOrCreate()

# 定义Python函数，将字符串转换为大写
def to_upper(s):
    return s.upper() if s else None

# 注册UDF，可以在SQL中使用
spark.udf.register("toUpperSQL", to_upper, StringType())

# 创建示例DataFrame
data = [(1, "Alice"), (2, "Bob"), (3, "Cathy")]
df = spark.createDataFrame(data, ["id", "name"])

# 将DataFrame创建为一个临时表，以便使用SQL查询
df.createOrReplaceTempView("people")

# 使用UDF在SQL查询中
result_df = spark.sql("SELECT id, toUpperSQL(name) AS name_upper FROM people")

# 展示结果
result_df.show()

输出：

+---+----------+
| id|name_upper|
+---+----------+
|  1|     ALICE|
|  2|       BOB|
|  3|     CATHY|
+---+----------+

代码详解：

  1. 定义函数并注册：我们定义了一个名为to_upper的函数，并使用spark.udf.register()注册为SQL可用的UDF，命名为toUpperSQL。
  2. 使用SQL查询：我们创建了一个临时表people，然后使用SQL查询时调用了toUpperSQL()来将名字转换为大写。

这种方法在处理数据并且想直接通过SQL查询时非常方便，例如，运行复杂SQL查询并希望在查询过程中使用自定义逻辑。

二、pyspark.sql.functions.udf() 注册UDF

pyspark.sql.functions.udf()是另一种注册UDF的方法，它主要用于DSL风格（DataFrame API）的操作。这种方法不支持直接在SQL查询中使用，而是通过DataFrame操作来处理。

1. 示例：将字符串转换为大写，并在DataFrame操作中使用

以下是详细步骤和代码示例：

  • 我们会定义一个Python函数，将名字转换为大写。
  • 使用pyspark.sql.functions.udf()注册这个函数，使它可以在DataFrame操作中应用。
from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# 初始化SparkSession
spark = SparkSession.builder.appName("UDF Example").getOrCreate()

# 创建示例DataFrame
data = [(1, "Alice"), (2, "Bob"), (3, "Cathy")]
df = spark.createDataFrame(data, ["id", "name"])

# 定义Python函数，将字符串转换为大写
def to_upper(s):
    return s.upper() if s else None

# 使用udf()注册Python函数
to_upper_udf = udf(to_upper, StringType())

# 使用UDF处理DataFrame中的列
df_with_upper = df.withColumn("name_upper", to_upper_udf(df["name"]))

# 展示结果
df_with_upper.show()

输出：

+---+-----+----------+
| id| name|name_upper|
+---+-----+----------+
|  1|Alice|     ALICE|
|  2|  Bob|       BOB|
|  3|Cathy|     CATHY|
+---+-----+----------+

代码详解：

  1. 定义和注册UDF：我们定义了一个名为to_upper的Python函数，并使用pyspark.sql.functions.udf()进行注册，将其转换为Spark中的UDF，这个UDF命名为to_upper_udf。
  2. 在DataFrame中应用UDF：然后我们使用withColumn()方法将to_upper_udf应用于DataFrame的name列，并生成一个新的列name_upper。

这种方法更适合在数据处理流中使用，当你想要使用DataFrame API来做流式的转换和计算时非常方便。

三、总结：spark.udf.register() vs pyspark.sql.functions.udf()

注册方式适用场景示例调用方式
spark.udf.register()SQL查询中调用自定义函数spark.sql("SELECT toUpperSQL(name) FROM people")
pyspark.sql.functions.udf()DataFrame API风格的操作df.withColumn("name_upper", to_upper_udf(df["name"]))

优缺点对比：

  1. spark.udf.register()：

    • 优点：可以直接在SQL查询中使用，方便复杂SQL操作。
    • 缺点：对SQL用户较为友好，但在处理大量流式数据或多个DataFrame转换时，可能不如DSL风格直观。

  2. pyspark.sql.functions.udf()：

    • 优点：适合DataFrame API操作，方便进行链式调用和分布式处理。
    • 缺点：不能直接在SQL语句中使用，因此在某些需要SQL查询的场景下不太适用。

四、UDF 返回值为数组的例子

在一些场景中，我们希望UDF的返回值为一个数组（例如，处理文本时将字符串拆分为多个单词）。

1. 示例：将字符串拆分为单词数组

以下代码演示了如何定义一个返回值为数组的UDF，并在DataFrame中使用它。

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType

# 初始化SparkSession
spark = SparkSession.builder.appName("UDF Array Example").getOrCreate()

# 创建示例DataFrame
data = [(1, "Hello world"), (2, "Spark UDF example"), (3, None)]
df = spark.createDataFrame(data, ["id", "text"])

# 定义Python函数，将字符串拆分为单词数组
def split_to_words(text):
    return text.split(" ") if text else []

# 使用udf()注册Python函数
split_to_words_udf = udf(split_to_words, ArrayType(StringType()))

# 使用UDF处理DataFrame中的列
df_with_words = df.withColumn("words", split_to_words_udf(df["text"]))

# 展示结果
df_with_words.show(truncate=False)

输出：

+---+-----------------+---------------------+
|id |text             |words                |
+---+-----------------+---------------------+
|1  |Hello world      |[Hello, world]       |
|2  |Spark UDF example|[Spark, UDF, example]|
|3  |null             |[]                   |
+---+-----------------+---------------------+

代码详解：

  1. 定义和注册UDF：我们定义了一个函数split_to_words，用于将字符串拆分为单词数组，并使用pyspark.sql.functions.udf()注册为UDF。返回类型指定为ArrayType(StringType())，表示返回值是一个字符串数组。
  2. 应用UDF：我们使用withColumn()方法将UDF应用于text列，生成新的列words。

五、UDF 返回值为字典的例子

在一些场景中，我们希望UDF返回字典（例如，进行文本分析并返回结果的多个特征）。

1. 示例：将字符串分析为字典，包含单词数和字符数

以下代码演示了如何定义一个返回值为字典的UDF，并在DataFrame中使用它。

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import MapType, StringType, IntegerType

# 初始化SparkSession
spark = SparkSession.builder.appName("UDF Dictionary Example").getOrCreate()

# 创建示例DataFrame
data = [(1, "Hello world"), (2, "Spark UDF example"), (3, "")]
df = spark.createDataFrame(data, ["id", "text"])

# 定义Python函数，将字符串分析为字典，包含单词数和字符数
def analyze_text(text):
    if text:
        return {
            "word_count": len(text.split(" ")),
            "char_count": len(text)
        }
    else:
        return {
            "word_count": 0,
            "char_count": 0
        }

# 使用udf()注册Python函数
analyze_text_udf = udf(analyze_text, MapType(StringType(), IntegerType()))

# 使用UDF处理DataFrame中的列
df_with_analysis = df.withColumn("text_analysis", analyze_text_udf(df["text"]))

# 展示结果
df_with_analysis.show(truncate=False)

输出：

+---+-----------------+-----------------------------------+
|id |text             |text_analysis                      |
+---+-----------------+-----------------------------------+
|1  |Hello world      |{char_count -> 11, word_count -> 2}|
|2  |Spark UDF example|{char_count -> 17, word_count -> 3}|
|3  |                 |{char_count -> 0, word_count -> 0} |
+---+-----------------+-----------------------------------+

代码详解：

  1. 定义和注册UDF：我们定义了一个函数analyze_text，用于返回一个字典，字典中包含word_count（单词数）和char_count（字符数）。使用pyspark.sql.functions.udf()注册这个函数，并将返回类型指定为MapType(StringType(), IntegerType())，表示返回的字典键是字符串类型，值是整数类型。
  2. 应用UDF：使用withColumn()方法将analyze_text_udf应用于text列，生成新的列text_analysis。

六、总结：返回值为数组和字典类型的UDF

  • 返回值为数组和字典的UDF可以帮助我们在Spark中实现更复杂的数据分析任务。
  • 当定义返回数组类型的UDF时，使用ArrayType()来定义返回值类型。
  • 当定义返回字典类型的UDF时，使用MapType(keyType, valueType)来定义键值对的类型。
  • 这些UDF适合在数据清洗、特征提取等复杂数据处理中使用，极大地扩展了Spark的能力。

在选择具体的UDF实现方式时，可以根据需求场景进行选择，例如：

  • 需要在SQL中直接使用自定义函数时，使用spark.udf.register()。
  • 需要链式处理并使用DataFrame API时，使用pyspark.sql.functions.udf()。

希望这些详细的例子能帮助你理解Spark中UDF的用法，并为你的项目选择合适的方法。

7 UDF 返回值为StructType

允许非常复杂的嵌套关系

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, ArrayType, MapType

# 初始化SparkSession
spark = SparkSession.builder.appName("Complex UDF Example").getOrCreate()

# 创建示例DataFrame
data = [(1, "Hello world"), (2, "Spark UDF example example"), (3, "")]
df = spark.createDataFrame(data, ["id", "text"])

# 定义Python函数，返回复杂的嵌套结构
def analyze_text_complex(text):
    if text:
        words = text.split(" ")
        word_count = len(words)
        char_count = len(text)
        word_frequency = {word: words.count(word) for word in set(words)}
        return {
            "word_count": word_count,
            "char_count": char_count,
            "words": words,
            "word_frequency": word_frequency
        }
    else:
        return {
            "word_count": 0,
            "char_count": 0,
            "words": [],
            "word_frequency": {}
        }

# 定义返回类型为StructType，包含嵌套的ArrayType和MapType
complex_schema = StructType([
    StructField("word_count", IntegerType(), True),
    StructField("char_count", IntegerType(), True),
    StructField("words", ArrayType(StringType()), True),
    StructField("word_frequency", MapType(StringType(), IntegerType()), True)
])

# 使用udf()注册Python函数
analyze_text_complex_udf = udf(analyze_text_complex, complex_schema)

# 使用UDF处理DataFrame中的列
df_with_analysis = df.withColumn("text_analysis", analyze_text_complex_udf(df["text"]))

# 展示结果
df_with_analysis.select('*').show(truncate=False)

输出：

+---+-------------------------+-----------------------------------------------------------------------------+
|id |text                     |text_analysis                                                                |
+---+-------------------------+-----------------------------------------------------------------------------+
|1  |Hello world              |{2, 11, [Hello, world], {Hello -> 1, world -> 1}}                            |
|2  |Spark UDF example example|{4, 25, [Spark, UDF, example, example], {Spark -> 1, UDF -> 1, example -> 2}}|
|3  |                         |{0, 0, [], {}}                                                               |
+---+-------------------------+-----------------------------------------------------------------------------+
上一篇 DataFrame通过JDBC读写Mysql
下一篇 Pyspark通过RDD+UDF实现UDAF