From c6c9c1220bfaced172bc55853947b1c1c9ac3206 Mon Sep 17 00:00:00 2001 From: Mahita Mahesh Date: Wed, 2 Nov 2022 22:29:57 -0400 Subject: [PATCH 1/7] Preprocess search request to send more search hits to Kendra Ranking than the default 10. * Add preprocess request interface to ResultTransformer * Override from and size parameters durig preprocess, to fetch more hits from OpenSearch * Limit the number of documents sent to Kendra Ranking (Will make this number configurable in a follow up) * Populate response documents from re-ranked hits (upto Kendra limit) and OpenSearch hits beyond the limit Signed-off-by: Mahita Mahesh Signed-off-by: Michael Froh --- .../actionfilter/SearchActionFilter.java | 407 +++++++++--------- .../transformer/ResultTransformer.java | 17 +- .../KendraIntelligentRanker.java | 90 +++- .../KendraIntelligentRankingException.java | 30 ++ .../preprocess/QueryParser.java | 1 + 5 files changed, 322 insertions(+), 223 deletions(-) create mode 100644 src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/KendraIntelligentRankingException.java 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..5336695 100644 --- a/src/main/java/org/opensearch/search/relevance/actionfilter/SearchActionFilter.java +++ b/src/main/java/org/opensearch/search/relevance/actionfilter/SearchActionFilter.java @@ -50,217 +50,222 @@ 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) { + + final long startTime = System.nanoTime(); + + if (!SearchAction.INSTANCE.name().equals(action)) { + chain.proceed(task, action, request, listener); + return; + } + + SearchRequest searchRequest = (SearchRequest) request; + // Make a copy of the search request that will not be modified by any transformer + SearchRequest originalSearchRequest = new SearchRequest(searchRequest.indices()); - // 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 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; } - 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); + 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, originalSearchRequest, config); + orderedTransformersAndConfigs.put(resultTransformer, config); } - } - - 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); - } - }; - } + + if (!orderedTransformersAndConfigs.isEmpty()) { + final ActionListener searchResponseListener = createSearchResponseListener( + listener, startTime, orderedTransformersAndConfigs, searchRequest, originalSearchRequest); + 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 originalSearchRequest original request 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 SearchRequest originalSearchRequest) { + 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 (originalSearchRequest.source() != null && + originalSearchRequest.source().fetchSource() != null && + !originalSearchRequest.source().fetchSource().fetchSource()) { + searchHitsList = searchHitsList.stream() + .map(hit -> hit.sourceRef(null)) + .collect(Collectors.toList()); + } + final int lastHitIndex = Math.min(searchHitsList.size(), + (originalSearchRequest.source().from() + originalSearchRequest.source().size())); + searchHitsList = searchHitsList.subList(originalSearchRequest.source().from(), lastHitIndex); + + // TODO: How to handle SearchHits.TotalHits when transformer modifies the hit count + hits = new SearchHits( + searchHitsList.toArray(new SearchHit[searchHitsList.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); + } + } + + 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..49b6b10 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,23 @@ 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 originalSearchRequest Search Request object in which to retain original input parameters + * @param configuration Configuration parameters for the transformer + * @return SearchRequest with updated attributes + */ + SearchRequest preprocessRequest(final SearchRequest request, + final SearchRequest originalSearchRequest, + final ResultTransformerConfiguration configuration); + /** * Rank hits based on the provided query * @param hits hits to be re-ranked @@ -36,5 +47,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..2e24ea2 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,14 +7,21 @@ */ 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.HashSet; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.PriorityQueue; +import java.util.Set; +import java.util.stream.Collectors; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -22,10 +29,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; @@ -46,6 +55,7 @@ public class KendraIntelligentRanker implements ResultTransformer { 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 int KENDRA_DOC_LIMIT = 25; private static final Logger logger = LogManager.getLogger(KendraIntelligentRanker.class); @@ -70,19 +80,42 @@ public List> getTransformerSettings() { * @return boolean decision on whether to re-rank */ @Override - public boolean shouldTransform(SearchRequest request, ResultTransformerConfiguration configuration) { + public boolean shouldTransform(final SearchRequest request, final ResultTransformerConfiguration configuration) { if (request.source() == null) { return false; } - // Skip if there is scroll, pagination, or sorting. - if (request.scroll() != null || request.source().from() > 0 || - (request.source().sorts() != null && !request.source().sorts().isEmpty())) { + // 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() >= KENDRA_DOC_LIMIT) { return false; } return true; } + @Override + public SearchRequest preprocessRequest(final SearchRequest request, + final SearchRequest originalSearchRequest, + 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().fetchSource() != null && !request.source().fetchSource().fetchSource()) { + originalSearchRequest.source().fetchSource(request.source().fetchSource()); + request.source().fetchSource(true); + } + + int from = request.source().from() == -1 ? SearchService.DEFAULT_FROM : request.source().from(); + int size = request.source().size() == -1 ? SearchService.DEFAULT_SIZE : request.source().size(); + originalSearchRequest.source().from(from); + originalSearchRequest.source().size(size); + + int sizeOverride = Math.max(KENDRA_DOC_LIMIT, from + size); + request.source().from(0); + request.source().size(sizeOverride); + return request; + } + /** * * @param hits Search hits to rerank with respect to query @@ -102,49 +135,66 @@ public SearchHits transform(final SearchHits hits, return hits; } try { - List originalHits = new ArrayList<>(); - for (SearchHit searchHit : hits.getHits()) { - Map docSourceMap = searchHit.getSourceAsMap(); + List originalHits = Arrays.asList(hits.getHits()); + final int numberOfHitsToRerank = Math.min(originalHits.size(), KENDRA_DOC_LIMIT); + 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); - List splitPassages = textSplitter.split(docSourceMap.get(queryParserResult.getBodyFieldName()).toString()); + 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 [].", + 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 (queryParserResult.getTitleFieldName() != 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++) { - originalHits.add( - new Document(searchHit.getId() + "@" + (i + 1), searchHit.getId(), tokenizedTitle, topPassages.get(i), searchHit.getScore()) + 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(), originalHits); + final RescoreRequest rescoreRequest = new RescoreRequest(queryParserResult.getQueryText(), originalHitsAsDocuments); 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.", + 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()); - continue; + 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 = KENDRA_DOC_LIMIT; 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 re-rank. Returning original search results without re-ranking.", ex); + logger.error("Failed to rescore. Returning original search results without rescore.", ex); return hits; } } 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/preprocess/QueryParser.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/preprocess/QueryParser.java index 5964038..742c53d 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/preprocess/QueryParser.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/preprocess/QueryParser.java @@ -18,6 +18,7 @@ import org.opensearch.index.query.MultiMatchQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants; +import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.KendraIntelligentRankingException; public class QueryParser { private static final Logger logger = LogManager.getLogger(QueryParser.class); From 6ffb32282974056fcb828d4626652b9f3d77680b Mon Sep 17 00:00:00 2001 From: Mahita Mahesh Date: Mon, 14 Nov 2022 09:23:50 -0500 Subject: [PATCH 2/7] Make Kendra document limit configurable Signed-off-by: Mahita Mahesh --- .../KendraIntelligentRanker.java | 15 ++-- .../configuration/Constants.java | 4 ++ .../KendraIntelligentRankerSettings.java | 68 +++++++++++++++---- ...KendraIntelligentRankingConfiguration.java | 48 +++++++++---- 4 files changed, 99 insertions(+), 36 deletions(-) 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 2e24ea2..38dbde6 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 @@ -14,14 +14,11 @@ import java.util.Collections; import java.util.Comparator; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.PriorityQueue; -import java.util.Set; -import java.util.stream.Collectors; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -55,7 +52,6 @@ public class KendraIntelligentRanker implements ResultTransformer { 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 int KENDRA_DOC_LIMIT = 25; private static final Logger logger = LogManager.getLogger(KendraIntelligentRanker.class); @@ -84,11 +80,12 @@ public boolean shouldTransform(final SearchRequest request, final ResultTransfor if (request.source() == null) { return false; } + 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() >= KENDRA_DOC_LIMIT) { + request.source().from() >= kendraConfiguration.getProperties().getDocLimit()) { return false; } return true; @@ -110,7 +107,8 @@ public SearchRequest preprocessRequest(final SearchRequest request, originalSearchRequest.source().from(from); originalSearchRequest.source().size(size); - int sizeOverride = Math.max(KENDRA_DOC_LIMIT, from + size); + KendraIntelligentRankingConfiguration kendraConfiguration = (KendraIntelligentRankingConfiguration) configuration; + int sizeOverride = Math.max(kendraConfiguration.getProperties().getDocLimit(), from + size); request.source().from(0); request.source().size(sizeOverride); return request; @@ -134,9 +132,10 @@ public SearchHits transform(final SearchHits hits, if (queryParserResult == null) { return hits; } + KendraIntelligentRankingConfiguration kendraConfiguration = (KendraIntelligentRankingConfiguration) configuration; try { List originalHits = Arrays.asList(hits.getHits()); - final int numberOfHitsToRerank = Math.min(originalHits.size(), KENDRA_DOC_LIMIT); + 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) { @@ -189,7 +188,7 @@ public SearchHits transform(final SearchHits hits, newSearchHits.add(searchHit); } // Add remaining hits to response, which are already sorted by OpenSearch score - for (int i = KENDRA_DOC_LIMIT; i < originalHits.size(); ++i) { + 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); 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; + } } } From ea1866101999a9c66c10e2d5ae6d3c99022a21d9 Mon Sep 17 00:00:00 2001 From: Mahita Mahesh Date: Mon, 14 Nov 2022 13:00:32 -0500 Subject: [PATCH 3/7] Use DTO object instead of copy of the SearchRequest Signed-off-by: Mahita Mahesh Signed-off-by: Michael Froh --- .../actionfilter/SearchActionFilter.java | 24 ++++++----- .../relevance/dto/OriginalRequestFields.java | 41 +++++++++++++++++++ .../transformer/ResultTransformer.java | 2 - .../KendraIntelligentRanker.java | 8 +--- .../preprocess/QueryParser.java | 1 - 5 files changed, 56 insertions(+), 20 deletions(-) create mode 100644 src/main/java/org/opensearch/search/relevance/dto/OriginalRequestFields.java 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 5336695..1551e1a 100644 --- a/src/main/java/org/opensearch/search/relevance/actionfilter/SearchActionFilter.java +++ b/src/main/java/org/opensearch/search/relevance/actionfilter/SearchActionFilter.java @@ -44,6 +44,7 @@ import org.opensearch.search.relevance.configuration.ConfigurationUtils; import org.opensearch.search.relevance.client.OpenSearchClient; import org.opensearch.search.relevance.configuration.ResultTransformerConfiguration; +import org.opensearch.search.relevance.dto.OriginalRequestFields; import org.opensearch.search.relevance.transformer.ResultTransformer; import org.opensearch.search.relevance.transformer.ResultTransformerType; import org.opensearch.search.suggest.Suggest; @@ -86,8 +87,10 @@ public void app } SearchRequest searchRequest = (SearchRequest) request; - // Make a copy of the search request that will not be modified by any transformer - SearchRequest originalSearchRequest = new SearchRequest(searchRequest.indices()); + + // TODO: Remove OriginalRequestFields and replace with a deep copy of the SearchRequest object + // once https://github.com/opensearch-project/OpenSearch/issues/869 is implemented + OriginalRequestFields originalRequestFields = new OriginalRequestFields(searchRequest); final String[] indices = searchRequest.indices(); // Skip if no, or more than 1, index is specified. @@ -104,14 +107,14 @@ public void app 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, originalSearchRequest, config); + searchRequest = resultTransformer.preprocessRequest(searchRequest, config); orderedTransformersAndConfigs.put(resultTransformer, config); } } if (!orderedTransformersAndConfigs.isEmpty()) { final ActionListener searchResponseListener = createSearchResponseListener( - listener, startTime, orderedTransformersAndConfigs, searchRequest, originalSearchRequest); + listener, startTime, orderedTransformersAndConfigs, searchRequest, originalRequestFields); chain.proceed(task, action, request, searchResponseListener); return; } @@ -164,7 +167,7 @@ private List getResultTransformerConfigurations( * @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 originalSearchRequest original request without any modifications made by transformers + * @param originalRequestFields original request without any modifications made by transformers * @param OpenSearch response type * @return ActionListener with override for onResponse method */ @@ -173,7 +176,7 @@ private ActionListener createSearchR final long startTime, final LinkedHashMap orderedTransformersAndConfigs, final SearchRequest searchRequest, - final SearchRequest originalSearchRequest) { + final OriginalRequestFields originalRequestFields) { return new ActionListener() { @Override @@ -199,16 +202,15 @@ public void onResponse(final Response response) { } List searchHitsList = Arrays.asList(hits.getHits()); - if (originalSearchRequest.source() != null && - originalSearchRequest.source().fetchSource() != null && - !originalSearchRequest.source().fetchSource().fetchSource()) { + if (originalRequestFields.fetchSourceContext() != null && + !originalRequestFields.fetchSourceContext().fetchSource()) { searchHitsList = searchHitsList.stream() .map(hit -> hit.sourceRef(null)) .collect(Collectors.toList()); } final int lastHitIndex = Math.min(searchHitsList.size(), - (originalSearchRequest.source().from() + originalSearchRequest.source().size())); - searchHitsList = searchHitsList.subList(originalSearchRequest.source().from(), lastHitIndex); + (originalRequestFields.from() + originalRequestFields.size())); + searchHitsList = searchHitsList.subList(originalRequestFields.from(), lastHitIndex); // TODO: How to handle SearchHits.TotalHits when transformer modifies the hit count hits = new SearchHits( diff --git a/src/main/java/org/opensearch/search/relevance/dto/OriginalRequestFields.java b/src/main/java/org/opensearch/search/relevance/dto/OriginalRequestFields.java new file mode 100644 index 0000000..b0f3f0b --- /dev/null +++ b/src/main/java/org/opensearch/search/relevance/dto/OriginalRequestFields.java @@ -0,0 +1,41 @@ +/* + * 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.dto; + +import org.opensearch.action.search.SearchRequest; +import org.opensearch.search.SearchService; +import org.opensearch.search.fetch.subphase.FetchSourceContext; + +/** + * This class is intended to be a temporary measure until https://github.com/opensearch-project/OpenSearch/issues/869 + * is implemented and a deep copy of a SearchRequest can easily be made. Meanwhile, all request fields that are over-written + * in preprocessRequest of any {@link org.opensearch.search.relevance.transformer.ResultTransformer} should be included here + */ +public class OriginalRequestFields { + private int from; + private int size; + private FetchSourceContext fetchSourceContext; + + public OriginalRequestFields(final SearchRequest searchRequest) { + this.from = searchRequest.source().from() == -1 ? SearchService.DEFAULT_FROM : searchRequest.source().from(); + this.size = searchRequest.source().size() == -1 ? SearchService.DEFAULT_SIZE : searchRequest.source().size(); + this.fetchSourceContext = searchRequest.source().fetchSource(); + } + + public int from() { + return this.from; + } + + public int size() { + return this.size; + } + + public FetchSourceContext fetchSourceContext() { + return this.fetchSourceContext; + } +} 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 49b6b10..6568828 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/ResultTransformer.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/ResultTransformer.java @@ -32,12 +32,10 @@ public interface ResultTransformer { /** * Preprocess the incoming Search Request to support the requirements of the transformer * @param request input Search Request - * @param originalSearchRequest Search Request object in which to retain original input parameters * @param configuration Configuration parameters for the transformer * @return SearchRequest with updated attributes */ SearchRequest preprocessRequest(final SearchRequest request, - final SearchRequest originalSearchRequest, 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 38dbde6..e0c743e 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 @@ -93,23 +93,19 @@ public boolean shouldTransform(final SearchRequest request, final ResultTransfor @Override public SearchRequest preprocessRequest(final SearchRequest request, - final SearchRequest originalSearchRequest, 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().fetchSource() != null && !request.source().fetchSource().fetchSource()) { - originalSearchRequest.source().fetchSource(request.source().fetchSource()); request.source().fetchSource(true); } int from = request.source().from() == -1 ? SearchService.DEFAULT_FROM : request.source().from(); int size = request.source().size() == -1 ? SearchService.DEFAULT_SIZE : request.source().size(); - originalSearchRequest.source().from(from); - originalSearchRequest.source().size(size); KendraIntelligentRankingConfiguration kendraConfiguration = (KendraIntelligentRankingConfiguration) configuration; int sizeOverride = Math.max(kendraConfiguration.getProperties().getDocLimit(), from + size); - request.source().from(0); + request.source().from(SearchService.DEFAULT_FROM); request.source().size(sizeOverride); return request; } @@ -145,7 +141,7 @@ public SearchHits transform(final SearchHits hits, 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 [].", + "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); diff --git a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/preprocess/QueryParser.java b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/preprocess/QueryParser.java index 742c53d..5964038 100644 --- a/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/preprocess/QueryParser.java +++ b/src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/preprocess/QueryParser.java @@ -18,7 +18,6 @@ import org.opensearch.index.query.MultiMatchQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants; -import org.opensearch.search.relevance.transformer.kendraintelligentranking.model.KendraIntelligentRankingException; public class QueryParser { private static final Logger logger = LogManager.getLogger(QueryParser.class); From 37c90ea7050b699fbb2e16d9a1803f2a4e8803cb Mon Sep 17 00:00:00 2001 From: Michael Froh Date: Wed, 23 Nov 2022 01:08:12 +0000 Subject: [PATCH 4/7] Add unit tests Add unit tests for the pagination logic in SearchActionFilter as well as the whole KendraIntelligentRanker class (including pagination). Also, removed OriginalRequestFields in favor of cloning the SearchSourceBuilder. Signed-off-by: Michael Froh --- .../actionfilter/SearchActionFilter.java | 48 ++-- .../relevance/dto/OriginalRequestFields.java | 41 ---- .../KendraIntelligentRanker.java | 10 +- .../model/dto/RescoreResult.java | 17 ++ .../model/dto/RescoreResultItem.java | 16 ++ .../actionfilter/SearchActionFilterTests.java | 227 +++++++++++++++++- .../KendraIntelligentRankerTests.java | 210 ++++++++++++++++ 7 files changed, 503 insertions(+), 66 deletions(-) delete mode 100644 src/main/java/org/opensearch/search/relevance/dto/OriginalRequestFields.java create mode 100644 src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java 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 1551e1a..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,12 +39,13 @@ 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; import org.opensearch.search.relevance.client.OpenSearchClient; import org.opensearch.search.relevance.configuration.ResultTransformerConfiguration; -import org.opensearch.search.relevance.dto.OriginalRequestFields; import org.opensearch.search.relevance.transformer.ResultTransformer; import org.opensearch.search.relevance.transformer.ResultTransformerType; import org.opensearch.search.suggest.Suggest; @@ -88,9 +89,18 @@ public void app SearchRequest searchRequest = (SearchRequest) request; - // TODO: Remove OriginalRequestFields and replace with a deep copy of the SearchRequest object + // TODO: Remove originalSearchSource and replace with a deep copy of the SearchRequest object // once https://github.com/opensearch-project/OpenSearch/issues/869 is implemented - OriginalRequestFields originalRequestFields = new OriginalRequestFields(searchRequest); + 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 String[] indices = searchRequest.indices(); // Skip if no, or more than 1, index is specified. @@ -114,7 +124,7 @@ public void app if (!orderedTransformersAndConfigs.isEmpty()) { final ActionListener searchResponseListener = createSearchResponseListener( - listener, startTime, orderedTransformersAndConfigs, searchRequest, originalRequestFields); + listener, startTime, orderedTransformersAndConfigs, searchRequest, originalSearchSource); chain.proceed(task, action, request, searchResponseListener); return; } @@ -167,7 +177,7 @@ private List getResultTransformerConfigurations( * @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 originalRequestFields original request without any modifications made by transformers + * @param originalSearchSource original search source without any modifications made by transformers * @param OpenSearch response type * @return ActionListener with override for onResponse method */ @@ -176,7 +186,7 @@ private ActionListener createSearchR final long startTime, final LinkedHashMap orderedTransformersAndConfigs, final SearchRequest searchRequest, - final OriginalRequestFields originalRequestFields) { + final SearchSourceBuilder originalSearchSource) { return new ActionListener() { @Override @@ -202,19 +212,27 @@ public void onResponse(final Response response) { } List searchHitsList = Arrays.asList(hits.getHits()); - if (originalRequestFields.fetchSourceContext() != null && - !originalRequestFields.fetchSourceContext().fetchSource()) { - searchHitsList = searchHitsList.stream() - .map(hit -> hit.sourceRef(null)) - .collect(Collectors.toList()); + 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); + } + } } - final int lastHitIndex = Math.min(searchHitsList.size(), - (originalRequestFields.from() + originalRequestFields.size())); - searchHitsList = searchHitsList.subList(originalRequestFields.from(), lastHitIndex); // TODO: How to handle SearchHits.TotalHits when transformer modifies the hit count hits = new SearchHits( - searchHitsList.toArray(new SearchHit[searchHitsList.size()]), + searchHitsList.toArray(new SearchHit[0]), hits.getTotalHits(), hits.getMaxScore()); diff --git a/src/main/java/org/opensearch/search/relevance/dto/OriginalRequestFields.java b/src/main/java/org/opensearch/search/relevance/dto/OriginalRequestFields.java deleted file mode 100644 index b0f3f0b..0000000 --- a/src/main/java/org/opensearch/search/relevance/dto/OriginalRequestFields.java +++ /dev/null @@ -1,41 +0,0 @@ -/* - * 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.dto; - -import org.opensearch.action.search.SearchRequest; -import org.opensearch.search.SearchService; -import org.opensearch.search.fetch.subphase.FetchSourceContext; - -/** - * This class is intended to be a temporary measure until https://github.com/opensearch-project/OpenSearch/issues/869 - * is implemented and a deep copy of a SearchRequest can easily be made. Meanwhile, all request fields that are over-written - * in preprocessRequest of any {@link org.opensearch.search.relevance.transformer.ResultTransformer} should be included here - */ -public class OriginalRequestFields { - private int from; - private int size; - private FetchSourceContext fetchSourceContext; - - public OriginalRequestFields(final SearchRequest searchRequest) { - this.from = searchRequest.source().from() == -1 ? SearchService.DEFAULT_FROM : searchRequest.source().from(); - this.size = searchRequest.source().size() == -1 ? SearchService.DEFAULT_SIZE : searchRequest.source().size(); - this.fetchSourceContext = searchRequest.source().fetchSource(); - } - - public int from() { - return this.from; - } - - public int size() { - return this.size; - } - - public FetchSourceContext fetchSourceContext() { - return this.fetchSourceContext; - } -} 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 e0c743e..559bf47 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 @@ -77,7 +77,7 @@ public List> getTransformerSettings() { */ @Override public boolean shouldTransform(final SearchRequest request, final ResultTransformerConfiguration configuration) { - if (request.source() == null) { + if (request.source() == null || request.source().query() == null) { return false; } KendraIntelligentRankingConfiguration kendraConfiguration = (KendraIntelligentRankingConfiguration) configuration; @@ -96,7 +96,8 @@ 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().fetchSource() != null && !request.source().fetchSource().fetchSource()) { + if (request.source() != null && request.source().fetchSource() != null && + !request.source().fetchSource().fetchSource()) { request.source().fetchSource(true); } @@ -126,6 +127,11 @@ public SearchHits transform(final SearchHits hits, kendraConfig.getProperties().getBodyFields(), kendraConfig.getProperties().getTitleFields()); if (queryParserResult == null) { + // Unknown query type or query does not reference body field + return hits; + } + if (hits.getHits().length == 0) { + // Avoid call to rerank empty results return hits; } KendraIntelligentRankingConfiguration kendraConfiguration = (KendraIntelligentRankingConfiguration) configuration; 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/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..51efe0f --- /dev/null +++ b/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java @@ -0,0 +1,210 @@ +/* + * 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(); + 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 = randomInt(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 From c3b9a8134eb8a050682c8046b73a87714baa3860 Mon Sep 17 00:00:00 2001 From: Michael Froh Date: Wed, 23 Nov 2022 02:07:55 +0000 Subject: [PATCH 5/7] Update integration test Updating integration test to match approach used by other recent plugins, in the hope of fixing test runs on Windows. Signed-off-by: Michael Froh --- .../relevance/SearchRelevancePluginIT.java | 22 +++++-------------- 1 file changed, 5 insertions(+), 17 deletions(-) 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")); } } From 87f344a0c487bc35a44298a5d83e21f089a884ef Mon Sep 17 00:00:00 2001 From: Michael Froh Date: Wed, 23 Nov 2022 07:08:39 +0000 Subject: [PATCH 6/7] Address PR feedback * Fix test case with non-null source but null query. * Skip transform on empty results before doing any work. Signed-off-by: Michael Froh --- .../KendraIntelligentRanker.java | 328 +++++++++--------- .../KendraIntelligentRankerTests.java | 3 +- 2 files changed, 166 insertions(+), 165 deletions(-) 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 559bf47..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 @@ -46,180 +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(final SearchRequest request, final ResultTransformerConfiguration configuration) { - if (request.source() == null || request.source().query() == 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); + + 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(); } - 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; - } - 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); + @Override + public List> getTransformerSettings() { + return KendraIntelligentRankerSettings.getAllSettings(); } - 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; - } - - /** - * - * @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) { - // Unknown query type or query does not reference body field - return hits; - } - if (hits.getHits().length == 0) { - // Avoid call to rerank empty results - return hits; + /** + * 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; + } + 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; + } + return true; } - 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); + + @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); } - 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; - } + + 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; + } + + /** + * @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; } - 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()) - ); + 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; } - // 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); + 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; } - 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; - } - } - - 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())); + 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/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java b/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java index 51efe0f..ece1589 100644 --- a/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java +++ b/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java @@ -81,7 +81,8 @@ public void testShouldNotTransformWithoutSource() { public void testShouldNotTransformWithoutQuery() { KendraIntelligentRanker ranker = new KendraIntelligentRanker(buildMockHttpClient()); - SearchRequest originalRequest = new SearchRequest(); + SearchRequest originalRequest = new SearchRequest() + .source(new SearchSourceBuilder().query(null)); boolean shouldTransform = ranker.shouldTransform(originalRequest, new KendraIntelligentRankingConfiguration()); assertFalse(shouldTransform); } From 989b22f57d0180835c2a673586a04466a3a133d5 Mon Sep 17 00:00:00 2001 From: Michael Froh Date: Wed, 23 Nov 2022 09:57:14 +0000 Subject: [PATCH 7/7] Fix subtle unit test failure * Earlier test covered empty hits, so we should assume that Kendra ranking handles at least 1 hit. Signed-off-by: Michael Froh --- .../kendraintelligentranking/KendraIntelligentRankerTests.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 index ece1589..d975f07 100644 --- a/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java +++ b/src/test/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/KendraIntelligentRankerTests.java @@ -161,7 +161,7 @@ public void testTransformHits() throws IOException { SearchRequest originalRequest = new SearchRequest() .source(new SearchSourceBuilder().query(new MatchQueryBuilder("body", "foo"))); - int docLimit = randomInt(20); + int docLimit = randomIntBetween(1, 20); KendraIntelligentRankingProperties properties = new KendraIntelligentRankingProperties(List.of("body"), List.of("title"), docLimit); ResultTransformerConfiguration configuration = new KendraIntelligentRankingConfiguration(1, properties); @@ -197,6 +197,7 @@ public void testTransformHits() throws IOException { 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++) {