This is an automated email from the ASF dual-hosted git repository.
xqhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 624b2fcd3fd Avoid unreasonably long stage names for @ptransform_fn.
(#35660)
624b2fcd3fd is described below
commit 624b2fcd3fdd249d3a13b6faba6e71521e9fd0c1
Author: Robert Bradshaw <[email protected]>
AuthorDate: Wed Aug 6 08:42:23 2025 -0400
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 <[email protected]>
* 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 <[email protected]>
---
sdks/python/apache_beam/pipeline.py | 10 +++++--
sdks/python/apache_beam/transforms/ptransform.py | 28 +++++++++++++++---
.../apache_beam/transforms/ptransform_test.py | 33 ++++++++++++++++++++++
3 files changed, 64 insertions(+), 7 deletions(-)
diff --git a/sdks/python/apache_beam/pipeline.py
b/sdks/python/apache_beam/pipeline.py
index 269b4acdc21..83a0bee8145 100644
--- a/sdks/python/apache_beam/pipeline.py
+++ b/sdks/python/apache_beam/pipeline.py
@@ -115,11 +115,11 @@ __all__ = ['Pipeline', 'transform_annotations']
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 @@ class Pipeline(HasDisplayData):
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)
diff --git a/sdks/python/apache_beam/transforms/ptransform.py
b/sdks/python/apache_beam/transforms/ptransform.py
index c4f0e3455d4..d2cf836713f 100644
--- a/sdks/python/apache_beam/transforms/ptransform.py
+++ b/sdks/python/apache_beam/transforms/ptransform.py
@@ -1002,6 +1002,7 @@ class _PTransformFnPTransform(PTransform):
self._fn = fn
self._args = args
self._kwargs = kwargs
+ self._use_backwards_compatible_label = True
def display_data(self):
res = {
@@ -1030,11 +1031,30 @@ class _PTransformFnPTransform(PTransform):
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):
diff --git a/sdks/python/apache_beam/transforms/ptransform_test.py
b/sdks/python/apache_beam/transforms/ptransform_test.py
index 78d5c3ef38b..e1c84c7dc9a 100644
--- a/sdks/python/apache_beam/transforms/ptransform_test.py
+++ b/sdks/python/apache_beam/transforms/ptransform_test.py
@@ -1157,6 +1157,39 @@ class PTransformLabelsTest(unittest.TestCase):
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:
