Skip to content

Commit

Permalink
Fix the source of concurrency issues in script sorting
Browse files Browse the repository at this point in the history 
We explicitly disable search concurrency when sorting by script, because there are
known issues with the script sorting implementation that assumes segments are searched
sequentially. This commit addresses that, which was a bad design relying on assumptions
that no longer hold.
  • Loading branch information
@javanna
javanna committed Feb 27, 2025
1 parent 79c388a commit 0c8037d
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 33 deletions.
40 changes: 21 additions & 19 deletions ...java/org/elasticsearch/index/fielddata/fieldcomparator/BytesRefFieldComparatorSource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOEx
return indexFieldData.load(context).getBytesValues();
}

protected void setScorer(Scorable scorer) {}
protected void setScorer(LeafReaderContext context, Scorable scorer) {}

@Override
public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning enableSkipping, boolean reversed) {
Expand Down Expand Up @@ -103,27 +103,29 @@ protected SortedDocValues getSortedDocValues(LeafReaderContext context, String f
}

return new FieldComparator.TermValComparator(numHits, null, sortMissingLast) {

@Override
protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, String field) throws IOException {
final SortedBinaryDocValues values = getValues(context);
final BinaryDocValues selectedValues;
if (nested == null) {
selectedValues = sortMode.select(values, missingBytes);
} else {
final BitSet rootDocs = nested.rootDocs(context);
final DocIdSetIterator innerDocs = nested.innerDocs(context);
final int maxChildren = nested.getNestedSort() != null ? nested.getNestedSort().getMaxChildren() : Integer.MAX_VALUE;
selectedValues = sortMode.select(values, missingBytes, rootDocs, innerDocs, maxChildren);
}
return selectedValues;
}
public LeafFieldComparator getLeafComparator(LeafReaderContext context) {
return new TermValComparator(numHits, null, sortMissingLast) {
protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, String field) throws IOException {
final SortedBinaryDocValues values = getValues(context);
final BinaryDocValues selectedValues;
if (nested == null) {
selectedValues = sortMode.select(values, missingBytes);
} else {
final BitSet rootDocs = nested.rootDocs(context);
final DocIdSetIterator innerDocs = nested.innerDocs(context);
final int maxChildren = nested.getNestedSort() != null ? nested.getNestedSort().getMaxChildren() : Integer.MAX_VALUE;
selectedValues = sortMode.select(values, missingBytes, rootDocs, innerDocs, maxChildren);
}
return selectedValues;
}

@Override
public void setScorer(Scorable scorer) {
BytesRefFieldComparatorSource.this.setScorer(scorer);
@Override
public void setScorer(Scorable scorer) {
BytesRefFieldComparatorSource.this.setScorer(context, scorer);
}
};
}

};
}

Expand Down
4 changes: 2 additions & 2 deletions .../java/org/elasticsearch/index/fielddata/fieldcomparator/DoubleValuesComparatorSource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ private NumericDoubleValues getNumericDocValues(LeafReaderContext context, doubl
}
}

protected void setScorer(Scorable scorer) {}
protected void setScorer(LeafReaderContext context, Scorable scorer) {}

@Override
public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning enableSkipping, boolean reversed) {
Expand All @@ -91,7 +91,7 @@ protected NumericDocValues getNumericDocValues(LeafReaderContext context, String

@Override
public void setScorer(Scorable scorer) {
DoubleValuesComparatorSource.this.setScorer(scorer);
DoubleValuesComparatorSource.this.setScorer(context, scorer);
}
};
}
Expand Down
57 changes: 45 additions & 12 deletions server/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.common.util.BigArrays;
import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
import org.elasticsearch.index.fielddata.AbstractBinaryDocValues;
import org.elasticsearch.index.fielddata.FieldData;
import org.elasticsearch.index.fielddata.IndexFieldData;
Expand Down Expand Up @@ -51,7 +52,9 @@
import org.elasticsearch.xcontent.XContentParser;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;

import static org.elasticsearch.search.sort.FieldSortBuilder.validateMaxChildrenExistOnlyInTopLevelNestedSort;
Expand Down Expand Up @@ -278,11 +281,11 @@ private IndexFieldData.XFieldComparatorSource fieldComparatorSource(SearchExecut
final StringSortScript.Factory factory = context.compile(script, StringSortScript.CONTEXT);
final StringSortScript.LeafFactory searchScript = factory.newFactory(script.getParams());
return new BytesRefFieldComparatorSource(null, null, valueMode, nested) {
StringSortScript leafScript;
final Map<Object, StringSortScript> leafScripts = ConcurrentCollections.newConcurrentMap();

@Override
protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOException {
leafScript = searchScript.newInstance(new DocValuesDocReader(searchLookup, context));
StringSortScript leafScript = getLeafScript(context);
final BinaryDocValues values = new AbstractBinaryDocValues() {
final BytesRefBuilder spare = new BytesRefBuilder();

Expand All @@ -302,8 +305,18 @@ public BytesRef binaryValue() {
}

@Override
protected void setScorer(Scorable scorer) {
leafScript.setScorer(scorer);
protected void setScorer(LeafReaderContext context, Scorable scorer) {
getLeafScript(context).setScorer(scorer);
}

StringSortScript getLeafScript(LeafReaderContext context) {
return leafScripts.computeIfAbsent(context.id(), o -> {
try {
return searchScript.newInstance(new DocValuesDocReader(searchLookup, context));
} catch (IOException e) {
throw new UncheckedIOException(e);
}
});
}

@Override
Expand All @@ -328,11 +341,11 @@ public BucketedSort newBucketedSort(
// searchLookup is unnecessary here, as it's just used for expressions
final NumberSortScript.LeafFactory numberSortScript = numberSortFactory.newFactory(script.getParams(), searchLookup);
return new DoubleValuesComparatorSource(null, Double.MAX_VALUE, valueMode, nested) {
NumberSortScript leafScript;
final Map<Object, NumberSortScript> leafScripts = ConcurrentCollections.newConcurrentMap();

@Override
protected SortedNumericDoubleValues getValues(LeafReaderContext context) throws IOException {
leafScript = numberSortScript.newInstance(new DocValuesDocReader(searchLookup, context));
NumberSortScript leafScript = getLeafScript(context);
final NumericDoubleValues values = new NumericDoubleValues() {
@Override
public boolean advanceExact(int doc) {
Expand All @@ -349,20 +362,30 @@ public double doubleValue() {
}

@Override
protected void setScorer(Scorable scorer) {
leafScript.setScorer(scorer);
protected void setScorer(LeafReaderContext context, Scorable scorer) {
getLeafScript(context).setScorer(scorer);
}

NumberSortScript getLeafScript(LeafReaderContext context) {
return leafScripts.computeIfAbsent(context.id(), o -> {
try {
return numberSortScript.newInstance(new DocValuesDocReader(searchLookup, context));
} catch (IOException e) {
throw new UncheckedIOException(e);
}
});
}
};
}
case VERSION -> {
final BytesRefSortScript.Factory factory = context.compile(script, BytesRefSortScript.CONTEXT);
final BytesRefSortScript.LeafFactory searchScript = factory.newFactory(script.getParams());
return new BytesRefFieldComparatorSource(null, null, valueMode, nested) {
BytesRefSortScript leafScript;
final Map<Object, BytesRefSortScript> leafScripts = ConcurrentCollections.newConcurrentMap();

@Override
protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOException {
leafScript = searchScript.newInstance(new DocValuesDocReader(searchLookup, context));
BytesRefSortScript leafScript = getLeafScript(context);
final BinaryDocValues values = new AbstractBinaryDocValues() {

@Override
Expand Down Expand Up @@ -391,8 +414,18 @@ public BytesRef binaryValue() {
}

@Override
protected void setScorer(Scorable scorer) {
leafScript.setScorer(scorer);
protected void setScorer(LeafReaderContext context, Scorable scorer) {
getLeafScript(context).setScorer(scorer);
}

BytesRefSortScript getLeafScript(LeafReaderContext context) {
return leafScripts.computeIfAbsent(context.id(), o -> {
try {
return searchScript.newInstance(new DocValuesDocReader(searchLookup, context));
} catch (IOException e) {
throw new UncheckedIOException(e);
}
});
}

@Override
Expand Down

0 comments on commit 0c8037d

Please sign in to comment.