UDTF

以下示例展示UDTF的注册和使用的流程。示例使用sequence_config表（包含start_val、end_val、step_val三个int32列），通过UDTF函数py_generate_series为每一行生成从start_val到end_val、步长为step_val的整数序列。

import os
import ibis
import aura_frame as aura
from aura_frame.multimodal import ai_lake
import ibis.expr.datatypes as dt

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）
)

# 定义 UDTF：输入 start、endnum、step，生成整数序列
def py_generate_series(start: int, endnum: int, step: int):
    current = start
    while current <= endnum:
        yield {"col1": current}
        current += step

try:
    # 创建测试表并写入数据
    table_name = "sequence_config"
    table_schema = ibis.schema({"start_val": "int32", "end_val": "int32", "step_val": "int32"})
    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}" (start_val, end_val, step_val) VALUES
       (1, 5, 1),
       (10, 20, 5),
       (0, 10, 2)
    ''')
    # 删除已存在的同名UDTF（支持重复执行）
    con.delete_function("py_generate_series", database=target_database, if_it_exists=True)
    func_signature = aura.Signature(
        parameters=[
            aura.Parameter(name="start", annotation=int),
            aura.Parameter(name="endnum", annotation=int),
            aura.Parameter(name="step", annotation=int),
        ],
        return_annotation=dt.Struct({"col1": int}),
    )
    con.create_table_function(py_generate_series, database=target_database, signature=func_signature)
    ds = con.load_dataset("sequence_config", database=target_database)
    py_generate_series_handler = con.get_function("py_generate_series", database=target_database)
    ds = ds.flat_map(
        fn=py_generate_series_handler,
        on=[ds.start_val, ds.end_val, ds.step_val],
        as_col="udtf_col",
    )
    res = ds.execute()
    print(res)
finally:
    con.close()