Problem with Module, Function, Association and Dataset

Question

I have struck a combination of the above-mentioned functions that causes evaluation to fail. I wonder whether someone can help me create a more minimal example in order to understand the issue better.

Assuming the definitions:

dataset1 = Dataset@<|"a" -> 1, "b" -> 2|>;
dataset2 = Dataset@<|"c" -> 3, "d" -> 4|>;
datasets = {dataset1, dataset2};
assoc1 = <|"aa" -> 8, "cc" -> 9|>;
test1[in1_] := in1
test2[in1_, in2_] := {in1, in2}

We can pass in that list of datasets to a Module that passes them, one by one, into another Module that calls a function on it:

Map[Function[localDataset, Module[{},
    Normal@Map[Function[localDatasetEntries, test1[localDatasetEntries]], localDataset]
  ]], datasets]
(* {<|"a" -> 1, "b" -> 2|>, <|"c" -> 3, "d" -> 4|>} *)

Upon switching the interior function to one that also uses the globally-defined association, it no longer works:

Map[Function[localDataset, Module[{},
   Normal@Map[Function[localDatasetEntries,test2[localDatasetEntries, assoc1]], localDataset]
   ]], datasets]
(* MapAt::partw: Part {All,2} of <|aa->8,cc->9|> does not exist. 
MapAt::partw: Part {All,2} of <|aa->8,cc->9|> does not exist.
{$Failed, $Failed} *)

although care should be taken because once the above statements have been passed once to the kernel, the evaluation does work. Assuming a new kernel, running the above statement twice gives:

Map[Function[localDataset, Module[{},
   Normal@Map[Function[localDatasetEntries, test2[localDatasetEntries, assoc1]], localDataset]
   ]], datasets]
Map[Function[localDataset, Module[{},
   Normal@Map[Function[localDatasetEntries, test2[localDatasetEntries, assoc1]], localDataset]
   ]], datasets]
(* MapAt::partw: Part {All,2} of <|aa->8,cc->9|> does not exist.
MapAt::partw: Part {All,2} of <|aa->8,cc->9|> does not exist.
{$Failed, $Failed}
{<|"a" -> {1, <|"aa" -> 8, "cc" -> 9|>}, "b" -> {2, <|"aa" -> 8, "cc" -> 9|>}|>, <|"c" -> {3, <|"aa" -> 8, "cc" -> 9|>}, "d" -> {4, <|"aa" -> 8, "cc" -> 9|>}|>} *)

Further, assuming a new kernel each time, certain ways to insert Normal[] into the expression solve the problem; others don't. This one works:

Map[Function[localDataset, Module[{},
    Normal@Map[Function[localDatasetEntries, test2[localDatasetEntries, assoc1]], Normal@localDataset]
   ]], datasets]
(* {{"a" -> {1, <|"aa" -> 8, "cc" -> 9|>}, "b" -> {2, <|"aa" -> 8, "cc" -> 9|>}}, {"c" -> {3, <|"aa" -> 8, "cc" -> 9|>}, "d" -> {4, <|"aa" -> 8, "cc" -> 9|>}}} *)

and this one does not:

Map[Function[localDataset, Module[{},
    Normal@Map[Function[localDatasetEntries,test2[Normal@localDatasetEntries, assoc1]], localDataset]
   ]], datasets]
(* MapAt::partw: Part {All,2} of <|aa->8,cc->9|> does not exist. 
MapAt::partw: Part {All,2} of <|aa->8,cc->9|> does not exist.
{$Failed, $Failed} *)

My particular issue is that I have partitioned a rather large dataset into smaller groups with some shared values, and I would like to operate one by one on the entries in each subgroup, using a list of globally-defined associations as parameters.

My question is, why is this happening? Why do some computations with structured datasets fail while their "normal" association forms work fine? Should we expect to generally encounter this type of issue in certain particular cases? Is there a general fix?

Answer 1

I am not quite sure why this is happening, but a quick solution is to wrap Association around assoc1. So this (using the OP's) example will work out without any error messages:

Map[ 
    Function[ localDataset, Module[ {}, 
       Normal @ Map[  
           Function[ localDatasetEntries, 
               test2[ localDatasetEntries, Association[assoc1] ] 
           ], 
           localDataset
       ]
      ]
    ], 
    datasets
]

{<|"a" -> {1, <|"aa" -> 8, "cc" -> 9|>}, "b" -> {2, <|"aa" -> 8, "cc" -> 9|>}|>, <|"c" -> {3, <|"aa" -> 8, "cc" -> 9|>}, "d" -> {4, <|"aa" -> 8, "cc" -> 9|>}|>}

But those Module, Map, Function and Normal concoctions look like the "most ugly code" I have seen in a long time :) So note, that something in this line will work (and may also be localized imo):

With[
    {assoc = assoc1},
    Dataset@datasets // Query[ All, All, List[#, Association[assoc]] & ]
]