Прикольный трюк: сжатие csv файлов 'на лету' в pandas

pandas - великолепный инструмент для работы с данными в python, а csv - де-факто стандартный формат хранения данных в Data Science (да и много где еще).

Однако, csv файлы могут занимать ооочень много места. Если Вы сохраняете какие-то промежуточные данные или регулярно делаете выгрузки из СУБД, то и количество этих файлов может быстро расти.

Если Вам приходится часто двигать файлы через сеть между различными окружениями - сервера/рабочая станция/Google Colab/Kaggle, то этот процесс может превратиться в настоящую головную боль. Большие файлы долго передаются по сети, дисковое пространство в сервисах быстро заканчивается и они начинают требовать от Вас апгрейдить аккаунт и расширять лимиты.

Но есть решение, причем удивительно простое и удобное!

Итак, у нас есть относительно большой csv файл.

user@d14 /tmp # ls -la data.csv
-rw-r--r-- 1 datascience datascience 226M Dec  5 16:07 data.csv

Откроем наш файл на 226MB в pandas:

import pandas

df = pandas.read_csv('data.csv', index_col=0)

df.info()

# <class 'pandas.core.frame.DataFrame'>
# Int64Index: 42367 entries, 0 to 42429
# Columns: 240 entries
# dtypes: bool(4), float64(178), int64(25), object(33)
# memory usage: 76.8+ MB

Как видно, данные тут очень разные: много интов, флоатов, есть также строки. Строки бывают как небольшие, так и приличные json объекты на несколько килобайт.

Теперь идем в документацию: pandas.read_csv? / pandas.pydata.org

compression : str or dict, default 'infer'
If str, represents compression mode. If dict, value at 'method' is
the compression mode. Compression mode may be any of the following
possible values: {'infer', 'gzip', 'bz2', 'zip', 'xz', None}. If
compression mode is 'infer' and path_or_buf is path-like, then
detect compression mode from the following extensions: '.gz',
'.bz2', '.zip' or '.xz'. (otherwise no compression).

То есть: можно на лету сжимать/разжимать csv файлы, и все что нужно - это всего лишь, чтобы файл имел правильное расширение ('.gz', '.bz2', '.zip' или '.xz'). Даже включать никакой флаг не нужно, это дефолтное поведение.

Пробуем!

exts = '', '.gz', '.bz2', '.zip', '.xz'

for ext in exts: df.to_csv(f'test_compression.csv{ext}')

Да, на сжатие ушло какое-то время. Смотрим результат:

user@d14 /tmp # ls -lh test_compression.csv*
-rw-r--r-- 1 user user 223M Dec  6 09:28 test_compression.csv
-rw-r--r-- 1 user user  38M Dec  6 09:29 test_compression.csv.bz2
-rw-r--r-- 1 user user  47M Dec  6 09:29 test_compression.csv.gz
-rw-r--r-- 1 user user  29M Dec  6 09:30 test_compression.csv.xz
-rw-r--r-- 1 user user  48M Dec  6 09:29 test_compression.csv.zip

Вау! Сжатие в 7.5 раз ! Сколько траффика, времени на скачивание/выкачивание, нервов и дискового пространства можно сэкономить!

Разумеется, открывается так же просто, как и сохраняется:

df = pandas.read_csv('test_compression.csv.xz', index_col=0)

А как же время открытия?

Должен же быть подвох! Может, надо ждать полчаса на каждое открытие? Давайте проверим:

%timeit pandas.read_csv('test_compression.csv', index_col=0)
1.58 s ± 2.59 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit pandas.read_csv('test_compression.csv.bz2', index_col=0)
6.16 s ± 5.89 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit pandas.read_csv('test_compression.csv.gz', index_col=0)
2.18 s ± 4.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit pandas.read_csv('test_compression.csv.xz', index_col=0)
3.14 s ± 6.34 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit pandas.read_csv('test_compression.csv.zip', index_col=0)
2.16 s ± 3.71 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Ждать полчаса не придется:)

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

Лучший способ не пропустить новые материалы - оформить подписку на телеграм канал!