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")
)
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),
)