-
Notifications
You must be signed in to change notification settings - Fork 25.2k
Add the ability to partition a scroll in multiple slices. #18237
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -42,6 +42,7 @@ | |
import org.elasticsearch.index.query.QueryShardContext; | ||
import org.elasticsearch.script.Script; | ||
import org.elasticsearch.search.aggregations.AggregationBuilder; | ||
import org.elasticsearch.search.slice.SliceBuilder; | ||
import org.elasticsearch.search.aggregations.AggregatorFactories; | ||
import org.elasticsearch.search.aggregations.AggregatorParsers; | ||
import org.elasticsearch.search.aggregations.PipelineAggregatorBuilder; | ||
|
@@ -98,6 +99,7 @@ public final class SearchSourceBuilder extends ToXContentToBytes implements Writ | |
public static final ParseField EXT_FIELD = new ParseField("ext"); | ||
public static final ParseField PROFILE_FIELD = new ParseField("profile"); | ||
public static final ParseField SEARCH_AFTER = new ParseField("search_after"); | ||
public static final ParseField SLICE = new ParseField("slice"); | ||
|
||
public static SearchSourceBuilder fromXContent(QueryParseContext context, AggregatorParsers aggParsers, | ||
Suggesters suggesters) throws IOException { | ||
|
@@ -138,6 +140,8 @@ public static HighlightBuilder highlight() { | |
|
||
private SearchAfterBuilder searchAfterBuilder; | ||
|
||
private SliceBuilder sliceBuilder; | ||
|
||
private Float minScore; | ||
|
||
private long timeoutInMillis = -1; | ||
|
@@ -175,9 +179,7 @@ public SearchSourceBuilder() { | |
* Read from a stream. | ||
*/ | ||
public SearchSourceBuilder(StreamInput in) throws IOException { | ||
if (in.readBoolean()) { | ||
aggregations = new AggregatorFactories.Builder(in); | ||
} | ||
aggregations = in.readOptionalWriteable(AggregatorFactories.Builder::new); | ||
explain = in.readOptionalBoolean(); | ||
fetchSourceContext = in.readOptionalStreamable(FetchSourceContext::new); | ||
boolean hasFieldDataFields = in.readBoolean(); | ||
|
@@ -206,15 +208,9 @@ public SearchSourceBuilder(StreamInput in) throws IOException { | |
indexBoost.put(in.readString(), in.readFloat()); | ||
} | ||
} | ||
if (in.readBoolean()) { | ||
minScore = in.readFloat(); | ||
} | ||
if (in.readBoolean()) { | ||
postQueryBuilder = in.readNamedWriteable(QueryBuilder.class); | ||
} | ||
if (in.readBoolean()) { | ||
queryBuilder = in.readNamedWriteable(QueryBuilder.class); | ||
} | ||
minScore = in.readOptionalFloat(); | ||
postQueryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class); | ||
queryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class); | ||
if (in.readBoolean()) { | ||
int size = in.readVInt(); | ||
rescoreBuilders = new ArrayList<>(); | ||
|
@@ -244,29 +240,20 @@ public SearchSourceBuilder(StreamInput in) throws IOException { | |
stats.add(in.readString()); | ||
} | ||
} | ||
if (in.readBoolean()) { | ||
suggestBuilder = new SuggestBuilder(in); | ||
} | ||
suggestBuilder = in.readOptionalWriteable(SuggestBuilder::new); | ||
terminateAfter = in.readVInt(); | ||
timeoutInMillis = in.readLong(); | ||
trackScores = in.readBoolean(); | ||
version = in.readOptionalBoolean(); | ||
if (in.readBoolean()) { | ||
ext = in.readBytesReference(); | ||
} | ||
ext = in.readOptionalBytesReference(); | ||
profile = in.readBoolean(); | ||
if (in.readBoolean()) { | ||
searchAfterBuilder = new SearchAfterBuilder(in); | ||
} | ||
searchAfterBuilder = in.readOptionalWriteable(SearchAfterBuilder::new); | ||
sliceBuilder = in.readOptionalWriteable(SliceBuilder::new); | ||
} | ||
|
||
@Override | ||
public void writeTo(StreamOutput out) throws IOException { | ||
boolean hasAggregations = aggregations != null; | ||
out.writeBoolean(hasAggregations); | ||
if (hasAggregations) { | ||
aggregations.writeTo(out); | ||
} | ||
out.writeOptionalWriteable(aggregations); | ||
out.writeOptionalBoolean(explain); | ||
out.writeOptionalStreamable(fetchSourceContext); | ||
boolean hasFieldDataFields = fieldDataFields != null; | ||
|
@@ -296,21 +283,9 @@ public void writeTo(StreamOutput out) throws IOException { | |
out.writeFloat(indexBoost.get(key.value)); | ||
} | ||
} | ||
boolean hasMinScore = minScore != null; | ||
out.writeBoolean(hasMinScore); | ||
if (hasMinScore) { | ||
out.writeFloat(minScore); | ||
} | ||
boolean hasPostQuery = postQueryBuilder != null; | ||
out.writeBoolean(hasPostQuery); | ||
if (hasPostQuery) { | ||
out.writeNamedWriteable(postQueryBuilder); | ||
} | ||
boolean hasQuery = queryBuilder != null; | ||
out.writeBoolean(hasQuery); | ||
if (hasQuery) { | ||
out.writeNamedWriteable(queryBuilder); | ||
} | ||
out.writeOptionalFloat(minScore); | ||
out.writeOptionalNamedWriteable(postQueryBuilder); | ||
out.writeOptionalNamedWriteable(queryBuilder); | ||
boolean hasRescoreBuilders = rescoreBuilders != null; | ||
out.writeBoolean(hasRescoreBuilders); | ||
if (hasRescoreBuilders) { | ||
|
@@ -344,26 +319,15 @@ public void writeTo(StreamOutput out) throws IOException { | |
out.writeString(stat); | ||
} | ||
} | ||
boolean hasSuggestBuilder = suggestBuilder != null; | ||
out.writeBoolean(hasSuggestBuilder); | ||
if (hasSuggestBuilder) { | ||
suggestBuilder.writeTo(out); | ||
} | ||
out.writeOptionalWriteable(suggestBuilder); | ||
out.writeVInt(terminateAfter); | ||
out.writeLong(timeoutInMillis); | ||
out.writeBoolean(trackScores); | ||
out.writeOptionalBoolean(version); | ||
boolean hasExt = ext != null; | ||
out.writeBoolean(hasExt); | ||
if (hasExt) { | ||
out.writeBytesReference(ext); | ||
} | ||
out.writeOptionalBytesReference(ext); | ||
out.writeBoolean(profile); | ||
boolean hasSearchAfter = searchAfterBuilder != null; | ||
out.writeBoolean(hasSearchAfter); | ||
if (hasSearchAfter) { | ||
searchAfterBuilder.writeTo(out); | ||
} | ||
out.writeOptionalWriteable(searchAfterBuilder); | ||
out.writeOptionalWriteable(sliceBuilder); | ||
} | ||
|
||
/** | ||
|
@@ -597,6 +561,22 @@ public SearchSourceBuilder searchAfter(Object[] values) { | |
return this; | ||
} | ||
|
||
/** | ||
* Sets a filter that will restrict the search hits, the top hits and the aggregations to a slice of the results | ||
* of the main query. | ||
*/ | ||
public SearchSourceBuilder slice(SliceBuilder builder) { | ||
this.sliceBuilder = builder; | ||
return this; | ||
} | ||
|
||
/** | ||
* Gets the slice used to filter the search hits, the top hits and the aggregations. | ||
*/ | ||
public SliceBuilder slice() { | ||
return sliceBuilder; | ||
} | ||
|
||
/** | ||
* Add an aggregation to perform as part of the search. | ||
*/ | ||
|
@@ -943,6 +923,7 @@ private SearchSourceBuilder shallowCopy(QueryBuilder queryBuilder, QueryBuilder | |
rewrittenBuilder.rescoreBuilders = rescoreBuilders; | ||
rewrittenBuilder.scriptFields = scriptFields; | ||
rewrittenBuilder.searchAfterBuilder = searchAfterBuilder; | ||
rewrittenBuilder.sliceBuilder = sliceBuilder; | ||
rewrittenBuilder.size = size; | ||
rewrittenBuilder.sorts = sorts; | ||
rewrittenBuilder.stats = stats; | ||
|
@@ -1039,6 +1020,8 @@ public void parseXContent(QueryParseContext context, AggregatorParsers aggParser | |
} else if (context.getParseFieldMatcher().match(currentFieldName, EXT_FIELD)) { | ||
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder().copyCurrentStructure(parser); | ||
ext = xContentBuilder.bytes(); | ||
} else if (context.getParseFieldMatcher().match(currentFieldName, SLICE)) { | ||
sliceBuilder = SliceBuilder.fromXContent(context); | ||
} else { | ||
throw new ParsingException(parser.getTokenLocation(), "Unknown key for a " + token + " in [" + currentFieldName + "].", | ||
parser.getTokenLocation()); | ||
|
@@ -1193,6 +1176,10 @@ public void innerToXContent(XContentBuilder builder, Params params) throws IOExc | |
builder.field(SEARCH_AFTER.getPreferredName(), searchAfterBuilder.getSortValues()); | ||
} | ||
|
||
if (sliceBuilder != null) { | ||
builder.field(SLICE.getPreferredName(), sliceBuilder); | ||
} | ||
|
||
if (indexBoost != null) { | ||
builder.startObject(INDICES_BOOST_FIELD.getPreferredName()); | ||
assert !indexBoost.containsKey(null); | ||
|
@@ -1355,7 +1342,7 @@ public boolean equals(Object obj) { | |
public int hashCode() { | ||
return Objects.hash(aggregations, explain, fetchSourceContext, fieldDataFields, fieldNames, from, | ||
highlightBuilder, indexBoost, minScore, postQueryBuilder, queryBuilder, rescoreBuilders, scriptFields, | ||
size, sorts, searchAfterBuilder, stats, suggestBuilder, terminateAfter, timeoutInMillis, trackScores, version, profile); | ||
size, sorts, searchAfterBuilder, sliceBuilder, stats, suggestBuilder, terminateAfter, timeoutInMillis, trackScores, version, profile); | ||
} | ||
|
||
@Override | ||
|
@@ -1383,6 +1370,7 @@ public boolean equals(Object obj) { | |
&& Objects.equals(size, other.size) | ||
&& Objects.equals(sorts, other.sorts) | ||
&& Objects.equals(searchAfterBuilder, other.searchAfterBuilder) | ||
&& Objects.equals(sliceBuilder, other.sliceBuilder) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. it should be used in the hashcode too |
||
&& Objects.equals(stats, other.stats) | ||
&& Objects.equals(suggestBuilder, other.suggestBuilder) | ||
&& Objects.equals(terminateAfter, other.terminateAfter) | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -29,5 +29,4 @@ public class ScrollContext { | |
public float maxScore; | ||
public ScoreDoc lastEmittedDoc; | ||
public Scroll scroll; | ||
|
||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,67 @@ | ||
/* | ||
* Licensed to Elasticsearch under one or more contributor | ||
* license agreements. See the NOTICE file distributed with | ||
* this work for additional information regarding copyright | ||
* ownership. Elasticsearch licenses this file to you under | ||
* the Apache License, Version 2.0 (the "License"); you may | ||
* not use this file except in compliance with the License. | ||
* You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, | ||
* software distributed under the License is distributed on an | ||
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
* KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations | ||
* under the License. | ||
*/ | ||
|
||
package org.elasticsearch.search.slice; | ||
|
||
import org.apache.lucene.index.LeafReaderContext; | ||
import org.apache.lucene.index.DocValues; | ||
import org.apache.lucene.index.SortedNumericDocValues; | ||
import org.apache.lucene.search.IndexSearcher; | ||
import org.apache.lucene.search.Weight; | ||
import org.apache.lucene.search.RandomAccessWeight; | ||
import org.apache.lucene.util.Bits; | ||
|
||
import java.io.IOException; | ||
|
||
/** | ||
* A {@link SliceQuery} that uses the numeric doc values of a field to do the slicing. | ||
* | ||
* <b>NOTE</b>: With deterministic field values this query can be used across different readers safely. | ||
* If updates are accepted on the field you must ensure that the same reader is used for all `slice` queries. | ||
*/ | ||
public final class DocValuesSliceQuery extends SliceQuery { | ||
public DocValuesSliceQuery(String field, int id, int max) { | ||
super(field, id, max); | ||
} | ||
|
||
@Override | ||
public Weight createWeight(IndexSearcher searcher, boolean needsScores) throws IOException { | ||
return new RandomAccessWeight(this) { | ||
@Override | ||
protected Bits getMatchingDocs(final LeafReaderContext context) throws IOException { | ||
final SortedNumericDocValues values = DocValues.getSortedNumeric(context.reader(), getField()); | ||
return new Bits() { | ||
@Override | ||
public boolean get(int doc) { | ||
values.setDocument(doc); | ||
for (int i = 0; i < values.count(); i++) { | ||
return contains(Long.hashCode(values.valueAt(i))); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. s/Long.hashCode/BitMixer.mix64/ ? otherwise we might still have the issue with doubles given that Long.hashCode is a bit naive There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks, I was naive too, I pushed another commit. |
||
} | ||
return contains(0); | ||
} | ||
|
||
@Override | ||
public int length() { | ||
return context.reader().maxDoc(); | ||
} | ||
}; | ||
} | ||
}; | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
+1 :)