Limit child collections in initial query sqlalchemy

This answer comes from Mike Bayer on the sqlalchemy google group. I'm posting it here to help folks: TLDR: I used version 1 of Mike's answer to solve my problem because, in this case, I do not have foreign keys involved in this relationship and so cannot make use of LATERAL. Version 1 worked great, but be sure to note the effect of offset. It threw me off during testing for a while because I didn't notice it was set to something other than 0.

Code Block for version 1:

subq = s.query(Messages.date).\
    filter(Messages.user_id == User.id).\
    order_by(Messages.date.desc()).\
    limit(1).offset(10).correlate(User).as_scalar()

q = s.query(User).join(
    Messages,
    and_(User.id == Messages.user_id, Messages.date > subq)
).options(contains_eager(User.messages))

Mike's Answer so you should ignore whether or not it uses "declarative", which has nothing to do with querying, and in fact at first ignore Query too, because first and foremost this is a SQL problem. You want one SQL statement that does this. What query in SQL would load lots of rows from the primary table, joined to the first ten rows of the secondary table for each primary?

LIMIT is tricky because it's not actually part of the usual "relational algebra" calculation. It's outside of that because it's an artificial limit on rows. For example, my first thought on how to do this was wrong:

    select * from users left outer join (select * from messages limit 10) as anon_1 on users.id = anon_1.user_id

This is wrong because it only gets the first ten messages in the aggregate, disregarding user. We want to get the first ten messages for each user, which means we need to do this "select from messages limit 10" individually for each user. That is, we need to correlate somehow. A correlated subquery though is not usually allowed as a FROM element, and is only allowed as a SQL expression, it can only return a single column and a single row; we can't normally JOIN to a correlated subquery in plain vanilla SQL. We can however, correlate inside the ON clause of the JOIN to make this possible in vanilla SQL.

But first, if we are on a modern Postgresql version, we can break that usual rule of correlation and use a keyword called LATERAL, which allows correlation in a FROM clause. LATERAL is only supported by modern Postgresql versions, and it makes this easy:

    select * from users left outer join lateral
    (select * from message where message.user_id = users.id order by messages.date desc limit 10) as anon1 on users.id = anon_1.user_id

we support the LATERAL keyword. The query above looks like this:

subq = s.query(Messages).\
    filter(Messages.user_id == User.id).\
    order_by(Messages.date.desc()).limit(10).subquery().lateral()

q = s.query(User).outerjoin(subq).\
     options(contains_eager(User.messages, alias=subq))

Note that above, in order to SELECT both users and messages and produce them into the User.messages collection, the "contains_eager()" option must be used and for that the "dynamic" has to go away. This is not the only option, you can for example build a second relationship for User.messages that doesn't have the "dynamic" or you can just load from query(User, Message) separately and organize the result tuples as needed.

if you aren't using Postgresql, or a version of Postgresql that doesn't support LATERAL, the correlation has to be worked into the ON clause of the join instead. The SQL looks like:

select * from users left outer join messages on
users.id = messages.user_id and messages.date > (select date from messages where messages.user_id = users.id order by date desc limit 1 offset 10)

Here, in order to jam the LIMIT in there, we are actually stepping through the first 10 rows with OFFSET and then doing LIMIT 1 to get the date that represents the lower bound date we want for each user. Then we have to join while comparing on that date, which can be expensive if this column isn't indexed and also can be inaccurate if there are duplicate dates.

This query looks like:

subq = s.query(Messages.date).\
    filter(Messages.user_id == User.id).\
    order_by(Messages.date.desc()).\
    limit(1).offset(10).correlate(User).as_scalar()

q = s.query(User).join(
    Messages,
    and_(User.id == Messages.user_id, Messages.date >= subq)
).options(contains_eager(User.messages))

These kinds of queries are the kind that I don't trust without a good test, so POC below includes both versions including a sanity check.

from sqlalchemy import *
from sqlalchemy.orm import *
from sqlalchemy.ext.declarative import declarative_base
import datetime

Base = declarative_base()


class User(Base):
    __tablename__ = 'user'
    id = Column(Integer, primary_key=True)
    messages = relationship(
        'Messages', order_by='desc(Messages.date)')

class Messages(Base):
    __tablename__ = 'message'
    id = Column(Integer, primary_key=True)
    user_id = Column(ForeignKey('user.id'))
    date = Column(Date)

e = create_engine("postgresql://scott:tiger@localhost/test", echo=True)
Base.metadata.drop_all(e)
Base.metadata.create_all(e)

s = Session(e)

s.add_all([
    User(id=i, messages=[
        Messages(id=(i * 20) + j, date=datetime.date(2017, 3, j))
        for j in range(1, 20)
    ]) for i in range(1, 51)
])

s.commit()

top_ten_dates = set(datetime.date(2017, 3, j) for j in range(10, 20))


def run_test(q):
    all_u = q.all()
    assert len(all_u) == 50
    for u in all_u:

        messages = u.messages
        assert len(messages) == 10

        for m in messages:
            assert m.user_id == u.id

        received = set(m.date for m in messages)

        assert received == top_ten_dates

# version 1.   no LATERAL

s.close()

subq = s.query(Messages.date).\
    filter(Messages.user_id == User.id).\
    order_by(Messages.date.desc()).\
    limit(1).offset(10).correlate(User).as_scalar()

q = s.query(User).join(
    Messages,
    and_(User.id == Messages.user_id, Messages.date > subq)
).options(contains_eager(User.messages))

run_test(q)

# version 2.  LATERAL

s.close()

subq = s.query(Messages).\
    filter(Messages.user_id == User.id).\
    order_by(Messages.date.desc()).limit(10).subquery().lateral()

q = s.query(User).outerjoin(subq).\
    options(contains_eager(User.messages, alias=subq))

run_test(q)