如何对长度不同的多列列表进行截断式Explode操作？

问题描述 #

我有如下Polars DataFrame：

data = {
    "a": [[1], [2], [3, 4], [5, 6, 7]],
    "b": [[], [8], [9, 10], [11, 12]],
}
df = pl.DataFrame(data)

对应的DataFrame展示为：
┌───────────┬───────────┐
│ a ┆ b │
│ --- ┆ --- │
│ list[i64] ┆ list[i64] │
╞═══════════╪═══════════╡
│ [1] ┆ [] │
│ [2] ┆ [8] │
│ [3, 4] ┆ [9, 10] │
│ [5, 6, 7] ┆ [11, 12] │
└───────────┴───────────┘
每一行的两个列表长度可能不一致，我希望按两个列表的最短长度进行截断式Explode，得到如下结果：
┌─────┬─────┐
│ a ┆ b │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 2 ┆ 8 │
│ 3 ┆ 9 │
│ 4 ┆ 10 │
│ 5 ┆ 11 │
│ 6 ┆ 12 │
└─────┴─────┘
我原本考虑将较短的列表用None填充至相同长度后再drop_nulls，但想知道是否有更直接的实现方式？

更直接的实现方式 #

方法一：计算最短长度后截断列表 #

先计算每行两个列表的最短长度，对列表进行截断后再执行explode，步骤简洁高效，无需填充Null再过滤：

import polars as pl

data = {
    "a": [[1], [2], [3, 4], [5, 6, 7]],
    "b": [[], [8], [9, 10], [11, 12]],
}
df = pl.DataFrame(data)

result = df.with_columns(
    # 计算每行两个列表的最短长度
    min_len=pl.min([pl.col("a").list.lengths(), pl.col("b").list.lengths()]),
    # 截断列表到最短长度
    a=pl.col("a").list.slice(0, pl.col("min_len")),
    b=pl.col("b").list.slice(0, pl.col("min_len")),
).drop("min_len")  # 移除临时计算列
# 执行explode并过滤空行
result = result.explode(["a", "b"]).drop_nulls()

print(result)

方法二：利用arr.zip_with自动配对 #

借助Polars的arr.zip_with函数，直接按两个列表的最短长度生成配对元组，再展开拆包，写法更简洁：

import polars as pl

data = {
    "a": [[1], [2], [3, 4], [5, 6, 7]],
    "b": [[], [8], [9, 10], [11, 12]],
}
df = pl.DataFrame(data)

result = df.with_columns(
    paired=pl.col("a").arr.zip_with(pl.col("b"), lambda x, y: (x, y))
).explode("paired").unnest("paired").drop_nulls()

print(result)

两种方法都能得到你想要的目标结果，且避免了填充Null的额外步骤，执行效率更高。

内容的提问来源于stack exchange，提问作者DJDuque