add script support for terms facets, allowing to control if a certain term will be included in the facets result (boolean result), or muging of the term

This commit is contained in:
kimchy 2010-08-08 10:27:06 +03:00
parent 566ae94478
commit ee1d50f8d8
5 changed files with 103 additions and 6 deletions

View File

@ -72,7 +72,7 @@ public class HistogramFacetCollectorParser implements FacetCollectorParser {
keyScript = parser.text();
} else if ("value_script".equals(fieldName) || "valueScript".equals(fieldName)) {
valueScript = parser.text();
} else if ("comparator".equals(fieldName)) {
} else if ("order".equals(fieldName) || "comparator".equals(fieldName)) {
comparatorType = HistogramFacet.ComparatorType.fromString(parser.text());
}
}

View File

@ -19,6 +19,7 @@
package org.elasticsearch.search.facets.terms;
import org.elasticsearch.common.collect.Maps;
import org.elasticsearch.common.regex.Regex;
import org.elasticsearch.common.xcontent.builder.XContentBuilder;
import org.elasticsearch.index.query.xcontent.XContentFilterBuilder;
@ -26,6 +27,7 @@ import org.elasticsearch.search.builder.SearchSourceBuilderException;
import org.elasticsearch.search.facets.AbstractFacetBuilder;
import java.io.IOException;
import java.util.Map;
/**
* @author kimchy (shay.banon)
@ -37,6 +39,8 @@ public class TermsFacetBuilder extends AbstractFacetBuilder {
private String regex;
private int regexFlags = 0;
private TermsFacet.ComparatorType comparatorType;
private String script;
private Map<String, Object> params;
public TermsFacetBuilder(String name) {
super(name);
@ -82,6 +86,19 @@ public class TermsFacetBuilder extends AbstractFacetBuilder {
return this;
}
public TermsFacetBuilder script(String script) {
this.script = script;
return this;
}
public TermsFacetBuilder param(String name, Object value) {
if (params == null) {
params = Maps.newHashMap();
}
params.put(name, value);
return this;
}
@Override public void toXContent(XContentBuilder builder, Params params) throws IOException {
if (fieldName == null) {
throw new SearchSourceBuilderException("field must be set on terms facet for facet [" + name + "]");
@ -107,6 +124,15 @@ public class TermsFacetBuilder extends AbstractFacetBuilder {
if (comparatorType != null) {
builder.field("order", comparatorType.name().toLowerCase());
}
if (script != null) {
builder.field("script", script);
if (this.params != null) {
builder.field("params");
builder.map(this.params);
}
}
builder.endObject();
addFilterFacetAndGlobal(builder, params);

View File

@ -23,19 +23,24 @@ import org.apache.lucene.index.IndexReader;
import org.elasticsearch.common.collect.BoundedTreeSet;
import org.elasticsearch.common.collect.ImmutableList;
import org.elasticsearch.common.collect.ImmutableSet;
import org.elasticsearch.common.collect.Maps;
import org.elasticsearch.common.thread.ThreadLocals;
import org.elasticsearch.common.trove.TObjectIntHashMap;
import org.elasticsearch.common.trove.TObjectIntIterator;
import org.elasticsearch.index.cache.field.data.FieldDataCache;
import org.elasticsearch.index.field.data.FieldData;
import org.elasticsearch.index.field.function.FieldsFunction;
import org.elasticsearch.index.field.function.script.ScriptFieldsFunction;
import org.elasticsearch.index.mapper.FieldMapper;
import org.elasticsearch.index.mapper.MapperService;
import org.elasticsearch.script.ScriptService;
import org.elasticsearch.search.facets.Facet;
import org.elasticsearch.search.facets.support.AbstractFacetCollector;
import java.io.IOException;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@ -73,7 +78,12 @@ public class TermsFacetCollector extends AbstractFacetCollector {
private final Pattern pattern;
public TermsFacetCollector(String facetName, String fieldName, int size, InternalTermsFacet.ComparatorType comparatorType, int numberOfShards, FieldDataCache fieldDataCache, MapperService mapperService, ImmutableSet<String> excluded, Pattern pattern) {
private final FieldsFunction scriptFunction;
private final Map<String, Object> params;
public TermsFacetCollector(String facetName, String fieldName, int size, InternalTermsFacet.ComparatorType comparatorType, int numberOfShards, FieldDataCache fieldDataCache, MapperService mapperService, ScriptService scriptService,
ImmutableSet<String> excluded, Pattern pattern, String script, Map<String, Object> params) {
super(facetName);
this.fieldDataCache = fieldDataCache;
this.size = size;
@ -91,15 +101,31 @@ public class TermsFacetCollector extends AbstractFacetCollector {
this.indexFieldName = fieldName;
this.fieldDataType = FieldData.Type.STRING;
}
if (excluded.isEmpty() && pattern == null) {
if (script != null) {
scriptFunction = new ScriptFieldsFunction(script, scriptService, mapperService, fieldDataCache);
if (params == null) {
this.params = Maps.newHashMapWithExpectedSize(1);
} else {
this.params = params;
}
} else {
this.params = null;
scriptFunction = null;
}
if (excluded.isEmpty() && pattern == null && scriptFunction == null) {
aggregator = new StaticAggregatorValueProc(popFacets());
} else {
aggregator = new AggregatorValueProc(popFacets(), excluded, pattern);
aggregator = new AggregatorValueProc(popFacets(), excluded, pattern, scriptFunction, params);
}
}
@Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException {
fieldData = fieldDataCache.cache(fieldDataType, reader, indexFieldName);
if (scriptFunction != null) {
scriptFunction.setNextReader(reader);
}
}
@Override protected void doCollect(int doc) throws IOException {
@ -147,10 +173,17 @@ public class TermsFacetCollector extends AbstractFacetCollector {
private final Matcher matcher;
public AggregatorValueProc(TObjectIntHashMap<String> facets, ImmutableSet<String> excluded, Pattern pattern) {
private final FieldsFunction scriptFunction;
private final Map<String, Object> params;
public AggregatorValueProc(TObjectIntHashMap<String> facets, ImmutableSet<String> excluded, Pattern pattern,
FieldsFunction scriptFunction, Map<String, Object> params) {
super(facets);
this.excluded = excluded;
this.matcher = pattern != null ? pattern.matcher("") : null;
this.scriptFunction = scriptFunction;
this.params = params;
}
@Override public void onValue(int docId, String value) {
@ -160,6 +193,17 @@ public class TermsFacetCollector extends AbstractFacetCollector {
if (matcher != null && !matcher.reset(value).matches()) {
return;
}
if (scriptFunction != null) {
params.put("term", value);
Object scriptValue = scriptFunction.execute(docId, params);
if (scriptValue instanceof Boolean) {
if (!((Boolean) scriptValue)) {
return;
}
} else {
value = scriptValue.toString();
}
}
super.onValue(docId, value);
}
}

View File

@ -27,6 +27,7 @@ import org.elasticsearch.search.facets.collector.FacetCollectorParser;
import org.elasticsearch.search.internal.SearchContext;
import java.io.IOException;
import java.util.Map;
import java.util.regex.Pattern;
/**
@ -50,9 +51,15 @@ public class TermsFacetCollectorParser implements FacetCollectorParser {
String regex = null;
String regexFlags = null;
TermsFacet.ComparatorType comparatorType = TermsFacet.ComparatorType.COUNT;
String script = null;
Map<String, Object> params = null;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
fieldName = parser.currentName();
} else if (token == XContentParser.Token.START_OBJECT) {
if ("params".equals(fieldName)) {
params = parser.map();
}
} else if (token == XContentParser.Token.START_ARRAY) {
if ("exclude".equals(fieldName)) {
ImmutableSet.Builder<String> builder = ImmutableSet.builder();
@ -72,6 +79,8 @@ public class TermsFacetCollectorParser implements FacetCollectorParser {
regexFlags = parser.text();
} else if ("order".equals(fieldName) || "comparator".equals(field)) {
comparatorType = TermsFacet.ComparatorType.fromString(parser.text());
} else if ("script".equals(fieldName)) {
script = parser.text();
}
}
}
@ -79,6 +88,7 @@ public class TermsFacetCollectorParser implements FacetCollectorParser {
if (regex != null) {
pattern = Regex.compile(regex, regexFlags);
}
return new TermsFacetCollector(facetName, field, size, comparatorType, context.numberOfShards(), context.fieldDataCache(), context.mapperService(), excluded, pattern);
return new TermsFacetCollector(facetName, field, size, comparatorType, context.numberOfShards(), context.fieldDataCache(), context.mapperService(), context.scriptService(),
excluded, pattern, script, params);
}
}

View File

@ -261,6 +261,23 @@ public class SimpleFacetsTests extends AbstractNodesTests {
assertThat(facet.entries().get(1).count(), equalTo(2));
assertThat(facet.entries().get(2).term(), equalTo("zzz"));
assertThat(facet.entries().get(2).count(), equalTo(1));
// Script
searchResponse = client.prepareSearch()
.setQuery(matchAllQuery())
.addFacet(termsFacet("facet1").field("tag").size(10).script("term + param1").param("param1", "a").order(TermsFacet.ComparatorType.TERM))
.execute().actionGet();
facet = searchResponse.facets().facet("facet1");
assertThat(facet.name(), equalTo("facet1"));
assertThat(facet.entries().size(), equalTo(3));
assertThat(facet.entries().get(0).term(), equalTo("xxxa"));
assertThat(facet.entries().get(0).count(), equalTo(1));
assertThat(facet.entries().get(1).term(), equalTo("yyya"));
assertThat(facet.entries().get(1).count(), equalTo(2));
assertThat(facet.entries().get(2).term(), equalTo("zzza"));
assertThat(facet.entries().get(2).count(), equalTo(1));
}
@Test public void testTermFacetWithEqualTermDistribution() throws Exception {