[Streaming] Fault Tolerance Implementation (#10008)

streaming/python/function.py

@@ -22,6 +22,12 @@ class Function(ABC):
     def close(self):
         pass
 
+    def save_checkpoint(self):
+        pass
+
+    def load_checkpoint(self, checkpoint_obj):
+        pass
+
 
 class EmptyFunction(Function):
     """Default function which does nothing"""
@@ -58,12 +64,15 @@ class SourceFunction(Function):
         pass
 
     @abstractmethod
-    def run(self, ctx: SourceContext):
+    def fetch(self, ctx: SourceContext):
         """Starts the source. Implementations can use the
          :class:`SourceContext` to emit elements.
         """
         pass
 
+    def close(self):
+        pass
+
 
 class MapFunction(Function):
     """
@@ -176,24 +185,29 @@ class CollectionSourceFunction(SourceFunction):
     def init(self, parallel, index):
         pass
 
-    def run(self, ctx: SourceContext):
+    def fetch(self, ctx: SourceContext):
         for v in self.values:
             ctx.collect(v)
+        self.values = []
 
 
 class LocalFileSourceFunction(SourceFunction):
     def __init__(self, filename):
         self.filename = filename
+        self.done = False
 
     def init(self, parallel, index):
         pass
 
-    def run(self, ctx: SourceContext):
+    def fetch(self, ctx: SourceContext):
+        if self.done:
+            return
         with open(self.filename, "r") as f:
             line = f.readline()
             while line != "":
                 ctx.collect(line[:-1])
                 line = f.readline()
+            self.done = True
 
 
 class SimpleMapFunction(MapFunction):