Выражения и контексты

Основным понятием, используемым в полярах, является выражение . Объекты типа pl.Expr не относятся напрямую к данным, а представляют собой просто выражение, которое можно оценить в контексте (например, в контексте with_columns, select или group_by/agg).

Простым примером является выражение pl.col(...), которое представляет столбец в фрейме данных.

df = pl.DataFrame({
    "x": [1, 2, 3],
    "y": [4, 5, 6],
})

expr = pl.col("x")

Обратите внимание, что в этом примере expr не содержит данных [1, 2, 3], а вместо этого относится к столбцу с именем x. Затем выражение можно оценить в контексте select, чтобы получить фактические данные.

df.select(expr)

shape: (3, 1)
┌─────┐
│ x   │
│ --- │
│ i64 │
╞═════╡
│ 1   │
│ 2   │
│ 3   │
└─────┘

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

expr_squared = expr.pow(2)
pl.select(expr_squared)

shape: (3, 1)
┌─────┐
│ x   │
│ --- │
│ i64 │
╞═════╡
│ 1   │
│ 4   │
│ 9   │
└─────┘

Пользовательская функция (UDF) является особым случаем этого. В качестве входных данных вы передаете выражение x и возвращаете 2 * x. Умножение для выражений, ссылающихся на данные числового типа, перегружено. Следовательно, ваша пользовательская функция возвращает допустимое выражение, которое вы затем оцениваете в контексте.

Структуры

pl.struct(...) — это «просто» выражение, относящееся к набору столбцов.

Вызовы неполярных функций

В частности, концепция выражений объясняет, почему вызовы других стандартных или внешних библиотек терпят неудачу при передаче полярных выражений - они обычно ожидают ввода конкретных данных, как np.random.binomial в вашем примере. Вместо этого вы передаете объекты pl.Expr, ведущие к TypeError.

Однако бывают случаи, когда необходимо использовать функции других библиотек. Для этого существует pl.Expr.map_elements. Обратите внимание, что map_elements вызывается для объекта выражения и возвращает выражение. В качестве параметра map_elements принимает пользовательскую/определяемую пользователем функцию. Когда возвращаемое выражение оценивается в контексте (with_columns или select), пользовательская функция применяется к каждому элементу выражения, к которому был вызван map_elements.

Следовательно, pl.Expr.map_elements предлагает способ применения UDF к данным в контексте. В качестве примера можно использовать следующее для воссоздания квадратного вывода, приведенного выше.

df.select(
    pl.col("x").map_elements(lambda x: x*x, return_dtype=pl.Int64)
)

Примечание. Это всего лишь пример: применение UDF обычно происходит намного медленнее, чем перефразирование операций с использованием API выражений поляров (если это возможно).

Структуры и пользовательские функции

В этом случае структуры являются полезными инструментами, поскольку они позволяют передавать данные из нескольких столбцов в UDF.

df.select(
    pl.struct("x", "y").map_elements(
        lambda s: s["x"] + s["y"],
        return_dtype=pl.Int64,
    )
)

shape: (3, 1)
┌─────┐
│ x   │
│ --- │
│ i64 │
╞═════╡
│ 5   │
│ 7   │
│ 9   │
└─────┘

Давайте вернемся к тому, что такое контекст и что он делает.

Polars DataFrames (или LazyFrame) имеют контексты, которые представляют собой общий способ обращения к with_columns, select, agg и group_by. Входными данными для контекстов являются выражения. В ограниченной степени полярная сторона Python может преобразовывать объекты Python в выражения полярных выражений. Например, datetime или int легко преобразуются в полярное выражение, и поэтому при вводе col('a')*2. Он преобразует это в выражение col('a').mul(lit(2)).

Функции, возвращающие выражения:

Вот ваша функция с аннотациями типов.

def local_double(x: pl.Expr) -> pl.Expr:
  return(2*x)

Он принимает выражение на входе и возвращает другое выражение на выходе. Он не выполняет никакой работы, а просто дает полярам новое выражение. Использование этой функции аналогично использованию df_ab.with_columns(rev_2x = 2*col("revenue")). Фактически, поляры ничего не делают с вашей функцией, когда вы это делаете, потому что порядок операций Python будет разрешать вашу функцию так, чтобы Python мог передавать полярам свои выходные данные в качестве входных данных для контекста поляров.

Зачем нужны df_ab.with_columns(rev_2x = local_double(col("revenue"))) и map_batches

Помните, что полярные контексты ожидают выражений. Одна из причин, по которой Polars настолько быстры и эффективны, заключается в том, что за API стоит собственный язык запросов и механизм обработки. Этот язык говорит выражениями. Чтобы «перевести» с Python то, что он еще не знает, вам нужно использовать одну из функций map_elements. Что они делают, так это преобразуют ваше выражение в значения. В случае map_* он сразу отдаст всю map_batches любой функции, которую вы выберете. В случае pl.Series функция будет получать по одному значению Python за раз. Они представляют собой слой трансляции, благодаря которому поляры могут взаимодействовать с произвольными функциями.

Зачем нам нужно переносить столбцы в структуру?

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

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

(бонус) Помимо функций, возвращающих Exprs, есть ли еще случаи, когда нам не нужны map_elements?

Да. В Numpy есть то, что они называют универсальными функциями или ufunc. Вы можете использовать ufunc непосредственно в контексте, передав ему свой map_* непосредственно в качестве входных данных. Например, вы можете сделать

df.with_columns(log_a = np.log(pl.col('a')))

и это просто сработает. Вы даже можете создать свой собственный ufunc с помощью numba, который тоже будет работать. Механизм работы ufuncs на самом деле такой же, как и у функций, возвращающих выражения, но с большим количеством скрытых шагов. Когда ufunc получает ввод, который не является col('a'), col('b'), вместо того, чтобы выдавать ошибку (как вы получили с np.array), он проверяет, имеет ли ввод np.random.binomial в качестве метода. Если да, то он запустит этот метод. Polars реализует этот метод в __array_ufunc__, поэтому приведенное выше преобразуется в

df.with_columns(log_a = pl.col('a').map_batches(np.log))

Если у вас есть ufunc, который принимает несколько входных данных, он даже автоматически преобразует все эти входные данные в структуру.

Почему вам иногда нужно использовать лямбду?

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

def binomial(x: dict) -> float:
    return np.random.binomial(n=x['daily_n'], p=x['prob'], size=x['size'])


df_sim.with_columns(
  pl.struct(["daily_n", "prob", "size"])
  .map_elements(binomial)
  .cast(pl.Int32)
  .alias('events')
  )

(бонус) Когда использовать pl.Expr, а когда map_elements?

Спойлер: ваш пример должен быть map_batches

Каждый раз, когда вы имеете дело с векторизованной функцией, лучший выбор — map_batches. Я считаю, что большинство (если не все) функций numpy векторизованы, как и функции scipy. Таким образом, ваш пример будет более производительным:

def binomial(x: pl.Series) -> np.array:
    return np.random.binomial(n=x.struct['daily_n'], p=x.struct['prob'])


df_sim.with_columns(
  pl.struct("daily_n", "prob")
  .map_batches(binomial)
  .cast(pl.Int32)
  .alias('events')
  )

Обратите внимание, что я убрал параметр map_batches, потому что numpy определяет размер вывода на основе размера size и daily_n.

Кроме того, когда вы делаете prob на Expr, он становится серией, а не диктом. Чтобы получить доступ к отдельным полям внутри структуры Series, вам нужно использовать пространство имен map_batches, так что между .struct и map_elements нужно учитывать немного другой синтаксис.

Вы также можете сделать это как лямбда, например

df_sim.with_columns(
  pl.struct("daily_n", "prob")
  .map_batches(lambda x: np.random.binomial(n=x.struct['daily_n'], p=x.struct['prob']))
  .cast(pl.Int32)
  .alias('events')
  )

И еще одна упущенная из виду вещь о map_batches

Функция, которую вы даете map_batches, должна возвращать map_batches, за исключением того, что в приведенном выше примере она возвращает pl.Series. У Polars довольно хорошая совместимость с numpy, поэтому он может автоматически конвертировать np.array в np.array. Одна из областей, где вы можете запутаться, — это использование функций pl.Series. Polars не преобразует это автоматически в pyarrow.compute, поэтому вам придется сделать это явно.

В качестве отступления:

Я изложил суть декоратора, который, в принципе, возьмет любую функцию и заставит ее искать метод ввода pl.Series, чтобы вам не приходилось использовать __array_ufunc__. Я говорю «в принципе», потому что я не тестировал его тщательно, поэтому не хочу его переоценивать.