@utaba/deep-memory-storage-cosmosdb 0.9.0 → 0.9.1

package/dist/index.cjs

@@ -919,6 +919,7 @@ async function exploreNeighborhood(conn, repositoryId, entityId, options) {
   for (let depth = 0; depth < options.depth; depth++) {
     const layer = {};
     const nextFrontier = /* @__PURE__ */ new Set();
+    const layerCache = /* @__PURE__ */ new Map();
     for (const currentId of frontier) {
       const rels = await getRelationshipsForEntity(conn, repositoryId, currentId, options);
       for (const rel of rels) {
@@ -932,20 +933,20 @@ async function exploreNeighborhood(conn, repositoryId, entityId, options) {
         } else {
           connectedId = rel.sourceEntityId;
         }
-        if (!visited.has(connectedId)) {
-          if (options.entityTypes && options.entityTypes.length > 0) {
-            const entity = await getEntityLight(conn, repositoryId, connectedId);
-            if (!entity || !options.entityTypes.includes(entity.entityType)) continue;
-            layer[relType].entities.push(entity);
+        if (visited.has(connectedId)) continue;
+        let entity = layerCache.get(connectedId);
+        if (entity === void 0) {
+          const loaded = await getEntityLight(conn, repositoryId, connectedId);
+          if (loaded && (!options.entityTypes || options.entityTypes.length === 0 || options.entityTypes.includes(loaded.entityType))) {
+            entity = loaded;
           } else {
-            const entity = await getEntityLight(conn, repositoryId, connectedId);
-            if (entity) {
-              layer[relType].entities.push(entity);
-            }
+            entity = null;
           }
-          visited.add(connectedId);
-          nextFrontier.add(connectedId);
+          layerCache.set(connectedId, entity);
+          if (entity) nextFrontier.add(connectedId);
         }
+        if (!entity) continue;
+        layer[relType].entities.push(entity);
         layer[relType].relationships.push(rel);
         layer[relType].total = layer[relType].entities.length;
       }
@@ -960,6 +961,9 @@ async function exploreNeighborhood(conn, repositoryId, entityId, options) {
     if (Object.keys(layer).length > 0) {
       layers.push(layer);
     }
+    for (const id of nextFrontier) {
+      visited.add(id);
+    }
     frontier = nextFrontier;
     if (frontier.size === 0) break;
   }