diff --git a/src/main/java/org/opensearch/search/relevance/actionfilter/SearchActionFilter.java b/src/main/java/org/opensearch/search/relevance/actionfilter/SearchActionFilter.java index 7ffbe39..1134b35 100644 --- a/src/main/java/org/opensearch/search/relevance/actionfilter/SearchActionFilter.java +++ b/src/main/java/org/opensearch/search/relevance/actionfilter/SearchActionFilter.java @@ -39,6 +39,8 @@ import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; import org.opensearch.search.aggregations.InternalAggregations; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.fetch.subphase.FetchSourceContext; import org.opensearch.search.internal.InternalSearchResponse; import org.opensearch.search.profile.SearchProfileShardResults; import org.opensearch.search.relevance.configuration.ConfigurationUtils; @@ -50,217 +52,240 @@ import org.opensearch.tasks.Task; public class SearchActionFilter implements ActionFilter { - private static final Logger logger = LogManager.getLogger(SearchActionFilter.class); - - private final int order; - - private final NamedWriteableRegistry namedWriteableRegistry; - private final Map supportedResultTransformers; - private final OpenSearchClient openSearchClient; - - public SearchActionFilter(Map supportedResultTransformers, OpenSearchClient openSearchClient) { - order = 10; // TODO: Finalize this value - namedWriteableRegistry = new NamedWriteableRegistry(Collections.emptyList()); - this.supportedResultTransformers = supportedResultTransformers; - this.openSearchClient = openSearchClient; - } - - @Override - public int order() { - return order; - } - - @Override - public void apply( - final Task task, - final String action, - final Request request, - final ActionListener listener, - final ActionFilterChain chain) { - - final long startTime = System.nanoTime(); - - if (!SearchAction.INSTANCE.name().equals(action)) { - chain.proceed(task, action, request, listener); - return; - } - - final SearchRequest searchRequest = (SearchRequest) request; + private static final Logger logger = LogManager.getLogger(SearchActionFilter.class); - final String[] indices = searchRequest.indices(); - // Skip if no, or more than 1, index is specified. - if (indices == null || indices.length != 1) { - chain.proceed(task, action, request, listener); - return; - } + private final int order; - List resultTransformerConfigurations = getResultTransformerConfigurations(indices[0], searchRequest); + private final NamedWriteableRegistry namedWriteableRegistry; + private final Map supportedResultTransformers; + private final OpenSearchClient openSearchClient; - LinkedHashMap orderedTransformersAndConfigs = new LinkedHashMap<>(); - for (ResultTransformerConfiguration config : resultTransformerConfigurations) { - ResultTransformer resultTransformer = supportedResultTransformers.get(config.getType()); - if (resultTransformer.shouldTransform(searchRequest, config)) { - orderedTransformersAndConfigs.put(resultTransformer, config); - } + public SearchActionFilter(Map supportedResultTransformers, OpenSearchClient openSearchClient) { + order = 10; // TODO: Finalize this value + namedWriteableRegistry = new NamedWriteableRegistry(Collections.emptyList()); + this.supportedResultTransformers = supportedResultTransformers; + this.openSearchClient = openSearchClient; } - if (!orderedTransformersAndConfigs.isEmpty()) { - // Source is returned in response hits by default. If disabled by the user, overwrite and enable - // in order to access document contents for reranking, then suppress at response time. - boolean suppressSourceOnResponse = false; - if (searchRequest.source() != null && searchRequest.source().fetchSource() != null && - !searchRequest.source().fetchSource().fetchSource()) { - searchRequest.source().fetchSource(true); - suppressSourceOnResponse = true; - } - - final ActionListener searchResponseListener = createSearchResponseListener( - listener, startTime, searchRequest, orderedTransformersAndConfigs, suppressSourceOnResponse); - chain.proceed(task, action, request, searchResponseListener); - return; + @Override + public int order() { + return order; } - chain.proceed(task, action, request, listener); - } - - /** - * Parse and return a list of result transformers from request and index level configurations - * Request level configuration takes precedence over index level - * @param indexName name of the OpenSearch index - * @param searchRequest input request - * @return ordered and validated list of result transformers, empty list if not specified at - * either request or index level - */ - private List getResultTransformerConfigurations( - final String indexName, - final SearchRequest searchRequest) { - - List configs = new ArrayList<>(); - - // Request level configuration takes precedence over index level - configs = ConfigurationUtils.getResultTransformersFromRequestConfiguration(searchRequest); - if (!configs.isEmpty()) { - return configs; - } + @Override + public void apply( + final Task task, + final String action, + final Request request, + final ActionListener listener, + final ActionFilterChain chain) { - // Fetch all index settings for this plugin - String[] settingNames = supportedResultTransformers.values() - .stream() - .map(t -> t.getTransformerSettings() - .stream() - .map(Setting::getKey) - .collect(Collectors.toList())) - .flatMap(Collection::stream) - .toArray(String[]::new); - - configs = ConfigurationUtils.getResultTransformersFromIndexConfiguration( - openSearchClient.getIndexSettings(indexName, settingNames)); - - return configs; - } - - /** - * Create a Listener that, during the OpenSearch response chain, - * calls external service Kendra Ranking to rerank OpenSearch hits - * @param listener default listened - * @param startTime time when request was received, used to calculate latency added by reranking - * @param searchRequest input search request - * @param orderedTransformersAndConfigs transformers to apply, with their corresponding configurations - * @param suppressSourceOnResponse boolean indicating whether to suppress the document source on response - * @param OpenSearch response type - * @return ActionListener with override for onResponse method - */ - private ActionListener createSearchResponseListener( - final ActionListener listener, - final long startTime, - final SearchRequest searchRequest, - final LinkedHashMap orderedTransformersAndConfigs, - final boolean suppressSourceOnResponse) { - return new ActionListener() { - - @Override - public void onResponse(final Response response) { - final SearchResponse searchResponse = (SearchResponse) response; - final long totalHits = searchResponse.getHits().getTotalHits().value; - if (totalHits == 0) { - logger.info("TotalHits = 0. Returning search response without re-ranking."); - listener.onResponse(response); - return; + final long startTime = System.nanoTime(); + + if (!SearchAction.INSTANCE.name().equals(action)) { + chain.proceed(task, action, request, listener); + return; } - logger.debug("Starting re-ranking for search response: {}", searchResponse); - try { - final BytesStreamOutput out = new BytesStreamOutput(); - searchResponse.writeTo(out); - - final StreamInput in = new NamedWriteableAwareStreamInput(out.bytes().streamInput(), - namedWriteableRegistry); - - SearchHits hits = new SearchHits(in); - for (Map.Entry entry : orderedTransformersAndConfigs.entrySet()) { - hits = entry.getKey().transform(hits, searchRequest, entry.getValue()); - } - - if (suppressSourceOnResponse) { - List hitsWithModifiedSource = Arrays.stream(hits.getHits()) - .map(hit -> hit.sourceRef(null)) - .collect(Collectors.toList()); - hits = new SearchHits( - hitsWithModifiedSource.toArray(new SearchHit[hitsWithModifiedSource.size()]), - hits.getTotalHits(), - hits.getMaxScore()); - } - - final InternalAggregations aggregations = - in.readBoolean() ? InternalAggregations.readFrom(in) : null; - final Suggest suggest = in.readBoolean() ? new Suggest(in) : null; - final boolean timedOut = in.readBoolean(); - final Boolean terminatedEarly = in.readOptionalBoolean(); - final SearchProfileShardResults profileResults = in.readOptionalWriteable( - SearchProfileShardResults::new); - final int numReducePhases = in.readVInt(); - - final SearchResponseSections internalResponse = new InternalSearchResponse(hits, - aggregations, suggest, - profileResults, timedOut, terminatedEarly, numReducePhases); - - final int totalShards = in.readVInt(); - final int successfulShards = in.readVInt(); - final int shardSearchFailureSize = in.readVInt(); - final ShardSearchFailure[] shardFailures; - if (shardSearchFailureSize == 0) { - shardFailures = ShardSearchFailure.EMPTY_ARRAY; - } else { - shardFailures = new ShardSearchFailure[shardSearchFailureSize]; - for (int i = 0; i < shardFailures.length; i++) { - shardFailures[i] = readShardSearchFailure(in); + SearchRequest searchRequest = (SearchRequest) request; + + // TODO: Remove originalSearchSource and replace with a deep copy of the SearchRequest object + // once https://github.com/opensearch-project/OpenSearch/issues/869 is implemented + SearchSourceBuilder originalSearchSource = null; + if (searchRequest.source() != null) { + originalSearchSource = searchRequest.source().shallowCopy(); + if (searchRequest.source().fetchSource() != null) { + FetchSourceContext fetchSourceContext = searchRequest.source().fetchSource(); + // Clone the fetchSource + originalSearchSource.fetchSource(new FetchSourceContext(fetchSourceContext.fetchSource(), + fetchSourceContext.includes(), fetchSourceContext.excludes())); } - } - - final SearchResponse.Clusters clusters = new SearchResponse.Clusters(in.readVInt(), - in.readVInt(), in.readVInt()); - final String scrollId = in.readOptionalString(); - final int skippedShards = in.readVInt(); - - final long tookInMillis = (System.nanoTime() - startTime) / 1000000; - final SearchResponse newResponse = new SearchResponse(internalResponse, scrollId, - totalShards, successfulShards, - skippedShards, tookInMillis, shardFailures, clusters); - listener.onResponse((Response) newResponse); - - // TODO: Change this to a metric - logger.info("Re-ranking overhead time: {}ms", - tookInMillis - searchResponse.getTook().getMillis()); - } catch (final Exception e) { - logger.error("Result transformer operations failed.", e); - throw new OpenSearchException("Result transformer operations failed.", e); } - } - - @Override - public void onFailure(final Exception e) { - listener.onFailure(e); - } - }; - } + + final String[] indices = searchRequest.indices(); + // Skip if no, or more than 1, index is specified. + if (indices == null || indices.length != 1) { + chain.proceed(task, action, request, listener); + return; + } + + List resultTransformerConfigurations = getResultTransformerConfigurations(indices[0], + searchRequest); + + LinkedHashMap orderedTransformersAndConfigs = new LinkedHashMap<>(); + for (ResultTransformerConfiguration config : resultTransformerConfigurations) { + ResultTransformer resultTransformer = supportedResultTransformers.get(config.getType()); + // TODO: Should transformers make a decision based on the original request or the request they receive in the chain + if (resultTransformer.shouldTransform(searchRequest, config)) { + searchRequest = resultTransformer.preprocessRequest(searchRequest, config); + orderedTransformersAndConfigs.put(resultTransformer, config); + } + } + + if (!orderedTransformersAndConfigs.isEmpty()) { + final ActionListener searchResponseListener = createSearchResponseListener( + listener, startTime, orderedTransformersAndConfigs, searchRequest, originalSearchSource); + chain.proceed(task, action, request, searchResponseListener); + return; + } + + chain.proceed(task, action, request, listener); + } + + /** + * Parse and return a list of result transformers from request and index level configurations + * Request level configuration takes precedence over index level + * + * @param indexName name of the OpenSearch index + * @param searchRequest input request + * @return ordered and validated list of result transformers, empty list if not specified at + * either request or index level + */ + private List getResultTransformerConfigurations( + final String indexName, + final SearchRequest searchRequest) { + + List configs = new ArrayList<>(); + + // Request level configuration takes precedence over index level + configs = ConfigurationUtils.getResultTransformersFromRequestConfiguration(searchRequest); + if (!configs.isEmpty()) { + return configs; + } + + // Fetch all index settings for this plugin + String[] settingNames = supportedResultTransformers.values() + .stream() + .map(t -> t.getTransformerSettings() + .stream() + .map(Setting::getKey) + .collect(Collectors.toList())) + .flatMap(Collection::stream) + .toArray(String[]::new); + + configs = ConfigurationUtils.getResultTransformersFromIndexConfiguration( + openSearchClient.getIndexSettings(indexName, settingNames)); + + return configs; + } + + /** + * Create a Listener that, during the OpenSearch response chain, + * calls external service Kendra Ranking to rerank OpenSearch hits + * + * @param listener default listened + * @param startTime time when request was received, used to calculate latency added by reranking + * @param searchRequest input search request + * @param orderedTransformersAndConfigs transformers to apply, with their corresponding configurations + * @param originalSearchSource original search source without any modifications made by transformers + * @param OpenSearch response type + * @return ActionListener with override for onResponse method + */ + private ActionListener createSearchResponseListener( + final ActionListener listener, + final long startTime, + final LinkedHashMap orderedTransformersAndConfigs, + final SearchRequest searchRequest, + final SearchSourceBuilder originalSearchSource) { + return new ActionListener() { + + @Override + public void onResponse(final Response response) { + final SearchResponse searchResponse = (SearchResponse) response; + final long totalHits = searchResponse.getHits().getTotalHits().value; + if (totalHits == 0) { + logger.info("TotalHits = 0. Returning search response without re-ranking."); + listener.onResponse(response); + return; + } + + logger.debug("Starting re-ranking for search response: {}", searchResponse); + try { + final BytesStreamOutput out = new BytesStreamOutput(); + searchResponse.writeTo(out); + + final StreamInput in = new NamedWriteableAwareStreamInput(out.bytes().streamInput(), namedWriteableRegistry); + + SearchHits hits = new SearchHits(in); + for (Map.Entry entry : orderedTransformersAndConfigs.entrySet()) { + hits = entry.getKey().transform(hits, searchRequest, entry.getValue()); + } + + List searchHitsList = Arrays.asList(hits.getHits()); + if (originalSearchSource != null) { + if (originalSearchSource.fetchSource() != null && + !originalSearchSource.fetchSource().fetchSource()) { + searchHitsList = searchHitsList.stream() + .map(hit -> hit.sourceRef(null)) + .collect(Collectors.toList()); + } + if (originalSearchSource.from() >= 0 && originalSearchSource.size() >= 0) { + final int lastHitIndex = Math.min(searchHitsList.size(), + (originalSearchSource.from() + originalSearchSource.size())); + if (originalSearchSource.from() > lastHitIndex) { + searchHitsList = Collections.emptyList(); + } else { + searchHitsList = searchHitsList.subList(originalSearchSource.from(), lastHitIndex); + } + } + } + + // TODO: How to handle SearchHits.TotalHits when transformer modifies the hit count + hits = new SearchHits( + searchHitsList.toArray(new SearchHit[0]), + hits.getTotalHits(), + hits.getMaxScore()); + + final InternalAggregations aggregations = + in.readBoolean() ? InternalAggregations.readFrom(in) : null; + final Suggest suggest = in.readBoolean() ? new Suggest(in) : null; + final boolean timedOut = in.readBoolean(); + final Boolean terminatedEarly = in.readOptionalBoolean(); + final SearchProfileShardResults profileResults = in.readOptionalWriteable( + SearchProfileShardResults::new); + final int numReducePhases = in.readVInt(); + + final SearchResponseSections internalResponse = new InternalSearchResponse(hits, + aggregations, suggest, + profileResults, timedOut, terminatedEarly, numReducePhases); + + final int totalShards = in.readVInt(); + final int successfulShards = in.readVInt(); + final int shardSearchFailureSize = in.readVInt(); + final ShardSearchFailure[] shardFailures; + if (shardSearchFailureSize == 0) { + shardFailures = ShardSearchFailure.EMPTY_ARRAY; + } else { + shardFailures = new ShardSearchFailure[shardSearchFailureSize]; + for (int i = 0; i < shardFailures.length; i++) { + shardFailures[i] = readShardSearchFailure(in); + } + } + + final SearchResponse.Clusters clusters = new SearchResponse.Clusters(in.readVInt(), + in.readVInt(), in.readVInt()); + final String scrollId = in.readOptionalString(); + final int skippedShards = in.readVInt(); + + final long tookInMillis = (System.nanoTime() - startTime) / 1000000; + final SearchResponse newResponse = new SearchResponse(internalResponse, scrollId, + totalShards, successfulShards, + skippedShards, tookInMillis, shardFailures, clusters); + listener.onResponse((Response) newResponse); + + // TODO: Change this to a metric + logger.info("Result transformer operations overhead time: {}ms", + tookInMillis - searchResponse.getTook().getMillis()); + } catch (final Exception e) { + logger.error("Result transformer operations failed.", e); + throw new OpenSearchException("Result transformer operations failed.", e); + } + } + + @Override + public void onFailure(final Exception e) { + listener.onFailure(e); + } + }; + } } diff --git a/src/main/java/org/opensearch/search/relevance/transformer/ResultTransformer.java b/src/main/java/org/opensearch/search/relevance/transformer/ResultTransformer.java index c5b77d8..6568828 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/ResultTransformer.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/ResultTransformer.java @@ -23,12 +23,21 @@ public interface ResultTransformer { /** * Decide whether to apply the transformer on the input request - * @param request input request + * @param request input Search Request * @param configuration Configuration parameters for the transformer * @return boolean decision on whether to apply the transformer */ boolean shouldTransform(final SearchRequest request, final ResultTransformerConfiguration configuration); + /** + * Preprocess the incoming Search Request to support the requirements of the transformer + * @param request input Search Request + * @param configuration Configuration parameters for the transformer + * @return SearchRequest with updated attributes + */ + SearchRequest preprocessRequest(final SearchRequest request, + final ResultTransformerConfiguration configuration); + /** * Rank hits based on the provided query * @param hits hits to be re-ranked @@ -36,5 +45,7 @@ public interface ResultTransformer { * @param configuration Configuration parameters for the transformer * @return SearchHits ordered by score generated by ranker */ - SearchHits transform(final SearchHits hits, final SearchRequest request, final ResultTransformerConfiguration configuration); + SearchHits transform(final SearchHits hits, + final SearchRequest request, + final ResultTransformerConfiguration configuration); } diff --git a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRanker.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRanker.java index d8fe8fc..6c5597c 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRanker.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRanker.java @@ -7,11 +7,15 @@ */ package org.opensearch.search.relevance.transformer.kendraintelligentranking; +import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants.BODY_FIELD; + import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.Comparator; import java.util.HashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.PriorityQueue; @@ -22,10 +26,12 @@ import org.opensearch.common.settings.Setting; import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; +import org.opensearch.search.SearchService; import org.opensearch.search.relevance.configuration.ResultTransformerConfiguration; import org.opensearch.search.relevance.transformer.ResultTransformer; import org.opensearch.search.relevance.transformer.kendraintelligentranking.client.KendraHttpClient; import org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankingConfiguration; +import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.KendraIntelligentRankingException; import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.PassageScore; import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.dto.Document; import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.dto.RescoreRequest; @@ -40,135 +46,180 @@ public class KendraIntelligentRanker implements ResultTransformer { - private static final int PASSAGE_SIZE_LIMIT = 600; - private static final int SLIDING_WINDOW_STEP = PASSAGE_SIZE_LIMIT - 50; - private static final int MAXIMUM_PASSAGES = 10; - private static final double BM25_B_VALUE = 0.75; - private static final double BM25_K1_VALUE = 1.6; - private static final int TOP_K_PASSAGES = 3; - - private static final Logger logger = LogManager.getLogger(KendraIntelligentRanker.class); - - private final KendraHttpClient kendraClient; - private final TextTokenizer textTokenizer; - private final QueryParser queryParser; - - public KendraIntelligentRanker(KendraHttpClient kendraClient) { - this.kendraClient = kendraClient; - this.textTokenizer = new TextTokenizer(); - this.queryParser = new QueryParser(); - } - - @Override - public List> getTransformerSettings() { - return KendraIntelligentRankerSettings.getAllSettings(); - } - - /** - * Check if search request is eligible for rescore - * @param request Search Request - * @return boolean decision on whether to re-rank - */ - @Override - public boolean shouldTransform(SearchRequest request, ResultTransformerConfiguration configuration) { - if (request.source() == null) { - return false; - } + private static final int PASSAGE_SIZE_LIMIT = 600; + private static final int SLIDING_WINDOW_STEP = PASSAGE_SIZE_LIMIT - 50; + private static final int MAXIMUM_PASSAGES = 10; + private static final double BM25_B_VALUE = 0.75; + private static final double BM25_K1_VALUE = 1.6; + private static final int TOP_K_PASSAGES = 3; + + private static final Logger logger = LogManager.getLogger(KendraIntelligentRanker.class); - // Skip if there is scroll, pagination, or sorting. - if (request.scroll() != null || request.source().from() > 0 || - (request.source().sorts() != null && !request.source().sorts().isEmpty())) { - return false; + private final KendraHttpClient kendraClient; + private final TextTokenizer textTokenizer; + private final QueryParser queryParser; + + public KendraIntelligentRanker(KendraHttpClient kendraClient) { + this.kendraClient = kendraClient; + this.textTokenizer = new TextTokenizer(); + this.queryParser = new QueryParser(); } - return true; - } - - /** - * - * @param hits Search hits to rerank with respect to query - * @param request Search request - * @return SearchHits reranked search hits - */ - @Override - public SearchHits transform(final SearchHits hits, - final SearchRequest request, - final ResultTransformerConfiguration configuration) { - KendraIntelligentRankingConfiguration kendraConfig = (KendraIntelligentRankingConfiguration) configuration; - QueryParserResult queryParserResult = queryParser.parse( - request.source().query(), - kendraConfig.getProperties().getBodyFields(), - kendraConfig.getProperties().getTitleFields()); - if (queryParserResult == null) { - return hits; + + @Override + public List> getTransformerSettings() { + return KendraIntelligentRankerSettings.getAllSettings(); } - try { - List originalHits = new ArrayList<>(); - for (SearchHit searchHit : hits.getHits()) { - Map docSourceMap = searchHit.getSourceAsMap(); - SlidingWindowTextSplitter textSplitter = new SlidingWindowTextSplitter(PASSAGE_SIZE_LIMIT, SLIDING_WINDOW_STEP, MAXIMUM_PASSAGES); - List splitPassages = textSplitter.split(docSourceMap.get(queryParserResult.getBodyFieldName()).toString()); - List> topPassages = getTopPassages(queryParserResult.getQueryText(), splitPassages); - List tokenizedTitle = null; - if (queryParserResult.getTitleFieldName() != null) { - tokenizedTitle = textTokenizer.tokenize(docSourceMap.get(queryParserResult.getTitleFieldName()).toString()); - // If tokens list is empty, use null - if (tokenizedTitle.isEmpty()) { - tokenizedTitle = null; - } + + /** + * Check if search request is eligible for rescore + * + * @param request Search Request + * @return boolean decision on whether to re-rank + */ + @Override + public boolean shouldTransform(final SearchRequest request, final ResultTransformerConfiguration configuration) { + if (request.source() == null || request.source().query() == null) { + return false; } - for (int i = 0; i < topPassages.size(); i++) { - originalHits.add( - new Document(searchHit.getId() + "@" + (i + 1), searchHit.getId(), tokenizedTitle, topPassages.get(i), searchHit.getScore()) - ); + KendraIntelligentRankingConfiguration kendraConfiguration = (KendraIntelligentRankingConfiguration) configuration; + + // Skip if there is scroll, sorting, or the start of the page is greater than the document limit for Kendra Ranking + if (request.scroll() != null || + (request.source().sorts() != null && !request.source().sorts().isEmpty()) || + request.source().from() >= kendraConfiguration.getProperties().getDocLimit()) { + return false; } - } - - final RescoreRequest rescoreRequest = new RescoreRequest(queryParserResult.getQueryText(), originalHits); - final RescoreResult rescoreResult = kendraClient.rescore(rescoreRequest); - Map idToSearchHitMap = new HashMap<>(); - for (SearchHit searchHit : hits.getHits()) { - idToSearchHitMap.put(searchHit.getId(), searchHit); - } - List newSearchHits = new ArrayList<>(); - float maxScore = 0; - for (RescoreResultItem rescoreResultItem : rescoreResult.getResultItems()) { - SearchHit searchHit = idToSearchHitMap.get(rescoreResultItem.getDocumentId()); - if (searchHit == null) { - logger.warn("Response from external service references hit id {}, which does not exist in original results. Skipping.", - rescoreResultItem.getDocumentId()); - continue; + return true; + } + + @Override + public SearchRequest preprocessRequest(final SearchRequest request, + final ResultTransformerConfiguration configuration) { + // Source is returned in response hits by default. If disabled by the user, overwrite and enable + // in order to access document contents for reranking, then suppress at response time. + if (request.source() != null && request.source().fetchSource() != null && + !request.source().fetchSource().fetchSource()) { + request.source().fetchSource(true); } - searchHit.score(rescoreResultItem.getScore()); - maxScore = Math.max(maxScore, rescoreResultItem.getScore()); - newSearchHits.add(searchHit); - } - return new SearchHits(newSearchHits.toArray(new SearchHit[newSearchHits.size()]), hits.getTotalHits(), maxScore); - } catch (Exception ex) { - logger.error("Failed to re-rank. Returning original search results without re-ranking.", ex); - return hits; + + int from = request.source().from() == -1 ? SearchService.DEFAULT_FROM : request.source().from(); + int size = request.source().size() == -1 ? SearchService.DEFAULT_SIZE : request.source().size(); + + KendraIntelligentRankingConfiguration kendraConfiguration = (KendraIntelligentRankingConfiguration) configuration; + int sizeOverride = Math.max(kendraConfiguration.getProperties().getDocLimit(), from + size); + request.source().from(SearchService.DEFAULT_FROM); + request.source().size(sizeOverride); + return request; } - } - - private List> getTopPassages(final String queryText, final List splitPassages) { - List query = textTokenizer.tokenize(queryText); - List> passages = textTokenizer.tokenize(splitPassages); - BM25Scorer bm25Scorer = new BM25Scorer(BM25_B_VALUE, BM25_K1_VALUE, passages); - PriorityQueue pq = new PriorityQueue<>(Comparator.comparingDouble(x -> x.getScore())); - - for (int i = 0; i < passages.size(); i++) { - double score = bm25Scorer.score(query, passages.get(i)); - pq.offer(new PassageScore(score, i)); - if (pq.size() > TOP_K_PASSAGES) { - // Maintain heap of top K passages - pq.poll(); - } + + /** + * @param hits Search hits to rerank with respect to query + * @param request Search request + * @return SearchHits reranked search hits + */ + @Override + public SearchHits transform(final SearchHits hits, + final SearchRequest request, + final ResultTransformerConfiguration configuration) { + if (hits.getHits().length == 0) { + // Avoid call to rerank empty results + return hits; + } + KendraIntelligentRankingConfiguration kendraConfig = (KendraIntelligentRankingConfiguration) configuration; + QueryParserResult queryParserResult = queryParser.parse( + request.source().query(), + kendraConfig.getProperties().getBodyFields(), + kendraConfig.getProperties().getTitleFields()); + if (queryParserResult == null) { + // Unknown query type or query does not reference body field + return hits; + } + KendraIntelligentRankingConfiguration kendraConfiguration = (KendraIntelligentRankingConfiguration) configuration; + try { + List originalHits = Arrays.asList(hits.getHits()); + final int numberOfHitsToRerank = Math.min(originalHits.size(), kendraConfiguration.getProperties().getDocLimit()); + List originalHitsAsDocuments = new ArrayList<>(); + Map idToSearchHitMap = new HashMap<>(); + for (int j = 0; j < numberOfHitsToRerank; ++j) { + Map docSourceMap = originalHits.get(j).getSourceAsMap(); + SlidingWindowTextSplitter textSplitter = new SlidingWindowTextSplitter(PASSAGE_SIZE_LIMIT, SLIDING_WINDOW_STEP, MAXIMUM_PASSAGES); + String bodyFieldName = queryParserResult.getBodyFieldName(); + String titleFieldName = queryParserResult.getTitleFieldName(); + if (docSourceMap.get(bodyFieldName) == null) { + String errorMessage = String.format(Locale.ENGLISH, + "Kendra Intelligent Ranking cannot be applied when documents are missing %s [%s]. Document ID [%s].", + BODY_FIELD, bodyFieldName, originalHits.get(j).getId()); + logger.error(errorMessage); + throw new KendraIntelligentRankingException(errorMessage); + } + List splitPassages = textSplitter.split(docSourceMap.get(bodyFieldName).toString()); + List> topPassages = getTopPassages(queryParserResult.getQueryText(), splitPassages); + List tokenizedTitle = null; + if (titleFieldName != null && docSourceMap.get(titleFieldName) != null) { + tokenizedTitle = textTokenizer.tokenize(docSourceMap.get(queryParserResult.getTitleFieldName()).toString()); + // If tokens list is empty, use null + if (tokenizedTitle.isEmpty()) { + tokenizedTitle = null; + } + } + for (int i = 0; i < topPassages.size(); ++i) { + originalHitsAsDocuments.add( + new Document(originalHits.get(j).getId() + "@" + (i + 1), originalHits.get(j).getId(), tokenizedTitle, topPassages.get(i), originalHits.get(j).getScore()) + ); + } + // Map search hits by their ID in order to map Kendra response documents back to hits later + idToSearchHitMap.put(originalHits.get(j).getId(), originalHits.get(j)); + } + + final RescoreRequest rescoreRequest = new RescoreRequest(queryParserResult.getQueryText(), originalHitsAsDocuments); + final RescoreResult rescoreResult = kendraClient.rescore(rescoreRequest); + + List newSearchHits = new ArrayList<>(); + float maxScore = 0; + for (RescoreResultItem rescoreResultItem : rescoreResult.getResultItems()) { + SearchHit searchHit = idToSearchHitMap.get(rescoreResultItem.getDocumentId()); + if (searchHit == null) { + String errorMessage = String.format(Locale.ENGLISH, + "Response from Kendra Intelligent Ranking service references document ID [%s], which does not exist in original results", + rescoreResultItem.getDocumentId()); + logger.error(errorMessage); + throw new KendraIntelligentRankingException(errorMessage); + } + searchHit.score(rescoreResultItem.getScore()); + maxScore = Math.max(maxScore, rescoreResultItem.getScore()); + newSearchHits.add(searchHit); + } + // Add remaining hits to response, which are already sorted by OpenSearch score + for (int i = numberOfHitsToRerank; i < originalHits.size(); ++i) { + newSearchHits.add(originalHits.get(i)); + } + return new SearchHits(newSearchHits.toArray(new SearchHit[newSearchHits.size()]), hits.getTotalHits(), maxScore); + } catch (Exception ex) { + logger.error("Failed to rescore. Returning original search results without rescore.", ex); + return hits; + } } - List> topPassages = new ArrayList<>(); - while (!pq.isEmpty()) { - topPassages.add(passages.get(pq.poll().getIndex())); + private List> getTopPassages(final String queryText, final List splitPassages) { + List query = textTokenizer.tokenize(queryText); + List> passages = textTokenizer.tokenize(splitPassages); + BM25Scorer bm25Scorer = new BM25Scorer(BM25_B_VALUE, BM25_K1_VALUE, passages); + PriorityQueue pq = new PriorityQueue<>(Comparator.comparingDouble(x -> x.getScore())); + + for (int i = 0; i < passages.size(); i++) { + double score = bm25Scorer.score(query, passages.get(i)); + pq.offer(new PassageScore(score, i)); + if (pq.size() > TOP_K_PASSAGES) { + // Maintain heap of top K passages + pq.poll(); + } + } + + List> topPassages = new ArrayList<>(); + while (!pq.isEmpty()) { + topPassages.add(passages.get(pq.poll().getIndex())); + } + Collections.reverse(topPassages); // reverse to order from highest to lowest score + return topPassages; } - Collections.reverse(topPassages); // reverse to order from highest to lowest score - return topPassages; - } } diff --git a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/Constants.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/Constants.java index 61e07cd..d4b9a66 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/Constants.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/Constants.java @@ -17,6 +17,7 @@ public class Constants { // Transformer properties public static final String BODY_FIELD = "body_field"; public static final String TITLE_FIELD = "title_field"; + public static final String DOC_LIMIT = "doc_limit"; public static final String KENDRA_SETTINGS_PREFIX = String.join(".", RESULT_TRANSFORMER_SETTING_PREFIX, KENDRA_INTELLIGENT_RANKING); @@ -27,5 +28,8 @@ public class Constants { String.join(".", KENDRA_SETTINGS_PREFIX, PROPERTIES, BODY_FIELD); public static final String TITLE_FIELD_SETTING_NAME = String.join(".", KENDRA_SETTINGS_PREFIX, PROPERTIES, TITLE_FIELD); + public static final String DOC_LIMIT_SETTING_NAME = + String.join(".", KENDRA_SETTINGS_PREFIX, PROPERTIES, DOC_LIMIT); + public static final int KENDRA_DEFAULT_DOC_LIMIT = 25; } diff --git a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankerSettings.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankerSettings.java index 756b344..1648ad9 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankerSettings.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankerSettings.java @@ -15,7 +15,6 @@ import org.opensearch.common.settings.SecureString; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Setting.Property; -import org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants; public class KendraIntelligentRankerSettings { @@ -26,24 +25,13 @@ public class KendraIntelligentRankerSettings { Property.Dynamic, Property.IndexScope); /** - * Document field to be considered as "body" when invoking Kendra. - */ - public static final Setting> KENDRA_BODY_FIELD_SETTING = Setting.listSetting(Constants.BODY_FIELD_SETTING_NAME, Collections.emptyList(), - Function.identity(), new FieldSettingValidator(Constants.BODY_FIELD_SETTING_NAME), - Property.Dynamic, Property.IndexScope); - - /** - * Document field to be considered as "title" when invoking Kendra. + * Validator for body and title field settings */ - public static final Setting> KENDRA_TITLE_FIELD_SETTING = Setting.listSetting(Constants.TITLE_FIELD_SETTING_NAME, Collections.emptyList(), - Function.identity(), new FieldSettingValidator(Constants.TITLE_FIELD_SETTING_NAME), - Property.Dynamic, Property.IndexScope); - - static final class FieldSettingValidator implements Setting.Validator> { + static final class FieldValidator implements Setting.Validator> { private String settingName; - public FieldSettingValidator(final String name) { + public FieldValidator(final String name) { this.settingName = name; } @@ -55,6 +43,55 @@ public void validate(List value) { } } + /** + * Validator for doc limit setting + */ + static final class DocLimitValidator implements Setting.Validator { + + private String settingName; + + public DocLimitValidator(final String name) { + this.settingName = name; + } + + @Override + public void validate(Integer value) { + if (value != null && value < Constants.KENDRA_DEFAULT_DOC_LIMIT) { + throw new IllegalArgumentException("Setting the value of [" + this.settingName + "] below " + + Constants.KENDRA_DEFAULT_DOC_LIMIT + " will affect ranking accuracy"); + } + } + } + + /** + * Validator objects + */ + public static final FieldValidator BODY_FIELD_VALIDATOR = new FieldValidator(Constants.BODY_FIELD); + public static final FieldValidator TITLE_FIELD_VALIDATOR = new FieldValidator(Constants.TITLE_FIELD); + public static final DocLimitValidator DOC_LIMIT_VALIDATOR = new DocLimitValidator(Constants.DOC_LIMIT); + + /** + * Document field to be considered as "body" when invoking Kendra. + */ + public static final Setting> KENDRA_BODY_FIELD_SETTING = Setting.listSetting(Constants.BODY_FIELD_SETTING_NAME, + Collections.emptyList(), Function.identity(), BODY_FIELD_VALIDATOR, + Property.Dynamic, Property.IndexScope); + + /** + * Document field to be considered as "title" when invoking Kendra. + */ + public static final Setting> KENDRA_TITLE_FIELD_SETTING = Setting.listSetting(Constants.TITLE_FIELD_SETTING_NAME, + Collections.emptyList(), Function.identity(), TITLE_FIELD_VALIDATOR, + Property.Dynamic, Property.IndexScope); + + + + public static final Setting KENDRA_DOC_LIMIT_SETTING = Setting.intSetting( + Constants.DOC_LIMIT_SETTING_NAME, Constants.KENDRA_DEFAULT_DOC_LIMIT, 1, + DOC_LIMIT_VALIDATOR, Property.Dynamic, Property.IndexScope); + + + /** * The access key (ie login id) for connecting to Kendra. */ @@ -83,6 +120,7 @@ public static final List> getAllSettings() { KENDRA_ORDER_SETTING, KENDRA_BODY_FIELD_SETTING, KENDRA_TITLE_FIELD_SETTING, + KENDRA_DOC_LIMIT_SETTING, ACCESS_KEY_SETTING, SECRET_KEY_SETTING, SESSION_TOKEN_SETTING, diff --git a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankingConfiguration.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankingConfiguration.java index e9dddc4..6bafb3d 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankingConfiguration.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankingConfiguration.java @@ -8,8 +8,14 @@ package org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration; import static org.opensearch.search.relevance.configuration.Constants.ORDER; +import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants.DOC_LIMIT_SETTING_NAME; +import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants.KENDRA_DEFAULT_DOC_LIMIT; +import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankerSettings.BODY_FIELD_VALIDATOR; +import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankerSettings.DOC_LIMIT_VALIDATOR; +import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankerSettings.TITLE_FIELD_VALIDATOR; import java.io.IOException; +import java.util.Collections; import java.util.List; import java.util.Objects; import org.opensearch.common.ParseField; @@ -55,7 +61,8 @@ public KendraIntelligentRankingConfiguration(Settings settings) { this.order = settings.getAsInt(ORDER, 0); this.properties = new KendraIntelligentRankingProperties( settings.getAsList("properties.body_field"), - settings.getAsList("properties.title_field")); + settings.getAsList("properties.title_field"), + settings.getAsInt("properties.doc_limit", KENDRA_DEFAULT_DOC_LIMIT)); } @Override @@ -117,6 +124,7 @@ public KendraIntelligentRankingProperties getProperties() { public static class KendraIntelligentRankingProperties implements Writeable, ToXContentObject { protected static final ParseField BODY_FIELD = new ParseField(Constants.BODY_FIELD); protected static final ParseField TITLE_FIELD = new ParseField(Constants.TITLE_FIELD); + protected static final ParseField DOC_LIMIT = new ParseField(Constants.DOC_LIMIT); private static final ObjectParser PARSER; @@ -124,41 +132,46 @@ public static class KendraIntelligentRankingProperties implements Writeable, ToX PARSER = new ObjectParser<>("kendra_intelligent_ranking_configuration", KendraIntelligentRankingProperties::new); PARSER.declareStringArray(KendraIntelligentRankingProperties::setBodyFields, BODY_FIELD); PARSER.declareStringArray(KendraIntelligentRankingProperties::setTitleFields, TITLE_FIELD); + PARSER.declareInt(KendraIntelligentRankingProperties::setDocLimit, DOC_LIMIT); } private List bodyFields; private List titleFields; + private int docLimit; - public KendraIntelligentRankingProperties() {} + public KendraIntelligentRankingProperties() { + bodyFields = Collections.emptyList(); + titleFields = Collections.emptyList(); + docLimit = KENDRA_DEFAULT_DOC_LIMIT; + } - public KendraIntelligentRankingProperties(final List bodyFields, final List titleFields) { + public KendraIntelligentRankingProperties(final List bodyFields, + final List titleFields, final int docLimit) { this.bodyFields = bodyFields; this.titleFields = titleFields; + this.docLimit = docLimit; } public KendraIntelligentRankingProperties(StreamInput input) throws IOException { this.bodyFields = input.readStringList(); this.bodyFields = input.readStringList(); + this.docLimit = input.readInt(); } @Override public void writeTo(StreamOutput out) throws IOException { out.writeStringCollection(this.bodyFields); out.writeStringCollection(this.titleFields); + out.writeInt(this.docLimit); } public static KendraIntelligentRankingProperties parse(XContentParser parser, Void context) throws IOException { try { KendraIntelligentRankingProperties properties = PARSER.parse(parser, null); if (properties != null) { - if (properties.getBodyFields() != null && !properties.getBodyFields().isEmpty() && properties.getBodyFields().size() > 1) { - throw new ParsingException(parser.getTokenLocation(), - "[" + BODY_FIELD + "] can have at most 1 element"); - } - if (properties.getTitleFields() != null && !properties.getTitleFields().isEmpty() && properties.getTitleFields().size() > 1) { - throw new ParsingException(parser.getTokenLocation(), - "[" + TITLE_FIELD + "] can have at most 1 element"); - } + BODY_FIELD_VALIDATOR.validate(properties.getBodyFields()); + TITLE_FIELD_VALIDATOR.validate(properties.getTitleFields()); + DOC_LIMIT_VALIDATOR.validate(properties.getDocLimit()); } return properties; } catch (IllegalArgumentException iae) { @@ -171,6 +184,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.startObject(); builder.field(BODY_FIELD.getPreferredName(), this.bodyFields); builder.field(TITLE_FIELD.getPreferredName(), this.titleFields); + builder.field(DOC_LIMIT.getPreferredName(), this.docLimit); return builder.endObject(); } @@ -181,8 +195,8 @@ public boolean equals(Object o) { KendraIntelligentRankingProperties properties = (KendraIntelligentRankingProperties) o; - if (bodyFields != properties.bodyFields) return false; - return (titleFields == properties.titleFields); + return (bodyFields == properties.bodyFields) && (titleFields == properties.titleFields) && + (docLimit == properties.docLimit); } @Override @@ -205,5 +219,13 @@ public List getTitleFields() { public void setTitleFields(final List titleFields) { this.titleFields = titleFields; } + + public int getDocLimit() { + return this.docLimit; + } + + public void setDocLimit(final int docLimit) { + this.docLimit = docLimit; + } } } diff --git a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/KendraIntelligentRankingException.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/KendraIntelligentRankingException.java new file mode 100644 index 0000000..d50b840 --- /dev/null +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/KendraIntelligentRankingException.java @@ -0,0 +1,30 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.search.relevance.transformer.kendraintelligentranking.model; + +import java.io.IOException; +import org.opensearch.OpenSearchException; +import org.opensearch.common.io.stream.StreamInput; + +public class KendraIntelligentRankingException extends OpenSearchException { + public KendraIntelligentRankingException(StreamInput in) throws IOException { + super(in); + } + + public KendraIntelligentRankingException(String message) { + super(message); + } + + public KendraIntelligentRankingException(String message, Throwable cause) { + super(message, cause); + } + + public KendraIntelligentRankingException(String message, Throwable cause, Object... args) { + super(message, cause, args); + } +} diff --git a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResult.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResult.java index 890ccc2..aece9ca 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResult.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResult.java @@ -24,4 +24,21 @@ public String getRescoreId() { public List getResultItems() { return resultItems; } + + /** + * Setter used for unit tests. + * @param rescoreId The identifier associated with the scores that Amazon Kendra Intelligent Ranking + * gives to the results. + */ + public void setRescoreId(String rescoreId) { + this.rescoreId = rescoreId; + } + + /** + * Setter used for unit tests. + * @param resultItems A list of result items for documents with new relevancy scores. + */ + public void setResultItems(List resultItems) { + this.resultItems = resultItems; + } } diff --git a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResultItem.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResultItem.java index 8d815c9..18c3060 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResultItem.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResultItem.java @@ -22,4 +22,20 @@ public String getDocumentId() { public Float getScore() { return score; } + + /** + * Setter for unit tests. + * @param documentId the ID of the rescored document. + */ + public void setDocumentId(String documentId) { + this.documentId = documentId; + } + + /** + * Setter for unit tests. + * @param score the updated score of the document. + */ + public void setScore(Float score) { + this.score = score; + } } diff --git a/src/test/java/org/opensearch/search/relevance/SearchRelevancePluginIT.java b/src/test/java/org/opensearch/search/relevance/SearchRelevancePluginIT.java index de33cea..fc54c71 100644 --- a/src/test/java/org/opensearch/search/relevance/SearchRelevancePluginIT.java +++ b/src/test/java/org/opensearch/search/relevance/SearchRelevancePluginIT.java @@ -7,33 +7,21 @@ */ package org.opensearch.search.relevance; -import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; import org.apache.http.util.EntityUtils; import org.opensearch.client.Request; import org.opensearch.client.Response; -import org.opensearch.plugins.Plugin; -import org.opensearch.test.OpenSearchIntegTestCase; +import org.opensearch.test.rest.OpenSearchRestTestCase; import java.io.IOException; -import java.util.Collection; -import java.util.Collections; -import static org.hamcrest.Matchers.containsString; - -@ThreadLeakScope(ThreadLeakScope.Scope.NONE) -@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE) -public class SearchRelevancePluginIT extends OpenSearchIntegTestCase { - - @Override - protected Collection> nodePlugins() { - return Collections.singletonList(SearchRelevancePlugin.class); - } +public class SearchRelevancePluginIT extends OpenSearchRestTestCase { public void testPluginInstalled() throws IOException { - Response response = createRestClient().performRequest(new Request("GET", "/_cat/plugins")); + Response response = client().performRequest(new Request("GET", "/_cat/plugins")); String body = EntityUtils.toString(response.getEntity()); logger.info("response body: {}", body); - assertThat(body, containsString("search-processor")); + assertNotNull(body); + assertTrue(body.contains("search-processor")); } } diff --git a/src/test/java/org/opensearch/search/relevance/actionfilter/SearchActionFilterTests.java b/src/test/java/org/opensearch/search/relevance/actionfilter/SearchActionFilterTests.java index fead167..7d748f5 100644 --- a/src/test/java/org/opensearch/search/relevance/actionfilter/SearchActionFilterTests.java +++ b/src/test/java/org/opensearch/search/relevance/actionfilter/SearchActionFilterTests.java @@ -25,11 +25,14 @@ import org.opensearch.action.search.ShardSearchFailure; import org.opensearch.action.support.ActionFilterChain; import org.opensearch.client.Client; +import org.opensearch.common.bytes.BytesReference; import org.opensearch.common.collect.ImmutableOpenMap; +import org.opensearch.common.document.DocumentField; import org.opensearch.common.io.stream.StreamOutput; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.common.xcontent.XContentBuilder; +import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; import org.opensearch.search.builder.SearchSourceBuilder; @@ -47,6 +50,8 @@ import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; @@ -54,6 +59,9 @@ public class SearchActionFilterTests extends OpenSearchTestCase { + /** + * This filter only operates on search requests. Other request types (e.g. Delete) will still pass through. + */ public void testIgnoresDelete() { Client client = Mockito.mock(Client.class); OpenSearchClient openSearchClient = new OpenSearchClient(client); @@ -68,6 +76,9 @@ public void testIgnoresDelete() { assertTrue(proceedCalled.get()); } + /** + * Test short-circuit code path where we skip the filter if no index is specified. + */ public void testIgnoresSearchRequestOnZeroIndices() { Client client = Mockito.mock(Client.class); OpenSearchClient openSearchClient = new OpenSearchClient(client); @@ -82,6 +93,9 @@ public void testIgnoresSearchRequestOnZeroIndices() { assertTrue(proceedCalled.get()); } + /** + * Test short-circuit code path where we skip the filter if multiple indices are specified. + */ public void testIgnoresSearchRequestOnMultipleIndices() { Client client = Mockito.mock(Client.class); OpenSearchClient openSearchClient = new OpenSearchClient(client); @@ -118,6 +132,9 @@ private static Client buildMockClient(String indexName, Settings... settings) { return client; } + /** + * Probe the code path where we have one index, but no transformers. + */ public void testOperatesOnSingleIndexWithNoTransformers() { Client client = buildMockClient("index"); OpenSearchClient openSearchClient = new OpenSearchClient(client); @@ -136,9 +153,19 @@ public void testOperatesOnSingleIndexWithNoTransformers() { private static class MockTransformer implements ResultTransformer { + public MockTransformer() { + requestTransformer = i -> {}; + } + + public MockTransformer(Consumer requestTransformer) { + this.requestTransformer = requestTransformer; + } + + private final Consumer requestTransformer; private boolean getTransformerSettingsWasCalled = false; private boolean shouldTransformWasCalled = false; private boolean transformWasCalled = false; + private boolean preproccessRequestWasCalled = false; @Override @@ -153,6 +180,12 @@ public boolean shouldTransform(SearchRequest request, ResultTransformerConfigura return true; } + @Override + public SearchRequest preprocessRequest(SearchRequest request, ResultTransformerConfiguration configuration) { + preproccessRequestWasCalled = true; + return request; + } + @Override public SearchHits transform(SearchHits hits, SearchRequest request, ResultTransformerConfiguration configuration) { @@ -184,6 +217,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws } }; + /** + * Even if a transformer is wired into the SearchActionFilter, if it's not enabled by search request or + * index setting, the transformer will not be called. + */ public void testTransformerDoesNotRunWhenNotEnabled() { Client client = buildMockClient("index"); OpenSearchClient openSearchClient = new OpenSearchClient(client); @@ -206,11 +243,15 @@ public void testTransformerDoesNotRunWhenNotEnabled() { assertTrue(proceedCalled.get()); // We should try to check for index-level settings assertTrue(mockTransformer.getTransformerSettingsWasCalled); + assertFalse(mockTransformer.preproccessRequestWasCalled); assertFalse(mockTransformer.transformWasCalled); assertFalse(mockTransformer.shouldTransformWasCalled); } - public void testTransformEnabledInRequest() { + /** + * Should be able to enable transformer explicitly in a search request. + */ + public void testTransformEnabledInRequest() throws IOException { Client client = buildMockClient("index"); OpenSearchClient openSearchClient = new OpenSearchClient(client); @@ -235,7 +276,7 @@ public void testTransformEnabledInRequest() { ).setIndices("index") .request(); AtomicBoolean proceedCalled = new AtomicBoolean(false); - SearchResponse searchResponse = buildMockSearchResponse(); + SearchResponse searchResponse = buildMockSearchResponse(randomInt(20)); ActionFilterChain searchFilterChain = (task1, action, request, listener) -> { @@ -244,7 +285,7 @@ public void testTransformEnabledInRequest() { }; AtomicBoolean onResponseCalled = new AtomicBoolean(false); AtomicBoolean onFailureCalled = new AtomicBoolean(false); - ActionListener downstreamListener = new ActionListener() { + ActionListener downstreamListener = new ActionListener<>() { @Override public void onResponse(SearchResponse searchResponse) { onResponseCalled.set(true); @@ -259,21 +300,37 @@ public void onFailure(Exception e) { assertTrue(proceedCalled.get()); // We should NOT try to check for index-level settings, because we saw request-level settings assertFalse(mockTransformer.getTransformerSettingsWasCalled); + assertTrue(mockTransformer.preproccessRequestWasCalled); assertTrue(mockTransformer.transformWasCalled); assertTrue(mockTransformer.shouldTransformWasCalled); assertTrue(onResponseCalled.get()); assertFalse(onFailureCalled.get()); } - private static SearchResponse buildMockSearchResponse() { + private static SearchResponse buildMockSearchResponse(int numHits) throws IOException { + SearchHit[] hitsArray = new SearchHit[numHits]; + for (int i = 0; i < numHits; i++) { + XContentBuilder sourceContent = JsonXContent.contentBuilder() + .startObject() + .field("_id", String.valueOf(i)) + .field("title", "doc" + i) + .endObject(); + hitsArray[i] = new SearchHit(i, String.valueOf(i), + Map.of("title", new DocumentField("title", List.of("doc" + i))), Map.of()); + hitsArray[i].sourceRef(BytesReference.bytes(sourceContent)); + } + return new SearchResponse(new InternalSearchResponse( - new SearchHits(new SearchHit[0], new TotalHits(100, TotalHits.Relation.EQUAL_TO), 1.0f), + new SearchHits(hitsArray, new TotalHits(100, TotalHits.Relation.EQUAL_TO), 1.0f), null, null, null, false, false, 1 ), null, 1, 1, 0, 0, new ShardSearchFailure[0], new SearchResponse.Clusters(1, 1, 0)); } - public void testTransformEnabledByIndexSetting() { + /** + * Should be able to enable transformer on all queries via index setting. + */ + public void testTransformEnabledByIndexSetting() throws IOException { String prefix = "index.plugin.searchrelevance.result_transformer." + ResultTransformerType.KENDRA_INTELLIGENT_RANKING; Settings enablePluginSettings = Settings.builder() @@ -294,7 +351,7 @@ public void testTransformEnabledByIndexSetting() { .setIndices("index") .request(); AtomicBoolean proceedCalled = new AtomicBoolean(false); - SearchResponse searchResponse = buildMockSearchResponse(); + SearchResponse searchResponse = buildMockSearchResponse(randomInt(20)); ActionFilterChain searchFilterChain = (task1, action, request, listener) -> { @@ -303,7 +360,7 @@ public void testTransformEnabledByIndexSetting() { }; AtomicBoolean onResponseCalled = new AtomicBoolean(false); AtomicBoolean onFailureCalled = new AtomicBoolean(false); - ActionListener downstreamListener = new ActionListener() { + ActionListener downstreamListener = new ActionListener<>() { @Override public void onResponse(SearchResponse searchResponse) { onResponseCalled.set(true); @@ -318,9 +375,163 @@ public void onFailure(Exception e) { assertTrue(proceedCalled.get()); // We should NOT try to check for index-level settings, because we saw request-level settings assertTrue(mockTransformer.getTransformerSettingsWasCalled); + assertTrue(mockTransformer.preproccessRequestWasCalled); + assertTrue(mockTransformer.transformWasCalled); + assertTrue(mockTransformer.shouldTransformWasCalled); + assertTrue(onResponseCalled.get()); + assertFalse(onFailureCalled.get()); + } + + /** + * Verify that even if the transformer overrides source, from, and fetchSource, the original values get applied + * in the end. + */ + public void testOutputUsesOriginalSourceParameters() throws IOException { + Client client = buildMockClient("index"); + OpenSearchClient openSearchClient = new OpenSearchClient(client); + + MockTransformer mockTransformer = new MockTransformer(request -> { + // Modify the request to always fetch source + request results 0-50 + request.source() + .from(0) + .size(50) + .fetchSource(true); + }); + + Map transformerMap = + Map.of(ResultTransformerType.KENDRA_INTELLIGENT_RANKING, mockTransformer); + + SearchActionFilter searchActionFilter = new SearchActionFilter(transformerMap, openSearchClient); + + Task task = Mockito.mock(Task.class); + SearchRequest searchRequest = new SearchRequestBuilder(null, SearchAction.INSTANCE) + .setSource( + new SearchSourceBuilder() + .from(10) + .size(10) + .fetchSource(false) + .ext( + Collections.singletonList(new SearchConfigurationExtBuilder() + .setResultTransformers( + Collections.singletonList(MOCK_TRANSFORMER_CONFIGURATION) + ) + ) + ) + ).setIndices("index") + .request(); + AtomicBoolean proceedCalled = new AtomicBoolean(false); + SearchResponse searchResponse = buildMockSearchResponse(50); + + ActionFilterChain searchFilterChain = + (task1, action, request, listener) -> { + proceedCalled.set(true); + listener.onResponse(searchResponse); + }; + AtomicBoolean onResponseCalled = new AtomicBoolean(false); + AtomicBoolean onFailureCalled = new AtomicBoolean(false); + AtomicReference returnedResponse = new AtomicReference<>(); + ActionListener downstreamListener = new ActionListener<>() { + @Override + public void onResponse(SearchResponse searchResponse) { + returnedResponse.set(searchResponse); + onResponseCalled.set(true); + } + + @Override + public void onFailure(Exception e) { + onFailureCalled.set(true); + } + }; + searchActionFilter.apply(task, SearchAction.NAME, searchRequest, downstreamListener, searchFilterChain); + assertTrue(proceedCalled.get()); + // We should NOT try to check for index-level settings, because we saw request-level settings + assertFalse(mockTransformer.getTransformerSettingsWasCalled); + assertTrue(mockTransformer.preproccessRequestWasCalled); assertTrue(mockTransformer.transformWasCalled); assertTrue(mockTransformer.shouldTransformWasCalled); assertTrue(onResponseCalled.get()); assertFalse(onFailureCalled.get()); + + assertNotNull(returnedResponse.get()); + SearchResponse response = returnedResponse.get(); + assertEquals(10, response.getHits().getHits().length); + for (int i = 0; i < 10; i++) { + assertEquals("doc" + (10 + i), response.getHits().getHits()[i].field("title").getValue()); + assertFalse(response.getHits().getHits()[i].hasSource()); + } + } + + /** + * Check that we handle the case where the transformer returns top N, but the "from" starts after that. + */ + public void testReturnEmptyWhenOriginalFromExceedsHitCount() throws IOException { + Client client = buildMockClient("index"); + OpenSearchClient openSearchClient = new OpenSearchClient(client); + + MockTransformer mockTransformer = new MockTransformer(request -> { + // Modify the request to always fetch source + request results 0-50 + request.source() + .from(0) + .size(50) + .fetchSource(true); + }); + + Map transformerMap = + Map.of(ResultTransformerType.KENDRA_INTELLIGENT_RANKING, mockTransformer); + + SearchActionFilter searchActionFilter = new SearchActionFilter(transformerMap, openSearchClient); + + Task task = Mockito.mock(Task.class); + SearchRequest searchRequest = new SearchRequestBuilder(null, SearchAction.INSTANCE) + .setSource( + new SearchSourceBuilder() + .from(50) + .size(10) + .fetchSource(false) + .ext( + Collections.singletonList(new SearchConfigurationExtBuilder() + .setResultTransformers( + Collections.singletonList(MOCK_TRANSFORMER_CONFIGURATION) + ) + ) + ) + ).setIndices("index") + .request(); + AtomicBoolean proceedCalled = new AtomicBoolean(false); + SearchResponse searchResponse = buildMockSearchResponse(50); + + ActionFilterChain searchFilterChain = + (task1, action, request, listener) -> { + proceedCalled.set(true); + listener.onResponse(searchResponse); + }; + AtomicBoolean onResponseCalled = new AtomicBoolean(false); + AtomicBoolean onFailureCalled = new AtomicBoolean(false); + AtomicReference returnedResponse = new AtomicReference<>(); + ActionListener downstreamListener = new ActionListener<>() { + @Override + public void onResponse(SearchResponse searchResponse) { + returnedResponse.set(searchResponse); + onResponseCalled.set(true); + } + + @Override + public void onFailure(Exception e) { + onFailureCalled.set(true); + } + }; + searchActionFilter.apply(task, SearchAction.NAME, searchRequest, downstreamListener, searchFilterChain); + assertTrue(proceedCalled.get()); + // We should NOT try to check for index-level settings, because we saw request-level settings + assertFalse(mockTransformer.getTransformerSettingsWasCalled); + assertTrue(mockTransformer.preproccessRequestWasCalled); + assertTrue(mockTransformer.transformWasCalled); + assertTrue(mockTransformer.shouldTransformWasCalled); + assertTrue(onResponseCalled.get()); + assertFalse(onFailureCalled.get()); + + assertNotNull(returnedResponse.get()); + SearchResponse response = returnedResponse.get(); + assertEquals(0, response.getHits().getHits().length); } } \ No newline at end of file diff --git a/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java b/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java new file mode 100644 index 0000000..d975f07 --- /dev/null +++ b/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java @@ -0,0 +1,212 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.search.relevance.transformer.kendraintelligentranking; + +import org.apache.lucene.search.TotalHits; +import org.mockito.Mockito; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.common.bytes.BytesReference; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.xcontent.XContentBuilder; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.index.query.MatchAllQueryBuilder; +import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.relevance.configuration.ResultTransformerConfiguration; +import org.opensearch.search.relevance.transformer.kendraintelligentranking.client.KendraHttpClient; +import org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankingConfiguration; +import org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankingConfiguration.KendraIntelligentRankingProperties; +import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.dto.RescoreRequest; +import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.dto.RescoreResult; +import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.dto.RescoreResultItem; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.stream.Collectors; + +public class KendraIntelligentRankerTests extends OpenSearchTestCase { + private static KendraHttpClient buildMockHttpClient(Function mockRescoreImpl) { + KendraHttpClient kendraHttpClient = Mockito.mock(KendraHttpClient.class); + Mockito.doAnswer(invocation -> { + RescoreRequest rescoreRequest = invocation.getArgument(0); + return mockRescoreImpl.apply(rescoreRequest); + }).when(kendraHttpClient).rescore(Mockito.any(RescoreRequest.class)); + return kendraHttpClient; + } + + private static KendraHttpClient buildMockHttpClient() { + return buildMockHttpClient(r -> new RescoreResult()); + } + + public void testGetSettings() { + List> settings = new KendraIntelligentRanker(buildMockHttpClient()).getTransformerSettings(); + assertNotNull(settings); + assertFalse(settings.isEmpty()); + } + + public void testPreprocess() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + KendraIntelligentRankingProperties properties = + new KendraIntelligentRankingProperties(List.of("body"), List.of("title"), 50); + ResultTransformerConfiguration configuration = new KendraIntelligentRankingConfiguration(1, properties); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder() + .fetchSource(false) + .from(5) + .size(10)); + SearchRequest transformedRequest = ranker.preprocessRequest(originalRequest, configuration); + assertTrue(transformedRequest.source().fetchSource().fetchSource()); + assertEquals(0, transformedRequest.source().from()); + assertEquals(50, transformedRequest.source().size()); + } + + public void testShouldNotTransformWithoutSource() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + SearchRequest originalRequest = new SearchRequest(); + boolean shouldTransform = ranker.shouldTransform(originalRequest, new KendraIntelligentRankingConfiguration()); + assertFalse(shouldTransform); + } + + public void testShouldNotTransformWithoutQuery() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder().query(null)); + boolean shouldTransform = ranker.shouldTransform(originalRequest, new KendraIntelligentRankingConfiguration()); + assertFalse(shouldTransform); + } + + public void testShouldNotTransformWithScroll() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder()) + .scroll("5h"); + boolean shouldTransform = ranker.shouldTransform(originalRequest, new KendraIntelligentRankingConfiguration()); + assertFalse(shouldTransform); + } + + public void testShouldNotTransformWithSort() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder() + .sort("foo")); + boolean shouldTransform = ranker.shouldTransform(originalRequest, new KendraIntelligentRankingConfiguration()); + assertFalse(shouldTransform); + } + + public void testShouldNotTransformIfFromExceedsDocLimit() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder() + .from(20)); + KendraIntelligentRankingProperties properties = + new KendraIntelligentRankingProperties(List.of("body"), List.of("title"), 10); + ResultTransformerConfiguration configuration = new KendraIntelligentRankingConfiguration(1, properties); + boolean shouldTransform = ranker.shouldTransform(originalRequest, configuration); + assertFalse(shouldTransform); + } + + public void testShouldTransformTrue() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder() + .query(new MatchAllQueryBuilder())); + KendraIntelligentRankingProperties properties = + new KendraIntelligentRankingProperties(List.of("body"), List.of("title"), 10); + ResultTransformerConfiguration configuration = new KendraIntelligentRankingConfiguration(1, properties); + boolean shouldTransform = ranker.shouldTransform(originalRequest, configuration); + assertTrue(shouldTransform); + } + + public void testTransformInvalidQueryType() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder().query(new MatchAllQueryBuilder())); + KendraIntelligentRankingProperties properties = + new KendraIntelligentRankingProperties(List.of("body"), List.of("title"), 10); + ResultTransformerConfiguration configuration = new KendraIntelligentRankingConfiguration(1, properties); + + SearchHits searchHits = new SearchHits(new SearchHit[0], new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1.0f); + + SearchHits transformedHits = ranker.transform(searchHits, originalRequest, configuration); + assertSame(searchHits, transformedHits); + } + + public void testTransformEmptyHits() { + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder().query(new MatchQueryBuilder("body", "foo"))); + KendraIntelligentRankingProperties properties = + new KendraIntelligentRankingProperties(List.of("body"), List.of("title"), 10); + ResultTransformerConfiguration configuration = new KendraIntelligentRankingConfiguration(1, properties); + + SearchHits searchHits = new SearchHits(new SearchHit[0], new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1.0f); + + SearchHits transformedHits = ranker.transform(searchHits, originalRequest, configuration); + assertSame(searchHits, transformedHits); + } + + public void testTransformHits() throws IOException { + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder().query(new MatchQueryBuilder("body", "foo"))); + + int docLimit = randomIntBetween(1, 20); + KendraIntelligentRankingProperties properties = + new KendraIntelligentRankingProperties(List.of("body"), List.of("title"), docLimit); + ResultTransformerConfiguration configuration = new KendraIntelligentRankingConfiguration(1, properties); + + int numHits = docLimit + randomInt(20); + SearchHit[] hitsArray = new SearchHit[numHits]; + for (int i = 0; i < numHits; i++) { + XContentBuilder sourceContent = JsonXContent.contentBuilder() + .startObject() + .field("_id", String.valueOf(i)) + .field("body", "Body text for document number " + i) + .field("title", "This is the title for document " + i) + .endObject(); + hitsArray[i] = new SearchHit(i, "doc" + i, Map.of(), Map.of()); + hitsArray[i].sourceRef(BytesReference.bytes(sourceContent)); + } + SearchHits searchHits = new SearchHits(hitsArray, new TotalHits(numHits, TotalHits.Relation.EQUAL_TO), 1.0f); + + AtomicReference rescoreRequestRef = new AtomicReference<>(); + KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient(req -> { + rescoreRequestRef.set(req); + // Return the top N results in reverse order. + List resultItems = req.getDocuments().stream() + .map(d -> { + RescoreResultItem item = new RescoreResultItem(); + item.setDocumentId(d.getGroupId()); + item.setScore(randomFloat()); + return item; + }).collect(Collectors.toList()); + Collections.reverse(resultItems); + RescoreResult result = new RescoreResult(); + result.setResultItems(resultItems); + return result; + })); + SearchHits transformedHits = ranker.transform(searchHits, originalRequest, configuration); + + assertNotSame(searchHits, transformedHits); + // The top N (according to doc limit) should be in reverse order + for (int i = 0; i < docLimit; i++) { + assertEquals("doc" + (docLimit - i - 1), transformedHits.getHits()[i].getId()); + } + // The remainder should be in the original order + for (int i = docLimit; i < numHits; i++) { + assertEquals("doc" + i, transformedHits.getHits()[i].getId()); + } + } + +} \ No newline at end of file