How can I create a boolean mask by summing row-wise (SQL Alchemy)?

user_7822853 0
Viewed 109

Given:

  • a number
  • a SQL Alchemy table (<class 'sqlalchemy.sql.schema.Table'>)
  • and a list (that can vary in length) of SQL Alchemy columns (<class 'sqlalchemy.sql.elements.ColumnClause'>)

I need to build a boolean mask that is True on the records where the row-wise sum of the columns is equal to the number.

Example:

  • number = 6
  • column_list = [sa.column(col_1), sa.column(col_2), sa.column(col_3)]
  • table:
col_1 col_2 col_3 col_4
1 2 3 3
2 0 2 8
6 0 0 1

Returns: [True, False, True]

Caveat: The computation must be performed prior to executing the query (e.g. in the where clause, case statement, etc), not after performing fetchall or similar.

I'm not even sure where to start. I know I can select the columns:

sa.select(column_list).select_from(table)

Also, a case statement will probably be needed:

sa.case([(ROW_WISE_SUM == number, True)], else=False)

Edit: I think I'm getting closer but I still can't figure out how to handle the list of column clauses:

sa.case((sa.func.sum(column_list) == number, True), else_=False)

psycopg2.ProgrammingError: can't adapt type 'ColumnClause'

.

.

.

[SQL: CASE WHEN (sum(%(sum_1)s) = %(sum_2)s) THEN %(param_1)s ELSE %(param_2)s END]
[parameters: {'sum_1': [<sqlalchemy.sql.elements.ColumnClause at 0x7fe62bf075d0; col_1>, <sqlalchemy.sql.elements.ColumnClause at 0x7fe62bf6e090; col_2>, <sqlalchemy.sql.elements.ColumnClause at 0x7fe62b0fabd0; col_3>], 'sum_2': 6, 'param_1': True, 'param_2': False}]
(Background on this error at: http://sqlalche.me/e/14/f405)
sqlalchemy python
1 Answers

One way to do it would be to use a subquery:

# create test table
tbl = sa.Table(
    "so68720845",
    sa.MetaData(),
    sa.Column("id", sa.Integer, primary_key=True, autoincrement=False),
    sa.Column("col_1", sa.Integer),
    sa.Column("col_2", sa.Integer),
    sa.Column("col_3", sa.Integer),
    sa.Column("col_4", sa.Integer),
)
tbl.drop(engine, checkfirst=True)
tbl.create(engine)

# insert test data
with engine.begin() as conn:
    conn.execute(
        tbl.insert(),
        [
            {"id": 1, "col_1": 1, "col_2": 2, "col_3": 3, "col_4": 3},
            {"id": 2, "col_1": 2, "col_2": 0, "col_3": 2, "col_4": 8},
            {"id": 3, "col_1": 6, "col_2": 0, "col_3": 0, "col_4": 1},
        ],
    )

# do the test
target_value = 6
col_list = [tbl.c.col_1, tbl.c.col_2, tbl.c.col_3]
col_sum = sum(col_list)

with engine.begin() as conn:
    subq = sa.select(tbl.c.id, col_sum.label("col_sum")).subquery()
    results = conn.execute(
        sa.select(subq.c.id, sa.case((subq.c.col_sum == 6, 1), else_=0))
    ).fetchall()
    print(results)  # [(1, 1), (2, 0), (3, 1)]
user_21443900
Related