Update

apache · Feb 28, 2023 · 1011eeb · 1011eeb
1 parent 9f6eaa0
commit 1011eeb
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/spark/src/main/scala/com/alibaba/graphar/GraphInfo.scala b/spark/src/main/scala/com/alibaba/graphar/GraphInfo.scala
@@ -203,6 +203,21 @@ class GraphInfo() {
   @BeanProperty var vertices = new java.util.ArrayList[String]()
   @BeanProperty var edges = new java.util.ArrayList[String]()
   @BeanProperty var version: String = ""
+
+  var vertexInfos: Map[String, VertexInfo] = Map[String, VertexInfo]()
+  var edgeInfos: Map[String, EdgeInfo] = Map[String, EdgeInfo]()
+
+  def addVertexInfo(label: String, vertexInfo: VertexInfo): Unit = {
+    vertexInfos += (label -> vertexInfo)
+  }
+
+  def addEdgeInfo(label: String, edgeInfo: EdgeInfo): Unit = {
+    edgeInfos += (label -> edgeInfo)
+  }
+
+  def getVertexInfo(label: String): VertexInfo = {
+    vertexInfos(label)
+  }
 }
 
 /** Helper object to load graph info files */
@@ -221,6 +236,26 @@ object GraphInfo {
         graph_info.setPrefix(prefix)
       }
     }
+    val prefix = graph_info.getPrefix
+    val vertices_yaml = graph_info.getVertices
+    val vertices_it = vertices_yaml.iterator
+    while (vertices_it.hasNext()) {
+      val file_name = vertices_it.next()
+      val path = prefix + file_name
+      println("path" + path)
+      val vertex_info = VertexInfo.loadVertexInfo(path, spark)
+      println(vertex_info.getLabel)
+      graph_info.addVertexInfo(vertex_info.getLabel, vertex_info)
+    }
+    val edges_yaml = graph_info.getEdges
+    val edges_it = edges_yaml.iterator
+    while (edges_it.hasNext()) {
+      val file_name = edges_it.next()
+      val path = prefix + file_name
+      val edge_info = EdgeInfo.loadEdgeInfo(path, spark)
+      val key = edge_info.getSrc_label + GeneralParams.regularSeperator + edge_info.getEdge_label + GeneralParams.regularSeperator + edge_info.getDst_label
+      graph_info.addEdgeInfo(key, edge_info)
+    }
     return graph_info
   }
 }
diff --git a/spark/src/test/scala/com/alibaba/graphar/TestGraphInfo.scala b/spark/src/test/scala/com/alibaba/graphar/TestGraphInfo.scala
@@ -34,6 +34,9 @@ class GraphInfoSuite extends AnyFunSuite {
     val prefix = getClass.getClassLoader.getResource("gar-test/ldbc_sample/csv/").getPath
     val graph_info = GraphInfo.loadGraphInfo(yaml_path, spark)
 
+    val vertex_info = graph_info.getVertexInfo("person")
+    assert(vertex_info.getLabel == "person")
+
     assert(graph_info.getName == "ldbc_sample")
     assert(graph_info.getPrefix == prefix )
     assert(graph_info.getVertices.size() == 1)