Skip to content

Commit b6a83fd

Browse files
s1monwkimchy
authored andcommitted
elastic#2332 support CustomScoreQuery highlighting and expose QueryBuilder on CustomScoreQuery via Java API
1 parent 2c78505 commit b6a83fd

File tree

5 files changed

+158
-29
lines changed

5 files changed

+158
-29
lines changed

src/main/java/org/elasticsearch/index/query/ConstantScoreQueryBuilder.java

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,10 @@
3232
public class ConstantScoreQueryBuilder extends BaseQueryBuilder implements BoostableQueryBuilder<ConstantScoreQueryBuilder> {
3333

3434
private final FilterBuilder filterBuilder;
35+
private final QueryBuilder queryBuilder;
3536

3637
private float boost = -1;
38+
3739

3840
/**
3941
* A query that wraps a filter and simply returns a constant score equal to the
@@ -43,7 +45,18 @@ public class ConstantScoreQueryBuilder extends BaseQueryBuilder implements Boost
4345
*/
4446
public ConstantScoreQueryBuilder(FilterBuilder filterBuilder) {
4547
this.filterBuilder = filterBuilder;
48+
this.queryBuilder = null;
4649
}
50+
/**
51+
* A query that wraps a query and simply returns a constant score equal to the
52+
* query boost for every document in the query.
53+
*
54+
* @param queryBuilder The query to wrap in a constant score query
55+
*/
56+
public ConstantScoreQueryBuilder(QueryBuilder queryBuilder) {
57+
this.filterBuilder = null;
58+
this.queryBuilder = queryBuilder;
59+
}
4760

4861
/**
4962
* Sets the boost for this query. Documents matching this query will (in addition to the normal
@@ -57,8 +70,15 @@ public ConstantScoreQueryBuilder boost(float boost) {
5770
@Override
5871
protected void doXContent(XContentBuilder builder, Params params) throws IOException {
5972
builder.startObject(ConstantScoreQueryParser.NAME);
60-
builder.field("filter");
61-
filterBuilder.toXContent(builder, params);
73+
if (queryBuilder != null) {
74+
assert filterBuilder == null;
75+
builder.field("query");
76+
queryBuilder.toXContent(builder, params);
77+
} else {
78+
builder.field("filter");
79+
filterBuilder.toXContent(builder, params);
80+
}
81+
6282
if (boost != -1) {
6383
builder.field("boost", boost);
6484
}

src/main/java/org/elasticsearch/index/query/QueryBuilders.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,16 @@ public static FilteredQueryBuilder filteredQuery(QueryBuilder queryBuilder, @Nul
454454
public static ConstantScoreQueryBuilder constantScoreQuery(FilterBuilder filterBuilder) {
455455
return new ConstantScoreQueryBuilder(filterBuilder);
456456
}
457+
458+
/**
459+
* A query that wraps another query and simply returns a constant score equal to the
460+
* query boost for every document in the query.
461+
*
462+
* @param queryBuilder The query to wrap in a constant score query
463+
*/
464+
public static ConstantScoreQueryBuilder constantScoreQuery(QueryBuilder queryBuilder) {
465+
return new ConstantScoreQueryBuilder(queryBuilder);
466+
}
457467

458468
/**
459469
* A query that simply applies the boost fact to the wrapped query (multiplies it).
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
/*
2+
* Licensed to ElasticSearch and Shay Banon under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. ElasticSearch licenses this
6+
* file to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.search.highlight;
21+
22+
import java.io.IOException;
23+
import java.util.Map;
24+
25+
import org.apache.lucene.index.IndexReader;
26+
import org.apache.lucene.search.ConstantScoreQuery;
27+
import org.apache.lucene.search.FilteredQuery;
28+
import org.apache.lucene.search.Query;
29+
import org.apache.lucene.search.highlight.QueryScorer;
30+
import org.apache.lucene.search.highlight.WeightedSpanTerm;
31+
import org.apache.lucene.search.highlight.WeightedSpanTermExtractor;
32+
import org.elasticsearch.common.lucene.search.function.FiltersFunctionScoreQuery;
33+
import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
34+
35+
public final class CustomQueryScorer extends QueryScorer {
36+
37+
public CustomQueryScorer(Query query, IndexReader reader, String field,
38+
String defaultField) {
39+
super(query, reader, field, defaultField);
40+
}
41+
42+
public CustomQueryScorer(Query query, IndexReader reader, String field) {
43+
super(query, reader, field);
44+
}
45+
46+
public CustomQueryScorer(Query query, String field, String defaultField) {
47+
super(query, field, defaultField);
48+
}
49+
50+
public CustomQueryScorer(Query query, String field) {
51+
super(query, field);
52+
}
53+
54+
public CustomQueryScorer(Query query) {
55+
super(query);
56+
}
57+
58+
public CustomQueryScorer(WeightedSpanTerm[] weightedTerms) {
59+
super(weightedTerms);
60+
}
61+
62+
@Override
63+
protected WeightedSpanTermExtractor newTermExtractor(String defaultField) {
64+
return defaultField == null ? new CustomWeightedSpanTermExtractor()
65+
: new CustomWeightedSpanTermExtractor(defaultField);
66+
}
67+
68+
private static class CustomWeightedSpanTermExtractor extends WeightedSpanTermExtractor {
69+
70+
public CustomWeightedSpanTermExtractor() {
71+
super();
72+
}
73+
74+
public CustomWeightedSpanTermExtractor(String defaultField) {
75+
super(defaultField);
76+
}
77+
78+
@Override
79+
protected void extractUnknownQuery(Query query,
80+
Map<String, WeightedSpanTerm> terms) throws IOException {
81+
if (query instanceof FunctionScoreQuery) {
82+
query = ((FunctionScoreQuery) query).getSubQuery();
83+
extract(query, terms);
84+
} else if (query instanceof FiltersFunctionScoreQuery) {
85+
query = ((FiltersFunctionScoreQuery) query).getSubQuery();
86+
extract(query, terms);
87+
} else if (query instanceof ConstantScoreQuery) {
88+
ConstantScoreQuery q = (ConstantScoreQuery) query;
89+
if (q.getQuery() != null) {
90+
query = q.getQuery();
91+
extract(query, terms);
92+
}
93+
} else if (query instanceof FilteredQuery) {
94+
query = ((FilteredQuery) query).getQuery();
95+
extract(query, terms);
96+
}
97+
}
98+
99+
}
100+
101+
}

src/main/java/org/elasticsearch/search/highlight/HighlightPhase.java

Lines changed: 1 addition & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -138,33 +138,8 @@ public void hitExecute(SearchContext context, HitContext hitContext) throws Elas
138138
// Don't use the context.query() since it might be rewritten, and we need to pass the non rewritten queries to
139139
// let the highlighter handle MultiTerm ones
140140

141-
// QueryScorer uses WeightedSpanTermExtractor to extract terms, but we can't really plug into
142-
// it, so, we hack here (and really only support top level queries)
143141
Query query = context.parsedQuery().query();
144-
while (true) {
145-
boolean extracted = false;
146-
if (query instanceof FunctionScoreQuery) {
147-
query = ((FunctionScoreQuery) query).getSubQuery();
148-
extracted = true;
149-
} else if (query instanceof FiltersFunctionScoreQuery) {
150-
query = ((FiltersFunctionScoreQuery) query).getSubQuery();
151-
extracted = true;
152-
} else if (query instanceof ConstantScoreQuery) {
153-
ConstantScoreQuery q = (ConstantScoreQuery) query;
154-
if (q.getQuery() != null) {
155-
query = q.getQuery();
156-
extracted = true;
157-
}
158-
} else if (query instanceof FilteredQuery) {
159-
query = ((FilteredQuery) query).getQuery();
160-
extracted = true;
161-
}
162-
if (!extracted) {
163-
break;
164-
}
165-
}
166-
167-
QueryScorer queryScorer = new QueryScorer(query, field.requireFieldMatch() ? mapper.names().indexName() : null);
142+
QueryScorer queryScorer = new CustomQueryScorer(query, field.requireFieldMatch() ? mapper.names().indexName() : null);
168143
queryScorer.setExpandMultiTermQuery(true);
169144
Fragmenter fragmenter;
170145
if (field.numberOfFragments() == 0) {

src/test/java/org/elasticsearch/test/integration/search/highlight/HighlighterSearchTests.java

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -312,8 +312,31 @@ public void testPlainHighlighter() throws Exception {
312312
assertThat(searchResponse.hits().totalHits(), equalTo(1l));
313313

314314
assertThat(searchResponse.hits().getAt(0).highlightFields().get("field2").fragments()[0].string(), equalTo("The <xxx>quick</xxx> brown fox jumps over the lazy dog"));
315-
}
315+
316+
logger.info("--> searching on _all with constant score, highlighting on field2");
317+
source = searchSource()
318+
.query(constantScoreQuery(prefixQuery("_all", "qui")))
319+
.from(0).size(60).explain(true)
320+
.highlight(highlight().field("field2").order("score").preTags("<xxx>").postTags("</xxx>"));
321+
322+
searchResponse = client.search(searchRequest("test").source(source).searchType(QUERY_THEN_FETCH).scroll(timeValueMinutes(10))).actionGet();
323+
assertThat("Failures " + Arrays.toString(searchResponse.shardFailures()), searchResponse.shardFailures().length, equalTo(0));
324+
assertThat(searchResponse.hits().totalHits(), equalTo(1l));
325+
326+
assertThat(searchResponse.hits().getAt(0).highlightFields().get("field2").fragments()[0].string(), equalTo("The <xxx>quick</xxx> brown fox jumps over the lazy dog"));
316327

328+
logger.info("--> searching on _all with constant score, highlighting on field2");
329+
source = searchSource()
330+
.query(boolQuery().should(constantScoreQuery(prefixQuery("_all", "qui"))))
331+
.from(0).size(60).explain(true)
332+
.highlight(highlight().field("field2").order("score").preTags("<xxx>").postTags("</xxx>"));
333+
334+
searchResponse = client.search(searchRequest("test").source(source).searchType(QUERY_THEN_FETCH).scroll(timeValueMinutes(10))).actionGet();
335+
assertThat("Failures " + Arrays.toString(searchResponse.shardFailures()), searchResponse.shardFailures().length, equalTo(0));
336+
assertThat(searchResponse.hits().totalHits(), equalTo(1l));
337+
assertThat(searchResponse.hits().getAt(0).highlightFields().get("field2").fragments()[0].string(), equalTo("The <xxx>quick</xxx> brown fox jumps over the lazy dog"));
338+
}
339+
317340
@Test
318341
public void testFastVectorHighlighter() throws Exception {
319342
try {

0 commit comments

Comments
 (0)