diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntitySearchProcessor.java b/repository/src/main/java/org/apache/atlas/discovery/EntitySearchProcessor.java index e002ecf7fe1..b110f195166 100644 --- a/repository/src/main/java/org/apache/atlas/discovery/EntitySearchProcessor.java +++ b/repository/src/main/java/org/apache/atlas/discovery/EntitySearchProcessor.java @@ -57,8 +57,9 @@ import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.SortOrder.DESC; public class EntitySearchProcessor extends SearchProcessor { - private static final Logger LOG = LoggerFactory.getLogger(EntitySearchProcessor.class); - private static final Logger PERF_LOG = AtlasPerfTracer.getPerfLogger("EntitySearchProcessor"); + private static final Logger LOG = LoggerFactory.getLogger(EntitySearchProcessor.class); + private static final Logger PERF_LOG = AtlasPerfTracer.getPerfLogger("EntitySearchProcessor"); + private static final int COUNT_BATCH_SIZE = 1000; private final AtlasIndexQuery indexQuery; private final AtlasGraphQuery graphQuery; @@ -316,7 +317,9 @@ public List execute() { @Override public long getResultCount() { - if (indexQuery != null) { + if (isEntityRootType()) { //when `typeName` is `_ALL_ENTITY_TYPES` (`isEntityRootType()` + return getFilteredResultCount(); + } else if (indexQuery != null) { //specific types like `hive_table` return indexQuery.vertexTotals(); } else if (graphQuery != null) { return StreamSupport.stream(graphQuery.vertexIds().spliterator(), false).count(); @@ -325,6 +328,82 @@ public long getResultCount() { } } + private long getFilteredResultCount() { + if (indexQuery != null) { + return countFilteredIndexResults(); + } else if (graphQuery != null) { + return countFilteredGraphResults(); + } + + return -1L; + } + + private long countFilteredIndexResults() { + long count = 0; + int qryOffset = 0; + LinkedHashMap offsetEntityVertexMap = new LinkedHashMap<>(); + + while (true) { + offsetEntityVertexMap.clear(); + + if (context.terminateSearch()) { + break; + } + + Iterator idxQueryResult = executeIndexQuery(context, indexQuery, qryOffset, COUNT_BATCH_SIZE); + + offsetEntityVertexMap = getVerticesFromIndexQueryResult(idxQueryResult, offsetEntityVertexMap, qryOffset); + + boolean isLastResultPage = offsetEntityVertexMap.size() < COUNT_BATCH_SIZE; + + offsetEntityVertexMap = applyCountFilters(offsetEntityVertexMap); + count += offsetEntityVertexMap.size(); + + if (isLastResultPage) { + break; + } + + qryOffset += COUNT_BATCH_SIZE; + } + + return count; + } + + private long countFilteredGraphResults() { + long count = 0; + int qryOffset = 0; + + while (true) { + LinkedHashMap offsetEntityVertexMap = new LinkedHashMap<>(); + Iterator queryResult = graphQuery.vertices(qryOffset, COUNT_BATCH_SIZE).iterator(); + + offsetEntityVertexMap = getVertices(queryResult, offsetEntityVertexMap, qryOffset); + + boolean isLastResultPage = offsetEntityVertexMap.size() < COUNT_BATCH_SIZE; + + offsetEntityVertexMap = applyCountFilters(offsetEntityVertexMap); + count += offsetEntityVertexMap.size(); + + if (isLastResultPage) { + break; + } + + qryOffset += COUNT_BATCH_SIZE; + } + + return count; + } + + private LinkedHashMap applyCountFilters(LinkedHashMap offsetEntityVertexMap) { + offsetEntityVertexMap = super.filter(offsetEntityVertexMap, inMemoryPredicate); + + if (graphQueryPredicate != null) { + offsetEntityVertexMap = super.filter(offsetEntityVertexMap, graphQueryPredicate); + } + + return super.filter(offsetEntityVertexMap); + } + @Override public LinkedHashMap filter(LinkedHashMap offsetEntityVertexMap) { LOG.debug("==> EntitySearchProcessor.filter({})", offsetEntityVertexMap.size());