Как преобразовать запрос к базе данных с предложением in в список списков в SQLAlchemy

Question

У меня следующий запрос базы данных SQLAlchemy

test_query = Unit.query.filter(Unit.id_1.in_(('3D0U|1|A|G|1', '3D0U|1|A|C|160')))

Я бы хотел, чтобы результатом этого запроса был список списков со строками , соответствующими каждому элементу в в пункте в виде отдельного списка .В настоящее время я получаю все строки в одном списке

Это то, что я пробовал

result = []
for row in test_query:
    result.append[(row.id_2)]

Когда я распечатываю результаты, это то, что я получаю

["3D0X|1|A|C|160", "4ERJ|1|A|C|160", "4ERL|1|A|C|160", "3D0X|1|A|G|1", "4ERJ|1|A|G|1", "4ERL|1|A|G|1"]

Желаемый результат:

[["3D0X|1|A|C|160", "4ERJ|1|A|C|160", "4ERL|1|A|C|160"], ["3D0X|1|A|G|1", "4ERJ|1|A|G|1", "4ERL|1|A|G|1"]]

Пример данных из таблицы единиц измерения приведен ниже

"id_1","chain_1","pdb_id_1","id_2","chain_2","pdb_id_2"
"3D0U|1|A|G|1","A","3D0U","3D0X|1|A|G|1","A","3D0X"
"3D0U|1|A|G|1","A","3D0U","4ERJ|1|A|G|1","A","4ERJ"
"3D0U|1|A|G|1","A","3D0U","4ERL|1|A|G|1","A","4ERL"
"3D0U|1|A|C|160","A","3D0U","3D0X|1|A|C|160","A","3D0X"
"3D0U|1|A|C|160","A","3D0U","4ERJ|1|A|C|160","A","4ERJ"
"3D0U|1|A|C|160","A","3D0U","4ERL|1|A|C|160","A","4ERL"

Любая помощь приветствуется.Спасибо

Answer 1

Это может быть решено следующим образом: сначала в запросе на Unit.id_1, а затем

result = [[unit.id_2 for unit in units] for id_1, units in itertools.groupby(test_query, lambda x: x.id_1)]

В случае postgresql есть также sqlalchemy.func.array_agg, который можно использовать для построениямассив id_2, сгруппированный по id_1.

Для справки документация python для itertools описывает, что делает itertools.groupby.