Partition By VS Group By

Test the performance of polars.DataFrame.partition_by and polars.DataFrame.group_by.

test_partition_by_vs_group_by.py

# -*- coding: utf-8 -*-

import polars as pl
import numpy as np
from dbsnaplake.vendor.timer import DateTimeTimer

# n_row = 1_000
n_row = 1_000_000
# n_row = 100_000_000
df = pl.DataFrame(
    {
        "year": np.random.randint(2001, 2010, n_row),
        "month": np.random.randint(1, 12, n_row),
        "value": np.random.randn(n_row),
    },
)
pkeys = ["year", "month"]
with DateTimeTimer("Method 1") as timer:
    for sub_df in df.partition_by(by=pkeys):
        kvs = sub_df.select(pkeys).head(1).to_dicts()[0]
        sub_df = sub_df.drop(pkeys)
        # sub_df.count()
        # print(kvs, sub_df)

with DateTimeTimer("Method 2") as timer:
    for pvalues, sub_df in df.group_by(pkeys):
        kvs = dict(zip(pkeys, pvalues))
        sub_df.count()
        # print(kvs, sub_df)

Conclusion

If you don’t do any operation in sub dataframe, and you need the partition key values, group by is faster.

If you just want to split the dataframe into sub dataframes, and you don’t need the partition key values, use partition by.