Skip to content

Commit

Permalink
Merge pull request #3804 from atlanhq/dev/janusprefetchwith0.6
Browse files Browse the repository at this point in the history 
Janus optimisation
  • Loading branch information
@aarshi0301
aarshi0301 authored Dec 2, 2024
2 parents 9b514cc + 786bd7a commit f719801
Showing 1 changed file with 40 additions and 9 deletions.
49 changes: 40 additions & 9 deletions ...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1010,7 +1010,7 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeader(AtlasVertex entityVertex)
return mapVertexToAtlasEntityHeader(entityVertex, Collections.<String>emptySet());
}

private Map<String, Object> preloadProperties(AtlasVertex entityVertex) {
private Map<String, Object> preloadProperties(AtlasVertex entityVertex, AtlasEntityType entityType, Set attributes) {
Map<String, Object> propertiesMap = new HashMap<>();

// Execute the traversal to fetch properties
Expand All @@ -1020,11 +1020,28 @@ private Map<String, Object> preloadProperties(AtlasVertex entityVertex) {
while (traversal.hasNext()) {
VertexProperty<Object> property = traversal.next();

if (property.isPresent()) { // Ensure the property exists
AtlasAttribute attribute = entityType.getAttribute(property.key());
TypeCategory typeCategory = attribute != null ? attribute.getAttributeType().getTypeCategory() : null;
TypeCategory elementTypeCategory = attribute != null && attribute.getAttributeType().getTypeCategory() == TypeCategory.ARRAY ? ((AtlasArrayType) attribute.getAttributeType()).getElementType().getTypeCategory() : null;

if (propertiesMap.get(property.key()) == null) {
propertiesMap.put(property.key(), property.value());
} else {
if (typeCategory == TypeCategory.ARRAY && elementTypeCategory == TypeCategory.PRIMITIVE) {
Object value = propertiesMap.get(property.key());
if (value instanceof List) {
((List) value).add(property.value());
} else {
List<Object> values = new ArrayList<>();
values.add(value);
values.add(property.value());
propertiesMap.put(property.key(), values);
}
} else {
propertiesMap.put(property.key(), property.value());
}
}
}

return propertiesMap;
}

Expand Down Expand Up @@ -1132,6 +1149,7 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeaderWithoutPrefetch(AtlasVerte
}
}


Object attrValue = getVertexAttribute(entityVertex, attribute);

if (attrValue != null) {
Expand All @@ -1152,9 +1170,10 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeaderWithPrefetch(AtlasVertex e
AtlasEntityHeader ret = new AtlasEntityHeader();
try {
//pre-fetching the properties
Map<String, Object> properties = preloadProperties(entityVertex);

String typeName = entityVertex.getProperty(Constants.TYPE_NAME_PROPERTY_KEY, String.class); //properties.get returns null
AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName); // this is not costly
Map<String, Object> properties = preloadProperties(entityVertex, typeRegistry.getEntityTypeByName(GraphHelper.getTypeName(entityVertex)), attributes);

String guid = (String) properties.get(Constants.GUID_PROPERTY_KEY);

Integer value = (Integer)properties.get(Constants.IS_INCOMPLETE_PROPERTY_KEY);
Expand Down Expand Up @@ -1190,7 +1209,6 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeaderWithPrefetch(AtlasVertex e
termAssignmentHeaders.stream().map(AtlasTermAssignmentHeader::getDisplayText)
.collect(Collectors.toList()));
}
AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName); // this is not costly

if (entityType != null) {
for (AtlasAttribute headerAttribute : entityType.getHeaderAttributes().values()) {
Expand Down Expand Up @@ -1864,14 +1882,27 @@ public Object getVertexAttributePreFetchCache(AtlasVertex vertex, AtlasAttribute
if (vertex == null || attribute == null) {
return null;
}
LOG.info("capturing property its category and value - {}: {} : {}", attribute.getName(), attribute.getAttributeType().getTypeCategory(), properties.get(attribute.getName()));

if (properties.get(attribute.getName()) != null) {
if ((properties.get(attribute.getName()) != null) &&
(attribute.getAttributeType().getTypeCategory().equals(TypeCategory.PRIMITIVE) ||
attribute.getAttributeType().getTypeCategory().equals(TypeCategory.ENUM) ||
attribute.getAttributeType().getTypeCategory().equals(TypeCategory.ARRAY)
)) {
return properties.get(attribute.getName());
}

if (AtlasConfiguration.ATLAS_INDEXSEARCH_ENABLE_FETCHING_NON_PRIMITIVE_ATTRIBUTES.getBoolean() && !attribute.getAttributeType().getTypeCategory().equals(TypeCategory.PRIMITIVE)) {
TypeCategory typeCategory = attribute.getAttributeType().getTypeCategory();
TypeCategory elementTypeCategory = typeCategory == TypeCategory.ARRAY ? ((AtlasArrayType) attribute.getAttributeType()).getElementType().getTypeCategory() : null;

if (elementTypeCategory == TypeCategory.PRIMITIVE) {
return new ArrayList<>();
}

Set<TypeCategory> ENRICH_PROPERTY_TYPES = new HashSet<>(Arrays.asList(TypeCategory.STRUCT, TypeCategory.OBJECT_ID_TYPE, TypeCategory.MAP));
if (ENRICH_PROPERTY_TYPES.contains(attribute.getAttributeType().getTypeCategory())) {
LOG.info("capturing excluded property set category and value - {}: {} : {}", attribute.getName(), attribute.getAttributeType().getTypeCategory(), properties.get(attribute.getName()));
AtlasPerfMetrics.MetricRecorder nonPrimitiveAttributes = RequestContext.get().startMetricRecord("processNonPrimitiveAttributes");
LOG.info("Fetching non primitive attributes for entity: {}", attribute.getName());
Object mappedVertex = mapVertexToAttribute(vertex, attribute, null, false);
properties.put(attribute.getName(), mappedVertex);
RequestContext.get().endMetricRecord(nonPrimitiveAttributes);
Expand Down

0 comments on commit f719801

Please sign in to comment.