Avoid unreasonably long stage names for @ptransform_fn. (#35660)

* Avoid unreasonably long stage names for @ptransform_fn.

This respects the update compatibility flag and truncates rather than elides
long first arguments to avoid issues with the prior attempt.

* yapf

* Update sdks/python/apache_beam/transforms/ptransform.py

Co-authored-by: Danny McCormick <dannymccormick@google.com>

* add type hint

* Slightly safer ane more unique.

Includes the tail as well as the prefix (useful for file paths) and ensures the result is a string (mock objects behave strangely).

* Bump compat version.

* Add a pair of tests.

---------

Co-authored-by: Danny McCormick <dannymccormick@google.com>

Author: robertwb

sdks/python/apache_beam/pipeline.py

@@ -115,11 +115,11 @@
 
 
 class Pipeline(HasDisplayData):
-  """A pipeline object that manages a DAG of 
-  :class:`~apache_beam.transforms.ptransform.PTransform` s 
+  """A pipeline object that manages a DAG of
+  :class:`~apache_beam.transforms.ptransform.PTransform` s
   and their :class:`~apache_beam.pvalue.PValue` s.
 
-  Conceptually the :class:`~apache_beam.transforms.ptransform.PTransform` s are 
+  Conceptually the :class:`~apache_beam.transforms.ptransform.PTransform` s are
   the DAG's nodes and the :class:`~apache_beam.pvalue.PValue` s are the edges.
 
   All the transforms applied to the pipeline must have distinct full labels.
@@ -722,6 +722,10 @@ def apply(
       return self.apply(
           transform.transform, pvalueish, label or transform.label)
 
+    if not label and isinstance(transform, ptransform._PTransformFnPTransform):
+      # This must be set before label is inspected.
+      transform.set_options(self._options)
+
     if not isinstance(transform, ptransform.PTransform):
       raise TypeError("Expected a PTransform object, got %s" % transform)
 

sdks/python/apache_beam/transforms/ptransform.py

@@ -1002,6 +1002,7 @@ def __init__(self, fn, *args, **kwargs):
     self._fn = fn
     self._args = args
     self._kwargs = kwargs
+    self._use_backwards_compatible_label = True
 
   def display_data(self):
     res = {
@@ -1030,11 +1031,30 @@ def expand(self, pcoll):
       pass
     return self._fn(pcoll, *args, **kwargs)
 
-  def default_label(self):
+  def set_options(self, options):
+    # Avoid circular import.
+    from apache_beam.transforms.util import is_compat_version_prior_to
+    self._use_backwards_compatible_label = is_compat_version_prior_to(
+        options, '2.68.0')
+
+  def default_label(self) -> str:
+    # Attempt to give a reasonable name to this transform.
+    # We want it to be reasonably unique, but also not sensitive to
+    # irrelevent parameters to minimize pipeline-to-pipeline variance.
+    # For now, use only the first argument (if any), iff it would not make
+    # the name unwieldy.
     if self._args:
-      return '%s(%s)' % (
-          label_from_callable(self._fn), label_from_callable(self._args[0]))
-    return label_from_callable(self._fn)
+      first_arg_string = label_from_callable(self._args[0])
+      if (self._use_backwards_compatible_label or
+          not isinstance(first_arg_string, str) or len(first_arg_string) <= 19):
+        suffix = '(%s)' % first_arg_string
+      else:
+        suffix = ('(%s...%s)' %
+                  (first_arg_string[:10], first_arg_string[-6:])).replace(
+                      '\n', ' ')
+    else:
+      suffix = ''
+    return label_from_callable(self._fn) + suffix
 
 
 def ptransform_fn(fn):

sdks/python/apache_beam/transforms/ptransform_test.py

@@ -1157,6 +1157,39 @@ def test_apply_ptransform_using_decorator(self):
     self.assertTrue('*Sample*/Group' in pipeline.applied_labels)
     self.assertTrue('*Sample*/Distinct' in pipeline.applied_labels)
 
+  def test_ptransformfn_default_label(self):
+    @beam.ptransform_fn
+    def MyTransform(self, suffix="xyz"):
+      return pcoll | beam.Map(lambda s: s + suffix)
+
+    pipeline = TestPipeline()
+    pcoll = pipeline | beam.Create(['a', 'b', 'c'])
+
+    _ = pcoll | MyTransform()
+    self.assertIn('MyTransform', pipeline.applied_labels)
+    _ = pcoll | MyTransform("suffix")
+    self.assertIn('MyTransform(suffix)', pipeline.applied_labels)
+    _ = pcoll | MyTransform("looooooooooooooooooooooooooooooooooooooooong")
+    self.assertIn('MyTransform(looooooooo...oooong)', pipeline.applied_labels)
+
+  def test_ptransformfn_legacy_default_label(self):
+    @beam.ptransform_fn
+    def MyTransform(self, suffix="xyz"):
+      return pcoll | beam.Map(lambda s: s + suffix)
+
+    pipeline = TestPipeline(
+        options=PipelineOptions(update_compatibility_version='2.67.0'))
+    pcoll = pipeline | beam.Create(['a', 'b', 'c'])
+
+    _ = pcoll | MyTransform()
+    self.assertIn('MyTransform', pipeline.applied_labels)
+    _ = pcoll | MyTransform("suffix")
+    self.assertIn('MyTransform(suffix)', pipeline.applied_labels)
+    _ = pcoll | MyTransform("looooooooooooooooooooooooooooooooooooooooong")
+    self.assertIn(
+        'MyTransform(looooooooooooooooooooooooooooooooooooooooong)',
+        pipeline.applied_labels)
+
   def test_combine_with_label(self):
     vals = [1, 2, 3, 4, 5, 6, 7]
     with TestPipeline() as pipeline: