github-actions[bot] commented on code in PR #41332:
URL: https://github.com/apache/beam/pull/41332#discussion_r4572912642
##########
sdks/python/apache_beam/runners/direct/transform_evaluator.py:
##########
@@ -930,9 +940,34 @@ def process_timer(self, timer_firing):
timer_firing.dynamic_timer_tag)
def process_element(self, element):
+ if self._batch_only:
+ key = (element.timestamp_micros, tuple(element.windows),
+ element.pane_info)
+ group = self._batch_groups.get(key)
+ if group is None:
+ group = self._batch_groups[key] = {
+ 'values': [],
+ 'timestamp': element.timestamp,
+ 'windows': element.windows,
+ 'pane_info': element.pane_info,
+ }
+ group['values'].append(element.value)
+ return
self.runner.process(element)
def finish_bundle(self):
+ if self._batch_only:
+ converter = getattr(self._batch_fn, 'input_batch_converter', None)
+ if converter is None:
+ raise NotImplementedError(
+ 'DoFn %s only defines process_batch() and no input batch '
+ 'converter could be constructed' % self._batch_fn)
+ for group in self._batch_groups.values():
+ batch = converter.produce_batch(group['values'])
+ windowed_batch = HomogeneousWindowedBatch.of(
+ batch, group['timestamp'], group['windows'], group['pane_info'])
+ self.runner.process_batch(windowed_batch)
Review Comment:
**🟠High**: `self.runner.process_batch(windowed_batch)` only succeeds when
`process_batch` is decorated with `@DoFn.yields_elements`
(`_process_batch_yields_elements == True`). For a standard `process_batch`-only
`DoFn` that yields batches (the default when `@DoFn.yields_elements` is not
applied), `_OutputHandler.handle_process_batch_outputs` dispatches outputs via
`receiver.receive_batch(windowed_batch)`. Because `_Bundle` and
`_TaggedReceivers.NullReceiver` do not override
`common.Receiver.receive_batch`, calling `self.runner.process_batch` still
raises `NotImplementedError`.
To support batch-yielding `DoFn`s (both default `process_batch` and
`@DoFn.yields_batches`), the receivers in `_tagged_receivers` (or `_Bundle` /
`NullReceiver`) should implement `receive_batch(self, windowed_batch)` by
exploding the batch with `dofn.output_batch_converter.explode_batch`:
```python
for wv in windowed_batch.as_windowed_values(
output_batch_converter.explode_batch):
self.receive(wv)
```
##########
sdks/python/apache_beam/runners/direct/direct_runner_test.py:
##########
@@ -166,6 +166,18 @@ def test_impulse(self):
with test_pipeline.TestPipeline(runner='BundleBasedDirectRunner') as p:
assert_that(p | beam.Impulse(), equal_to([b'']))
+ def test_process_batch_only_dofn(self):
+ class BatchDoFn(beam.DoFn):
+
+ @beam.DoFn.yields_elements
Review Comment:
**🟡 Medium**: Consider also testing a `process_batch`-only `DoFn` without
`@beam.DoFn.yields_elements` (e.g., `def process_batch(self, batch: list[int])
-> Iterator[list[int]]: yield batch`) to verify default batch-emitting
`process_batch` behavior under `BundleBasedDirectRunner`.
##########
sdks/python/apache_beam/runners/direct/transform_evaluator.py:
##########
@@ -930,9 +940,34 @@ def process_timer(self, timer_firing):
timer_firing.dynamic_timer_tag)
def process_element(self, element):
+ if self._batch_only:
Review Comment:
**🟡 Medium**: `WindowedBatch.from_windowed_values(windowed_values,
produce_fn=converter.produce_batch)` in `apache_beam.utils.windowed_value`
already implements grouping `WindowedValue` elements by `(timestamp_micros,
windows, pane_info)` and constructing `HomogeneousWindowedBatch` instances (as
used in `GeneralPurposeConsumerSet.flush` in
`sdks/python/apache_beam/runners/worker/operations.py`).
Instead of manually maintaining `self._batch_groups` dicts in
`process_element`, you can append `element` to a `self._batched_elements:
list[WindowedValue]` list in `process_element` and iterate over
`WindowedBatch.from_windowed_values(self._batched_elements,
produce_fn=converter.produce_batch)` in `finish_bundle`.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
