dataframe. js - Можно ли суммировать несколько столбцов в сгруппированном кадре данных?

Question

Этот вопрос задается от c до dataframe.js.

Вот тестовые данные, которые я использую

let data = [
    {
        year : 2020,
        v : 0.1,
        cnt_1 : 1,
        cnt_2 : 20
    },
    {
        year : 2020,
        v : 0.1,
        cnt_1 : 3,
        cnt_2 : 20
    },
    {
        year : 2020,
        v : 0.1,
        cnt_1 : 5,
        cnt_2 : 4
    },
    {
        year : 2020,
        v : 0.1,
        cnt_1 : 7,
        cnt_2 : 20
    },
    {
        year : 2020,
        v : 0.2,
        cnt_1 : 9,
        cnt_2 : 20
    },
    {
        year : 2020,
        v : 0.2,
        cnt_1 : 11,
        cnt_2 : 20
    },
    {
        year : 2021,
        v : 0.2,
        cnt_1 : 13,
        cnt_2 : 20
    },
    {
        year : 2020,
        v : 0.1,
        cnt_1 : 15,
        cnt_2 : 20
    },
    {
        year : 2021,
        v : 0.1,
        cnt_1 : 17,
        cnt_2 : 20
    }
];

И ожидаемый мной результат выглядит так ...

| year      | v         | cnt_1_sum | cnt_2_sum |
    ------------------------------------
    | 2020      | 0.1       | 31        | 84        |
    | 2020      | 0.2       | 20        | 40        |
    | 2021      | 0.2       | 13        | 20        |
    | 2021      | 0.1       | 17        | 20        |

Я мог бы сделать это с одним столбцом, как показано ниже. Но я понятия не имел о нескольких столбцах. (В данном случае cnt_1 и cnt_2)

let df = new DataFrame(data);
let grouped = df.groupBy('year', 'v');
let cnt1_sum = grouped.aggregate(grpObj => grpObj.stat.sum('cnt_1')).rename('aggregation', 'cnt_1_sum');
cnt1_sum.show();
// and shows below
| year      | v         | cnt_1_sum |
------------------------------------
| 2020      | 0.1       | 31        |
| 2020      | 0.2       | 20        |
| 2021      | 0.2       | 13        |
| 2021      | 0.1       | 17        |

Единственный известный мне способ - это объединить 2 фрейма данных с year и v. Но это так ... неэффективно, когда есть несколько сгруппированных столбцов. (Если у меня 8 столбцов, то должен ли я присоединяться к 8 фреймам данных?)

Итак, вот вопрос. Так или иначе, чтобы

• применить функцию stat к нескольким столбцам?
• добавить столбец с данными? (withColumn API не работает с простым массивом)

Answer 1

Этого можно достичь, написав что-нибудь, аналогичное тому, как работает функция aggregate(). Вот код, с которого можно начать:

// assuming that sourceDf has columns ['key', 'col1', 'col2']
const groupedDf = sourceDf.groupBy('key');
const complexAggregateDf = new DataFrame([...groups].map(({groupKey, group}) => ({
  ...groupKey,
  'sum1': group.stat.sum('col1'),
  'sum2': group.stat.sum('col2'),
}), [...groupedDf.on, 'sum1', 'sum2']);