Início / user-14744492

pinpss's questions

Martin Hope
pinpss
Asked: 2025-04-04 07:13:12 +0800 CST
  • 6

Estou tentando amostrar aleatoriamente n IDs para cada combinação de group_id e date em um Polars DataFrame. No entanto, notei que a função sample está produzindo o mesmo conjunto de IDs para cada date, não importa o grupo.

Como preciso definir uma semente para fins de replicação, acredito que o problema esteja ocorrendo porque o mesmo valor de semente está sendo aplicado em todas as combinações. Tentei resolver isso criando uma semente exclusiva para cada combinação gerando uma coluna "group_date_int" combinando group_id e date casted como Int64, mas encontrei o seguinte erro:

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

Para cada data, estou obtendo o mesmo conjunto de IDs, em vez de ter uma amostra aleatória diferente para cada combinação de group_id e data.

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"})
)

Além disso, tentei usar a função shuffle, mas os resultados são os mesmos: 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   │
└──────────┴────────────┴─────┘

Fui encaminhado para a seguinte questão nos comentários: Sample from each group in polars dataframe? , onde um problema semelhante foi levantado. No entanto, a solução não inclui uma semente, que é necessária para replicação.

python