[Python] How to randomly sample n IDs for each combination of group_id and date in a Polars...

I am trying to randomly sample n IDs for each combination of group_id and date in a Polars DataFrame. However, I noticed that the sample function is producing the same set of IDs for each date no matter the group.

Since I need to set a seed for replication purposes, I believe the issue is occurring because the same seed value is being applied across all combinations. I tried to resolve this by creating a unique seed for each combination by generating a "group_date_int" column by combining group_id and date casted as Int64, but I encountered the following error:

.sample(n=n_samples, shuffle=True, seed=pl.col("group_date_int"))
TypeError: argument 'seed': 'Expr' object cannot be interpreted as an integer


For each date, I am getting the same set of IDs, rather than having a different random sample for each combination of group_id and date.

import polars as pl

df = pl.DataFrame(
{
"date": pl.date_range(
pl.date(2010, 1, 1), pl.date(2025, 12, 1), "1mo", eager=True
).implode(),
"group_id": [["bd01", "bd02", "bd03"]],
"ids": [list(range(10))],
}
).explode("date").explode("group_id").explode("ids")

# Parameters
n_samples = 3 # Number of random samples to pick for each group
SEED = 42 # The seed used for sampling

# Create `selected_samples` by sampling `n_samples` IDs per (group_id, date) combination
selected_samples = (
df
.group_by(['group_id', 'date'])
.agg(
pl.col("id")
.sample(n=n_samples, shuffle=True, seed=SEED)
.alias("random_ids")
)
.explode("random_ids")
.select(["group_id", "date", "random_ids"])
.rename({"random_ids": "id"})
)


Additionally, I tried using the shuffle function, but the results are the same: 1,6,5...1,6,5

┌──────────┬────────────┬─────┐
│ group_id ┆ date ┆ id │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 │
╞══════════╪════════════╪═════╡
│ bd01 ┆ 2025-07-01 ┆ 1 │
│ bd01 ┆ 2025-07-01 ┆ 6 │
│ bd01 ┆ 2025-07-01 ┆ 5 │
│ bd01 ┆ 2012-03-01 ┆ 1 │
│ bd01 ┆ 2012-03-01 ┆ 6 │
│ … ┆ … ┆ … │
│ bd03 ┆ 2024-10-01 ┆ 6 │
│ bd03 ┆ 2024-10-01 ┆ 5 │
│ bd01 ┆ 2010-08-01 ┆ 1 │
│ bd01 ┆ 2010-08-01 ┆ 6 │
│ bd01 ┆ 2010-08-01 ┆ 5 │
└──────────┴────────────┴─────┘


I was referred to the following question in the comments: Sample from each group in polars dataframe?, where a similar issue was raised. However, the solution does not include a seed, which is needed for replication.

Continue reading...