[BEAM-8019] Updates Python SDK to handle remote SDK coders and preserve tags added by remote SDKs and propagate restriction coders. #11185

chamikaramj · 2020-03-20T21:55:49Z

These are needed for for runners that need to build a Python object graph from a runner API proto with external transforms (for example, Dataflow).

Testing - cross-language test suite [1] works for Dataflow with these changes (will be enabled separately).

[1] https://github.com/apache/beam/blob/master/sdks/python/apache_beam/transforms/validate_runner_xlang_test.py#L51

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Choose reviewer(s) and mention them in a comment (R: @username).
Format the pull request title like [BEAM-XXX] Fixes bug in ApproximateQuantiles, where you replace BEAM-XXX with the appropriate JIRA issue, if applicable. This will automatically link the pull request to the issue.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

Post-Commit Tests Status (on master branch)

Lang	SDK	Apex	Dataflow	Gearpump	Samza
Go		---	---	---	---
Java
Python		---		---	---
XLang	---	---	---	---	---

Pre-Commit Tests Status (on master branch)

---	Java	Python	Go	Website
Non-portable
Portable	---		---	---

See .test-infra/jenkins/README for trigger phrase, status and link of all Jenkins jobs.

chamikaramj · 2020-03-20T21:56:36Z

cc: @robertwb @lukecwik

Still addressing some unit test failures but sharing for any early comments.

robertwb · 2020-03-21T00:47:35Z

sdks/python/apache_beam/coders/coders.py

+      raise ValueError((
+          'Expected an instance of ElementTypeHolder'
+          ', but got a %s' % typehint))
+


yapf -ir path/to/files/...

robertwb · 2020-03-21T00:49:58Z

sdks/python/apache_beam/pipeline.py

-    ]
-    side_inputs = [si for _, si in sorted(indexed_side_inputs)]
+    if python_indexed_side_inputs:
+      # Ordering is important here.


This all seems rather fragile. Would it be possible to just make side_inputs a dict everywhere in the internal SDK representation? (Or is there introspection with the legacy worker code that would make this hard?)

I slightly reverted some of the code here to preserve the old behavior for Python. New changes are basically to preserve input tags for remote SDKs instead of letting Python override tags (which breaks Java).

Thanks. This is still has to work around the existing ugliness, but looks much better now.

chamikaramj

R: @robertwb

Thanks Robert. PTAL.

chamikaramj · 2020-03-23T17:21:48Z

Run XVR_Flink PostCommit

chamikaramj · 2020-03-23T17:22:49Z

Run XVR_Spark PostCommit

chamikaramj · 2020-03-23T19:53:11Z

Run Python2_PVR_Flink PreCommit

chamikaramj · 2020-03-23T21:04:07Z

All tests pass now. Thanks.

robertwb

Some initial comments. It would be helpful if you broke the separate fixes into separate commits (e.g. via git commit -p, or at least enumerated them in the description. "Some generalizations" doesn't give much context as to what the specific changes you're trying to make here are. Thanks.

robertwb · 2020-03-25T00:18:10Z

sdks/python/apache_beam/coders/coders.py


-  def proto(self):
-    return self._proto
+  coder_count = 0


I see this referenced half a dozen times below. Perhaps encapsulate the increment+return in a method _unique_id or similar?

This was just added for logging/debugging. Removed.

robertwb · 2020-03-25T00:18:40Z

sdks/python/apache_beam/coders/coders.py

-class RunnerAPICoderHolder(Coder):
+class ElementTypeHolder(typehints.TypeConstraint):
+  """A dummy element type for external coders that cannot be parsed in Python"""
+  def __init__(self, coder, context):


s/coder/coder_proto/ to be clear what it is?

robertwb · 2020-03-25T00:22:05Z

sdks/python/apache_beam/coders/coders.py

+    kind_str = 'kind:external' + str(ExternalCoder.coder_count)
+    ExternalCoder.coder_count = ExternalCoder.coder_count + 1
+    component_encodings = []
+    if coder_proto.spec.urn == 'beam:coder:kv:v1':


What about correctly handling kind:stream? Perhaps others? It seems it would be preferable to instead try to only use external_coder for the unknown leafs rather than duplicate this special-dataflow logic here.

This does not seems to be needed for now (at least to get the test suite working). So I added a TODO here. Also, note that 'kind:extenal' was used to represent Java only Void coder here, so we'll probably continue to need that even if we handle 'kind:stream' here. Lemme know if you think this is inadequate for this PR.

I anticipate kind:stream will be needed to handle GBK of unknown types. Others may be needed for other cases, or in the future, and it seems risky to enumerate them here and in the dataflow runner. There may also be cases where we have to go more than one level deep. We should try to return the same thing the external SDK would have returned just to be safe, and that means wrapping only the leaves as external coders. I think that'll clean stuff up as well (e.g. no need for _coerce_to_kv_type_from_external_type).

Cleaned this up. Now we only handle unknown types here.

robertwb · 2020-03-25T00:24:48Z

sdks/python/apache_beam/coders/coders.py

+
+  def to_runner_api_parameter(self, context):
+    if self.element_type_holder.coder.component_coder_ids:
+      raise NotImplementedError


Why raise NotImplementedError here instead of returning this in place of the empty tuple below?

robertwb · 2020-03-25T00:28:30Z

sdks/python/apache_beam/pvalue.py

+    for component_coder_id in coder_proto.component_coder_ids:
+      component_coder_proto = (
+          old_context.coders.get_id_to_proto_map()[component_coder_id])
+      new_context.coders.get_id_to_proto_map()[component_coder_id] = (


new_context.put_proto?

robertwb · 2020-03-25T00:29:05Z

sdks/python/apache_beam/pvalue.py

+      self, new_context, old_context, coder_proto):
+    for component_coder_id in coder_proto.component_coder_ids:
+      component_coder_proto = (
+          old_context.coders.get_id_to_proto_map()[component_coder_id])


You could add a get_proto(id) method.

robertwb · 2020-03-25T00:30:31Z

sdks/python/apache_beam/pvalue.py

+      self._recursively_add_external_coders(
+          context, self.element_type.context, coder_proto)
+    else:
+      coder_id = context.coder_id_from_element_type(self.element_type)


Could you make this work generically rather than having an isinstance(self.element_type, ElementTypeHolder) check above?

Maybe this'd require populating the map above with actual Coders as well as their protos.

robertwb · 2020-03-25T00:33:04Z

sdks/python/apache_beam/pvalue.py

@@ -208,9 +233,23 @@ def from_runner_api(proto, context):
    # deserialization.  It will be populated soon after this call, in
    # Pipeline.from_runner_api(). This brief period is the only time that
    # PCollection.pipeline is allowed to be None.
+
+    try:
+      element_type = context.element_type_from_coder_id(proto.coder_id)


Have this method do the ExternalCoder stuff rather than catching an exception here. Possibly it would not even be needed if the coders map is properly populated beforehand.

robertwb · 2020-03-25T00:33:52Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

@@ -380,7 +381,12 @@ def visit_transform(self, transform_node):
        # pylint: disable=wrong-import-order, wrong-import-position
        from apache_beam import Flatten
        if isinstance(transform_node.transform, Flatten):
-          output_pcoll = transform_node.outputs[None]
+          from apache_beam.runners.portability.fn_api_runner.translations \


If you want to use this, let's put it in a shared utils file.

Implementation seems to be pretty straightforward. So just added a version to this file.

robertwb · 2020-03-25T00:35:03Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

+            import only_element
+          output_pcoll = (
+              transform_node.outputs[None] if None in transform_node.outputs
+              else only_element(transform_node.outputs.values()))


only_element is fine to use everywhere, no need to branch on the key being None.

chamikaramj

Thanks. PTAL.

I updated the description but let me know if you need additional details.

chamikaramj · 2020-03-25T15:29:26Z

sdks/python/apache_beam/coders/coders.py

-class RunnerAPICoderHolder(Coder):
+class ElementTypeHolder(typehints.TypeConstraint):
+  """A dummy element type for external coders that cannot be parsed in Python"""
+  def __init__(self, coder, context):


chamikaramj · 2020-03-25T15:42:36Z

sdks/python/apache_beam/coders/coders.py


-  def proto(self):
-    return self._proto
+  coder_count = 0


This was just added for logging/debugging. Removed.

chamikaramj · 2020-03-25T19:36:55Z

sdks/python/apache_beam/coders/coders.py

+
+  def to_runner_api_parameter(self, context):
+    if self.element_type_holder.coder.component_coder_ids:
+      raise NotImplementedError


chamikaramj · 2020-03-25T20:51:56Z

sdks/python/apache_beam/pvalue.py

+      self, new_context, old_context, coder_proto):
+    for component_coder_id in coder_proto.component_coder_ids:
+      component_coder_proto = (
+          old_context.coders.get_id_to_proto_map()[component_coder_id])


chamikaramj · 2020-03-25T20:52:03Z

sdks/python/apache_beam/pvalue.py

+    for component_coder_id in coder_proto.component_coder_ids:
+      component_coder_proto = (
+          old_context.coders.get_id_to_proto_map()[component_coder_id])
+      new_context.coders.get_id_to_proto_map()[component_coder_id] = (


chamikaramj · 2020-03-26T22:08:18Z

sdks/python/apache_beam/pvalue.py

@@ -208,9 +233,23 @@ def from_runner_api(proto, context):
    # deserialization.  It will be populated soon after this call, in
    # Pipeline.from_runner_api(). This brief period is the only time that
    # PCollection.pipeline is allowed to be None.
+
+    try:
+      element_type = context.element_type_from_coder_id(proto.coder_id)


chamikaramj · 2020-03-26T22:34:29Z

sdks/python/apache_beam/pvalue.py

+      self._recursively_add_external_coders(
+          context, self.element_type.context, coder_proto)
+    else:
+      coder_id = context.coder_id_from_element_type(self.element_type)


chamikaramj · 2020-03-26T22:49:01Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

+            import only_element
+          output_pcoll = (
+              transform_node.outputs[None] if None in transform_node.outputs
+              else only_element(transform_node.outputs.values()))


chamikaramj · 2020-03-27T14:47:50Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

@@ -380,7 +381,12 @@ def visit_transform(self, transform_node):
        # pylint: disable=wrong-import-order, wrong-import-position
        from apache_beam import Flatten
        if isinstance(transform_node.transform, Flatten):
-          output_pcoll = transform_node.outputs[None]
+          from apache_beam.runners.portability.fn_api_runner.translations \


Implementation seems to be pretty straightforward. So just added a version to this file.

chamikaramj · 2020-03-27T15:28:58Z

sdks/python/apache_beam/coders/coders.py

+    kind_str = 'kind:external' + str(ExternalCoder.coder_count)
+    ExternalCoder.coder_count = ExternalCoder.coder_count + 1
+    component_encodings = []
+    if coder_proto.spec.urn == 'beam:coder:kv:v1':


This does not seems to be needed for now (at least to get the test suite working). So I added a TODO here. Also, note that 'kind:extenal' was used to represent Java only Void coder here, so we'll probably continue to need that even if we handle 'kind:stream' here. Lemme know if you think this is inadequate for this PR.

robertwb

OK, I've finished going through the full set of changes.

robertwb · 2020-03-27T22:13:02Z

sdks/python/apache_beam/coders/coders.py

@@ -370,7 +370,8 @@ def from_runner_api(cls, coder_proto, context):
    except Exception:
      if context.allow_proto_holders:
        # ignore this typing scenario for now, since it can't be easily tracked
-        return RunnerAPICoderHolder(coder_proto)  # type: ignore
+        return ExternalCoder(


I'd rather not do this for arbitrary exceptions, instead let's do this iff coder_proto.spec.urn is not in cls._known_urns.

Thanks. I think it'd be cleaner to just put this if instead of the try, and not have a broadly-catching "except Exception" block.

robertwb · 2020-03-27T22:18:13Z

sdks/python/apache_beam/coders/coders.py

+    kind_str = 'kind:external' + str(ExternalCoder.coder_count)
+    ExternalCoder.coder_count = ExternalCoder.coder_count + 1
+    component_encodings = []
+    if coder_proto.spec.urn == 'beam:coder:kv:v1':


I anticipate kind:stream will be needed to handle GBK of unknown types. Others may be needed for other cases, or in the future, and it seems risky to enumerate them here and in the dataflow runner. There may also be cases where we have to go more than one level deep. We should try to return the same thing the external SDK would have returned just to be safe, and that means wrapping only the leaves as external coders. I think that'll clean stuff up as well (e.g. no need for _coerce_to_kv_type_from_external_type).

robertwb · 2020-03-27T22:19:50Z

sdks/python/apache_beam/pipeline.py

@@ -1102,6 +1105,10 @@ def transform_to_runner_api(transform,  # type: Optional[ptransform.PTransform]
        (transform_urn in Pipeline.sdk_transforms_with_environment())):
      environment_id = context.default_environment_id()

+    def _may_be_preserve_tag(new_tag, pc, input_tags_to_preserve):


maybe is one word.

robertwb · 2020-03-27T22:22:05Z

sdks/python/apache_beam/pipeline.py

      # type: (str) -> bool
      # As per named_inputs() above.
-      return tag.startswith('side')
+      return re.match(SIDE_INPUT_REGEX, tag)


This'll break if a java user names their side inputs coincidentally something that matches.

I changed the Python side input prefix from 'side' to 'python_side_input' (SIDE_INPUT_PREFIX) to make this tighter. This is still not 100% fail safe (an external SDK might still start side input tags with 'python_side_input') but I don't think this will be an issue in practice.

robertwb · 2020-03-27T22:23:11Z

sdks/python/apache_beam/pipeline.py

@@ -1110,7 +1117,8 @@ def transform_to_runner_api(transform,  # type: Optional[ptransform.PTransform]
            for part in self.parts
        ],
        inputs={
-            tag: context.pcollections.get_id(pc)
+            _may_be_preserve_tag(tag, pc, self.input_tags_to_preserve):
+            context.pcollections.get_id(pc)
            for tag,
            pc in sorted(self.named_inputs().items())


While we're here, did yapf not like this on the previous line? Maybe write (tag, pc) (unless it tries to split that too).

robertwb · 2020-03-27T22:53:17Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

+      # parsed in Python SDK.
+      if isinstance(element_type, ElementTypeHolder):
+        coder = self._get_encoding_for_external_environment(element_type)
+        if window_value:


Don't duplicate this logic, instead augment _get_typehint_based_encoding (well, really coders.registry.get_coder) to do the natural thing for ElementTypeHolders.

robertwb · 2020-03-27T22:56:38Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

@@ -925,8 +970,11 @@ def run_ParDo(self, transform_node, options):
        (transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn or
         use_unified_worker)):
      # Patch side input ids to be unique across a given pipeline.
-      if (label_renames and
-          transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn):
+      # This should not be done for external transforms since external SDKs may


It seems cleaner to simply not populate the label_renames map in this case rather than add a branch here.

robertwb · 2020-03-27T22:59:43Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

@@ -995,10 +1044,13 @@ def run_ParDo(self, transform_node, options):
      # The assumption here is that all outputs will have the same typehint
      # and coder as the main output. This is certainly the case right now
      # but conceivably it could change in the future.
+      encoding = self._get_encoded_output_coder(
+          transform_node,
+          output_tag=side_tag) if external_transform else step.encoding


Why do we need to branch on external_transform her?

We don't have to. Previously we were using same encoding for all outputs for Python but seems like we can just generalize this without issues.

robertwb · 2020-03-27T23:02:34Z

sdks/python/apache_beam/runners/pipeline_context.py

    # type: (str, message.Message) -> str
-    if id in self._id_to_proto:
+    if not ignore_duplicates and id in self._id_to_proto:


If ignore_duplicates=True, should we ensure that the protos are the same?

robertwb · 2020-03-27T23:05:55Z

sdks/python/apache_beam/coders/coders.py

@@ -1383,22 +1384,69 @@ def from_runner_api_parameter(payload, components, context):
        write_state_threshold=int(payload))


-class RunnerAPICoderHolder(Coder):
+class ElementTypeHolder(typehints.TypeConstraint):


Should this be an ExternalElementType?

Even better, perhaps this should be a CoderElementType that holds a Coder (external or not), and then an ExternalCoder would be a Coder that holds a proto.

chamikaramj

Thanks Robert. PTAL.

chamikaramj · 2020-03-29T21:15:46Z

sdks/python/apache_beam/coders/coders.py

@@ -370,7 +370,8 @@ def from_runner_api(cls, coder_proto, context):
    except Exception:
      if context.allow_proto_holders:
        # ignore this typing scenario for now, since it can't be easily tracked
-        return RunnerAPICoderHolder(coder_proto)  # type: ignore
+        return ExternalCoder(


chamikaramj · 2020-03-29T21:29:50Z

sdks/python/apache_beam/coders/coders.py

@@ -1383,22 +1384,69 @@ def from_runner_api_parameter(payload, components, context):
        write_state_threshold=int(payload))


-class RunnerAPICoderHolder(Coder):
+class ElementTypeHolder(typehints.TypeConstraint):


chamikaramj · 2020-03-29T21:40:23Z

sdks/python/apache_beam/pipeline.py

@@ -1102,6 +1105,10 @@ def transform_to_runner_api(transform,  # type: Optional[ptransform.PTransform]
        (transform_urn in Pipeline.sdk_transforms_with_environment())):
      environment_id = context.default_environment_id()

+    def _may_be_preserve_tag(new_tag, pc, input_tags_to_preserve):


chamikaramj · 2020-03-29T23:20:14Z

sdks/python/apache_beam/pipeline.py

+      payload = (
+          proto_utils.parse_Bytes(
+              proto.spec.payload, beam_runner_api_pb2.ParDoPayload))
+      for tag, si in payload.side_inputs.items():


chamikaramj · 2020-03-30T00:08:16Z

sdks/python/apache_beam/pipeline.py

    side_inputs = [si for _, si in sorted(indexed_side_inputs)]
+
+    input_tags_to_preserve = {}


chamikaramj · 2020-03-31T02:30:36Z

sdks/python/apache_beam/pipeline.py

@@ -1110,7 +1117,8 @@ def transform_to_runner_api(transform,  # type: Optional[ptransform.PTransform]
            for part in self.parts
        ],
        inputs={
-            tag: context.pcollections.get_id(pc)
+            _may_be_preserve_tag(tag, pc, self.input_tags_to_preserve):
+            context.pcollections.get_id(pc)
            for tag,
            pc in sorted(self.named_inputs().items())


chamikaramj · 2020-03-31T15:32:21Z

sdks/python/apache_beam/pipeline.py

+    side_input_tags = []
+    if common_urns.primitives.PAR_DO.urn == proto.spec.urn:
+      # Preserving side input tags.
+      from apache_beam.utils import proto_utils


Moved proto_utils but left beam_runner_api_pb2 to prevent adding a dependency that to all runners (seems like this is imported similarly in other places in this file I suspect for the same reason).

chamikaramj · 2020-03-31T20:49:59Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

+      raise Exception(
+          'Expected to receive an object of type ElementTypeHolder '
+          'but received %r' % element)
+    from apache_beam.coders.coders import ExternalCoder


chamikaramj · 2020-03-31T20:50:42Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

+    If output_tag is not specified, we assume all outputs to have the same
+    encoding.
+    """
+    from apache_beam.transforms.core import RunnerAPIPTransformHolder


chamikaramj · 2020-03-31T21:26:42Z

sdks/python/apache_beam/coders/coders.py

+    kind_str = 'kind:external' + str(ExternalCoder.coder_count)
+    ExternalCoder.coder_count = ExternalCoder.coder_count + 1
+    component_encodings = []
+    if coder_proto.spec.urn == 'beam:coder:kv:v1':


Cleaned this up. Now we only handle unknown types here.

chamikaramj · 2020-04-02T17:38:14Z

Run Python PreCommit

chamikaramj · 2020-04-02T18:52:36Z

Run Portable_Python PreCommit

robertwb

Getting there; pretty minor comments at this point.

robertwb · 2020-04-03T00:21:05Z

sdks/python/apache_beam/coders/coders.py

@@ -370,7 +370,8 @@ def from_runner_api(cls, coder_proto, context):
    except Exception:
      if context.allow_proto_holders:
        # ignore this typing scenario for now, since it can't be easily tracked
-        return RunnerAPICoderHolder(coder_proto)  # type: ignore
+        return ExternalCoder(


Thanks. I think it'd be cleaner to just put this if instead of the try, and not have a broadly-catching "except Exception" block.

robertwb · 2020-04-03T00:32:29Z

sdks/python/apache_beam/coders/coders.py

+class CoderElementType(typehints.TypeConstraint):
+  """An element type that just holds a coder proto."""
+  def __init__(self, coder_proto, context):
+    self.coder_proto = coder_proto


Logically the ElementType should hold the ExternalCoder, and the ExternalCoder should hold the proto. (This will also help some issues where coders are requested manually and then not respected.)

robertwb · 2020-04-03T00:34:12Z

sdks/python/apache_beam/pipeline.py

-        for ix, si in enumerate(self.side_inputs)
-    }
+    side_inputs = {(SIDE_INPUT_PREFIX + '%s') % ix: si.pvalue
+                   for ix,


Put (ix, si) to help yapf do the right thing.

robertwb · 2020-04-03T00:37:37Z

sdks/python/apache_beam/pipeline.py

@@ -1107,6 +1111,10 @@ def transform_to_runner_api(transform,  # type: Optional[ptransform.PTransform]
        (transform_urn in Pipeline.sdk_transforms_with_environment())):
      environment_id = context.default_environment_id()

+    def _maybe_preserve_tag(new_tag, pc, input_tags_to_preserve):


Add a TODO(BEAM-1833) to get rid of this.

robertwb · 2020-04-03T00:39:31Z

sdks/python/apache_beam/pipeline.py

-    ]
-    side_inputs = [si for _, si in sorted(indexed_side_inputs)]
+    if python_indexed_side_inputs:
+      # Ordering is important here.


Thanks. This is still has to work around the existing ugliness, but looks much better now.

robertwb · 2020-04-03T00:43:30Z

sdks/python/apache_beam/pipeline.py

+    if isinstance(transform, RunnerAPIPTransformHolder):
+      # For external transforms that are ParDos, we have to set side-inputs
+      # manually and preserve input tags.
+      transform.side_inputs = [pvalue.AsMultiMap(pc) for pc in side_inputs]


That's out of date: https://github.com/apache/beam/blob/master/model/pipeline/src/main/proto/beam_runner_api.proto#L395

We should be able to use pvalue.SideInputData.from_runner_api here.

robertwb · 2020-04-03T00:47:45Z

sdks/python/apache_beam/pipeline.py

+          for tag,
+          id in proto.inputs.items()
+      }
+    else:


Where does transform.side_inputs in this path?

That's done in respective from_runner_api methods. For example,
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/transforms/core.py#L1361

But RunnerAPIPTransformHolder is a holder type that is constructed directly so setting it here seemed like the best option (we may be able to move this to constructor but that would require passing in additional parameters to get access to inputs etc. since using Python only pvalue.SideInputData.from_runner_api is not an option).

robertwb · 2020-04-03T00:48:34Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

-                None if None in transform_node.outputs.keys() else only_element(
-                    transform_node.outputs.keys()))
+                None if None in transform_node.outputs.keys() else
+                DataflowRunner._only_element(transform_node.outputs.keys()))


The conditional is a tautology; just return _only_element(transform_node.outputs.keys()).

Done. I assume this is because there can be only one output here.

robertwb · 2020-04-03T00:49:03Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

-      only_element
-    if len(transform_node.outputs) == 1:
-      output_tag = only_element(transform_node.outputs.keys())
+    external_transform = isinstance(


Nit: is_external_transform.

robertwb · 2020-04-03T00:52:52Z

sdks/python/apache_beam/typehints/typehints.py

@@ -1180,6 +1180,28 @@ def get_yielded_type(type_hint):
  raise ValueError('%s is not iterable' % type_hint)


+def _coerce_to_kv_type_from_external_type(element_type_holder):


Do we still need this?

We don't. Removed :)

These are needed for for runners that need to build a Python object graph from a runner API proto with external transforms (for example, Dataflow).

chamikaramj

Thanks. PTAL.

chamikaramj · 2020-04-03T19:05:16Z

sdks/python/apache_beam/coders/coders.py

@@ -370,7 +370,8 @@ def from_runner_api(cls, coder_proto, context):
    except Exception:
      if context.allow_proto_holders:
        # ignore this typing scenario for now, since it can't be easily tracked
-        return RunnerAPICoderHolder(coder_proto)  # type: ignore
+        return ExternalCoder(


chamikaramj · 2020-04-03T20:41:19Z

sdks/python/apache_beam/coders/coders.py

+class CoderElementType(typehints.TypeConstraint):
+  """An element type that just holds a coder proto."""
+  def __init__(self, coder_proto, context):
+    self.coder_proto = coder_proto


chamikaramj · 2020-04-03T20:43:18Z

sdks/python/apache_beam/pipeline.py

-        for ix, si in enumerate(self.side_inputs)
-    }
+    side_inputs = {(SIDE_INPUT_PREFIX + '%s') % ix: si.pvalue
+                   for ix,


chamikaramj · 2020-04-03T20:47:03Z

sdks/python/apache_beam/pipeline.py

@@ -1107,6 +1111,10 @@ def transform_to_runner_api(transform,  # type: Optional[ptransform.PTransform]
        (transform_urn in Pipeline.sdk_transforms_with_environment())):
      environment_id = context.default_environment_id()

+    def _maybe_preserve_tag(new_tag, pc, input_tags_to_preserve):


chamikaramj · 2020-04-03T21:46:07Z

sdks/python/apache_beam/pipeline.py

+          for tag,
+          id in proto.inputs.items()
+      }
+    else:


That's done in respective from_runner_api methods. For example,
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/transforms/core.py#L1361

But RunnerAPIPTransformHolder is a holder type that is constructed directly so setting it here seemed like the best option (we may be able to move this to constructor but that would require passing in additional parameters to get access to inputs etc. since using Python only pvalue.SideInputData.from_runner_api is not an option).

chamikaramj · 2020-04-03T21:53:01Z

sdks/python/apache_beam/pipeline.py

+    if isinstance(transform, RunnerAPIPTransformHolder):
+      # For external transforms that are ParDos, we have to set side-inputs
+      # manually and preserve input tags.
+      transform.side_inputs = [pvalue.AsMultiMap(pc) for pc in side_inputs]


Unfortunately 'pvalue.SideInputData.from_runner_api' is Python only and fails for external side inputs: https://github.com/apache/beam/blob/master/sdks/python/apache_beam/pvalue.py#L463

So enumerated the two supported types here and added a TODO.

chamikaramj · 2020-04-04T00:06:10Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

-                None if None in transform_node.outputs.keys() else only_element(
-                    transform_node.outputs.keys()))
+                None if None in transform_node.outputs.keys() else
+                DataflowRunner._only_element(transform_node.outputs.keys()))


Done. I assume this is because there can be only one output here.

chamikaramj · 2020-04-04T00:08:24Z

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

-      only_element
-    if len(transform_node.outputs) == 1:
-      output_tag = only_element(transform_node.outputs.keys())
+    external_transform = isinstance(


chamikaramj · 2020-04-04T00:17:53Z

sdks/python/apache_beam/typehints/typehints.py

@@ -1180,6 +1180,28 @@ def get_yielded_type(type_hint):
  raise ValueError('%s is not iterable' % type_hint)


+def _coerce_to_kv_type_from_external_type(element_type_holder):


We don't. Removed :)

robertwb

Thank you. LGTM.

chamikaramj · 2020-04-07T14:04:43Z

Thanks Robert.

probot-autolabeler bot added the python label Mar 20, 2020

robertwb reviewed Mar 21, 2020

View reviewed changes

chamikaramj force-pushed the xlang_generalizations branch from f46ce1d to e84f6e8 Compare March 23, 2020 17:17

chamikaramj commented Mar 23, 2020

View reviewed changes

robertwb reviewed Mar 25, 2020

View reviewed changes

chamikaramj changed the title ~~[BEAM-8019] Some generalizations to support cross-language transforms.~~ [BEAM-8019] Updates Python SDK to handle pass-through coders from remote SDKs during expansion Mar 27, 2020

chamikaramj changed the title ~~[BEAM-8019] Updates Python SDK to handle pass-through coders from remote SDKs during expansion~~ [BEAM-8019] Updates Python SDK to handle remote SDK coders and preserve tags added by remote SDKs Mar 27, 2020

chamikaramj commented Mar 27, 2020

View reviewed changes

robertwb requested changes Mar 27, 2020

View reviewed changes

robertwb changed the title ~~[BEAM-8019] Updates Python SDK to handle remote SDK coders and preserve tags added by remote SDKs~~ [BEAM-8019] Updates Python SDK to handle remote SDK coders and preserve tags added by remote SDKs and propagate restriction coders. Mar 27, 2020

chamikaramj force-pushed the xlang_generalizations branch 2 times, most recently from e88eba0 to e5dfdc9 Compare April 2, 2020 01:25

chamikaramj commented Apr 2, 2020

View reviewed changes

robertwb reviewed Apr 3, 2020

View reviewed changes

chamikaramj added 7 commits April 3, 2020 17:43

Some generalizations to support cross-language transforms.

1459f8a

These are needed for for runners that need to build a Python object graph from a runner API proto with external transforms (for example, Dataflow).

Fixes to tests and lint.

0dc9b29

Test and lint fixes.

2dd8458

Addresses reviewer comments.

96c3a81

Addresses reviewer comments.

2be17e6

Fixes lint.

33a18ec

Addressed reviewer comments.

ac308c6

chamikaramj force-pushed the xlang_generalizations branch from 44d99a4 to ac308c6 Compare April 4, 2020 01:57

chamikaramj commented Apr 4, 2020

View reviewed changes

robertwb approved these changes Apr 6, 2020

View reviewed changes

robertwb merged commit 58c51bf into apache:master Apr 6, 2020

chamikaramj mentioned this pull request Dec 12, 2020

[BEAM-11360] Updates Dataflow Python multi-language pipelines to use portable job submission by default #13536

Merged

4 tasks

		side_inputs = [si for _, si in sorted(indexed_side_inputs)]

		input_tags_to_preserve = {}

		@@ -1180,6 +1180,28 @@ def get_yielded_type(type_hint):
		raise ValueError('%s is not iterable' % type_hint)


		def _coerce_to_kv_type_from_external_type(element_type_holder):

[BEAM-8019] Updates Python SDK to handle remote SDK coders and preserve tags added by remote SDKs and propagate restriction coders. #11185

[BEAM-8019] Updates Python SDK to handle remote SDK coders and preserve tags added by remote SDKs and propagate restriction coders. #11185

Conversation

chamikaramj commented Mar 20, 2020 • edited

Post-Commit Tests Status (on master branch)

Pre-Commit Tests Status (on master branch)

chamikaramj commented Mar 20, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chamikaramj left a comment

Choose a reason for hiding this comment

chamikaramj commented Mar 23, 2020

chamikaramj commented Mar 23, 2020

chamikaramj commented Mar 23, 2020

chamikaramj commented Mar 23, 2020

robertwb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chamikaramj left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

robertwb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chamikaramj left a comment

Choose a reason for hiding this comment

chamikaramj commented Mar 20, 2020 •

edited