Merge pull request #3276 from atlanhq/feat/add-lineage-vertical-node-…

…count [beta] LIN-974: add horizontal pagination node count
atlanhq · Jun 24, 2024 · 05a40f5 · 05a40f5
2 parents 85c4106 + 2f633b9
commit 05a40f5
Showing 1 changed file with 10 additions and 23 deletions.
diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java
@@ -612,7 +612,7 @@ private boolean incrementAndCheckIfRelationsLimitReached(AtlasEdge atlasEdge, bo
 
         setHorizontalPaginationFlags(isInput, atlasLineageOnDemandContext, ret, depth, entitiesTraversed, inVertex, inGuid, outVertex, outGuid, inLineageInfo, outLineageInfo, visitedVertices);
 
-        boolean hasRelationsLimitReached = setVerticalPaginationFlags(depth, atlasLineageOnDemandContext, isInput, entitiesTraversed, inLineageInfo, outLineageInfo, inVertex, outVertex);
+        boolean hasRelationsLimitReached = setVerticalPaginationFlags(entitiesTraversed, inLineageInfo, outLineageInfo);
         if (!hasRelationsLimitReached) {
             ret.getRelationsOnDemand().put(inGuid, inLineageInfo);
             ret.getRelationsOnDemand().put(outGuid, outLineageInfo);
@@ -622,29 +622,12 @@ private boolean incrementAndCheckIfRelationsLimitReached(AtlasEdge atlasEdge, bo
         return hasRelationsLimitReached;
     }
 
-    private boolean setVerticalPaginationFlags(int depth, AtlasLineageOnDemandContext atlasLineageOnDemandContext, boolean isInput, AtomicInteger entitiesTraversed, LineageInfoOnDemand inLineageInfo, LineageInfoOnDemand outLineageInfo, AtlasVertex inVertex, AtlasVertex outVertex) {
+    private boolean setVerticalPaginationFlags(AtomicInteger entitiesTraversed, LineageInfoOnDemand inLineageInfo, LineageInfoOnDemand outLineageInfo) {
         boolean hasRelationsLimitReached = false;
         if (inLineageInfo.isInputRelationsReachedLimit() || outLineageInfo.isOutputRelationsReachedLimit() || isEntityTraversalLimitReached(entitiesTraversed)) {
             inLineageInfo.setHasMoreInputs(true);
             outLineageInfo.setHasMoreOutputs(true);
             hasRelationsLimitReached = true;
-
-            if (depth >= 0) {
-                boolean inIsProcess = Objects.equals(AtlasGraphUtilsV2.getTypeName(inVertex), PROCESS_SUPER_TYPE);
-                AtlasEdgeDirection inDirection = inIsProcess ? OUT: IN;     // Process Node edges always points outwards
-                AtlasEdgeDirection outDirection = !inIsProcess ? OUT: IN;   // Data Node edges always points inwards
-                String edgeLabel = isInput ? PROCESS_INPUTS_EDGE : PROCESS_OUTPUTS_EDGE;
-
-                if (outLineageInfo.getTotalOutputRelationsCount() == 0) {
-                    List<AtlasEdge> outFilteredEdges = getFilteredAtlasEdges(outVertex, outDirection, edgeLabel, atlasLineageOnDemandContext, false);
-                    outLineageInfo.setTotalOutputRelationsCount(outFilteredEdges.size());
-                }
-
-                if (inLineageInfo.getTotalInputRelationsCount() == 0) {
-                    List<AtlasEdge> inFilteredEdges = getFilteredAtlasEdges(inVertex, inDirection, edgeLabel, atlasLineageOnDemandContext, false);
-                    inLineageInfo.setTotalInputRelationsCount(inFilteredEdges.size());
-                }
-            }
         }
 
         if (!hasRelationsLimitReached) {
@@ -666,15 +649,19 @@ else if (!isInput && ! isInVertexVisited)
     }
 
     private void setHasDownstream(AtlasLineageOnDemandContext atlasLineageOnDemandContext, AtlasVertex inVertex, LineageInfoOnDemand inLineageInfo) {
-        List<AtlasEdge> filteredEdges = getFilteredAtlasEdges(inVertex, IN, PROCESS_INPUTS_EDGE, atlasLineageOnDemandContext, true);
-        if (!filteredEdges.isEmpty())
+        List<AtlasEdge> filteredEdges = getFilteredAtlasEdges(inVertex, IN, PROCESS_INPUTS_EDGE, atlasLineageOnDemandContext, false);
+        if (!filteredEdges.isEmpty()) {
             inLineageInfo.setHasDownstream(true);
+            inLineageInfo.setTotalOutputRelationsCount(filteredEdges.size());
+        }
     }
 
     private void setHasUpstream(AtlasLineageOnDemandContext atlasLineageOnDemandContext, AtlasVertex outVertex, LineageInfoOnDemand outLineageInfo) {
-        List<AtlasEdge> filteredEdges = getFilteredAtlasEdges(outVertex, IN, PROCESS_OUTPUTS_EDGE, atlasLineageOnDemandContext, true);
-        if (!filteredEdges.isEmpty())
+        List<AtlasEdge> filteredEdges = getFilteredAtlasEdges(outVertex, IN, PROCESS_OUTPUTS_EDGE, atlasLineageOnDemandContext, false);
+        if (!filteredEdges.isEmpty()) {
             outLineageInfo.setHasUpstream(true);
+            outLineageInfo.setTotalInputRelationsCount(filteredEdges.size());
+        }
     }
 
     private List<AtlasEdge> getFilteredAtlasEdges(AtlasVertex outVertex, AtlasEdgeDirection direction, String processEdgeLabel, AtlasLineageOnDemandContext atlasLineageOnDemandContext, boolean hasAnyCheck) {