Parallelize[MapIndexed[...,Association[...]]] returns broken result

Question

Bug introduced in 10.0.2 or earlier and persisting through 12.0

---

Confirmed: CASE:4041910

Consider the following example:

MapIndexed[f[##, $KernelID]&, <|a -> 1, b -> 2, c -> 3|>]
(* <|a -> f[1, {Key[a]}, 0], b -> f[2, {Key[b]}, 0], c -> f[3, {Key[c]}, 0]|> *)

Now lets try to Parallelize this:

Parallelize@MapIndexed[f[##, $KernelID]&, <|a -> 1, b -> 2, c -> 3|>]
(* ---- Output in Mathematica 11.2 - 11.3 ---- *)
(* Association[f[a -> 1, {1}, 2], f[b -> 2, {2}, 2], f[c -> 3, {3}, 1]] *)

(* ---- Output in Mathematica 10 - 11.1.1 ---- *)
(* Transpose::nmtx: The first two levels of {<|a->1,b->2,c->3|>,{{1},{2},{3}}} cannot be transposed. >> *)
(* Transpose[f[<|a -> 1, b -> 2, c -> 3|>, {{1}, {2}, {3}}, 4]] *)

Clearly, this is completely broken - it doesn't simply switch to sequential evaluation, the output is actually useless.

How can one work around this issue until it gets fixed?

Answer 1

Workaround

Wolfram support suggested the following workaround until the issue is resolved:

Parallelize@Map[First, MapIndexed[Hold[f[##, $KernelID]]&, <|a -> 1, b -> 2, c -> 3|>]]
(* <|a -> f[1, {Key[a]}, 2], b -> f[2, {Key[b]}, 2], c -> f[3, {Key[c]}, 1]|> *)

which is working. This is because the expressions to be evaluated are built sequentially (but wrapped in Hold). The evaluation (by means of First/ReleaseHold) is then done using ParallelMap[…]/Parallelize[Map[…]]

Fix

While the above is a nice workaround, it's still something that has to be remembered. The following code fixes the issue globally.

Mathematica 11 & 12

Begin["Parallel`Evaluate`Private`"];
Unprotect@MapIndexed;
HoldPattern[tryCombine[MapIndexed[f_, str_Association], opts___]] ^:= 
 ParallelCombine[MapIndexed[f], str, opts]
Protect@MapIndexed;
End[];

To verify:

Parallelize@MapIndexed[f[##, $KernelID] &, <|a -> 1, b -> 2, c -> 3|>]
(* <|a -> f[1, {Key[a]}, 4], b -> f[2, {Key[b]}, 3], c -> f[3, {Key[c]}, 2]|> *)

This works by adding a special case for the problematic case (Parallelize[MapIndexed[…,<|…|>]]) to the code that handles the parallelization internally. Since ParallelCombine supports associations and associations can be Joined, the actual implementation is trivial (since we don't have to add the position specification manually).

Mathematica 10 (works also for 11)

Begin["Parallel`Evaluate`Private`"];
Unprotect@MapIndexed;
tryCombine[MapIndexed[f_, str_Association], opts___] ^:= 
 ParallelCombine[MapIndexed[f, <|#|>] &, Normal@str, opts]
Protect@MapIndexed;
End[];

The only difference to the above fix is that ParallelCombine doesn't seem to support associations in 10, so we have to convert it to a list first. Of course this will also work for Mathematica 11+, but it feels like an unnecessary step to covert to a list first.