Search: Allow to pass a search filter, applying only on the query (and not on facets for example), closes #650.

modules/elasticsearch/src/main/java/org/elasticsearch/client/action/search/SearchRequestBuilder.java

@@ -28,6 +28,7 @@ import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.client.action.support.BaseRequestBuilder;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.index.query.xcontent.XContentFilterBuilder;
 import org.elasticsearch.index.query.xcontent.XContentQueryBuilder;
 import org.elasticsearch.search.Scroll;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
@@ -201,6 +202,33 @@ public class SearchRequestBuilder extends BaseRequestBuilder<SearchRequest, Sear
         return this;
     }
 
+    /**
+     * Sets a filter on the query executed that only applies to the search query
+     * (and not facets for example).
+     */
+    public SearchRequestBuilder setFilter(XContentFilterBuilder filter) {
+        sourceBuilder().filter(filter);
+        return this;
+    }
+
+    /**
+     * Sets a filter on the query executed that only applies to the search query
+     * (and not facets for example).
+     */
+    public SearchRequestBuilder setFilter(String filter) {
+        sourceBuilder().filter(filter);
+        return this;
+    }
+
+    /**
+     * Sets a filter on the query executed that only applies to the search query
+     * (and not facets for example).
+     */
+    public SearchRequestBuilder setFilter(byte[] filter) {
+        sourceBuilder().filter(filter);
+        return this;
+    }
+
     /**
      * From index to start the search from. Defaults to <tt>0</tt>.
      */

modules/elasticsearch/src/main/java/org/elasticsearch/common/lucene/search/FilteredCollector.java

@@ -0,0 +1,64 @@
+/*
+ * 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.common.lucene.search;
+
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.Filter;
+import org.apache.lucene.search.Scorer;
+import org.elasticsearch.common.lucene.docset.DocSet;
+import org.elasticsearch.common.lucene.docset.DocSets;
+
+import java.io.IOException;
+
+/**
+ * @author kimchy (shay.banon)
+ */
+public class FilteredCollector extends Collector {
+
+    private final Collector collector;
+
+    private final Filter filter;
+
+    private DocSet docSet;
+
+    public FilteredCollector(Collector collector, Filter filter) {
+        this.collector = collector;
+        this.filter = filter;
+    }
+
+    @Override public void setScorer(Scorer scorer) throws IOException {
+        collector.setScorer(scorer);
+    }
+
+    @Override public void collect(int doc) throws IOException {
+        if (docSet.get(doc)) {
+            collector.collect(doc);
+        }
+    }
+
+    @Override public void setNextReader(IndexReader reader, int docBase) throws IOException {
+        docSet = DocSets.convert(reader, filter.getDocIdSet(reader));
+    }
+
+    @Override public boolean acceptsDocsOutOfOrder() {
+        return collector.acceptsDocsOutOfOrder();
+    }
+}

modules/elasticsearch/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -30,6 +30,7 @@ import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.index.query.xcontent.XContentFilterBuilder;
 import org.elasticsearch.index.query.xcontent.XContentQueryBuilder;
 import org.elasticsearch.search.facet.AbstractFacetBuilder;
 import org.elasticsearch.search.highlight.HighlightBuilder;
@@ -69,6 +70,10 @@ public class SearchSourceBuilder implements ToXContent {
 
     private byte[] queryBinary;
 
+    private XContentFilterBuilder filterBuilder;
+
+    private byte[] filterBinary;
+
     private int from = -1;
 
     private int size = -1;
@@ -122,6 +127,33 @@ public class SearchSourceBuilder implements ToXContent {
         return this;
     }
 
+    /**
+     * Sets a filter on the query executed that only applies to the search query
+     * (and not facets for example).
+     */
+    public SearchSourceBuilder filter(XContentFilterBuilder filter) {
+        this.filterBuilder = filter;
+        return this;
+    }
+
+    /**
+     * Sets a filter on the query executed that only applies to the search query
+     * (and not facets for example).
+     */
+    public SearchSourceBuilder filter(String filterString) {
+        this.filterBinary = Unicode.fromStringAsBytes(filterString);
+        return this;
+    }
+
+    /**
+     * Sets a filter on the query executed that only applies to the search query
+     * (and not facets for example).
+     */
+    public SearchSourceBuilder filter(byte[] filter) {
+        this.filterBinary = filter;
+        return this;
+    }
+
     /**
      * From index to start the search from. Defaults to <tt>0</tt>.
      */
@@ -357,6 +389,19 @@ public class SearchSourceBuilder implements ToXContent {
             }
         }
 
+        if (filterBuilder != null) {
+            builder.field("filter");
+            filterBuilder.toXContent(builder, params);
+        }
+
+        if (filterBinary != null) {
+            if (XContentFactory.xContentType(queryBinary) == builder.contentType()) {
+                builder.rawField("filter", filterBinary);
+            } else {
+                builder.field("filter_binary", queryBinary);
+            }
+        }
+
         if (explain != null) {
             builder.field("explain", explain);
         }

modules/elasticsearch/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java

@@ -24,6 +24,7 @@ import org.elasticsearch.common.collect.Lists;
 import org.elasticsearch.common.collect.Maps;
 import org.elasticsearch.common.lucene.MultiCollector;
 import org.elasticsearch.common.lucene.search.ExtendedIndexSearcher;
+import org.elasticsearch.common.lucene.search.FilteredCollector;
 import org.elasticsearch.index.engine.Engine;
 import org.elasticsearch.search.dfs.CachedDfSource;
 
@@ -116,6 +117,11 @@ public class ContextIndexSearcher extends ExtendedIndexSearcher {
     }
 
     @Override public void search(Weight weight, Filter filter, Collector collector) throws IOException {
+        if (searchContext.parsedFilter() != null) {
+            // this will only get applied to the actual search collector and not
+            // to any scoped collectors
+            collector = new FilteredCollector(collector, searchContext.parsedFilter());
+        }
         if (searchContext.timeout() != null) {
             collector = new TimeLimitingCollector(collector, searchContext.timeout().millis());
         }

modules/elasticsearch/src/main/java/org/elasticsearch/search/internal/SearchContext.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.search.internal;
 
+import org.apache.lucene.search.Filter;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.Sort;
 import org.elasticsearch.ElasticSearchException;
@@ -115,6 +116,8 @@ public class SearchContext implements Releasable {
 
     private Query query;
 
+    private Filter filter;
+
     private int[] docIdsToLoad;
 
     private int docsIdsToLoadFrom;
@@ -291,6 +294,15 @@ public class SearchContext implements Releasable {
         return this.sort;
     }
 
+    public SearchContext parsedFilter(Filter filter) {
+        this.filter = filter;
+        return this;
+    }
+
+    public Filter parsedFilter() {
+        return this.filter;
+    }
+
     public String queryParserName() {
         return queryParserName;
     }

modules/elasticsearch/src/main/java/org/elasticsearch/search/query/FilterBinaryParseElement.java

@@ -0,0 +1,43 @@
+/*
+ * 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.query;
+
+import org.elasticsearch.common.xcontent.XContentFactory;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.query.xcontent.XContentIndexQueryParser;
+import org.elasticsearch.search.SearchParseElement;
+import org.elasticsearch.search.internal.SearchContext;
+
+/**
+ * @author kimchy (shay.banon)
+ */
+public class FilterBinaryParseElement implements SearchParseElement {
+
+    @Override public void parse(XContentParser parser, SearchContext context) throws Exception {
+        XContentIndexQueryParser indexQueryParser = (XContentIndexQueryParser) context.queryParser();
+        byte[] filterSource = parser.binaryValue();
+        XContentParser fSourceParser = XContentFactory.xContent(filterSource).createParser(filterSource);
+        try {
+            context.parsedFilter(indexQueryParser.parseInnerFilter(fSourceParser));
+        } finally {
+            fSourceParser.close();
+        }
+    }
+}

modules/elasticsearch/src/main/java/org/elasticsearch/search/query/FilterParseElement.java

@@ -0,0 +1,36 @@
+/*
+ * 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.query;
+
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.query.xcontent.XContentIndexQueryParser;
+import org.elasticsearch.search.SearchParseElement;
+import org.elasticsearch.search.internal.SearchContext;
+
+/**
+ * @author kimchy (shay.banon)
+ */
+public class FilterParseElement implements SearchParseElement {
+
+    @Override public void parse(XContentParser parser, SearchContext context) throws Exception {
+        XContentIndexQueryParser indexQueryParser = (XContentIndexQueryParser) context.queryParser();
+        context.parsedFilter(indexQueryParser.parseInnerFilter(parser));
+    }
+}

modules/elasticsearch/src/main/java/org/elasticsearch/search/query/QueryBinaryParseElement.java

@@ -34,6 +34,10 @@ public class QueryBinaryParseElement implements SearchParseElement {
         XContentIndexQueryParser indexQueryParser = (XContentIndexQueryParser) context.queryParser();
         byte[] querySource = parser.binaryValue();
         XContentParser qSourceParser = XContentFactory.xContent(querySource).createParser(querySource);
+        try {
             context.parsedQuery(indexQueryParser.parse(qSourceParser));
+        } finally {
+            qSourceParser.close();
+        }
     }
 }

modules/elasticsearch/src/main/java/org/elasticsearch/search/query/QueryPhase.java

@@ -57,6 +57,9 @@ public class QueryPhase implements SearchPhase {
                 .put("query", new QueryParseElement())
                 .put("queryBinary", new QueryBinaryParseElement())
                 .put("query_binary", new QueryBinaryParseElement())
+                .put("filter", new FilterParseElement())
+                .put("filterBinary", new FilterBinaryParseElement())
+                .put("filter_binary", new FilterBinaryParseElement())
                 .put("sort", new SortParseElement())
                 .putAll(facetPhase.parseElements());
         return parseElements.build();
@@ -73,6 +76,7 @@ public class QueryPhase implements SearchPhase {
     }
 
     public void execute(SearchContext searchContext) throws QueryPhaseExecutionException {
+        // set the filter on the searcher
         if (searchContext.parsedQuery().scopePhases().length > 0) {
             // we have scoped queries, refresh the id cache
             try {

modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/facet/SimpleFacetsTests.java

@@ -73,6 +73,58 @@ public class SimpleFacetsTests extends AbstractNodesTests {
         return client("server1");
     }
 
+    @Test public void testSearchFilter() throws Exception {
+        try {
+            client.admin().indices().prepareDelete("test").execute().actionGet();
+        } catch (Exception e) {
+            // ignore
+        }
+        client.admin().indices().prepareCreate("test").execute().actionGet();
+        client.admin().cluster().prepareHealth().setWaitForGreenStatus().execute().actionGet();
+
+        client.admin().cluster().prepareHealth().setWaitForGreenStatus().execute().actionGet();
+
+        client.prepareIndex("test", "type1").setSource(jsonBuilder().startObject()
+                .field("tag", "green")
+                .endObject()).execute().actionGet();
+        client.admin().indices().prepareFlush().setRefresh(true).execute().actionGet();
+
+        client.prepareIndex("test", "type1").setSource(jsonBuilder().startObject()
+                .field("tag", "blue")
+                .endObject()).execute().actionGet();
+
+        client.admin().indices().prepareRefresh().execute().actionGet();
+
+        SearchResponse searchResponse = client.prepareSearch()
+                .setQuery(matchAllQuery())
+                .addFacet(termsFacet("facet1").field("tag").size(10))
+                .execute().actionGet();
+
+        assertThat(searchResponse.hits().hits().length, equalTo(2));
+        TermsFacet facet = searchResponse.facets().facet("facet1");
+        assertThat(facet.name(), equalTo("facet1"));
+        assertThat(facet.entries().size(), equalTo(2));
+        assertThat(facet.entries().get(0).term(), anyOf(equalTo("green"), equalTo("blue")));
+        assertThat(facet.entries().get(0).count(), equalTo(1));
+        assertThat(facet.entries().get(1).term(), anyOf(equalTo("green"), equalTo("blue")));
+        assertThat(facet.entries().get(1).count(), equalTo(1));
+
+        searchResponse = client.prepareSearch()
+                .setQuery(matchAllQuery())
+                .setFilter(termFilter("tag", "blue"))
+                .addFacet(termsFacet("facet1").field("tag").size(10))
+                .execute().actionGet();
+
+        assertThat(searchResponse.hits().hits().length, equalTo(1));
+        facet = searchResponse.facets().facet("facet1");
+        assertThat(facet.name(), equalTo("facet1"));
+        assertThat(facet.entries().size(), equalTo(2));
+        assertThat(facet.entries().get(0).term(), anyOf(equalTo("green"), equalTo("blue")));
+        assertThat(facet.entries().get(0).count(), equalTo(1));
+        assertThat(facet.entries().get(1).term(), anyOf(equalTo("green"), equalTo("blue")));
+        assertThat(facet.entries().get(1).count(), equalTo(1));
+    }
+
     @Test public void testFacetsWithSize0() throws Exception {
         try {
             client.admin().indices().prepareDelete("test").execute().actionGet();