会话范围的 Scala 和 Java UDF

本页介绍如何在Azure Databricks中创建会话范围的 Scala 和 Java UDF。 会话范围的 UDF 在笔记本或作业中定义，仅适用于当前的 SparkSession。 有关 SQL 语言参考，请参阅外部用户定义的标量函数（UDF）。

选择你的方法

可以通过以下方式定义 Scala 或 Java UDF。 若要比较不同语言、治理和计算环境下的所有 UDF 类型，请参阅 受 Unity Catalog 治理的 UDF 与会话范围 UDF。

要求

• 在启用了 Unity Catalog 且使用标准访问模式的计算资源上，Scala UDF 需要 Databricks Runtime 14.2 或更高版本。
• 在启用 Unity Catalog 的群集上使用 Scala UDF 所需的 ARM 实例支持需要 Databricks Runtime 15.2 或更高版本。
• 通过 spark.udf.registerJavaFunction 从 JAR 注册 Java UDF 要求使用 Databricks Runtime 18 LTS 或更高版本。 请参阅从 JAR 注册Java UDF。

将函数注册为 UDF

使用 spark.udf.register 将 Scala 函数注册为 UDF：

val squared = (s: Long) => {
  s * s
}
spark.udf.register("square", squared)

在 Spark SQL 中调用 UDF

创建临时视图，然后在 SQL 查询中调用 UDF：

spark.range(1, 20).createOrReplaceTempView("test")

%sql select id, square(id) as id_squared from test

将 UDF 与数据帧配合使用

还可以使用数据帧 API 调用 UDF：

import org.apache.spark.sql.functions.{col, udf}
val squared = udf((s: Long) => s * s)
display(spark.range(1, 20).select(squared(col("id")) as "id_squared"))

通过 JAR 包注册 Java UDF

将 UDF 打包为 JAR，使用 spark.addArtifact 将其添加到当前会话中，并使用 spark.udf.registerJavaFunction 注册 UDF 类。

以下步骤介绍了如何创建项目、编写 UDF 类、构建胖 JAR 并将其注册。

步骤 1：创建项目

在 Scala 或 Java 中设置项目。

Scala

使用 sbt 创建一个新的 Scala 项目：

sbt new scala/scala-seed.g8

将你的 build.sbt 文件的内容替换为以下内容。 将 scalaVersion 和 spark-sql 版本设置为与您的计算环境相匹配：

scalaVersion := "2.13.16"

ThisBuild / organization := "com.example"

lazy val myUDF = (project in file("."))
  .settings(
    name := "my-udf",
    libraryDependencies += "org.apache.spark" %% "spark-sql" % "4.0.0" % "provided"
  )

启用 sbt-assembly 插件以构建 fat JAR 包。 创建或编辑 project/assembly.sbt 并添加：

addSbtPlugin("com.eed3si9n" % "sbt-assembly" % "2.0.0")

Java

使用快速入门原型创建新的 Maven 项目：

mvn archetype:generate \
  -DgroupId=com.example \
  -DartifactId=my-udf \
  -DarchetypeArtifactId=maven-archetype-quickstart \
  -DinteractiveMode=false

此命令会创建包含 src/main/java 和 src/test/java 目录的标准 Maven 项目结构。

在生成的 pom.xml 中，在 <project></project> 标签内添加一个 <properties> 块，并将 maven-shade-plugin 配置为构建胖 JAR：

<properties>
  <maven.compiler.source>17</maven.compiler.source>
  <maven.compiler.target>17</maven.compiler.target>
  <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
</properties>

<build>
  <plugins>
    <plugin>
      <groupId>org.apache.maven.plugins</groupId>
      <artifactId>maven-shade-plugin</artifactId>
      <version>3.5.0</version>
      <executions>
        <execution>
          <phase>package</phase>
          <goals>
            <goal>shade</goal>
          </goals>
        </execution>
      </executions>
    </plugin>
  </plugins>
</build>

步骤 2：编写 UDF 类

UDF 类必须实现其中 org.apache.spark.sql.api.java.UDF 一个接口（UDF1 通过 UDF22），其中数字指示 UDF 采用的输入参数数。 使用您的逻辑实现 call() 方法。

处理程序必须是Java类。 spark.udf.registerJavaFunction 通过反射加载类，因此它必须是具有公共 no-arg 构造函数的顶级（或 static 嵌套）公共类。 Scala class 或 object 不符合此要求，并在调用时失败。 可以使用 sbt 生成 JAR，但必须使用 Java 编写 UDF 类本身。

创建 src/main/java/com/example/MyIntegerUDF.java：

package com.example;

import org.apache.spark.sql.api.java.UDF1;

public class MyIntegerUDF implements UDF1<Integer, Integer> {
  @Override
  public Integer call(Integer x) {
    return x + 1;
  }
}

步骤 3：构建 fat JAR

将编译后的 UDF 打包成胖 JAR 包。

Scala

从项目根目录运行：

sbt clean assembly

fat JAR 会在 target/scala-2.13/ 中创建，名称类似于 my-udf-assembly-0.1.0-SNAPSHOT.jar。

Java

从项目根目录运行：

mvn clean package

fat JAR 会在 target/ 中创建，名称类似于 my-udf-1.0-SNAPSHOT.jar。

步骤 4：将 JAR 上传到 Unity 目录卷

将 JAR 上传到 Unity 目录卷 ，以便计算可以访问它。 如果还没有卷，请创建一个：

CREATE VOLUME IF NOT EXISTS my_catalog.my_schema.udf_jars
COMMENT 'Storage for UDF JAR files';

使用 目录资源管理器将 JAR 文件上传到卷：

1. 在 Azure Databricks 工作区中，单击 Catalog打开目录资源管理器。
2. 选择目录，然后选择包含您的卷的架构。
3. 单击卷名称。
4. 单击“ 上传到此卷 ”并选择 JAR 文件。
5. 单击“上载” 。
6. 上传完成后，单击 JAR 文件的名称，然后单击 “复制路径 ”以复制卷路径。 例如，/Volumes/my_catalog/my_schema/udf_jars/my-udf-assembly-0.1.0-SNAPSHOT.jar。 在下一步中需要此路径。

步骤 5：注册并调用 UDF

使用卷路径将 JAR 添加到会话，注册 UDF 类，并从 Spark SQL 调用它：

# Add the JAR containing your UDF class to the session
spark.addArtifact("/Volumes/my_catalog/my_schema/udf_jars/my-udf-assembly-0.1.0-SNAPSHOT.jar")

# Register the UDF class, providing the SQL function name,
# the fully qualified class name, and the return type
from pyspark.sql.types import IntegerType

spark.udf.registerJavaFunction(
    "my_udf",
    "com.example.MyIntegerUDF",
    IntegerType(),
)

# Call the UDF from Spark SQL
spark.sql("SELECT my_udf(21)").show()

在无服务器和标准访问模式计算中，必须传递显式返回类型。 省略返回类型会因 UC_COMMAND_NOT_SUPPORTED_IN_SHARED_ACCESS_MODE 而失败。 不支持将用户定义的聚合函数（UDAF）与 registerJavaFunction 一起使用。

该查询返回 UDF 输出，确认函数已注册且可调用：

+----------+
| my_udf(21)|
+----------+
|        22|
+----------+

评估顺序和空值检查

Spark SQL（包括 SQL 和数据帧和数据集 API）不能保证子表达式计算的顺序。 Spark 不计算运算符或函数从左到右的输入。 逻辑 AND 和 OR 表达式没有从左到右的短路语义。

不要依赖布尔表达式的副作用或求值顺序，也不要依赖 WHERE 和 HAVING 子句的顺序。 查询优化器可以重新排序这些表达式和子句。 如果 UDF 依赖短路求值语义来进行 null 检查，Spark 不保证该 null 检查会在 UDF 之前执行。 例如：

spark.udf.register("strlen", (s: String) => s.length)
spark.sql("select s from test1 where s is not null and strlen(s) > 1") // no guarantee

此 WHERE 子句不保证 Spark 在筛选出 null 后调用 strlen UDF。

为了处理 null 检查，Databricks 建议以下任一操作：

• 使 UDF 本身支持处理 null 值，并在 UDF 内部进行 null 检查
• 使用 IF 或 CASE WHEN 表达式来执行 NULL 检查并在条件分支中调用 UDF

spark.udf.register("strlen_nullsafe", (s: String) => if (s != null) s.length else -1)
spark.sql("select s from test1 where s is not null and strlen_nullsafe(s) > 1") // ok
spark.sql("select s from test1 where if(s is not null, strlen(s), null) > 1")   // ok

类型化数据集 API

使用类型化数据集 API 通过用户定义的函数在数据集上运行映射、筛选器和聚合等转换。

以下示例使用 map() API 将结果列中的数字修改为前缀字符串：

spark.range(3).map(f => s"row-$f").show()

此示例使用map()，但相同的模式适用于其他类型化数据集 API，例如filter()，mapPartitions()、、foreach()foreachPartition()和reduce()flatMap()。

Scala UDF 功能与 Databricks 运行时的兼容性

以下功能要求在启用了 Unity Catalog 且采用标准（共享）访问模式的集群上，Databricks Runtime 至少达到相应的最低版本。