SQLAlchemy中如何封装sum()内自定义公式简化查询

你可以直接在Python侧封装普通函数实现逻辑复用，不需要额外注册数据库自定义函数——SQLAlchemy的列、计算表达式本身就是可传递、可组合的Python对象，函数只要返回合法的SQL表达式对象，就能直接嵌入查询语句中。

基础实现 #

首先抽离公共公式逻辑，定义可复用函数：

def my_custom_formula(a_col, b_col, x_col, y_col):
    return a_col * (x_col * (1 + (b_col / (1 - b_col))) + (y_col - x_col))

定义完成后就可以直接按照你预期的形式调用，和原生写法生成的SQL完全一致：

from sqlalchemy import select, func

query = select(
    col_1,
    col_2,
    func.sum(my_custom_formula(a, b, x_1, y_1)).over(partition_by=[col_1, col_2]).label("sum1"),
    func.sum(my_custom_formula(a, b, x_2, y_2)).over(partition_by=[col_1, col_2]).label("sum2"),
    func.sum(my_custom_formula(a, b, x_3, y_3)).over(partition_by=[col_1, col_2]).label("sum3"),
    func.sum(my_custom_formula(a, b, x_4, y_4)).over(partition_by=[col_1, col_2]).label("sum4"),
    func.sum(my_custom_formula(a, b, x_5, y_5)).over(partition_by=[col_1, col_2]).label("sum5")
)

批量生成优化（适配更多同规则列） #

如果后续同规则的x_i、y_i列数量增加，还可以通过循环自动生成聚合字段，进一步减少重复代码：

# 初始化查询字段列表，先放入固定分区列
select_fields = [col_1, col_2]

# 循环生成所有sum聚合字段，这里i的范围可以根据实际列数调整
for i in range(1, 6):
    # 如果你是用表对象定义的列，通过getattr动态取对应列
    current_x = getattr(your_table, f"x_{i}")
    current_y = getattr(your_table, f"y_{i}")
    sum_field = func.sum(
        my_custom_formula(a, b, current_x, current_y)
    ).over(partition_by=[col_1, col_2]).label(f"sum{i}")
    select_fields.append(sum_field)

query = select(*select_fields)

注意：你贴出的初始代码中，公式内使用的是c/(1-c)，和需求描述里的b/(1-b)不一致，封装时注意核对入参和实际使用的列，避免列名写错。

内容的提问来源于stack exchange，提问作者user3557405