Remove workaround for beam 2.40 not supporting FlatMap(builtin), and modify apache-beam[gcp] dependency to >=2.41.

PiperOrigin-RevId: 476083873

Author: zoyahav

README.md

@@ -96,7 +96,7 @@ other *untested* combinations may also work.
 
 tensorflow-transform                                                            | apache-beam[gcp] | pyarrow | tensorflow        | tensorflow-metadata | tfx-bsl |
 ------------------------------------------------------------------------------- | -----------------| --------|-------------------|---------------------|---------|
-[GitHub master](https://github.com/tensorflow/transform/blob/master/RELEASE.md) | 2.40.0           | 6.0.0   | nightly (1.x/2.x) | 1.10.0              | 1.10.0  |
+[GitHub master](https://github.com/tensorflow/transform/blob/master/RELEASE.md) | 2.41.0           | 6.0.0   | nightly (1.x/2.x) | 1.10.0              | 1.10.0  |
 [1.10.0](https://github.com/tensorflow/transform/blob/v1.10.0/RELEASE.md)       | 2.40.0           | 6.0.0   | 1.15.5 / 2.9      | 1.10.0              | 1.10.0  |
 [1.9.0](https://github.com/tensorflow/transform/blob/v1.9.0/RELEASE.md)         | 2.38.0           | 5.0.0   | 1.15.5 / 2.9      | 1.9.0               | 1.9.0   |
 [1.8.0](https://github.com/tensorflow/transform/blob/v1.8.0/RELEASE.md)         | 2.38.0           | 5.0.0   | 1.15.5 / 2.8      | 1.8.0               | 1.8.0   |

RELEASE.md

@@ -13,6 +13,7 @@
 ## Bug Fixes and Other Changes
 
 *   Depends on `tensorflow>=1.15.5,<2` or `tensorflow>=2.10,<2.11`
+*   Depends on `apache-beam[gcp]>=2.41,<3`.
 
 ## Breaking Changes
 

docs/install.md

@@ -95,7 +95,7 @@ other *untested* combinations may also work.
 
 tensorflow-transform                                                            | apache-beam[gcp] | pyarrow | tensorflow        | tensorflow-metadata | tfx-bsl
 ------------------------------------------------------------------------------- | ---------------- | ------- | ----------------- | ------------------- | -------
-[GitHub master](https://github.com/tensorflow/transform/blob/master/RELEASE.md) | 2.40.0           | 6.0.0   | nightly (1.x/2.x) | 1.10.0              | 1.10.0
+[GitHub master](https://github.com/tensorflow/transform/blob/master/RELEASE.md) | 2.41.0           | 6.0.0   | nightly (1.x/2.x) | 1.10.0              | 1.10.0
 [1.10.0](https://github.com/tensorflow/transform/blob/v1.10.0/RELEASE.md)       | 2.40.0           | 6.0.0   | 1.15.5 / 2.9      | 1.10.0              | 1.10.0
 [1.9.0](https://github.com/tensorflow/transform/blob/v1.9.0/RELEASE.md)         | 2.38.0           | 5.0.0   | 1.15.5 / 2.9      | 1.9.0               | 1.9.0
 [1.8.0](https://github.com/tensorflow/transform/blob/v1.8.0/RELEASE.md)         | 2.38.0           | 5.0.0   | 1.15.5 / 2.8      | 1.8.0               | 1.8.0

setup.py

@@ -43,7 +43,7 @@ def _make_required_install_packages():
   # protobuf) with TF and pyarrow version with tfx-bsl.
   return [
       'absl-py>=0.9,<2.0.0',
-      'apache-beam[gcp]>=2.40,<3',
+      'apache-beam[gcp]>=2.41,<3',
       'numpy>=1.16,<2',
       'protobuf>=3.13,<4',
       'pyarrow>=6,<7',

tensorflow_transform/beam/cached_impl_test.py

@@ -1115,13 +1115,10 @@ 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])