更新时间:2026-10-16 GMT+08:00
分享

UDAF

UDAF是输入多行返回一组聚合值的函数,适用于典型的多模态数据处理场景,包括视频/音频/图像的统计特征聚合,多模态标签总结等。

UDAF整体开发流程可以参考Scalar UDF,需要注意以下几点:

  • Handler指定的必须是Python Class,且Python Class中必须包括accumulate、aggregate_state、merge、finish 4个方法。UDAF要求用户实现的accumulate、aggregate_state、merge方法是无副作用的(side-effect free),即这些方法的返回结果只取决于输入参数和聚合状态本身,而不应依赖于方法被调用的次数、顺序或外部环境。
  • UDAF的finish方法不论是单CN模式,还是多DN模式执行,对于每个分组(由GROUP BY从句指定)都只会在全局调用1次。
  • UDAF要求用户在定义非STRICT时,能够在accumulate方法中正确处理NULL值输入,即Python中的None值。
  • UDAF不支持作为窗口函数(window function)使用,即不支持在OVER从句中使用UDAF。
  • UDAF不支持开启并发和向量化计算。

UDAF注册

UDAF注册传入的必须是Python类,对于__init__、aggregate_state、accumulate、merge、finish、 __del__ 6个实例方法有特殊的含义认定,详情请参见下表;其他的Python类和实例方法不做限制,用户可以任意添加。

表1 注册UDAF时的特殊实例方法

Python Class方法

是否必须

参数

含义

适用场景

__init__(self, *args, **kwargs)

否

通过with_arguments方法传入参数,只允许传入标量Scalar。

UDAF的构造方法。

初始化UDAF属性(如保存参数、打开文件、建立连接等)。

aggregate_state(self)

是

无。

可序列化的聚合中间态,用于跨分区传递。

accumulate与merge间共享状态。

accumulate(self, *args, **kwargs)

是

通过UDAF算子传入参数,可以传入标量Scalar和列名Column。

增量更新聚合中间态。

正常数据扫过阶段。

merge(self, other_state)

是

来自其它分区的中间态。

合并其它分区状态到当前中间态。

分布式reduce阶段,多分区聚合结果归并。

finish(self)

是

无。

输出最终聚合结果。

聚合收尾。

__del__(self)

否

不支持传入参数。

UDAF的析构方法。

UDAF资源清理(如关闭文件、断开网络连接等)。

示例

import os
import json
import datetime as pydt
import ibis
import ibis.expr.datatypes as dt
import pandas as pd
from typing import Any, Dict
import aura_frame as aura
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="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"),
    obs_bucket_name=os.getenv("obs_bucket_name"),
    obs_server=os.getenv("obs_server"),
    access_key=os.getenv("access_key"),
    secret_key=os.getenv("secret_key"),
)
class _AggState:
    def __init__(self):
        self.sum_px_qty = 0.0
        self.sum_qty = 0.0
        self.rows_seen = 0
        self.rows_used = 0
class PythonVWAP:
    def __init__(self) -> None:
        self._state = _AggState()
    def accumulate(self, quantity: int, price: float, date: pydt.date, symbol: str) -> None:
        self._state.rows_seen += 1
        if quantity is None or price is None or date is None or symbol is None:
            return
        self._state.sum_px_qty += price * quantity
        self._state.sum_qty += quantity
        self._state.rows_used += 1
    @property
    def aggregate_state(self) -> Dict[str, Any]:
        return {
            "sum_px_qty": self._state.sum_px_qty,
            "sum_qty":    self._state.sum_qty,
            "rows_seen":  self._state.rows_seen,
            "rows_used":  self._state.rows_used,
        }
    def merge(self, other_state: Dict[str, Any]) -> None:
        self._state.sum_px_qty += other_state.get("sum_px_qty", 0.0)
        self._state.sum_qty    += other_state.get("sum_qty",    0.0)
        self._state.rows_seen   += other_state.get("rows_seen",   0)
        self._state.rows_used   += other_state.get("rows_used",   0)
    def finish(self) -> str:
        vwap = (self._state.sum_px_qty / self._state.sum_qty
                if self._state.sum_qty != 0 else None)
        result = {
            "vwap":           round(vwap, 4) if vwap is not None else None,
            "total_quantity": self._state.sum_qty,
            "total_value":    round(self._state.sum_px_qty, 2),
            "rows_seen":      self._state.rows_seen,
            "rows_used":      self._state.rows_used,
        }
        return json.dumps(result)

try:
    # 创建测试表并写入数据
    table_name = "stock_trades"
    table_schema = ibis.schema({
        "quantity": "int32",
        "price": "float64",
        "date": "date",
        "symbol": "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}" (quantity, price, date, symbol) VALUES
        (100, 150.25, '2024-01-15', 'AAPL'),
        (200, 152.30, '2024-01-16', 'AAPL'),
        (150, 148.50, '2024-01-17', 'AAPL'),
        (300, 145.75, '2024-01-15', 'GOOG'),
        (250, 147.80, '2024-01-16', 'GOOG'),
        (100, 143.20, '2024-01-17', 'GOOG'),
        (500, 380.50, '2024-01-15', 'MSFT'),
        (300, 385.40, '2024-01-16', 'MSFT'),
        (200, 382.10, '2024-01-17', 'MSFT')
    ''')
    con.delete_function("pythonvwap", database=target_database, if_it_exists=True)
    con.create_agg_function(
        PythonVWAP,
        database=target_database,
        signature=aura.Signature(
            parameters=[
                aura.Parameter(name="quantity", annotation=int),
                aura.Parameter(name="price", annotation=float),
                aura.Parameter(name="date", annotation=dt.date),
                aura.Parameter(name="symbol", annotation=str),
            ],
            return_annotation=str,
        ),
        volatility=aura.udf.VolatilityType.IMMUTABLE,
        strict=False,
    )
    ds = con.load_dataset("stock_trades", database=target_database)
    vwap_handler = con.get_function("pythonvwap", database=target_database)
    udaf_builder = AggregateFnBuilder(
        fn=vwap_handler,
        on=[ds.quantity, ds.price, ds.date, ds.symbol],
        as_col="VWAP_column",
        num_cpus=0.5,
    )
    ds = ds.aggregate([udaf_builder], by=[ds.symbol])
    res = ds.execute()
    df = res.copy()
    parsed = df["VWAP_column"].apply(json.loads).apply(pd.Series)
    df = pd.concat([df.drop(columns=["VWAP_column"]), parsed], axis=1)
    print(df)
finally:
    con.close()

相关文档