ATLAS-1527 : Batch entity retrievals - DefaultMetadataService.loadEntities

release-log.txt

@@ -9,6 +9,7 @@ ATLAS-1060 Add composite indexes for exact match performance improvements for al
 ATLAS-1127 Modify creation and modification timestamps to Date instead of Long(sumasai)
 
 ALL CHANGES:
+ATLAS-1527 Batch entity retrievals - DefaultMetadataService.loadEntities
 ATLAS-1385 Add configuration property to disable full text mapper (wwojcik via jnhagelb)
 ATLAS-746  After updating a set of entities, response contains only the first entity definition (jnhagelb)
 ATLAS-1510 Consolidate/batch calls to GraphBackedTypeStore.findVertex() (jnhagelb)

repository/src/main/java/org/apache/atlas/repository/MetadataRepository.java

@@ -104,10 +104,21 @@ public interface MetadataRepository {
      * @param guid globally unique identifier for the entity
      * @return entity (typed instance) definition
      * @throws RepositoryException
+     * @throws EntityNotFoundException
      */
     ITypedReferenceableInstance getEntityDefinition(String guid) throws RepositoryException, EntityNotFoundException;
 
     /**
+     * Fetch the complete entity definitions for the entities with the given GUIDs
+     *
+     * @param guids globally unique identifiers for the entities
+     * @return entity (typed instance) definitions list
+     * @throws RepositoryException
+     * @throws EntityNotFoundException
+     */
+    List<ITypedReferenceableInstance> getEntityDefinitions(String... guids) throws RepositoryException, EntityNotFoundException;
+
+    /**
      * Gets the list of entities for a given entity type.
      *
      * @param entityType name of a type which is unique

repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepository.java

@@ -19,11 +19,15 @@
 package org.apache.atlas.repository.graph;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
@@ -169,19 +173,70 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
     }
 
     @Override
-    @GraphTransaction
     public ITypedReferenceableInstance getEntityDefinition(String guid) throws RepositoryException, EntityNotFoundException {
+        return getEntityDefinitions(guid).get(0);
+    }
+
+    @Override
+    @GraphTransaction
+    public List<ITypedReferenceableInstance> getEntityDefinitions(String... guids) throws RepositoryException, EntityNotFoundException {
         if (LOG.isDebugEnabled()) {
-            LOG.debug("Retrieving entity with guid={}", guid);
+            LOG.debug("Retrieving entities with guids={}", Arrays.toString(guids));
         }
 
-        AtlasVertex instanceVertex = graphHelper.getVertexForGUID(guid);
+        RequestContext context = RequestContext.get();
+        ITypedReferenceableInstance[] result = new ITypedReferenceableInstance[guids.length];
+
+        // Map of the guids of instances not in the cache to their index(es) in the result.
+        // This is used to put the loaded instances into the location(s) corresponding
+        // to their guid in the result.  Note that a set is needed since guids can
+        // appear more than once in the list.
+        Map<String, Set<Integer>> uncachedGuids = new HashMap<>();
+
+        for (int i = 0; i < guids.length; i++) {
+            String guid = guids[i];
+
+            // First, check the cache.
+            ITypedReferenceableInstance cached = context.getInstance(guid);
+            if (cached != null) {
+                result[i] = cached;
+            } else {
+                Set<Integer> indices = uncachedGuids.get(guid);
+                if (indices == null) {
+                    indices = new HashSet<>(1);
+                    uncachedGuids.put(guid, indices);
+                }
+                indices.add(i);
+            }
+        }
 
-        try {
-            return graphToInstanceMapper.mapGraphToTypedInstance(guid, instanceVertex);
-        } catch (AtlasException e) {
-            throw new RepositoryException(e);
+        List<String> guidsToFetch = new ArrayList<>(uncachedGuids.keySet());
+        Map<String, AtlasVertex> instanceVertices = graphHelper.getVerticesForGUIDs(guidsToFetch);
+
+        // search for missing entities
+        if (instanceVertices.size() != guidsToFetch.size()) {
+            Set<String> missingGuids = new HashSet<String>(guidsToFetch);
+            missingGuids.removeAll(instanceVertices.keySet());
+            if (!missingGuids.isEmpty()) {
+                if (LOG.isDebugEnabled()) {
+                    LOG.debug("Failed to find guids={}", missingGuids);
+                }
+                throw new EntityNotFoundException(
+                    "Could not find entities in the repository with guids: " + missingGuids.toString());
+            }
+        }
+
+        for (String guid : guidsToFetch) {
+            try {
+                ITypedReferenceableInstance entity = graphToInstanceMapper.mapGraphToTypedInstance(guid, instanceVertices.get(guid));
+                for(int index : uncachedGuids.get(guid)) {
+                    result[index] = entity;
+                }
+            } catch (AtlasException e) {
+                throw new RepositoryException(e);
+            }
         }
+        return Arrays.asList(result);
     }
 
     @Override

repository/src/main/java/org/apache/atlas/services/DefaultMetadataService.java

@@ -705,13 +705,8 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
         }
     }
 
-    private List<ITypedReferenceableInstance> loadEntities(List<String> guids) throws EntityNotFoundException,
-            RepositoryException {
-        List<ITypedReferenceableInstance> entities = new ArrayList<>();
-        for (String guid : guids) {
-            entities.add(repository.getEntityDefinition(guid));
-        }
-        return entities;
+    private List<ITypedReferenceableInstance> loadEntities(List<String> guids) throws RepositoryException, EntityNotFoundException {
+        return repository.getEntityDefinitions(guids.toArray(new String[guids.size()]));
     }
 
     private void onTypesUpdated(Map<String, IDataType> typesUpdated) throws AtlasException {

repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedRepositoryHardDeleteTest.java

@@ -80,7 +80,7 @@ public class GraphBackedRepositoryHardDeleteTest extends GraphBackedMetadataRepo
             repositoryService.getEntityDefinition(id);
             fail("Expected EntityNotFoundException");
         } catch(EntityNotFoundException e) {
-            //expected
+            // expected
         }
     }