如何通过逐行求和（SQL Alchemy）来创建布尔掩码？

Posted 2023-04-17

【中文标题】如何通过逐行求和（SQL Alchemy）来创建布尔掩码？

【英文标题】：How can I create a boolean mask by summing row-wise (SQL Alchemy)?

【发布时间】：2021-08-10 03:41:14

【问题描述】：

给定：

一个数字 SQL Alchemy 表 (<class 'sqlalchemy.sql.schema.Table'>) 以及 SQL Alchemy 列 (<class 'sqlalchemy.sql.elements.ColumnClause'>) 的列表（长度可能不同）

我需要在列的逐行总和等于数字的记录上构建一个布尔掩码 True。

例子：

数字 = 6 column_list = [sa.column(col_1), sa.column(col_2), sa.column(col_3)] 表：

col_1	col_2	col_3	col_4
1	2	3	3
2	0	2	8
6	0	0	1

返回：[真，假，真]

警告：必须在执行查询之前执行计算（例如，在where 子句、case 语句等中），而不是在执行 fetchall 或类似操作之后。

我什至不知道从哪里开始。我知道我可以选择列：

sa.select(column_list).select_from(table)

另外，可能还需要一个案例陈述：

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

编辑：我想我越来越近了，但我仍然不知道如何处理列子句列表：

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)

【参考方案1】：

一种方法是使用子查询：

# 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)]