直接使用已注册的Scalar UDF

场景描述

在团队协作开发中，UDF的注册和使用往往是分离的：一部分人负责编写和注册UDF，另一部分人负责在数据处理流程中调用已注册的UDF。本示例展示如何查看数据库中已注册的UDF列表，获取指定UDF算子，并在DataFrame表达式中直接使用。

约束与限制

• 使用前需确保目标UDF已注册到指定数据库中，函数名在数据库中存储为全小写。
• Scalar UDF和Class UDF均通过map方法调用，on参数传入__call__的输入列。
• Class UDF的__init__参数通过map方法的关键字参数传入（如model_file、eos）。

前提条件

• 已创建表并写入数据。具体操作，请参见创建表并写入数据。
• 运行时环境要求Python版本为3.11。

示例

本示例展示如何查看数据库中已注册的UDF列表，获取已注册的Scalar UDF和Class UDF算子，并在Aura DataFrame中直接调用，无需重复注册。

1. 连接Aura服务并设置UDF工作空间。
2. 查看目标数据库中已注册的UDF名称列表。
3. 若calc_final_price存在，获取该Scalar UDF算子并加载sales_line_items数据集，通过map方法计算final_price列。
4. 若spmanager存在，获取该Class UDF算子并加载text_samples数据集，通过map方法对content列执行分词（传入model_file和eos参数）。
5. 关闭Aura连接。

import os
from aura_frame.multimodal import ai_lake

target_database = "test"  #在LakeFormation中创建数据库，例如：test。具体操作，请参见创建LakeFormation实例并规划元数据
#连接AuraJobV2端点，具体参数说明，请参见建立与AuraJobV2端点的连接。
con = ai_lake.connect(
    aura_endpoint=os.getenv("aura_endpoint"),         # AuraJobV2端点所在的终端节点（Endpoint）
    aura_endpoint_id=os.getenv("aura_endpoint_id"),    # AuraJobV2端点ID
    aura_endpoint_name=os.getenv("aura_endpoint_name"),# AuraJobV2端点名称
    aura_workspace_id=os.getenv("aura_workspace_id"),  # 工作空间ID
    lf_catalog_name=os.getenv("lf_catalog_name"),      # LakeFormation Catalog名称
    access_key=os.getenv("access_key"),                # 访问密钥ID（AK）
    secret_key=os.getenv("secret_key"),                # 访问密钥（SK）
    default_database=target_database,                       # LakeFormation数据库名称
    use_single_cn_mode=True,   # 是否开启单CN模式
    logging_level=logging.INFO # 设置日志级别
)
#配置UDF部署工作区，具体参数说明，请参见配置UDF部署工作区。
con.set_function_staging_workspace(
    obs_directory_base=os.getenv("obs_directory_base"),  # 并行文件系统中UDF的存储路径
    obs_bucket_name=os.getenv("obs_bucket_name"),       # 并行文件系统名称
    obs_server=os.getenv("obs_server"),                  # OBS终端节点（Endpoint）
    access_key=os.getenv("access_key"),                  # 访问密钥ID（AK）
    secret_key=os.getenv("secret_key"),                  # 访问密钥（SK）
)
try:
    # 查看已注册的所有UDF名称
    udf_names = con._conn.udf.names(database=target_database)
    print(f"已注册的UDF: {udf_names}")
    if "calc_final_price" in udf_names:
        # 获取已注册的Scalar UDF
        calc_final_price_udf = con.get_function("calc_final_price", database=target_database)
        ds = con.load_dataset("sales_line_items", database=target_database)
        # 在DataFrame中直接使用已注册的UDF
        ds = ds.map(
            fn=calc_final_price_udf,
            on=[ds.unit_price, ds.discount_rate, ds.tax_rate],
            as_col="final_price",
        )
        ds = ds.select_columns(ds.unit_price, ds.discount_rate, ds.tax_rate, ds.final_price)
        df = ds.execute()
        print(df)
    if "spmanager" in udf_names:
        # 获取已注册的Class UDF
        sp_udf = con.get_function("spmanager", database=target_database)
        ds = con.load_dataset("text_samples", database=target_database)
        # 通过map方法对content列执行分词
        ds = ds.map(
            fn=sp_udf, on=[ds.content], as_col="pieces_column",
            model_file="spiece.model", eos=True,
        )
        ds = ds.select_columns(ds.content, ds.pieces_column)
        df = ds.execute()
        print(df)
finally:
    con.close()

相关文档

• Session APIs
• Dataset APIs
• Table APIs
• User-Defined Function APIs