Как разбить несколько столбцов List[_] с пропущенными значениями в полях Python?

Вы были на правильном пути, пытаясь заполнить недостающие значения списком нулевых значений правильной длины.

Чтобы pl.Expr.repeat_by работал с нулевым значением, нам нужно убедиться, что базовое выражение имеет ненулевой тип. Этого можно добиться, установив аргумент dtype явности pl.lit.

Затем столбец списка (списков) нулей можно использовать для заполнения нулевых значений в y. После этого одновременное внесение x и y работает как обычно.

(
    df
    .with_columns(
        pl.col("y").fill_null(
            pl.lit(None, dtype=pl.Boolean).repeat_by(pl.col("x").list.len())
        )
    )
)

shape: (3, 2)
┌───────────┬─────────────────────┐
│ x         ┆ y                   │
│ ---       ┆ ---                 │
│ list[i64] ┆ list[bool]          │
╞═══════════╪═════════════════════╡
│ [1]       ┆ [true]              │
│ [1, 2]    ┆ [null, null]        │
│ [1, 2, 3] ┆ [true, false, true] │
└───────────┴─────────────────────┘

Отсюда df.explode("x", "y") должно работать как положено.

Примечание. Если имеется более двух столбцов, каждый из которых может содержать нулевые значения, можно объединить приведенный выше ответ с этим ответом, чтобы получить правильное решение.

Примечание.

См. производительность внизу:

Мне нравится элегантность и интуитивность подхода repeat_by, но я обожаю наказания, поэтому вот подход, который разбивает данные по условию, а затем снова объединяет их. Это хуже, чем простой подход, но может быть полезен для другой операции/варианта использования.

pl.concat(
    [
        part.lazy().select("i", "x", pl.lit(None, pl.Boolean).alias("y")).explode("x")
        if isnull[0]
        else part.lazy().explode("x", "y")
        for isnull, part in df.with_row_index("i").group_by(
            pl.col("y").is_null(), maintain_order=True
        )
    ]
).sort("i").drop("i").collect()

У этого есть добавленный with_row_index, поэтому вы можете сохранить исходный порядок, но если порядок не важен, вы можете удалить его, а также последующую сортировку/отбросить. Это также делает part ленивыми и в конце собирает их. Это связано с тем, что если вы объедините несколько ленивых фреймов, каждый из их планов будет выполняться параллельно. Опять же, если это не важно, вы можете удалить 2 .lazy() и .collect().

Если вы начинаете с ленивого фрейма, вы не можете напрямую использовать group_by в качестве итератора, но можете использовать map_groups, чтобы получить тот же эффект.

Вам нужно создать такую ​​функцию, как:

def part_explode(part: pl.DataFrame):
    if part.select(pl.col('x').first().list.len()==pl.col('y').first().list.len()).item():
        return part.explode('x','y')
    else:
        return part.with_columns(pl.lit(None, pl.Boolean).alias('y')).explode('x')

и тогда ты делаешь

df.group_by(pl.col("y").is_null(), maintain_order=True).map_groups(
    part_explode, schema = {"i": pl.UInt32, "x": pl.Int64, "y": pl.Boolean}
).sort('i').drop('i').collect()

Я не думаю, что map_groups будет распараллеливать части, поскольку он основан на выполнении функции Python, поэтому не используйте этот подход, если вы не начинаете с лени и у вас нет памяти для первой реализации.

Производительность

Настройка с помощью

import polars as pl
import numpy as np
n=1_000_000
df=pl.DataFrame({
    'x':np.random.randint(0,10,n),
    'y':np.random.randint(0,2,n),
    'group':np.random.randint(0, 100_000, n),
}).with_columns(pl.col('y').cast(pl.Boolean)).group_by('group').agg('x','y').with_columns(
    y=pl.when(pl.col('group').mod(20)==0).then(pl.lit(None)).otherwise('y')
).drop('group')

а потом тесты

%%timeit
(
    df
    .with_columns(
        pl.col("y").fill_null(
            pl.lit(None, dtype=pl.Boolean).repeat_by(pl.col("x").list.len())
        )
    ).explode('x','y')
)
31.7 ms ± 5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

против конката сверху

84.1 ms ± 6.59 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)