This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 5dc12703dc9 [Prism] Keep a stateful key in progress while a split
primary holds it (#40624)
5dc12703dc9 is described below
commit 5dc12703dc95fa123456fa843807c174c0843208
Author: pablo rodriguez defino <[email protected]>
AuthorDate: Mon Oct 19 11:30:05 2026 -0700
[Prism] Keep a stateful key in progress while a split primary holds it
(#40624)
When the SDK accepts a channel split of a bundle of a stateful stage,
stageState.splitBundle removed the key of every residual element from
inprogressKeys and inprogressKeysByBundle, even when elements of the
primary still carry that key. The residual of that key was then
scheduled at once in a second bundle that ran concurrently with the
primary. Prism gives each bundle a copy of the key's state at bundle
start and replaces the key's state with the bundle's copy at commit,
so the last bundle to commit wins. State updates, and the outputs that
depend on them, are lost without any error.
Serializing the processing of a key is the runner's job. The model
defines state per key and window: ReadModifyWriteState "holds a single
value per key and window. The value can be read, modified, and written
back atomically" (beam_runner_api.proto). The stateful processing
post says the runner "cannot parallelize processing of elements for a
particular key+window", and the programming guide says "State is
always fully scoped only to the current processing key". Splitting is
valid SDK behaviour (beam_fn_api.proto), and the Java and Python
harnesses accept channel splits on stateful stages.
splitBundle now collects the keys of the primary elements and releases
a residual's key only if no primary element has it. A held key is
released when the primary bundle is persisted, as before. The handling
of residual timer holds is unchanged.
Reproduction with the Python SDK: all elements share one key, and a
per-element sleep makes Prism split the bundle.
class CountOnOneKey(beam.DoFn):
COUNT = ReadModifyWriteStateSpec('count', VarIntCoder())
DONE = TimerSpec('done', TimeDomain.WATERMARK)
def process(self, element, count=beam.DoFn.StateParam(COUNT),
done=beam.DoFn.TimerParam(DONE)):
time.sleep(0.25)
count.write((count.read() or 0) + 1)
done.set(GlobalWindow().max_timestamp())
@on_timer(DONE)
def report(self, count=beam.DoFn.StateParam(COUNT)):
yield count.read()
options = PipelineOptions(['--runner=PrismRunner',
'--prism_location=<prism binary>',
'--environment_type=LOOPBACK'])
with beam.Pipeline(options=options) as p:
(p | beam.Create(range(40)) | beam.Map(lambda i: ('k', i))
| beam.ParDo(CountOnOneKey()) | beam.Map(print))
A correct runner prints 40. Prism built from master printed 20 in 10
of 10 runs; with this change it printed 40 in 10 of 10 runs. With
--prism_log_level=debug, the same pipeline on master logs one split at
element 20 of 40, and the residual bundle starts before the primary
commits.
TestElementManager_SplitStatefulBundle covers a split inside the only
key, a split inside one of two keys, and a split between keys. The
first two fail without this change.
Not fixed here, follow-ups:
- splitBundle slices the bundle's elements by the data channel index.
A stateful bundle can hold data of one key and timers of another, and
timers are not counted in the data channel, so the slice can
misattribute elements. See #35771.
- startTriggeredBundle, for aggregate stages, does not check
inprogressKeys before it injects a bundle.
---
.../prism/internal/engine/elementmanager.go | 15 +-
.../internal/engine/elementmanager_split_test.go | 194 +++++++++++++++++++++
2 files changed, 207 insertions(+), 2 deletions(-)
diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
index 80b4a75ac72..2ed73e2c663 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
@@ -2318,9 +2318,20 @@ func (ss *stageState) splitBundle(rb RunBundle,
firstResidual int, em *ElementMa
es.es = prim
+ // A key stays assigned to this bundle while any primary element still
+ // carries it. Releasing it would let a residual of the same key run in
a
+ // second bundle while the primary is still processing it, so both
bundles
+ // would start from the same state and the last commit would drop the
+ // other's state changes.
+ primKeys := set[string]{}
+ for _, e := range prim {
+ primKeys.insert(string(e.keyBytes))
+ }
for _, e := range res {
- delete(ss.inprogressKeysByBundle[rb.BundleID],
string(e.keyBytes))
- delete(ss.inprogressKeys, string(e.keyBytes))
+ if k := string(e.keyBytes); !primKeys.present(k) {
+ delete(ss.inprogressKeysByBundle[rb.BundleID], k)
+ delete(ss.inprogressKeys, k)
+ }
if e.IsTimer() {
slog.Warn("Unexpected split on a bundle with timers.
See https://github.com/apache/beam/issues/35771 for information.")
diff --git
a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go
b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go
new file mode 100644
index 00000000000..d13552b0a36
--- /dev/null
+++
b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go
@@ -0,0 +1,194 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package engine
+
+import (
+ "fmt"
+ "testing"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
+)
+
+// splitInsideKey returns a split point that keeps the first keep elements of
+// key in the primary and returns the remaining elements of key, and every
+// element after them, as the residual.
+func splitInsideKey(key string, keep int) func(t *testing.T, bundleKeys
[]string) int {
+ return func(t *testing.T, bundleKeys []string) int {
+ t.Helper()
+ for i, k := range bundleKeys {
+ if k == key {
+ return i + keep
+ }
+ }
+ t.Fatalf("key %q not in bundle %v", key, bundleKeys)
+ return 0
+ }
+}
+
+// splitAfterFirstKey returns a split point that keeps exactly the elements of
+// the first key of the bundle in the primary.
+func splitAfterFirstKey(t *testing.T, bundleKeys []string) int {
+ t.Helper()
+ i := 0
+ for i < len(bundleKeys) && bundleKeys[i] == bundleKeys[0] {
+ i++
+ }
+ if i == len(bundleKeys) {
+ t.Fatalf("bundle %v has a single key, want at least two",
bundleKeys)
+ }
+ return i
+}
+
+// TestElementManager_SplitStatefulBundle checks that a channel split of a
+// stateful stage's bundle keeps every key held by a primary element assigned
+// to the primary bundle, so no other bundle processes that key concurrently.
+//
+// Prism gives each bundle a copy of the key's state at bundle start and
+// replaces the key's state with the bundle's copy at commit. If two bundles
+// process the same key at the same time, the last commit discards the state
+// changes of the other.
+func TestElementManager_SplitStatefulBundle(t *testing.T) {
+ tests := []struct {
+ name string
+ // keys has one entry per input element, all in the global
window.
+ keys []string
+ // firstResidual picks the split point from the bundle's key
order,
+ // which depends on map iteration order.
+ firstResidual func(t *testing.T, bundleKeys []string) int
+ }{
+ {
+ name: "split inside the only key",
+ keys: []string{"K", "K", "K"},
+ firstResidual: splitInsideKey("K", 1),
+ }, {
+ name: "split inside one of two keys",
+ keys: []string{"A", "B", "B"},
+ firstResidual: splitInsideKey("B", 1),
+ }, {
+ name: "split between keys",
+ keys: []string{"A", "A", "B", "B"},
+ firstResidual: splitAfterFirstKey,
+ },
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ em := NewElementManager(Config{})
+ em.AddStage("impulse", nil, []string{"input"}, nil)
+ em.AddStage("stateful", []string{"input"}, nil, nil)
+ em.StageStateful("stateful", nil)
+ ss := em.stages["stateful"]
+
+ var es []element
+ for i, k := range test.keys {
+ es = append(es, element{
+ window: window.GlobalWindow{},
+ timestamp: mtime.Time(i),
+ pane: typex.NoFiringPane(),
+ elmBytes: []byte(fmt.Sprintf("%s%d",
k, i)),
+ keyBytes: []byte(k),
+ sequence: i,
+ })
+ }
+ em.addPending(ss.AddPending(em, es))
+
+ var nextID int
+ genBundID := func() string {
+ defer func() { nextID++ }()
+ return fmt.Sprintf("inst%03d", nextID)
+ }
+ // startBundle builds the next event time bundle, if
any, and
+ // returns it with the key of each of its elements.
+ startBundle := func() (RunBundle, []string, bool) {
+ bundID, ok, _, _ :=
ss.startEventTimeBundle(mtime.MaxTimestamp, genBundID)
+ if !ok {
+ return RunBundle{}, nil, false
+ }
+ var keys []string
+ for _, e := range ss.inprogress[bundID].es {
+ keys = append(keys, string(e.keyBytes))
+ }
+ return RunBundle{StageID: ss.ID, BundleID:
bundID, Watermark: mtime.MaxTimestamp}, keys, true
+ }
+
+ primary, bundleKeys, ok := startBundle()
+ if !ok {
+ t.Fatal("startEventTimeBundle() built no
bundle, want one")
+ }
+ if got, want := len(bundleKeys), len(test.keys); got !=
want {
+ t.Fatalf("bundle %v has %v elements, want %v",
bundleKeys, got, want)
+ }
+ split := test.firstResidual(t, bundleKeys)
+ wantHeld := set[string]{}
+ for _, k := range bundleKeys[:split] {
+ wantHeld.insert(k)
+ }
+
+ em.ReturnResiduals(primary, split, PColInfo{},
Residuals{})
+
+ checked := set[string]{}
+ for _, k := range bundleKeys {
+ if checked.present(k) {
+ continue
+ }
+ checked.insert(k)
+ want := wantHeld.present(k)
+ if got := ss.inprogressKeys.present(k); got !=
want {
+ t.Errorf("after split of %v at %v: key
%v in inprogressKeys = %v, want %v", bundleKeys, split, k, got, want)
+ }
+ if got :=
ss.inprogressKeysByBundle[primary.BundleID].present(k); got != want {
+ t.Errorf("after split of %v at %v: key
%v in inprogressKeysByBundle[%v] = %v, want %v", bundleKeys, split, k,
primary.BundleID, got, want)
+ }
+ }
+
+ // While the primary runs, only residuals of released
keys may
+ // be scheduled.
+ processed := split
+ for {
+ rb, keys, ok := startBundle()
+ if !ok {
+ break
+ }
+ for _, k := range keys {
+ if wantHeld.present(k) {
+ t.Fatalf("bundle %v has key %v
while primary bundle %v still processes it", rb.BundleID, k, primary.BundleID)
+ }
+ }
+ processed += len(keys)
+ em.PersistBundle(rb, nil, TentativeData{},
PColInfo{}, Residuals{})
+ }
+
+ // Once the primary commits, the residuals of its keys
run.
+ em.PersistBundle(primary, nil, TentativeData{},
PColInfo{}, Residuals{})
+ for {
+ rb, keys, ok := startBundle()
+ if !ok {
+ break
+ }
+ processed += len(keys)
+ em.PersistBundle(rb, nil, TentativeData{},
PColInfo{}, Residuals{})
+ }
+
+ if got, want := processed, len(test.keys); got != want {
+ t.Errorf("processed %v elements, want %v", got,
want)
+ }
+ if got := em.livePending.Load(); got != 0 {
+ t.Errorf("live pending = %v after all bundles
persisted, want 0", got)
+ }
+ })
+ }
+}
