String query = "foreach (cogroup (load 'a') by $1, (load 'b') by $1) generate org.apache.pig.builtin.AVG($1) ;";
LogicalPlan lp = planTester.buildPlan(query);
//check that the loads' projection map is null
LOLoad loada = (LOLoad) lp.getRoots().get(0);
ProjectionMap loadaProjectionMap = loada.getProjectionMap();
assertTrue(loadaProjectionMap.changes() == false);
LOLoad loadb = (LOLoad) lp.getRoots().get(1);
ProjectionMap loadbProjectionMap = loadb.getProjectionMap();
assertTrue(loadbProjectionMap.changes() == false);
//check cogroup projection map
LOCogroup cogroup = (LOCogroup)lp.getSuccessors(loada).get(0);
ProjectionMap cogroupProjectionMap = cogroup.getProjectionMap();
assertTrue(cogroupProjectionMap.changes() == true);
MultiMap<Integer, ProjectionMap.Column> cogroupMapFields = cogroupProjectionMap.getMappedFields();
assertTrue(cogroupMapFields != null);
List<ProjectionMap.Column> mapValues = (ArrayList<ProjectionMap.Column>)cogroupMapFields.get(0);
assertTrue(mapValues.get(0).getInputColumn().first == 0);
assertTrue(mapValues.get(0).getInputColumn().second == 1);
assertTrue(mapValues.get(1).getInputColumn().first == 1);
assertTrue(mapValues.get(1).getInputColumn().second == 1);
//check the cogroup removed fields is null
assertTrue(cogroupProjectionMap.getRemovedFields() == null);
//check that cogroup added fields contain [1, 2]
List<Integer> cogroupAddedFields = cogroupProjectionMap.getAddedFields();
assertTrue(cogroupAddedFields.size() == 2);
assertTrue(cogroupAddedFields.get(0) == 1);
assertTrue(cogroupAddedFields.get(1) == 2);
//check that the foreach projection map has null mappedFields
LOForEach foreach = (LOForEach)lp.getLeaves().get(0);
ProjectionMap foreachProjectionMap = foreach.getProjectionMap();
assertTrue(foreachProjectionMap.changes() == true);
assertTrue(foreachProjectionMap.getMappedFields() == null);
//check that removed fields has all the columns from the input cogroup
List<Pair<Integer, Integer>> foreachRemovedFields = foreachProjectionMap.getRemovedFields();
assertTrue(foreachProjectionMap.getRemovedFields().size() == 3);
int expectedColumn = 0;
for(Pair<Integer, Integer> removedField: foreachRemovedFields) {
assertTrue(removedField.first == 0);
assertTrue(removedField.second == expectedColumn++);
}
//check that added fields contain [0]
List<Integer> foreachAddedFields = foreachProjectionMap.getAddedFields();
assertTrue(foreachAddedFields.size() == 1);
assertTrue(foreachAddedFields.get(0) == 0);
}