Workaround in test for an issue with beam 2.40 not allowing a built-in function to be passed to transformations.

This change also adds the TFT label to the raised error when constructing the ptransform fails.

PiperOrigin-RevId: 457773563

Author: zoyahav

tensorflow_transform/beam/cached_impl_test.py

@@ -1104,10 +1104,13 @@ def test_cached_ptransform_analyzer(self, use_tf_compat_v1):
     class _AnalyzerMakeAccumulators(beam.PTransform):
 
       def expand(self, pcoll):
+        # TODO(b/237367328): Use sum directly when beam>=2.40 allows it.
+        def _sum(x):
+          return sum(x)
         input_sum = pcoll | beam.FlatMap(
-            sum) | 'ReduceSum' >> beam.CombineGlobally(sum)
+            _sum) | 'ReduceSum' >> beam.CombineGlobally(_sum)
         size = pcoll | beam.Map(
-            np.size) | 'ReduceCount' >> beam.CombineGlobally(sum)
+            np.size) | 'ReduceCount' >> beam.CombineGlobally(_sum)
 
         return (pcoll.pipeline
                 | beam.Create([None])

tensorflow_transform/beam/common.py

@@ -190,8 +190,11 @@ def visit(self, operation, inputs):
       tagged_label = operation.label
     else:
       tagged_label = '{label}[{tag}]'.format(label=operation.label, tag=tag)
-    outputs = ((inputs or beam.pvalue.PBegin(self._extra_args.pipeline))
-               | tagged_label >> ptransform(operation, self._extra_args))
+    try:
+      outputs = ((inputs or beam.pvalue.PBegin(self._extra_args.pipeline))
+                 | tagged_label >> ptransform(operation, self._extra_args))
+    except Exception as e:
+      raise RuntimeError('Failed to apply: {}'.format(tagged_label)) from e
 
     if isinstance(outputs, beam.pvalue.PCollection):
       return (outputs,)