ClusterService clusterService = cluster().getInstance(ClusterService.class);
Functions functions = cluster().getInstance(Functions.class);
AggregationFunction countAggregation =
(AggregationFunction)functions.get(new FunctionIdent(CountAggregation.NAME, ImmutableList.<DataType>of()));
TransportMergeNodeAction transportMergeNodeAction = cluster().getInstance(TransportMergeNodeAction.class);
Set<String> nodes = new HashSet<>();
for (DiscoveryNode discoveryNode : clusterService.state().nodes()) {
nodes.add(discoveryNode.getId());
}
// select count(*), user ... group by user
MergeNode mergeNode = new MergeNode("merge1", 2);
mergeNode.contextId(UUID.randomUUID());
mergeNode.executionNodes(nodes);
mergeNode.inputTypes(Arrays.<DataType>asList(DataTypes.UNDEFINED, DataTypes.STRING));
GroupProjection groupProjection = new GroupProjection();
groupProjection.keys(Arrays.<DataTypeSymbol>asList(new InputColumn(1, DataTypes.STRING)));
groupProjection.values(Arrays.asList(
new Aggregation(
countAggregation.info(),
ImmutableList.<Symbol>of(new InputColumn(0)),
Aggregation.Step.PARTIAL,
Aggregation.Step.FINAL
)
));
TopNProjection topNProjection = new TopNProjection(
10, 0, Arrays.<Symbol>asList(new InputColumn(1)), new boolean[] { false }, new Boolean[] { null });
topNProjection.outputs(Arrays.<Symbol>asList(new InputColumn(0), new InputColumn(1)));
mergeNode.projections(Arrays.asList(groupProjection, topNProjection));
mergeNode.outputTypes(Arrays.<DataType>asList(DataTypes.STRING, DataTypes.LONG));
Streamer<?>[] mapperOutputStreamer = new Streamer[] {
new AggregationStateStreamer(countAggregation),
DataTypes.STRING.streamer()
};
NoopListener noopListener = new NoopListener();
DistributedMergeTask task = new DistributedMergeTask(transportMergeNodeAction, mergeNode);
task.start();
Iterator<String> iterator = nodes.iterator();
String firstNode = iterator.next();
DistributedResultRequest request1 = new DistributedResultRequest(mergeNode.contextId(), mapperOutputStreamer);
request1.rows(new Object[][] {
new Object[] { new CountAggregation.CountAggState() {{ value = 1; }}, new BytesRef("bar") },
});
DistributedResultRequest request2 = new DistributedResultRequest(mergeNode.contextId(), mapperOutputStreamer);
request2.rows(new Object[][] {
new Object[] { new CountAggregation.CountAggState() {{ value = 1; }}, new BytesRef("bar") },
new Object[] { new CountAggregation.CountAggState() {{ value = 3; }}, new BytesRef("bar") },
new Object[] { new CountAggregation.CountAggState() {{ value = 3; }}, new BytesRef("foobar") },
});
transportMergeNodeAction.mergeRows(firstNode, request1, noopListener);
transportMergeNodeAction.mergeRows(firstNode, request2, noopListener);
DistributedResultRequest request3 = new DistributedResultRequest(mergeNode.contextId(), mapperOutputStreamer);
request3.rows(new Object[][] {
new Object[] { new CountAggregation.CountAggState() {{ value = 10; }}, new BytesRef("foo") },
new Object[] { new CountAggregation.CountAggState() {{ value = 20; }}, new BytesRef("foo") },
});
DistributedResultRequest request4 = new DistributedResultRequest(mergeNode.contextId(), mapperOutputStreamer);
request4.rows(new Object[][] {
new Object[] { new CountAggregation.CountAggState() {{ value = 10; }}, new BytesRef("foo") },
new Object[] { new CountAggregation.CountAggState() {{ value = 14; }}, new BytesRef("test") },
});
String secondNode = iterator.next();
transportMergeNodeAction.mergeRows(secondNode, request3, noopListener);
transportMergeNodeAction.mergeRows(secondNode, request4, noopListener);
assertThat(task.result().size(), is(2)); // 2 reducer nodes
// results from first node
Object[][] rows = task.result().get(0).get().rows();