Polars中按组重采样后展开列表的实现方法求助

用Polars构建蒙特卡洛模拟器，需求为：按colA列分组，对分组进行重采样，再将聚合得到的列表按组内原始顺序展开，仅打乱组的顺序。目前已完成分组重采样步骤，但无法实现列表展开以得到目标结果。

原始数据 #

生成代码 #

df_original = pl.DataFrame({
   'colA': ['A','A','B','B','C','C'],
   'colB': [11,12,13,14,15,16],
   'colC': [21,22,23,24,25,26]})

数据结构 #

shape: (6, 3)
┌──────┬──────┬──────┐
│ colA ┆ colB ┆ colC │
│ ---  ┆ ---  ┆ ---  │
│ str  ┆ i64  ┆ i64  │
╞══════╪══════╪══════╡
│ A    ┆ 11   ┆ 21   │
│ A    ┆ 12   ┆ 22   │
│ B    ┆ 13   ┆ 23   │
│ B    ┆ 14   ┆ 24   │
│ C    ┆ 15   ┆ 25   │
│ C    ┆ 16   ┆ 26   │
└──────┴──────┴──────┘

已完成的分组重采样代码 #

df_resampled = (
   df_original
      .group_by('colA', maintain_order=True)
      .agg(pl.all())
      .sample(fraction=1.0, shuffle=True, seed=9)
)

重采样后的数据结构 #

shape: (3, 3)
┌──────┬───────────┬───────────┐
│ colA ┆ colB      ┆ colC      │
│ ---  ┆ ---       ┆ ---       │
│ str  ┆ list[i64] ┆ list[i64] │
╞══════╪═══════════╪═══════════╡
│ B    ┆ [13, 14]  ┆ [23, 24]  │
│ C    ┆ [15, 16]  ┆ [25, 26]  │
│ A    ┆ [11, 12]  ┆ [21, 22]  │
└──────┴───────────┴───────────┘

期望展开结果 #

shape: (6, 3)
┌──────┬──────┬──────┐
│ colA ┆ colB ┆ colC │
│ ---  ┆ ---  ┆ ---  │
│ str  ┆ i64  ┆ i64  │
╞══════╪══════╪══════╡
│ B    ┆ 13   ┆ 23   │
│ B    ┆ 14   ┆ 24   │
│ C    ┆ 15   ┆ 25   │
│ C    ┆ 16   ┆ 26   │
│ A    ┆ 11   ┆ 21   │
│ A    ┆ 12   ┆ 22   │
└──────┴──────┴──────┘

解决方案 #

在重采样后的DataFrame上调用explode方法，指定需要展开的列表列即可。该方法会自动保留组内元素的原始顺序，同时维持sample操作打乱后的组顺序。

完整代码 #

df_resampled = (
   df_original
      .group_by('colA', maintain_order=True)
      .agg(pl.all())
      .sample(fraction=1.0, shuffle=True, seed=9)
      .explode(['colB', 'colC'])  # 新增行：展开列表列
)

执行上述代码后，得到的结果与期望完全一致。

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