Search code examples
sqlalchemywindow-functions

SQLAlchemy custom formula in sum()


I have a table with columns: a, b, x_1, x_2, x_3, x_4, x_5, y_1, y_2, y_3, y_4, y_5 and need to apply a formula such as: a * (x_i * (1 + (b / (1 - b)) ) + (y_i - x_i)) where i = {1,2,3,4,5}.

My problem is that I need to sum() the results of this formula partitioned w.r.t two the other columns a, b. So the select query will be like:

select(
    col_1,
    col_2,
    func.sum(
        a * (x_1 * (1 + (c / (1 - c)) ) + (y_1 - x_1))
    ).over(partition_by=[col_1, col_2]).label("sum1"),
    func.sum(
        a * (x_2 * (1 + (c / (1 - c)) ) + (y_2 - x_2))
    ).over(partition_by=[col_1, col_2]).label("sum2"),
    ...
    func.sum(
        a * (x_5 * (1 + (c / (1 - c)) ) + (y_5 - x_5))
    ).over(partition_by=[col_1, col_2]).label("sum5")
)

As you can see it is cumbersome to maintain and keep repeating the formula is error prone.

Is there a way to wrap this formula so that I can simply refer to formula like below or some way of abstracting it?

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_5, y_5)).over(partition_by=[col_1, col_2]).label("sum5")
)

Solution

  • You can hide even more of the implementation detail if you want. The below should give you directional idea:

    def get_expression(T: MyTable, pos: int, c):
        xn = getattr(T, f"x_{pos}")
        yn = getattr(T, f"y_{pos}")
        _expr = T.a * (xn * (1 + (c / (1 - c)) ) + (yn - xn))
        expr = func.sum(_expr).over(partition_by=[T.col_1, T.col_2]).label(f"sum{pos}")
        return expr
    
    q = select(
        T.col_1,
        T.col_2,
        get_expression(T, 1, c),
        get_expression(T, 2, c),
        get_expression(T, 3, c),
        get_expression(T, 4, c),
        get_expression(T, 5, c),
    )