ATLAS-1961: Basic search improvement in use of Solr index for attribute filtering (# 2)

repository/src/main/java/org/apache/atlas/discovery/EntitySearchProcessor.java

@@ -20,6 +20,7 @@ package org.apache.atlas.discovery;
 import org.apache.atlas.model.discovery.SearchParameters.FilterCriteria;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graphdb.*;
+import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.utils.AtlasPerfTracer;
 import org.apache.commons.collections.CollectionUtils;
@@ -39,18 +40,21 @@ public class EntitySearchProcessor extends SearchProcessor {
     public EntitySearchProcessor(SearchContext context) {
         super(context);
 
-        final AtlasEntityType entityType         = context.getEntityType();
-        final FilterCriteria  filterCriteria     = context.getSearchParameters().getEntityFilters();
-        final Set<String>     typeAndSubTypes    = entityType.getTypeAndAllSubTypes();
-        final Set<String>     solrAttributes     = new HashSet<>();
-        final Set<String>     gremlinAttributes  = new HashSet<>();
-        final Set<String>     allAttributes      = new HashSet<>();
+        final AtlasEntityType entityType        = context.getEntityType();
+        final FilterCriteria  filterCriteria    = context.getSearchParameters().getEntityFilters();
+        final Set<String>     typeAndSubTypes   = entityType.getTypeAndAllSubTypes();
+        final Set<String>     solrAttributes    = new HashSet<>();
+        final Set<String>     gremlinAttributes = new HashSet<>();
+        final Set<String>     allAttributes     = new HashSet<>();
+
+        final AtlasClassificationType classificationType   = context.getClassificationType();
+        final boolean                 filterClassification = classificationType != null && !context.needClassificationProcessor();
 
 
         processSearchAttributes(entityType, filterCriteria, solrAttributes, gremlinAttributes, allAttributes);
 
-        final boolean typeSearchBySolr = typeAndSubTypes.size() <= MAX_ENTITY_TYPES_IN_INDEX_QUERY;
-        final boolean attrSearchBySolr = canApplySolrFilter(entityType, filterCriteria, false);
+        final boolean typeSearchBySolr = !filterClassification && typeAndSubTypes.size() <= MAX_ENTITY_TYPES_IN_INDEX_QUERY;
+        final boolean attrSearchBySolr = !filterClassification && CollectionUtils.isNotEmpty(solrAttributes) && canApplySolrFilter(entityType, filterCriteria, false);
 
         StringBuilder solrQuery = new StringBuilder();
 
@@ -82,6 +86,10 @@ public class EntitySearchProcessor extends SearchProcessor {
                 query.in(Constants.TYPE_NAME_PROPERTY_KEY, typeAndSubTypes);
             }
 
+            if (filterClassification) {
+                query.in(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationType.getTypeAndAllSubTypes());
+            }
+
             graphQuery = toGremlinFilterQuery(entityType, filterCriteria, gremlinAttributes, query);
 
             if (context.getSearchParameters().getExcludeDeletedEntities() && indexQuery == null) {
@@ -93,6 +101,10 @@ public class EntitySearchProcessor extends SearchProcessor {
 
         AtlasGraphQuery query = context.getGraph().query().in(Constants.TYPE_NAME_PROPERTY_KEY, typeAndSubTypes);
 
+        if (filterClassification) {
+            query.in(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationType.getTypeAndAllSubTypes());
+        }
+
         filterGraphQuery = toGremlinFilterQuery(entityType, filterCriteria, allAttributes, query);
 
         if (context.getSearchParameters().getExcludeDeletedEntities()) {
@@ -115,7 +127,7 @@ public class EntitySearchProcessor extends SearchProcessor {
         }
 
         try {
-            int qryOffset = (nextProcessor == null) ? context.getSearchParameters().getOffset() : 0;
+            int qryOffset = (nextProcessor == null && (graphQuery == null || indexQuery == null)) ? context.getSearchParameters().getOffset() : 0;
             int limit     = context.getSearchParameters().getLimit();
             int resultIdx = qryOffset;
 

repository/src/main/java/org/apache/atlas/discovery/FullTextSearchProcessor.java

@@ -22,12 +22,14 @@ import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graphdb.AtlasIndexQuery;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.utils.AtlasPerfTracer;
+import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Set;
 
 
 public class FullTextSearchProcessor extends SearchProcessor {
@@ -40,9 +42,39 @@ public class FullTextSearchProcessor extends SearchProcessor {
         super(context);
 
         SearchParameters searchParameters = context.getSearchParameters();
-        String           queryString      = String.format("v.\"%s\":(%s)", Constants.ENTITY_TEXT_PROPERTY_KEY, searchParameters.getQuery());
+        StringBuilder    queryString      = new StringBuilder();
 
-        indexQuery = context.getGraph().indexQuery(Constants.FULLTEXT_INDEX, queryString);
+        queryString.append("v.\"").append(Constants.ENTITY_TEXT_PROPERTY_KEY).append("\":(").append(searchParameters.getQuery());
+
+        // if search includes entity-type criteria, adding a filter here can help avoid unnecessary
+        // processing (and rejection) by subsequent EntitySearchProcessor
+        if (context.getEntityType() != null) {
+            Set<String> typeAndSubTypeNames = context.getEntityType().getTypeAndAllSubTypes();
+
+            if (typeAndSubTypeNames.size() <= MAX_ENTITY_TYPES_IN_INDEX_QUERY) {
+                queryString.append(AND_STR).append("(").append(StringUtils.join(typeAndSubTypeNames, SPACE_STRING)).append(")");
+            } else {
+                LOG.warn("'{}' has too many subtypes ({}) to include in index-query; might cause poor performance",
+                         context.getEntityType().getTypeName(), typeAndSubTypeNames.size());
+            }
+        }
+
+        // if search includes classification criteria, adding a filter here can help avoid unnecessary
+        // processing (and rejection) by subsequent ClassificationSearchProcessor or EntitySearchProcessor
+        if (context.getClassificationType() != null) {
+            Set<String> typeAndSubTypeNames = context.getClassificationType().getTypeAndAllSubTypes();
+
+            if (typeAndSubTypeNames.size() <= MAX_CLASSIFICATION_TYPES_IN_INDEX_QUERY) {
+                queryString.append(AND_STR).append("(").append(StringUtils.join(typeAndSubTypeNames, SPACE_STRING)).append(")");
+            } else {
+                LOG.warn("'{}' has too many subtypes ({}) to include in index-query; might cause poor performance",
+                        context.getEntityType().getTypeName(), typeAndSubTypeNames.size());
+            }
+        }
+
+        queryString.append(")");
+
+        indexQuery = context.getGraph().indexQuery(Constants.FULLTEXT_INDEX, queryString.toString());
     }
 
     @Override

repository/src/main/java/org/apache/atlas/discovery/SearchContext.java

@@ -61,7 +61,6 @@ public class SearchContext {
 
         if (needEntityProcessor()) {
             addProcessor(new EntitySearchProcessor(this));
-
         }
     }
 
@@ -104,15 +103,15 @@ public class SearchContext {
         return toString(new StringBuilder()).toString();
     }
 
-    public boolean needFullTextrocessor() {
+    boolean needFullTextrocessor() {
         return StringUtils.isNotEmpty(searchParameters.getQuery());
     }
 
-    public boolean needClassificationProcessor() {
-        return classificationType != null;
+    boolean needClassificationProcessor() {
+        return classificationType != null && (entityType == null || hasAttributeFilter(searchParameters.getTagFilters()));
     }
 
-    public boolean needEntityProcessor() {
+    boolean needEntityProcessor() {
         return entityType != null;
     }
 

repository/src/main/java/org/apache/atlas/discovery/SearchProcessor.java

@@ -181,7 +181,7 @@ public abstract class SearchProcessor {
     protected void constructTypeTestQuery(StringBuilder solrQuery, Set<String> typeAndAllSubTypes) {
         String typeAndSubtypesString = StringUtils.join(typeAndAllSubTypes, SPACE_STRING);
 
-        solrQuery.append("v.\"__typeName\": (")
+        solrQuery.append("v.\"").append(Constants.TYPE_NAME_PROPERTY_KEY).append("\": (")
                 .append(typeAndSubtypesString)
                 .append(")");
     }
@@ -206,7 +206,7 @@ public abstract class SearchProcessor {
                 solrQuery.append(AND_STR);
             }
 
-            solrQuery.append("v.\"__state\":").append("ACTIVE");
+            solrQuery.append("v.\"").append(Constants.STATE_PROPERTY_KEY).append("\":ACTIVE");
         }
     }