使用聚合统计的UDAF

场景描述

在销售数据分析中，需要按商品类别（category）统计每组中金额（amount）最高的前K个值。内置的MAX函数只能返回最大值，无法返回Top-K列表。通过自定义UDAF PythonTopK，在 accumulate 中收集每行的金额值并保持降序排序，在 merge 中合并各分区的中间状态，在 finish 中输出JSON格式的Top-K列表，实现灵活的自定义聚合统计。

约束与限制

• UDAF必须实现aggregate_state、accumulate、merge、finish 4个方法。
• UDAF要求用户实现的accumulate、aggregate_state、merge方法是无副作用的（side-effect free），即这些方法的返回结果只取决于输入参数和聚合状态本身，而不应依赖于方法被调用的次数、顺序或外部环境。
• UDAF不支持作为窗口函数使用，即不支持在OVER从句中使用。
• UDAF不支持开启并发和向量化计算。
• 注册UDAF前建议先调用delete_function清理同名函数，避免重复执行时报错。

前提条件

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

示例

本示例定义一个统计每组中金额最高的前K个值的UDAF，将其注册到数据库后，在category_sales表上按商品类别分组聚合，取每组金额前3的记录。

1. 定义UDAF类PythonTopK，实现aggregate_state、accumulate、merge、finish四个方法，用于统计每组中amount值最大的前K个记录。
2. 连接AuraJobV2端点并设置UDF工作空间。
3. 删除已存在的同名UDF（支持重复执行）。
4. 注册UDAF到目标数据库。
5. 加载category_sales数据集。
6. 通过AggregateFnBuilder构建聚合表达式，按category分组取amount前3。
7. 执行聚合查询并打印结果。
8. 关闭AuraJobV2端点连接。

import os
import json
from aura_frame.multimodal import ai_lake
from aura_frame.multimodal.function import AggregateFnBuilder

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

# 定义UDAF：统计每组金额最高的前K个值
class PythonTopK:
    def __init__(self, k: int):
        self.k = k
        self._agg_state = []
    @property
    def aggregate_state(self):
        return self._agg_state
    def accumulate(self, value: float) -> None:
        if value is not None:
            self._agg_state.append(value)
            self._agg_state.sort(reverse=True)
            self._agg_state = self._agg_state[: self.k]
    def merge(self, other_state) -> None:
        merged = self._agg_state + other_state
        merged.sort(reverse=True)
        self._agg_state = merged[: self.k]
    def finish(self) -> str:
        return json.dumps(self._agg_state)
try:
    # 删除已存在的同名UDF（支持重复执行）
    con.delete_function("PythonTopK", database=target_database, if_it_exists=True)
    # 注册UDAF
    udaf = con.create_agg_function(PythonTopK, database=target_database)
    # 在DataFrame中使用UDAF：按category分组，取每组amount最高的前3个
    ds = con.load_dataset("category_sales", database=target_database)
    udaf_builder = AggregateFnBuilder(fn=udaf, on=[ds.amount], as_col="top3", constructor_kwargs={"k": 3})
    summary = ds.aggregate([udaf_builder], by=[ds.category])
    df = summary.execute()
    print(df)
finally:
    con.close()

相关文档

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