Skip to content

Commit 934711e

Browse files
dnhatndnhatn
committed
Keep query inside LuceneScorer
1 parent 39e63c7 commit 934711e

File tree

4 files changed

+33
-38
lines changed

4 files changed

+33
-38
lines changed

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneOperator.java‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -206,15 +206,23 @@ private PartialLeafReaderContext nextPartialLeaf() {
206206

207207
private void loadScorerForNewPartialLeaf(PartialLeafReaderContext partialLeaf) {
208208
final LeafReaderContext leaf = partialLeaf.leafReaderContext();
209-
// the current Weight can be reused with the current slice
210-
if (currentScorer != null && currentSlice.isWeightCompatible(currentScorer.weight)) {
209+
if (currentScorer != null
210+
&& currentScorer.query == currentSlice.query()
211+
&& currentScorer.shardContext == currentSlice.shardContext()) {
212+
assert currentScorer.tags == currentSlice.tags() : "different tags: " + currentScorer.tags + " != " + currentSlice.tags();
211213
if (currentScorer.leafReaderContext != leaf) {
212-
currentScorer = new LuceneScorer(currentSlice.shardContext(), currentScorer.weight, currentSlice.tags(), leaf);
214+
currentScorer = new LuceneScorer(
215+
currentSlice.shardContext(),
216+
currentScorer.weight,
217+
currentSlice.query(),
218+
currentSlice.tags(),
219+
leaf
220+
);
213221
}
214222
} else {
215223
final var weight = currentSlice.createWeight();
216-
processedQueries.add(weight.getQuery());
217-
currentScorer = new LuceneScorer(currentSlice.shardContext(), weight, currentSlice.tags(), leaf);
224+
currentScorer = new LuceneScorer(currentSlice.shardContext(), weight, currentSlice.query(), currentSlice.tags(), leaf);
225+
processedQueries.add(currentScorer.query);
218226
}
219227
assert currentScorer.maxPosition <= partialLeaf.maxDoc() : currentScorer.maxPosition + ">" + partialLeaf.maxDoc();
220228
currentScorer.maxPosition = partialLeaf.maxDoc();
@@ -234,6 +242,7 @@ ShardRefCounted currentScorerShardRefCounted() {
234242
static final class LuceneScorer {
235243
private final ShardContext shardContext;
236244
private final Weight weight;
245+
private final Query query;
237246
private final LeafReaderContext leafReaderContext;
238247
private final List<Object> tags;
239248

@@ -242,9 +251,10 @@ static final class LuceneScorer {
242251
private int maxPosition;
243252
private Thread executingThread;
244253

245-
LuceneScorer(ShardContext shardContext, Weight weight, List<Object> tags, LeafReaderContext leafReaderContext) {
254+
LuceneScorer(ShardContext shardContext, Weight weight, Query query, List<Object> tags, LeafReaderContext leafReaderContext) {
246255
this.shardContext = shardContext;
247256
this.weight = weight;
257+
this.query = query;
248258
this.tags = tags;
249259
this.leafReaderContext = leafReaderContext;
250260
reinitialize();

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSlice.java‎

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77

88
package org.elasticsearch.compute.lucene;
99

10-
import org.apache.lucene.search.FilterWeight;
1110
import org.apache.lucene.search.Query;
1211
import org.apache.lucene.search.ScoreMode;
1312
import org.apache.lucene.search.Weight;
@@ -24,8 +23,8 @@ public record LuceneSlice(
2423
boolean queryHead,
2524
ShardContext shardContext,
2625
List<PartialLeafReaderContext> leaves,
27-
Query query,
2826
ScoreMode scoreMode,
27+
Query query,
2928
List<Object> tags
3029
) {
3130

@@ -40,23 +39,9 @@ PartialLeafReaderContext getLeaf(int index) {
4039
Weight createWeight() {
4140
var searcher = shardContext.searcher();
4241
try {
43-
Weight w = searcher.createWeight(query, scoreMode, 1);
44-
return new OwningWeight(query, w);
42+
return searcher.createWeight(query, scoreMode, 1);
4543
} catch (IOException e) {
4644
throw new UncheckedIOException(e);
4745
}
4846
}
49-
50-
private static class OwningWeight extends FilterWeight {
51-
final Query originalQuery;
52-
53-
protected OwningWeight(Query originalQuery, Weight weight) {
54-
super(weight);
55-
this.originalQuery = originalQuery;
56-
}
57-
}
58-
59-
boolean isWeightCompatible(Weight weight) {
60-
return weight instanceof OwningWeight ow && ow.originalQuery == query;
61-
}
6247
}

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSliceQueue.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,7 @@ public static LuceneSliceQueue create(
212212
for (List<PartialLeafReaderContext> group : groups) {
213213
if (group.isEmpty() == false) {
214214
final int slicePosition = nextSliceId++;
215-
slices.add(new LuceneSlice(slicePosition, queryHead, ctx, group, query, scoreMode, queryAndExtra.tags));
215+
slices.add(new LuceneSlice(slicePosition, queryHead, ctx, group, scoreMode, query, queryAndExtra.tags));
216216
queryHead = false;
217217
}
218218
}

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSliceQueueTests.java‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -56,23 +56,23 @@ public void testBasics() {
5656
var scoreMode = ScoreMode.COMPLETE_NO_SCORES;
5757
List<LuceneSlice> sliceList = List.of(
5858
// query1: new segment
59-
new LuceneSlice(0, true, null, List.of(new PartialLeafReaderContext(leaf1, 0, 10)), null, scoreMode, t1),
60-
new LuceneSlice(1, false, null, List.of(new PartialLeafReaderContext(leaf2, 0, 10)), null, scoreMode, t1),
61-
new LuceneSlice(2, false, null, List.of(new PartialLeafReaderContext(leaf2, 10, 20)), null, scoreMode, t1),
59+
new LuceneSlice(0, true, null, List.of(new PartialLeafReaderContext(leaf1, 0, 10)), scoreMode, null, t1),
60+
new LuceneSlice(1, false, null, List.of(new PartialLeafReaderContext(leaf2, 0, 10)), scoreMode, null, t1),
61+
new LuceneSlice(2, false, null, List.of(new PartialLeafReaderContext(leaf2, 10, 20)), scoreMode, null, t1),
6262
// query1: new segment
63-
new LuceneSlice(3, false, null, List.of(new PartialLeafReaderContext(leaf3, 0, 20)), null, scoreMode, t1),
64-
new LuceneSlice(4, false, null, List.of(new PartialLeafReaderContext(leaf3, 10, 20)), null, scoreMode, t1),
65-
new LuceneSlice(5, false, null, List.of(new PartialLeafReaderContext(leaf3, 20, 30)), null, scoreMode, t1),
63+
new LuceneSlice(3, false, null, List.of(new PartialLeafReaderContext(leaf3, 0, 20)), scoreMode, null, t1),
64+
new LuceneSlice(4, false, null, List.of(new PartialLeafReaderContext(leaf3, 10, 20)), scoreMode, null, t1),
65+
new LuceneSlice(5, false, null, List.of(new PartialLeafReaderContext(leaf3, 20, 30)), scoreMode, null, t1),
6666
// query1: new segment
67-
new LuceneSlice(6, false, null, List.of(new PartialLeafReaderContext(leaf4, 0, 10)), null, scoreMode, t1),
68-
new LuceneSlice(7, false, null, List.of(new PartialLeafReaderContext(leaf4, 10, 20)), null, scoreMode, t1),
67+
new LuceneSlice(6, false, null, List.of(new PartialLeafReaderContext(leaf4, 0, 10)), scoreMode, null, t1),
68+
new LuceneSlice(7, false, null, List.of(new PartialLeafReaderContext(leaf4, 10, 20)), scoreMode, null, t1),
6969
// query2: new segment
70-
new LuceneSlice(8, true, null, List.of(new PartialLeafReaderContext(leaf2, 0, 10)), null, scoreMode, t2),
71-
new LuceneSlice(9, false, null, List.of(new PartialLeafReaderContext(leaf2, 10, 20)), null, scoreMode, t2),
70+
new LuceneSlice(8, true, null, List.of(new PartialLeafReaderContext(leaf2, 0, 10)), scoreMode, null, t2),
71+
new LuceneSlice(9, false, null, List.of(new PartialLeafReaderContext(leaf2, 10, 20)), scoreMode, null, t2),
7272
// query1: new segment
73-
new LuceneSlice(10, false, null, List.of(new PartialLeafReaderContext(leaf3, 0, 20)), null, scoreMode, t2),
74-
new LuceneSlice(11, false, null, List.of(new PartialLeafReaderContext(leaf3, 10, 20)), null, scoreMode, t2),
75-
new LuceneSlice(12, false, null, List.of(new PartialLeafReaderContext(leaf3, 20, 30)), null, scoreMode, t2)
73+
new LuceneSlice(10, false, null, List.of(new PartialLeafReaderContext(leaf3, 0, 20)), scoreMode, null, t2),
74+
new LuceneSlice(11, false, null, List.of(new PartialLeafReaderContext(leaf3, 10, 20)), scoreMode, null, t2),
75+
new LuceneSlice(12, false, null, List.of(new PartialLeafReaderContext(leaf3, 20, 30)), scoreMode, null, t2)
7676
);
7777
// single driver
7878
{
@@ -141,8 +141,8 @@ public void testRandom() throws Exception {
141141
false,
142142
mock(ShardContext.class),
143143
List.of(new PartialLeafReaderContext(leafContext, minDoc, maxDoc)),
144-
null,
145144
ScoreMode.COMPLETE_NO_SCORES,
145+
null,
146146
null
147147
);
148148
sliceList.add(slice);

0 commit comments

Comments
 (0)