prodriguezdefino opened a new pull request, #41074:
URL: https://github.com/apache/beam/pull/41074
Problem:
In TentativeData.ClearMultimapState
(sdks/go/pkg/beam/runners/prism/internal/engine/data.go), clearing one map
entry sets a nil slice on userMap.Multimap instead of deleting the key:
userMap.Multimap[string(mapKey)] = nil
TentativeData.GetMultimapKeysState iterates over userMap.Multimap without
skipping nil or empty slices:
for k := range userMap.Multimap {
keys = append(keys, []byte(k))
}
StateForBundle
(sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go) also copies
nil entries across bundles. When an SDK removes a single entry from map, set,
or multimap state and reads MultimapKeysUserState in a later bundle (or without
in-bundle transaction filtering), Prism still returns the removed key even
though GetMultimapState returns empty data for it. Go's existing
TestMapStateClear runs all elements in one bundle, where the Go harness filters
Remove transactions in memory before calling Clear() on the whole map.
Fix:
Delete string(mapKey) from userMap.Multimap in ClearMultimapState, and skip
empty slices in GetMultimapKeysState and StateForBundle. Add unit tests in
data_test.go and a multi-bundle map-state removal test in engine_test.go.
Reproduction (Python SDK on PrismRunner):
Python's userstate module has no MapStateSpec or MultimapStateSpec, so this
script sends MultimapUserState and MultimapKeysUserState requests through the
worker state handler.
```python
import apache_beam as beam
from apache_beam.coders import StrUtf8Coder
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.portability.api import beam_fn_api_pb2
from apache_beam.testing.util import assert_that, equal_to
from apache_beam.transforms.userstate import BagStateSpec
class RemoveMultimapKeyDoFn(beam.DoFn):
ANCHOR = BagStateSpec("anchor", StrUtf8Coder())
def process(self, element, anchor=beam.DoFn.StateParam(ANCHOR)):
_, op = element
base = anchor._state_key.bag_user_state
handler = anchor._state_handler
coder = StrUtf8Coder()
def mm_key(mk):
return beam_fn_api_pb2.StateKey(
multimap_user_state=beam_fn_api_pb2.StateKey.MultimapUserState(
transform_id=base.transform_id,
user_state_id="mm",
window=base.window,
key=base.key,
map_key=coder.encode_nested(mk),
)
)
keys_key = beam_fn_api_pb2.StateKey(
multimap_keys_user_state=beam_fn_api_pb2.StateKey.MultimapKeysUserState(
transform_id=base.transform_id,
user_state_id="mm",
window=base.window,
key=base.key,
)
)
if op == "PUT":
handler.extend(mm_key("a"), coder.get_impl(), ["1"]).get()
handler.extend(mm_key("b"), coder.get_impl(), ["2"]).get()
elif op == "REMOVE_A":
handler.clear(mm_key("a")).get()
yield sorted(handler.blocking_get(keys_key, coder.get_impl()))
if __name__ == "__main__":
options = PipelineOptions(["--runner=PrismRunner",
"--environment_type=LOOPBACK"])
with beam.Pipeline(options=options) as p:
out = (
p
| beam.Create([("k", "PUT"), ("k", "REMOVE_A")])
| beam.ParDo(RemoveMultimapKeyDoFn())
)
assert_that(out, equal_to([["b"]]))
```
Before the fix: Prism returns `[['a', 'b']]`, and `assert_that` fails.
After the fix: Prism returns `[['b']]`.
<!--
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]
