rust - 使用Arrow/Datafusion/Polars(如python panda的groupby)按列值分区？

Question

我正在寻找方便的 python panda 语法的等价物:

#df is a pandas dataframe

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

它基本上是一个 groupby，其中每个组都可以作为单个数据框与其标签(分组列中的公共(public)值)一起访问。
我查看了数据融合，但我无法重现此行为，而无需首先选择所有唯一值，然后执行一个选择 par 值，这会导致多次重新解析整个文件。
我看了看很有希望但也无法达到我的目标的 Polars crate 。
你将如何以与 python 代码相似/更好的性能来做到这一点？我对任何允许我通过固定列的值有效地对 Parquet 文件进行分区的语法/库/方法持开放态度。

这是一个 rust 示例代码，使用 polar 作为我正在处理的输入类型的示例:

    let s0 = Series::new("fruits", ["Apple", "Apple", "Pear", "Pear", "Pear", "Pear"].as_ref());
    let s1 = Series::new("maturity", ["A", "B", "A", "C", "A", "D"].as_ref());
    let s1 = Series::new("N", [1, 2, 2, 4, 2, 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.

Answer 1

您可以执行以下操作:

    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]" |
+---------+--------------------------------------+-----------------------------+-----------------+