# 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() ```