Fix default session for processing vineyard operation. (alibaba#230)

Author: siyuan0322

python/graphscope/framework/graph.py

@@ -734,7 +734,11 @@ def add_vertices(self, vertices, label="_", properties=None, vid_field=0):
             raise ValueError("Cannot manually add vertices after inferred vertices.")
         unsealed_vertices = deepcopy(self._unsealed_vertices)
         unsealed_vertices[label] = VertexLabel(
-            label=label, loader=vertices, properties=properties, vid_field=vid_field
+            label=label,
+            loader=vertices,
+            properties=properties,
+            vid_field=vid_field,
+            session_id=self._session.session_id,
         )
         v_labels = deepcopy(self._v_labels)
         v_labels.append(label)
@@ -845,7 +849,7 @@ def add_edges(
         else:
             e_labels.append(label)
             relations.append([(src_label, dst_label)])
-            cur_label = EdgeLabel(label)
+            cur_label = EdgeLabel(label, self._session.session_id)
         cur_label.add_sub_label(
             EdgeSubLabel(edges, properties, src_label, dst_label, src_field, dst_field)
         )

python/graphscope/framework/graph_utils.py

@@ -58,6 +58,7 @@ def __init__(
         loader: Any,
         properties: Sequence = None,
         vid_field: Union[str, int] = 0,
+        session_id=None,
     ):
         self.label = label
         if isinstance(loader, Loader):
@@ -68,7 +69,7 @@ def __init__(
         self.raw_properties = properties
         self.properties = []
         self.vid_field = vid_field
-
+        self._session_id = session_id
         self._finished = False
 
     def finish(self, id_type: str = "int64_t"):
@@ -84,7 +85,7 @@ def finish(self, id_type: str = "int64_t"):
         self.loader.select_columns(
             self.properties, include_all=bool(self.raw_properties is None)
         )
-        self.loader.finish()
+        self.loader.finish(self._session_id)
         self._finished = True
 
     def __str__(self) -> str:
@@ -153,7 +154,7 @@ def __init__(
                 "Source vid and destination vid must have same formats, both use name or both use index"
             )
 
-    def finish(self, id_type: str = "int64_t"):
+    def finish(self, id_type: str = "int64_t", session_id=None):
         if self._finished:
             return
         self.add_property(str(self.src_field), id_type)
@@ -165,7 +166,7 @@ def finish(self, id_type: str = "int64_t"):
         self.loader.select_columns(
             self.properties, include_all=bool(self.raw_properties is None)
         )
-        self.loader.finish()
+        self.loader.finish(session_id)
         self._finished = True
 
     def __str__(self) -> str:
@@ -221,9 +222,12 @@ class EdgeLabel(object):
          src_label3 -> edge_label -> dst_label3
     """
 
-    def __init__(self, label: str):
+    def __init__(self, label: str, session_id=None):
         self.label = label
         self.sub_labels = {}
+
+        self._session_id = session_id
+
         self._finished = False
 
     def __str__(self):
@@ -251,7 +255,7 @@ def finish(self, id_type: str = "int64_t"):
         if self._finished:
             return
         for sub_label in self.sub_labels.values():
-            sub_label.finish(id_type)
+            sub_label.finish(id_type, self._session_id)
         self._finished = True
 
 

python/graphscope/framework/loader.py

@@ -278,17 +278,24 @@ def func(source, storage_options, read_options, sess):
         self.source = source
         self.preprocessor = func
 
-    def finish(self):
-        from graphscope.client.session import get_default_session
-
+    def finish(self, session_id=None):
         if self.finished:
             return
         if self.preprocessor is not None:
+            if session_id is None:
+                from graphscope.client.session import get_default_session
+
+                sess = get_default_session()
+            else:
+                from graphscope.client.session import get_session_by_id
+
+                sess = get_session_by_id(session_id)
+
             self.protocol, self.source = self.preprocessor(
                 self.source,
                 self.storage_options,
                 self.options.to_dict(),
-                get_default_session(),
+                sess,
             )
             logger.debug(
                 f"processed protocol = {self.protocol}, source = {self.source}"