使用Arrow / Datafusion / Polars在锈中按列值进行分区例如python panda的groupby？

问题描述

我正在寻找一种等效的方便的python panda语法：

#df is a pandas dataframe

for fruit,sub_df in df.groupby('fruits'):
    # Do some stuff with sub_df and fruit

基本上是一个groupby，每个组都可以在其标签（分组列中的公共值）旁边作为单个数据框进行访问。

我看过数据融合，但是如果不首先选择所有唯一值，然后执行一个select par值，就无法重现此行为，这将导致多次重新解析整个文件。 我看了看Polars板条箱，该板条箱很有希望，但也无法实现我的目标。

您将如何以与python代码类似/更好的性能执行此操作？我愿意接受任何语法/库/方法，这些语法/库/方法都可以使我根据固定列的值有效地对镶木地板文件进行分区。

---

这里是一个使用Polar作为我正在处理的输入类型的示例的生锈示例代码：

    let s0 = Series::new("fruits",["Apple","Apple","Pear","Pear"].as_ref());
    let s1 = Series::new("maturity",["A","B","A","C","D"].as_ref());
    let s1 = Series::new("N",[1,2,4,8].as_ref());
    // create a new DataFrame
    let df = DataFrame::new(vec![s0,s1,s2]).unwrap();

    // I would like to loop on all fruits values,each time with a dataframe containing only the records with this fruit.

解决方法

您可以执行以下操作：

    let columns = df.columns();
    if let Ok(grouped) = df.groupby("fruits") {
        let sub_df = grouped.select(columns).agg_list()?;
        dbg!(sub_df);
    }

从技术上讲，这将为您留一个数据框。然后的问题是，数据框的列是具有每个水果的所有值的数组，这可能不是您想要的。

+---------+--------------------------------------+-----------------------------+-----------------+
| fruits  | fruits_agg_list                      | maturity_agg_list           | N_agg_list      |
| ---     | ---                                  | ---                         | ---             |
| str     | list [str]                           | list [str]                  | list [i32]      |
+=========+======================================+=============================+=================+
| "Apple" | "[\"Apple\",\"Apple\"]"             | "[\"A\",\"B\"]"            | "[1,2]"        |
+---------+--------------------------------------+-----------------------------+-----------------+
| "Pear"  | "[\"Pear\",\"Pear\",... \"Pear\"]" | "[\"A\",\"C\",... \"D\"]" | "[2,4,... 8]" |
+---------+--------------------------------------+-----------------------------+-----------------+

,

您可以使用get_groups方法获取组的first索引以及元组中同一组的所有indexes。

    fn so_example() -> Result<()> {
        let df = df! {
            "fruits" => &["Apple","Apple","Pear","Pear"],"maturity" => &["A","B","A","C","D"],"N" => &[1,2,8]

        }.unwrap();

        // Unwrap the fruits Series as Utf8Chunked
        let fruits = df.column("fruits")?.utf8()?;

        for (group_first,group_indexes) in df.groupby("fruits")?.get_groups() {
            let fruit_type = fruits.get(*group_first);
            let sub_df = df.take(group_indexes);

            dbg!(&fruit_type,&sub_df);
        };
        Ok(())
    }

这将输出：

[polars/src/frame/mod.rs:1672] &sub_df = shape: (4,3)
╭────────┬──────────┬─────╮
│ fruits ┆ maturity ┆ N   │
│ ---    ┆ ---      ┆ --- │
│ str    ┆ str      ┆ i32 │
╞════════╪══════════╪═════╡
│ Pear   ┆ A        ┆ 2   │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
│ Pear   ┆ C        ┆ 4   │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
│ Pear   ┆ A        ┆ 2   │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
│ Pear   ┆ D        ┆ 8   │
╰────────┴──────────┴─────╯
[polars/src/frame/mod.rs:1672] &fruit_type = Some(
    "Apple",)
[polars/src/frame/mod.rs:1672] &sub_df = shape: (2,3)
╭────────┬──────────┬─────╮
│ fruits ┆ maturity ┆ N   │
│ ---    ┆ ---      ┆ --- │
│ str    ┆ str      ┆ i32 │
╞════════╪══════════╪═════╡
│ Apple  ┆ A        ┆ 1   │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
│ Apple  ┆ B        ┆ 2   │
╰────────┴──────────┴─────╯