UDF错误处理与重试

场景描述

在AI数据工程中，UDF常涉及网络调用、外部模型推理、数据解析等易出错操作。当UDF执行遇到异常时，默认行为是直接抛出异常并中断整个查询。但在实际业务中，用户往往需要更灵活的错误处理策略，例如：

• 静默忽略：部分数据异常不影响整体结果，希望跳过错误行继续执行。
• 日志记录：记录错误信息用于事后排查，但不中断查询。
• 自动重试：网络超时、服务临时不可用等场景，希望自动重试以提升成功率。
• 按异常类型区分：不同类型的异常采用不同策略，例如超时重试、数据格式错误忽略。

UDF错误处理与重试功能允许用户在不修改源UDF代码的前提下，通过wrap_function API为已注册的UDF动态包装错误处理能力，支持忽略、日志记录、指定类型重试、退避策略等组合配置。

使用方式

1. 注册源UDF：按照正常流程注册Python UDF（Scalar/Class/PyarrowVector/PandasVector均可）。
2. 调用wrap_function包装：通过AuraConnection.wrap_function()为源UDF添加错误处理配置，返回一个包装后的可调用对象。

   wrapped = con.wrap_function(
       source_udf_name="my_func",       # 源UDF名称
       source_database="my_db",          # 源UDF所在数据库
       config=ErrorHandlerConfig(...),   # 错误处理配置
       wrapper_name="my_func_wrapped",   # 包装后UDF名称（可选）
       imports=["file1", "file2"],  # 包装UDF的额外导入文件
       packages=["module1", "module2"]  # 包装UDF的额外依赖包
    )

3. 使用包装后的UDF执行查询：将包装后的UDF用于ds.map() 或 df.select()等操作。

   ds = con.load_dataset("my_table", database="my_db")
   result_ds = ds.map(wrapped, on="input_col", as_col="result")
   df = result_ds.execute()
   con.close()

配置项说明

通过ErrorHandlerConfig配置错误处理行为，各参数说明如下：

错误处理优先级

当UDF执行遇到异常时，按以下优先级依次匹配：

┌─ 优先级 1: ignore_on（最高）
│   if ignore_on and isinstance(e, ignore_on):
│       return None                    ← 静默忽略
│
├─ 优先级 2: log_on
│   if log_on and isinstance(e, log_on):
│       _log_error(e, ...)
│       return None                    ← 记录日志 + 返回 None
│
├─ 优先级 3: retry_on（指定类型重试）
│   if retry_on and isinstance(e, retry_on)
│           and retry_count < max_retries:
│       delay = _compute_delay(...)
│       sleep(delay)
│       continue                       ← 重试
│
├─ 优先级 4: on_error=RETRY（兜底全重试）
│   if on_error==RETRY and not retry_on
│           and retry_count < max_retries:
│       delay = _compute_delay(...)
│       sleep(delay)
│       continue                       ← 重试所有异常
│
├─ 优先级 5: on_error=IGNORE
│   if on_error == IGNORE:
│       return None                    ← 静默返回 None
│
├─ 优先级 6: on_error=LOG
│   if on_error == LOG:
│       _log_error(e, ...)
│       return None                    ← 记录日志 + 返回 None
│
└─ 优先级 7: raise（默认 / 重试耗尽）
    raise                              ← 抛出原始异常

关键规则

• ignore_on优先级最高，匹配后直接返回None，不进入重试或日志逻辑。
• log_on优先级高于retry_on，匹配后记录日志并返回None，不进入重试。
• retry_on指定类型重试；on_error=RETRY不指定类型，对所有异常重试。两者共存时retry_on优先。
• 重试耗尽后，落入on_error决定最终行为（RAISE/IGNORE/LOG）。

退避策略

重试间隔由backoff策略和initial_delay_ms计算，所有策略最终受max_delay_ms约束：

所有策略最终受max_delay_ms约束：min(delay, max_delay_ms)。

约束限制

• 自定义异常限制：不支持定义在类中的自定义异常，定义在同名模块中的同名异常会被认为是相同异常。
• 不支持包装BUILTIN函数：wrap_function 仅支持包装Python UDF，不支持数据库内置函数（如 ABS()、UPPER() 等）。
• WRAPPER默认为TEMPORARY：包装后的UDF默认为临时函数，会话结束后自动销毁。
• Vectorized UDF整批失败：对于PyarrowVector / PandasVector 类型的UDF，当批次中任一行触发异常时，整个批次返回None。
• 敏感数据安全：用户自定义日志字段和UDF参数中如果存在敏感数据（如账号、密码等信息），需要用户自行脱敏。

示例

示例1：指定异常类型重试

网络调用UDF，遇到TimeoutError时自动指数退避重试，重试最多3次：

from aura_frame import ErrorHandlerConfig, BackoffStrategy
from aura_frame.multimodal import ai_lake

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 # 设置日志级别
)

try:
    wrapped = con.wrap_function(
        source_udf_name="call_api",
        source_database="my_db",
        config=ErrorHandlerConfig(
            retry_on=[TimeoutError],
            max_retries=3,
            backoff=BackoffStrategy.EXPONENTIAL,
            initial_delay_ms=200,
            max_delay_ms=5000,
        ),
    )
    ds = con.load_dataset("api_input", database="my_db")
    result_ds = ds.map(wrapped, on="input", as_col="result")
    df = result_ds.execute()
finally:
    con.close()

示例2：多策略组合

源UDF同时支持重试、日志降级和静默忽略，异常匹配按ignore_on、log_on、retry_on、on_error(默认RAISE) 的优先级依次判断：

1. IndexError：静默忽略返回None，不记录日志、不重试。
2. ValueError/TypeError：记录日志后返回None，因配置了local_log_path="/tmp/udf_errors"，写日志到该目录下。
3. ConnectionError/TimeoutError：
   1. 重试最多5次，每次重试前：记录日志（同上）->计算延迟->sleep->重试。
   2. 当5次重试全部失败后再次异常：on_error默认为RAISE，抛出原始异常，该行执行失败。
4. 其他未配置的异常（如 RuntimeError、KeyError 等）：on_error默认为RAISE，抛出原始异常，该行执行失败。

wrapped = con.wrap_function(
    source_udf_name="complex_model",
    source_database="my_db",
    config=ErrorHandlerConfig(
        retry_on=[ConnectionError, TimeoutError],
        log_on=[ValueError, TypeError],
        ignore_on=[IndexError],
        max_retries=5,
        backoff=BackoffStrategy.LINEAR,
        initial_delay_ms=500,
        local_log_path="/tmp/udf_errors.log",
    ),
)

import random
import pandas
from aura_frame.error_handler import (
    BackoffStrategy,
    ErrorHandlerConfig,
    ObsLogConfig,
    OnErrorAction,
)
from aura_frame.multimodal import ai_lake
class InvalidInputError(Exception):
    """无效输入"""
class ExceedsLength(RuntimeError):
    """字符串超过长度"""
class WriteFileError(Exception):
    """写文件失败"""
class DataTransformer:
    """调用外部 HTTP 服务进行数据转换的类 UDF"""
    def __init__(self):
        self.handler = open("/tmp/demo.txt", "a+")  # noqa: SIM115
    def __call__(self, col: str) -> int:
        import os
        log_extra = globals().get("log_extra", None)
        if isinstance(log_extra, dict):
            log_extra.update(
                {"ak": os.getenv("UDF_TEMP_AK"), "sk": os.getenv("UDF_TEMP_SK")}
            )
        if not isinstance(col, str) or not col:
            raise InvalidInputError("The input cannot be an empty string.")
        if len(col) > 10:
            raise ExceedsLength("The input string is exceeds length.")
        if random.random() < 0.5:  # 模拟 50% 的概率发生写文件失败
            raise WriteFileError("Failed to write to file.")
        try:
            self.handler.write(col + "\n")
        except Exception as e:  # noqa: BLE001
            raise WriteFileError(f"Failed to write to file: {e}")
        return id(self.handler)
    def __del__(self):
        self.handler.close()
session = ai_lake.connect(
    aura_endpoint="...",
    aura_endpoint_name="...",
    aura_workspace_id="...",
    lf_catalog_name="...",
    access_key="...",
    secret_key="...",
    security_token="...",
    default_database="err_handler",
    use_single_cn_mode=True,
)
session.sql("create schema if not exists err_handler;")
session.sql("set current_schema to err_handler;")
session.set_function_staging_workspace(
    obs_directory_base="...",
    obs_bucket_name="...",
    obs_server="...",
    access_key="...",
    secret_key="...",
    security_token="...",
)
session.delete_function("data_transformer_src", if_it_exists=True)
source_udf = session.create_scalar_function(
    DataTransformer,
    name="data_transformer_src",
)
def on_retry_cleanup(exc, instance):
    """重试前根据异常类型执行不同清理逻辑"""
    if instance is None:
        return  # __init__ 阶段，对象尚未创建
    # is_exception_match 由框架自动注入到 __globals__
    if globals()["is_exception_match"](exc, WriteFileError):
        instance.handler.close()
    instance.handler = open("/tmp/demo.txt", "a+")  # noqa: SIM115
def custom_log_fn(exc, log_context):
    """自定义日志格式：脱敏 + 添加业务标签"""
    import traceback
    record = dict(log_context)
    record["exc_msg"] = traceback.extract_tb(exc.__traceback__)
    record["sk"] = "******"  # 敏感信息脱敏
    return record
# 包装源 UDF
wrapped_udf = session.wrap_function(
    source_udf_name="data_transformer_src",
    wrapper_name="data_transformer",
    config=ErrorHandlerConfig(
        # 异常分类处理
        on_error=OnErrorAction.LOG,
        ignore_on=[InvalidInputError],  # 无效输入 -> 静默跳过
        log_on=[ExceedsLength],  # 格式错误 -> 日志 + 跳过
        retry_on=[WriteFileError],  # 瞬时错误 -> 重试
        # 重试退避策略
        backoff=BackoffStrategy.EXPONENTIAL,
        # 调用阶段重试控制
        max_retries=10,  # 最多重试10次
        initial_delay_ms=100,
        max_delay_ms=5000,
        retry_cleanup=on_retry_cleanup,
        # 初始化阶段重试
        init_on_error=OnErrorAction.RETRY,
        init_max_retries=2,
        init_retry_on=[ConnectionError],
        init_retry_cleanup=on_retry_cleanup,
        # 日志定义
        log_fn=custom_log_fn,  # 自定义日志格式
        local_log_path="/tmp/...",  # 本地日志目录
        obs_log_config=ObsLogConfig(
            obs_server="...",
            obs_bucket_name="...",
            obs_directory_base="...",
            max_buffer_size=4096,
            flush_interval_s=300.0,
            obs_log_format="parquet",
        ),
    ),
)
# 使用包装后的UDF查询（使用方式与源UDF完全一致）
ds = session.from_pandas(pandas.DataFrame({"col": ["", "123", "1234", "12345", "123456", "1234567", "123456789012345"]}))
ds = ds.map(wrapped_udf, on=[ds.col])
ds.show()
session.close()