Как оптимизировать pandas при работе с большими datasetами (очерк)

Когда памяти вагоны и/или dataset небольшой можно смело закидывать его в pandas безо всяких оптимизаций. Однако, если данные большие, остро встает вопрос, как их обрабатывать или хотя бы считать.

Предлагается взглянуть на оптимизацию в миниатюре, дабы не вытаскивать из сети гигантские датасеты.

В качестве датасета будем использовать хабрастатистику с комментариями пользователей за 2019 г., которая является общедоступной благодаря одному трудолюбивому пользователю:
dataset

В качестве инфо-основы будет использоваться ранее переведенная статья с Хабра, в которой намешано много интересного.

Вместо вступления

Датасет хабрастатистики считается небольшим, хотя и занимает 288 Мб и состоит из 448533 строк.
Разумеется, можно найти и побольше данных, но, чтобы не вешать машину, остановимся на нем.

Для удобства операций внесем (просто запишем в файл первую строку) названия столбцов:

a,b,c,d

Теперь, если напрямую загрузить dataset в pandas и проверить, сколько он использует памяти

import os
import time
import pandas as pd
import numpy as np
gl = pd.read_csv('habr_2019_comments.csv',encoding='UTF')
def mem_usage(pandas_obj):
    if isinstance(pandas_obj,pd.DataFrame):
        usage_b = pandas_obj.memory_usage(deep=True).sum()
    else: # исходим из предположения о том, что если это не DataFrame, то это Series
        usage_b = pandas_obj.memory_usage(deep=True)
    usage_mb = usage_b / 1024 ** 2 # преобразуем байты в мегабайты
    return "{:03.2f} MB".format(usage_mb)
print (gl.info(memory_usage='deep'))

увидим, что он «кушает» 436.1 MB:

RangeIndex: 448533 entries, 0 to 448532
Data columns (total 4 columns):
a    448533 non-null object
b    448533 non-null object
c    448533 non-null object
d    448528 non-null object
dtypes: object(4)
memory usage: 436.1 MB

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

В столбце b содержатся даты, и, для удобства дальнейших вычислений и наглядности лучше отправить их в index датасета. Для этого изменим код, используемый при считывании датасета:

gl = pd.read_csv('habr_2019_comments.csv', parse_dates=['b'], encoding='UTF')

Теперь даты считываются как index датасета и потребление памяти немного снизилось:

memory usage: 407.0 MB

Теперь оптимизируем данные в самом датасете вне столбцов и индекса

Оптимизация называется: «Оптимизация хранения данных объектных типов с использованием категориальных переменных».

Если перевести на русский язык, то нам необходимо объединить данные в столбцах по категориям, где это эффективно.

Чтобы определить эффективность, необходимо узнать количество уникальных значений в столбцах и если оно будет меньше 50% от общего числа значений в столбце, то объединение значений в категории будет эффективно.

Посмотрим на датасет:

gl_obj=gl.select_dtypes(include=['object']).copy()
gl_obj.describe()

:

            a       c         d
count   448533  448533    448528
unique   25100     185    447059
top      VolCh       0  Спасибо!
freq      3377  260438       184

*столбец с датами в индексе и не отображен

Как видно, из строки unique, в столбцах a и с эффективно объединение в категории. Для столбца а — это 25100 пользователей (явно меньше 448533), для с — 185 значений шкалы с "+" и "-" (тоже значительно меньше 448533).

Оптимизируем столбцы:

for col in gl_obj.columns:
    num_unique_values = len(gl_obj[col].unique())
    num_total_values = len(gl_obj[col])
    if num_unique_values / num_total_values < 0.5:
        converted_obj.loc[:,col] = gl_obj[col].astype('category')        
    else:
        converted_obj.loc[:,col] = gl_obj[col]

Чтобы понять сколько памяти используется для удобства введем функцию:

def mem_usage(pandas_obj):
    if isinstance(pandas_obj,pd.DataFrame):
        usage_b = pandas_obj.memory_usage(deep=True).sum()
    else: # исходим из предположения о том, что если это не DataFrame, то это Series
        usage_b = pandas_obj.memory_usage(deep=True)
    usage_mb = usage_b / 1024 ** 2 # преобразуем байты в мегабайты
    return "{:03.2f} MB".format(usage_mb)

И проверим, была ли оптимизация эффективна:

>>> print('До оптимизации столбцов: '+mem_usage(gl_obj))
До оптимизации столбцов: 407.14 MB
>>> print('После оптимизации столбцов: '+mem_usage(converted_obj))
После оптимизации столбцов: 356.40 MB
>>> 

Как видно, был получен выигрыш еще 50 МВ.

Теперь, поняв, что оптимизация пошла на пользу (бывает и наоборот), зададим параметры датасета при считывании, чтобы сразу учитывать данные, с которыми имеем дело:

gl = pd.read_csv('habr_2019_comments.csv', parse_dates=['b'],index_col='b',dtype ={'c':'category','a':'category','d':'object'}, encoding='UTF')

Желаем быстрой работы с датасетами!

Код для скачивания — здесь.