Search: Allow to specify script fields to be loaded, closes #221.

This commit is contained in:
kimchy 2010-06-14 12:50:33 +03:00
parent c087bbe804
commit 0a1bc874c3
10 changed files with 404 additions and 3 deletions

View File

@ -38,6 +38,7 @@ import org.elasticsearch.search.facets.histogram.HistogramFacet;
import org.elasticsearch.util.TimeValue;
import javax.annotation.Nullable;
import java.util.Map;
/**
* A search action request builder.
@ -232,6 +233,16 @@ public class SearchRequestBuilder {
return this;
}
public SearchRequestBuilder addScriptField(String name, String script) {
sourceBuilder().scriptField(name, script);
return this;
}
public SearchRequestBuilder addScriptField(String name, String script, Map<String, Object> params) {
sourceBuilder().scriptField(name, script, params);
return this;
}
/**
* Adds a sort against the given field name and the sort ordering.
*

View File

@ -29,6 +29,7 @@ import org.elasticsearch.index.shard.ShardId;
import org.elasticsearch.index.shard.service.IndexShard;
import org.elasticsearch.indices.IndicesLifecycle;
import org.elasticsearch.indices.IndicesService;
import org.elasticsearch.script.ScriptService;
import org.elasticsearch.search.dfs.CachedDfSource;
import org.elasticsearch.search.dfs.DfsPhase;
import org.elasticsearch.search.dfs.DfsSearchResult;
@ -76,6 +77,8 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
private final TimerService timerService;
private final ScriptService scriptService;
private final DfsPhase dfsPhase;
private final QueryPhase queryPhase;
@ -95,11 +98,12 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
private final ImmutableMap<String, SearchParseElement> elementParsers;
@Inject public SearchService(Settings settings, ClusterService clusterService, IndicesService indicesService, TimerService timerService,
DfsPhase dfsPhase, QueryPhase queryPhase, FetchPhase fetchPhase) {
ScriptService scriptService, DfsPhase dfsPhase, QueryPhase queryPhase, FetchPhase fetchPhase) {
super(settings);
this.clusterService = clusterService;
this.indicesService = indicesService;
this.timerService = timerService;
this.scriptService = scriptService;
this.dfsPhase = dfsPhase;
this.queryPhase = queryPhase;
this.fetchPhase = fetchPhase;
@ -285,7 +289,7 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
SearchShardTarget shardTarget = new SearchShardTarget(clusterService.state().nodes().localNodeId(), request.index(), request.shardId());
SearchContext context = new SearchContext(idGenerator.incrementAndGet(), shardTarget, request.timeout(), request.types(), engineSearcher, indexService);
SearchContext context = new SearchContext(idGenerator.incrementAndGet(), shardTarget, request.timeout(), request.types(), engineSearcher, indexService, scriptService);
context.scroll(request.scroll());

View File

@ -20,6 +20,7 @@
package org.elasticsearch.search.builder;
import org.elasticsearch.index.query.xcontent.XContentQueryBuilder;
import org.elasticsearch.util.collect.Lists;
import org.elasticsearch.util.gnu.trove.TObjectFloatHashMap;
import org.elasticsearch.util.gnu.trove.TObjectFloatIterator;
import org.elasticsearch.util.io.FastByteArrayOutputStream;
@ -32,6 +33,7 @@ import org.elasticsearch.util.xcontent.builder.XContentBuilder;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import static org.elasticsearch.util.collect.Lists.*;
@ -84,6 +86,8 @@ public class SearchSourceBuilder implements ToXContent {
private List<String> fieldNames;
private List<ScriptField> scriptFields;
private SearchSourceFacetsBuilder facetsBuilder;
private SearchSourceHighlightBuilder highlightBuilder;
@ -255,6 +259,18 @@ public class SearchSourceBuilder implements ToXContent {
return this;
}
public SearchSourceBuilder scriptField(String name, String script) {
return scriptField(name, script, null);
}
public SearchSourceBuilder scriptField(String name, String script, Map<String, Object> params) {
if (scriptFields == null) {
scriptFields = Lists.newArrayList();
}
scriptFields.add(new ScriptField(name, script, params));
return this;
}
/**
* Sets the boost a specific index will receive when the query is executeed against it.
*
@ -332,6 +348,20 @@ public class SearchSourceBuilder implements ToXContent {
}
}
if (scriptFields != null) {
builder.startObject("script_fields");
for (ScriptField scriptField : scriptFields) {
builder.startObject(scriptField.fieldName());
builder.field("script", scriptField.script());
if (scriptField.params() != null) {
builder.field("params");
builder.map(scriptField.params());
}
builder.endObject();
}
builder.endObject();
}
if (sortFields != null) {
builder.field("sort");
builder.startObject();
@ -369,6 +399,30 @@ public class SearchSourceBuilder implements ToXContent {
builder.endObject();
}
private static class ScriptField {
private final String fieldName;
private final String script;
private final Map<String, Object> params;
private ScriptField(String fieldName, String script, Map<String, Object> params) {
this.fieldName = fieldName;
this.script = script;
this.params = params;
}
public String fieldName() {
return fieldName;
}
public String script() {
return script;
}
public Map<String, Object> params() {
return params;
}
}
private static class SortTuple {
private final String fieldName;
private final boolean reverse;

View File

@ -22,10 +22,13 @@ package org.elasticsearch.search.fetch;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.FieldSelector;
import org.apache.lucene.document.Fieldable;
import org.apache.lucene.index.IndexReader;
import org.elasticsearch.index.mapper.*;
import org.elasticsearch.search.SearchHitField;
import org.elasticsearch.search.SearchParseElement;
import org.elasticsearch.search.SearchPhase;
import org.elasticsearch.search.fetch.script.ScriptFieldsContext;
import org.elasticsearch.search.fetch.script.ScriptFieldsParseElement;
import org.elasticsearch.search.highlight.HighlightPhase;
import org.elasticsearch.search.internal.InternalSearchHit;
import org.elasticsearch.search.internal.InternalSearchHitField;
@ -54,6 +57,8 @@ public class FetchPhase implements SearchPhase {
ImmutableMap.Builder<String, SearchParseElement> parseElements = ImmutableMap.builder();
parseElements.put("explain", new ExplainParseElement())
.put("fields", new FieldsParseElement())
.put("script_fields", new ScriptFieldsParseElement())
.put("scriptFields", new ScriptFieldsParseElement())
.putAll(highlightPhase.parseElements());
return parseElements.build();
}
@ -109,6 +114,29 @@ public class FetchPhase implements SearchPhase {
}
hitField.values().add(value);
}
if (context.scriptFields() != null) {
int readerIndex = context.searcher().readerIndex(docId);
IndexReader subReader = context.searcher().subReaders()[readerIndex];
int subDoc = docId - context.searcher().docStarts()[readerIndex];
for (ScriptFieldsContext.ScriptField scriptField : context.scriptFields().fields()) {
scriptField.scriptFieldsFunction().setNextReader(subReader);
Object value = scriptField.scriptFieldsFunction().execute(subDoc, scriptField.params());
if (searchHit.fields() == null) {
searchHit.fields(new HashMap<String, SearchHitField>(2));
}
SearchHitField hitField = searchHit.fields().get(scriptField.name());
if (hitField == null) {
hitField = new InternalSearchHitField(scriptField.name(), new ArrayList<Object>(2));
searchHit.fields().put(scriptField.name(), hitField);
}
hitField.values().add(value);
}
}
doExplanation(context, docId, searchHit);
}
context.fetchResult().hits(new InternalSearchHits(hits, context.queryResult().topDocs().totalHits));

View File

@ -0,0 +1,65 @@
/*
* Licensed to Elastic Search and Shay Banon under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. Elastic Search licenses this
* file to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.elasticsearch.search.fetch.script;
import org.elasticsearch.index.field.function.script.ScriptFieldsFunction;
import java.util.List;
import java.util.Map;
/**
* @author kimchy (shay.banon)
*/
public class ScriptFieldsContext {
public static class ScriptField {
private final String name;
private final ScriptFieldsFunction scriptFieldsFunction;
private final Map<String, Object> params;
public ScriptField(String name, ScriptFieldsFunction scriptFieldsFunction, Map<String, Object> params) {
this.name = name;
this.scriptFieldsFunction = scriptFieldsFunction;
this.params = params;
}
public String name() {
return name;
}
public ScriptFieldsFunction scriptFieldsFunction() {
return scriptFieldsFunction;
}
public Map<String, Object> params() {
return params;
}
}
private List<ScriptField> fields;
public ScriptFieldsContext(List<ScriptField> fields) {
this.fields = fields;
}
public List<ScriptField> fields() {
return this.fields;
}
}

View File

@ -0,0 +1,75 @@
/*
* Licensed to Elastic Search and Shay Banon under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. Elastic Search licenses this
* file to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.elasticsearch.search.fetch.script;
import org.elasticsearch.index.field.function.script.ScriptFieldsFunction;
import org.elasticsearch.search.SearchParseElement;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.util.collect.Lists;
import org.elasticsearch.util.xcontent.XContentParser;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/**
* <pre>
* "script_fields" : {
* "test1" : {
* "script" : "doc['field_name'].value"
* },
* "test2" : {
* "script" : "..."
* }
* }
* </pre>
*
* @author kimchy (shay.banon)
*/
public class ScriptFieldsParseElement implements SearchParseElement {
@Override public void parse(XContentParser parser, SearchContext context) throws Exception {
XContentParser.Token token;
String currentFieldName = null;
List<ScriptFieldsContext.ScriptField> scriptFields = Lists.newArrayList();
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentFieldName = parser.currentName();
} else if (token == XContentParser.Token.START_OBJECT) {
String fieldName = currentFieldName;
String script = null;
Map<String, Object> params = null;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentFieldName = parser.currentName();
} else if (token == XContentParser.Token.START_OBJECT) {
params = parser.map();
} else if (token.isValue()) {
script = parser.text();
}
}
scriptFields.add(new ScriptFieldsContext.ScriptField(fieldName,
new ScriptFieldsFunction(script, context.scriptService(), context.mapperService(), context.fieldDataCache()),
params == null ? new HashMap<String, Object>() : params));
}
}
context.scriptFields(new ScriptFieldsContext(scriptFields));
}
}

View File

@ -48,6 +48,37 @@ public class ContextIndexSearcher extends IndexSearcher {
this.searchContext = searchContext;
}
public IndexReader[] subReaders() {
return this.subReaders;
}
public int[] docStarts() {
return this.docStarts;
}
// taken from DirectoryReader#readerIndex
public int readerIndex(int doc) {
int lo = 0; // search starts array
int hi = subReaders.length - 1; // for first element less
while (hi >= lo) {
int mid = (lo + hi) >>> 1;
int midValue = docStarts[mid];
if (doc < midValue)
hi = mid - 1;
else if (doc > midValue)
lo = mid + 1;
else { // found a match
while (mid + 1 < subReaders.length && docStarts[mid + 1] == midValue) {
mid++; // scan to last match
}
return mid;
}
}
return hi;
}
public void dfSource(CachedDfSource dfSource) {
this.dfSource = dfSource;
}

View File

@ -32,11 +32,13 @@ import org.elasticsearch.index.query.IndexQueryParserMissingException;
import org.elasticsearch.index.query.IndexQueryParserService;
import org.elasticsearch.index.service.IndexService;
import org.elasticsearch.index.similarity.SimilarityService;
import org.elasticsearch.script.ScriptService;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchShardTarget;
import org.elasticsearch.search.dfs.DfsSearchResult;
import org.elasticsearch.search.facets.SearchContextFacets;
import org.elasticsearch.search.fetch.FetchSearchResult;
import org.elasticsearch.search.fetch.script.ScriptFieldsContext;
import org.elasticsearch.search.highlight.SearchContextHighlight;
import org.elasticsearch.search.query.QuerySearchResult;
import org.elasticsearch.util.TimeValue;
@ -57,6 +59,8 @@ public class SearchContext implements Releasable {
private final Engine.Searcher engineSearcher;
private final ScriptService scriptService;
private final IndexService indexService;
private final ContextIndexSearcher searcher;
@ -102,6 +106,8 @@ public class SearchContext implements Releasable {
private SearchContextHighlight highlight;
private ScriptFieldsContext scriptFields;
private boolean queryRewritten;
private volatile TimeValue keepAlive;
@ -111,12 +117,13 @@ public class SearchContext implements Releasable {
private volatile Timeout keepAliveTimeout;
public SearchContext(long id, SearchShardTarget shardTarget, TimeValue timeout,
String[] types, Engine.Searcher engineSearcher, IndexService indexService) {
String[] types, Engine.Searcher engineSearcher, IndexService indexService, ScriptService scriptService) {
this.id = id;
this.shardTarget = shardTarget;
this.timeout = timeout;
this.types = types;
this.engineSearcher = engineSearcher;
this.scriptService = scriptService;
this.dfsResult = new DfsSearchResult(id, shardTarget);
this.queryResult = new QuerySearchResult(id, shardTarget);
this.fetchResult = new FetchSearchResult(id, shardTarget);
@ -187,6 +194,14 @@ public class SearchContext implements Releasable {
this.highlight = highlight;
}
public ScriptFieldsContext scriptFields() {
return this.scriptFields;
}
public void scriptFields(ScriptFieldsContext scriptFields) {
this.scriptFields = scriptFields;
}
public ContextIndexSearcher searcher() {
return this.searcher;
}
@ -214,6 +229,10 @@ public class SearchContext implements Releasable {
return indexService.similarityService();
}
public ScriptService scriptService() {
return scriptService;
}
public FilterCache filterCache() {
return indexService.cache().filter();
}

View File

@ -0,0 +1,108 @@
/*
* Licensed to Elastic Search and Shay Banon under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. Elastic Search licenses this
* file to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.elasticsearch.test.integration.search.scriptfield;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.test.integration.AbstractNodesTests;
import org.elasticsearch.util.MapBuilder;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
import java.util.Map;
import static org.elasticsearch.client.Requests.*;
import static org.elasticsearch.index.query.xcontent.QueryBuilders.*;
import static org.elasticsearch.search.builder.SearchSourceBuilder.*;
import static org.elasticsearch.util.xcontent.XContentFactory.*;
import static org.hamcrest.MatcherAssert.*;
import static org.hamcrest.Matchers.*;
/**
* @author kimchy (shay.banon)
*/
@Test
public class ScriptFieldSearchTests extends AbstractNodesTests {
private Client client;
@BeforeMethod public void createNodes() throws Exception {
startNode("server1");
client = getClient();
}
@AfterMethod public void closeNodes() {
client.close();
closeAllNodes();
}
protected Client getClient() {
return client("server1");
}
@Test
public void testCustomScriptBoost() throws Exception {
client.admin().indices().prepareCreate("test").execute().actionGet();
client.prepareIndex("test", "type1", "1")
.setSource(jsonBuilder().startObject().field("test", "value beck").field("num1", 1.0f).endObject())
.execute().actionGet();
client.admin().indices().prepareFlush().execute().actionGet();
client.prepareIndex("test", "type1", "2")
.setSource(jsonBuilder().startObject().field("test", "value beck").field("num1", 2.0f).endObject())
.execute().actionGet();
client.admin().indices().prepareFlush().execute().actionGet();
client.prepareIndex("test", "type1", "3")
.setSource(jsonBuilder().startObject().field("test", "value beck").field("num1", 3.0f).endObject())
.execute().actionGet();
client.admin().indices().refresh(refreshRequest()).actionGet();
logger.info("running doc['num1'].value");
SearchResponse response = client.prepareSearch()
.setQuery(matchAllQuery())
.addSort("num1", Order.ASC)
.addScriptField("sNum1", "doc['num1'].value")
.execute().actionGet();
assertThat(response.hits().totalHits(), equalTo(3l));
assertThat(response.hits().getAt(0).id(), equalTo("1"));
assertThat((Double) response.hits().getAt(0).fields().get("sNum1").values().get(0), equalTo(1.0));
assertThat(response.hits().getAt(1).id(), equalTo("2"));
assertThat((Double) response.hits().getAt(1).fields().get("sNum1").values().get(0), equalTo(2.0));
assertThat(response.hits().getAt(2).id(), equalTo("3"));
assertThat((Double) response.hits().getAt(2).fields().get("sNum1").values().get(0), equalTo(3.0));
logger.info("running doc['num1'].value * factor");
Map<String, Object> params = MapBuilder.<String, Object>newMapBuilder().put("factor", 2.0).map();
response = client.prepareSearch()
.setQuery(matchAllQuery())
.addSort("num1", Order.ASC)
.addScriptField("sNum1", "doc['num1'].value * factor", params)
.execute().actionGet();
assertThat(response.hits().totalHits(), equalTo(3l));
assertThat(response.hits().getAt(0).id(), equalTo("1"));
assertThat((Double) response.hits().getAt(0).fields().get("sNum1").values().get(0), equalTo(2.0));
assertThat(response.hits().getAt(1).id(), equalTo("2"));
assertThat((Double) response.hits().getAt(1).fields().get("sNum1").values().get(0), equalTo(4.0));
assertThat(response.hits().getAt(2).id(), equalTo("3"));
assertThat((Double) response.hits().getAt(2).fields().get("sNum1").values().get(0), equalTo(6.0));
}
}

View File

@ -0,0 +1,6 @@
cluster:
routing:
schedule: 100ms
index:
number_of_shards: 1
number_of_replicas: 0