Глубокое погружение в режим Copy-on-Write в pandas. Часть 2

Первая часть статьи.

Введение

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

Я являюсь членом основной команды pandas и принимал активное участие в реализации и улучшении механизма CoW. Я инженер со специализацией по открытому исходному коду в Coiled, где работаю над Dask, включая налаживание интеграции с pandas и обеспечение соответствия Dask требованиям CoW.

Удаление защитных копий

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

df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
df2 = df.reset_index()
df2.iloc[0, 0] = 100

В reset_index нет необходимости копировать данные, но возврат представления привел бы к побочным эффектам при модификации результата (к примеру, также обновился бы df). Поэтому в reset_index выполняется защитное копирование.

При включении Copy-on-Write всех этих защитных копий уже нет. Такое действие затрагивает многие методы (с полным списком можно ознакомиться здесь).

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

Посмотрим, как все это выглядит с точки зрения производительности, когда мы комбинируем некоторые из этих методов:

import pandas as pd
import numpy as np

N = 2_000_000
int_df = pd.DataFrame(
np.random.randint(1, 100, (N, 10)),
columns=[f"col_{i}" for i in range(10)],
)
float_df = pd.DataFrame(
np.random.random((N, 10)),
columns=[f"col_{i}" for i in range(10, 20)],
)
str_df = pd.DataFrame(
"a",
index=range(N),
columns=[f"col_{i}" for i in range(20, 30)],
)

df = pd.concat([int_df, float_df, str_df], axis=1)

Создается DataFrame с 30 столбцами, 3 разными типами данных и 2 миллионами строк. Выполним следующую цепочку методов для этого DataFrame:

%%timeit
(
df.rename(columns={"col_1": "new_index"})
.assign(sum_val=df["col_1"] + df["col_2"])
.drop(columns=["col_10", "col_20"])
.astype({"col_5": "int32"})
.reset_index()
.set_index("new_index")
)

Все указанные методы выполняют защитное копирование без включения CoW.

Производительность без CoW:

2.45 s ± 293 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Производительность при включении CoW:

13.7 ms ± 286 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Производительность улучшилась примерно в 200 раз. Я специально выбрал этот пример, чтобы проиллюстрировать потенциальные преимущества CoW. Не каждым методом можно добиться такой же скорости.

Оптимизация копий, вызванная изменениями в памяти

В предыдущем разделе было продемонстрировано множество методов, при которых защитное копирование больше не требуется. При использовании CoW вы не сможете изменять два объекта одновременно. Это означает, что мы должны ввести копию, когда на одни и те же данные ссылаются два DataFrame. Рассмотрим методы, позволяющие сделать эти копии максимально эффективными.

В первой части мы показали, что к запуску копии может привести следующее:

df.iloc[0, 0] = 100

Копия запускается, если на данные, которые поддерживают df, ссылается другой DataFrame. Наш DataFrame предположительно содержит n столбцов с целочисленными данными (например, поддерживается одним Блоком).

Изображение автора

Объект отслеживания ссылок (Reference tracking object) также ссылается на другой Блок, поэтому мы не можем изменить DataFrame в памяти, не изменив другой объект. Наивным подходом было бы скопировать весь блок целиком и покончить с этим.

Изображение автора

Это действие установило бы новый объект отслеживания ссылок и создало бы новый Блок, поддерживаемый новым массивом NumPy. У этого Блока больше нет ссылок, поэтому другая операция смогла бы опять изменить его в памяти. Данный подход предполагает копирование n-1 столбцов, которые нам не обязательно копировать. Чтобы избежать этого, используем технику, которая называется “разделение блоков”.

Изображение автора

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

У этого метода есть один недостаток. Исходный массив содержит n столбцов. Мы создали представление для столбцов со 2 по n, но это поддерживает весь массив в “живом” состоянии. Мы также добавили новый массив с одним столбцом для первого столбца. Такое действие позволит сохранить немного больше памяти, чем необходимо.

Эта система напрямую приводит к нескольким DataFrame различных типов. Все Блоки, которые вообще не модифицировались, возвращаются как есть. Разделяются только те Блоки, которые были изменены в памяти.

Изображение автора

Теперь мы устанавливаем новое значение в столбец n+1 Блока с float, чтобы создать представление для столбцов с n+2 по m. Новый Блок будет поддерживать только столбец n+1.

df.iloc[0, n+1] = 100.5
Изображение автора

Методы, которые могут работать в памяти

Операции индексации, которые мы рассматривали, обычно не создают новый объект. Они изменяют существующий объект в памяти, включая данные указанного объекта. Другая группа методов pandas вообще не затрагивает данные DataFrame. Одним из ярких примеров этого является rename (переименование). Переименование изменяет только метки. Эти методы имеют возможность использовать механизм отложенного (ленивого) копирования, упомянутый выше.

Существует еще и третья группа методов, которые могут быть выполнены в памяти, например replace и fillna. Они всегда будут запускать копирование.

df2 = df.replace(...)

Изменение данных в памяти без запуска копирования привело бы к изменению df и df2, что нарушает правила CoW. Это одна из причин, по которой мы рассматриваем возможность сохранения ключевого слова inplace для этих методов.

df.replace(..., inplace=True)

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

Вывод

Мы исследуем то, как CoW меняет внутреннее поведение pandas и как это приводит к улучшениям в коде. Многие методы становятся быстрее с CoW, и в то же время мы замечаем замедление в работе некоторых операций, связанных с индексацией. Раньше эти операции всегда проводились в памяти, что могло привести к побочным эффектам. Эти побочные эффекты исчезли с введением CoW, и модификация одного объекта DataFrame никогда не повлияет на другой.

В следующей части  —  о том, как обновить код в соответствии с CoW и каких шаблонов следует избегать в будущем.

Читайте также:

Читайте нас в Telegram, VK и Дзен


Перевод статьи Patrick Hoefler: Deep Dive into pandas Copy-on-Write Mode — Part II

Предыдущая статьяКак разделить монолитное приложение на микрофронтенды
Следующая статьяТехнология составления промптов для модели ИИ на примере одного чат-бота