Разделение по значению столбца в rust с помощью Arrow / Datafusion / Polars (например, groupby в python panda)?

#rust #apache-arrow

#Ржавчина #apache-arrow

Вопрос:

Я ищу эквивалент удобного синтаксиса 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? Я открыт для любого синтаксиса / библиотеки / подхода, который позволил бы мне эффективно разбивать файл 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.
  

Ответ №1:

Если вы активируете эту partition_by функцию, polars открывается

DataFrame::partition_by и DataFrame::partition_by_stable .

 use polars::prelude::*;

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

    }?.partition_by(["fruits"])?;

    dbg!(partitioned);
    Ok(())
}
  

Запуск этого печатает:

 [src/main.rs:17] partitioned = [
    shape: (4, 3)
    ┌────────┬──────────┬─────┐
    │ fruits ┆ maturity ┆ N   │
    │ ---    ┆ ---      ┆ --- │
    │ strstri32 │
    ╞════════╪══════════╪═════╡
    │ Pear   ┆ A        ┆ 2   │
    ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
    │ Pear   ┆ C        ┆ 4   │
    ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
    │ Pear   ┆ A        ┆ 2   │
    ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
    │ Pear   ┆ D        ┆ 8   │
    └────────┴──────────┴─────┘,
    shape: (2, 3)
    ┌────────┬──────────┬─────┐
    │ fruits ┆ maturity ┆ N   │
    │ ---    ┆ ---      ┆ --- │
    │ strstri32 │
    ╞════════╪══════════╪═════╡
    │ Apple  ┆ A        ┆ 1   │
    ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
    │ Apple  ┆ B        ┆ 2   │
    └────────┴──────────┴─────┘,
]

  

Комментарии:

1. Привет, @ritchie46, поведение get_groups изменилось: df.groupby(["groups"])?.get_groups().iter().next() возвращает GroupIndicator . Как перейти от GroupIndicator к фактическому sub_df? Кстати, я рад обновить ваш ответ (если смогу)

2. В настоящее partition_by время существует метод. Это должно быть намного проще и быстрее.

3. Большое спасибо за быстрый ответ! Я не смог найти упоминания о «partition_by» в документации?

4. Я соответствующим образом обновил свой ответ.

Ответ №2:

Вы можете сделать что-то вроде:

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

И это технически оставит вас с фреймом данных. Тогда проблема в том, что столбцы dataframe представляют собой массивы со всеми значениями для каждого фрукта, что может быть не тем, что вы хотите.

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

Комментарии:

1. Спасибо за ваш ответ. Я запустил ваш код, но не могу понять: 1) почему вы группируете по каждому столбцу? 2) как я предполагаю извлечь один фрейм данных для каждого значения из столбца группировки. Я обновил вопрос образцом rust, чтобы сделать его немного понятнее.

2. Ах, я неправильно понял. Я отредактировал свой ответ, но я не знаю, помогает ли это.