Skip to content

Commit

Permalink
Added facet support to the percolate api.
Browse files Browse the repository at this point in the history
  • Loading branch information
martijnvg committed Oct 16, 2013
1 parent b7c4ade commit 294ceb9
Show file tree
Hide file tree
Showing 9 changed files with 454 additions and 97 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.FilterBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.search.facet.FacetBuilder;
import org.elasticsearch.search.highlight.HighlightBuilder;

import java.util.Map;
Expand Down Expand Up @@ -104,16 +105,68 @@ public PercolateRequestBuilder setSize(int size) {
return this;
}

/**
* Similar as {@link #setScore(boolean)}, but also sort by the score.
*/
public PercolateRequestBuilder setSort(boolean sort) {
sourceBuilder().setSort(sort);
return this;
}

/**
* Whether to compute a score for each match and include it in the response. The score is based on
* {@link #setPercolateQuery(QueryBuilder)}}.
*/
public PercolateRequestBuilder setScore(boolean score) {
sourceBuilder().setScore(score);
return this;
}

/**
* Sets a query to reduce the number of percolate queries to be evaluated and score the queries that match based
* on this query.
*/
public PercolateRequestBuilder setPercolateDoc(PercolateSourceBuilder.DocBuilder docBuilder) {
sourceBuilder().setDoc(docBuilder);
return this;
}

/**
* Sets a query to reduce the number of percolate queries to be evaluated and score the queries that match based
* on this query.
*/
public PercolateRequestBuilder setPercolateQuery(QueryBuilder queryBuilder) {
sourceBuilder().setQueryBuilder(queryBuilder);
return this;
}

/**
* Sets a filter to reduce the number of percolate queries to be evaluated.
*/
public PercolateRequestBuilder setPercolateFilter(FilterBuilder filterBuilder) {
sourceBuilder().setFilterBuilder(filterBuilder);
return this;
}

/**
* Enables highlighting for the percolate document. Per matched percolate query highlight the percolate document.
*/
public PercolateRequestBuilder setHighlightBuilder(HighlightBuilder highlightBuilder) {
sourceBuilder().setHighlightBuilder(highlightBuilder);
return this;
}

/**
* Add a facet definition.
*/
public PercolateRequestBuilder addFacet(FacetBuilder facetBuilder) {
sourceBuilder().addFacet(facetBuilder);
return this;
}

/**
* Sets the raw percolate request body.
*/
public PercolateRequestBuilder setSource(PercolateSourceBuilder source) {
sourceBuilder = source;
return this;
Expand Down Expand Up @@ -164,26 +217,6 @@ public PercolateRequestBuilder setSource(byte[] source, int offset, int length,
return this;
}

public PercolateRequestBuilder setPercolateDoc(PercolateSourceBuilder.DocBuilder docBuilder) {
sourceBuilder().setDoc(docBuilder);
return this;
}

public PercolateRequestBuilder setPercolateQuery(QueryBuilder queryBuilder) {
sourceBuilder().setQueryBuilder(queryBuilder);
return this;
}

public PercolateRequestBuilder setPercolateFilter(FilterBuilder filterBuilder) {
sourceBuilder().setFilterBuilder(filterBuilder);
return this;
}

public PercolateRequestBuilder setHighlightBuilder(HighlightBuilder highlightBuilder) {
sourceBuilder().setHighlightBuilder(highlightBuilder);
return this;
}

private PercolateSourceBuilder sourceBuilder() {
if (sourceBuilder == null) {
sourceBuilder = new PercolateSourceBuilder();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import org.elasticsearch.common.xcontent.XContentBuilderString;
import org.elasticsearch.percolator.PercolatorService;
import org.elasticsearch.rest.action.support.RestActions;
import org.elasticsearch.search.facet.InternalFacets;
import org.elasticsearch.search.highlight.HighlightField;

import java.io.IOException;
Expand All @@ -46,20 +47,15 @@ public class PercolateResponse extends BroadcastOperationResponse implements Ite
private long tookInMillis;
private Match[] matches;
private long count;
private InternalFacets facets;

public PercolateResponse(int totalShards, int successfulShards, int failedShards, List<ShardOperationFailedException> shardFailures,
Match[] matches, long count, long tookInMillis) {
Match[] matches, long count, long tookInMillis, InternalFacets facets) {
super(totalShards, successfulShards, failedShards, shardFailures);
this.tookInMillis = tookInMillis;
this.matches = matches;
this.count = count;
}

public PercolateResponse(int totalShards, int successfulShards, int failedShards, List<ShardOperationFailedException> shardFailures, long count, long tookInMillis) {
super(totalShards, successfulShards, failedShards, shardFailures);
this.tookInMillis = tookInMillis;
this.matches = EMPTY;
this.count = count;
this.facets = facets;
}

public PercolateResponse(int totalShards, int successfulShards, int failedShards, List<ShardOperationFailedException> shardFailures, long tookInMillis) {
Expand Down Expand Up @@ -97,6 +93,10 @@ public long getCount() {
return count;
}

public InternalFacets getFacets() {
return facets;
}

@Override
public Iterator<Match> iterator() {
return Arrays.asList(matches).iterator();
Expand Down Expand Up @@ -163,6 +163,9 @@ public void readFrom(StreamInput in) throws IOException {
matches[i] = new Match();
matches[i].readFrom(in);
}
if (in.readBoolean()) {
facets = InternalFacets.readFacets(in);
}
}

@Override
Expand All @@ -174,6 +177,12 @@ public void writeTo(StreamOutput out) throws IOException {
for (Match match : matches) {
match.writeTo(out);
}
if (facets != null) {
out.writeBoolean(true);
facets.writeTo(out);
} else {
out.writeBoolean(false);
}
}

public static class Match implements Streamable {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.percolator.PercolateContext;
import org.elasticsearch.search.facet.InternalFacets;
import org.elasticsearch.search.highlight.HighlightField;

import java.io.IOException;
Expand All @@ -45,6 +46,8 @@ public class PercolateShardResponse extends BroadcastShardOperationResponse {
private byte percolatorTypeId;
private int requestedSize;

private InternalFacets facets;

PercolateShardResponse() {
}

Expand All @@ -56,6 +59,7 @@ public PercolateShardResponse(BytesRef[] matches, List<Map<String, HighlightFiel
this.scores = scores;
this.percolatorTypeId = context.percolatorTypeId;
this.requestedSize = context.size;
buildFacets(context);
}

public PercolateShardResponse(BytesRef[] matches, long count, float[] scores, PercolateContext context, String index, int shardId) {
Expand All @@ -65,6 +69,7 @@ public PercolateShardResponse(BytesRef[] matches, long count, float[] scores, Pe
this.scores = scores;
this.percolatorTypeId = context.percolatorTypeId;
this.requestedSize = context.size;
buildFacets(context);
}

public PercolateShardResponse(BytesRef[] matches, List<Map<String, HighlightField>> hls, long count, PercolateContext context, String index, int shardId) {
Expand All @@ -75,6 +80,7 @@ public PercolateShardResponse(BytesRef[] matches, List<Map<String, HighlightFiel
this.count = count;
this.percolatorTypeId = context.percolatorTypeId;
this.requestedSize = context.size;
buildFacets(context);
}

public PercolateShardResponse(long count, PercolateContext context, String index, int shardId) {
Expand All @@ -84,6 +90,7 @@ public PercolateShardResponse(long count, PercolateContext context, String index
this.scores = new float[0];
this.percolatorTypeId = context.percolatorTypeId;
this.requestedSize = context.size;
buildFacets(context);
}

public PercolateShardResponse(PercolateContext context, String index, int shardId) {
Expand Down Expand Up @@ -113,6 +120,10 @@ public List<Map<String, HighlightField>> hls() {
return hls;
}

public InternalFacets facets() {
return facets;
}

public byte percolatorTypeId() {
return percolatorTypeId;
}
Expand Down Expand Up @@ -144,6 +155,9 @@ public void readFrom(StreamInput in) throws IOException {
}
hls.add(fields);
}
if (in.readBoolean()) {
facets = InternalFacets.readFacets(in);
}
}

@Override
Expand All @@ -168,5 +182,17 @@ public void writeTo(StreamOutput out) throws IOException {
entry.getValue().writeTo(out);
}
}
if (facets != null) {
out.writeBoolean(true);
facets.writeTo(out);
} else {
out.writeBoolean(false);
}
}

private void buildFacets(PercolateContext context) {
if (context.queryResult() != null && context.queryResult().facets() != null) {
this.facets = new InternalFacets(context.queryResult().facets().facets());
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,20 +19,24 @@

package org.elasticsearch.action.percolate;

import com.google.common.collect.Lists;
import org.elasticsearch.ElasticSearchGenerationException;
import org.elasticsearch.common.bytes.BytesArray;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.xcontent.*;
import org.elasticsearch.index.query.FilterBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilderException;
import org.elasticsearch.search.facet.FacetBuilder;
import org.elasticsearch.search.highlight.HighlightBuilder;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
* Builder to create the percolate request body.
*/
public class PercolateSourceBuilder implements ToXContent {

Expand All @@ -43,6 +47,7 @@ public class PercolateSourceBuilder implements ToXContent {
private Boolean sort;
private Boolean score;
private HighlightBuilder highlightBuilder;
private List<FacetBuilder> facets;

public DocBuilder percolateDocument() {
if (docBuilder == null) {
Expand All @@ -55,6 +60,9 @@ public DocBuilder getDoc() {
return docBuilder;
}

/**
* Sets the document to run the percolate queries against.
*/
public PercolateSourceBuilder setDoc(DocBuilder docBuilder) {
this.docBuilder = docBuilder;
return this;
Expand All @@ -64,6 +72,10 @@ public QueryBuilder getQueryBuilder() {
return queryBuilder;
}

/**
* Sets a query to reduce the number of percolate queries to be evaluated and score the queries that match based
* on this query.
*/
public PercolateSourceBuilder setQueryBuilder(QueryBuilder queryBuilder) {
this.queryBuilder = queryBuilder;
return this;
Expand All @@ -73,31 +85,58 @@ public FilterBuilder getFilterBuilder() {
return filterBuilder;
}

/**
* Sets a filter to reduce the number of percolate queries to be evaluated.
*/
public PercolateSourceBuilder setFilterBuilder(FilterBuilder filterBuilder) {
this.filterBuilder = filterBuilder;
return this;
}

/**
* Limits the maximum number of percolate query matches to be returned.
*/
public PercolateSourceBuilder setSize(int size) {
this.size = size;
return this;
}

/**
* Similar as {@link #setScore(boolean)}, but also sort by the score.
*/
public PercolateSourceBuilder setSort(boolean sort) {
this.sort = sort;
return this;
}

/**
* Whether to compute a score for each match and include it in the response. The score is based on
* {@link #setQueryBuilder(QueryBuilder)}.
*/
public PercolateSourceBuilder setScore(boolean score) {
this.score = score;
return this;
}

/**
* Enables highlighting for the percolate document. Per matched percolate query highlight the percolate document.
*/
public PercolateSourceBuilder setHighlightBuilder(HighlightBuilder highlightBuilder) {
this.highlightBuilder = highlightBuilder;
return this;
}

/**
* Add a facet definition.
*/
public PercolateSourceBuilder addFacet(FacetBuilder facetBuilder) {
if (facets == null) {
facets = Lists.newArrayList();
}
facets.add(facetBuilder);
return this;
}

public BytesReference buildAsBytes(XContentType contentType) throws SearchSourceBuilderException {
try {
XContentBuilder builder = XContentFactory.contentBuilder(contentType);
Expand Down Expand Up @@ -134,6 +173,14 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
if (highlightBuilder != null) {
highlightBuilder.toXContent(builder, params);
}
if (facets != null) {
builder.field("facets");
builder.startObject();
for (FacetBuilder facet : facets) {
facet.toXContent(builder, params);
}
builder.endObject();
}
builder.endObject();
return builder;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,7 @@ public static PercolateResponse reduce(PercolateRequest request, AtomicReference
long tookInMillis = System.currentTimeMillis() - request.startTime;
return new PercolateResponse(
shardsResponses.length(), successfulShards, failedShards, shardFailures,
result.matches(), result.count(), tookInMillis
result.matches(), result.count(), tookInMillis, result.reducedFacets()
);
}
}
Expand Down
Loading

0 comments on commit 294ceb9

Please sign in to comment.