Skip to content

Commit

Permalink
Consume retrievers during the query rewrite
Browse files Browse the repository at this point in the history
  • Loading branch information
@jimczi
jimczi committed Jul 4, 2024
1 parent 7aa49d1 commit 83a0fe4
Show file tree
Hide file tree
Showing 5 changed files with 240 additions and 176 deletions.
131 changes: 4 additions & 127 deletions server/src/main/java/org/elasticsearch/action/search/SearchRequest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,20 +17,16 @@
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.util.CollectionUtils;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.mapper.SourceLoader;
import org.elasticsearch.index.query.QueryRewriteContext;
import org.elasticsearch.index.query.Rewriteable;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchService;
import org.elasticsearch.search.builder.PointInTimeBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.rescore.RescorerBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.ShardDocSortField;
import org.elasticsearch.search.sort.SortBuilder;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.tasks.TaskId;
Expand Down Expand Up @@ -324,124 +320,15 @@ public void writeTo(StreamOutput out) throws IOException {
public ActionRequestValidationException validate() {
ActionRequestValidationException validationException = null;
boolean scroll = scroll() != null;

if (source != null) {
validationException = source.validate(validationException, scroll);
}
if (scroll) {
if (source != null) {
if (source.trackTotalHitsUpTo() != null && source.trackTotalHitsUpTo() != SearchContext.TRACK_TOTAL_HITS_ACCURATE) {
validationException = addValidationError(
"disabling [track_total_hits] is not allowed in a scroll context",
validationException
);
}
if (source.from() > 0) {
validationException = addValidationError("using [from] is not allowed in a scroll context", validationException);
}
if (source.size() == 0) {
validationException = addValidationError("[size] cannot be [0] in a scroll context", validationException);
}
if (source.rescores() != null && source.rescores().isEmpty() == false) {
validationException = addValidationError("using [rescore] is not allowed in a scroll context", validationException);
}
if (CollectionUtils.isEmpty(source.searchAfter()) == false) {
validationException = addValidationError("[search_after] cannot be used in a scroll context", validationException);
}
if (source.collapse() != null) {
validationException = addValidationError("cannot use `collapse` in a scroll context", validationException);
}
}
if (requestCache != null && requestCache) {
validationException = addValidationError("[request_cache] cannot be used in a scroll context", validationException);
}
}
if (source != null) {
if (source.slice() != null) {
if (source.pointInTimeBuilder() == null && (scroll == false)) {
validationException = addValidationError(
"[slice] can only be used with [scroll] or [point-in-time] requests",
validationException
);
}
}
if (source.from() > 0 && CollectionUtils.isEmpty(source.searchAfter()) == false) {
validationException = addValidationError(
"[from] parameter must be set to 0 when [search_after] is used",
validationException
);
}
if (source.storedFields() != null) {
if (source.storedFields().fetchFields() == false) {
if (source.fetchSource() != null && source.fetchSource().fetchSource()) {
validationException = addValidationError(
"[stored_fields] cannot be disabled if [_source] is requested",
validationException
);
}
if (source.fetchFields() != null) {
validationException = addValidationError(
"[stored_fields] cannot be disabled when using the [fields] option",
validationException
);
}

}
}
if (source.subSearches().size() >= 2 && source.rankBuilder() == null) {
validationException = addValidationError("[sub_searches] requires [rank]", validationException);
}
if (source.aggregations() != null) {
validationException = source.aggregations().validate(validationException);
}
if (source.rankBuilder() != null) {
int size = source.size() == -1 ? SearchService.DEFAULT_SIZE : source.size();
if (size == 0) {
validationException = addValidationError("[rank] requires [size] greater than [0]", validationException);
}
if (size > source.rankBuilder().rankWindowSize()) {
validationException = addValidationError(
"[rank] requires [rank_window_size: "
+ source.rankBuilder().rankWindowSize()
+ "]"
+ " be greater than or equal to [size: "
+ size
+ "]",
validationException
);
}
int queryCount = source.subSearches().size() + source.knnSearch().size();
if (source.rankBuilder().isCompoundBuilder() && queryCount < 2) {
validationException = addValidationError(
"[rank] requires a minimum of [2] result sets using a combination of sub searches and/or knn searches",
validationException
);
}
if (scroll) {
validationException = addValidationError("[rank] cannot be used in a scroll context", validationException);
}
if (source.rescores() != null && source.rescores().isEmpty() == false) {
validationException = addValidationError("[rank] cannot be used with [rescore]", validationException);
}
if (source.sorts() != null && source.sorts().isEmpty() == false) {
validationException = addValidationError("[rank] cannot be used with [sort]", validationException);
}
if (source.collapse() != null) {
validationException = addValidationError("[rank] cannot be used with [collapse]", validationException);
}
if (source.suggest() != null && source.suggest().getSuggestions().isEmpty() == false) {
validationException = addValidationError("[rank] cannot be used with [suggest]", validationException);
}
if (source.highlighter() != null) {
validationException = addValidationError("[rank] cannot be used with [highlighter]", validationException);
}
if (source.pointInTimeBuilder() != null) {
validationException = addValidationError("[rank] cannot be used with [point in time]", validationException);
}
}
if (source.rescores() != null) {
for (@SuppressWarnings("rawtypes")
RescorerBuilder rescoreBuilder : source.rescores()) {
validationException = rescoreBuilder.validate(this, validationException);
}
}
}
if (pointInTimeBuilder() != null) {
if (scroll) {
validationException = addValidationError("using [point in time] is not allowed in a scroll context", validationException);
Expand All @@ -461,16 +348,6 @@ public ActionRequestValidationException validate() {
if (preference() != null) {
validationException = addValidationError("[preference] cannot be used with point in time", validationException);
}
} else if (source != null && source.sorts() != null) {
for (SortBuilder<?> sortBuilder : source.sorts()) {
if (sortBuilder instanceof FieldSortBuilder
&& ShardDocSortField.NAME.equals(((FieldSortBuilder) sortBuilder).getFieldName())) {
validationException = addValidationError(
"[" + FieldSortBuilder.SHARD_DOC_FIELD_NAME + "] sort field cannot be used without [point in time]",
validationException
);
}
}
}
if (minCompatibleShardNode() != null) {
if (isCcsMinimizeRoundtrips()) {
Expand Down
Loading

0 comments on commit 83a0fe4

Please sign in to comment.