Writing a query in SQLAlchemy to count occurrences and store IDs

user_12932447 0
Viewed 89

I'm working with a postgres db using SQLAlchemy.

I have a table like this

class Author(Base):
    __tablename__ = "Author"
    id = Column(BIGINT, primary_key=True)
    name = Column(Unicode)

and I want to identify all homonymous authors and save their id in a list.

For example if in the database there are 2 authors named "John" and 3 named "Jack", with ID respectively 11, 22, 33, 44 a 55, I want my query to return

[("John", [11,22]), ("Jack", [33,44,55])]

For now I've been able to write

[x for x in db_session.query(
          func.count(Author.name),
          Author.name
          ).group_by(Author.name) if x[0]>1]

but this just gives me back occurrences

[(2,"John"),(3,"Jack")]

Thank you very much for the help!

sqlalchemy postgresql
1 Answers

The way to do this in SQL would be to use PostgreSQL's array_agg function to group the ids into an array:

SELECT
    name,
    array_agg(id) AS ids
FROM
    my_table
GROUP BY
    name
HAVING
    count(name) > 1;

The array_agg function collects the ids for each name, and the HAVING clause excludes those with only a single row. The output of the query would look like this:

 name  │        ids         
═══════╪════════════════════
 Alice │ {2,4,9,10,16}
 Bob   │ {1,6,11,12,13}
 Carol │ {3,5,7,8,14,15,17}

Translated into SQLAlchemy, the query would look like this:

import sqlalchemy as sa
...
q = (
    db_session.query(Author.name, sa.func.array_agg(Author.id).label('ids'))
    .group_by(Author.name)
    .having(sa.func.count(Author.name) > 1)
)

Calling q.all() will return a list of (name, [ids]) tuples like this:

[
    ('Alice', [2, 4, 9, 10, 16]),
    ('Bob', [1, 6, 11, 12, 13]),
    ('Carol', [3, 5, 7, 8, 14, 15, 17]),
]

In SQLAlchemy 1.4/2.0-style syntax equivalent would be:

with Session() as s:
    q = (
        sa.select(Author.name, sa.func.array_agg(Author.id).label('ids'))
        .group_by(Author.name)
        .having(sa.func.count(Author.name) > 1)
    )
    res = s.execute(q)
user_53209060
Related