public static void sortAndDedup(LongArrayList list) {
list.elementsCount = sortAndDedup(list.buffer, list.elementsCount);
}
@Test
public void testFetchAction() throws Exception {
setUpCharacters();
Analysis analysis = analyze("select id, name from characters");
QueryThenFetchConsumer queryThenFetchConsumer =
internalCluster().getInstance(QueryThenFetchConsumer.class);
Planner.Context plannerContext = new Planner.Context(clusterService(), UUID.randomUUID());
ConsumerContext consumerContext = new ConsumerContext(analysis.rootRelation(), plannerContext);
QueryThenFetch plan =
(QueryThenFetch)
queryThenFetchConsumer.consume(analysis.rootRelation(), consumerContext).plan();
List<Bucket> results = getBuckets(plan.collectNode());
TransportFetchNodeAction transportFetchNodeAction =
internalCluster().getInstance(TransportFetchNodeAction.class);
// extract docIds by nodeId and jobSearchContextId
Map<String, LongArrayList> jobSearchContextDocIds = new HashMap<>();
for (Bucket rows : results) {
long docId = (long) rows.iterator().next().get(0);
// unpack jobSearchContextId and reader doc id from docId
int jobSearchContextId = (int) (docId >> 32);
String nodeId = plannerContext.nodeId(jobSearchContextId);
LongArrayList docIdsPerNode = jobSearchContextDocIds.get(nodeId);
if (docIdsPerNode == null) {
docIdsPerNode = new LongArrayList();
jobSearchContextDocIds.put(nodeId, docIdsPerNode);
}
docIdsPerNode.add(docId);
}
Iterable<Projection> projections =
Iterables.filter(
plan.mergeNode().projections(), Predicates.instanceOf(FetchProjection.class));
FetchProjection fetchProjection = (FetchProjection) Iterables.getOnlyElement(projections);
RowInputSymbolVisitor rowInputSymbolVisitor =
new RowInputSymbolVisitor(internalCluster().getInstance(Functions.class));
RowInputSymbolVisitor.Context context =
rowInputSymbolVisitor.extractImplementations(fetchProjection.outputSymbols());
final CountDownLatch latch = new CountDownLatch(jobSearchContextDocIds.size());
final List<Row> rows = new ArrayList<>();
for (Map.Entry<String, LongArrayList> nodeEntry : jobSearchContextDocIds.entrySet()) {
NodeFetchRequest nodeFetchRequest = new NodeFetchRequest();
nodeFetchRequest.jobId(plan.collectNode().jobId());
nodeFetchRequest.executionPhaseId(plan.collectNode().executionPhaseId());
nodeFetchRequest.toFetchReferences(context.references());
nodeFetchRequest.closeContext(true);
nodeFetchRequest.jobSearchContextDocIds(nodeEntry.getValue());
transportFetchNodeAction.execute(
nodeEntry.getKey(),
nodeFetchRequest,
new ActionListener<NodeFetchResponse>() {
@Override
public void onResponse(NodeFetchResponse nodeFetchResponse) {
for (Row row : nodeFetchResponse.rows()) {
rows.add(row);
}
latch.countDown();
}
@Override
public void onFailure(Throwable e) {
latch.countDown();
fail(e.getMessage());
}
});
}
latch.await();
assertThat(rows.size(), is(2));
for (Row row : rows) {
assertThat((Integer) row.get(0), anyOf(is(1), is(2)));
assertThat(
(BytesRef) row.get(1), anyOf(is(new BytesRef("Arthur")), is(new BytesRef("Ford"))));
}
}
public static void sort(LongArrayList list) {
sort(list.buffer, list.size());
}