首页 > 解决方案 > 使用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,其中每个组都可以作为单个数据框与其标签(分组列中的公共值)一起访问。

我查看了数据融合,但我无法重现此行为,而无需首先选择所有唯一值,然后执行一个选择 par 值,这会导致多次重新解析整个文件。我看了看很有希望但也无法达到我的目标的 Polars 板条箱。

你将如何以与 python 代码相似/更好的性能来做到这一点?我对任何允许我通过固定列的值有效地对镶木地板文件进行分区的语法/库/方法持开放态度。


这是一个 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.

标签: rustapache-arrow

解决方案


您可以执行以下操作:

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

推荐阅读