Início / coding / Perguntas / 79556592
Accepted
DJDuque
Asked: 2025-04-05 13:34:27 +0800 CST

Como repetir e truncar elementos de lista para um comprimento fixo

  • 772

Tenho dados que parecem:

lf = pl.LazyFrame(
    {
        "points": [
            [
                [1.0, 2.0],
            ],
            [
                [3.0, 4.0],
                [5.0, 6.0],
            ],
            [
                [7.0, 8.0],
                [9.0, 10.0],
                [11.0, 12.0],
            ],
        ],
        "other": ["foo", "bar", "baz"],
    },
    schema={
        "points": pl.List(pl.Array(pl.Float32, 2)),
        "other": pl.String,
    },
)

E eu quero fazer com que todas as listas tenham o mesmo número de elementos. Se atualmente tiver mais do que eu preciso, ela deve truncar. Se tiver menos do que eu preciso, ela deve se repetir em ordem até que tenha o suficiente.

Consegui fazer funcionar, mas sinto que estou pulando obstáculos. Existe uma maneira mais limpa de fazer isso? Talvez com gather?

target_length = 3

result = (
    lf.with_columns(
        needed=pl.lit(target_length).truediv(pl.col("points").list.len()).ceil()
    )
    .with_columns(
        pl.col("points")
        .repeat_by("needed")
        .list.eval(pl.element().explode())
        .list.head(target_length)
    )
    .drop("needed")
)

EDITAR

O método acima funciona para exemplos de brinquedo, mas quando tento usá-lo em meu conjunto de dados real, ele falha com:

pyo3_runtime.PanicException: Polars' maximum length reached. Consider installing 'polars-u64-idx'.

Não consegui fazer um MRE para isso, mas meus dados têm 4 milhões de linhas, e a lista de "pontos" em cada linha tem entre 1 e 8000 elementos (e estou tentando preencher/truncar para 800 elementos). Tudo isso parece bem pequeno, não vejo como um u32comprimento máximo é alcançado.

Agradeço qualquer abordagem alternativa que eu possa tentar.

O mais próximo que tenho (o que não causa pânico) é:

Mas isso não preenche a lista em ordem. Ele apenas preenche a repetição do último elemento.

target_length = 3

result = (
    lf.with_columns(
        pl.col("points")
        .list.gather(
            pl.int_range(target_length),
            null_on_oob=True,
        )
        .list.eval(pl.element().forward_fill())
    )
    .drop("needed")
)
python

1 respostas

  1. Best Answer

    Os padrões repr para listas são bem pequenos, então vamos aumentá-los para o exemplo.

    pl.Config(fmt_table_cell_list_len=8, fmt_str_lengths=120)

    Se você usar pl.int_ranges()(plural) e aritmética de módulo, você pode gerar os índices.

    target_length = 5

lf.select(pl.int_ranges(target_length) % pl.col("points").list.len()).collect()

    shape: (3, 1)
┌─────────────────┐
│ literal         │
│ ---             │
│ list[i64]       │
╞═════════════════╡
│ [0, 0, 0, 0, 0] │
│ [0, 1, 0, 1, 0] │
│ [0, 1, 2, 0, 1] │
└─────────────────┘

    Que você pode passar para.list.gather()

    lf.with_columns(
    pl.col("points").list.gather(
        pl.int_ranges(target_length) % pl.col("points").list.len()
    )    
).collect()

    shape: (3, 2)
┌──────────────────────────────────────────────────────────────────┬───────┐
│ points                                                           ┆ other │
│ ---                                                              ┆ ---   │
│ list[array[f32, 2]]                                              ┆ str   │
╞══════════════════════════════════════════════════════════════════╪═══════╡
│ [[1.0, 2.0], [1.0, 2.0], [1.0, 2.0], [1.0, 2.0], [1.0, 2.0]]     ┆ foo   │
│ [[3.0, 4.0], [5.0, 6.0], [3.0, 4.0], [5.0, 6.0], [3.0, 4.0]]     ┆ bar   │
│ [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0], [7.0, 8.0], [9.0, 10.0]] ┆ baz   │
└──────────────────────────────────────────────────────────────────┴───────┘
    • 2

relate perguntas