prodriguezdefino opened a new pull request, #41182:
URL: https://github.com/apache/beam/pull/41182
…estStream
Problem:
In stageState.bundleReady
(sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go), a stateful
stage whose upstream watermark has not changed is marked ready only after a
stage has emitted a self-checkpointing residual:
case isStatefulStage && len(ss.sides) == 0 && em.sawResidual.Load() &&
ss.hasBuildableDataLocked(upstreamW):
case inputW == upstreamW && previousInputW == inputW:
return mtime.MinTimestamp, false, ptimeEventsReady, injectedReady
PersistBundle also skips tracking changedConsumers when
em.sawResidual.Load() is false. When TestStream emits elements with
AddElementEvent without advancing the watermark, upstreamW, inputW, and
previousInputW all stay at MinTimestamp and em.sawResidual is false. Prism
leaves the stateful stage unscheduled, decides the pipeline is quiescent, and
runs the next TestStream AdvanceProcessingTime event before the DoFn processes
the elements and sets its processing-time timer. In addition,
makeTestStreamHandler (teststream.go) leaves processingTime at the zero
time.Time instead of initializing it as its comment states.
Fix:
Add ElementManager.canHoldWatermark (true when sawResidual is set or
testStreamHandler is present) and use it in PersistBundle and
stageState.bundleReady so stateful stages with no side inputs and buildable
pending data run before TestStream advances processing time. Initialize
processingTime in makeTestStreamHandler, and add unit and pipeline tests in
teststream_test.go and engine_test.go.
Reproduction (Python SDK on PrismRunner):
```python
import time
import apache_beam as beam
from apache_beam.coders import VarIntCoder
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.testing.test_stream import TestStream
from apache_beam.testing.util import assert_that, equal_to
from apache_beam.transforms.timeutil import TimeDomain
from apache_beam.transforms.userstate import BagStateSpec, TimerSpec,
on_timer
class BufferWithTimeoutDoFn(beam.DoFn):
BUF = BagStateSpec("buf", VarIntCoder())
FLUSH = TimerSpec("flush", TimeDomain.REAL_TIME)
def process(
self,
element,
buf=beam.DoFn.StateParam(BUF),
flush=beam.DoFn.TimerParam(FLUSH),
):
_, val = element
if not list(buf.read()):
flush.set(int(time.time()) + 10)
buf.add(val)
@on_timer(FLUSH)
def on_flush(self, key=beam.DoFn.KeyParam,
buf=beam.DoFn.StateParam(BUF)):
vals = sorted(buf.read())
if vals:
buf.clear()
yield (key, vals)
if __name__ == "__main__":
options = PipelineOptions(["--runner=PrismRunner",
"--environment_type=LOOPBACK"])
with beam.Pipeline(options=options) as p:
out = (
p
| TestStream()
.add_elements([("k", 1), ("k", 2)])
.advance_processing_time(60)
.add_elements([("k", 3)])
.advance_processing_time(60)
.advance_watermark_to_infinity()
| beam.ParDo(BufferWithTimeoutDoFn())
)
assert_that(out, equal_to([("k", [1, 2]), ("k", [3])]))
```
Before the fix: Prism runs `advance_processing_time(60)` before `process`,
emits `[('k', [1, 2, 3])]`, and `assert_that` fails.
After the fix: Prism emits `[('k', [1, 2]), ('k', [3])]`.
<!--
Thank you for your contribution! Follow this checklist to help us
incorporate your contribution quickly and easily:
- [ ] If this contribution is large, please file an Apache [Individual
Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
- [ ] Replace these instructions with a description of your change and why
you have made it.
- [ ] Mention the appropriate issue in your description (for example:
`addresses #123` or `fixes #456`).
- [ ] Update `CHANGES.md` with noteworthy changes.
See [CONTRIBUTING.md]() for more details.
-->
--
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]
