Skip to content

Commit 3c6ab94

Browse files
dnhatndnhatn
committed
Hold query and tags
1 parent 71495f5 commit 3c6ab94

File tree

4 files changed

+44
-37
lines changed

4 files changed

+44
-37
lines changed

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

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -207,22 +207,15 @@ private PartialLeafReaderContext nextPartialLeaf() {
207207
private void loadScorerForNewPartialLeaf(PartialLeafReaderContext partialLeaf) {
208208
final LeafReaderContext leaf = partialLeaf.leafReaderContext();
209209
if (currentScorer != null
210-
&& currentScorer.query == currentSlice.query()
210+
&& currentScorer.query() == currentSlice.query()
211211
&& currentScorer.shardContext == currentSlice.shardContext()) {
212-
assert currentScorer.tags == currentSlice.tags() : "different tags: " + currentScorer.tags + " != " + currentSlice.tags();
213212
if (currentScorer.leafReaderContext != leaf) {
214-
currentScorer = new LuceneScorer(
215-
currentSlice.shardContext(),
216-
currentScorer.weight,
217-
currentSlice.query(),
218-
currentSlice.tags(),
219-
leaf
220-
);
213+
currentScorer = new LuceneScorer(currentSlice.shardContext(), currentScorer.weight, currentSlice.queryAndTags(), leaf);
221214
}
222215
} else {
223216
final var weight = currentSlice.createWeight();
224-
currentScorer = new LuceneScorer(currentSlice.shardContext(), weight, currentSlice.query(), currentSlice.tags(), leaf);
225-
processedQueries.add(currentScorer.query);
217+
currentScorer = new LuceneScorer(currentSlice.shardContext(), weight, currentSlice.queryAndTags(), leaf);
218+
processedQueries.add(currentScorer.query());
226219
}
227220
assert currentScorer.maxPosition <= partialLeaf.maxDoc() : currentScorer.maxPosition + ">" + partialLeaf.maxDoc();
228221
currentScorer.maxPosition = partialLeaf.maxDoc();
@@ -242,20 +235,23 @@ ShardRefCounted currentScorerShardRefCounted() {
242235
static final class LuceneScorer {
243236
private final ShardContext shardContext;
244237
private final Weight weight;
245-
private final Query query;
238+
private final LuceneSliceQueue.QueryAndTags queryAndTags;
246239
private final LeafReaderContext leafReaderContext;
247-
private final List<Object> tags;
248240

249241
private BulkScorer bulkScorer;
250242
private int position;
251243
private int maxPosition;
252244
private Thread executingThread;
253245

254-
LuceneScorer(ShardContext shardContext, Weight weight, Query query, List<Object> tags, LeafReaderContext leafReaderContext) {
246+
LuceneScorer(
247+
ShardContext shardContext,
248+
Weight weight,
249+
LuceneSliceQueue.QueryAndTags queryAndTags,
250+
LeafReaderContext leafReaderContext
251+
) {
255252
this.shardContext = shardContext;
256253
this.weight = weight;
257-
this.query = query;
258-
this.tags = tags;
254+
this.queryAndTags = queryAndTags;
259255
this.leafReaderContext = leafReaderContext;
260256
reinitialize();
261257
}
@@ -305,7 +301,11 @@ int position() {
305301
* Tags to add to the data returned by this query.
306302
*/
307303
List<Object> tags() {
308-
return tags;
304+
return queryAndTags.tags();
305+
}
306+
307+
Query query() {
308+
return queryAndTags.query();
309309
}
310310
}
311311

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,7 @@ public record LuceneSlice(
2424
ShardContext shardContext,
2525
List<PartialLeafReaderContext> leaves,
2626
ScoreMode scoreMode,
27-
Query query,
28-
List<Object> tags
27+
LuceneSliceQueue.QueryAndTags queryAndTags
2928
) {
3029

3130
int numLeaves() {
@@ -36,10 +35,18 @@ PartialLeafReaderContext getLeaf(int index) {
3635
return leaves.get(index);
3736
}
3837

38+
Query query() {
39+
return queryAndTags.query();
40+
}
41+
42+
List<Object> tags() {
43+
return queryAndTags.tags();
44+
}
45+
3946
Weight createWeight() {
4047
var searcher = shardContext.searcher();
4148
try {
42-
return searcher.createWeight(query, scoreMode, 1);
49+
return searcher.createWeight(queryAndTags.query(), scoreMode, 1);
4350
} catch (IOException e) {
4451
throw new UncheckedIOException(e);
4552
}

‎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, scoreMode, query, queryAndExtra.tags));
215+
slices.add(new LuceneSlice(slicePosition, queryHead, ctx, group, scoreMode, queryAndExtra));
216216
queryHead = false;
217217
}
218218
}

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

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.apache.lucene.index.TermVectors;
2525
import org.apache.lucene.index.Terms;
2626
import org.apache.lucene.search.KnnCollector;
27+
import org.apache.lucene.search.MatchAllDocsQuery;
2728
import org.apache.lucene.search.ScoreMode;
2829
import org.apache.lucene.util.Bits;
2930
import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
@@ -51,28 +52,28 @@ public void testBasics() {
5152
LeafReaderContext leaf2 = new MockLeafReader(1000).getContext();
5253
LeafReaderContext leaf3 = new MockLeafReader(1000).getContext();
5354
LeafReaderContext leaf4 = new MockLeafReader(1000).getContext();
54-
List<Object> t1 = List.of("1");
55-
List<Object> t2 = List.of("q2");
55+
LuceneSliceQueue.QueryAndTags t1 = new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of("q1"));
56+
LuceneSliceQueue.QueryAndTags t2 = new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of("q2"));
5657
var scoreMode = ScoreMode.COMPLETE_NO_SCORES;
5758
List<LuceneSlice> sliceList = List.of(
5859
// query1: new segment
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),
60+
new LuceneSlice(0, true, null, List.of(new PartialLeafReaderContext(leaf1, 0, 10)), scoreMode, t1),
61+
new LuceneSlice(1, false, null, List.of(new PartialLeafReaderContext(leaf2, 0, 10)), scoreMode, t1),
62+
new LuceneSlice(2, false, null, List.of(new PartialLeafReaderContext(leaf2, 10, 20)), scoreMode, t1),
6263
// query1: new segment
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),
64+
new LuceneSlice(3, false, null, List.of(new PartialLeafReaderContext(leaf3, 0, 20)), scoreMode, t1),
65+
new LuceneSlice(4, false, null, List.of(new PartialLeafReaderContext(leaf3, 10, 20)), scoreMode, t1),
66+
new LuceneSlice(5, false, null, List.of(new PartialLeafReaderContext(leaf3, 20, 30)), scoreMode, t1),
6667
// query1: new segment
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),
68+
new LuceneSlice(6, false, null, List.of(new PartialLeafReaderContext(leaf4, 0, 10)), scoreMode, t1),
69+
new LuceneSlice(7, false, null, List.of(new PartialLeafReaderContext(leaf4, 10, 20)), scoreMode, t1),
6970
// query2: new segment
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),
71+
new LuceneSlice(8, true, null, List.of(new PartialLeafReaderContext(leaf2, 0, 10)), scoreMode, t2),
72+
new LuceneSlice(9, false, null, List.of(new PartialLeafReaderContext(leaf2, 10, 20)), scoreMode, t2),
7273
// query1: new segment
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)
74+
new LuceneSlice(10, false, null, List.of(new PartialLeafReaderContext(leaf3, 0, 20)), scoreMode, t2),
75+
new LuceneSlice(11, false, null, List.of(new PartialLeafReaderContext(leaf3, 10, 20)), scoreMode, t2),
76+
new LuceneSlice(12, false, null, List.of(new PartialLeafReaderContext(leaf3, 20, 30)), scoreMode, t2)
7677
);
7778
// single driver
7879
{
@@ -142,7 +143,6 @@ public void testRandom() throws Exception {
142143
mock(ShardContext.class),
143144
List.of(new PartialLeafReaderContext(leafContext, minDoc, maxDoc)),
144145
ScoreMode.COMPLETE_NO_SCORES,
145-
null,
146146
null
147147
);
148148
sliceList.add(slice);

0 commit comments

Comments
 (0)