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

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

SQL в Pandas. У вас есть некоторые данные в реляционной базе данных, и вы хотите обработать их с помощью Pandas. Поэтому вы используете удобный API Pandas read_sql() для получения DataFrame и быстро расходуете память.

Проблема: вы загружаете в память сразу все данные. Если у вас достаточно строк в результатах SQL-запроса, они просто не поместятся в оперативной памяти.

В Pandas есть опция пакетной обработки для read_sql(), которая может уменьшить использование памяти, но она все же не идеальна: она также загружает все данные в память сразу!

Так как же обрабатывать запросы, объем которых превышает объем памяти, с помощью Pandas? Давайте разбираться.

Итерация #1: Просто загрузите данные

В качестве отправной точки рассмотрим наивный, но часто достаточный метод загрузки данных из базы данных SQL в Pandas DataFrame. Вы можете использовать функцию pandas.read_sql() для преобразования SQL-запроса в DataFrame:

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 станет практически нулевым; все потребление памяти связано с импортом библиотек:

профилировщик памяти 3
профилировщик памяти 3

Проблема #3: Разве Pandas не должен делать это по умолчанию?

Pandas, вероятно, должен устанавливать эту опцию автоматически, если задан chunksize , чтобы уменьшить использование памяти. Есть открытый вопрос об этом; надеюсь, кто-то – возможно, вы! – отправит PR.

Экономия памяти с помощью пакетной обработки

С помощью пакетной обработки и серверных курсоров вы можете обрабатывать произвольно большие результаты SQL в виде серии DataFrames, не исчерпывая памяти. Независимо от того, получите ли вы 1000 строк или 10 000 000 000, у вас не закончится память, пока вы храните в памяти только один пакет за раз.

Правда, вы не сможете загрузить все данные сразу. Но довольно часто пакетной обработки достаточно, если не для всей обработки, то хотя бы для начального прохода, суммирующего данные настолько, что вы сможете затем загрузить всю информацию в память.

# SQL в Pandas