add reshuffle as a first class yaml transform and a test (#38046)

Author: derrickaw

sdks/python/apache_beam/yaml/tests/reshuffle.yaml

@@ -0,0 +1,48 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+pipelines:
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {id: 1, value: "a"}
+              - {id: 2, value: "b"}
+              - {id: 3, value: "c"}
+              - {id: 4, value: "d"}
+              - {id: 5, value: "e"}
+              - {id: 6, value: "f"}
+              - {id: 7, value: "g"}
+              - {id: 8, value: "h"}
+              - {id: 9, value: "i"}
+              - {id: 10, value: "j"}
+        - type: Reshuffle
+        - type: AssertEqual
+          config:
+            elements:
+              - {id: 1, value: "a"}
+              - {id: 2, value: "b"}
+              - {id: 3, value: "c"}
+              - {id: 4, value: "d"}
+              - {id: 5, value: "e"}
+              - {id: 6, value: "f"}
+              - {id: 7, value: "g"}
+              - {id: 8, value: "h"}
+              - {id: 9, value: "i"}
+              - {id: 10, value: "j"}

sdks/python/apache_beam/yaml/yaml_provider.py

@@ -1207,6 +1207,22 @@ def log_and_return(x):
 
     return pcoll | "LogForTesting" >> beam.Map(log_and_return)
 
+  class Reshuffle(beam.PTransform):
+    """Reshuffles the elements of a PCollection.
+
+    Redistributes the elements of a PCollection to prevent fusion or balance
+    load.
+
+    Args:
+      num_buckets: (optional) Specifies the maximum random keys that would be
+        generated. If not set, a default value is used.
+    """
+    def __init__(self, num_buckets: Optional[int] = None):
+      self.num_buckets = num_buckets
+
+    def expand(self, pcoll):
+      return pcoll | beam.Reshuffle(num_buckets=self.num_buckets)
+
   @staticmethod
   def create_builtin_provider():
     return InlineProvider({
@@ -1216,6 +1232,7 @@ def create_builtin_provider():
         'PyTransform': YamlProviders.fully_qualified_named_transform,
         'Flatten': YamlProviders.Flatten,
         'WindowInto': YamlProviders.WindowInto,
+        'Reshuffle': YamlProviders.Reshuffle,
     },
                           no_input_transforms=('Create', ))