2 个分解极柱的左连接

Question

考虑中

import polars as pl

df = pl.DataFrame({"a": [
    [1, 2],
    [3]],
    "b": [
        [{"id": 1, "x": 1}, {"id": 3, "x": 3}],
        [{"id": 3, "x": 4}]]})

看起来像：

+------+---------------------+
|a     |b                    |
+------+---------------------+
|[1, 2]|[{1,1}, {3,3}]|
|[3]   |[{3,4}]              |
+------+---------------------+

如何

• 为每个展平

  a

  元素获取一行并且
• 如果

  dict

  中的

  b

  列表包含

  a

  元素作为

  id

• 然后将其放入列中

  b

• 否则

  b

  应该是

  null

---

目前的做法

.explode

a

和

b

和

.filter

（内连接）：

df.explode("a").explode("b").filter(
    pl.col("a") == pl.col("b").struct.field('id')
).select(
    pl.col("a"),
    pl.col("b").struct.field("x")
)

不幸的是我只得到（预期的）：

+-+----+
|a|b   |
+-+----+
|1|1   |
|3|4   |
+-+----+

我的目标不是完整的“左连接”：

+-+----+
|a|b   |
+-+----+
|1|1   |
|2|null|
|3|4   |
+-+----+

当DataFrame是这样的结构时，如何高效地得到想要的结果？

Answer 1

您可以执行以下操作。

1. 分别分解列

   a

   和

   b

   。
2. 取消嵌套列

   b

   以获得列

   id

   和

   x

   。
3. 对于

   a

   定义的每个组，计算

   x

   对应的值。

   • pl.when(pl.col("a") == pl.col("id")).then("x")

     创建一个列表，对于每一行（

     a

     、

     id

     和

     x

     对），如果

     x

     与

     a

     匹配，则包含

     id

     ，否则

     None

   • pl.Expr.sort

     将首先放置空值。因此，我们可以使用

     pl.Expr.last

     选择非空值（如果存在），否则使用

     None

     。

(
    df
    .explode("a").explode("b").unnest("b")
    .group_by("a", maintain_order=True)
    .agg(
        pl.when(pl.col("a") == pl.col("id")).then("x").sort().last()
    )
)

输出。

shape: (3, 2)
┌─────┬──────┐
│ a   ┆ x    │
│ --- ┆ ---  │
│ i64 ┆ i64  │
╞═════╪══════╡
│ 1   ┆ 1    │
│ 2   ┆ null │
│ 3   ┆ 4    │
└─────┴──────┘