Scalar UDF

开发流程

开发Python Scalar UDF，总体流程如下：

1. 根据用户的业务编写和调试Python代码。
   代码示例如下：

   import numpy as np
   
   def pysum(arg1: int, arg2: int) -> int:
       printLog(arg1, arg2)
       return np.sum(np.array([arg1, arg2]))
   

2. 注册Python UDF。
   使用ibis-aura创建会话。

   import aura_frame as aura
   from aura_frame.multimodal import ai_lake
   
   target_database = "test"  #在LakeFormation中创建数据库，例如：test。具体操作，请参见创建LakeFormation实例并规划元数据
   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 # 设置日志级别
   )

   注册Python UDF。

   try:
       udf = con.create_scalar_function(pysum, packages=["numpy"])
   finally:
       con.close()

3. 查看Python UDF。

   Aura提供show functions查看特定schema下的所有函数，且提供了describe function function_name方法查看特定函数的详细信息，具体的使用方式请参考SHOW和DESCRIBE。

4. 调用Python UDF。

   使用Aura任意客户端，下发查询即可，例如：

   select "test".pysum(3,4);

5. 查看Python UDF性能指标及调优。

   Aura支持了对Python UDF运行时的性能监控，以及提供了接口，可配置UDF运行时的资源规格和并发度，具体使用方式请参考Python UDF性能调优。

示例

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=target_database,                  # 默认数据库
    use_single_cn_mode=True,                           # 是否开启单CN模式
)
#配置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）
)
# 定义 Scalar UDF：两数求和
def add_udf(lhs: float, rhs: int) -> float:
    return lhs + rhs
try:
    # 创建测试表并写入数据
    table_name = "product_orders"
    table_schema = ibis.schema({"price": "float64", "quantity": "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}" (price, quantity) VALUES
        (100.0, 10),
        (200.0, 20),
        (150.0, 15)
    ''')
    # 删除已存在的同名UDF（支持重复执行）
    con.delete_function("add_udf", database=target_database, if_it_exists=True)
    # 注册 Scalar UDF
    udf = con.create_scalar_function(
        add_udf,
        database=target_database,
        comment="To test create_scalar_function",
    )
    # 使用 UDF 处理 product_orders 表数据
    ds = con.load_dataset("product_orders", database=target_database)
    ds = ds.map(fn=udf, on=[ds.price, ds.quantity], as_col="sum_column")
    ds = ds.select_columns(ds.price, ds.quantity, ds.sum_column)
    print(ds.execute())
finally:
    con.close()