razvanculea commented on code in PR #34777:
URL: https://github.com/apache/beam/pull/34777#discussion_r2074890016
##########
sdks/python/apache_beam/io/iobase.py:
##########
@@ -1127,47 +1130,155 @@ def __init__(self, sink: Sink) -> None:
self.sink = sink
def expand(self, pcoll):
- do_once = pcoll.pipeline | 'DoOnce' >> core.Create([None])
- init_result_coll = do_once | 'InitializeWrite' >> core.Map(
- lambda _, sink: sink.initialize_write(), self.sink)
+ if (pcoll.is_bounded):
+ do_once = pcoll.pipeline | 'DoOnce' >> core.Create([None])
+ init_result_coll = do_once | 'InitializeWrite' >> core.Map(
+ lambda _, sink: sink.initialize_write(), self.sink)
if getattr(self.sink, 'num_shards', 0):
min_shards = self.sink.num_shards
- if min_shards == 1:
- keyed_pcoll = pcoll | core.Map(lambda x: (None, x))
- else:
- keyed_pcoll = pcoll | core.ParDo(_RoundRobinKeyFn(), count=min_shards)
- write_result_coll = (
+
+ if (pcoll.is_bounded):
+ if min_shards == 1:
+ keyed_pcoll = pcoll | core.Map(lambda x: (None, x))
+ else:
+ keyed_pcoll = pcoll | core.ParDo(_RoundRobinKeyFn(),
count=min_shards)
+ write_result_coll = (
+ keyed_pcoll
+ | core.WindowInto(window.GlobalWindows())
+ | core.GroupByKey()
+ | 'WriteBundles' >> core.ParDo(
+ _WriteKeyedBundleDoFn(self.sink),
AsSingleton(init_result_coll)))
+ else: #unbounded PCollection needes to be written per window
+ if isinstance(pcoll.windowing.windowfn, window.GlobalWindows):
+ widowed_pcoll = (
+ pcoll
+ |
core.WindowInto(window.FixedWindows(self.sink.triggering_frequency),
+
trigger=beam.transforms.trigger.AfterWatermark(),
+
accumulation_mode=beam.transforms.trigger.AccumulationMode.DISCARDING,
+
allowed_lateness=beam.utils.timestamp.Duration(seconds=0))
+ )
Review Comment:
that would be an alternative way to implement.
The windowing will help have elements of the same window in the same file,
vs group into batches & trigger will have elements as they come (possibly mixed)
--
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.
To unsubscribe, e-mail: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
