[ 
https://issues.apache.org/jira/browse/BEAM-10703?focusedWorklogId=521528&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-521528
 ]

ASF GitHub Bot logged work on BEAM-10703:
-----------------------------------------

                Author: ASF GitHub Bot
            Created on: 08/Dec/20 04:32
            Start Date: 08/Dec/20 04:32
    Worklog Time Spent: 10m 
      Work Description: nehsyc commented on a change in pull request #13405:
URL: https://github.com/apache/beam/pull/13405#discussion_r538027069



##########
File path: sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
##########
@@ -329,3 +330,46 @@ 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.

Review comment:
       Done.

##########
File path: sdks/python/apache_beam/transforms/util.py
##########
@@ -825,7 +827,60 @@ def expand(self, pcoll):
               key_value[1]))
       return (
           sharded_pcoll
-          | GroupIntoBatches(self.batch_size, 
self.max_buffering_duration_secs))
+          | GroupIntoBatches(
+              self.params.batch_size, self.params.max_buffering_duration_secs))
+
+    def to_runner_api_parameter(
+        self,
+        unused_context  # type: PipelineContext
+    ):  # type: (...) -> Tuple[str, 
beam_runner_api_pb2.GroupIntoBatchesPayload]
+      return (
+          common_urns.composites.GROUP_INTO_BATCHES_WITH_SHARDED_KEY.urn,
+          self.params.get_payload())

Review comment:
       The main reason I extracted the params to a class was to reuse the code 
as much as possible since the params of GroupIntoBatches and WithShardedKey are 
exactly the same. The class also does parameter validation other than 
conversion from/to runner api, and perhaps it would make it easier to add more 
params in the future. But let me know if you have opinions.




----------------------------------------------------------------
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:
[email protected]


Issue Time Tracking
-------------------

    Worklog Id:     (was: 521528)
    Time Spent: 15h 10m  (was: 15h)

> Add support for auto-sharded GroupIntoBatches in Dataflow runner
> ----------------------------------------------------------------
>
>                 Key: BEAM-10703
>                 URL: https://issues.apache.org/jira/browse/BEAM-10703
>             Project: Beam
>          Issue Type: Improvement
>          Components: runner-dataflow
>            Reporter: Siyuan Chen
>            Assignee: Siyuan Chen
>            Priority: P1
>          Time Spent: 15h 10m
>  Remaining Estimate: 0h
>
> The proposal of improving GroupIntoBatches transform is in BEAM-10475
> This tracks the support in Cloud Dataflow Runner.



--
This message was sent by Atlassian Jira
(v8.3.4#803005)

Reply via email to