«Переходи на Spark», говорили они. Сравнил pandas, Polars, DuckDB и PySpark на слабой машине

Сценарий знакомый: открываешь в pandas файл побольше, ноутбук задумывается, вентилятор набирает обороты, и через минуту всё падает с MemoryError. Следом обычно приходит совет от старших коллег или из интернета: pandas для маленьких данных, для больших есть Spark. Мне захотелось проверить этот совет цифрами. Не на кластере в облаке, где Spark чувствует себя как дома, а на машине, которая больше похожа на рабочий ноутбук аналитика.

Для этого я написал небольшую программу TableBench. Она сама генерирует данные, прогоняет одни и те же задачи на разных библиотеках, меряет время и память и следит, чтобы все посчитали одно и то же. Код открыт, ссылка в конце, и повторить замеры у себя можно одной командой.

Статья для тех, кто уверенно пишет на pandas и слышал про Polars и DuckDB, но руки до них пока не дошли.

Сравнений этих библиотек в интернете хватает, самое известное из них db-benchmark, который сейчас поддерживает команда DuckDB. Он гоняет задачи на мощном сервере и отвечает на вопрос, кто быстрее. Меня интересовало другое: что происходит, когда памяти мало, на каком объёме каждая библиотека сдаётся и какие ошибки легко допустить по дороге.

Кого и на чём сравнивал

Участников шесть. Обычный pandas и pandas с dtype_backend="pyarrow", когда данные внутри хранятся в формате Arrow. Polars в ленивом режиме, где запрос сначала описывается целиком и выполняется только на collect(), и тот же Polars с потоковым движком collect(engine="streaming"), который обрабатывает данные кусками. DuckDB, встраиваемая аналитическая база, в которой пишешь обычный SQL прямо по parquet-файлам. И PySpark в локальном режиме на всех ядрах.

Задач пять, и все они из повседневной работы:

Код для каждой библиотеки я старался писать так, как на ней пишут в жизни. Никакого экзотического тюнинга, но и никаких заведомо медленных приёмов. Из настроек задал только потолок памяти для Spark и DuckDB (60% доступной памяти) и число shuffle-партиций для Spark по числу ядер.

Версии такие: pandas 3.0.6, pyarrow 25.0.1, Polars 1.44.2, DuckDB 1.5.5, PySpark 4.2.0 на Java 21, Python 3.11.

Данные и машина

Данные устроены по мотивам схемы TPC-H, известного бенчмарка для баз данных. Есть таблица заказов orders и таблица позиций заказов lineitem, основная нагрузка приходится на вторую. Генератор я написал свой: так ничего не нужно скачивать, а один и тот же масштаб на любом компьютере даёт одинаковый набор данных до последней строки.

Машину выбрал намеренно слабую. Это облачная виртуалка с двумя ядрами Intel Xeon 2,1 ГГц, и программе доступно 5,8 ГБ памяти. Примерно так выглядит старый офисный ноутбук. Зато все ограничения здесь вылезают быстро, и хорошо видно, кто на каком объёме сдаётся.

Как мерил, чтобы цифрам можно было верить

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

С памятью пришлось повозиться. Её меряет отдельный поток раз в 50 мс, причём по процессу вместе со всеми дочерними: у PySpark основную память съедает JVM, а не Python, и без этого Spark выглядел бы подозрительно экономным. Короткий всплеск между двумя опросами можно пропустить, поэтому процесс замера в конце ещё и сам сообщает свой пик по данным ОС. На замер отведено не больше 5 ГБ. Кто вылез за потолок, получает пометку «не хватило памяти», иначе машина ушла бы в своп и время перестало бы что-то значить.

Старт Spark-сессии занимает около 5 секунд, и его я вынес за скобки, иначе на маленьких данных Spark проигрывал бы только из-за него. Если библиотека не справилась с операцией на меньшем объёме, на большем её уже не запускаю.

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

Что получилось

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

* DuckDB упёрся не в память, а в свободное место на диске, куда он сбрасывал данные.

Загрузка и сортировка у обоих вариантов Polars выполняются одним и тем же кодом, так что для потокового их отдельно не мерил.

Обе оси на графиках логарифмические. Крестик сверху означает объём, на котором библиотека не справилась.

На маленьких данных медленный только Spark

На 6 млн строк всё, кроме PySpark, укладывается в секунды:

Spark здесь медленнее pandas от 4 до 22 раз, и это без учёта старта сессии. Дело не в том, что Spark плохой. Он создан для распределённых вычислений и любой запрос превращает в план из стадий и задач, которые можно раздать по машинам кластера. Когда данных немного, это планирование обходится дороже самой работы.

Polars и DuckDB на группировке обгоняют pandas примерно в 4 и 6 раз. Разницу между секундой и пятой долей секунды в ноутбуке почти не замечаешь, но дальше она только растёт.

Где заканчивается pandas

Сначала цифра, которая меня удивила. Таблица, которая на диске занимает 149 МБ, после загрузки в pandas съедает около 1,9 ГБ памяти, примерно в 13 раз больше. Parquet хранит данные сжатыми и по колонкам, а в памяти каждое значение разворачивается целиком. Отсюда и результат: при 5 ГБ памяти pandas не смог загрузить таблицу уже на 18 млн строк, хотя на диске это меньше полугигабайта.

Дальше интереснее. На 60 млн строк группировка в pandas упала, а join и оконная функция на том же объёме прошли. Всё решает число колонок: группировке нужно шесть, join три, окну две. Читать только нужные колонки через columns=[...] выглядит банальным советом, но на практике это самое дешёвое ускорение, которое есть у pandas. Все замеры в бенчмарке написаны именно так.

Вариант с pyarrow сдвинул границу загрузки с 18 до 60 млн строк и в целом ест меньше памяти. Но на группировке и сортировке он сломался там же, где обычный pandas.

Polars: быстрый, пока хватает памяти

Polars оказался самым быстрым на загрузке и сортировке и обогнал pandas везде, кроме оконной функции, где idxmax в pandas держался наравне или чуть впереди. Но у обычного collect() есть особенность: он выполняет весь запрос в памяти. Поэтому на 180 млн строк группировка и join в Polars упали.

Лечится одним аргументом. С collect(engine="streaming") Polars обрабатывает данные кусками, и группировка на 180 млн строк прошла за 7,35 с при 2,3 ГБ памяти. Разницу хорошо видно уже на 60 млн строк: обычный режим занял 4,2 ГБ, потоковый 1,2 ГБ.

С сортировкой потоковый режим не помог. Она у Polars и так идёт через потоковый sink_parquet, но всё равно упала по памяти на 18 млн строк.

DuckDB удивил сильнее всех

Я ожидал, что DuckDB будет где-то рядом с Polars. Вышло иначе: он справился со всеми операциями, где ему хватило места на диске, и почти везде был самым быстрым. Группировка на 180 млн строк заняла у него 5,84 с, и памяти на это ушло 0,4 ГБ. Напомню, что pandas на той же операции сдался ещё на 60 млн строк.

Всё потому, что DuckDB не тащит таблицу в память. Он читает из parquet только нужные колонки и считает агрегаты на лету. От пользователя для этого ничего не требуется, достаточно обычного SQL:

import duckdb

con = duckdb.connect()
con.execute("SET memory_limit = '3.5GB'")
rows = con.execute("""
    SELECT l_returnflag, l_linestatus,
           sum(l_quantity), sum(l_extendedprice * (1 - l_discount)), count(*)
    FROM read_parquet('lineitem.parquet')
    WHERE l_shipdate <= DATE '1998-09-02'
    GROUP BY ALL
""").fetchall()

Строчка с memory_limit появилась не сразу. В первом прогоне её не было, и сортировка на 60 млн строк упала по памяти. По умолчанию DuckDB рассчитывает на 80% всей RAM машины и, судя по всему, не учёл лимит памяти виртуалки. С явным потолком он начал сбрасывать лишнее на диск и отсортировал 60 млн строк за 63 с. На 180 млн строк для этого уже не хватило свободного места.

Проигрывает DuckDB только на загрузке в память: 1,81 с против 0,37 с у Polars на 6 млн строк. Оно и понятно, при загрузке он строит полноценную таблицу своей базы данных. Только в реальной работе грузить таблицу целиком в DuckDB обычно незачем, запросы и так идут прямо по файлам.

PySpark медленный, но упрямый

На одной машине Spark был медленнее всех на каждой операции. На 180 млн строк группировка заняла у него 24,1 с против 5,8 с у DuckDB, оконная функция 56,4 с против 8,3 с.

При этом падает он неохотно. Загрузить в кэш 180 млн строк не смог никто, кроме него, хотя ушло на это 426 секунд. Правда, загрузка у разных библиотек означает немного разное: pandas и Polars строят датафрейм в памяти, DuckDB создаёт таблицу своей базы, а Spark материализует кэш, который при нехватке памяти частично уходит на диск. Поэтому эту строку таблицы я бы читал скорее как «кто вообще справился», чем как честную гонку на время. Оконная функция на 180 млн строк тоже прошла. Сломался Spark на join и сортировке при 60 млн строк, причём join прошёл в двух повторах из трёх, то есть работал на самой границе.

Для меня вывод отсюда такой. Spark нужен, когда данные не помещаются на одну машину или когда вокруг уже есть кластер и инфраструктура под него. Если же pandas просто начал тормозить на ноутбуке, Spark, скорее всего, не лучший выход: Polars и DuckDB решают ту же проблему быстрее и проще.

Ниже сравнение с pandas по каждой операции. Для каждой взят самый большой объём, на котором pandas ещё справлялся.

Грабли первые: Spark потерял целый день данных

В первом же прогоне сверка показала, что группировка в PySpark даёт другую сумму, чем остальные пять участников. Расхождение около 0,03%, глазами такое не заметишь.

Причина нашлась в фильтре по дате. Я передавал в него обычный datetime:

df.where(F.col("l_shipdate") <= F.lit(datetime(1998, 9, 2)))

PySpark переводит такой литерал с учётом часового пояса Python-процесса. В Москве это UTC+3, так что граница уезжала на три часа назад, на 21:00 первого сентября. Даты в parquet лежат без часового пояса, и все строки за 2 сентября выпадали из фильтра. На 6 млн строк это 1882 строки.

Исправление простое: передать дату строкой и явно привести к типу без часового пояса.

F.lit("1998-09-02 00:00:00").cast("timestamp_ntz")

Самое неприятное, что join с таким же фильтром по датам сошёлся с остальными. Там сдвигались обе границы диапазона, а данные хранятся с точностью до дня, поэтому набор строк случайно остался прежним. Без автоматической сверки ошибка в группировке могла бы тоже спокойно дожить до публикации.

Грабли вторые: один результат, разница в 9 раз

Оконную функцию «самая дорогая позиция в заказе» на pandas можно написать дословным переводом SQL:

place = df.groupby("l_orderkey")["l_extendedprice"].rank(method="first", ascending=False)
top = df[place == 1]

А можно так:

top = df.loc[df.groupby("l_orderkey")["l_extendedprice"].idxmax()]

Результат совпадает до копейки. Время нет. На 6 млн строк получилось вот что:

В бенчмарке для pandas и Polars стоят самые быстрые варианты, иначе сравнение было бы нечестным. Проверить у себя можно скриптом scripts/window_variants.py.

Выходит, что внутри одного pandas разница между двумя способами записи больше, чем между pandas и DuckDB на этой же операции. Так что прежде чем менять библиотеку, стоит посмотреть, как написан код на текущей.

Грабли третьи: pandas 3.0 уже другой

Долгое время популярным советом было читать данные с dtype_backend="pyarrow". Тогда строки хранятся в формате Arrow, а не как объекты Python, и всё работает заметно быстрее. В pandas 3.0 строки по умолчанию и так хранятся в Arrow, и совет во многом потерял смысл. Я сравнил загрузку 6 млн строк на двух версиях:

На pandas 2 опция ускоряла загрузку почти вдвое, на pandas 3 разница небольшая. В пробном прогоне на старом ноутбуке с Windows и pandas 2.1 она доходила примерно до четырёх раз. Если вы ещё сидите на pandas 2, обновление или одна опция при чтении могут дать заметный эффект без переписывания кода. Скрипт для такого сравнения лежит в scripts/pandas_versions.py.

Так что же выбрать

Если данные спокойно помещаются в память, а это примерно до пары гигабайт в памяти или пары сотен мегабайт в parquet, pandas вполне хватит. Читайте только нужные колонки и присмотритесь к тому, как написан код.

Если удобнее думать на SQL, берите DuckDB. В этих замерах он оказался самым экономным по памяти и почти везде самым быстрым, а начать можно в том же ноутбуке, рядом с pandas.

Кто хочет остаться в привычном синтаксисе датафреймов, тому подойдёт Polars. Только для больших объёмов не забывайте про engine="streaming".

Spark имеет смысл, когда данные не помещаются на одну машину или кластер уже есть. Вот там он на своём месте.

И отдельно про память. Если pandas падает на загрузке, это ещё не повод менять инструмент. На этих данных таблица в памяти весила в 13 раз больше, чем на диске, и самым простым лекарством оказалось не читать лишнего.

Ограничения

Всё мерилось на одной слабой машине. На мощном компьютере границы сдвинутся, а соотношения между библиотеками могут измениться. Больше всех от этого страдает Spark: он создавался для кластеров и умеет распределять работу, а здесь ему досталось два ядра. На машине с десятком ядер отставание от DuckDB и Polars, скорее всего, заметно сократится, так что выводы про Spark относятся именно к скромному железу.

Облачная виртуалка шумит. В большинстве замеров разброс между повторами оказался меньше 10%, медиана около 4%, но самые короткие замеры, в доли секунды, иногда давали разовый выброс почти вдвое. Поэтому в результатах медиана, а различия в пределах 10-15% я бы всерьёз не воспринимал. Сырые замеры по каждому повтору лежат в results/results.csv.

Данные синтетические. По схеме они похожи на TPC-H, но это не официальный TPC-H, и сравнивать цифры с опубликованными результатами TPC-H нельзя. Ключи в них распределены равномерно, а в реальных данных часто бывают перекосы, когда на один ключ приходится огромная доля строк, и join на таких данных ведёт себя иначе. Все данные хранятся в parquet, чтение CSV, на котором новички страдают чаще всего, я не мерил. И в целом пять операций не покрывают всего, что встречается в реальной работе.

Как повторить

Нужны Python 3.11 или новее и, для PySpark, Java 17 или 21. Дальше всё просто:

git clone https://github.com/Gregory-Bondarenko/tablebench
cd tablebench
./setup.sh                                 # на Windows: setup.bat
source .venv/bin/activate                  # на Windows: .venv\Scripts\activate
python run.py --scales 0.1 --repeats 1     # быстрая проверка, пара минут
python run.py                              # полный прогон

Программа сама сгенерирует данные, прогонит замеры и соберёт отчёт results/REPORT.md с графиками и таблицами. Если нужна только часть библиотек или операций:

python run.py --scales 1 10 --engines polars duckdb --ops groupby join

На Windows для записи файлов из PySpark понадобятся winutils.exe и hadoop.dll, как их поставить, написано в README.

Про такие разборы, работу с данными и в целом про свою жизнь, пишу в телеграм-канале ЛОГОВО.DATA, заглядывайте, если было интересно