Class UDF

Class UDF注册

Class UDF注册传入的是Python类，对于__init__、__call__、 __del__ 3个实例方法有特殊的含义认定，详情请参见下表；其他的Python类和实例方法不做限制，用户可以任意添加。

注意事项：__init__的生命周期是每个batch调用一次，即Aura内核每处理一个batch，都会重新初始化一次变量。

示例

本示例定义一个带初始化参数的Class UDF，在__init__中接收姓名参数，在__call__中结合城市名生成欢迎语，将其注册到数据库后通过map方法对city_list表的每行数据生成问候消息。

1. 定义Class UDF类Hello，在__init__中接收firstname和lastname参数，在__call__中结合城市名生成欢迎语。
2. 连接Aura服务并设置UDF工作空间。
3. 创建city_list表（包含city列），并插入5条城市数据。
4. 删除已存在的同名UDF（支持重复执行）。
5. 注册Class UDF到目标数据库。
6. 加载city_list数据集，通过map方法调用UDF对city列生成问候语，传入firstname和lastname参数。
7. 选择city和new_column列并执行查询打印结果。
8. 关闭Aura连接。

import os
import ibis
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="postgres",                       # 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）
)
# 定义 Class UDF：带初始化参数的问候生成器
class Hello:
    def __init__(self, firstname: str, lastname: str):
        self.firstname = firstname
        self.lastname = lastname
    def __call__(self, location: str) -> str:
        return f"Hello {self.firstname} {self.lastname}! Welcome to {location}"

try:
    # 创建测试表并写入数据
    table_name = "city_list"
    table_schema = ibis.schema({"city": "string"})
    con.sql(f'DROP TABLE IF EXISTS "{target_database}"."{table_name}";')
    con.create_table(
        table_name,
        schema=table_schema,
        database=target_database,
        table_format="orc",
        if_not_exists=True,
    )
    con.sql(f'''
    INSERT INTO "{target_database}"."{table_name}" (city) VALUES
        ('Beijing'),
        ('Shanghai'),
        ('New York'),
        ('London'),
        ('Tokyo')
    ''')
    # 删除已存在的同名UDF（支持重复执行）
    con.delete_function("hello", database=target_database, if_it_exists=True)
    # 注册 Class UDF
    hello_udf = con.create_scalar_function(
        Hello,
        name="hello",
        database=target_database,
    )
    # 使用 UDF 处理 city_list 表数据
    ds = con.load_dataset("city_list", database=target_database)
    ds = ds.map(
        fn=hello_udf,
        on=[ds.city],
        as_col="new_column",
        firstname="Ethan",
        lastname="Carter",
    )
    ds = ds.select_columns(ds.city, ds.new_column)
    print(ds.execute())
finally:
    con.close()