Stream async iterator very slow

[comodore31]

Hello,

Parsing the iterator returned by the async function stream()
https://www.psycopg.org/psycopg3/docs/api/cursors.html?highlight=stream#psycopg.AsyncCursor.stream
is 100 times slower than parsing the returning of fetchall()

eg:

from datetime import datetime
import sys
import asyncio
import psycopg

async def main():

    async with await psycopg.AsyncConnection.connect(conninfo='postgresql://xxxxxxxxx', connect_timeout=7200, autocommit=True) as con:
        async with con.cursor() as cur:

            t = datetime.now()
            i = 0
            for row in (await (await cur.execute("select 'test' a, t b from generate_series(1, 100000) t")).fetchall()):
                i += 1
            print(datetime.now() - t)
            # 0.08s

            t = datetime.now()
            i = 0
            async for row in cur.stream("select 'test' a, t b from generate_series(1, 100000) t"):
                i += 1
            print(datetime.now() - t)
            # 8.35s !!!!!!! x100 slower

if __name__ == '__main__':

    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

    asyncio.run(main())

Results:

0:00:00.088759
0:00:08.105972

[dvarrazzo]

I'm not entirely surprised that stream() is slower, but maybe 100x is highlighting some unneeded work. Maybe something we should do only once per query in the Transformer and instead we do it for every row. We can take a look at the execution profile.

[dvarrazzo]

I made a quick test with a modified version of your test program. Here I'm playing sync objects rather than async ones, but the ratio of slowdown is the same. For me stream is 50x slower than fetchall (on Linux, connecting to localhost).

Test program

from datetime import datetime
from argparse import ArgumentParser
import sys
import asyncio
import psycopg

DSN = "host=localhost"


def main() -> None:

    with psycopg.connect(conninfo=DSN, connect_timeout=7200, autocommit=True) as con:
        with con.cursor() as cur:

            if arg.method == "fetch":
                t = datetime.now()
                i = 0
                for row in cur.execute(
                    "select 'test' a, t b from generate_series(1, 100000) t"
                ).fetchall():
                    i += 1
                d1 = datetime.now() - t
                print("sync fetchall", d1)
            else:
                t = datetime.now()
                i = 0
                for row in cur.stream(
                    "select 'test' a, t b from generate_series(1, 100000) t"
                ):
                    i += 1
                d2 = datetime.now() - t
                print("sync stream", d2)


async def amain() -> None:

    async with await psycopg.AsyncConnection.connect(
        conninfo=DSN, connect_timeout=7200, autocommit=True
    ) as con:
        async with con.cursor() as cur:

            if arg.method == "fetch":
                t = datetime.now()
                i = 0
                for row in await (
                    await cur.execute(
                        "select 'test' a, t b from generate_series(1, 100000) t"
                    )
                ).fetchall():
                    i += 1
                d1 = datetime.now() - t
                print("async fetchall", d1)

            else:
                t = datetime.now()
                i = 0
                async for row in cur.stream(
                    "select 'test' a, t b from generate_series(1, 100000) t"
                ):
                    i += 1
                d2 = datetime.now() - t
                print("async stream", d2)


if __name__ == "__main__":

    parser = ArgumentParser()
    parser.add_argument("sync", choices=["sync", "async"])
    parser.add_argument("method", choices=["fetch", "stream"])
    arg = parser.parse_args()
    if arg.sync == "sync":
        main()
    else:
        if sys.platform == "win32":
            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

        asyncio.run(amain())

The result of running a profiler on fetchall() is the following flame graph:

Profiling stream() we get instead:

Psycopg conversion work is all in the two spikes on the right. So the overhead seems to be all in the libpq. stream() is implemented using the libpq's single-row mode, so we call PQconsumeInput() once per row instead of once per result.

I don't know if this overhead is expected. Will ask PostgreSQL devs.

[dvarrazzo]

Before writing them, I made a step-by-step run.

Sequence diagram of `cursor.stream()`

sequenceDiagram
    Caller ->> Cursor: stream(query)
    Cursor ->> Stream: <<create>>
    Cursor -->> Caller: it := Stream()

    Caller ->> Stream: next(it)
    activate Stream

    Stream ->> Cursor: _stream_send_gen()
    Note over Cursor, Stream: query sent

    Stream ->> Cursor: _stream_fetchone_gen(first=True)
    Cursor ->> FetchGen: <<create>>
    Cursor -->> Stream: gen = FetchGen()
    Stream ->> Conn: wait(gen)
    Conn ->> FetchGen: next()
    activate FetchGen

    FetchGen ->> PGConn: PQconsumeInput()
    PGConn ->> Kernel: recvfrom()
    Kernel -->> PGConn: 117 bytes
    
    FetchGen ->> PGConn: PQisBusy()
    PGConn -->> FetchGen: 0

    FetchGen ->> PGConn: PQgetResult()
    PGConn -->> FetchGen: res = PGresult(SINGLE_TUPLE)
    FetchGen -->> Conn: res

    deactivate FetchGen

    Conn ->> FetchGen: next()
    FetchGen -->> Conn: StopIteration
    Conn -->> Stream: res
    Stream -->> Caller: Record
    
    loop from second record

        Caller ->> Stream: next(it)
        Stream ->> Cursor: _stream_fetchone_gen(first=False)
        Cursor ->> FetchGen: <<create>>
        Cursor -->> Stream: gen = FetchGen()
        Stream ->> Conn: wait(gen)
        Conn ->> FetchGen: next()
        activate FetchGen

        FetchGen ->> PGConn: PQconsumeInput()
        PGConn ->> Kernel: recvfrom()
        Kernel -->> PGConn: -1 EAGAIN
        
        FetchGen ->> PGConn: PQisBusy()
        PGConn -->> FetchGen: 0

        FetchGen ->> PGConn: PQgetResult()
        PGConn -->> FetchGen: res = PGresult(SINGLE_TUPLE)
        FetchGen -->> Conn: res

        deactivate FetchGen

        Conn ->> FetchGen: next()
        FetchGen -->> Conn: StopIteration
        Conn -->> Stream: res
        Stream -->> Caller: Record
    end

    note right of Caller: records finished
    Caller ->> Stream: next(it)
    Stream ->> Cursor: _stream_fetchone_gen(first=False)
    Cursor ->> FetchGen: <<create>>
    Cursor -->> Stream: gen = FetchGen()

    Stream ->> Conn: wait(gen)
    Conn ->> FetchGen: next()
    activate FetchGen

    FetchGen ->> PGConn: PQconsumeInput()
    PGConn ->> Kernel: recvfrom()
    Kernel -->> PGConn: -1 EAGAIN
    
    FetchGen ->> PGConn: PQisBusy()
    PGConn -->> FetchGen: 0

    FetchGen ->> PGConn: PQgetResult()
    PGConn -->> FetchGen: res = PGresult(TUPLES_OK)
    FetchGen -->> Conn: res

    deactivate FetchGen

    Conn ->> FetchGen: next()
    FetchGen -->> Conn: StopIteration
    Conn -->> Stream: None

    Stream ->> Conn: wait(gen)
    Conn ->> FetchGen: next()
    activate FetchGen

    FetchGen ->> PGConn: PQconsumeInput()
    PGConn ->> Kernel: recvfrom()
    Kernel -->> PGConn: -1 EAGAIN
    
    FetchGen ->> PGConn: PQisBusy()
    PGConn -->> FetchGen: 0

    FetchGen ->> PGConn: PQgetResult()
    PGConn -->> FetchGen: res = None
    FetchGen -->> Conn: None

    deactivate FetchGen

    Conn ->> FetchGen: next()
    FetchGen -->> Conn: StopIteration
    Conn -->> Stream: None
    Stream -->> Caller: StopIteration

    deactivate Stream

Loading

(Diagram also here)

I don't see anything abnormal in this flow; all the data is read in the first PQconsumeInput(), but, if I understand the docs correctly, we have to keep on calling it, otherwise PQgetResult() might block.

The strace is consistent: we have a recvfrom() returning data, then one returning EAGAIN for every row from the second one, finally other two: one for the TUPLES_OK result, one for the NULL result signalling the end of the stream. The strace -tt for a 2-records iteration is:

21:36:53.659529 sendto(3, "P\0\0\0>\0select 'test' a, t b from "..., 108, MSG_NOSIGNAL, NULL, 0) = 108
21:36:53.660236 recvfrom(3, 0x1f6a870, 16384, 0, NULL, NULL) = -1 EAGAIN (Resource temporarily unavailable)
21:36:53.660589 epoll_create1(EPOLL_CLOEXEC) = 4
21:36:53.660848 epoll_ctl(4, EPOLL_CTL_ADD, 3, {EPOLLIN|EPOLLONESHOT, {u32=3, u64=3}}) = 0
21:36:53.661099 epoll_wait(4, [{EPOLLIN, {u32=3, u64=3}}], 1023, 100000) = 1
21:36:53.661941 recvfrom(3, "1\0\0\0\0042\0\0\0\4T\0\0\0.\0\2a\0\0\0\0\0\0\0\0\0\0\31\377\377\377"..., 16384, 0, NULL, NULL) = 117
21:36:53.662506 close(4)                = 0
21:36:53.662830 recvfrom(3, 0x1f6a898, 16344, 0, NULL, NULL) = -1 EAGAIN (Resource temporarily unavailable)
21:36:53.663162 recvfrom(3, 0x1f6a884, 16364, 0, NULL, NULL) = -1 EAGAIN (Resource temporarily unavailable)
21:36:53.663359 recvfrom(3, 0x1f6a876, 16378, 0, NULL, NULL) = -1 EAGAIN (Resource temporarily unavailable)

The only thing that gives me some doubt is this grumble in the libpq. As strace above shows, we actually epoll the connection only once. That's because, in our waiting loop, the single-row generator doesn't report the need to wait. I hope this doesn't converts to do excessive work in PQconsumeInput()... Note that the send() at line 345 is probably wrong and to be removed, but it doesn't seem to make any difference.

[dvarrazzo]

Tom Lane suggested a better fetch pattern that avoids the extra calls to recvfrom(). Now the overhead of stream() versus fetchall() is around 3x, which is reasonable as fetchall() is almost entirely a C function, with Python calls only to create objects, whereas stream() calls much more into Python and generators.

The flame graph with the fixed fetch function looks like:

[dvarrazzo]

Fix released in psycopg 3.0.13. Thank you for the report!