From 3858e4fabe1b20835fbecbd81f53e895daf9dc6f Mon Sep 17 00:00:00 2001
From: naveenaechan <naveen.aechan@atlan.com>
Date: Wed, 15 Nov 2023 20:15:52 +0530
Subject: [PATCH] PLT-191: enable mapVertexToAtlasEntityHeader2

---
 .../java/org/apache/atlas/discovery/EntityDiscoveryService.java | 2 +-
 .../atlas/repository/store/graph/v2/EntityGraphRetriever.java   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
index 112029cb88..34236c7044 100644
--- a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
+++ b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
@@ -1132,7 +1132,7 @@ private void prepareSearchResult2(AtlasSearchResult ret, DirectIndexQueryResult
     }
 
     private AtlasEntityHeader enrichVertex(AtlasSearchResult ret, Set<String> resultAttributes, Result result, AtlasVertex vertex, boolean fetchCollapsedResults) throws AtlasBaseException {
-        AtlasEntityHeader header = entityRetriever.toAtlasEntityHeader(vertex, resultAttributes);
+        AtlasEntityHeader header = entityRetriever.toAtlasEntityHeader2(vertex, resultAttributes);
         if(RequestContext.get().includeClassifications()){
             header.setClassifications(entityRetriever.getAllClassifications(vertex));
         }
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
index 47d16210e1..8e4a1777fe 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -1034,7 +1034,7 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeader2(AtlasVertex entityVertex
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
 
         if (entityType != null) {
-            entityType.getHeaderAttributes().values().forEach(headerAttribute->{
+            entityType.getHeaderAttributes().values().stream().parallel().forEach(headerAttribute->{
                 try {
                     Object attrValue = getVertexAttribute(entityVertex, headerAttribute);
                     if (attrValue != null) {