Skip to content

Commit 0a58652

Browse files
KunjueYuKunjueYu
committed
Script sort support parallel collection
1 parent 49f7cfb commit 0a58652

6 files changed

+427
-199
lines changed

‎server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/BytesRefFieldComparatorSource.java

+25-51
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import org.apache.lucene.index.SortedSetDocValues;
1616
import org.apache.lucene.search.DocIdSetIterator;
1717
import org.apache.lucene.search.FieldComparator;
18-
import org.apache.lucene.search.LeafFieldComparator;
1918
import org.apache.lucene.search.Pruning;
2019
import org.apache.lucene.search.Scorable;
2120
import org.apache.lucene.search.SortField;
@@ -70,6 +69,20 @@ protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOEx
7069

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

72+
protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, BytesRef missingBytes, SortedBinaryDocValues values)
73+
throws IOException {
74+
final BinaryDocValues selectedValues;
75+
if (nested == null) {
76+
selectedValues = sortMode.select(values, missingBytes);
77+
} else {
78+
final BitSet rootDocs = nested.rootDocs(context);
79+
final DocIdSetIterator innerDocs = nested.innerDocs(context);
80+
final int maxChildren = nested.getNestedSort() != null ? nested.getNestedSort().getMaxChildren() : Integer.MAX_VALUE;
81+
selectedValues = sortMode.select(values, missingBytes, rootDocs, innerDocs, maxChildren);
82+
}
83+
return selectedValues;
84+
}
85+
7386
@Override
7487
public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning enableSkipping, boolean reversed) {
7588
assert indexFieldData == null || fieldname.equals(indexFieldData.getFieldName());
@@ -102,61 +115,22 @@ protected SortedDocValues getSortedDocValues(LeafReaderContext context, String f
102115

103116
};
104117
}
118+
return newComparatorWithoutOrdinal(fieldname, numHits, enableSkipping, reversed, missingBytes, sortMissingLast);
119+
}
105120

121+
protected FieldComparator<?> newComparatorWithoutOrdinal(
122+
String fieldname,
123+
int numHits,
124+
Pruning enableSkipping,
125+
boolean reversed,
126+
BytesRef missingBytes,
127+
boolean sortMissingLast
128+
) {
106129
return new FieldComparator.TermValComparator(numHits, null, sortMissingLast) {
107130

108131
@Override
109132
protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, String field) throws IOException {
110-
final SortedBinaryDocValues values = getValues(context);
111-
final BinaryDocValues selectedValues;
112-
if (nested == null) {
113-
selectedValues = sortMode.select(values, missingBytes);
114-
} else {
115-
final BitSet rootDocs = nested.rootDocs(context);
116-
final DocIdSetIterator innerDocs = nested.innerDocs(context);
117-
final int maxChildren = nested.getNestedSort() != null ? nested.getNestedSort().getMaxChildren() : Integer.MAX_VALUE;
118-
selectedValues = sortMode.select(values, missingBytes, rootDocs, innerDocs, maxChildren);
119-
}
120-
return selectedValues;
121-
}
122-
123-
@Override
124-
public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
125-
LeafFieldComparator leafComparator = super.getLeafComparator(context);
126-
// TopFieldCollector interacts with inter-segment concurrency by creating a FieldValueHitQueue per slice, each one with a
127-
// specific instance of the FieldComparator. This ensures sequential execution across LeafFieldComparators returned by
128-
// the same parent FieldComparator. That allows for effectively sharing the same instance of leaf comparator, like in this
129-
// case in the Lucene code. That's fine dealing with sorting by field, but not when using script sorting, because we then
130-
// need to set to Scorer to the specific leaf comparator, to make the _score variable available in sort scripts. The
131-
// setScorer call happens concurrently across slices and needs to target the specific leaf context that is being searched.
132-
return new LeafFieldComparator() {
133-
@Override
134-
public void setBottom(int slot) throws IOException {
135-
leafComparator.setBottom(slot);
136-
}
137-
138-
@Override
139-
public int compareBottom(int doc) throws IOException {
140-
return leafComparator.compareBottom(doc);
141-
}
142-
143-
@Override
144-
public int compareTop(int doc) throws IOException {
145-
return leafComparator.compareTop(doc);
146-
}
147-
148-
@Override
149-
public void copy(int slot, int doc) throws IOException {
150-
leafComparator.copy(slot, doc);
151-
}
152-
153-
@Override
154-
public void setScorer(Scorable scorer) {
155-
// this ensures that the scorer is set for the specific leaf comparator
156-
// corresponding to the leaf context we are scoring
157-
BytesRefFieldComparatorSource.this.setScorer(context, scorer);
158-
}
159-
};
133+
return BytesRefFieldComparatorSource.this.getBinaryDocValues(context, missingBytes, getValues(context));
160134
}
161135
};
162136
}

‎server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/DoubleValuesComparatorSource.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,11 @@ protected SortedNumericDoubleValues getValues(LeafReaderContext context) throws
6060
}
6161

6262
private NumericDoubleValues getNumericDocValues(LeafReaderContext context, double missingValue) throws IOException {
63-
final SortedNumericDoubleValues values = getValues(context);
63+
return getNumericDocValues(context, missingValue, getValues(context));
64+
}
65+
66+
protected NumericDoubleValues getNumericDocValues(LeafReaderContext context, double missingValue, SortedNumericDoubleValues values)
67+
throws IOException {
6468
if (nested == null) {
6569
return FieldData.replaceMissing(sortMode.select(values), missingValue);
6670
} else {
@@ -78,6 +82,10 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
7882
assert indexFieldData == null || fieldname.equals(indexFieldData.getFieldName());
7983

8084
final double dMissingValue = (Double) missingObject(missingValue, reversed);
85+
return newComparator(numHits, enableSkipping, reversed, dMissingValue);
86+
}
87+
88+
protected FieldComparator<?> newComparator(int numHits, Pruning enableSkipping, boolean reversed, double dMissingValue) {
8189
// NOTE: it's important to pass null as a missing value in the constructor so that
8290
// the comparator doesn't check docsWithField since we replace missing values in select()
8391
return new DoubleComparator(numHits, null, null, reversed, Pruning.NONE) {

‎server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/ScriptDoubleValuesComparatorSource.java

+132
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.index.fielddata.fieldcomparator;
11+
12+
import org.apache.lucene.index.LeafReaderContext;
13+
import org.apache.lucene.index.NumericDocValues;
14+
import org.apache.lucene.search.FieldComparator;
15+
import org.apache.lucene.search.LeafFieldComparator;
16+
import org.apache.lucene.search.Pruning;
17+
import org.apache.lucene.search.Scorable;
18+
import org.apache.lucene.search.comparators.DoubleComparator;
19+
import org.elasticsearch.common.util.BigArrays;
20+
import org.elasticsearch.core.CheckedFunction;
21+
import org.elasticsearch.core.Nullable;
22+
import org.elasticsearch.index.fielddata.FieldData;
23+
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
24+
import org.elasticsearch.index.fielddata.NumericDoubleValues;
25+
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
26+
import org.elasticsearch.script.NumberSortScript;
27+
import org.elasticsearch.search.DocValueFormat;
28+
import org.elasticsearch.search.MultiValueMode;
29+
import org.elasticsearch.search.sort.BucketedSort;
30+
import org.elasticsearch.search.sort.SortOrder;
31+
32+
import java.io.IOException;
33+
34+
/**
35+
* Script comparator source for double values.
36+
*/
37+
public class ScriptDoubleValuesComparatorSource extends DoubleValuesComparatorSource {
38+
39+
private final CheckedFunction<LeafReaderContext, NumberSortScript, IOException> scriptSupplier;
40+
41+
public ScriptDoubleValuesComparatorSource(
42+
CheckedFunction<LeafReaderContext, NumberSortScript, IOException> scriptSupplier,
43+
IndexNumericFieldData indexFieldData,
44+
@Nullable Object missingValue,
45+
MultiValueMode sortMode,
46+
Nested nested
47+
) {
48+
super(indexFieldData, missingValue, sortMode, nested);
49+
this.scriptSupplier = scriptSupplier;
50+
}
51+
52+
private SortedNumericDoubleValues getValues(NumberSortScript leafScript) throws IOException {
53+
final NumericDoubleValues values = new NumericDoubleValues() {
54+
@Override
55+
public boolean advanceExact(int doc) {
56+
leafScript.setDocument(doc);
57+
return true;
58+
}
59+
60+
@Override
61+
public double doubleValue() {
62+
return leafScript.execute();
63+
}
64+
};
65+
return FieldData.singleton(values);
66+
}
67+
68+
private NumericDoubleValues getNumericDocValues(LeafReaderContext context, double missingValue, NumberSortScript leafScript)
69+
throws IOException {
70+
return getNumericDocValues(context, missingValue, getValues(leafScript));
71+
}
72+
73+
@Override
74+
protected FieldComparator<?> newComparator(int numHits, Pruning enableSkipping, boolean reversed, double dMissingValue) {
75+
// NOTE: it's important to pass null as a missing value in the constructor so that
76+
// the comparator doesn't check docsWithField since we replace missing values in select()
77+
return new DoubleComparator(numHits, null, null, reversed, Pruning.NONE) {
78+
@Override
79+
public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
80+
NumberSortScript leafScript = scriptSupplier.apply(context);
81+
return new DoubleLeafComparator(context) {
82+
@Override
83+
protected NumericDocValues getNumericDocValues(LeafReaderContext context, String field) throws IOException {
84+
return ScriptDoubleValuesComparatorSource.this.getNumericDocValues(context, dMissingValue, leafScript)
85+
.getRawDoubleValues();
86+
}
87+
88+
@Override
89+
public void setScorer(Scorable scorer) {
90+
leafScript.setScorer(scorer);
91+
}
92+
};
93+
}
94+
};
95+
}
96+
97+
@Override
98+
public BucketedSort newBucketedSort(
99+
BigArrays bigArrays,
100+
SortOrder sortOrder,
101+
DocValueFormat format,
102+
int bucketSize,
103+
BucketedSort.ExtraData extra
104+
) {
105+
return new BucketedSort.ForDoubles(bigArrays, sortOrder, format, bucketSize, extra) {
106+
private final double dMissingValue = (Double) missingObject(missingValue, sortOrder == SortOrder.DESC);
107+
108+
@Override
109+
public Leaf forLeaf(LeafReaderContext ctx) throws IOException {
110+
NumberSortScript leafScript = scriptSupplier.apply(ctx);
111+
return new Leaf(ctx) {
112+
private final NumericDoubleValues docValues = getNumericDocValues(ctx, dMissingValue, leafScript);
113+
private double docValue;
114+
115+
@Override
116+
protected boolean advanceExact(int doc) throws IOException {
117+
if (docValues.advanceExact(doc)) {
118+
docValue = docValues.doubleValue();
119+
return true;
120+
}
121+
return false;
122+
}
123+
124+
@Override
125+
protected double docValue() {
126+
return docValue;
127+
}
128+
};
129+
}
130+
};
131+
}
132+
}

‎server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/ScriptStringFieldComparatorSource.java

+112
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.index.fielddata.fieldcomparator;
11+
12+
import org.apache.lucene.index.BinaryDocValues;
13+
import org.apache.lucene.index.LeafReaderContext;
14+
import org.apache.lucene.search.FieldComparator;
15+
import org.apache.lucene.search.Pruning;
16+
import org.apache.lucene.search.Scorable;
17+
import org.apache.lucene.util.BytesRef;
18+
import org.apache.lucene.util.BytesRefBuilder;
19+
import org.elasticsearch.common.util.BigArrays;
20+
import org.elasticsearch.core.CheckedFunction;
21+
import org.elasticsearch.index.fielddata.AbstractBinaryDocValues;
22+
import org.elasticsearch.index.fielddata.FieldData;
23+
import org.elasticsearch.index.fielddata.IndexFieldData;
24+
import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
25+
import org.elasticsearch.script.StringSortScript;
26+
import org.elasticsearch.search.DocValueFormat;
27+
import org.elasticsearch.search.MultiValueMode;
28+
import org.elasticsearch.search.sort.BucketedSort;
29+
import org.elasticsearch.search.sort.ScriptSortBuilder;
30+
import org.elasticsearch.search.sort.SortOrder;
31+
32+
import java.io.IOException;
33+
34+
/**
35+
* Script comparator source for string/binary values.
36+
*/
37+
public class ScriptStringFieldComparatorSource extends BytesRefFieldComparatorSource {
38+
39+
final CheckedFunction<LeafReaderContext, StringSortScript, IOException> scriptSupplier;
40+
41+
public ScriptStringFieldComparatorSource(
42+
CheckedFunction<LeafReaderContext, StringSortScript, IOException> scriptSupplier,
43+
IndexFieldData<?> indexFieldData,
44+
Object missingValue,
45+
MultiValueMode sortMode,
46+
Nested nested
47+
) {
48+
super(indexFieldData, missingValue, sortMode, nested);
49+
this.scriptSupplier = scriptSupplier;
50+
}
51+
52+
private SortedBinaryDocValues getValues(StringSortScript leafScript) throws IOException {
53+
final BinaryDocValues values = new AbstractBinaryDocValues() {
54+
final BytesRefBuilder spare = new BytesRefBuilder();
55+
56+
@Override
57+
public boolean advanceExact(int doc) {
58+
leafScript.setDocument(doc);
59+
return true;
60+
}
61+
62+
@Override
63+
public BytesRef binaryValue() {
64+
spare.copyChars(leafScript.execute());
65+
return spare.get();
66+
}
67+
};
68+
return FieldData.singleton(values);
69+
}
70+
71+
@Override
72+
protected FieldComparator<?> newComparatorWithoutOrdinal(
73+
String fieldname,
74+
int numHits,
75+
Pruning enableSkipping,
76+
boolean reversed,
77+
BytesRef missingBytes,
78+
boolean sortMissingLast
79+
) {
80+
return new FieldComparator.TermValComparator(numHits, null, sortMissingLast) {
81+
82+
StringSortScript leafScript;
83+
84+
@Override
85+
protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, String field) throws IOException {
86+
leafScript = scriptSupplier.apply(context);
87+
return ScriptStringFieldComparatorSource.this.getBinaryDocValues(context, missingBytes, getValues(leafScript));
88+
}
89+
90+
@Override
91+
public void setScorer(Scorable scorer) {
92+
leafScript.setScorer(scorer);
93+
}
94+
};
95+
}
96+
97+
@Override
98+
public BucketedSort newBucketedSort(
99+
BigArrays bigArrays,
100+
SortOrder sortOrder,
101+
DocValueFormat format,
102+
int bucketSize,
103+
BucketedSort.ExtraData extra
104+
) {
105+
throw new IllegalArgumentException(
106+
"error building sort for [_script]: "
107+
+ "script sorting only supported on [numeric] scripts but was ["
108+
+ ScriptSortBuilder.ScriptSortType.STRING
109+
+ "]"
110+
);
111+
}
112+
}

0 commit comments

Comments
 (0)