[go: nahoru, domu]

Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

aggregate multiple columns #127

Open
wants to merge 1 commit into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 18 additions & 6 deletions src/group.js
Original file line number Diff line number Diff line change
Expand Up @@ -136,19 +136,31 @@ class GroupedDataFrame {

/**
* Create an aggregation from a function.
* @param {Function} func The aggregation function.
* @param {Function | Object} func The aggregation function or an Object with keys representing
* the output column, and values as aggregation functions.
* @param {String} [columnName='aggregation'] The column name created by the aggregation.
* @returns {DataFrame} A new DataFrame with a column 'aggregation' containing the result.
* @example
* groupedDF.aggregate(group => group.stat.sum('column1'));
* groupedDF.aggregate({ 'column1': group => group.stat.sum('column1') })
*/
aggregate(func, columnName = "aggregation") {
// Convert into an object, where keys are the target column name, and
// values are the aggregation function.
const mappings = (typeof func == 'function') ? { [columnName]: func } : func;

return this.df.__newInstance__(
[...this].map(({ group, groupKey }) => ({
...groupKey,
[columnName]: func(group, groupKey)
})),
[...this.on, columnName]
[...this].map(({ group, groupKey }) => {
var transformed = { ...groupKey };

for (const column in mappings) {
const aggregationFunc = mappings[column];
transformed[column] = aggregationFunc(group, groupKey);
}

return transformed;
}),
[...this.on].concat(Object.keys(mappings))
);
}

Expand Down
28 changes: 28 additions & 0 deletions tests/dataframe.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -488,6 +488,18 @@ test("DataFrame rows can be ", (assert) => {
"groupBy and compute (by aggregation) the count by group."
);

assert.deepEqual(
df3
.groupBy("id")
.aggregate({ output: (group) => group.count() })
.toDict(),
{
id: [3, 6, 8, 1],
output: [2, 1, 2, 2]
},
"groupBy and compute (by aggregation) the count by group via object params."
);

assert.deepEqual(
df4
.groupBy("id", "id2")
Expand All @@ -501,6 +513,22 @@ test("DataFrame rows can be ", (assert) => {
"groupBy on multiple columns and compute the count by group."
);

assert.deepEqual(
df4b
.groupBy('test1')
.aggregate({
test2: group => group.stat.mean('test2'),
test3: group => group.stat.max('test3')
})
.toDict(),
{
test1: [95, 94],
test2: [91, 96],
test3: [76, 99]
},
"groupBy on multiple columns via Object and compute independent aggregations."
);

assert.deepEqual(
new DataFrame(
{
Expand Down