В чем разница между Polars.collect_all и Polars.LazyFrame.collect ⇐ Python

[Anonymous]

Начнем с примера ниже:

Код: Выделить всё

import time
import numpy as np
import polars as pl

n_index = 1000
n_a = 10
n_b = 500
n_obs = 5000000

df = pl.DataFrame(
{
"id": np.random.randint(0, n_index, size=n_obs),
"a": np.random.randint(0, n_a, size=n_obs),
"b": np.random.randint(0, n_b, size=n_obs),
"x": np.random.normal(0, 1, n_obs),
}
).lazy()

dfs = [
pl.DataFrame(
{
"id": np.random.randint(0, n_index, size=n_obs),
"a": np.random.randint(0, n_a, size=n_obs),
f"b_{i}": np.random.randint(0, n_b, size=n_obs),
"x": np.random.normal(0, 1, n_obs),
}
).lazy()
for i in range(50)
]

res = [
df.join(
dfs[i], left_on=["id", "a", "b"], right_on=["id", "a", f"b_{i}"], how="inner"
)
.group_by("a", "b")
.agg((pl.col("x") * pl.col("x_right")).sum().alias(f"x_{i}"))
for i in range(50)
]

На самом деле задача состоит в том, чтобы обработать разные кадры данных, выполнить над ними некоторые вычисления, а затем объединить все результаты. Приведенный выше код создает res, который содержит все результаты в виде списка.
Что касается объединения результатов, я попробовал два следующих варианта.
Вариант 1:

Код: Выделить всё

start = time.perf_counter()
res2 = pl.collect_all(res)
res3 = res2[0]
for i in range(1, 50):
res3 = res3.join(res2[i], on=["a", "b"])
time.perf_counter() - start

Вариант 2:

Код: Выделить всё

start = time.perf_counter()
res4 = res[0]
for i in range(1, 50):
res4 = res4.join(res[i], on=["a", "b"])
res4 = res4.collect()
time.perf_counter() - start

Вариант 1 сначала выполняет сбор_все, а затем объединяет все отдельные кадры данных.
Вариант 2 просто выполняет все действия совершенно ленивым способом и выполняет сбор в самый конец.
Насколько я знаю, Collect будет выполнять внутреннюю оптимизацию, и я должен ожидать, что вариант 1 и вариант 2 будут иметь одинаковую производительность. Однако результаты моего сравнительного анализа показывают, что вариант 2 занимает вдвое больше времени, чем вариант 1 (21 с против 10 с в моей системе с 32 ядрами).
Итак, < Strong>Соответствует ли такое поведение ожиданиям? Или есть какие-то неэффективные подходы, которые я использовал?
Одна хорошая вещь в варианте 2 заключается в том, что он совершенно ленив, и это предпочтительный подход в случае, когда мы хотим иметь полностью ленивый API, возвращающий ленивый фрейм данных и позволяющий пользователям определять, что делать дальше. Но, согласно моему эксперименту, производительность во многом приносится в жертву. Итак, интересно, есть ли способ сделать что-то вроде варианта 2, не жертвуя при этом производительностью (производительностью, сравнимой с вариантом 1)?

Подробнее здесь: https://stackoverflow.com/questions/759 ... me-collect