Форматы датасетов

Обзорный пост для тех, кто работает с большими датасетами. Недавно мы взяли в работу один датасет объёмом несколько терабайт для задачи Speech‑to‑Text. Датасет представляет собой набор колонок (аудио, текст, длительность и др.) и десятки миллионов строк. Такая здоровенная таблица.

Что мы делаем с датасетом:
🔹 Загружаем с диска в RAM. Причём читаем гораздо чаще, чем сохраняем.
🔹 Берём определённую строку датасета по индексу. В ходе обучения мы собираем батч из рандомных индексов.
🔹 Ведём подсчёт статистики для определённой колонки. Некоторые форматы позволяют загружать не все колонки из таблицы, а только определённые.
🔹 Проводим трансформацию, например, фильтрацию или преобразование строк.
🔹 Храним на долгосрочную перспективу, и здесь нам очень пригодится сжатие для экономии места.

Мы столкнулись с тем, что банальная загрузка в память и подсчёт статистики могут занимать часы. Поэтому решили рассмотреть следующие варианты:
🔹 Форматы хранения на диске:
parquet, pickle, ORC, feather.
🔹 Библиотеки для загрузки и их табличные типы:
polars, pyarrow, pandas, dask.

Лучшие результаты для фрагмента датасета объёмом 4 млн строк:
🔹 Чтение с диска в RAM:
- feather + pyarrow: < 1 сек.;
- parquet + polars: чуть дольше — 2 сек.
🔹 Доступ по индексу:
- у всех: < 0.1 сек.
🔹Подсчёт статистики и базовые математические операции:
- pyarrow.Table — 0,1 сек.

Ещё не в полной мере исследован dask, который обеспечивает быструю распределённую обработку. Если его результаты превзойдут описанные выше, я добавлю их в комментарии. На текущий момент мы предпочитаем остановиться на формате parquet из‑за его большей популярности и библиотеке pyarrow.

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

Еще подобный обзорный пост