Merge pull request #11283 [BEAM-9322] [BEAM-1833] Better naming for composite transform output tags.

Author: robertwb

CHANGES.md

@@ -131,14 +131,14 @@ Schema Options, it will be removed in version `2.23.0`. ([BEAM-9704](https://iss
  used to be deployed to
  [apachebeam](https://hub.docker.com/search?q=apachebeam&type=image) repository.
  ([BEAM-9063](https://issues.apache.org/jira/browse/BEAM-9093))
+* PCollections now have tags inferred from the result type (e.g. the keys of a dict or index of a tuple).  Users may expect the old implementation which gave PCollection output ids a monotonically increasing id. To go back to the old implementation, use the `force_generated_pcollection_output_ids` experiment.
 
 ## Deprecations
 
 ## Bugfixes
 
 * Fixed numpy operators in ApproximateQuantiles (Python) ([BEAM-9579](https://issues.apache.org/jira/browse/BEAM-9579)).
 * Fixed exception when running in IPython notebook (Python) ([BEAM-X9277](https://issues.apache.org/jira/browse/BEAM-9277)).
-* Fixed 1833 (Python) ([BEAM-1833](https://issues.apache.org/jira/browse/BEAM-1833))
 * Fixed Flink uberjar job termination bug. ([BEAM-9225](https://issues.apache.org/jira/browse/BEAM-9225))
 * Fixed SyntaxError in process worker startup ([BEAM-9503](https://issues.apache.org/jira/browse/BEAM-9503))
 

sdks/python/apache_beam/pipeline.py

@@ -636,7 +636,7 @@ def apply(
       if type_options is not None and type_options.pipeline_type_check:
         transform.type_check_outputs(pvalueish_result)
 
-      for result in ptransform.get_nested_pvalues(pvalueish_result):
+      for tag, result in ptransform.get_named_nested_pvalues(pvalueish_result):
         assert isinstance(result, (pvalue.PValue, pvalue.DoOutputsTuple))
 
         # Make sure we set the producer only for a leaf node in the transform
@@ -669,12 +669,16 @@ def apply(
           current.add_output(result, tag)
           continue
 
-        # TODO(BEAM-9322): Find the best auto-generated tags for nested
-        # PCollections.
-        # If the user wants the old implementation of always generated
-        # PCollection output ids, then set the tag to None first, then count up
-        # from 1.
-        tag = len(current.outputs) if None in current.outputs else None
+        if self._options.view_as(DebugOptions).lookup_experiment(
+            'force_generated_pcollection_output_ids', default=False):
+          tag = len(current.outputs) if None in current.outputs else None
+        else:
+          base = tag
+          counter = 0
+          while tag in current.outputs:
+            counter += 1
+            tag = '%s_%d' % (base, counter)
+
         current.add_output(result, tag)
 
       if (type_options is not None and
@@ -1252,7 +1256,8 @@ def is_python_side_input(tag):
     return result
 
 
-class PTransformOverride(with_metaclass(abc.ABCMeta, object)):  # type: ignore[misc]
+class PTransformOverride(with_metaclass(abc.ABCMeta,
+                                        object)):  # type: ignore[misc]
   """For internal use only; no backwards-compatibility guarantees.
 
   Gives a matcher and replacements for matching PTransforms.

sdks/python/apache_beam/transforms/ptransform.py

@@ -256,21 +256,29 @@ def get_nested_pvalues(pvalueish):
   return pvalues
 
 
-def get_nested_pvalues0(pvalueish):
-  if isinstance(pvalueish, (tuple, list)):
+def get_named_nested_pvalues(pvalueish):
+  if isinstance(pvalueish, tuple):
+    # Check to see if it's a named tuple.
+    fields = getattr(pvalueish, '_fields', None)
+    if fields and len(fields) == len(pvalueish):
+      tagged_values = zip(fields, pvalueish)
+    else:
+      tagged_values = enumerate(pvalueish)
+  elif isinstance(pvalueish, list):
     tagged_values = enumerate(pvalueish)
-  if isinstance(pvalueish, dict):
+  elif isinstance(pvalueish, dict):
     tagged_values = pvalueish.items()
   else:
-    yield None, pvalueish
+    if isinstance(pvalueish, (pvalue.PValue, pvalue.DoOutputsTuple)):
+      yield None, pvalueish
     return
 
   for tag, subvalue in tagged_values:
-    for subtag, subsubvalue in get_nested_pvalues(subvalue):
+    for subtag, subsubvalue in get_named_nested_pvalues(subvalue):
       if subtag is None:
         yield tag, subsubvalue
       else:
-        yield '%s.%s' % (tag, subsubvalue), subsubvalue
+        yield '%s.%s' % (tag, subtag), subsubvalue
 
 
 class _ZipPValues(object):