[BEAM-8823] Make FnApiRunner work by executing ready elements instead of stages

sdks/python/apache_beam/io/iobase.py

@@ -1148,11 +1148,13 @@ def expand(self, pcoll):
           | 'Extract' >> core.FlatMap(lambda x: x[1]))
     # PreFinalize should run before FinalizeWrite, and the two should not be
     # fused.
-    pre_finalize_coll = do_once | 'PreFinalize' >> core.FlatMap(
-        _pre_finalize,
-        self.sink,
-        AsSingleton(init_result_coll),
-        AsIter(write_result_coll))
+    pre_finalize_coll = (
+        do_once
+        | 'PreFinalize' >> core.FlatMap(
+            _pre_finalize,
+            self.sink,
+            AsSingleton(init_result_coll),
+            AsIter(write_result_coll)))
     return do_once | 'FinalizeWrite' >> core.FlatMap(
         _finalize_write,
         self.sink,

sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py

@@ -294,6 +294,11 @@ def cross_product(elem, sides):
           pcoll | beam.FlatMap(cross_product, beam.pvalue.AsList(pcoll)),
           equal_to([('a', 'a'), ('a', 'b'), ('b', 'a'), ('b', 'b')]))
 
+  def test_pardo_unfusable_side_inputs_with_separation(self):
+    def cross_product(elem, sides):
+      for side in sides:
+        yield elem, side
+
     with self.create_pipeline() as p:
       pcoll = p | beam.Create(['a', 'b'])
       derived = ((pcoll, ) | beam.Flatten()

sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py

@@ -25,6 +25,9 @@
 import itertools
 import logging
 import operator
+from builtins import object
+from typing import TYPE_CHECKING
+from typing import Any
 from typing import Callable
 from typing import Collection
 from typing import Container
@@ -34,6 +37,8 @@
 from typing import Iterable
 from typing import Iterator
 from typing import List
+from typing import MutableMapping
+from typing import NamedTuple
 from typing import Optional
 from typing import Set
 from typing import Tuple
@@ -44,11 +49,17 @@
 from apache_beam.internal import pickler
 from apache_beam.portability import common_urns
 from apache_beam.portability import python_urns
+from apache_beam.portability.api import beam_fn_api_pb2
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.runners.worker import bundle_processor
 from apache_beam.transforms import combiners
 from apache_beam.transforms import core
 from apache_beam.utils import proto_utils
+from apache_beam.utils import timestamp
+
+if TYPE_CHECKING:
+  from apache_beam.runners.portability.fn_api_runner.execution import ListBuffer
+  from apache_beam.runners.portability.fn_api_runner.execution import PartitionableBuffer
 
 T = TypeVar('T')
 
@@ -77,14 +88,15 @@
 IMPULSE_BUFFER = b'impulse'
 
 # TimerFamilyId is identified by transform name + timer family
+# TODO(pabloem): Rename this type to express this name is unique per pipeline.
 TimerFamilyId = Tuple[str, str]
 
+BufferId = bytes
+
 # SideInputId is identified by a consumer ParDo + tag.
 SideInputId = Tuple[str, str]
 SideInputAccessPattern = beam_runner_api_pb2.FunctionSpec
 
-DataOutput = Dict[str, bytes]
-
 # A map from a PCollection coder ID to a Safe Coder ID
 # A safe coder is a coder that can be used on the runner-side of the FnApi.
 # A safe coder receives a byte string, and returns a type that can be
@@ -96,6 +108,27 @@
 # (MultiMap / Iterable).
 DataSideInput = Dict[SideInputId, Tuple[bytes, SideInputAccessPattern]]
 
+DataOutput = Dict[str, BufferId]
+
+# A map of [Transform ID, Timer Family ID] to [Buffer ID, Time Domain for timer]
+# The time domain comes from beam_runner_api_pb2.TimeDomain. It may be
+# EVENT_TIME or PROCESSING_TIME.
+OutputTimers = MutableMapping[TimerFamilyId, Tuple[BufferId, Any]]
+
+# A map of [Transform ID, Timer Family ID] to [Buffer CONTENTS, Timestamp]
+OutputTimerData = MutableMapping[TimerFamilyId,
+                                 Tuple['PartitionableBuffer',
+                                       timestamp.Timestamp]]
+
+BundleProcessResult = Tuple[beam_fn_api_pb2.InstructionResponse,
+                            List[beam_fn_api_pb2.ProcessBundleSplitResponse]]
+
+
+# TODO(pabloem): Change tha name to a more representative one
+class DataInput(NamedTuple):
+  data: MutableMapping[str, 'PartitionableBuffer']
+  timers: MutableMapping[TimerFamilyId, 'PartitionableBuffer']
+
 
 class Stage(object):
   """A set of Transforms that can be sent to the worker for processing."""

sdks/python/apache_beam/runners/portability/fn_api_runner/watermark_manager.py

@@ -50,7 +50,7 @@ def set_watermark(self, wm: timestamp.Timestamp):
 
     def upstream_watermark(self):
       if self.producers:
-        return min(p.input_watermark() for p in self.producers)
+        return min(p.output_watermark() for p in self.producers)
       else:
         return timestamp.MAX_TIMESTAMP
 
@@ -71,13 +71,16 @@ def __init__(self, name):
 
     def __str__(self):
       return 'StageNode<inputs=%s,side_inputs=%s' % (
-          [i.name for i in self.inputs], [i.name for i in self.side_inputs])
+          [
+              '%s(%s, upstream:%s)' %
+              (i.name, i.watermark(), i.upstream_watermark())
+              for i in self.inputs
+          ], ['%s(%s)' % (i.name, i.watermark()) for i in self.side_inputs])
 
     def output_watermark(self):
-      if not self.outputs:
-        return self.input_watermark()
-      else:
-        return min(o.watermark() for o in self.outputs)
+      if not self.inputs:
+        return timestamp.MAX_TIMESTAMP
+      return min(i.watermark() for i in self.inputs)
 
     def input_watermark(self):
       if not self.inputs:

sdks/python/apache_beam/runners/portability/fn_api_runner/worker_handlers.py

@@ -964,6 +964,9 @@ def append(self, item):
         self._overlay[self._key] = list(self._underlying[self._key])
       self._overlay[self._key].append(item)
 
+    def extend(self, other: Buffer) -> None:
+      raise NotImplementedError()
+
   StateType = Union[CopyOnWriteState, DefaultDict[bytes, Buffer]]
 
   def __init__(self):

sdks/python/apache_beam/transforms/util.py

@@ -44,6 +44,7 @@
 from apache_beam.transforms import window
 from apache_beam.transforms.combiners import CountCombineFn
 from apache_beam.transforms.core import CombinePerKey
+from apache_beam.transforms.core import Create
 from apache_beam.transforms.core import DoFn
 from apache_beam.transforms.core import FlatMap
 from apache_beam.transforms.core import Flatten
@@ -159,6 +160,8 @@ def _extract_input_pvalues(self, pvalueish):
       return pcolls, pcolls
 
   def expand(self, pcolls):
+    if not pcolls:
+      pcolls = (self.pipeline | Create([]), )
     if isinstance(pcolls, dict):
       if all(isinstance(tag, str) and len(tag) < 10 for tag in pcolls.keys()):
         # Small, string tags. Pass them as data.