Script Filter: Support providing a custom script as a filter, closes #226.

modules/elasticsearch/src/main/java/org/elasticsearch/index/query/xcontent/FilterBuilders.java

@@ -175,6 +175,10 @@ public abstract class FilterBuilders {
         return new QueryFilterBuilder(queryBuilder);
     }
 
+    public static ScriptFilterBuilder scriptFilter(String script) {
+        return new ScriptFilterBuilder(script);
+    }
+
     public static BoolFilterBuilder boolFilter() {
         return new BoolFilterBuilder();
     }

modules/elasticsearch/src/main/java/org/elasticsearch/index/query/xcontent/ScriptFilterBuilder.java

@@ -0,0 +1,68 @@
+/*
+ * 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.index.query.xcontent;
+
+import org.elasticsearch.common.xcontent.builder.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Map;
+
+import static org.elasticsearch.common.collect.Maps.*;
+
+/**
+ * @author kimchy (shay.banon)
+ */
+public class ScriptFilterBuilder extends BaseFilterBuilder {
+
+    private final String script;
+
+    private Map<String, Object> params;
+
+    public ScriptFilterBuilder(String script) {
+        this.script = script;
+    }
+
+    public ScriptFilterBuilder addParam(String name, Object value) {
+        if (params == null) {
+            params = newHashMap();
+        }
+        params.put(name, value);
+        return this;
+    }
+
+    public ScriptFilterBuilder params(Map<String, Object> params) {
+        if (params == null) {
+            this.params = params;
+        } else {
+            this.params.putAll(params);
+        }
+        return this;
+    }
+
+    @Override protected void doXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject(ScriptFilterParser.NAME);
+        builder.field("script", script);
+        if (this.params != null) {
+            builder.field("params");
+            builder.map(this.params);
+        }
+        builder.endObject();
+    }
+}

modules/elasticsearch/src/main/java/org/elasticsearch/index/query/xcontent/ScriptFilterParser.java

@@ -0,0 +1,197 @@
+/*
+ * 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.index.query.xcontent;
+
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.search.DocIdSet;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.Filter;
+import org.elasticsearch.common.collect.Maps;
+import org.elasticsearch.common.inject.Inject;
+import org.elasticsearch.common.lucene.docset.DocSet;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.AbstractIndexComponent;
+import org.elasticsearch.index.Index;
+import org.elasticsearch.index.cache.field.data.FieldDataCache;
+import org.elasticsearch.index.field.function.script.ScriptFieldsFunction;
+import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.query.QueryParsingException;
+import org.elasticsearch.index.settings.IndexSettings;
+import org.elasticsearch.script.ScriptService;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * @author kimchy (shay.banon)
+ */
+public class ScriptFilterParser extends AbstractIndexComponent implements XContentFilterParser {
+
+    public static final String NAME = "script";
+
+    @Inject public ScriptFilterParser(Index index, @IndexSettings Settings settings) {
+        super(index, settings);
+    }
+
+    @Override public String[] names() {
+        return new String[]{NAME};
+    }
+
+    @Override public Filter parse(QueryParseContext parseContext) throws IOException, QueryParsingException {
+        XContentParser parser = parseContext.parser();
+
+        XContentParser.Token token;
+
+        String script = null;
+        Map<String, Object> params = null;
+
+        String currentFieldName = 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) {
+                if ("params".equals(currentFieldName)) {
+                    params = parser.map();
+                }
+            } else if (token.isValue()) {
+                if ("script".equals(currentFieldName)) {
+                    script = parser.text();
+                }
+            }
+        }
+
+        if (script == null) {
+            throw new QueryParsingException(index, "script must be provided with a [script] filter");
+        }
+        if (params == null) {
+            params = Maps.newHashMap();
+        }
+
+        return new ScriptFilter(script, params, parseContext.mapperService(), parseContext.indexCache().fieldData(), parseContext.scriptService());
+    }
+
+    public static class ScriptFilter extends Filter {
+
+        private final String script;
+
+        private final Map<String, Object> params;
+
+        private final MapperService mapperService;
+
+        private final FieldDataCache fieldDataCache;
+
+        private final ScriptService scriptService;
+
+        private ScriptFilter(String script, Map<String, Object> params,
+                             MapperService mapperService, FieldDataCache fieldDataCache, ScriptService scriptService) {
+            this.script = script;
+            this.params = params;
+            this.mapperService = mapperService;
+            this.fieldDataCache = fieldDataCache;
+            this.scriptService = scriptService;
+        }
+
+        @Override public String toString() {
+            StringBuilder buffer = new StringBuilder();
+            buffer.append("ScriptFilter(");
+            buffer.append(script);
+            buffer.append(")");
+            return buffer.toString();
+        }
+
+        @Override public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+
+            ScriptFilter that = (ScriptFilter) o;
+
+            if (params != null ? !params.equals(that.params) : that.params != null) return false;
+            if (script != null ? !script.equals(that.script) : that.script != null) return false;
+
+            return true;
+        }
+
+        @Override public int hashCode() {
+            int result = script != null ? script.hashCode() : 0;
+            result = 31 * result + (params != null ? params.hashCode() : 0);
+            return result;
+        }
+
+        @Override public DocIdSet getDocIdSet(final IndexReader reader) throws IOException {
+            final ScriptFieldsFunction function = new ScriptFieldsFunction(script, scriptService, mapperService, fieldDataCache);
+            function.setNextReader(reader);
+            final int maxDoc = reader.maxDoc();
+            return new DocSet() {
+                @Override public boolean isCacheable() {
+                    return false; // though it is, we want to cache it into in memory rep so it will be faster
+                }
+
+                @Override public boolean get(int doc) throws IOException {
+                    Object val = function.execute(doc, params);
+                    if (val == null) {
+                        return false;
+                    }
+                    if (val instanceof Boolean) {
+                        return (Boolean) val;
+                    }
+                    if (val instanceof Number) {
+                        return ((Number) val).longValue() != 0;
+                    }
+                    throw new IOException("Can't handle type [" + val + "] in script filter");
+                }
+
+                @Override public DocIdSetIterator iterator() throws IOException {
+                    return new DocIdSetIterator() {
+                        private int doc = -1;
+
+                        @Override public int docID() {
+                            return doc;
+                        }
+
+                        @Override public int nextDoc() throws IOException {
+                            do {
+                                doc++;
+                                if (doc >= maxDoc) {
+                                    return doc = NO_MORE_DOCS;
+                                }
+                            } while (!get(doc));
+                            return doc;
+                        }
+
+                        @Override public int advance(int target) throws IOException {
+                            if (target >= maxDoc) {
+                                return doc = NO_MORE_DOCS;
+                            }
+                            doc = target;
+                            while (!get(doc)) {
+                                doc++;
+                                if (doc >= maxDoc) {
+                                    return doc = NO_MORE_DOCS;
+                                }
+                            }
+                            return doc;
+                        }
+                    };
+                }
+            };
+        }
+    }
+}

modules/elasticsearch/src/main/java/org/elasticsearch/index/query/xcontent/XContentQueryParserRegistry.java

@@ -85,6 +85,7 @@ public class XContentQueryParserRegistry {
         add(filterParsersMap, new TermsFilterParser(index, indexSettings));
         add(filterParsersMap, new RangeFilterParser(index, indexSettings));
         add(filterParsersMap, new PrefixFilterParser(index, indexSettings));
+        add(filterParsersMap, new ScriptFilterParser(index, indexSettings));
         add(filterParsersMap, new QueryFilterParser(index, indexSettings));
         add(filterParsersMap, new BoolFilterParser(index, indexSettings));
         add(filterParsersMap, new AndFilterParser(index, indexSettings));

modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/scriptfilter/ScriptFilterSearchTests.java

@@ -0,0 +1,114 @@
+/*
+ * 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.scriptfilter;
+
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.test.integration.AbstractNodesTests;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.elasticsearch.client.Requests.*;
+import static org.elasticsearch.common.xcontent.XContentFactory.*;
+import static org.elasticsearch.index.query.xcontent.FilterBuilders.*;
+import static org.elasticsearch.index.query.xcontent.QueryBuilders.*;
+import static org.elasticsearch.search.builder.SearchSourceBuilder.*;
+import static org.hamcrest.MatcherAssert.*;
+import static org.hamcrest.Matchers.*;
+
+/**
+ * @author kimchy (shay.banon)
+ */
+@Test
+public class ScriptFilterSearchTests 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 > 1");
+        SearchResponse response = client.prepareSearch()
+                .setQuery(filtered(matchAllQuery(), scriptFilter("doc['num1'].value > 1")))
+                .addSort("num1", Order.ASC)
+                .addScriptField("sNum1", "doc['num1'].value")
+                .execute().actionGet();
+
+        assertThat(response.hits().totalHits(), equalTo(2l));
+        assertThat(response.hits().getAt(0).id(), equalTo("2"));
+        assertThat((Double) response.hits().getAt(0).fields().get("sNum1").values().get(0), equalTo(2.0));
+        assertThat(response.hits().getAt(1).id(), equalTo("3"));
+        assertThat((Double) response.hits().getAt(1).fields().get("sNum1").values().get(0), equalTo(3.0));
+
+        logger.info("running doc['num1'].value > param1");
+        response = client.prepareSearch()
+                .setQuery(filtered(matchAllQuery(), scriptFilter("doc['num1'].value > param1").addParam("param1", 2)))
+                .addSort("num1", Order.ASC)
+                .addScriptField("sNum1", "doc['num1'].value")
+                .execute().actionGet();
+
+        assertThat(response.hits().totalHits(), equalTo(1l));
+        assertThat(response.hits().getAt(0).id(), equalTo("3"));
+        assertThat((Double) response.hits().getAt(0).fields().get("sNum1").values().get(0), equalTo(3.0));
+
+        logger.info("running doc['num1'].value > param1");
+        response = client.prepareSearch()
+                .setQuery(filtered(matchAllQuery(), scriptFilter("doc['num1'].value > param1").addParam("param1", -1)))
+                .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));
+    }
+}

modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/scriptfilter/ScriptFilterSearchTests.yml

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