lukecwik commented on a change in pull request #12289:
URL: https://github.com/apache/beam/pull/12289#discussion_r456708526
##########
File path: sdks/python/apache_beam/runners/worker/operations.py
##########
@@ -743,6 +743,32 @@ def pcollection_count_monitoring_infos(self,
tag_to_pcollection_id):
return infos
+class SdfTruncateSizedRestrictions(DoOperation):
+ def __init__(self, *args, **kwargs):
+ super(SdfTruncateSizedRestrictions, self).__init__(*args, **kwargs)
+ self.sdf_process_op = None
+ self.consumers_count = 0
+
+ def current_element_progress(self):
+ # type: () -> Optional[iobase.RestrictionProgress]
+ if self.consumers_count == 1 and self.sdf_process_op is None:
+ return None
+ return self.sdf_process_op.current_element_progress()
+
+ def try_split(self, fraction_of_remainder): # type: (...) -> Optional[Any]
+ if self.consumers_count == 1 and self.sdf_process_op is None:
+ return None
+ return self.sdf_process_op.try_split(fraction_of_remainder)
+
+ def add_receiver(self, operation, output_index=0):
Review comment:
We shouldn't need to record which is the sdf_process_op, we should be
able to use `self.receivers[0].try_split` and
`self.receivers[0].current_element_progress` regardless. It will only forward
the split/progress if there is exactly one consumer and it supports splitting.
##########
File path: sdks/python/apache_beam/runners/worker/operations.py
##########
@@ -743,6 +743,32 @@ def pcollection_count_monitoring_infos(self,
tag_to_pcollection_id):
return infos
+class SdfTruncateSizedRestrictions(DoOperation):
+ def __init__(self, *args, **kwargs):
+ super(SdfTruncateSizedRestrictions, self).__init__(*args, **kwargs)
+ self.sdf_process_op = None
+ self.consumers_count = 0
+
+ def current_element_progress(self):
+ # type: () -> Optional[iobase.RestrictionProgress]
+ if self.consumers_count == 1 and self.sdf_process_op is None:
+ return None
+ return self.sdf_process_op.current_element_progress()
+
+ def try_split(self, fraction_of_remainder): # type: (...) -> Optional[Any]
Review comment:
Please fix the type, after #12275 it is:
Optional[Tuple[Iterable[SdfSplitResultsPrimary],
Iterable[SdfSplitResultsResidual]]
----------------------------------------------------------------
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]
