Pandas: от хаоса к красоте кода

Работа с pandas.DataFrame может превратиться в неловкую кучу старого (не очень) доброго спагетти-кода. Я и мои коллеги часто используем эту библиотеку, и хотя мы стараемся придерживаться хороших практик программирования, таких как разделение кода на модули и модульное тестирование, иногда мы все равно мешаем друг другу, создавая запутанный код.

Я собрала несколько советов и подводных камней, которых следует избегать, чтобы сделать код на pandas чистым. Надеюсь, вам они тоже будут полезны. Также я буду ссылаться на классическую книгу Роберта Мартина «Чистый код: создание, анализ и рефакторинг».

TL;DR:

Не существует единственно правильного способа написания кода, но вот несколько советов для работы с pandas:

«Нет»

• не изменяйте DataFrame слишком сильно внутри функций, потому что так можно потерять контроль над тем, что и где будет добавлено/удалено из него;

• не пишите методы, которые изменяют DataFrame и не возвращают его, потому что это сбивает с толку.

«Да»

• Создавайте новые объекты вместо того, чтобы изменять исходныйDataFrame, и не забывайте делать глубокую копию, когда это необходимо;

• выполняйте только операции аналогичного уровня внутри одной функции;

• разрабатывайте функции с учетом возможности переиспользования;

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

Не надо так

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

Мутабельность

Самое главное, что тут нужно вспомнить: pandas.DataFrame — это изменяемые объекты [2, 3]. Когда вы изменяете мутабельный объект, это затрагивает тот же самый экземпляр, который вы изначально создали, и его физическое расположение в памяти остается неизменным. В отличие от этого, когда вы изменяете неизменяемый объект (например, строку), Python создает новый объект в новом месте памяти и меняет ссылку на новый объект.

Это очень важный момент: в Python объекты передаются в функцию путем присваивания [4, 5]. Посмотрите на картинку ниже: значение df было присвоено переменной in_df, когда она была передана в функцию в качестве аргумента. И исходное значение df, и in_df внутри функции указывают на одну и ту же область памяти (числовое значение в круглых скобках), даже если они имеют разные имена переменных. Во время модификации атрибутов расположение изменяемого объекта остается неизменным.

Изменение мутабельного объекта в памяти

На самом деле, поскольку мы изменили исходный экземпляр, возвращать DataFrame и присваивать его переменной избыточно. Этот код дает точно такой же эффект:

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

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

Напротив, если объект неизменяемый, он будет менять место в памяти в процессе модификации, как в примере ниже. На картинке ниже, поскольку красная строка не может быть изменена (строки неизменяемы), зеленая строка создается как новый объект, занимающий новое место в памяти. Возвращаемая методом строка не является той же самой строкой, тогда как в случае с DataFrame возвращаемый объект был бы ровно тем же DataFrame.

Работа с неизменяемым объектом в памяти

Дело в том, что изменениеDataFrame внутри функций имеет глобальный эффект. Если вы не будете помнить об этом, вы можете:

• случайно изменить или удалить часть данных, думая, что действие происходит только внутри области видимости функции, а это не так;

• потерять контроль над тем, что и когда добавляется в DataFrame, например, при вызове вложенных функций.

Выходные аргументы

Мы исправим эту проблему позже, а сейчас — еще одно «нет», прежде чем мы перейдем к «да».

Конструкция из предыдущего раздела на самом деле является антипаттерном, называемым выходным аргументом [1 стр.45]. Как правило, входные данные функции используются для создания выходного значения. Если единственным смыслом передачи аргумента в функцию является его модификация, то есть входной аргумент меняет свое состояние, то это бросает вызов нашей интуиции. Такое поведение называется побочным эффектом (англ. side effect) [1 стр.44] функции, и оно должно быть хорошо задокументировано, а лучше — сведено к минимуму, поскольку заставляет программиста помнить о том, что происходит в фоновом режиме, а значит, повышает вероятность ошибиться.

When we read a function, we are used to the idea of information going in to the function through arguments and out through the return value. We don’t usually expect information to be going out through the arguments. [1 p.41]

Когда мы читаем функцию, мы привыкли к тому, что информация поступает в функцию через аргументы, а выходит через возвращаемое значение. Обычно мы не ожидаем, что информация будет возвращена через аргументы. [1 p.41]

Ситуация становится еще хуже, если функция несет двойную ответственность: и изменяет входные данные, и возвращает выходные. Рассмотрим эту функцию:

def find_max_name_length(df: pd.DataFrame) -> int:
    df["name_len"] = df["name"].str.len()  # side effect
    return max(df["name_len"])

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

Надо вот так!

Минимизируем модификации объектов

Чтобы устранить побочный эффект, в приведенном ниже коде мы создали новую временную переменную вместо того, чтобы модифицировать исходный DataFrame. Обозначение lengths: pd.Series указывает на тип данных переменной.

def find_max_name_length(df: pd.DataFrame) -> int:
    lengths: pd.Series = df["name"].str.len()
    return max(lengths)

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

Еще одно предупреждение: пожалуйста, помните о различиях между глубоким и поверхностным копированием [6] элементов из DataFrame. В приведенном выше примере мы изменили каждый элемент исходной серии df["name"], поэтому старый DataFrame и новая переменная не имеют общих элементов. Однако если вы напрямую присвоите один из исходных столбцов новой переменной, базовые элементы по-прежнему будут иметь одинаковые ссылки в памяти. Вот примеры:

df = pd.DataFrame({"name": ["bert", "albert"]})

series = df["name"]     # поверхностная копия
series[0] = "roberta"   # <-- изначальный DataFrame изменяется

series = df["name"].copy(deep=True)
series[0] = "roberta"   # <-- изначальный DataFrame не изменяется

series = df["name"].str.title()  # в любом случае не копия
series[0] = "roberta"   # <-- изначальный DataFrame не изменяется

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

Группируем похожие операции

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

Мне нравится правило One Level of Abstraction per Function, которое гласит:

We need to make sure that the statements within our function are all at the same level of abstraction.

Mixing levels of abstraction within a function is always confusing. Readers may not be able to tell whether a particular expression is an essential concept or a detail. [1 p.36]

Нам нужно убедиться, что все действия внутри нашей функции находятся на одном уровне абстракции.

Смешение уровней абстракции в функции всегда приводит к путанице. Читатель может не понять, является ли конкретное выражение существенной концепцией или деталью. [1 стр.36]

Также давайте воспользуемся принципом единственной ответственности [1 стр.138] из ООП, хотя сейчас мы не сосредоточены на объектно-ориентированном коде. (И в принципе ООП даже в контексте Python — это последнее, с чем ассоциируется анализ данных с использованием pandas. (примечание автора перевода))

Почему бы не подготовить данные заранее? Давайте разделим подготовку данных и собственно вычисления на отдельные функции:

def create_name_len_col(series: pd.Series) -> pd.Series:
    return series.str.len()

def find_max_element(collection: Collection) -> int:
    return max(collection) if len(collection) else 0

df = pd.DataFrame({"name": ["bert", "albert"]})
df["name_len"] = create_name_len_col(df.name)
max_name_len = find_max_element(df.name_len)

Отдельная задача создания столбца name_len была передана другой функции. Она не изменяет исходный DataFrame и выполняет одну задачу за раз. Позже мы получим максимальный элемент, передав новый столбец другой специальной функции.

Приведем код в порядок с помощью следующих шагов:

• Мы можем использовать функцию concat и перенести ее в отдельную функцию prepare_data, которая сгруппировала бы все шаги подготовки данных в одном месте;

• Мы также можем воспользоваться методом apply и работать с отдельными текстами, а не с сериями текстов;

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

def compute_length(word: str) -> int:
    return len(word)

def prepare_data(df: pd.DataFrame) -> pd.DataFrame:
    return pd.concat([
        df.copy(deep=True),  # deep copy
        df.name.apply(compute_length).rename("name_len"),
        ...
    ], axis=1)

Переиспользуем код

То, как мы разделили код, позволяет легко вернуться к скрипту позже, взять всю функцию и повторно использовать ее в другом скрипте, и это замечательно!

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

def create_name_len_col(df: pd.DataFrame, orig_col: str, target_col: str) -> pd.Series:
    return df[orig_col].str.len().rename(target_col)

name_label, name_len_label = "name", "name_len"
pd.concat([
    df,
    create_name_len_col(df, name_label, name_len_label)
], axis=1)

Делаем код тестируемым

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

Итак, важные скрипты должны быть протестированы [1, с. 121, 7]. Даже если скрипт — всего лишь помощник, теперь я стараюсь тестировать хотя бы важнейшие, самые низкоуровневые функции. Давайте вернемся к шагам, которые мы сделали с самого начала:

1. Здесь тестируется куча разных функций: вычисление длины имени и агрегирование результата для элемента max . А тест падает c AttributeError: Can only use .str accessor with string values!, хотя мы этого не ожидали, не так ли?

def find_max_name_length(df: pd.DataFrame) -> int:
    df["name_len"] = df["name"].str.len()  # побочный эффект
    return max(df["name_len"])


@pytest.mark.parametrize("df, result", [
    (pd.DataFrame({"name": []}), 0),  # упс, здесь тест упадет
    (pd.DataFrame({"name": ["bert"]}), 4),
    (pd.DataFrame({"name": ["bert", "roberta"]}), 7),
])
def test_find_max_name_length(df: pd.DataFrame, result: int):
    assert find_max_name_length(df) == result

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

def create_name_len_col(series: pd.Series) -> pd.Series:
    return series.str.len()


@pytest.mark.parametrize("series1, series2", [
    (pd.Series([]), pd.Series([])),
    (pd.Series(["bert"]), pd.Series([4])),
    (pd.Series(["bert", "roberta"]), pd.Series([4, 7]))
])
def test_create_name_len_col(series1: pd.Series, series2: pd.Series):
    pd.testing.assert_series_equal(create_name_len_col(series1), series2, check_dtype=False)

3. Здесь мы навели порядок и тестируем саму вычислительную функцию без обёртки в виде pandas. Легче придумать крайние случаи, если сосредоточиться на чем-то одном. Я поняла, что хочу проверить значения None, которые могут появиться в DataFrame, и в итоге мне пришлось усовершенствовать свою функцию, чтобы тест прошел. Баг пойман!

def compute_length(word: Optional[str]) -> int:
    return len(word) if word else 0


@pytest.mark.parametrize("word, length", [
    ("", 0),
    ("bert", 4),
    (None, 0)
])
def test_compute_length(word: str, length: int):
    assert compute_length(word) == length

4. Нам не хватает только теста для find_max_element:

def find_max_element(collection: Collection) -> int:
    return max(collection) if len(collection) else 0


@pytest.mark.parametrize("collection, result", [
    ([], 0),
    ([4], 4),
    ([4, 7], 7),
    (pd.Series([4, 7]), 7),
])
def test_find_max_element(collection: Collection, result: int):
    assert find_max_element(collection) == result

Еще одно преимущество модульного тестирования, о котором я никогда не забываю упомянуть: это способ документирования вашего кода, поскольку тот, кто не знает, что там происходит (например, вы сам из будущего), может легко определить входные данные и ожидаемые результаты, включая крайние случаи, просто взглянув на тесты.

Заключение

Это несколько приемов, которые я сочла полезными при написании кода и чтении чужого кода. Я не утверждаю, что тот или иной способ кодирования является единственно верным — только вы решаете, нужна ли вам быстрая работа или отполированная и протестированная кодовая база. Но надеюсь, эта статья поможет вам структурировать свои скрипты так, чтобы они были красивее и надёжнее.

Буду рада фидбеку и другим комментариям. Счастливого кодинга!

Cписок источников

• [1] Robert C. Martin, Clean code A Handbook of Agile Software Craftsmanship (2009), Pearson Education, Inc.

• [2] pandas documentation - Package overview — Mutability and copying of data

• [3] Python’s Mutable vs Immutable Types: What’s the Difference?

• [4] 5 Levels of Understanding the Mutability of Python Objects

• [5] Pass by Reference in Python: Background and Best Practices

• [6] Shallow vs Deep Copying of Python Objects

• [7] Brian Okken, Python Testing with pytest, Second Edition (2022), The Pragmatic Programmers, LLC.

Иллюстрации были созданы с помощью Miro.