使用Scalar UDF处理数据

场景描述

在销售数据分析场景中，需要根据商品原价、折扣率和税率计算含税最终价格。该计算逻辑涉及多个字段的组合运算，使用内置SQL函数表达不够直观，且难以复用。通过将计算逻辑封装为Scalar UDF并注册到数据库，可以在DataFrame中像调用内置函数一样使用，代码更清晰，逻辑更易维护。

约束与限制

• 仅支持Python语言编写的用户自定义函数。
• Python环境依赖管理由用户自行负责，需确保各依赖库版本兼容。
• 运行时环境要求Python版本为3.11。
• Scalar UDF的输入参数和返回值必须显式指定数据类型。
• 注册UDF前建议先调用delete_function清理同名函数，避免重复执行时报错。

前提条件

已创建表并写入数据。具体操作，请参见创建表并写入数据。

示例

本示例定义一个计算含税最终价格的Scalar UDF，将其注册到数据库后，在Aura DataFrame中通过map方法对sales_line_items表的每行数据执行行级计算。

1. 定义Scalar UDF函数calc_final_price，根据单价、折扣率和税率计算含税最终价格。
2. 连接AuraJobV2端点并设置UDF工作空间。
3. 删除已存在的同名UDF（支持重复执行）。
4. 注册Scalar UDF到目标数据库。
5. 加载sales_line_items数据集。
6. 通过map方法调用UDF对每行数据计算final_price列。
7. 选择相关列并执行查询打印结果。
8. 关闭AuraJobV2端点连接。

import os
from aura_frame.multimodal import ai_lake

def calc_final_price(price: float, discount: float, tax: float) -> float:
    return price * (1 - discount) * (1 + tax)

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（支持重复执行）
    con.delete_function("calc_final_price", database=target_database, if_it_exists=True)
    # 注册Scalar UDF
    udf = con.create_scalar_function(
        calc_final_price,
        database=target_database,
    )
    # 在DataFrame中使用UDF
    ds = con.load_dataset("sales_line_items", database=target_database)
    ds = ds.map(fn=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)
finally:
    con.close()

相关文档

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