Skip to content

Commit 9c6cc9c

Browse files
iveraseiverase
committed
Implement our own nested clause counter in ContextIndexSearcher
1 parent cc02ac8 commit 9c6cc9c

File tree

2 files changed

+228
-1
lines changed

2 files changed

+228
-1
lines changed

‎server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,17 +13,21 @@
1313
import org.apache.lucene.index.IndexReader;
1414
import org.apache.lucene.index.LeafReaderContext;
1515
import org.apache.lucene.index.Term;
16+
import org.apache.lucene.search.BooleanClause;
17+
import org.apache.lucene.search.BooleanQuery;
1618
import org.apache.lucene.search.BulkScorer;
1719
import org.apache.lucene.search.CollectionStatistics;
1820
import org.apache.lucene.search.CollectionTerminatedException;
1921
import org.apache.lucene.search.Collector;
2022
import org.apache.lucene.search.CollectorManager;
2123
import org.apache.lucene.search.ConjunctionUtils;
2224
import org.apache.lucene.search.ConstantScoreQuery;
25+
import org.apache.lucene.search.DisjunctionMaxQuery;
2326
import org.apache.lucene.search.DocIdSetIterator;
2427
import org.apache.lucene.search.IndexSearcher;
2528
import org.apache.lucene.search.LeafCollector;
2629
import org.apache.lucene.search.MatchNoDocsQuery;
30+
import org.apache.lucene.search.MultiPhraseQuery;
2731
import org.apache.lucene.search.Query;
2832
import org.apache.lucene.search.QueryCache;
2933
import org.apache.lucene.search.QueryCachingPolicy;
@@ -202,7 +206,12 @@ public Query rewrite(Query original) throws IOException {
202206
rewriteTimer = profiler.startRewriteTime();
203207
}
204208
try {
205-
return super.rewrite(original);
209+
Query query = original;
210+
for (Query rewrittenQuery = query.rewrite(this); rewrittenQuery != query; rewrittenQuery = query.rewrite(this)) {
211+
query = rewrittenQuery;
212+
}
213+
checkNumNestedClauses(query, new int[1]);
214+
return query;
206215
} catch (TimeExceededException e) {
207216
timeExceeded = true;
208217
return new MatchNoDocsQuery("rewrite timed out");
@@ -215,6 +224,23 @@ public Query rewrite(Query original) throws IOException {
215224
}
216225
}
217226

227+
private void checkNumNestedClauses(Query query, int[] numClauses) {
228+
switch (query) {
229+
case BooleanQuery booleanQuery -> {
230+
// recurse into nested boolean queries
231+
for (BooleanClause clause : booleanQuery.clauses()) {
232+
checkNumNestedClauses(clause.query(), numClauses);
233+
}
234+
}
235+
case MultiPhraseQuery multiPhraseQuery -> numClauses[0] += multiPhraseQuery.getTermArrays().length;
236+
case DisjunctionMaxQuery disjunctionMaxQuery -> numClauses[0] += disjunctionMaxQuery.getDisjuncts().size();
237+
default -> numClauses[0]++;
238+
}
239+
if (numClauses[0] > getMaxClauseCount()) {
240+
throw new TooManyNestedClauses();
241+
}
242+
}
243+
218244
@Override
219245
public Weight createWeight(Query query, ScoreMode scoreMode, float boost) throws IOException {
220246
if (profiler != null) {

‎server/src/test/java/org/elasticsearch/search/internal/MaxClauseLimitTests.java‎

Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,201 @@
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.search.internal;
11+
12+
import org.apache.lucene.document.Document;
13+
import org.apache.lucene.document.LongField;
14+
import org.apache.lucene.document.LongPoint;
15+
import org.apache.lucene.document.SortedNumericDocValuesField;
16+
import org.apache.lucene.index.IndexReader;
17+
import org.apache.lucene.index.MultiReader;
18+
import org.apache.lucene.index.Term;
19+
import org.apache.lucene.search.BooleanClause;
20+
import org.apache.lucene.search.BooleanQuery;
21+
import org.apache.lucene.search.DisjunctionMaxQuery;
22+
import org.apache.lucene.search.IndexOrDocValuesQuery;
23+
import org.apache.lucene.search.IndexSearcher;
24+
import org.apache.lucene.search.MultiPhraseQuery;
25+
import org.apache.lucene.search.PhraseQuery;
26+
import org.apache.lucene.search.Query;
27+
import org.apache.lucene.search.TermQuery;
28+
import org.apache.lucene.store.Directory;
29+
import org.apache.lucene.tests.index.RandomIndexWriter;
30+
import org.elasticsearch.lucene.document.NumericField;
31+
import org.elasticsearch.test.ESTestCase;
32+
33+
import java.io.IOException;
34+
import java.util.Arrays;
35+
36+
import static org.hamcrest.Matchers.instanceOf;
37+
38+
/**
39+
* Mostly copied from lucene test code, but adapted to use with {@link ContextIndexSearcher}.
40+
*/
41+
public class MaxClauseLimitTests extends ESTestCase {
42+
43+
public void testFlattenInnerDisjunctionsWithMoreThan1024Terms() throws IOException {
44+
IndexSearcher searcher = newContextIndexSearcher(new MultiReader());
45+
46+
BooleanQuery.Builder builder1024 = new BooleanQuery.Builder();
47+
for (int i = 0; i < 1024; i++) {
48+
builder1024.add(new TermQuery(new Term("foo", "bar-" + i)), BooleanClause.Occur.SHOULD);
49+
}
50+
Query inner = builder1024.build();
51+
Query query = new BooleanQuery.Builder().add(inner, BooleanClause.Occur.SHOULD)
52+
.add(new TermQuery(new Term("foo", "baz")), BooleanClause.Occur.SHOULD)
53+
.build();
54+
55+
IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> { searcher.rewrite(query); });
56+
assertFalse(
57+
"Should have been caught during flattening and not required full nested walk",
58+
ex.getCause() instanceof IndexSearcher.TooManyNestedClauses
59+
);
60+
}
61+
62+
public void testLargeTermsNestedFirst() throws IOException {
63+
IndexSearcher searcher = newContextIndexSearcher(new MultiReader());
64+
BooleanQuery.Builder nestedBuilder = new BooleanQuery.Builder();
65+
66+
nestedBuilder.setMinimumNumberShouldMatch(5);
67+
68+
for (int i = 0; i < 600; i++) {
69+
nestedBuilder.add(new TermQuery(new Term("foo", "bar-" + i)), BooleanClause.Occur.SHOULD);
70+
}
71+
Query inner = nestedBuilder.build();
72+
BooleanQuery.Builder builderMixed = new BooleanQuery.Builder().add(inner, BooleanClause.Occur.SHOULD);
73+
74+
builderMixed.setMinimumNumberShouldMatch(5);
75+
76+
for (int i = 0; i < 600; i++) {
77+
builderMixed.add(new TermQuery(new Term("foo", "bar")), BooleanClause.Occur.SHOULD);
78+
}
79+
80+
Query query = builderMixed.build();
81+
82+
// Can't be flattened, but high clause count should still be cause during nested walk...
83+
IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> { searcher.rewrite(query); });
84+
assertThat(ex.getCause(), instanceOf(IndexSearcher.TooManyNestedClauses.class));
85+
}
86+
87+
public void testLargeTermsNestedLast() throws IOException {
88+
IndexSearcher searcher = newContextIndexSearcher(new MultiReader());
89+
BooleanQuery.Builder nestedBuilder = new BooleanQuery.Builder();
90+
91+
nestedBuilder.setMinimumNumberShouldMatch(5);
92+
93+
for (int i = 0; i < 600; i++) {
94+
nestedBuilder.add(new TermQuery(new Term("foo", "bar-" + i)), BooleanClause.Occur.SHOULD);
95+
}
96+
Query inner = nestedBuilder.build();
97+
BooleanQuery.Builder builderMixed = new BooleanQuery.Builder();
98+
99+
builderMixed.setMinimumNumberShouldMatch(5);
100+
101+
for (int i = 0; i < 600; i++) {
102+
builderMixed.add(new TermQuery(new Term("foo", "bar")), BooleanClause.Occur.SHOULD);
103+
}
104+
105+
builderMixed.add(inner, BooleanClause.Occur.SHOULD);
106+
107+
Query query = builderMixed.build();
108+
109+
// Can't be flattened, but high clause count should still be cause during nested walk...
110+
IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> { searcher.rewrite(query); });
111+
assertThat(ex.getCause(), instanceOf(IndexSearcher.TooManyNestedClauses.class));
112+
}
113+
114+
public void testLargeDisjunctionMaxQuery() throws IOException {
115+
IndexSearcher searcher = newContextIndexSearcher(new MultiReader());
116+
Query[] clausesQueryArray = new Query[1050];
117+
118+
for (int i = 0; i < 1049; i++) {
119+
clausesQueryArray[i] = new TermQuery(new Term("field", "a"));
120+
}
121+
122+
PhraseQuery pq = new PhraseQuery("field", new String[0]);
123+
124+
clausesQueryArray[1049] = pq;
125+
126+
DisjunctionMaxQuery dmq = new DisjunctionMaxQuery(Arrays.asList(clausesQueryArray), 0.5f);
127+
128+
// Can't be flattened, but high clause count should still be cause during nested walk...
129+
IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> { searcher.rewrite(dmq); });
130+
assertThat(ex.getCause(), instanceOf(IndexSearcher.TooManyNestedClauses.class));
131+
}
132+
133+
public void testMultiExactWithRepeats() throws IOException {
134+
IndexSearcher searcher = newContextIndexSearcher(new MultiReader());
135+
MultiPhraseQuery.Builder qb = new MultiPhraseQuery.Builder();
136+
137+
for (int i = 0; i < 1050; i++) {
138+
qb.add(new Term[] { new Term("foo", "bar-" + i), new Term("foo", "bar+" + i) }, 0);
139+
}
140+
141+
// Can't be flattened, but high clause count should still be cause during nested walk...
142+
IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> { searcher.rewrite(qb.build()); });
143+
assertThat(ex.getCause(), instanceOf(IndexSearcher.TooManyNestedClauses.class));
144+
}
145+
146+
public void testIndexOrDocValues() throws IOException {
147+
Directory dir = newDirectory();
148+
RandomIndexWriter writer = new RandomIndexWriter(random(), dir);
149+
Document d = new Document();
150+
d.add(new LongField("foo", 0L, LongField.Store.NO));
151+
writer.addDocument(d);
152+
d = new Document();
153+
d.add(new LongField("foo", Long.MAX_VALUE, LongField.Store.NO));
154+
writer.addDocument(d);
155+
156+
IndexReader reader = writer.getReader();
157+
IndexSearcher searcher = newContextIndexSearcher(reader);
158+
writer.close();
159+
int maxClauseCount = IndexSearcher.getMaxClauseCount();
160+
BooleanQuery.Builder qb = new BooleanQuery.Builder();
161+
162+
for (int i = 0; i < maxClauseCount; i++) {
163+
qb.add(LongPoint.newRangeQuery("foo", 0, i), BooleanClause.Occur.SHOULD);
164+
}
165+
// should not throw an exception, because it is below the limit
166+
searcher.rewrite(qb.build());
167+
168+
qb = new BooleanQuery.Builder();
169+
for (int i = 0; i < maxClauseCount; i++) {
170+
qb.add(NumericField.newRangeLongQuery("foo", 0, i), BooleanClause.Occur.SHOULD);
171+
}
172+
// should not throw an exception, because it is below the limit
173+
searcher.rewrite(qb.build());
174+
175+
qb = new BooleanQuery.Builder();
176+
for (int i = 0; i < maxClauseCount; i++) {
177+
qb.add(
178+
new IndexOrDocValuesQuery(LongPoint.newRangeQuery("foo", 0, i), SortedNumericDocValuesField.newSlowRangeQuery("foo", 0, i)),
179+
BooleanClause.Occur.SHOULD
180+
);
181+
}
182+
// should not throw an exception, because it is below the limit
183+
searcher.rewrite(qb.build());
184+
185+
reader.close();
186+
dir.close();
187+
}
188+
189+
private static IndexSearcher newContextIndexSearcher(IndexReader reader) throws IOException {
190+
return new ContextIndexSearcher(
191+
reader,
192+
IndexSearcher.getDefaultSimilarity(),
193+
IndexSearcher.getDefaultQueryCache(),
194+
IndexSearcher.getDefaultQueryCachingPolicy(),
195+
false,
196+
null,
197+
-1,
198+
1
199+
);
200+
}
201+
}

0 commit comments

Comments
 (0)