Python标量用户自定义函数(UDFs)

Python标量UDF允许你在Azure Databricks上的SQL查询中运行自定义Python逻辑。 本页展示了如何注册和调用这些信息,使用服务凭据和秘密,以及处理 Spark SQL 中子表达式评估顺序的注意事项。

要求

  • 在 Databricks Runtime 12.2 LTS 及更低版本中,使用标准访问模式的 Unity 目录计算不支持 Python UDF 和 Pandas UDF。

  • Databricks Runtime 13.3 LTS 及更高版本支持标量 Python UDF 和 Pandas UDF,适用于所有访问模式。

  • 在启用 Unity Catalog 的群集上,Python UDF 对 ARM 实例支持需要 Databricks Runtime 15.2 或更高版本。

在 Databricks Runtime 14.0 及更低版本中,使用标准访问模式的 Unity 目录群集不支持 Python UDF 和 Pandas UDF。 Databricks Runtime 14.1 及更高版本中的所有访问模式都支持标量 Python UDF 和 Pandas UDF。

在 Databricks Runtime 14.1 及更高版本中,可以使用 SQL 语法将标量 Python UDF 注册到 Unity 目录。 请参阅 Unity 目录中的 SQL 和Python用户定义的函数(UDF)。

将函数注册为 UDF

def squared(s):
  return s * s
spark.udf.register("squaredWithPython", squared)

可以选择设置 UDF 的返回类型。 默认返回类型为 StringType。

from pyspark.sql.types import LongType
def squared_typed(s):
  return s * s
spark.udf.register("squaredWithPython", squared_typed, LongType())

在 Spark SQL 中调用 UDF

spark.range(1, 20).createOrReplaceTempView("test")
%sql select id, squaredWithPython(id) as id_squared from test

将 UDF 与数据帧配合使用

from pyspark.sql.functions import udf
from pyspark.sql.types import LongType
squared_udf = udf(squared, LongType())
df = spark.table("test")
display(df.select("id", squared_udf("id").alias("id_squared")))

另外,还可以使用注释语法声明同一 UDF:

from pyspark.sql.functions import udf

@udf("long")
def squared_udf(s):
  return s * s
df = spark.table("test")
display(df.select("id", squared_udf("id").alias("id_squared")))

使用 UDF 的变体

变体的 PySpark 类型是 VariantType 且值为类型 VariantVal。 有关变体的信息,请参阅 查询变体数据。

from pyspark.sql.functions import col, lit, udf
from pyspark.sql.types import VariantType, VariantVal

# Return Variant
@udf(returnType = VariantType())
def toVariant(jsonString):
  return VariantVal.parseJson(jsonString)

spark.range(1).select(lit('{"a" : 1}').alias("json")).select(toVariant(col("json"))).display()
+---------------+
|toVariant(json)|
+---------------+
|        {"a":1}|
+---------------+
from pyspark.sql.functions import col, lit, udf
from pyspark.sql.types import StructField, StructType, VariantType, VariantVal

# Return Struct<Variant>
@udf(returnType = StructType([StructField("v", VariantType(), True)]))
def toStructVariant(jsonString):
  return {"v": VariantVal.parseJson(jsonString)}

spark.range(1).select(lit('{"a" : 1}').alias("json")).select(toStructVariant(col("json"))).display()
+---------------------+
|toStructVariant(json)|
+---------------------+
|        {"v":{"a":1}}|
+---------------------+
from pyspark.sql.functions import col, lit, udf
from pyspark.sql.types import ArrayType, VariantType, VariantVal

# Return Array<Variant>
@udf(returnType = ArrayType(VariantType()))
def toArrayVariant(jsonString):
  return [VariantVal.parseJson(jsonString)]

spark.range(1).select(lit('{"a" : 1}').alias("json")).select(toArrayVariant(col("json"))).display()
+--------------------+
|toArrayVariant(json)|
+--------------------+
|           [{"a":1}]|
+--------------------+
from pyspark.sql.functions import col, lit, udf
from pyspark.sql.types import MapType, StringType, VariantType, VariantVal

# Return Map<String, Variant>
@udf(returnType = MapType(StringType(), VariantType(), True))
def toMapVariant(jsonString):
  return {"v1": VariantVal.parseJson(jsonString), "v2": VariantVal.parseJson("[" + jsonString + "]")}

spark.range(1).select(lit('{"a" : 1}').alias("json")).select(toMapVariant(col("json"))).display()
+-----------------------------+
|           toMapVariant(json)|
+-----------------------------+
|{"v2":[{"a":1}],"v1":{"a":1}}|
+-----------------------------+

带有UDF的文件

Important

此功能在 Beta 版中。

文件的 PySpark 类型为 FileType。 在UDF中将其用作参数或返回类型,可以作为顶层类型或嵌套。 关于类型、嵌套规则和 FileRef API,请参见 文件类型。

要在 UDF 中读取文件内容,可以调用 file.as_local_file() 获取一个可打开的本地路径,或者调用 file.open() 以流的形式读取其字节数据。 关于Python、Scala和SQL中的示例,包括图像处理、文件类型检测和视频帧提取,请参见带UDF的进程文件。 关于类型参考,请参见 FILE 类型。

已在 Unity Catalog 中注册的 CREATE FUNCTION UDF 可以读取 FILE 的元数据,但不能读取其内容,也无法创建文件。 使用会话范围的 UDF 读取文件内容(open、from_bytes)或创建文件(as_local_file、from_local_file)。

在UDF主体中,每个 FILE 值是一个 FileRef 对象:

from pyspark.sql.functions import col, udf
from pyspark.sql.types import FileRef, StringType

@udf(returnType=StringType())
def file_content_type(file: FileRef) -> str:
  return file.content_type

df = spark.table("documents")
display(df.select(col("file").uri, file_content_type(col("file"))))

计算顺序和 NULL 检查

Spark SQL(包括 SQL、数据帧和数据集 API)不保证子表达式的计算顺序。 具体而言,运算符或函数的输入不一定按从左到右的顺序或任何其他固定顺序进行计算。 例如,逻辑 AND 和 OR 表达式没有从左到右的“短路”语义。

因此,依赖于布尔表达式计算的副作用或顺序以及 WHERE 和 HAVING 子句的顺序是危险的,因为在查询优化和规划过程中,这些表达式和子句可能重新排序。 具体而言,如果 UDF 依赖于 SQL 中的短路语义进行 NULL 检查,则不能保证在调用 UDF 之前执行 NULL 检查。 例如,

spark.udf.register("strlen", lambda s: len(s), "int")
spark.sql("select s from test1 where s is not null and strlen(s) > 1") # no guarantee

这个 WHERE 子句并不保证在筛选掉 NULL 后调用 strlen UDF。

若要执行正确的 NULL 检查,建议执行以下操作之一:

  • 使 UDF 自身能够识别 NULL,在 UDF 自身内部进行 NULL 检查
  • 使用 IF 或 CASE WHEN 表达式来执行 NULL 检查并在条件分支中调用 UDF
spark.udf.register("strlen_nullsafe", lambda s: len(s) if not s is None else -1, "int")
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

访问 Unity Catalog 机密

要从会话范围的 Python UDF 访问 Unity 目录秘密,请参见“在会话范围的 Python UDF 中使用秘密。 要从标量或批处理的 Unity Catalog Python UDF 访问已声明的密钥,请参见“在 Python UDF 中使用密钥”。

Python UDF中的服务凭证

会话级标量 Python UDF 和 Unity Catalog 标量 Python UDF 可以使用 Unity Catalog 服务凭证安全地访问外部云服务。 这可用于将基于云的令牌化、加密或机密管理等作直接集成到数据转换中。

需求因UDF类型和计算量而异。 参见“在 Python UDF 中使用服务凭证”。

若要创建服务凭据,请参阅 “创建服务凭据”。

在会话级标量 Python UDF 中使用服务凭据

若要访问服务凭据,请使用 databricks.service_credentials.getServiceCredentialsProvider() UDF 逻辑中的实用工具通过相应的凭据初始化云 SDK。 所有代码都必须封装在 UDF 正文中。

@udf
def use_service_credential():
    from azure.mgmt.web import WebSiteManagementClient

    # Assuming there is a service credential named 'testcred' set up in Unity Catalog
    web_client = WebSiteManagementClient(subscription_id, credential = getServiceCredentialsProvider('testcred'))
    # Use web_client to perform operations

服务凭据权限

会话范围的UDF使用调用者的权限。 请参阅 Python UDF 中使用服务凭证获取所需权限。

会话范围UDF的计算级默认凭证

在标量 Python UDF 中使用时,Databricks 会自动使用计算环境变量中的默认服务凭证。 此行为允许在不显式管理 UDF 代码中的凭据别名的情况下安全地引用外部服务。 请参阅 为计算资源指定默认服务凭据

默认凭据支持仅适用于标准和专用访问模式群集。 它不存在于 SQL 仓库中。

必须安装包 azure-identity 才能使用 DefaultAzureCredential 提供程序。 若要安装包,请参阅 笔记本级别的 Python 库 或 计算级别的库。

@udf
def use_service_credential():
    from azure.identity import DefaultAzureCredential
    from azure.mgmt.web import WebSiteManagementClient

    # DefaultAzureCredential is automatically using the default service credential for the compute
    web_client_default = WebSiteManagementClient(DefaultAzureCredential(), subscription_id)

    # Use web_client to perform operations

在标量Unity Catalog Python UDF中使用服务凭据

在 UDF 定义的 CREDENTIALS 子句中指定服务凭证。 你可以把一个凭据标记为 DEFAULT ,这样经过补丁的云SDK会自动使用它。 在经典计算中,此功能需要Databricks Runtime 18.1及以上版本。 在无服务器计算以及专业版和无服务器 SQL 仓库中,显式地将 UDF 的 environment_version 设置为 6 或更高。 关于完整的计算、网络和权限要求,请参见“在 Python UDF 中使用服务凭证”。

获取任务执行上下文

使用 TaskContext PySpark API 获取上下文信息,例如用户的标识、群集标记、spark 作业 ID 等。 请参阅 UDF 中的“获取任务上下文”。

限制

以下限制适用于 PySpark UDF:

  • 文件访问限制: 在 Databricks Runtime 14.2 及更低版本上,共享群集上的 PySpark UDF 无法访问 Git 文件夹、工作区文件或 Unity 目录卷。

  • 广播变量: 标准访问模式群集和无服务器计算上的 PySpark UDF 不支持广播变量。

  • 无服务器内存限制:无服务器计算上的 PySpark UDF 的内存限制为每个 PySpark UDF 限制 1GB。 超出此限制会导致类型 UDF_PYSPARK_USER_CODE_ERROR.MEMORY_LIMIT_SERVERLESS 错误。
  • 标准访问模式上的内存限制:标准访问模式上的 PySpark UDF 基于所选实例类型的可用内存限制。 超出可用内存会导致 UDF_PYSPARK_USER_CODE_ERROR.MEMORY_LIMIT 类型错误。
  • 无服务器 SQL 仓库中的网络访问:默认情况下,无服务器 SQL 仓库中的 Python UDF 无法发起出站网络请求,并且尝试进行网络调用的查询会无限期挂起。 要启用出站网络访问,请在工作区的预览页中启用为无服务器 SQL 仓库中的隔离工作负载启用网络公共预览功能。 否则,对需要网络访问的 UDF 使用无服务器计算或经典计算。