Add loadResolvedEntities by id and entity cache support

Author: collado-mike

polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisEntity.java

@@ -23,13 +23,14 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
+import org.apache.polaris.core.persistence.dao.entity.EntityResult;
+
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
-import org.apache.polaris.core.persistence.dao.entity.EntityResult;
 
 /**
  * For legacy reasons, this class is only a thin facade over PolarisBaseEntity's members/methods. No

polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java

@@ -1804,26 +1804,53 @@ public ResolvedEntitiesResult loadResolvedEntities(
                     return entities.get(i);
                   }
                 })
-            .map(
-                e -> {
-                  if (e == null) {
+            .map(e -> toResolvedPolarisEntity(callCtx, e, ms))
+            .collect(Collectors.toList());
+    return new ResolvedEntitiesResult(ret);
+  }
+
+  @Nonnull
+  @Override
+  public ResolvedEntitiesResult loadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull List<PolarisEntityId> entityIds) {
+    BasePersistence ms = callCtx.getMetaStore();
+    List<PolarisBaseEntity> entities = ms.lookupEntities(callCtx, entityIds);
+
+    // mimic the behavior of loadEntity above, return null if not found or type mismatch
+    List<ResolvedPolarisEntity> ret =
+        IntStream.range(0, entityIds.size())
+            .mapToObj(
+                i -> {
+                  if (entities.get(i) != null && !entities.get(i).getType().equals(entityType)) {
                     return null;
                   } else {
-                    // load the grant records
-                    final List<PolarisGrantRecord> grantRecordsAsSecurable =
-                        ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId());
-                    final List<PolarisGrantRecord> grantRecordsAsGrantee =
-                        e.getType().isGrantee()
-                            ? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId())
-                            : List.of();
-                    return new ResolvedPolarisEntity(
-                        PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable);
+                    return entities.get(i);
                   }
                 })
+            .map(e -> toResolvedPolarisEntity(callCtx, e, ms))
             .collect(Collectors.toList());
     return new ResolvedEntitiesResult(ret);
   }
 
+  private static ResolvedPolarisEntity toResolvedPolarisEntity(
+      PolarisCallContext callCtx, PolarisBaseEntity e, BasePersistence ms) {
+    if (e == null) {
+      return null;
+    } else {
+      // load the grant records
+      final List<PolarisGrantRecord> grantRecordsAsSecurable =
+          ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId());
+      final List<PolarisGrantRecord> grantRecordsAsGrantee =
+          e.getType().isGrantee()
+              ? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId())
+              : List.of();
+      return new ResolvedPolarisEntity(
+          PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable);
+    }
+  }
+
   /** {@inheritDoc} */
   @Override
   public @Nonnull ResolvedEntityResult refreshResolvedEntity(

polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java

@@ -137,7 +137,7 @@ ListEntitiesResult listEntities(
   /**
    * Load full entities matching the given criteria with pagination. If only the entity name/id/type
    * is required, use {@link #listEntities} instead. If no pagination is required, use {@link
-   * #loadFullEntitiesAll} instead.
+   * #listFullEntitiesAll} instead.
    *
    * @param callCtx call context
    * @param catalogPath path inside a catalog. If null or empty, the entities to list are top-level,
@@ -166,7 +166,7 @@ Page<PolarisBaseEntity> listFullEntities(
    * @param entitySubType subType of entities to list (or ANY_SUBTYPE)
    * @return list of all matching entities
    */
-  default @Nonnull List<PolarisBaseEntity> loadFullEntitiesAll(
+  default @Nonnull List<PolarisBaseEntity> listFullEntitiesAll(
       @Nonnull PolarisCallContext callCtx,
       @Nullable List<PolarisEntityCore> catalogPath,
       @Nonnull PolarisEntityType entityType,
@@ -434,6 +434,23 @@ ResolvedEntitiesResult loadResolvedEntities(
       @Nonnull PolarisCallContext callCtx,
       @Nonnull List<EntityNameLookupRecord> entityLookupRecords);
 
+  /**
+   * Load a batch of resolved entities of a specified entity type given their {@link
+   * PolarisEntityId}. Will return an empty list if the input list is empty. Order in that returned
+   * list is the same as the input list. Some elements might be NULL if the entity has been dropped.
+   *
+   * @param callCtx call context
+   * @param entityType the type of entities to load
+   * @param entityIds the list of entity ids to load
+   * @return a non-null list of entities corresponding to the lookup keys. Some elements might be
+   *     NULL if the entity has been dropped.
+   */
+  @Nonnull
+  ResolvedEntitiesResult loadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull List<PolarisEntityId> entityIds);
+
   /**
    * Refresh a resolved entity from the backend store. Will return NULL if the entity does not
    * exist, i.e. has been purged or dropped. Else, will determine what has changed based on the

polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java

@@ -389,6 +389,16 @@ public ResolvedEntitiesResult loadResolvedEntities(
     return null;
   }
 
+  @Nonnull
+  @Override
+  public ResolvedEntitiesResult loadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull List<PolarisEntityId> entityIds) {
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadResolvedEntities");
+    return null;
+  }
+
   @Override
   public ResolvedEntityResult refreshResolvedEntity(
       @Nonnull PolarisCallContext callCtx,

polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java

@@ -20,8 +20,11 @@
 
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
+import java.util.List;
 import org.apache.polaris.core.PolarisCallContext;
+import org.apache.polaris.core.entity.EntityNameLookupRecord;
 import org.apache.polaris.core.entity.PolarisBaseEntity;
+import org.apache.polaris.core.entity.PolarisEntityId;
 import org.apache.polaris.core.entity.PolarisEntityType;
 import org.apache.polaris.core.persistence.ResolvedPolarisEntity;
 
@@ -80,4 +83,36 @@ EntityCacheLookupResult getOrLoadEntityById(
   @Nullable
   EntityCacheLookupResult getOrLoadEntityByName(
       @Nonnull PolarisCallContext callContext, @Nonnull EntityCacheByNameKey entityNameKey);
+
+  /**
+   * Load multiple entities by id, returning those found in the cache and loading those not found.
+   *
+   * @param callCtx the Polaris call context
+   * @param entityType the entity type
+   * @param entityIds the list of entity ids to load
+   * @return the list of resolved entities, in the same order as the requested entity ids. As in
+   *     {@link
+   *     org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadResolvedEntities(PolarisCallContext,
+   *     PolarisEntityType, List)}, elements in the returned list may be null if the corresponding
+   *     entity id does not exist.
+   */
+  List<EntityCacheLookupResult> getOrLoadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull List<PolarisEntityId> entityIds);
+
+  /**
+   * Load multiple entities by {@link EntityNameLookupRecord}, returning those found in the cache
+   * and loading those not found.
+   *
+   * @param callCtx the Polaris call context
+   * @param lookupRecords the list of entity name to load
+   * @return the list of resolved entities, in the same order as the requested entity records. As in
+   *     {@link
+   *     org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadResolvedEntities(PolarisCallContext,
+   *     PolarisEntityType, List)}, elements in the returned list may be null if the corresponding
+   *     entity id does not exist.
+   */
+  List<EntityCacheLookupResult> getOrLoadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx, @Nonnull List<EntityNameLookupRecord> lookupRecords);
 }

polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java

@@ -24,24 +24,39 @@
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
 import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 import org.apache.polaris.core.PolarisCallContext;
 import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.config.BehaviorChangeConfiguration;
 import org.apache.polaris.core.config.FeatureConfiguration;
 import org.apache.polaris.core.config.RealmConfig;
+import org.apache.polaris.core.entity.EntityNameLookupRecord;
 import org.apache.polaris.core.entity.PolarisBaseEntity;
+import org.apache.polaris.core.entity.PolarisChangeTrackingVersions;
+import org.apache.polaris.core.entity.PolarisEntityId;
 import org.apache.polaris.core.entity.PolarisEntityType;
 import org.apache.polaris.core.entity.PolarisGrantRecord;
 import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
 import org.apache.polaris.core.persistence.ResolvedPolarisEntity;
+import org.apache.polaris.core.persistence.dao.entity.ChangeTrackingResult;
+import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult;
 import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** An in-memory entity cache with a limit of 100k entities and a 1h TTL. */
 public class InMemoryEntityCache implements EntityCache {
-
+  private static final Logger LOGGER = LoggerFactory.getLogger(InMemoryEntityCache.class);
   private EntityCacheMode cacheMode;
   private final PolarisDiagnostics diagnostics;
   private final PolarisMetaStoreManager polarisMetaStoreManager;
@@ -473,4 +488,152 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) {
     // return what we found
     return new EntityCacheLookupResult(entry, cacheHit);
   }
+
+  @Override
+  public List<EntityCacheLookupResult> getOrLoadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull List<PolarisEntityId> entityIds) {
+    // use a map to collect cached entries to avoid concurrency problems in case a second thread is
+    // trying to populate
+    // the cache from a different snapshot
+    Map<PolarisEntityId, ResolvedPolarisEntity> resolvedEntities = new HashMap<>();
+    for (int i = 0; i < 100; i++) {
+      Function<List<PolarisEntityId>, ResolvedEntitiesResult> loaderFunc =
+          idsToLoad -> polarisMetaStoreManager.loadResolvedEntities(callCtx, entityType, idsToLoad);
+      if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) {
+        break;
+      }
+    }
+
+    return entityIds.stream()
+        .map(
+            id -> {
+              ResolvedPolarisEntity entity = resolvedEntities.get(id);
+              return entity == null ? null : new EntityCacheLookupResult(entity, true);
+            })
+        .collect(Collectors.toList());
+  }
+
+  @Override
+  public List<EntityCacheLookupResult> getOrLoadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx, @Nonnull List<EntityNameLookupRecord> lookupRecords) {
+    Map<PolarisEntityId, EntityNameLookupRecord> entityIdMap =
+        lookupRecords.stream()
+            .collect(
+                Collectors.toMap(
+                    e -> new PolarisEntityId(e.getCatalogId(), e.getId()),
+                    Function.identity(),
+                    (a, b) -> a));
+    Function<List<PolarisEntityId>, ResolvedEntitiesResult> loaderFunc =
+        idsToLoad ->
+            polarisMetaStoreManager.loadResolvedEntities(
+                callCtx, idsToLoad.stream().map(entityIdMap::get).collect(Collectors.toList()));
+
+    // use a map to collect cached entries to avoid concurrency problems in case a second thread is
+    // trying to populate
+    // the cache from a different snapshot
+    Map<PolarisEntityId, ResolvedPolarisEntity> resolvedEntities = new HashMap<>();
+    List<PolarisEntityId> entityIds =
+        lookupRecords.stream()
+            .map(e -> new PolarisEntityId(e.getCatalogId(), e.getId()))
+            .collect(Collectors.toList());
+    for (int i = 0; i < 100; i++) {
+      if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) {
+        break;
+      }
+    }
+
+    return lookupRecords.stream()
+        .map(
+            lookupRecord -> {
+              ResolvedPolarisEntity entity =
+                  resolvedEntities.get(
+                      new PolarisEntityId(lookupRecord.getCatalogId(), lookupRecord.getId()));
+              return entity == null ? null : new EntityCacheLookupResult(entity, true);
+            })
+        .collect(Collectors.toList());
+  }
+
+  private boolean isCacheStateValid(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull Map<PolarisEntityId, ResolvedPolarisEntity> resolvedEntities,
+      @Nonnull List<PolarisEntityId> entityIds,
+      @Nonnull Function<List<PolarisEntityId>, ResolvedEntitiesResult> loaderFunc) {
+    ChangeTrackingResult changeTrackingResult =
+        polarisMetaStoreManager.loadEntitiesChangeTracking(callCtx, entityIds);
+    List<PolarisEntityId> idsToLoad = new ArrayList<>();
+    if (changeTrackingResult.isSuccess()) {
+      idsToLoad.addAll(validateCacheEntries(entityIds, resolvedEntities, changeTrackingResult));
+    } else {
+      idsToLoad.addAll(entityIds);
+    }
+    if (!idsToLoad.isEmpty()) {
+      ResolvedEntitiesResult resolvedEntitiesResult = loaderFunc.apply(idsToLoad);
+      if (resolvedEntitiesResult.isSuccess()) {
+        LOGGER.debug("Resolved entities - validating cache");
+        resolvedEntitiesResult.getResolvedEntities().stream()
+            .filter(Objects::nonNull)
+            .forEach(
+                e -> {
+                  this.cacheNewEntry(e);
+                  resolvedEntities.put(
+                      new PolarisEntityId(e.getEntity().getCatalogId(), e.getEntity().getId()), e);
+                });
+      }
+    }
+
+    // the loader function should always return a batch of results from the same "snapshot" of the
+    // persistence, so
+    // if the changeTracking call above failed, we should have loaded the entire batch in one shot.
+    // There should be no
+    // need to revalidate the entities.
+    List<PolarisEntityId> idsToReload =
+        changeTrackingResult.isSuccess()
+            ? validateCacheEntries(entityIds, resolvedEntities, changeTrackingResult)
+            : List.of();
+    return idsToReload.isEmpty();
+  }
+
+  private List<PolarisEntityId> validateCacheEntries(
+      List<PolarisEntityId> entityIds,
+      Map<PolarisEntityId, ResolvedPolarisEntity> resolvedEntities,
+      ChangeTrackingResult changeTrackingResult) {
+    List<PolarisEntityId> idsToReload = new ArrayList<>();
+    Iterator<PolarisEntityId> idIterator = entityIds.iterator();
+    Iterator<PolarisChangeTrackingVersions> changeTrackingIterator =
+        changeTrackingResult.getChangeTrackingVersions().iterator();
+    while (idIterator.hasNext() && changeTrackingIterator.hasNext()) {
+      PolarisEntityId entityId = idIterator.next();
+      PolarisChangeTrackingVersions changeTrackingVersions = changeTrackingIterator.next();
+      if (changeTrackingVersions == null) {
+        // entity has been purged
+        ResolvedPolarisEntity cachedEntity = getEntityById(entityId.getId());
+        if (cachedEntity != null || resolvedEntities.containsKey(entityId)) {
+          LOGGER.debug("Entity {} has been purged, removing from cache", entityId);
+          Optional.ofNullable(cachedEntity).ifPresent(this::removeCacheEntry);
+          resolvedEntities.remove(entityId);
+        }
+        continue;
+      }
+      // compare versions using equals rather than less than so we can use the same function to
+      // validate that the cache
+      // entries are consistent with a single call to the change tracking table, rather than some
+      // grants ahead and some
+      // grants behind
+      ResolvedPolarisEntity cachedEntity =
+          resolvedEntities.computeIfAbsent(entityId, id -> this.getEntityById(id.getId()));
+      if (cachedEntity == null
+          || cachedEntity.getEntity().getEntityVersion()
+              != changeTrackingVersions.getEntityVersion()
+          || cachedEntity.getEntity().getGrantRecordsVersion()
+              != changeTrackingVersions.getGrantRecordsVersion()) {
+        idsToReload.add(entityId);
+      } else {
+        resolvedEntities.put(entityId, cachedEntity);
+      }
+    }
+    LOGGER.debug("Cache entries {} need to be reloaded", idsToReload);
+    return idsToReload;
+  }
 }