如何根据索引DataFrame从聚合DataFrame的嵌套列表中提取数据？

给定聚合后的DataFrame df（包含嵌套列表类型的列）和索引DataFrame df_idx，需要从df的列表列中提取df_idx指定位置的元素。

原始数据 #

df（聚合后的数据）

┌──────┬────────────┬───────────┐
│ read ┆ region     ┆ cov       │
│ ---  ┆ ---        ┆ ---       │
│ str  ┆ list[str]  ┆ list[i32] │
╞══════╪════════════╪═══════════╡
│ a    ┆ ["x", "y"] ┆ [25, 10]  │
│ b    ┆ ["x", "z"] ┆ [15, 30]  │
└──────┴────────────┴───────────┘

df_idx（索引数据）

┌──────┬─────────┐
│ read ┆ cov_idx │
│ ---  ┆ ---     │
│ str  ┆ u32     │
╞══════╪═════════╡
│ a    ┆ 0       │
│ b    ┆ 1       │
└──────┴─────────┘

生成数据的示例代码 #

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

fn main() -> PolarsResult<()> {
    let df0 = df![
        "read" => ["a", "a", "b", "b"],
        "region" => ["x", "y", "x", "y"],
        "cov" => [25, 10, 15, 30]
    ]?;
    let df = df0.lazy()
        .group_stable([col("read")])
        .agg([col("*")])
        .collect()?;
    let df_idx = df![
        "read" => ["a", "b"],
        "cov_idx" => [0, 1]
    ]?;
    Ok(())
}

预期结果 #

┌──────┬────────┬─────┐
│ read ┆ region ┆ cov │
│ ---  ┆ ---    ┆ --- │
│ str  ┆ str    ┆ i32 │
╞══════╪════════╪═════╡
│ a    ┆ "x"    ┆ 25  │
│ b    ┆ "z"    ┆ 30  │
└──────┴────────┴─────┘

解决方法 #

通过合并两个DataFrame，再利用Polars的arr.get()方法根据索引提取列表元素，具体实现如下：

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

fn main() -> PolarsResult<()> {
    let df0 = df![
        "read" => ["a", "a", "b", "b"],
        "region" => ["x", "y", "x", "y"],
        "cov" => [25, 10, 15, 30]
    ]?;
    let df = df0.lazy()
        .group_stable([col("read")])
        .agg([col("*")])
        .collect()?;
    let df_idx = df![
        "read" => ["a", "b"],
        "cov_idx" => [0, 1]
    ]?;

    // 合并DataFrame并提取对应索引的元素
    let result_df = df.lazy()
        .join(df_idx.lazy(), on = "read", how = Inner)
        .with_columns([
            col("region").arr.get(col("cov_idx")).alias("region"),
            col("cov").arr.get(col("cov_idx")).alias("cov")
        ])
        .collect()?;

    println!("{}", result_df);
    Ok(())
}

代码说明 #

1. 合并DataFrame：通过join方法以read为关联键，将df和df_idx合并，确保每条数据匹配到对应的索引值。
2. 提取列表元素：调用arr.get()方法，传入cov_idx列作为索引，从region和cov的嵌套列表中提取对应位置的元素，并重命名覆盖原列。

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