Unity Catalog 中的批量 Python 用户定义函数（UDF）

Unity Catalog 批量 Python UDF 已正式发布。 它们处理的是批量数据，而不是一行一行。

要求

在经典计算资源中，用于批处理的 Unity Catalog Python UDF 需要 Databricks Runtime 16.3 及更高版本。 它们也受支持于无服务器计算以及专业版和无服务器 SQL 仓库。

额外功能则有各自的计算和版本需求。 请参见 Python UDF 功能要求。

创建一个批处理的 Unity 目录 Python UDF

创建 Batch Unity 目录 Python UDF 类似于创建常规 Unity 目录 UDF，并添加以下内容：

  • PARAMETER STYLE PANDAS：这指定 UDF 使用 pandas 迭代器批量处理数据。
  • HANDLER 'handler_function'： 这指定了处理批次的处理程序函数。

以下示例在 Unity 目录中创建持久化的批处理 Python UDF。 用你的目录和模式替换 my_catalog 和 my_schema ：

%sql
CREATE OR REPLACE FUNCTION my_catalog.my_schema.calculate_bmi_pandas(weight_kg DOUBLE, height_m DOUBLE)
RETURNS DOUBLE
LANGUAGE PYTHON
DETERMINISTIC
PARAMETER STYLE PANDAS
HANDLER 'handler_function'
AS $$
import pandas as pd
from typing import Iterator, Tuple

def handler_function(batch_iter: Iterator[Tuple[pd.Series, pd.Series]]) -> Iterator[pd.Series]:
  for weight_series, height_series in batch_iter:
    yield weight_series / (height_series ** 2)
$$;

注册 UDF 后，可以使用 SQL 或 Python 调用它：

SELECT person_id, my_catalog.my_schema.calculate_bmi_pandas(weight_kg, height_m) AS bmi
FROM (
  SELECT 1 AS person_id, CAST(70.0 AS DOUBLE) AS weight_kg, CAST(1.75 AS DOUBLE) AS height_m UNION ALL
  SELECT 2 AS person_id, CAST(80.0 AS DOUBLE) AS weight_kg, CAST(1.80 AS DOUBLE) AS height_m
);

批量 UDF 处理器函数

Batch Unity Catalog Python UDF 需要处理批并生成结果的处理程序函数。 使用 HANDLER 子句创建 UDF 时，必须指定处理程序函数的名称。

处理程序函数执行以下任务：

  1. 接受一个迭代子参数，该参数对一个或多个 pandas.Series进行迭代。 每个 pandas.Series 参数都包含 UDF 的输入参数。
  2. 循环访问生成器并处理数据。
  3. 返回 生成器迭代器。

批量 Unity Catalog Python UDF 返回的行数必须与输入相同。 处理程序函数通过为每个批生成一个与输入序列长度相同的 pandas.Series 来确保这一点。

安装自定义依赖项

可以通过定义外部库的自定义依赖项，将 Batch Unity 目录 Python UDF 的功能扩展到 Databricks Runtime 环境之外。

请参阅 使用自定义依赖项扩展 UDF。

访问 Unity Catalog 机密

批处理 Unity Catalog Python UDF 可以访问在 SECRETS 子句中声明的机密。 你必须将 environment_version 显式地设置为 6 或更高。 专用访问模式计算不支持直接调用使用该条款的UDF。 有关列掩码例外情况，请参阅 在专用计算上将启用机密的 UDF 用于列掩码。

Batch UDF 可以接受单个或多个参数

单参数： 当处理函数使用单个输入参数时，它会接收针对每个批次的 pandas.Series 迭代器。

%sql
CREATE OR REPLACE TEMPORARY FUNCTION one_parameter_udf(value INT)
RETURNS STRING
LANGUAGE PYTHON
DETERMINISTIC
PARAMETER STYLE PANDAS
HANDLER 'handler_func'
AS $$
import pandas as pd
from typing import Iterator
def handler_func(batch_iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
  for value_batch in batch_iter:
    d = {"min": value_batch.min(), "max": value_batch.max()}
    yield pd.Series([str(d)] * len(value_batch))
$$;
SELECT one_parameter_udf(id), count(*) from range(0, 100000, 3, 8) GROUP BY ALL;

多个参数： 对于多个输入参数，处理程序函数接收一个迭代器，用于迭代pandas.Series中的多个元素。 序列中的值与输入参数的顺序相同。

%sql
CREATE OR REPLACE TEMPORARY FUNCTION two_parameter_udf(p1 INT, p2 INT)
RETURNS INT
LANGUAGE PYTHON
DETERMINISTIC
PARAMETER STYLE PANDAS
HANDLER 'handler_function'
AS $$
import pandas as pd
from typing import Iterator, Tuple

def handler_function(batch_iter: Iterator[Tuple[pd.Series, pd.Series]]) -> Iterator[pd.Series]:
  for p1, p2 in batch_iter: # same order as arguments above
    yield p1 + p2
$$;
SELECT two_parameter_udf(id , id + 1) from range(0, 100000, 3, 8);

通过区分贵重的操作来优化性能

可以通过将这些计算成本高昂的操作与处理函数分离来进行优化。 这可确保它们仅执行一次，而不是在每次对批处理数据进行迭代期间执行。

以下示例演示如何确保仅执行一次昂贵的计算：

%sql
CREATE OR REPLACE TEMPORARY FUNCTION expensive_computation_udf(value INT)
RETURNS INT
LANGUAGE PYTHON
DETERMINISTIC
PARAMETER STYLE PANDAS
HANDLER 'handler_func'
AS $$
def compute_value():
  # expensive computation...
  return 1

expensive_value = compute_value()
def handler_func(batch_iter):
  for batch in batch_iter:
    yield batch * expensive_value
$$;
SELECT expensive_computation_udf(id), count(*) from range(0, 100000, 3, 8) GROUP BY ALL

环境隔离

注释

共享隔离环境需要 Databricks Runtime 17.1 及更高版本。 在早期版本中，所有 Batch Unity 目录 Python UDF 都以严格的隔离模式运行。

默认情况下，具有相同所有者和会话的 Batch Unity 目录 Python UDF 可以共享隔离环境。 这可以通过减少需要启动的独立环境的数量来提高性能和减少内存使用率。

严格隔离

若要确保 UDF 始终在其自己的完全隔离环境中运行，请添加 STRICT ISOLATION 特征子句。

大多数 UDF 不需要严格的隔离。 标准数据处理 UDF 受益于默认共享隔离环境，并且运行速度更快，内存消耗量较低。

将 STRICT ISOLATION 特征子句添加到满足以下条件的 UDF：

  • 使用eval()exec()或类似函数以代码形式运行输入
  • 将文件写入本地文件系统
  • 修改全局变量或系统状态
  • 修改环境变量

以下示例演示一个 UDF，该 UDF 以代码的形式执行输入，并且需要严格隔离：

CREATE OR REPLACE TEMPORARY FUNCTION eval_string(input STRING)
RETURNS STRING
LANGUAGE PYTHON
PARAMETER STYLE PANDAS
HANDLER 'handler_func'
STRICT ISOLATION
AS $$
import pandas as pd
from typing import Iterator

def handler_func(batch_iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
  for code_series in batch_iter:
    def eval_func(code):
      try:
        return str(eval(code))
      except Exception as e:
        return f"Error: {e}"
    yield code_series.apply(eval_func)
$$;

Batch Unity Catalog Python UDF 中的服务凭据

批处理的Unity Catalog Python UDF可以使用Unity Catalog服务凭证访问外部云服务。 这对于将安全令牌化器等云函数集成到数据处理工作流中特别有用。

注释

UDF 专用的服务凭据 API：
在 UDF 中，使用 databricks.service_credentials.getServiceCredentialsProvider() 访问服务凭据。

这不同于 dbutils.credentials.getServiceCredentialsProvider() 笔记本中使用的函数，该函数在 UDF 执行上下文中不可用。

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

指定要在 UDF 定义中的 CREDENTIALS 子句中使用的服务凭据：

CREATE OR REPLACE TEMPORARY FUNCTION example_udf(data STRING)
RETURNS STRING
LANGUAGE PYTHON
PARAMETER STYLE PANDAS
HANDLER 'handler_function'
CREDENTIALS (
  `credential-name` DEFAULT,
  `complicated-credential-name` AS short_name,
  `simple-cred`,
  cred_no_quotes
)
AS $$
# Python code here
$$;

服务凭据权限

关于跨计算类型的创建和调用权限要求，请参见“在 Python UDF 中使用服务凭证”。

默认凭据和别名

可以在子句中包含 CREDENTIALS 多个凭据，但只能将一个凭据标记为 DEFAULT。 可以使用关键字对非默认凭据 AS 进行别名。 每个凭证必须具有唯一的别名。

已打补丁的云 SDK 会自动获取默认凭据。 默认凭据优先于计算的 Spark 配置中指定的任何默认值，并保留在 Unity 目录 UDF 定义中。

必须安装包 azure-identity 才能使用 DefaultAzureCredential 提供程序。 使用 ENVIRONMENT 子句安装外部库。 若要详细了解如何安装外部库，请参阅 使用自定义依赖项扩展 UDF。

服务凭据示例 - Azure Blob 存储

以下示例使用服务凭据从 Batch Unity 目录 Python UDF 访问 Azure Blob 存储：

%sql
CREATE OR REPLACE FUNCTION main.test.read_azure_blob(blob_name STRING) RETURNS STRING LANGUAGE PYTHON
PARAMETER STYLE PANDAS
HANDLER 'batchhandler'
CREDENTIALS (
  `batch-udf-service-creds-example-cred` DEFAULT
)
ENVIRONMENT (
  dependencies = '["azure-identity", "azure-storage-blob"]', environment_version = 'None'
)
AS $$
import pandas as pd
from azure.identity import DefaultAzureCredential
from azure.storage.blob import BlobServiceClient


def batchhandler(it):
  # DefaultAzureCredential automatically uses the DEFAULT service credential
  credential = DefaultAzureCredential()
  blob_service = BlobServiceClient(
    account_url="https://your-storage-account.blob.core.windows.net",
    credential=credential
  )
  container = blob_service.get_container_client("your-container")

  for blob_names in it:
    results = []
    for name in blob_names:
      blob_client = container.get_blob_client(name)
      try:
        content = blob_client.download_blob().readall().decode("utf-8")
        results.append(content)
      except Exception as e:
        results.append(f"Error: {e}")
    yield pd.Series(results)
$$;

注册 UDF 后调用它：

SELECT main.test.read_azure_blob(blob_name)
FROM VALUES
('config/settings.json'),
('data/input.txt')
AS t(blob_name)

获取任务执行上下文

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

如果您的函数产生一致的结果，请设置DETERMINISTIC

如果您的函数在相同输入下生成相同输出，请将 DETERMINISTIC 添加到函数定义中。 这允许查询优化来提高性能。

默认情况下，Batch Unity Catalog Python UDTF 被假定为非确定性的，除非已被显式声明。 非确定性函数的示例包括：生成随机值、访问当前时间或日期或进行外部 API 调用。

请参阅 CREATE FUNCTION （SQL、Python、Scala 和 Java）

局限性

  • Python 函数必须独立处理 NULL 值，并且所有类型映射都必须遵循 Azure Databricks SQL 语言映射。
  • Batch Unity 目录 Python UDF 在安全隔离的环境中运行，并且无权访问共享文件系统或内部服务。
  • 在一个阶段中，多个 UDF 调用被序列化，中间结果会被保存并可能溢出到磁盘。
  • 若要在无服务器笔记本或作业计算上进行 Batch Unity Catalog Python UDF 调用，必须配置 无服务器出口控制
  • 上次更新时间为