Início / coding / Perguntas / 79554176
Accepted
pinpss
Asked: 2025-04-04 07:13:12 +0800 CST

Como amostrar aleatoriamente n IDs para cada combinação de group_id e date em um Polars DataFrame

  • 772

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

1 respostas

  1. Best Answer

    Se você precisa que cada grupo seja aleatório, mas também precisa ser capaz de definir uma semente para obter resultados previsíveis, então use numpy para gerar números aleatórios e então escolha sua amostra com base neles como este. (Tecnicamente, você poderia usar python base para gerar os números aleatórios, mas é mais lento)

    Primeira abordagem

    n_samples = 3 
SEED = 46
np.random.seed(SEED)
(
    df
    .with_columns(
        pl.col("ids")
        .sort_by(pl.Series(np.random.normal(0,1,df.shape[0]))))
    .group_by("group_id","date",maintain_order=True)
    .agg(pl.col("ids").gather(range(n_samples)))
    .explode("ids")
)

    Observe que também configurei maintain_order=Trueo group_bycomo se isso fosse aleatório.

    Segunda abordagem

    Ter que fazer uma ordenação sobre toda a série pode ser desnecessariamente caro. Se usarmos numpy para criar um array 2d que é ordenado por linha, então usar isso para escolher nossos índices, em teoria, deveria ser mais eficiente.

    No entanto, isso só funciona se você tiver um número fixo de membros por grupo e souber quantos serão com antecedência.

    Primeiro, faça esta função

    def keep_args(members_per_group: int, n_samples: int, rows: int):
    return pl.Series(
        np.argsort(
            np.random.normal(0, 1, (rows, members_per_group)), 
            axis=1)[:, :n_samples],
        dtype=pl.List(pl.Int32),
    )

    Ele vai gerar um array 2d onde cada linha tem uma lista aleatória de índices para escolher. Nós usamos isso com nosso df assim

    np.random.seed(SEED)
(
    df
    .group_by("group_id","date",maintain_order=True)
    .agg(pl.col("ids"))
    .with_columns(
        pl.col("ids").map_batches(lambda s: (
            s.list.gather(keep_args(10, n_samples, s.len()))
        ))
    )
    .explode("ids")
)

    Nesta versão, fazemos o group_byprimeiro, o que significa que precisamos usar map_batchespara obter o novo len de ids. Se preferir, você pode fazer a pipee usar o novo, df.heightmas não acho que faria uma grande diferença de qualquer maneira.

    Diferença de desempenho

    Ao testar esses dois, o primeiro teve 10,4 ms e o segundo teve 9,97 ms, então basicamente o mesmo.

    Terceira abordagem

    Aqui está uma abordagem somente polar que é cerca de 60x mais lenta que a acima. Basicamente, ela apenas corta seu df em grupos individuais e então os amostra.

    pl.concat([
    g.sample(n_samples, seed=SEED) 
    for (_, g) in df.group_by("group_id","date",maintain_order=True)
    ])

    Quarta abordagem

    Você pode converter cada um dos grupos para lazy para obter paralelismo, o que reduz o tempo em 33%, tornando-o apenas 40x mais lento do que as abordagens numpy

    (
    pl.concat([
    g.lazy()
    .select(
        pl.col("group_id","date").first(), 
        pl.col("ids")
        .sample(n_samples, seed=SEED)
        .implode()
        )
    for (_, g) in df.group_by("group_id","date",maintain_order=True)
    ])
.explode("ids")
.collect()
)

    Nota sobre semente

    Talvez isso seja desnecessário dizer, mas, por precaução, o resultado entre cada abordagem será diferente, mesmo com a mesma semente. Os resultados são consistentes apenas dentro de uma abordagem específica. Além disso, apenas para reiterar, você deve usar maintain_order=Truenas duas primeiras abordagens para obter resultados consistentes.

    • 1

relate perguntas