Rust Polars：是否可以将列表列分解为多列？

Question

我有一个返回列表类型列的函数。因此，我的专栏之一是一个列表。我想将此列表列变成多列。例如：

use polars::prelude::*;
use polars::df;

fn main() {
    let s0 = Series::new("a", &[1i64, 2, 3]);
    let s1 = Series::new("b", &[1i64, 1, 1]);
    let s2 = Series::new("c", &[Some(2i64), None, None]);
    // construct a new ListChunked for a slice of Series.
    let list = Series::new("foo", &[s0, s1, s2]);

    // construct a few more Series.
    let s0 = Series::new("Group", ["A", "B", "A"]);
    let s1 = Series::new("Cost", [1, 1, 1]);
    let df = DataFrame::new(vec![s0, s1, list]).unwrap();

    dbg!(df);

现阶段 DF 看起来像这样：

┌───────┬──────┬─────────────────┐
│ Group ┆ Cost ┆ foo             │
│ ---   ┆ ---  ┆ ---             │
│ str   ┆ i32  ┆ list [i64]      │
╞═══════╪══════╪═════════════════╡
│ A     ┆ 1    ┆ [1, 2, 3]       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ B     ┆ 1    ┆ [1, 1, 1]       │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ A     ┆ 1    ┆ [2, null, null] │

问题从这里，我想得到：

┌───────┬──────┬─────┬──────┬──────┐
│ Group ┆ Cost ┆ a   ┆ b    ┆ c    │
│ ---   ┆ ---  ┆ --- ┆ ---  ┆ ---  │
│ str   ┆ i32  ┆ i64 ┆ i64  ┆ i64  │
╞═══════╪══════╪═════╪══════╪══════╡
│ A     ┆ 1    ┆ 1   ┆ 2    ┆ 3    │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ B     ┆ 1    ┆ 1   ┆ 1    ┆ 1    │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ A     ┆ 1    ┆ 2   ┆ null ┆ null │

所以我需要类似 .explode() 但按列定向的东西。是否有现有的功能或潜在的解决方法？

非常感谢

Answer 1

是的，你可以。通过 Polars Lazy，我们可以访问表达式 API，并且可以使用

list()

命名空间，通过索引获取元素。

let out = df
    .lazy()
    .select([
        all().exclude(["foo"]),
        col("foo").list().get(0).alias("a"),
        col("foo").list().get(1).alias("b"),
        col("foo").list().get(2).alias("c"),
    ])
    .collect()?;
dbg!(out);

┌───────┬──────┬─────┬──────┬──────┐
│ Group ┆ Cost ┆ a   ┆ b    ┆ c    │
│ ---   ┆ ---  ┆ --- ┆ ---  ┆ ---  │
│ str   ┆ i32  ┆ i64 ┆ i64  ┆ i64  │
╞═══════╪══════╪═════╪══════╪══════╡
│ A     ┆ 1    ┆ 1   ┆ 2    ┆ 3    │
│ B     ┆ 1    ┆ 1   ┆ 1    ┆ 1    │
│ A     ┆ 1    ┆ 2   ┆ null ┆ null │
└───────┴──────┴─────┴──────┴──────┘

Answer 2

此代码在 Rust v1.67 上针对 v0.27.2 中的极坐标进行了测试。

    let out = df
    .lazy()
    .with_columns([
        col("foo").arr().get(lit(0)).alias("a"),
        col("foo").arr().get(lit(1)).alias("b"),
        col("foo").arr().get(lit(2)).alias("c"),
    ])
    .drop_columns(["foo"])
    .collect()?;

    println!("out:\n{out}");

另一种使用for循环的方法：

    let mut lazyframe = df.lazy();
    let column_name: Vec<char> = ('a'..='z').into_iter().collect();

    for (index, ch) in column_name.iter().enumerate().take(3) {
        lazyframe = lazyframe
        .with_columns([
            // split list into new columns
            col("foo").arr().get(lit(index as i64)).alias(&ch.to_string()),
        ])
    }

    let out = lazyframe
    .drop_columns(["foo"])
    .collect()?;

    println!("out:\n{out}");