Загрузка данных SQL в Pandas без потери памяти

import pandas as pd
from sqlalchemy import create_engine

def process_sql_using_pandas():
    engine = create_engine(
        "postgresql://postgres:pass@localhost/example"
    )
    dataframe = pd.read_sql("SELECT * FROM users", engine)
    print(f"Got dataframe with {len(dataframe)} entries")
    # ... do something with dataframe ...


if __name__ == '__main__':
    process_sql_using_pandas()

Если мы запустим это, то увидим, что для данного примера загружается 1 000 000 строк:

$ python pandas_sql_1.py 
Got dataframe with 1000000 entries

Проблема #1: все данные в памяти, несколько раз!

Сколько памяти используется? И откуда берется потребление памяти? Чтобы узнать это, мы можем использовать профилировщик памяти Fil для измерения пикового использования памяти.

$ fil-profile run pandas_sql_1.py
...

Вот как выглядит результат:

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

1. dbapi_cursor.fetchall() извлекает все строки.
2. SQLAlchemy выполняет некоторые дополнительные манипуляции со строками.
3. Pandas преобразует данные в кортежи.
4. Pandas преобразует некоторые данные (кортежи?) в массивы.

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

Загрузка четырех копий данных в память – это слишком много, поэтому давайте посмотрим, можно ли сделать лучше.

Итерация #2: Несовершенное пакетирование

Следующий шаг – использование одной из основных техник уменьшения объема памяти: пакетная обработка или разбивка на части. Во многих случаях вам не нужны все строки в памяти сразу. Если вы можете загрузить данные в виде фрагментов, вы часто можете обрабатывать данные по одному фрагменту за раз, что означает, что вам потребуется только столько памяти, сколько нужно для одного фрагмента.

На самом деле,  pandas.read_sql() имеет API для разбивки на куски, передавая параметр chunksize. Результатом является итерабельность DataFrames:

import pandas as pd
from sqlalchemy import create_engine

def process_sql_using_pandas():
    engine = create_engine(
        "postgresql://postgres:pass@localhost/example"
    )
    for chunk_dataframe in pd.read_sql(
            "SELECT * FROM users", engine, chunksize=1000):
        print(
            f"Got dataframe w/{len(chunk_dataframe)} rows"
        )
        # ... do something with dataframe ...

if __name__ == '__main__':
    process_sql_using_pandas()

Если мы запустим его, то увидим, что код загружает 1000 строк за раз:

$ python pandas_sql_2.py 
Got dataframe w/1000 rows
Got dataframe w/1000 rows
Got dataframe w/1000 rows
...

Проблема #2: все данные в памяти, по-прежнему

Итак, уменьшили ли мы использование памяти? Мы можем снова запустить программу с Fil, получив следующий результат:

С одной стороны, это большое улучшение: мы сократили использование памяти с ~400 МБ до ~100 МБ. С другой стороны, мы, очевидно, все еще загружаем все данные в память в  cursor.execute()!

SQLAlchemy загружает все данные в память, а затем передает Pandas API по 1000 строк за раз, но из локальной памяти. Если наши данные достаточно велики, они все равно не поместятся в памяти.

Итерация # 3: Реальный пакетный метод

Чтобы получить настоящую пакетную обработку, нужно указать SQLAlchemy использовать на стороне сервера, так называемую потоковую обработку. Вместо того чтобы загружать все строки в память, он будет загружать строки из базы данных только тогда, когда они запрашиваются пользователем, в данном случае Pandas. Это работает с различными СУБД, такими как Oracle и MySQL, а не только с PostgreSQL.

Чтобы использовать эту возможность, нам нужно написать наш код немного по-другому:

import pandas as pd
from sqlalchemy import create_engine

def process_sql_using_pandas():
    engine = create_engine(
        "postgresql://postgres:pass@localhost/example"
    )
    conn = engine.connect().execution_options(
        stream_results=True)

    for chunk_dataframe in pd.read_sql(
            "SELECT * FROM users", conn, chunksize=1000):
        print(f"Got dataframe w/{len(chunk_dataframe)} rows")
        # ... do something with dataframe ...

if __name__ == '__main__':
    process_sql_using_pandas()

Как только мы сделаем это изменение, потребление памяти строками базы данных и DataFrame станет практически нулевым; все потребление памяти связано с импортом библиотек: