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类和实例方法不做限制,用户可以任意添加。
| 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()