[Python] What is the difference between polars.collect_all and polars.LazyFrame.collect

Starting with the example below:

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, 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)
]


The task is really processing different dataframes, do some computations on them and then join back together all results. The code above constructs res which contains all results as a list.

As for joining back together the results, I tried two options as follows.

Option 1:

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


Option 2:

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


Option 1 does collect_all first and then joins all individual dataframes. Option 2 just does all things in an entirely lazy way and performs collect at the very end.

As far as I know, collect will do optimizations under the hood and I should expect option 1 and option 2 have similar performance. However, my benchmarking results show that option 2 takes twice as long as option 1 (21s vs. 10s on my system with 32 cores).

So, is this behavior kind of as expected? Or are there some inefficiencies about the approach I took?

One good thing about option 2 is that it is entirely lazy and it is a preferable approach in a case where we want to have an API which is entirely lazy and return a lazy dataframe and let users determine what to do next. But, from my experiment, performance is sacrificed by a lot. So, wondering if is there a way to do something like option 2 without sacrificing performance (performance comparable to option 1)?

Continue reading...