Group by multiple keys, merge at specific keys in dataset

Question

I want to join associations in a dataset which share identical keys specified by the list sameAt and merge them by taking a Mean of keys specified by meanAt and taking First of the rest of keys. Below is a minimum example, not even doing what I really want (see later).

ds = {
   <|"a" -> 0, "b" -> "b1", "c" -> "c1", "x" -> 1, "y" -> 7|>,
   <|"a" -> 1, "b" -> "b1", "c" -> "c2", "x" -> 2, "y" -> 8|>,
   <|"a" -> 0, "b" -> "b2", "c" -> "c3", "x" -> 3, "y" -> 9|>,
   <|"a" -> 1, "b" -> "b2", "c" -> "c4", "x" -> 4, "y" -> 10|>,

   <|"a" -> 0, "b" -> "b1", "c" -> "c5", "x" -> 2, "y" -> 10|>,
   <|"a" -> 1, "b" -> "b1", "c" -> "c6", "x" -> 3, "y" -> 11|>,
   <|"a" -> 0, "b" -> "b2", "c" -> "c7", "x" -> 4, "y" -> 12|>,
   <|"a" -> 1, "b" -> "b2", "c" -> "c8", "x" -> 5, "y" -> 13|>
   };
sameAt = {"a", "b"}; (* associations with identical values at these keys should be merged *)
meanAt = {"x", "y"}; (* merging should be Mean at these keys, First at rest *)

Query[N @* Mean] /@ GatherBy[ds, (#a && #b) &]

{<|"a" -> 0., "b" -> "b1", "c" -> 0.5 ("c1" + "c5"), "x" -> 1.5, "y" -> 8.5|>,
 <|"a" -> 1., "b" -> "b1", "c" -> 0.5 ("c2" + "c6"), "x" -> 2.5, "y" -> 9.5|>,
 <|"a" -> 0., "b" -> "b2", "c" -> 0.5 ("c3" + "c7"), "x" -> 3.5, "y" -> 10.5|>,
 <|"a" -> 1., "b" -> "b2", "c" -> 0.5 ("c4" + "c8"), "x" -> 4.5, "y" -> 11.5|>}

As can be seen, my method is not general enough: had to manually specify grouping criterion #a && #b &; and averaging is not restricted for the keys "x" and "y". In reality, I would like to make this work with:

• any number of keys in sameAt and meanAt (but see last point);
• any number of mergable "copies", not just two (so JoinAcross is no good);
• possibly specifying a different merging operator for each key, not just First and Mean;

---

A general approach, where each key has a specific aggregator defined:

sameAt = {"a", "b"};
op = <|"a" -> First, "b" -> First, "c" -> Last, "x" -> N@*Mean, "y" -> Mean|>; 
MapThread[Apply, {op, Merge[#, List]}] & /@ GatherBy[ds, Query[sameAt]]

{<|"a" -> 0, "b" -> "b1", "c" -> "c5", "x" -> 1.5, "y" -> 17/2|>,
 <|"a" -> 1, "b" -> "b1", "c" -> "c6", "x" -> 2.5, "y" -> 19/2|>,
 <|"a" -> 0, "b" -> "b2", "c" -> "c7", "x" -> 3.5, "y" -> 21/2|>,
 <|"a" -> 1, "b" -> "b2", "c" -> "c8", "x" -> 4.5, "y" -> 23/2|>}

---

QUESTION

Can this be solved elegantly (being simpler than my above example) within the Dataset-framework, with a single Query? (If not, what is missing from the SQL-like Wolfram language that could make this work easily? Perhaps we can expect new functionality to be added, as it is hinted e.g. here.)

Answer 1

Dataset[ds][GroupBy[Query[sameAt]], 
  Merge[Identity] /* 
   Query[Thread[meanAt -> Mean]~ Join~ 
     Thread[{"a", "b", "c"} -> First]]][Values]

{<|a->0,b->b1,c->c1,x->3/2,y->17/2|>, <|a->1,b->b1,c->c2,x->5/2,y->19/2|>, <|a->0,b->b2,c->c3,x->7/2,y->21/2|>, <|a->1,b->b2,c->c4,x->9/2,y->23/2|>} . (* Normal'ed *)

Obvious limitation is having to specify keys complementary to meanAt. A 'KeyComplement' query would be great here.

EDIT

Here's an alternative using KeyDrop and post-Join:

ds[GroupBy[Query[sameAt]], 
  Merge[Identity] /* 
   Query[{meanAt /* Map[Mean], KeyDrop[meanAt] /* Map[First]} /* 
     Apply[Join]]][Values]

Both of these can be wrapped in a single Query op (here the grouping keys are dropped with Values as a separate op for clarity).

Merge will work with any number of Association arguments (I thought so did JoinAcross?)

I've already suggested to WRI to streamline the Thread[meanAt -> Mean] syntax to just meanAt -> Mean as composite keys can always be specified as Key[{"d","e"}]-> f (if such a key existed) to distinguish from mapping at multiple keys.

Answer 2

Here is one way:

ds // Query[
  GroupBy[#[[sameAt]]&] /* Values
, Transpose /* MapIndexed[#2[[1, 1]] /. {Alternatives@@meanAt -> N@Mean@#, _ -> First@#}&]
]

(*
  {<|"a" -> 0, "b" -> "b1", "c" -> "c1", "x" -> 1.5, "y" -> 8.5|>, 
   <|"a" -> 1, "b" -> "b1", "c" -> "c2", "x" -> 2.5, "y" -> 9.5|>, 
   <|"a" -> 0, "b" -> "b2", "c" -> "c3", "x" -> 3.5, "y" -> 10.5|>, 
   <|"a" -> 1, "b" -> "b2", "c" -> "c4", "x" -> 4.5, "y" -> 11.5|>}
*)

Or, rephrased to use GroupBy directly:

GroupBy[ds
, #[[sameAt]]&
, Query[Transpose] /*
  MapIndexed[#2[[1, 1]] /. {Alternatives@@meanAt -> N@Mean@#, _ -> First@#}&]
] // Values

Answer 3

GroupBy[ 
    ds, Lookup[sameAt], Merge[Identity]
] //  Query[Values, Thread[meanAt -> (List@*Total)]
] //  Map[First, #, {2}] &

Answer 4

Query[All, Join[Thread[Complement[Keys[ds][[1]], meanAt]->First], Thread[meanAt->Mean]]]@ 
 (Merge[Identity] /@ GatherBy[ds, (And@@(Function[{x},#[x]] /@ sameAt))&]) 

{<|"a" -> 0, "b" -> "b1", "c" -> "c1", "x" -> 3/2, "y" -> 17/2|>,
<|"a" -> 1, "b" -> "b1", "c" -> "c2", "x" -> 5/2, "y" -> 19/2|>,
<|"a" -> 0, "b" -> "b2", "c" -> "c3", "x" -> 7/2, "y" -> 21/2|>,
<|"a" -> 1, "b" -> "b2", "c" -> "c4", "x" -> 9/2, "y" -> 23/2|>}

Answer 5

Just for the record, I ended up using the following packaged functions:

• associationMapAt[$a$, {$key_1$ → $f_1$, $key_2$ → $f_2$, ...}] applies $f_i$ to the value of $key_i$ in association $a$. The values of keys not listed in the second argument are wrapped in Identity.

• associationMapAt[$a$, {$key_1$ → $f_1$, $key_2$ → $f_2$, ...}, $def$] applies the default function $def$ to values of keys not listed in the second argument.

• mergeByMapAt[{$a_1$, $a_2$, ...}, $by$, $at$] gathers associations $a_i$ that share identical values of the keys listed in $by$ and merges each group to a single association by applying Join to matching keys, then it applies Mean to keys listed in $at$ and First to keys not listed in $at$.

• mergeByMapAt[$ds$, $by$ → $f$, $at$ → $g$, $def$] flattens nested lists in $ds$ then applies $f$ instead of Join when merging, $g$ instead of Mean to keys listed in $at$ and $def$ instead of First to keys not listed in $at$.

Function definitions:

ClearAll[associationMapAt, mergeByMapAt];

associationMapAt[a_Association, f:_Rule|_RuleDelayed, def_:Identity] := 
   associationMapAt[a, {f}, def];
associationMapAt[a_Association, f:{(_Rule|_RuleDelayed)...}|_Association,
   def_: Identity] := AssociationMap[((# /. Append[Normal@f, _ -> def])@a@#) &, Keys@a];

mergeByMapAt[ds_List, by_] := 
  mergeByMapAt[ds, by -> Join, Union@Catenate[Keys /@ ds]];
mergeByMapAt[ds_List, by_, at_, def_: First] := 
  mergeByMapAt[ds, by -> Join, at -> Mean, def];
mergeByMapAt[ds_List, by_ -> f_, at_, def_: First] := 
  mergeByMapAt[ds, by -> f, at -> Mean, def];
mergeByMapAt[ds_List, by_, at_ -> g_, def_: First] := 
  mergeByMapAt[ds, by -> Join, at -> g, def];
mergeByMapAt[ds_List, by_ -> f_, at_ -> g_, def_: First] := Module[
   {gathered = GatherBy[Flatten@ds, Query@by]},
   associationMapAt[Merge[#, f], Thread[at -> g], def] & /@ gathered];

Original example:

ds = {
 <|"a" -> 0, "b" -> "b1", "c" -> "c1", "x" -> 1, "y" -> 7|>,
 <|"a" -> 1, "b" -> "b1", "c" -> "c2", "x" -> 2, "y" -> 8|>,
 <|"a" -> 0, "b" -> "b2", "c" -> "c3", "x" -> 3, "y" -> 9|>,
 <|"a" -> 1, "b" -> "b2", "c" -> "c4", "x" -> 4, "y" -> 10|>,
 <|"a" -> 0, "b" -> "b1", "c" -> "c5", "x" -> 2, "y" -> 10|>,
 <|"a" -> 1, "b" -> "b1", "c" -> "c6", "x" -> 3, "y" -> 11|>,
 <|"a" -> 0, "b" -> "b2", "c" -> "c7", "x" -> 4, "y" -> 12|>,
 <|"a" -> 1, "b" -> "b2", "c" -> "c8", "x" -> 5, "y" -> 13|>};

sameBy = {"a", "b"}; (* merge associations sharing identical values at these keys *)
meanAt = {"x", "y"}; (* apply Mean to these keys at the end *)

mergeByMapAt[ds, sameBy, meanAt] // N

   {
    <|"a" -> 0., "b" -> "b1", "c" -> "c1", "x" -> 1.5, "y" -> 8.5|>,
    <|"a" -> 1., "b" -> "b1", "c" -> "c2", "x" -> 2.5, "y" -> 9.5|>,
    <|"a" -> 0., "b" -> "b2", "c" -> "c3", "x" -> 3.5, "y" -> 10.5|>,
    <|"a" -> 1., "b" -> "b2", "c" -> "c4", "x" -> 4.5, "y" -> 11.5|>
   }