@Test
public void testCollectWithPartitionedColumns() throws Throwable {
Routing routing =
docSchemaInfo.getTableInfo(PARTITIONED_TABLE_NAME).getRouting(WhereClause.MATCH_ALL, null);
TableIdent tableIdent = new TableIdent(Schemas.DEFAULT_SCHEMA_NAME, PARTITIONED_TABLE_NAME);
CollectPhase collectNode =
getCollectNode(
Arrays.<Symbol>asList(
new Reference(
new ReferenceInfo(
new ReferenceIdent(tableIdent, "id"),
RowGranularity.DOC,
DataTypes.INTEGER)),
new Reference(
new ReferenceInfo(
new ReferenceIdent(tableIdent, "date"),
RowGranularity.SHARD,
DataTypes.TIMESTAMP))),
routing,
WhereClause.MATCH_ALL);
Bucket result = collect(collectNode);
for (Row row : result) {
System.out.println("Row:" + Arrays.toString(row.materialize()));
}
assertThat(result, containsInAnyOrder(isRow(1, 0L), isRow(2, 1L)));
}
protected boolean synchronizedSetNextRow(Row row) {
if (paused) {
pauseFifo.add(row.materialize());
return true;
} else {
Object[] bufferedCells;
while ((bufferedCells = pauseFifo.poll()) != null) {
sharedRow.cells(bufferedCells);
boolean wantMore = delegate.setNextRow(sharedRow);
if (!wantMore) {
return false;
}
if (paused) {
pauseFifo.add(row.materialize());
return true;
}
}
return delegate.setNextRow(row);
}
}
@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"))));
}
}