Rust Porars：是否可以将列表列爆炸到多个列中？

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（）的东西，但是列的东方。是否有可能有可能的funciton？

非常感谢

Answer 1

是的，你可以。通过Polars Lazy，我们可以使用list（）名称空间访问表达式API，以通过索引获取元素。

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

该代码在V0.27.2中的Porars的Rust V1.67上进行了测试。

    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}");

另一种使用循环的方式：

    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}");