HSEARCH-2477 Implement shard filtering for Elasticsearch queries

hibernate · May 12, 2017 · 51fc60d · 51fc60d
1 parent 33beda0
commit 51fc60d
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 43 deletions.
diff --git a/...earch/src/main/java/org/hibernate/search/elasticsearch/impl/ElasticsearchHSQueryImpl.java b/...earch/src/main/java/org/hibernate/search/elasticsearch/impl/ElasticsearchHSQueryImpl.java
@@ -281,6 +281,20 @@ protected Set<String> getSupportedProjectionConstants() {
 		return SUPPORTED_PROJECTION_CONSTANTS;
 	}
 
+	@Override
+	protected List<IndexManager> getIndexManagers(EntityIndexBinding binding) {
+		List<IndexManager> indexManagers = super.getIndexManagers( binding );
+		for ( IndexManager indexManager : indexManagers ) {
+			if ( !( indexManager instanceof ElasticsearchIndexManager ) ) {
+				throw LOG.cannotRunEsQueryTargetingEntityIndexedWithNonEsIndexManager(
+					binding.getDocumentBuilder().getBeanClass(),
+					rawSearchPayload.toString()
+				);
+			}
+		}
+		return indexManagers;
+	}
+
 	private void execute() {
 		IndexSearcher searcher = getOrCreateSearcher();
 		if ( searcher != null ) {
@@ -304,16 +318,9 @@ private IndexSearcher getOrCreateSearcher() {
 		for ( Map.Entry<String, EntityIndexBinding> entry: targetedEntityBindingsByName.entrySet() ) {
 			EntityIndexBinding binding = entry.getValue();
 
-			IndexManager[] indexManagers = binding.getIndexManagers();
+			List<IndexManager> indexManagers = getIndexManagers( binding );
 
 			for ( IndexManager indexManager : indexManagers ) {
-				if ( !( indexManager instanceof ElasticsearchIndexManager ) ) {
-					throw LOG.cannotRunEsQueryTargetingEntityIndexedWithNonEsIndexManager(
-						binding.getDocumentBuilder().getBeanClass(),
-						rawSearchPayload.toString()
-					);
-				}
-
 				ElasticsearchIndexManager esIndexManager = (ElasticsearchIndexManager) indexManager;
 				indexNames.add( URLEncodedString.fromString( esIndexManager.getActualIndexName() ) );
 				if ( elasticsearchService == null ) {

diff --git a/engine/src/main/java/org/hibernate/search/query/engine/impl/AbstractHSQuery.java b/engine/src/main/java/org/hibernate/search/query/engine/impl/AbstractHSQuery.java
@@ -11,6 +11,7 @@
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -33,8 +34,10 @@
 import org.hibernate.search.engine.spi.EntityIndexBinding;
 import org.hibernate.search.exception.SearchException;
 import org.hibernate.search.filter.FullTextFilter;
+import org.hibernate.search.filter.FullTextFilterImplementor;
 import org.hibernate.search.filter.ShardSensitiveOnlyFilter;
 import org.hibernate.search.filter.impl.FullTextFilterImpl;
+import org.hibernate.search.indexes.spi.IndexManager;
 import org.hibernate.search.metadata.NumericFieldSettingsDescriptor.NumericEncodingType;
 import org.hibernate.search.query.engine.spi.HSQuery;
 import org.hibernate.search.query.engine.spi.TimeoutExceptionFactory;
@@ -63,6 +66,8 @@ public abstract class AbstractHSQuery implements HSQuery, Serializable {
 	 */
 	public static final String HSEARCH_PROJECTION_FIELD_PREFIX = "__HSearch_";
 
+	private static final FullTextFilterImplementor[] EMPTY_FULL_TEXT_FILTER_IMPLEMENTOR_ARRAY = new FullTextFilterImplementor[0];
+
 	protected transient ExtendedSearchIntegrator extendedIntegrator;
 	protected transient TimeoutExceptionFactory timeoutExceptionFactory;
 	protected transient TimeoutManagerImpl timeoutManager;
@@ -530,6 +535,25 @@ private FacetMetadata findFacetMetadata(FacetingRequest facetingRequest, Iterabl
 		return facetMetadata;
 	}
 
+	protected List<IndexManager> getIndexManagers(EntityIndexBinding binding) {
+		FullTextFilterImplementor[] fullTextFilters = getFullTextFilterImplementors();
+		List<IndexManager> indexManagers = Arrays.asList( binding.getSelectionStrategy().getIndexManagersForQuery( fullTextFilters ) );
+		return indexManagers;
+	}
+
+	private FullTextFilterImplementor[] getFullTextFilterImplementors() {
+		FullTextFilterImplementor[] fullTextFilters;
+
+		if ( filterDefinitions != null && !filterDefinitions.isEmpty() ) {
+			fullTextFilters = filterDefinitions.values().toArray( new FullTextFilterImplementor[filterDefinitions.size()] );
+		}
+		else {
+			// no filter get all shards
+			fullTextFilters = EMPTY_FULL_TEXT_FILTER_IMPLEMENTOR_ARRAY;
+		}
+		return fullTextFilters;
+	}
+
 	// hooks to be implemented by specific sub-classes
 
 	protected abstract void extractFacetResults();

diff --git a/engine/src/main/java/org/hibernate/search/query/engine/impl/LuceneHSQuery.java b/engine/src/main/java/org/hibernate/search/query/engine/impl/LuceneHSQuery.java
@@ -12,7 +12,6 @@
 import java.io.IOException;
 import java.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
@@ -42,7 +41,6 @@
 import org.hibernate.search.exception.AssertionFailure;
 import org.hibernate.search.exception.SearchException;
 import org.hibernate.search.filter.FilterKey;
-import org.hibernate.search.filter.FullTextFilterImplementor;
 import org.hibernate.search.filter.StandardFilterKey;
 import org.hibernate.search.filter.impl.CachingWrapperQuery;
 import org.hibernate.search.filter.impl.DefaultFilterKey;
@@ -66,7 +64,6 @@
 public class LuceneHSQuery extends AbstractHSQuery implements HSQuery {
 
 	static final Log log = LoggerFactory.make();
-	private static final FullTextFilterImplementor[] EMPTY_FULL_TEXT_FILTER_IMPLEMENTOR = new FullTextFilterImplementor[0];
 
 	private static final Set<String> SUPPORTED_PROJECTION_CONSTANTS = Collections.unmodifiableSet(
 			CollectionHelper.asSet(
@@ -282,6 +279,25 @@ protected Set<String> getSupportedProjectionConstants() {
 		return SUPPORTED_PROJECTION_CONSTANTS;
 	}
 
+	@Override
+	protected List<IndexManager> getIndexManagers(EntityIndexBinding binding) {
+		List<IndexManager> indexManagers = super.getIndexManagers( binding );
+
+		for ( IndexManager indexManager : indexManagers ) {
+			if ( !( indexManager instanceof DirectoryBasedIndexManager ) ) {
+				//TODO clarify intent:
+				// A) improve the error message so that people understand what they should do
+				// B) Is the point really to not support "directory-based" or rather non-Lucene native based ones?
+				throw log.cannotRunLuceneQueryTargetingEntityIndexedWithNonDirectoryBasedIndexManager(
+						binding.getDocumentBuilder().getBeanClass(),
+						luceneQuery.toString()
+				);
+			}
+		}
+
+		return indexManagers;
+	}
+
 	/**
 	 * Execute the lucene search and return the matching hits.
 	 *
@@ -584,38 +600,6 @@ else if ( !similarity.getClass().equals( entitySimilarity.getClass() ) ) {
 		return similarity;
 	}
 
-	private List<IndexManager> getIndexManagers(EntityIndexBinding binding) {
-		FullTextFilterImplementor[] fullTextFilters = getFullTextFilters();
-		List<IndexManager> indexManagers = Arrays.asList( binding.getSelectionStrategy().getIndexManagersForQuery( fullTextFilters ) );
-
-		for ( IndexManager indexManager : indexManagers ) {
-			if ( !( indexManager instanceof DirectoryBasedIndexManager ) ) {
-				//TODO clarify intent:
-				// A) improve the error message so that people understand what they should do
-				// B) Is the point really to not support "directory-based" or rather non-Lucene native based ones?
-				throw log.cannotRunLuceneQueryTargetingEntityIndexedWithNonDirectoryBasedIndexManager(
-						binding.getDocumentBuilder().getBeanClass(),
-						luceneQuery.toString()
-				);
-			}
-		}
-
-		return indexManagers;
-	}
-
-	private FullTextFilterImplementor[] getFullTextFilters() {
-		FullTextFilterImplementor[] fullTextFilters;
-
-		if ( filterDefinitions != null && !filterDefinitions.isEmpty() ) {
-			fullTextFilters = filterDefinitions.values().toArray( new FullTextFilterImplementor[filterDefinitions.size()] );
-		}
-		else {
-			// no filter get all shards
-			fullTextFilters = EMPTY_FULL_TEXT_FILTER_IMPLEMENTOR;
-		}
-		return fullTextFilters;
-	}
-
 	private QueryFilters createFilters() {
 		List<Query> filterQueries = new ArrayList<>();
 		if ( !filterDefinitions.isEmpty() ) {