Skip to content

Commit 847ef04

Browse files
atrisatris
committed
Update writers and get tests to pass
Signed-off-by: Atri Sharma <atri.jiit@gmail.com>
1 parent 478e22a commit 847ef04

File tree

6 files changed

+405
-171
lines changed

6 files changed

+405
-171
lines changed

‎server/src/main/java/org/opensearch/search/query/BooleanFlatteningRewriter.java‎

Lines changed: 110 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -37,16 +37,73 @@ public QueryBuilder rewrite(QueryBuilder query, QueryShardContext context) {
3737
}
3838

3939
BoolQueryBuilder boolQuery = (BoolQueryBuilder) query;
40+
41+
// First check if flattening is needed
42+
if (!needsFlattening(boolQuery)) {
43+
return query;
44+
}
45+
4046
return flattenBoolQuery(boolQuery);
4147
}
4248

49+
private boolean needsFlattening(BoolQueryBuilder boolQuery) {
50+
// Check all clause types for nested bool queries that can be flattened
51+
if (hasFlattenableBool(boolQuery.must(), ClauseType.MUST)
52+
|| hasFlattenableBool(boolQuery.filter(), ClauseType.FILTER)
53+
|| hasFlattenableBool(boolQuery.should(), ClauseType.SHOULD)
54+
|| hasFlattenableBool(boolQuery.mustNot(), ClauseType.MUST_NOT)) {
55+
return true;
56+
}
57+
58+
// Check if any nested bool queries need flattening
59+
return hasNestedBoolThatNeedsFlattening(boolQuery);
60+
}
61+
62+
private boolean hasFlattenableBool(List<QueryBuilder> clauses, ClauseType parentType) {
63+
for (QueryBuilder clause : clauses) {
64+
if (clause instanceof BoolQueryBuilder) {
65+
BoolQueryBuilder nestedBool = (BoolQueryBuilder) clause;
66+
// Can flatten if nested bool only has one type of clause matching parent
67+
if (canFlatten(nestedBool, parentType)) {
68+
return true;
69+
}
70+
}
71+
}
72+
return false;
73+
}
74+
75+
private boolean hasNestedBoolThatNeedsFlattening(BoolQueryBuilder boolQuery) {
76+
for (QueryBuilder clause : boolQuery.must()) {
77+
if (clause instanceof BoolQueryBuilder && needsFlattening((BoolQueryBuilder) clause)) {
78+
return true;
79+
}
80+
}
81+
for (QueryBuilder clause : boolQuery.filter()) {
82+
if (clause instanceof BoolQueryBuilder && needsFlattening((BoolQueryBuilder) clause)) {
83+
return true;
84+
}
85+
}
86+
for (QueryBuilder clause : boolQuery.should()) {
87+
if (clause instanceof BoolQueryBuilder && needsFlattening((BoolQueryBuilder) clause)) {
88+
return true;
89+
}
90+
}
91+
for (QueryBuilder clause : boolQuery.mustNot()) {
92+
if (clause instanceof BoolQueryBuilder && needsFlattening((BoolQueryBuilder) clause)) {
93+
return true;
94+
}
95+
}
96+
return false;
97+
}
98+
4399
private BoolQueryBuilder flattenBoolQuery(BoolQueryBuilder original) {
44100
BoolQueryBuilder flattened = new BoolQueryBuilder();
45101

46102
flattened.boost(original.boost());
47103
flattened.queryName(original.queryName());
48104
flattened.minimumShouldMatch(original.minimumShouldMatch());
49105
flattened.adjustPureNegative(original.adjustPureNegative());
106+
50107
flattenClauses(original.must(), flattened, ClauseType.MUST);
51108
flattenClauses(original.filter(), flattened, ClauseType.FILTER);
52109
flattenClauses(original.should(), flattened, ClauseType.SHOULD);
@@ -55,101 +112,86 @@ private BoolQueryBuilder flattenBoolQuery(BoolQueryBuilder original) {
55112
return flattened;
56113
}
57114

58-
private void flattenClauses(List<QueryBuilder> clauses, BoolQueryBuilder target, ClauseType type) {
115+
private void flattenClauses(List<QueryBuilder> clauses, BoolQueryBuilder target, ClauseType clauseType) {
59116
for (QueryBuilder clause : clauses) {
60-
if (clause instanceof BoolQueryBuilder && canFlatten((BoolQueryBuilder) clause, type)) {
61-
BoolQueryBuilder innerBool = (BoolQueryBuilder) clause;
62-
if (type == ClauseType.FILTER) {
63-
for (QueryBuilder innerClause : innerBool.filter()) {
64-
target.filter(processClause(innerClause));
65-
}
66-
if (!innerBool.must().isEmpty() || !innerBool.should().isEmpty() || !innerBool.mustNot().isEmpty()) {
67-
BoolQueryBuilder preservedBool = new BoolQueryBuilder();
68-
preservedBool.boost(innerBool.boost());
69-
flattenClauses(innerBool.must(), preservedBool, ClauseType.MUST);
70-
flattenClauses(innerBool.should(), preservedBool, ClauseType.SHOULD);
71-
flattenClauses(innerBool.mustNot(), preservedBool, ClauseType.MUST_NOT);
72-
if (preservedBool.hasClauses()) {
73-
target.filter(preservedBool);
117+
if (clause instanceof BoolQueryBuilder) {
118+
BoolQueryBuilder nestedBool = (BoolQueryBuilder) clause;
119+
120+
if (canFlatten(nestedBool, clauseType)) {
121+
// Flatten the nested bool query by extracting its clauses
122+
List<QueryBuilder> nestedClauses = getClausesForType(nestedBool, clauseType);
123+
for (QueryBuilder nestedClause : nestedClauses) {
124+
// Recursively flatten if needed
125+
if (nestedClause instanceof BoolQueryBuilder) {
126+
nestedClause = flattenBoolQuery((BoolQueryBuilder) nestedClause);
74127
}
128+
addClauseBasedOnType(target, nestedClause, clauseType);
75129
}
76130
} else {
77-
addClausesBasedOnType(innerBool, target, type);
131+
// Can't flatten this bool, but recursively flatten its contents
132+
BoolQueryBuilder flattenedNested = flattenBoolQuery(nestedBool);
133+
addClauseBasedOnType(target, flattenedNested, clauseType);
78134
}
79135
} else {
80-
addClauseToTarget(target, processClause(clause), type);
136+
// Non-boolean clause, add as-is
137+
addClauseBasedOnType(target, clause, clauseType);
81138
}
82139
}
83140
}
84141

85-
private QueryBuilder processClause(QueryBuilder clause) {
86-
if (clause instanceof BoolQueryBuilder) {
87-
return flattenBoolQuery((BoolQueryBuilder) clause);
88-
}
89-
return clause;
90-
}
142+
private boolean canFlatten(BoolQueryBuilder nestedBool, ClauseType parentType) {
143+
// Can only flatten if:
144+
// 1. The nested bool has the same properties as default (boost=1, no queryName, etc.)
145+
// 2. The nested bool only has clauses of the same type as the parent
91146

92-
private boolean canFlatten(BoolQueryBuilder boolQuery, ClauseType parentType) {
93-
if (boolQuery.boost() != 1.0f || boolQuery.minimumShouldMatch() != null) {
147+
if (nestedBool.boost() != 1.0f || nestedBool.queryName() != null) {
94148
return false;
95149
}
96150

97-
if (parentType == ClauseType.FILTER) {
98-
return !boolQuery.filter().isEmpty();
99-
}
100-
int clauseTypeCount = 0;
101-
if (!boolQuery.must().isEmpty()) clauseTypeCount++;
102-
if (!boolQuery.filter().isEmpty()) clauseTypeCount++;
103-
if (!boolQuery.should().isEmpty()) clauseTypeCount++;
104-
if (!boolQuery.mustNot().isEmpty()) clauseTypeCount++;
105-
106-
return clauseTypeCount == 1;
107-
}
108-
109-
private void addClausesBasedOnType(BoolQueryBuilder source, BoolQueryBuilder target, ClauseType type) {
110-
List<QueryBuilder> relevantClauses = getClausesForType(source, type);
111-
112-
boolean hasOnlyThisClauseType = (type == ClauseType.MUST
113-
&& !source.must().isEmpty()
114-
&& source.filter().isEmpty()
115-
&& source.should().isEmpty()
116-
&& source.mustNot().isEmpty())
117-
|| (type == ClauseType.SHOULD
118-
&& source.must().isEmpty()
119-
&& source.filter().isEmpty()
120-
&& !source.should().isEmpty()
121-
&& source.mustNot().isEmpty())
122-
|| (type == ClauseType.MUST_NOT
123-
&& source.must().isEmpty()
124-
&& source.filter().isEmpty()
125-
&& source.should().isEmpty()
126-
&& !source.mustNot().isEmpty());
127-
128-
if (hasOnlyThisClauseType && !relevantClauses.isEmpty()) {
129-
for (QueryBuilder clause : relevantClauses) {
130-
addClauseToTarget(target, processClause(clause), type);
131-
}
132-
} else {
133-
addClauseToTarget(target, flattenBoolQuery(source), type);
151+
// Check if only has clauses matching parent type
152+
switch (parentType) {
153+
case MUST:
154+
return !nestedBool.must().isEmpty()
155+
&& nestedBool.filter().isEmpty()
156+
&& nestedBool.should().isEmpty()
157+
&& nestedBool.mustNot().isEmpty();
158+
case FILTER:
159+
return nestedBool.must().isEmpty()
160+
&& !nestedBool.filter().isEmpty()
161+
&& nestedBool.should().isEmpty()
162+
&& nestedBool.mustNot().isEmpty();
163+
case SHOULD:
164+
return nestedBool.must().isEmpty()
165+
&& nestedBool.filter().isEmpty()
166+
&& !nestedBool.should().isEmpty()
167+
&& nestedBool.mustNot().isEmpty()
168+
&& nestedBool.minimumShouldMatch() == null;
169+
case MUST_NOT:
170+
return nestedBool.must().isEmpty()
171+
&& nestedBool.filter().isEmpty()
172+
&& nestedBool.should().isEmpty()
173+
&& !nestedBool.mustNot().isEmpty();
174+
default:
175+
return false;
134176
}
135177
}
136178

137-
private List<QueryBuilder> getClausesForType(BoolQueryBuilder source, ClauseType type) {
179+
private List<QueryBuilder> getClausesForType(BoolQueryBuilder bool, ClauseType type) {
138180
switch (type) {
139181
case MUST:
140-
return source.must();
182+
return bool.must();
141183
case FILTER:
142-
return source.filter();
184+
return bool.filter();
143185
case SHOULD:
144-
return source.should();
186+
return bool.should();
145187
case MUST_NOT:
146-
return source.mustNot();
188+
return bool.mustNot();
147189
default:
148190
return new ArrayList<>();
149191
}
150192
}
151193

152-
private void addClauseToTarget(BoolQueryBuilder target, QueryBuilder clause, ClauseType type) {
194+
private void addClauseBasedOnType(BoolQueryBuilder target, QueryBuilder clause, ClauseType type) {
153195
switch (type) {
154196
case MUST:
155197
target.must(clause);

0 commit comments

Comments
 (0)