Issue: Massive overhead when iterating over 1k+ rows in postgres even with server side cursors
Author: ParthGandhi
State: closed
Created: 2017-11-15 17:21:30

I'm seeing an inexplicably large overhead when iterating over a postgres table. 

I profiled the code, and also did a smoke test with `SQLAlchemy` to make sure it wasn't a slow connection or the underlying driver (`psycopg2`).

Running this against a postgres table of ~1M records but fetching only a tiny fraction of that.

```python
import time

import peewee
import sqlalchemy
from playhouse import postgres_ext
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.engine.url import URL as AlchemyURL
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker as alchemy_sessionmaker

user = 'XXX'
password = 'XXX'
database = 'XXX'
host = 'XXX'
port = 5432

table = 'person'
limit = 1000

peewee_db = postgres_ext.PostgresqlExtDatabase(
    database=database,
    host=host, port=port,
    user=user, password=password,
    use_speedups=True,
    server_side_cursors=True,
    register_hstore=False,
)

alchemy_engine = sqlalchemy.create_engine(AlchemyURL('postgresql', username=user, password=password,
                                                     database=database, host=host, port=port))
alchemy_session = alchemy_sessionmaker(bind=alchemy_engine)()


class PeeweePerson(peewee.Model):
    class Meta:
        database = peewee_db
        db_table = table

    id = peewee.CharField(primary_key=True, max_length=64)
    data = postgres_ext.BinaryJSONField(index=True, index_type='GIN')


class SQLAlchemyPerson(declarative_base()):
    __tablename__ = table

    id = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True)
    data = sqlalchemy.Column(JSONB)


def run_raw_query():
    ids = list(peewee_db.execute_sql(f"SELECT id from {table} order by id desc limit {limit}"))
    return ids


def run_peewee_query():
    query = PeeweePerson.select(PeeweePerson.id).order_by(PeeweePerson.id.desc()).limit(limit)
    ids = list(query.tuples())
    return ids


def run_sqlalchemy_query():
    query = alchemy_session.query(SQLAlchemyPerson.id).order_by(sqlalchemy.desc(SQLAlchemyPerson.id)).limit(limit)
    ids = list(query)
    return ids


if __name__ == '__main__':
    t0 = time.time()
    raw_result = run_raw_query()
    t1 = time.time()
    print(f'Raw: {t1 - t0}')

    t2 = time.time()
    sqlalchemy_result = run_sqlalchemy_query()
    t3 = time.time()
    print(f'SQLAlchemy: {t3 - t2}')

    t4 = time.time()
    peewee_result = run_peewee_query()
    t5 = time.time()
    print(f'peewee: {t5 - t4}')

    assert raw_result == sqlalchemy_result == peewee_result
```

#### Outputs

- With `limit = 1000`:
```
Raw: 0.02643609046936035
SQLAlchemy: 0.03697466850280762
peewee: 1.0509874820709229
```
- With `limit = 10000`
```
Raw: 0.15931344032287598
SQLAlchemy: 0.07229042053222656
peewee: 10.82826042175293
```

Both examples use server side cursors.

I briefly profiled this, and looks like 95%+ of the time is spent calling `cursor.fetchone` https://github.com/coleifer/peewee/blob/d8e34b0682d87bd56c1a3636445d9c0fccf2b1e2/peewee.py#L2340

I'll continue profiling this, but was wondering if you knew what was up?

My sense is that this has to do with using server-side cursors, but I couldn't be more specific. A further consideration is that when you iterate over a large list of rows, Peewee will cache the results in memory to make subsequent iterations fast. If you only need to do a single-pass, you can add ".iterator()" to your query:

```python
def run_peewee_query():
    query = PeeweePerson.select(PeeweePerson.id).order_by(PeeweePerson.id.desc()).limit(limit)
    ids = list(query.iterator().tuples())
    return ids
```

It's possible that SQLAlchemy may be calling `fetchmany()` or `fetchall()`? I'm not really sure. If the time is being spent in `cursor.fetchone()` then that tells me that any processing in Peewee itself is probably not the cause of the slowness. It might simply be the way peewee is handling the cursor?

Leaving link to psycopg2 docs for handy reference: http://initd.org/psycopg/docs/usage.html#server-side-cursors

I tried the iterator (fixed to `list(query.tuples().iterator())`) but found no improvement.

As far as I can tell you're right - the cursor is not being handled correctly. I'll continue digging into this.

FYI also asked this on SO. Ref: https://stackoverflow.com/q/47313869/2533895

How's the performance when using client-side cursors?

@coleifer Ok, performance with client-side cursors is great - i should've tried that in the first place.

Definitely a bug with the server-side cursor implementation.

I found this issue, too. I try to use this solution in our company:

```python
class QueryResultCursor(object):
    def __init__(self, cursor):
        super(QueryResultCursor, self).__init__()
        self.rows = cursor.fetchall()
        self.row_iterator = iter(self.rows)
        self.close = cursor.close
        self.description = cursor.descrition
        self.name = getattr(cursor, 'name', None)

    def fetchone(self):
        try:
            return self.row_iterator.next()
        except StopIteration:
            return None

class QueryResultWrapper(object):
    def __init__(self, model, cursor, meta=None):
        self.cursor = QueryResultCursor(cursor)
        ...
```

Yes, actually this will fetch all rows at once to memory, and I don't need to change your code any more :)

But you are right, don't use server-side cursors is a better choice

I'll see about fixing this for 3.0a.

@ParthGandhi -- out of curiosity, what version of psycopg2 are you using?

If you're interested, you might want to try out the latest commit in the 3.0a branch, which has a rewritten server-side implementation: 0ae17c519475c935d9db3c338f36ef058a3f879c

`psycopg2==2.7.3.2`

Oooh, thanks. Will check that out.

How'd things go with the new implementation?

Haven't had a chance to try it out yet - got pulled into other work.