ATLAS-1379 Avoid object query overhead when report query selects class type…

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-1379 Avoid object query overhead when report query selects class type alias (guptaneeru via dkantor)
 ATLAS-1419 update entity-update API impl to preserve value of entity attribute when no value is provided
 ATLAS-1346 Search API to return empty list/container object instead of exception (apoorvnaik via mneethiraj)
 ATLAS-1428 Create of entityDef type fails with type already exists exception (sarath.kum4r@gmail.com via mneethiraj)

repository/src/main/java/org/apache/atlas/discovery/DataSetLineageService.java

@@ -207,7 +207,7 @@ public class DataSetLineageService implements LineageService {
                                              .vertices().iterator();
         while (results.hasNext()) {
             AtlasVertex vertex = results.next();
-            return TypeUtils.Pair.of(GraphHelper.getTypeName(vertex), GraphHelper.getIdFromVertex(vertex));
+            return TypeUtils.Pair.of(GraphHelper.getTypeName(vertex), GraphHelper.getGuid(vertex));
         }
         throw new EntityNotFoundException("Dataset with name = " + datasetName + " does not exist");
     }

repository/src/main/java/org/apache/atlas/discovery/graph/DefaultGraphPersistenceStrategy.java

@@ -36,10 +36,12 @@ import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.graphdb.GremlinVersion;
+import org.apache.atlas.typesystem.IReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedStruct;
 import org.apache.atlas.typesystem.persistence.Id;
 import org.apache.atlas.typesystem.types.AttributeInfo;
+import org.apache.atlas.typesystem.types.ClassType;
 import org.apache.atlas.typesystem.types.DataTypes;
 import org.apache.atlas.typesystem.types.IDataType;
 import org.apache.atlas.typesystem.types.Multiplicity;
@@ -110,6 +112,19 @@ public class DefaultGraphPersistenceStrategy implements GraphPersistenceStrategi
     }
 
     @Override
+    public ITypedReferenceableInstance constructClassInstanceId(ClassType classType, Object value) {
+        try {
+            AtlasVertex classVertex = (AtlasVertex) value;
+            ITypedReferenceableInstance classInstance = classType.createInstance(GraphHelper.getIdFromVertex(classVertex),
+                    new String[0]);
+            return classType.convert(classInstance, Multiplicity.OPTIONAL);
+        } catch (AtlasException e) {
+            LOG.error("error while constructing an instance", e);
+        }
+        return null;
+    }
+
+    @Override
     public <U> U constructInstance(IDataType<U> dataType, Object value) {
         try {
             switch (dataType.getTypeCategory()) {
@@ -136,7 +151,6 @@ public class DefaultGraphPersistenceStrategy implements GraphPersistenceStrategi
                 AtlasVertex structVertex = (AtlasVertex) value;
                 StructType structType = (StructType) dataType;
                 ITypedStruct structInstance = structType.createInstance();
-
                 TypeSystem.IdType idType = TypeSystem.getInstance().getIdType();
 
                 if (dataType.getName().equals(idType.getName())) {
@@ -146,6 +160,7 @@ public class DefaultGraphPersistenceStrategy implements GraphPersistenceStrategi
                     if (stateValue != null) {
                         structInstance.set(idType.stateAttrName(), stateValue);
                     }
+                    structInstance.set(idType.versionAttrName(), structVertex.getProperty(versionAttributeName(), Integer.class));
                 } else {
                     metadataRepository.getGraphToInstanceMapper()
                         .mapVertexToInstance(structVertex, structInstance, structType.fieldMapping().fields);
@@ -227,6 +242,11 @@ public class DefaultGraphPersistenceStrategy implements GraphPersistenceStrategi
     }
 
     @Override
+    public String versionAttributeName() {
+        return metadataRepository.getVersionAttributeName();
+    }
+
+    @Override
     public boolean collectTypeInstancesIntoVar() {
         return GraphPersistenceStrategies$class.collectTypeInstancesIntoVar(this);
     }

repository/src/main/java/org/apache/atlas/discovery/graph/GraphBackedDiscoveryService.java

@@ -95,12 +95,12 @@ public class GraphBackedDiscoveryService implements DiscoveryService {
         }
 
         while (results.hasNext() && response.length() < queryParams.limit()) {
-            
+
             AtlasIndexQuery.Result<?,?> result = results.next();
             AtlasVertex<?,?> vertex = result.getVertex();
 
             JSONObject row = new JSONObject();
-            String guid = GraphHelper.getIdFromVertex(vertex);
+            String guid = GraphHelper.getGuid(vertex);
             if (guid != null) { //Filter non-class entities
                 try {
                     row.put("guid", guid);

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

@@ -54,6 +54,11 @@ public interface MetadataRepository {
      * @return
      */
     String getStateAttributeName();
+    /**
+     * Returns the attribute name used for entity version
+     * @return
+     */
+    String getVersionAttributeName();
 
     /**
      * Return the property key used to store a given traitName in the repository.

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

@@ -76,7 +76,7 @@ public abstract class DeleteHandler {
        Set<AtlasVertex> deletionCandidateVertices = new HashSet<>();
 
        for (AtlasVertex instanceVertex : instanceVertices) {
-            String guid = GraphHelper.getIdFromVertex(instanceVertex);
+            String guid = GraphHelper.getGuid(instanceVertex);
             Id.EntityState state = GraphHelper.getState(instanceVertex);
             if (requestContext.getDeletedEntityIds().contains(guid) || state == Id.EntityState.DELETED) {
                    LOG.debug("Skipping deletion of {} as it is already deleted", guid);
@@ -274,7 +274,7 @@ public abstract class DeleteHandler {
         LOG.debug("Removing edge from {} to {} with attribute name {}", string(outVertex), string(inVertex),
                 attributeName);
         String typeName = GraphHelper.getTypeName(outVertex);
-        String outId = GraphHelper.getIdFromVertex(outVertex);
+        String outId = GraphHelper.getGuid(outVertex);
         Id.EntityState state = GraphHelper.getState(outVertex);
         if ((outId != null && RequestContext.get().isDeletedEntity(outId)) || state == Id.EntityState.DELETED) {
             //If the reference vertex is marked for deletion, skip updating the reference

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

@@ -52,7 +52,7 @@ public class FullTextMapper {
 
     @Monitored
     public String mapRecursive(AtlasVertex instanceVertex, boolean followReferences) throws AtlasException {
-        String guid = GraphHelper.getIdFromVertex(instanceVertex);
+        String guid = GraphHelper.getGuid(instanceVertex);
         ITypedReferenceableInstance typedReference;
         if (instanceCache.containsKey(guid)) {
             typedReference = instanceCache.get(guid);

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

@@ -105,6 +105,10 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
     public String getIdAttributeName() {
         return Constants.GUID_PROPERTY_KEY;
     }
+    @Override
+    public String getVersionAttributeName() {
+        return Constants.VERSION_PROPERTY_KEY;
+    }
 
     @Override
     public String getTraitLabel(IDataType<?> dataType, String traitName) {
@@ -169,7 +173,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
                 Constants.ENTITY_TYPE_PROPERTY_KEY, entityType,
                 Constants.STATE_PROPERTY_KEY, Id.EntityState.ACTIVE.name());
 
-        String guid = GraphHelper.getIdFromVertex(instanceVertex);
+        String guid = GraphHelper.getGuid(instanceVertex);
         return graphToInstanceMapper.mapGraphToTypedInstance(guid, instanceVertex);
     }
 
@@ -186,7 +190,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
         ArrayList<String> entityList = new ArrayList<>();
         while (results.hasNext()) {
             AtlasVertex vertex = results.next();
-            entityList.add(GraphHelper.getIdFromVertex(vertex));
+            entityList.add(GraphHelper.getGuid(vertex));
         }
 
         return entityList;

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

@@ -505,12 +505,17 @@ public final class GraphHelper {
     }
 
     public static Id getIdFromVertex(String dataTypeName, AtlasVertex vertex) {
-        return new Id(getIdFromVertex(vertex),
-            vertex.getProperty(Constants.VERSION_PROPERTY_KEY, Integer.class), dataTypeName, getStateAsString(vertex));
+        return new Id(getGuid(vertex),
+                getVersion(vertex), dataTypeName,
+                getStateAsString(vertex));
     }
 
-    public static String getIdFromVertex(AtlasVertex vertex) {
-        return vertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class);
+    public static Id getIdFromVertex(AtlasVertex vertex) {
+        return getIdFromVertex(getTypeName(vertex), vertex);
+    }
+
+    public static String getGuid(AtlasVertex vertex) {
+        return vertex.<String>getProperty(Constants.GUID_PROPERTY_KEY, String.class);
     }
 
     public static String getTypeName(AtlasVertex instanceVertex) {
@@ -522,6 +527,10 @@ public final class GraphHelper {
         return state == null ? null : Id.EntityState.valueOf(state);
     }
 
+    public static Integer getVersion(AtlasElement element) {
+        return element.getProperty(Constants.VERSION_PROPERTY_KEY, Integer.class);
+    }
+
     public static String getStateAsString(AtlasElement element) {
         return element.getProperty(Constants.STATE_PROPERTY_KEY, String.class);
     }
@@ -629,7 +638,7 @@ public final class GraphHelper {
         while (vertices.size() > 0) {
             AtlasVertex vertex = vertices.pop();
             String typeName = GraphHelper.getTypeName(vertex);
-            String guid = GraphHelper.getIdFromVertex(vertex);
+            String guid = GraphHelper.getGuid(vertex);
             Id.EntityState state = GraphHelper.getState(vertex);
             if (state == Id.EntityState.DELETED) {
                 //If the reference vertex is marked for deletion, skip it
@@ -814,7 +823,7 @@ public final class GraphHelper {
     public static String getVertexDetails(AtlasVertex<?,?> vertex) {
 
         return String.format("vertex[id=%s type=%s guid=%s]", vertex.getIdForDisplay(), getTypeName(vertex),
-                getIdFromVertex(vertex));
+                getGuid(vertex));
     }
 
 

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

@@ -627,7 +627,7 @@ public final class TypedInstanceToGraphMapper {
 
         if (id.isUnassigned()) {
             AtlasVertex classVertex = idToVertexMap.get(id);
-            String guid = GraphHelper.getIdFromVertex(classVertex);
+            String guid = GraphHelper.getGuid(classVertex);
             id = new Id(guid, 0, typedReference.getTypeName());
         }
         return id;
@@ -641,7 +641,7 @@ public final class TypedInstanceToGraphMapper {
         LOG.debug("Updating {} for reference attribute {}", string(currentEdge), attributeInfo.name);
         // Update edge if it exists
         AtlasVertex currentVertex = currentEdge.getInVertex();
-        String currentEntityId = GraphHelper.getIdFromVertex(currentVertex);
+        String currentEntityId = GraphHelper.getGuid(currentVertex);
         String newEntityId = getId(newAttributeValue).id;
         AtlasEdge newEdge = currentEdge;
         if (!currentEntityId.equals(newEntityId)) {

repository/src/main/scala/org/apache/atlas/query/GraphPersistenceStrategies.scala

@@ -76,6 +76,10 @@ trait GraphPersistenceStrategies {
       * Name of attribute used to store state in vertex
       */
     def stateAttributeName : String
+    /**
+     * Name of attribute used to store version in vertex
+     */
+    def versionAttributeName : String
 
     /**
      * Given a dataType and a reference attribute, how is edge labeled
@@ -123,6 +127,8 @@ trait GraphPersistenceStrategies {
 
     def constructInstance[U](dataType: IDataType[U], v: java.lang.Object): U
 
+    def constructClassInstanceId[U](dataType: ClassType, v: java.lang.Object): ITypedReferenceableInstance
+
     def addGraphVertexPrefix(preStatements : Traversable[GroovyExpression]) = !collectTypeInstancesIntoVar
 
     /**
@@ -162,6 +168,7 @@ case class GraphPersistenceStrategy1(g: AtlasGraph[_,_]) extends GraphPersistenc
     val superTypeAttributeName = "superTypeNames"
     val idAttributeName = "guid"
     val stateAttributeName = "state"
+    val versionAttributeName = "version"
 
     override def getGraph() : AtlasGraph[_,_] =  {
         return g;
@@ -178,6 +185,9 @@ case class GraphPersistenceStrategy1(g: AtlasGraph[_,_]) extends GraphPersistenc
     def getIdFromVertex(dataTypeNm: String, v: AtlasVertex[_,_]): Id =
         new Id(v.getId.toString, 0, dataTypeNm)
 
+    def getIdFromVertex(v: AtlasVertex[_,_]): Id =
+        getIdFromVertex(v.getProperty(typeAttributeName, classOf[java.lang.String]), v)
+
     def traitNames(v: AtlasVertex[_,_]): java.util.List[String] = {
         val s = v.getProperty("traitNames", classOf[String])
         if (s != null) {
@@ -186,7 +196,12 @@ case class GraphPersistenceStrategy1(g: AtlasGraph[_,_]) extends GraphPersistenc
             Seq()
         }
     }
-
+    def constructClassInstanceId[U](classType: ClassType, v: AnyRef): ITypedReferenceableInstance = {
+        val vertex = v.asInstanceOf[AtlasVertex[_,_]];
+        val id = getIdFromVertex(vertex)
+        val cInstance = classType.createInstance(id)
+        classType.convert(cInstance, Multiplicity.OPTIONAL)
+    }
     def constructInstance[U](dataType: IDataType[U], v: AnyRef): U = {
         dataType.getTypeCategory match {
             case DataTypes.TypeCategory.PRIMITIVE => dataType.convert(v, Multiplicity.OPTIONAL)

repository/src/main/scala/org/apache/atlas/query/GremlinEvaluator.scala

@@ -28,6 +28,7 @@ import org.json4s._
 import org.json4s.native.Serialization._
 import scala.language.existentials
 import org.apache.atlas.query.Expressions._
+import org.apache.atlas.typesystem.types.DataTypes.TypeCategory
 
 case class GremlinQueryResult(query: String,
                               resultDataType: IDataType[_],
@@ -109,7 +110,12 @@ class GremlinEvaluator(qry: GremlinQuery, persistenceStrategy: GraphPersistenceS
                         val cName = aE.alias
                         val (src, idx) = qry.resultMaping(cName)
                         val v = getColumnValue(rV, src, idx)
-                        sInstance.set(cName, persistenceStrategy.constructInstance(aE.dataType, v))
+                        //if select clause is selecting the entire object then return only the instance id (guid, version, state and typeName)
+                        if (aE.dataType.getTypeCategory == TypeCategory.CLASS) {
+                             sInstance.set(cName, persistenceStrategy.constructClassInstanceId(aE.dataType.asInstanceOf[ClassType], v))
+                        } else {
+                            sInstance.set(cName, persistenceStrategy.constructInstance(aE.dataType, v))
+                        }
                     }
                 }
                 else if(qry.isGroupBy) {

repository/src/test/java/org/apache/atlas/discovery/GraphBackedDiscoveryServiceTest.java

@@ -81,7 +81,7 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
     @Inject
     private GraphBackedDiscoveryService discoveryService;
     private QueryParams queryParams = new QueryParams(40, 0);
-
+    private static final String idType = "idType";
     @Override
     @BeforeClass
     public void setUp() throws Exception {
@@ -387,6 +387,7 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
     @DataProvider(name = "dslQueriesProvider")
     private Object[][] createDSLQueries() {
         return new Object[][]{
+                {"from hive_db as h select h as id", 3},
                 {"from hive_db", 3},
                 {"hive_db", 3},
                 {"hive_db where hive_db.name=\"Reporting\"", 1},
@@ -863,6 +864,16 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
             };
     }
 
+    @DataProvider(name = "dslObjectQueriesReturnIdProvider")
+    private Object[][] createDSLObjectIdQueries() {
+        return new Object[][] { {
+                "from hive_db as h select h as id",
+                new FieldValueValidator().withFieldNames("id")
+                        .withExpectedValues(idType).withExpectedValues(idType)
+                        .withExpectedValues(idType) }
+        };
+    }
+
     @Test(dataProvider = "dslOrderByQueriesProvider")
     public void  testSearchByDSLQueriesWithOrderBy(String dslQuery, Integer expectedNumRows, String orderBy, boolean ascending) throws Exception {
         System.out.println("Executing dslQuery = " + dslQuery);
@@ -1023,6 +1034,10 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
     }
 
     private void runCountGroupByQuery(String dslQuery, ResultChecker checker) throws Exception {
+        runAndValidateQuery(dslQuery, checker);
+    }
+
+    private void runAndValidateQuery(String dslQuery, ResultChecker checker) throws Exception {
         System.out.println("Executing dslQuery = " + dslQuery);
         String jsonResults = searchByDSL(dslQuery);
         assertNotNull(jsonResults);
@@ -1046,6 +1061,12 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
         runCountGroupByQuery(dslQuery, checker);
     }
 
+    @Test(dataProvider = "dslObjectQueriesReturnIdProvider")
+    public void testSearchObjectQueriesReturnId(String dslQuery,
+            ResultChecker checker) throws Exception {
+        runAndValidateQuery(dslQuery, checker);
+    }
+
     private interface ResultChecker {
         void validateResult(String dslQuery, JSONArray foundRows) throws JSONException;
     }
@@ -1053,6 +1074,9 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
     static class FieldValueValidator implements ResultChecker {
         static class ResultObject {
 
+            private static String[] idTypeAttributes = { "id", "$typeName$",
+                "state", "version" };
+
             @Override
             public String toString() {
                 return "ResultObject [fieldValues_=" + fieldValues_ + "]";
@@ -1072,6 +1096,8 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
                     Object foundValue = null;
                     if (expectedValue.getClass() == Integer.class) {
                         foundValue = object.getInt(fieldName);
+                    } else if (expectedValue == idType) {
+                        return validateObjectIdType(object, fieldName);
                     } else {
                         foundValue = object.get(fieldName);
                     }
@@ -1081,6 +1107,17 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
                 }
                 return true;
             }
+            // validates that returned object id contains all the required attributes.
+            private boolean validateObjectIdType(JSONObject object,
+                    String fieldName) throws JSONException {
+                JSONObject foundJson = object.getJSONObject(fieldName);
+                for (String idAttr : idTypeAttributes) {
+                    if (foundJson.get(idAttr) == null) {
+                        return false;
+                    }
+                }
+                return true;
+            }
         }
 
         private String[] fieldNames_;

typesystem/src/main/java/org/apache/atlas/typesystem/types/TypeSystem.java

@@ -736,6 +736,7 @@ public class TypeSystem {
         private static final String ID_ATTRNAME = "guid";
         private static final String TYPENAME_ATTRNAME = "typeName";
         private static final String STATE_ATTRNAME = "state";
+        private static final String VERSION_ATTRNAME = "version";
         private static final String TYP_NAME = "__IdType";
 
         private StructType type;
@@ -750,11 +751,15 @@ public class TypeSystem {
             AttributeDefinition stateAttr =
                     new AttributeDefinition(STATE_ATTRNAME, DataTypes.STRING_TYPE.getName(), Multiplicity.REQUIRED,
                             false, null);
+            AttributeDefinition versionAttr =
+                    new AttributeDefinition(VERSION_ATTRNAME, DataTypes.INT_TYPE.getName(), Multiplicity.REQUIRED,
+                            false, null);
             try {
-                AttributeInfo[] infos = new AttributeInfo[3];
+                AttributeInfo[] infos = new AttributeInfo[4];
                 infos[0] = new AttributeInfo(TypeSystem.this, idAttr, null);
                 infos[1] = new AttributeInfo(TypeSystem.this, typNmAttr, null);
                 infos[2] = new AttributeInfo(TypeSystem.this, stateAttr, null);
+                infos[3] = new AttributeInfo(TypeSystem.this, versionAttr, null);
 
                 type = new StructType(TypeSystem.this, TYP_NAME, null, infos);
             } catch (AtlasException me) {
@@ -781,6 +786,10 @@ public class TypeSystem {
         public String stateAttrName() {
             return STATE_ATTRNAME;
         }
+
+        public String versionAttrName() {
+            return VERSION_ATTRNAME;
+        }
     }
 
     public static final String ID_STRUCT_ID_ATTRNAME = IdType.ID_ATTRNAME;