Extracting values with multiple keys

Question

I want to extract sorted data from a huge database based upon two (or more) keys in a very timely manner. Here is a reproducible toy example for two keys only:

n = 10^5;
keys = RandomInteger[{1, 100}, n];
vals = RandomReal[{0, 1}, n];

data = Transpose[{keys, vals}];

The fastest "traditional" way I' ve found:

result1 = Sort @ Cases[data, {25 | 73, r_} :> r];

Much much faster is a V10 solution:

assoc = Merge[Association /@ Rule @@@ data, Identity];

(I use Mergeto allow for duplicate keys, and the time cost of getting assoc is not important to me).

result2 = Sort[assoc[25] ~ Join ~ assoc[73]];

result1 == result2

True

Speed comparison:

Do[Sort @ Cases[data, {25 | 73, r_} :> r], {100}]; // AbsoluteTiming // First

2.017115

Do[Sort[assoc[25] ~ Join ~ assoc[73]], {100}]; // AbsoluteTiming // First

0.030002

Certainly one reason to upgrade, but two or more questions remain:

(a) Could this code be improved ?

(b) And, passing to n = 10^6, result1 still works, but result2 runs forever and has to be aborted.

Answer 1

In response to a), you can write Merge[Rule @@@ data, Identity], which is slightly simpler.

In response to b), there are two different ways to do this nicely. One of which works in 10.0.0, the other will have to wait for 10.0.1.

In 10.0.0 we can use GroupBy to associate each unique key with the set of corresponding values:

grouped = GroupBy[data, First -> Last];

Now do lookups by writing, e.g.:

grouped[25]

or

Catenate @ Lookup[grouped, {25, 73}]

You can also use PositionIndex. Unfortunately 10.0.0 has a slow implementation of PositionIndex as described here. But this is fixed in 10.0.1, takes a fraction of a second to complete on your example for n=6. It's actually faster than the GroupBy code above.

keyindex = PositionIndex[keys];

Now we can easily look up the positions for which the key was 25, and from that get the corresponding values:

Part[vals, keyindex[25]]

To look up both 25 and 73 we just do:

Part[vals, Catenate @ Lookup[keyindex, {25, 73}]]

Answer 2

You can use the ResourceFunction GroupByList to do this. Your example:

n = 10^5;
keys = RandomInteger[{1, 100}, n];
vals = RandomReal[{0, 1}, n];

data = Transpose[{keys, vals}];

Comparison:

r1 = GroupBy[data, First -> Last]; //AbsoluteTiming
r2 = Merge[Rule @@@ data, Identity]; //AbsoluteTiming
r3 = ResourceFunction["GroupByList"][vals, keys]; //AbsoluteTiming

r1 === r2 === r3

{0.053079, Null}

{0.395158, Null}

{0.002779, Null}

True

Also, GroupByList is faster than PositionIndex:

PositionIndex[keys]; //AbsoluteTiming

{0.010389, Null}

Answer 3

There were (are) many methods faster than your Cases code for earlier versions.

Two of the best are simple DownValues assignments, e.g. fn[key] = {val1, val2, ...} and Rules optimized with Dispatch. I believe the latter is usually slightly faster so I will illustrate that one.

dsp = Dispatch[ #[[1, 1]] -> #[[All, 2]] & /@ GatherBy[data, First] ];

Then:

Join[25, 73] /. dsp

By the way, for your particular data it would be better to store the list of values as a packed array by using this in the line above:

#[[1, 1]] -> Developer`ToPackedArray @ #[[All, 2]] &