Skip to content

Commit 621e535

Browse files
collado-mikecollado-mike
committed
Add loadResolvedEntities by id and entity cache support
1 parent eee219b commit 621e535

File tree

10 files changed

+1011
-171
lines changed

10 files changed

+1011
-171
lines changed

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

Lines changed: 39 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1804,26 +1804,53 @@ public ResolvedEntitiesResult loadResolvedEntities(
18041804
return entities.get(i);
18051805
}
18061806
})
1807-
.map(
1808-
e -> {
1809-
if (e == null) {
1807+
.map(e -> toResolvedPolarisEntity(callCtx, e, ms))
1808+
.collect(Collectors.toList());
1809+
return new ResolvedEntitiesResult(ret);
1810+
}
1811+
1812+
@Nonnull
1813+
@Override
1814+
public ResolvedEntitiesResult loadResolvedEntities(
1815+
@Nonnull PolarisCallContext callCtx,
1816+
@Nonnull PolarisEntityType entityType,
1817+
@Nonnull List<PolarisEntityId> entityIds) {
1818+
BasePersistence ms = callCtx.getMetaStore();
1819+
List<PolarisBaseEntity> entities = ms.lookupEntities(callCtx, entityIds);
1820+
1821+
// mimic the behavior of loadEntity above, return null if not found or type mismatch
1822+
List<ResolvedPolarisEntity> ret =
1823+
IntStream.range(0, entityIds.size())
1824+
.mapToObj(
1825+
i -> {
1826+
if (entities.get(i) != null && !entities.get(i).getType().equals(entityType)) {
18101827
return null;
18111828
} else {
1812-
// load the grant records
1813-
final List<PolarisGrantRecord> grantRecordsAsSecurable =
1814-
ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId());
1815-
final List<PolarisGrantRecord> grantRecordsAsGrantee =
1816-
e.getType().isGrantee()
1817-
? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId())
1818-
: List.of();
1819-
return new ResolvedPolarisEntity(
1820-
PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable);
1829+
return entities.get(i);
18211830
}
18221831
})
1832+
.map(e -> toResolvedPolarisEntity(callCtx, e, ms))
18231833
.collect(Collectors.toList());
18241834
return new ResolvedEntitiesResult(ret);
18251835
}
18261836

1837+
private static ResolvedPolarisEntity toResolvedPolarisEntity(
1838+
PolarisCallContext callCtx, PolarisBaseEntity e, BasePersistence ms) {
1839+
if (e == null) {
1840+
return null;
1841+
} else {
1842+
// load the grant records
1843+
final List<PolarisGrantRecord> grantRecordsAsSecurable =
1844+
ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId());
1845+
final List<PolarisGrantRecord> grantRecordsAsGrantee =
1846+
e.getType().isGrantee()
1847+
? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId())
1848+
: List.of();
1849+
return new ResolvedPolarisEntity(
1850+
PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable);
1851+
}
1852+
}
1853+
18271854
/** {@inheritDoc} */
18281855
@Override
18291856
public @Nonnull ResolvedEntityResult refreshResolvedEntity(

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

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ ListEntitiesResult listEntities(
137137
/**
138138
* Load full entities matching the given criteria with pagination. If only the entity name/id/type
139139
* is required, use {@link #listEntities} instead. If no pagination is required, use {@link
140-
* #loadFullEntitiesAll} instead.
140+
* #listFullEntitiesAll} instead.
141141
*
142142
* @param callCtx call context
143143
* @param catalogPath path inside a catalog. If null or empty, the entities to list are top-level,
@@ -166,7 +166,7 @@ Page<PolarisBaseEntity> listFullEntities(
166166
* @param entitySubType subType of entities to list (or ANY_SUBTYPE)
167167
* @return list of all matching entities
168168
*/
169-
default @Nonnull List<PolarisBaseEntity> loadFullEntitiesAll(
169+
default @Nonnull List<PolarisBaseEntity> listFullEntitiesAll(
170170
@Nonnull PolarisCallContext callCtx,
171171
@Nullable List<PolarisEntityCore> catalogPath,
172172
@Nonnull PolarisEntityType entityType,
@@ -434,6 +434,23 @@ ResolvedEntitiesResult loadResolvedEntities(
434434
@Nonnull PolarisCallContext callCtx,
435435
@Nonnull List<EntityNameLookupRecord> entityLookupRecords);
436436

437+
/**
438+
* Load a batch of resolved entities of a specified entity type given their {@link
439+
* PolarisEntityId}. Will return an empty list if the input list is empty. Order in that returned
440+
* list is the same as the input list. Some elements might be NULL if the entity has been dropped.
441+
*
442+
* @param callCtx call context
443+
* @param entityType the type of entities to load
444+
* @param entityIds the list of entity ids to load
445+
* @return a non-null list of entities corresponding to the lookup keys. Some elements might be
446+
* NULL if the entity has been dropped.
447+
*/
448+
@Nonnull
449+
ResolvedEntitiesResult loadResolvedEntities(
450+
@Nonnull PolarisCallContext callCtx,
451+
@Nonnull PolarisEntityType entityType,
452+
@Nonnull List<PolarisEntityId> entityIds);
453+
437454
/**
438455
* Refresh a resolved entity from the backend store. Will return NULL if the entity does not
439456
* 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‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,16 @@ public ResolvedEntitiesResult loadResolvedEntities(
389389
return null;
390390
}
391391

392+
@Nonnull
393+
@Override
394+
public ResolvedEntitiesResult loadResolvedEntities(
395+
@Nonnull PolarisCallContext callCtx,
396+
@Nonnull PolarisEntityType entityType,
397+
@Nonnull List<PolarisEntityId> entityIds) {
398+
diagnostics.fail("illegal_method_in_transaction_workspace", "loadResolvedEntities");
399+
return null;
400+
}
401+
392402
@Override
393403
public ResolvedEntityResult refreshResolvedEntity(
394404
@Nonnull PolarisCallContext callCtx,

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,11 @@
2020

2121
import jakarta.annotation.Nonnull;
2222
import jakarta.annotation.Nullable;
23+
import java.util.List;
2324
import org.apache.polaris.core.PolarisCallContext;
25+
import org.apache.polaris.core.entity.EntityNameLookupRecord;
2426
import org.apache.polaris.core.entity.PolarisBaseEntity;
27+
import org.apache.polaris.core.entity.PolarisEntityId;
2528
import org.apache.polaris.core.entity.PolarisEntityType;
2629
import org.apache.polaris.core.persistence.ResolvedPolarisEntity;
2730

@@ -80,4 +83,36 @@ EntityCacheLookupResult getOrLoadEntityById(
8083
@Nullable
8184
EntityCacheLookupResult getOrLoadEntityByName(
8285
@Nonnull PolarisCallContext callContext, @Nonnull EntityCacheByNameKey entityNameKey);
86+
87+
/**
88+
* Load multiple entities by id, returning those found in the cache and loading those not found.
89+
*
90+
* @param callCtx the Polaris call context
91+
* @param entityType the entity type
92+
* @param entityIds the list of entity ids to load
93+
* @return the list of resolved entities, in the same order as the requested entity ids. As in
94+
* {@link
95+
* org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadResolvedEntities(PolarisCallContext,
96+
* PolarisEntityType, List)}, elements in the returned list may be null if the corresponding
97+
* entity id does not exist.
98+
*/
99+
List<EntityCacheLookupResult> getOrLoadResolvedEntities(
100+
@Nonnull PolarisCallContext callCtx,
101+
@Nonnull PolarisEntityType entityType,
102+
@Nonnull List<PolarisEntityId> entityIds);
103+
104+
/**
105+
* Load multiple entities by {@link EntityNameLookupRecord}, returning those found in the cache
106+
* and loading those not found.
107+
*
108+
* @param callCtx the Polaris call context
109+
* @param lookupRecords the list of entity name to load
110+
* @return the list of resolved entities, in the same order as the requested entity records. As in
111+
* {@link
112+
* org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadResolvedEntities(PolarisCallContext,
113+
* PolarisEntityType, List)}, elements in the returned list may be null if the corresponding
114+
* entity id does not exist.
115+
*/
116+
List<EntityCacheLookupResult> getOrLoadResolvedEntities(
117+
@Nonnull PolarisCallContext callCtx, @Nonnull List<EntityNameLookupRecord> lookupRecords);
83118
}

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

Lines changed: 164 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,24 +24,39 @@
2424
import jakarta.annotation.Nonnull;
2525
import jakarta.annotation.Nullable;
2626
import java.util.AbstractMap;
27+
import java.util.ArrayList;
28+
import java.util.HashMap;
29+
import java.util.Iterator;
2730
import java.util.List;
31+
import java.util.Map;
32+
import java.util.Objects;
33+
import java.util.Optional;
2834
import java.util.concurrent.ConcurrentHashMap;
2935
import java.util.concurrent.TimeUnit;
36+
import java.util.function.Function;
37+
import java.util.stream.Collectors;
3038
import org.apache.polaris.core.PolarisCallContext;
3139
import org.apache.polaris.core.PolarisDiagnostics;
3240
import org.apache.polaris.core.config.BehaviorChangeConfiguration;
3341
import org.apache.polaris.core.config.FeatureConfiguration;
3442
import org.apache.polaris.core.config.RealmConfig;
43+
import org.apache.polaris.core.entity.EntityNameLookupRecord;
3544
import org.apache.polaris.core.entity.PolarisBaseEntity;
45+
import org.apache.polaris.core.entity.PolarisChangeTrackingVersions;
46+
import org.apache.polaris.core.entity.PolarisEntityId;
3647
import org.apache.polaris.core.entity.PolarisEntityType;
3748
import org.apache.polaris.core.entity.PolarisGrantRecord;
3849
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
3950
import org.apache.polaris.core.persistence.ResolvedPolarisEntity;
51+
import org.apache.polaris.core.persistence.dao.entity.ChangeTrackingResult;
52+
import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult;
4053
import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult;
54+
import org.slf4j.Logger;
55+
import org.slf4j.LoggerFactory;
4156

4257
/** An in-memory entity cache with a limit of 100k entities and a 1h TTL. */
4358
public class InMemoryEntityCache implements EntityCache {
44-
59+
private static final Logger LOGGER = LoggerFactory.getLogger(InMemoryEntityCache.class);
4560
private EntityCacheMode cacheMode;
4661
private final PolarisDiagnostics diagnostics;
4762
private final PolarisMetaStoreManager polarisMetaStoreManager;
@@ -473,4 +488,152 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) {
473488
// return what we found
474489
return new EntityCacheLookupResult(entry, cacheHit);
475490
}
491+
492+
@Override
493+
public List<EntityCacheLookupResult> getOrLoadResolvedEntities(
494+
@Nonnull PolarisCallContext callCtx,
495+
@Nonnull PolarisEntityType entityType,
496+
@Nonnull List<PolarisEntityId> entityIds) {
497+
// use a map to collect cached entries to avoid concurrency problems in case a second thread is
498+
// trying to populate
499+
// the cache from a different snapshot
500+
Map<PolarisEntityId, ResolvedPolarisEntity> resolvedEntities = new HashMap<>();
501+
for (int i = 0; i < 100; i++) {
502+
Function<List<PolarisEntityId>, ResolvedEntitiesResult> loaderFunc =
503+
idsToLoad -> polarisMetaStoreManager.loadResolvedEntities(callCtx, entityType, idsToLoad);
504+
if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) {
505+
break;
506+
}
507+
}
508+
509+
return entityIds.stream()
510+
.map(
511+
id -> {
512+
ResolvedPolarisEntity entity = resolvedEntities.get(id);
513+
return entity == null ? null : new EntityCacheLookupResult(entity, true);
514+
})
515+
.collect(Collectors.toList());
516+
}
517+
518+
@Override
519+
public List<EntityCacheLookupResult> getOrLoadResolvedEntities(
520+
@Nonnull PolarisCallContext callCtx, @Nonnull List<EntityNameLookupRecord> lookupRecords) {
521+
Map<PolarisEntityId, EntityNameLookupRecord> entityIdMap =
522+
lookupRecords.stream()
523+
.collect(
524+
Collectors.toMap(
525+
e -> new PolarisEntityId(e.getCatalogId(), e.getId()),
526+
Function.identity(),
527+
(a, b) -> a));
528+
Function<List<PolarisEntityId>, ResolvedEntitiesResult> loaderFunc =
529+
idsToLoad ->
530+
polarisMetaStoreManager.loadResolvedEntities(
531+
callCtx, idsToLoad.stream().map(entityIdMap::get).collect(Collectors.toList()));
532+
533+
// use a map to collect cached entries to avoid concurrency problems in case a second thread is
534+
// trying to populate
535+
// the cache from a different snapshot
536+
Map<PolarisEntityId, ResolvedPolarisEntity> resolvedEntities = new HashMap<>();
537+
List<PolarisEntityId> entityIds =
538+
lookupRecords.stream()
539+
.map(e -> new PolarisEntityId(e.getCatalogId(), e.getId()))
540+
.collect(Collectors.toList());
541+
for (int i = 0; i < 100; i++) {
542+
if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) {
543+
break;
544+
}
545+
}
546+
547+
return lookupRecords.stream()
548+
.map(
549+
lookupRecord -> {
550+
ResolvedPolarisEntity entity =
551+
resolvedEntities.get(
552+
new PolarisEntityId(lookupRecord.getCatalogId(), lookupRecord.getId()));
553+
return entity == null ? null : new EntityCacheLookupResult(entity, true);
554+
})
555+
.collect(Collectors.toList());
556+
}
557+
558+
private boolean isCacheStateValid(
559+
@Nonnull PolarisCallContext callCtx,
560+
@Nonnull Map<PolarisEntityId, ResolvedPolarisEntity> resolvedEntities,
561+
@Nonnull List<PolarisEntityId> entityIds,
562+
@Nonnull Function<List<PolarisEntityId>, ResolvedEntitiesResult> loaderFunc) {
563+
ChangeTrackingResult changeTrackingResult =
564+
polarisMetaStoreManager.loadEntitiesChangeTracking(callCtx, entityIds);
565+
List<PolarisEntityId> idsToLoad = new ArrayList<>();
566+
if (changeTrackingResult.isSuccess()) {
567+
idsToLoad.addAll(validateCacheEntries(entityIds, resolvedEntities, changeTrackingResult));
568+
} else {
569+
idsToLoad.addAll(entityIds);
570+
}
571+
if (!idsToLoad.isEmpty()) {
572+
ResolvedEntitiesResult resolvedEntitiesResult = loaderFunc.apply(idsToLoad);
573+
if (resolvedEntitiesResult.isSuccess()) {
574+
LOGGER.debug("Resolved entities - validating cache");
575+
resolvedEntitiesResult.getResolvedEntities().stream()
576+
.filter(Objects::nonNull)
577+
.forEach(
578+
e -> {
579+
this.cacheNewEntry(e);
580+
resolvedEntities.put(
581+
new PolarisEntityId(e.getEntity().getCatalogId(), e.getEntity().getId()), e);
582+
});
583+
}
584+
}
585+
586+
// the loader function should always return a batch of results from the same "snapshot" of the
587+
// persistence, so
588+
// if the changeTracking call above failed, we should have loaded the entire batch in one shot.
589+
// There should be no
590+
// need to revalidate the entities.
591+
List<PolarisEntityId> idsToReload =
592+
changeTrackingResult.isSuccess()
593+
? validateCacheEntries(entityIds, resolvedEntities, changeTrackingResult)
594+
: List.of();
595+
return idsToReload.isEmpty();
596+
}
597+
598+
private List<PolarisEntityId> validateCacheEntries(
599+
List<PolarisEntityId> entityIds,
600+
Map<PolarisEntityId, ResolvedPolarisEntity> resolvedEntities,
601+
ChangeTrackingResult changeTrackingResult) {
602+
List<PolarisEntityId> idsToReload = new ArrayList<>();
603+
Iterator<PolarisEntityId> idIterator = entityIds.iterator();
604+
Iterator<PolarisChangeTrackingVersions> changeTrackingIterator =
605+
changeTrackingResult.getChangeTrackingVersions().iterator();
606+
while (idIterator.hasNext() && changeTrackingIterator.hasNext()) {
607+
PolarisEntityId entityId = idIterator.next();
608+
PolarisChangeTrackingVersions changeTrackingVersions = changeTrackingIterator.next();
609+
if (changeTrackingVersions == null) {
610+
// entity has been purged
611+
ResolvedPolarisEntity cachedEntity = getEntityById(entityId.getId());
612+
if (cachedEntity != null || resolvedEntities.containsKey(entityId)) {
613+
LOGGER.debug("Entity {} has been purged, removing from cache", entityId);
614+
Optional.ofNullable(cachedEntity).ifPresent(this::removeCacheEntry);
615+
resolvedEntities.remove(entityId);
616+
}
617+
continue;
618+
}
619+
// compare versions using equals rather than less than so we can use the same function to
620+
// validate that the cache
621+
// entries are consistent with a single call to the change tracking table, rather than some
622+
// grants ahead and some
623+
// grants behind
624+
ResolvedPolarisEntity cachedEntity =
625+
resolvedEntities.computeIfAbsent(entityId, id -> this.getEntityById(id.getId()));
626+
if (cachedEntity == null
627+
|| cachedEntity.getEntity().getEntityVersion()
628+
!= changeTrackingVersions.getEntityVersion()
629+
|| cachedEntity.getEntity().getGrantRecordsVersion()
630+
!= changeTrackingVersions.getGrantRecordsVersion()) {
631+
idsToReload.add(entityId);
632+
} else {
633+
resolvedEntities.put(entityId, cachedEntity);
634+
}
635+
}
636+
LOGGER.debug("Cache entries {} need to be reloaded", idsToReload);
637+
return idsToReload;
638+
}
476639
}

0 commit comments

Comments
 (0)