ATLAS-3367: updated search-result to include approximate-total-count if available

intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java

@@ -52,6 +52,7 @@ public class AtlasSearchResult implements Serializable {
     private AttributeSearchResult          attributes;
     private List<AtlasFullTextResult>      fullTextResult;
     private Map<String, AtlasEntityHeader> referredEntities;
+    private long                           approximateCount = -1;
 
     public AtlasSearchResult() {}
 
@@ -126,6 +127,10 @@ public class AtlasSearchResult implements Serializable {
         this.referredEntities = referredEntities;
     }
 
+    public long getApproximateCount() { return approximateCount; }
+
+    public void setApproximateCount(long approximateCount) { this.approximateCount = approximateCount; }
+
     @Override
     public int hashCode() { return Objects.hash(queryType, searchParameters, queryText, type, classification, entities, attributes, fullTextResult, referredEntities); }
 
@@ -184,6 +189,7 @@ public class AtlasSearchResult implements Serializable {
                 ", attributes=" + attributes +
                 ", fullTextResult=" + fullTextResult +
                 ", referredEntities=" + referredEntities +
+                ", approximateCount=" + approximateCount +
                 '}';
     }
 

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

@@ -347,4 +347,9 @@ public class ClassificationSearchProcessor extends SearchProcessor {
             LOG.debug("<== ClassificationSearchProcessor.filter(): ret.size()={}", entityVertices.size());
         }
     }
+
+    @Override
+    public long getResultCount() {
+        return (indexQuery != null) ? indexQuery.vertexTotals() : -1;
+    }
 }

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

@@ -180,6 +180,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
             LOG.debug("Executing Full text query: {}", fullTextQuery);
         }
         ret.setFullTextResult(getIndexQueryResults(idxQuery, params, excludeDeletedEntities));
+        ret.setApproximateCount(idxQuery.vertexTotals());
 
         scrubSearchResults(ret);
 
@@ -285,7 +286,8 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
             int          resultIdx  = 0;
 
             for (int indexQueryOffset = 0; ; indexQueryOffset += getMaxResultSetSize()) {
-                final Iterator<Result<?, ?>> qryResult = graph.indexQuery(Constants.FULLTEXT_INDEX, idxQuery, indexQueryOffset).vertices();
+                final AtlasIndexQuery        qry       = graph.indexQuery(Constants.FULLTEXT_INDEX, idxQuery, indexQueryOffset);
+                final Iterator<Result<?, ?>> qryResult = qry.vertices();
 
                 if (LOG.isDebugEnabled()) {
                     LOG.debug("indexQuery: query=" + idxQuery + "; offset=" + indexQueryOffset);
@@ -344,6 +346,10 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
                     }
                 }
 
+                if (ret.getApproximateCount() < 0) {
+                    ret.setApproximateCount(qry.vertexTotals());
+                }
+
                 if (ret.getEntities() != null && ret.getEntities().size() == resultSize) {
                     break;
                 }
@@ -393,7 +399,6 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
                     if (firstElement instanceof AtlasVertex) {
                         for (Object element : queryResult) {
                             if (element instanceof AtlasVertex) {
-
                                 ret.addEntity(entityRetriever.toAtlasEntityHeader((AtlasVertex) element));
                             } else {
                                 LOG.warn("searchUsingBasicQuery({}): expected an AtlasVertex; found unexpected entry in result {}", basicQuery, element);
@@ -465,6 +470,8 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
         try {
             List<AtlasVertex> resultList = searchContext.getSearchProcessor().execute();
 
+            ret.setApproximateCount(searchContext.getSearchProcessor().getResultCount());
+
             // By default any attribute that shows up in the search parameter should be sent back in the response
             // If additional values are requested then the entityAttributes will be a superset of the all search attributes
             // and the explicitly requested attribute(s)

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

@@ -311,4 +311,9 @@ public class EntitySearchProcessor extends SearchProcessor {
             LOG.debug("<== EntitySearchProcessor.filter(): ret.size()={}", entityVertices.size());
         }
     }
+
+    @Override
+    public long getResultCount() {
+        return (indexQuery != null) ? indexQuery.vertexTotals() : -1;
+    }
 }

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

@@ -174,4 +174,9 @@ public class FreeTextSearchProcessor extends SearchProcessor {
 
         return ret;
     }
+
+    @Override
+    public long getResultCount() {
+        return indexQuery.vertexTotals();
+    }
 }

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

@@ -166,4 +166,9 @@ public class FullTextSearchProcessor extends SearchProcessor {
 
         return ret;
     }
+
+    @Override
+    public long getResultCount() {
+        return indexQuery.vertexTotals();
+    }
 }

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

@@ -128,6 +128,7 @@ public abstract class SearchProcessor {
     }
 
     public abstract List<AtlasVertex> execute();
+    public abstract long getResultCount();
 
     protected int collectResultVertices(final List<AtlasVertex> ret, final int startIdx, final int limit, int resultIdx, final List<AtlasVertex> entityVertices) {
         for (AtlasVertex entityVertex : entityVertices) {

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

@@ -110,4 +110,9 @@ public class TermSearchProcessor extends SearchProcessor {
             LOG.debug("<== TermSearchProcessor.filter(): ret.size()={}", entityVertices.size());
         }
     }
+
+    @Override
+    public long getResultCount() {
+        return -1;
+    }
 }