[GitHub] [beam] boyuanzz commented on a change in pull request #13405: [BEAM-10703, BEAM-10475] Add GroupIntoBatches to runner API; add Dataflow override in Python SDK

[GitBox <gi...@apache.org>]

boyuanzz commented on a change in pull request #13405:
URL: https://github.com/apache/beam/pull/13405#discussion_r534539958



##########
File path: sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
##########
@@ -1112,6 +1128,17 @@ def run_ParDo(self, transform_node, options):
       if is_stateful_dofn:
         step.add_property(PropertyNames.USES_KEYED_STATE, 'true')
 
+        # Also checks whether the step allows shardable keyed states.
+        for pcoll in transform_node.outputs.values():

Review comment:
       Are you going to take care of portable job submission as well? If fo, maybe we need a TODO here.

##########
File path: sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
##########
@@ -329,3 +330,44 @@ def expand(self, pcoll):
         return {key: out for key in self.outputs}
 
     return WriteToBigQuery(ptransform, self.outputs)
+
+
+class GroupIntoBatchesWithShardedKeyPTransformOverride(PTransformOverride):
+  """A ``PTransformOverride`` for ``GroupIntoBatches.WithShardedKey``.
+
+  This override simply returns the original transform but additionally records
+  the output PCollection in order to append required step properties during
+  graph translation.
+  """
+  def __init__(self, dataflow_runner, options):
+    self.dataflow_runner = dataflow_runner
+    self.options = options
+
+  def matches(self, applied_ptransform):
+    # Imported here to avoid circular dependencies.
+    # pylint: disable=wrong-import-order, wrong-import-position
+    from apache_beam import util
+
+    transform = applied_ptransform.transform
+
+    if not isinstance(transform, util.GroupIntoBatches.WithShardedKey):
+      return False
+
+    # The replacement is only valid for portable Streaming Engine jobs.
+    standard_options = self.options.view_as(StandardOptions)
+    if not standard_options.streaming:
+      return False
+    google_cloud_options = self.options.view_as(GoogleCloudOptions)
+    if not google_cloud_options.enable_streaming_engine:
+      return False
+    experiments = self.options.view_as(DebugOptions).experiments or []
+    if 'beam_fn_api' not in experiments:

Review comment:
       I think we also want to check `use_runner_v2` as well since it's only supported in runner v2




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org