use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.
the class DfsQueryPhase method run.
@Override
public void run() throws IOException {
// TODO we can potentially also consume the actual per shard results from the initial phase here in the aggregateDfs
// to free up memory early
final AggregatedDfs dfs = searchPhaseController.aggregateDfs(dfsSearchResults);
final CountedCollector<QuerySearchResultProvider> counter = new CountedCollector<>(queryResult::consumeResult, dfsSearchResults.asList().size(), () -> {
context.executeNextPhase(this, nextPhaseFactory.apply(queryResult));
}, context);
for (final AtomicArray.Entry<DfsSearchResult> entry : dfsSearchResults.asList()) {
DfsSearchResult dfsResult = entry.value;
final int shardIndex = entry.index;
final SearchShardTarget searchShardTarget = dfsResult.shardTarget();
Transport.Connection connection = context.getConnection(searchShardTarget.getNodeId());
QuerySearchRequest querySearchRequest = new QuerySearchRequest(context.getRequest(), dfsResult.id(), dfs);
searchTransportService.sendExecuteQuery(connection, querySearchRequest, context.getTask(), ActionListener.wrap(result -> counter.onResult(shardIndex, result, searchShardTarget), exception -> {
try {
if (context.getLogger().isDebugEnabled()) {
context.getLogger().debug((Supplier<?>) () -> new ParameterizedMessage("[{}] Failed to execute query phase", querySearchRequest.id()), exception);
}
counter.onFailure(shardIndex, searchShardTarget, exception);
} finally {
context.sendReleaseSearchContext(querySearchRequest.id(), connection);
}
}));
}
}
use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.
the class SearchPhaseController method aggregateDfs.
public AggregatedDfs aggregateDfs(AtomicArray<DfsSearchResult> results) {
ObjectObjectHashMap<Term, TermStatistics> termStatistics = HppcMaps.newNoNullKeysMap();
ObjectObjectHashMap<String, CollectionStatistics> fieldStatistics = HppcMaps.newNoNullKeysMap();
long aggMaxDoc = 0;
for (AtomicArray.Entry<DfsSearchResult> lEntry : results.asList()) {
final Term[] terms = lEntry.value.terms();
final TermStatistics[] stats = lEntry.value.termStatistics();
assert terms.length == stats.length;
for (int i = 0; i < terms.length; i++) {
assert terms[i] != null;
TermStatistics existing = termStatistics.get(terms[i]);
if (existing != null) {
assert terms[i].bytes().equals(existing.term());
// totalTermFrequency is an optional statistic we need to check if either one or both
// are set to -1 which means not present and then set it globally to -1
termStatistics.put(terms[i], new TermStatistics(existing.term(), existing.docFreq() + stats[i].docFreq(), optionalSum(existing.totalTermFreq(), stats[i].totalTermFreq())));
} else {
termStatistics.put(terms[i], stats[i]);
}
}
assert !lEntry.value.fieldStatistics().containsKey(null);
final Object[] keys = lEntry.value.fieldStatistics().keys;
final Object[] values = lEntry.value.fieldStatistics().values;
for (int i = 0; i < keys.length; i++) {
if (keys[i] != null) {
String key = (String) keys[i];
CollectionStatistics value = (CollectionStatistics) values[i];
assert key != null;
CollectionStatistics existing = fieldStatistics.get(key);
if (existing != null) {
CollectionStatistics merged = new CollectionStatistics(key, existing.maxDoc() + value.maxDoc(), optionalSum(existing.docCount(), value.docCount()), optionalSum(existing.sumTotalTermFreq(), value.sumTotalTermFreq()), optionalSum(existing.sumDocFreq(), value.sumDocFreq()));
fieldStatistics.put(key, merged);
} else {
fieldStatistics.put(key, value);
}
}
}
aggMaxDoc += lEntry.value.maxDoc();
}
return new AggregatedDfs(termStatistics, fieldStatistics, aggMaxDoc);
}
use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.
the class DfsQueryPhaseTests method testFailPhaseOnException.
public void testFailPhaseOnException() throws IOException {
AtomicArray<DfsSearchResult> results = new AtomicArray<>(2);
AtomicReference<AtomicArray<QuerySearchResultProvider>> responseRef = new AtomicReference<>();
results.set(0, new DfsSearchResult(1, new SearchShardTarget("node1", new Index("test", "na"), 0)));
results.set(1, new DfsSearchResult(2, new SearchShardTarget("node2", new Index("test", "na"), 0)));
results.get(0).termsStatistics(new Term[0], new TermStatistics[0]);
results.get(1).termsStatistics(new Term[0], new TermStatistics[0]);
SearchPhaseController controller = new SearchPhaseController(Settings.EMPTY, BigArrays.NON_RECYCLING_INSTANCE, null);
SearchTransportService searchTransportService = new SearchTransportService(Settings.builder().put("search.remote.connect", false).build(), null, null) {
@Override
public void sendExecuteQuery(Transport.Connection connection, QuerySearchRequest request, SearchTask task, ActionListener<QuerySearchResult> listener) {
if (request.id() == 1) {
QuerySearchResult queryResult = new QuerySearchResult(123, new SearchShardTarget("node1", new Index("test", "na"), 0));
queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(42, 1.0F) }, 2.0F), new DocValueFormat[0]);
// the size of the result set
queryResult.size(2);
listener.onResponse(queryResult);
} else if (request.id() == 2) {
throw new UncheckedIOException(new MockDirectoryWrapper.FakeIOException());
} else {
fail("no such request ID: " + request.id());
}
}
};
MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
mockSearchPhaseContext.searchTransport = searchTransportService;
DfsQueryPhase phase = new DfsQueryPhase(results, controller, (response) -> new SearchPhase("test") {
@Override
public void run() throws IOException {
responseRef.set(response.results);
}
}, mockSearchPhaseContext);
assertEquals("dfs_query", phase.getName());
expectThrows(UncheckedIOException.class, () -> phase.run());
// phase execution will clean up on the contexts
assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
}
use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.
the class CountedCollectorTests method testCollect.
public void testCollect() throws InterruptedException {
AtomicArray<SearchPhaseResult> results = new AtomicArray<>(randomIntBetween(1, 100));
List<Integer> state = new ArrayList<>();
int numResultsExpected = randomIntBetween(1, results.length());
MockSearchPhaseContext context = new MockSearchPhaseContext(results.length());
CountDownLatch latch = new CountDownLatch(1);
boolean maybeFork = randomBoolean();
Executor executor = (runnable) -> {
if (randomBoolean() && maybeFork) {
new Thread(runnable).start();
} else {
runnable.run();
}
};
CountedCollector<SearchPhaseResult> collector = new CountedCollector<>(results::set, numResultsExpected, latch::countDown, context);
for (int i = 0; i < numResultsExpected; i++) {
int shardID = i;
switch(randomIntBetween(0, 2)) {
case 0:
state.add(0);
executor.execute(() -> collector.countDown());
break;
case 1:
state.add(1);
executor.execute(() -> collector.onResult(shardID, new DfsSearchResult(shardID, null), new SearchShardTarget("foo", new Index("bar", "baz"), shardID)));
break;
case 2:
state.add(2);
executor.execute(() -> collector.onFailure(shardID, new SearchShardTarget("foo", new Index("bar", "baz"), shardID), new RuntimeException("boom")));
break;
default:
fail("unknown state");
}
}
latch.await();
assertEquals(numResultsExpected, state.size());
for (int i = 0; i < numResultsExpected; i++) {
switch(state.get(i)) {
case 0:
assertNull(results.get(i));
break;
case 1:
assertNotNull(results.get(i));
assertEquals(i, results.get(i).id());
break;
case 2:
final int shardId = i;
assertEquals(1, context.failures.stream().filter(f -> f.shardId() == shardId).count());
break;
default:
fail("unknown state");
}
}
for (int i = numResultsExpected; i < results.length(); i++) {
assertNull("index: " + i, results.get(i));
}
}
use of org.elasticsearch.search.dfs.DfsSearchResult in project elasticsearch by elastic.
the class DfsQueryPhaseTests method testDfsWith1ShardFailed.
public void testDfsWith1ShardFailed() throws IOException {
AtomicArray<DfsSearchResult> results = new AtomicArray<>(2);
AtomicReference<AtomicArray<QuerySearchResultProvider>> responseRef = new AtomicReference<>();
results.set(0, new DfsSearchResult(1, new SearchShardTarget("node1", new Index("test", "na"), 0)));
results.set(1, new DfsSearchResult(2, new SearchShardTarget("node2", new Index("test", "na"), 0)));
results.get(0).termsStatistics(new Term[0], new TermStatistics[0]);
results.get(1).termsStatistics(new Term[0], new TermStatistics[0]);
SearchPhaseController controller = new SearchPhaseController(Settings.EMPTY, BigArrays.NON_RECYCLING_INSTANCE, null);
SearchTransportService searchTransportService = new SearchTransportService(Settings.builder().put("search.remote.connect", false).build(), null, null) {
@Override
public void sendExecuteQuery(Transport.Connection connection, QuerySearchRequest request, SearchTask task, ActionListener<QuerySearchResult> listener) {
if (request.id() == 1) {
QuerySearchResult queryResult = new QuerySearchResult(123, new SearchShardTarget("node1", new Index("test", "na"), 0));
queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(42, 1.0F) }, 2.0F), new DocValueFormat[0]);
// the size of the result set
queryResult.size(2);
listener.onResponse(queryResult);
} else if (request.id() == 2) {
listener.onFailure(new MockDirectoryWrapper.FakeIOException());
} else {
fail("no such request ID: " + request.id());
}
}
};
MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
mockSearchPhaseContext.searchTransport = searchTransportService;
DfsQueryPhase phase = new DfsQueryPhase(results, controller, (response) -> new SearchPhase("test") {
@Override
public void run() throws IOException {
responseRef.set(response.results);
}
}, mockSearchPhaseContext);
assertEquals("dfs_query", phase.getName());
phase.run();
mockSearchPhaseContext.assertNoFailure();
assertNotNull(responseRef.get());
assertNotNull(responseRef.get().get(0));
assertNull(responseRef.get().get(0).fetchResult());
assertEquals(1, responseRef.get().get(0).queryResult().topDocs().totalHits);
assertEquals(42, responseRef.get().get(0).queryResult().topDocs().scoreDocs[0].doc);
assertNull(responseRef.get().get(1));
assertEquals(1, mockSearchPhaseContext.numSuccess.get());
assertEquals(1, mockSearchPhaseContext.failures.size());
assertTrue(mockSearchPhaseContext.failures.get(0).getCause() instanceof MockDirectoryWrapper.FakeIOException);
assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(2L));
assertNull(responseRef.get().get(1));
}
Aggregations