This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch release-2.4
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/release-2.4 by this push:
new ca602e4c1fc [FLINK-34625][python] Fix state TTL read cache disabling
(#29576)
ca602e4c1fc is described below
commit ca602e4c1fca5d3dfc920b13555fd43aa558a257
Author: wuyifan <[email protected]>
AuthorDate: Mon Oct 26 15:40:16 2026 +0800
[FLINK-34625][python] Fix state TTL read cache disabling (#29576)
Use per-state uncached handlers for TTL policies that require backend reads,
preserving other states' shared caches and the existing write-cache policy.
Merge pending MapState mutations in empty checks without flushing active
iterators' caches. Query backend emptiness without TTL cleanup or timestamp
refresh, excluding pending deleted keys and retaining TTL visibility.
Restore the request's key context without record-triggered cleanup while
preserving storage key prefixes across the supported backends.
Cover mixed expiration, paginated iterator continuation, pending mutations,
shared-cache isolation and backend key switching with regression tests.
---
flink-python/pyflink/fn_execution/state_impl.py | 84 +++-
.../pyflink/fn_execution/tests/test_state_impl.py | 436 +++++++++++++++++++++
.../python/beam/state/BeamKeyedStateStore.java | 17 +-
.../python/beam/state/BeamMapStateHandler.java | 44 +++
.../streaming/api/utils/PythonOperatorUtils.java | 12 +-
.../python/beam/state/BeamMapStateHandlerTest.java | 425 ++++++++++++++++++++
.../runtime/state/AbstractKeyedStateBackend.java | 7 +
.../flink/runtime/state/KeyedStateBackend.java | 9 +
.../runtime/state/internal/InternalMapState.java | 24 +-
.../state/metrics/MetricsTrackingMapState.java | 17 +
.../flink/runtime/state/ttl/TtlMapState.java | 19 +
.../flink/runtime/state/StateBackendTestBase.java | 31 ++
.../changelog/ChangelogKeyedStateBackend.java | 5 +
.../restore/ChangelogMigrationRestoreTarget.java | 5 +
.../forst/sync/ForStSyncKeyedStateBackend.java | 6 +
.../state/rocksdb/RocksDBKeyedStateBackend.java | 6 +
16 files changed, 1122 insertions(+), 25 deletions(-)
diff --git a/flink-python/pyflink/fn_execution/state_impl.py
b/flink-python/pyflink/fn_execution/state_impl.py
index 257c6e2c118..5e79f50f97b 100644
--- a/flink-python/pyflink/fn_execution/state_impl.py
+++ b/flink-python/pyflink/fn_execution/state_impl.py
@@ -21,6 +21,8 @@ from abc import ABC, abstractmethod
from apache_beam.coders import coder_impl
from apache_beam.portability.api import beam_fn_api_pb2
from apache_beam.runners.worker.bundle_processor import
SynchronousBagRuntimeState
+from apache_beam.runners.worker.sdk_worker import GlobalCachingStateHandler
+from apache_beam.runners.worker.statecache import StateCache
from apache_beam.transforms import userstate
from enum import Enum
from functools import partial
@@ -120,10 +122,6 @@ class SynchronousKvRuntimeState(InternalKvState, ABC):
elif ttl_config.get_update_type() ==
StateTtlConfig.UpdateType.OnReadAndWrite:
self._cache_type =
SynchronousKvRuntimeState.CacheType.ENABLE_WRITE_CACHE
- if self._cache_type !=
SynchronousKvRuntimeState.CacheType.ENABLE_READ_WRITE_CACHE:
- # disable read cache
-
self._remote_state_backend._state_handler._state_cache._cache._max_entries = 0
-
@abstractmethod
def get_internal_state(self):
pass
@@ -145,7 +143,7 @@ class
SynchronousBagKvRuntimeState(SynchronousKvRuntimeState, ABC):
def get_internal_state(self):
if self._internal_state is None:
self._internal_state =
self._remote_state_backend._get_internal_bag_state(
- self.name, self.namespace, self._value_coder, self._ttl_config)
+ self.name, self.namespace, self._value_coder,
self._ttl_config, self._cache_type)
return self._internal_state
def _maybe_clear_write_cache(self):
@@ -365,6 +363,7 @@ class CachingMapStateHandler(object):
GET_FLAG = 0
ITERATE_FLAG = 1
CHECK_EMPTY_FLAG = 2
+ CHECK_EMPTY_EXCLUDING_KEYS_FLAG = 3
# GET response flags
EXIST_FLAG = 0
IS_NONE_FLAG = 1
@@ -555,6 +554,18 @@ class CachingMapStateHandler(object):
return False
return self._check_empty_raw(state_key)
+ def check_empty_excluding(self, state_key, excluded_keys, map_key_encoder):
+ output_stream = coder_impl.create_OutputStream()
+ output_stream.write_byte(self.CHECK_EMPTY_EXCLUDING_KEYS_FLAG)
+ output_stream.write_bigendian_int32(len(excluded_keys))
+ for key in excluded_keys:
+ key_stream = coder_impl.create_OutputStream()
+ map_key_encoder(key, key_stream)
+ encoded_key = key_stream.get()
+ output_stream.write_bigendian_int32(len(encoded_key))
+ output_stream.write(encoded_key)
+ return self._check_empty_raw(state_key, output_stream.get())
+
def clear(self, state_key):
self.clear_read_cache(state_key)
return self._underlying.clear(state_key)
@@ -577,10 +588,11 @@ class CachingMapStateHandler(object):
def reset_cached_iterators_num(self):
self._cached_iterator_num = 0
- def _check_empty_raw(self, state_key):
- output_stream = coder_impl.create_OutputStream()
- output_stream.write_byte(self.CHECK_EMPTY_FLAG)
- continuation_token = output_stream.get()
+ def _check_empty_raw(self, state_key, continuation_token=None):
+ if continuation_token is None:
+ output_stream = coder_impl.create_OutputStream()
+ output_stream.write_byte(self.CHECK_EMPTY_FLAG)
+ continuation_token = output_stream.get()
data, response_token = self._underlying.get_raw(state_key,
continuation_token)
if data[0] == self.IS_EMPTY_FLAG:
return True
@@ -736,8 +748,10 @@ class InternalSynchronousMapRuntimeState(object):
state_key,
map_key_coder,
map_value_coder,
- max_write_cache_entries):
+ max_write_cache_entries,
+ read_cache_enabled=True):
self._map_state_handler = map_state_handler
+ self._read_cache_enabled = read_cache_enabled
self._state_key = state_key
self._map_key_coder = map_key_coder
if isinstance(map_key_coder, FieldCoder):
@@ -815,9 +829,22 @@ class InternalSynchronousMapRuntimeState(object):
return True
def is_empty(self):
- if self._is_empty is None:
- if len(self._write_cache) > 0:
- self.commit()
+ if self._is_empty:
+ return True
+ if self._write_cache:
+ if any(exists for exists, _ in self._write_cache.values()):
+ self._is_empty = False
+ elif self._cleared:
+ self._is_empty = True
+ else:
+ self._is_empty = self._map_state_handler.check_empty_excluding(
+ self._state_key, self._write_cache, self._map_key_encoder)
+ elif self._cleared:
+ self._is_empty = True
+ elif not self._read_cache_enabled:
+ self._is_empty = self._map_state_handler.check_empty_excluding(
+ self._state_key, (), self._map_key_encoder)
+ elif self._is_empty is None:
self._is_empty =
self._map_state_handler.check_empty(self._state_key)
return self._is_empty
@@ -960,6 +987,11 @@ class RemoteKeyedStateBackend(object):
self._state_handler = state_handler
self._map_state_handler = CachingMapStateHandler(
state_handler, map_state_read_cache_size)
+ # TTL reads must reach the state backend without disabling other
states' caches.
+ self._uncached_state_handler = GlobalCachingStateHandler(
+ StateCache(0), state_handler._underlying)
+ self._uncached_map_state_handler = CachingMapStateHandler(
+ self._uncached_state_handler, map_state_read_cache_size)
self._key_coder_impl = key_coder.get_impl()
self.namespace_coder = namespace_coder
if namespace_coder:
@@ -1052,7 +1084,7 @@ class RemoteKeyedStateBackend(object):
self._all_states[name] = wrapped_state
return wrapped_state
- def _get_internal_bag_state(self, name, namespace, element_coder,
ttl_config):
+ def _get_internal_bag_state(self, name, namespace, element_coder,
ttl_config, cache_type):
encoded_namespace = self._encode_namespace(namespace)
cached_state = self._internal_state_cache.get(
(name, self._encoded_current_key, encoded_namespace))
@@ -1065,7 +1097,8 @@ class RemoteKeyedStateBackend(object):
if isinstance(element_coder, FieldCoder):
element_coder = FlinkCoder(element_coder)
state_spec = userstate.BagStateSpec(name, element_coder)
- internal_state = self._create_bag_state(state_spec, encoded_namespace,
ttl_config)
+ internal_state = self._create_bag_state(
+ state_spec, encoded_namespace, ttl_config, cache_type)
return internal_state
def _get_internal_map_state(
@@ -1079,11 +1112,15 @@ class RemoteKeyedStateBackend(object):
name, encoded_namespace, map_key_coder, map_value_coder,
ttl_config, cache_type)
return internal_map_state
- def _create_bag_state(self, state_spec: userstate.StateSpec,
encoded_namespace, ttl_config) \
+ def _create_bag_state(
+ self, state_spec: userstate.StateSpec, encoded_namespace,
ttl_config, cache_type) \
-> userstate.AccumulatingRuntimeState:
if isinstance(state_spec, userstate.BagStateSpec):
+ state_handler = self._state_handler
+ if cache_type !=
SynchronousKvRuntimeState.CacheType.ENABLE_READ_WRITE_CACHE:
+ state_handler = self._uncached_state_handler
bag_state = SynchronousBagRuntimeState(
- self._state_handler,
+ state_handler,
state_key=self.get_bag_state_key(
state_spec.name, self._encoded_current_key,
encoded_namespace, ttl_config),
value_coder=state_spec.coder)
@@ -1109,12 +1146,18 @@ class RemoteKeyedStateBackend(object):
write_cache_size = 0
else:
write_cache_size = self._map_state_write_cache_size
+ read_cache_enabled = \
+ cache_type ==
SynchronousKvRuntimeState.CacheType.ENABLE_READ_WRITE_CACHE
+ map_state_handler = self._map_state_handler
+ if not read_cache_enabled:
+ map_state_handler = self._uncached_map_state_handler
return InternalSynchronousMapRuntimeState(
- self._map_state_handler,
+ map_state_handler,
state_key,
map_key_coder,
map_value_coder,
- write_cache_size)
+ write_cache_size,
+ read_cache_enabled)
def _encode_namespace(self, namespace):
if namespace is not None:
@@ -1154,7 +1197,8 @@ class RemoteKeyedStateBackend(object):
self.commit_internal_state(state._internal_state)
def clear_cached_iterators(self):
- if self._map_state_handler.get_cached_iterators_num() > 0:
+ if self._map_state_handler.get_cached_iterators_num() > 0 or \
+ self._uncached_map_state_handler.get_cached_iterators_num() >
0:
self._clear_iterator_mark.multimap_side_input.key =
self._encoded_current_key
self._map_state_handler.clear(self._clear_iterator_mark)
diff --git a/flink-python/pyflink/fn_execution/tests/test_state_impl.py
b/flink-python/pyflink/fn_execution/tests/test_state_impl.py
new file mode 100644
index 00000000000..73a3f2ada5b
--- /dev/null
+++ b/flink-python/pyflink/fn_execution/tests/test_state_impl.py
@@ -0,0 +1,436 @@
+################################################################################
+# 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.
+################################################################################
+
+import contextlib
+import unittest
+from unittest import mock
+
+from apache_beam.coders import VarIntCoder, StrUtf8Coder
+from apache_beam.runners.worker.sdk_worker import GlobalCachingStateHandler,
StateHandler
+from apache_beam.runners.worker.statecache import StateCache
+
+from pyflink.common import Time, Types
+from pyflink.datastream.state import StateTtlConfig, MapStateDescriptor
+from pyflink.fn_execution.state_impl import RemoteKeyedStateBackend,
RemoteOperatorStateBackend, \
+ CachingMapStateHandler
+
+
+class StateTtlReadCacheTests(unittest.TestCase):
+
+ @staticmethod
+ def _ttl_configs():
+ for update_type in (StateTtlConfig.UpdateType.OnCreateAndWrite,
+ StateTtlConfig.UpdateType.OnReadAndWrite):
+ yield StateTtlConfig.new_builder(Time.seconds(5)) \
+ .set_update_type(update_type) \
+
.set_state_visibility(StateTtlConfig.StateVisibility.NeverReturnExpired) \
+ .build()
+ yield StateTtlConfig.new_builder(Time.seconds(5)) \
+ .update_ttl_on_read_and_write() \
+
.set_state_visibility(StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp)
\
+ .build()
+
+ @staticmethod
+ def _create_backend(cache_size=1 << 20):
+ # Only remote I/O is mocked. Both Beam caches and Flink runtime states
are real.
+ remote = mock.Mock(spec=StateHandler)
+ remote.process_instruction_id.return_value = contextlib.nullcontext()
+ remote.append_raw.return_value = None
+ remote.clear.return_value = None
+ handler = GlobalCachingStateHandler(StateCache(cache_size), remote)
+ backend = RemoteKeyedStateBackend(
+ handler, StrUtf8Coder(), StrUtf8Coder(), 16, 16, 16)
+ backend.set_current_key('key')
+ return backend, handler, remote
+
+ @staticmethod
+ def _encode(value):
+ return VarIntCoder().get_impl().encode_nested(value)
+
+ @classmethod
+ def _map_entry(cls, key, value):
+ return cls._encode(key) + b'\x01' + cls._encode(value)
+
+ @classmethod
+ def _check_empty_excluding_request(cls, *keys):
+ encoded_keys = [cls._encode(key) for key in keys]
+ return b'\x03' + len(keys).to_bytes(4, 'big') + b''.join(
+ len(key).to_bytes(4, 'big') + key for key in encoded_keys)
+
+ def
test_value_ttl_bypasses_read_cache_without_invalidating_other_states(self):
+ for ttl in self._ttl_configs():
+ with self.subTest(ttl=ttl):
+ backend, handler, remote = self._create_backend()
+ cache = handler._state_cache
+ remote.get_raw.side_effect = [
+ (self._encode(42), b''), (self._encode(100), b''), (b'',
b'')]
+ with handler.process_instruction_id('bundle', []):
+ cached = backend.get_value_state('cached', VarIntCoder())
+ self.assertEqual(cached.value(), 42)
+ state = backend.get_value_state('ttl', VarIntCoder(), ttl)
+ self.assertEqual(state.value(), 100)
+ # An empty backend response represents expiration, without
a timing dependency.
+ self.assertIsNone(state.value())
+ self.assertEqual(cached.value(), 42)
+ self.assertEqual(remote.get_raw.call_count, 3)
+ self.assertIs(backend._state_handler, handler)
+ self.assertIs(handler._state_cache, cache)
+ self.assertTrue(cache.is_cache_enabled())
+ self.assertIs(cached.get_internal_state()._state_handler,
handler)
+ self.assertIsNot(state.get_internal_state()._state_handler,
handler)
+
+ def test_cache_permitting_states_keep_read_cache(self):
+ ttl = StateTtlConfig.new_builder(Time.seconds(5)) \
+
.set_state_visibility(StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp)
\
+ .build()
+ for config in (None, ttl):
+ with self.subTest(ttl=config):
+ backend, handler, remote = self._create_backend()
+ remote.get_raw.side_effect = [(self._encode(100), b'')]
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_value_state('value', VarIntCoder(),
config)
+ self.assertEqual(state.value(), 100)
+ self.assertEqual(state.value(), 100)
+ remote.get_raw.assert_called_once()
+ self.assertIs(state.get_internal_state()._state_handler,
handler)
+
+ def test_ttl_write_cache_policy_is_preserved(self):
+ for ttl in self._ttl_configs():
+ with self.subTest(ttl=ttl):
+ backend, handler, remote = self._create_backend()
+ with handler.process_instruction_id('bundle', []):
+ value = backend.get_value_state('value', VarIntCoder(),
ttl)
+ map_state = backend.get_map_state('map', VarIntCoder(),
VarIntCoder(), ttl)
+ value.update(100)
+ map_state.put(1, 100)
+ if ttl.get_state_visibility() == \
+ StateTtlConfig.StateVisibility.NeverReturnExpired:
+ self.assertEqual(remote.append_raw.call_count, 2)
+ else:
+ remote.append_raw.assert_not_called()
+ backend.commit()
+ self.assertEqual(remote.append_raw.call_count, 2)
+
+ def
test_map_ttl_bypasses_read_cache_without_invalidating_other_states(self):
+ for ttl in self._ttl_configs():
+ with self.subTest(ttl=ttl):
+ backend, handler, remote = self._create_backend()
+ exists = bytes([CachingMapStateHandler.EXIST_FLAG])
+ missing = bytes([CachingMapStateHandler.NOT_EXIST_FLAG])
+ remote.get_raw.side_effect = [
+ (exists + self._encode(42), b''),
+ (exists + self._encode(100), b''), (missing, b'')]
+ with handler.process_instruction_id('bundle', []):
+ cached = backend.get_map_state('cached', VarIntCoder(),
VarIntCoder())
+ self.assertEqual(cached.get(1), 42)
+ state = backend.get_map_state('ttl', VarIntCoder(),
VarIntCoder(), ttl)
+ self.assertEqual(state.get(1), 100)
+ self.assertIsNone(state.get(1))
+ self.assertEqual(cached.get(1), 42)
+ self.assertEqual(remote.get_raw.call_count, 3)
+
self.assertIs(cached.get_internal_state()._map_state_handler._state_cache,
+ handler._state_cache)
+ self.assertFalse(state.get_internal_state()._map_state_handler
+ ._state_cache.is_cache_enabled())
+
+ def test_map_ttl_is_empty_after_committed_value_expires(self):
+ for ttl in self._ttl_configs():
+ with self.subTest(ttl=ttl):
+ backend, handler, remote = self._create_backend()
+ remote.get_raw.side_effect = [
+ (bytes([CachingMapStateHandler.NOT_EXIST_FLAG]), b''),
+ (bytes([CachingMapStateHandler.IS_EMPTY_FLAG]), b'')]
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('ttl', VarIntCoder(),
VarIntCoder(), ttl)
+ state.put(1, 100)
+ backend.commit()
+ self.assertIsNone(state.get(1))
+ self.assertTrue(state.is_empty())
+ self.assertEqual(remote.get_raw.call_count, 2)
+
+ def test_map_ttl_is_empty_refreshes_within_bundle(self):
+ for ttl in self._ttl_configs():
+ with self.subTest(ttl=ttl):
+ backend, handler, remote = self._create_backend()
+ responses = [
+ (bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b''),
+ (bytes([CachingMapStateHandler.IS_EMPTY_FLAG]), b'')]
+ writes_are_uncached = ttl.get_state_visibility() == \
+ StateTtlConfig.StateVisibility.NeverReturnExpired
+ if writes_are_uncached:
+
responses.append((bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b''))
+ remote.get_raw.side_effect = responses
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('ttl', VarIntCoder(),
VarIntCoder(), ttl)
+ mod_count = state.get_internal_state()._mod_count
+ self.assertFalse(state.is_empty())
+ self.assertTrue(state.is_empty())
+ self.assertEqual(state.get_internal_state()._mod_count,
mod_count)
+ self.assertTrue(state.is_empty())
+ self.assertEqual(remote.get_raw.call_count, 2)
+ for call in remote.get_raw.call_args_list:
+ self.assertEqual(call.args[1],
self._check_empty_excluding_request())
+ state.put(1, 100)
+ self.assertFalse(state.is_empty())
+ self.assertEqual(remote.get_raw.call_count, 3 if
writes_are_uncached else 2)
+
+ def test_map_is_empty_merges_pending_mutations(self):
+ ttl = list(self._ttl_configs())[-1]
+ for config in (None, ttl):
+ for operation in ('put', 'remove', 'clear_then_put', 'clear'):
+ with self.subTest(ttl=config, operation=operation):
+ backend, handler, remote = self._create_backend()
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('map', VarIntCoder(),
VarIntCoder(), config)
+ if operation == 'remove':
+ state.put(1, 100)
+ backend.commit()
+ remote.reset_mock()
+ state.remove(1)
+ remote.get_raw.return_value = (
+ bytes([CachingMapStateHandler.IS_EMPTY_FLAG]),
b'')
+ elif operation == 'clear':
+ state.clear()
+ else:
+ if operation == 'clear_then_put':
+ state.clear()
+ state.put(1, 100)
+ self.assertEqual(state.is_empty(), operation in
('remove', 'clear'))
+ remote.append_raw.assert_not_called()
+ remote.clear.assert_not_called()
+ if operation == 'remove':
+ remote.get_raw.assert_called_once()
+ self.assertEqual(remote.get_raw.call_args.args[1],
+
self._check_empty_excluding_request(1))
+ else:
+ remote.get_raw.assert_not_called()
+
+ def test_map_ttl_is_empty_preserves_buffered_iterators(self):
+ ttl = list(self._ttl_configs())[-1]
+ expected = {'items': [(1, 100), (2, 200)], 'keys': [1, 2], 'values':
[100, 200]}
+ for method in expected:
+ for cleared, remove in ((False, False), (False, True), (True,
False)):
+ with self.subTest(method=method, cleared=cleared,
remove=remove):
+ backend, handler, remote = self._create_backend()
+ remote.get_raw.side_effect = lambda _, request: (
+ bytes([CachingMapStateHandler.IS_EMPTY_FLAG]) if
request[0] == 3 else b'',
+ b'')
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('ttl', VarIntCoder(),
VarIntCoder(), ttl)
+ if cleared:
+ state.clear()
+ state.put(1, 100)
+ state.put(2, 200)
+ iterator = getattr(state, method)()
+ self.assertEqual(next(iterator), expected[method][0])
+ if remove:
+ iterator.remove()
+ self.assertFalse(state.is_empty())
+ self.assertEqual(next(iterator), expected[method][1])
+ if remove:
+ iterator.remove()
+ self.assertTrue(state.is_empty())
+ self.assertEqual(list(iterator), [])
+ remote.append_raw.assert_not_called()
+ remote.clear.assert_not_called()
+
+ def test_map_ttl_iterator_still_rejects_external_mutations(self):
+ ttl = list(self._ttl_configs())[-1]
+ for method in ('items', 'keys', 'values'):
+ for operation in ('put', 'remove', 'clear'):
+ with self.subTest(method=method, operation=operation):
+ backend, handler, remote = self._create_backend()
+ remote.get_raw.return_value = (b'', b'')
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('ttl', VarIntCoder(),
VarIntCoder(), ttl)
+ state.put(1, 100)
+ state.put(2, 200)
+ iterator = getattr(state, method)()
+ next(iterator)
+ if operation == 'put':
+ state.put(3, 300)
+ elif operation == 'remove':
+ state.remove(2)
+ else:
+ state.clear()
+ with self.assertRaisesRegex(Exception, 'Concurrent
modification detected'):
+ next(iterator)
+
+ def test_map_ttl_is_empty_excludes_pending_deletions(self):
+ ttl = list(self._ttl_configs())[-1]
+ for backend_keys in ((1,), (1, 2)):
+ with self.subTest(backend_keys=backend_keys):
+ backend, handler, remote = self._create_backend()
+ response = CachingMapStateHandler.IS_EMPTY_FLAG if
backend_keys == (1,) else \
+ CachingMapStateHandler.NOT_EMPTY_FLAG
+ remote.get_raw.return_value = (bytes([response]), b'')
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('ttl', VarIntCoder(),
VarIntCoder(), ttl)
+ state.remove(1)
+ self.assertEqual(state.is_empty(), backend_keys == (1,))
+ remote.append_raw.assert_not_called()
+ remote.get_raw.assert_called_once()
+ self.assertEqual(remote.get_raw.call_args.args[1],
+ self._check_empty_excluding_request(1))
+
+ def test_map_ttl_is_empty_preserves_paged_remote_iterators(self):
+ ttl = list(self._ttl_configs())[-1]
+ expected = {'items': [(1, 100), (2, 200)], 'keys': [1, 2], 'values':
[100, 200]}
+ for method in expected:
+ with self.subTest(method=method):
+ backend, handler, remote = self._create_backend()
+ encode_page = self._encode if method == 'keys' else \
+ lambda key: self._map_entry(key, key * 100)
+ remote.get_raw.side_effect = [
+ (encode_page(1), b'original-next'),
+ (bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b''),
+ (encode_page(2), b'')]
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('ttl', VarIntCoder(),
VarIntCoder(), ttl)
+ iterator = getattr(state, method)()
+ self.assertEqual(next(iterator), expected[method][0])
+ iterator.remove()
+ self.assertFalse(state.is_empty())
+ self.assertEqual(next(iterator), expected[method][1])
+ self.assertEqual(list(iterator), [])
+ self.assertEqual(
+
backend._uncached_map_state_handler.get_cached_iterators_num(), 0)
+ self.assertEqual(remote.get_raw.call_count, 3)
+ self.assertEqual(remote.get_raw.call_args_list[1].args[1],
+ self._check_empty_excluding_request(1))
+ self.assertIn(b'original-next',
remote.get_raw.call_args_list[2].args[1])
+ remote.append_raw.assert_not_called()
+
+ def test_map_ttl_is_empty_does_not_start_another_ttl_iterator(self):
+ ttl = list(self._ttl_configs())[-1]
+ backend, handler, remote = self._create_backend()
+ remote.get_raw.return_value = (self._map_entry(1, 100), b'active')
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(),
ttl)
+ iterator = state.items()
+ self.assertEqual(next(iterator), (1, 100))
+ iterator.remove()
+ internal = state.get_internal_state()
+ pending = dict(internal._write_cache)
+ mod_count = internal._mod_count
+
+ def check_empty(_, request):
+ # Iterating mixed expired/live backend entries could
invalidate the active token.
+ self.assertEqual(request,
self._check_empty_excluding_request(1))
+ return bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b''
+
+ remote.get_raw.side_effect = check_empty
+ self.assertFalse(state.is_empty())
+ self.assertEqual(internal._write_cache, pending)
+ self.assertEqual(internal._mod_count, mod_count)
+
self.assertEqual(backend._uncached_map_state_handler.get_cached_iterators_num(),
1)
+ remote.append_raw.assert_not_called()
+ remote.clear.assert_not_called()
+ remote.get_raw.side_effect = None
+ remote.get_raw.return_value = (self._map_entry(2, 200), b'')
+ self.assertEqual(list(iterator), [(2, 200)])
+
self.assertEqual(backend._uncached_map_state_handler.get_cached_iterators_num(),
0)
+
+ def test_cache_permitting_map_states_keep_is_empty_cache(self):
+ ttl = StateTtlConfig.new_builder(Time.seconds(5)) \
+
.set_state_visibility(StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp)
\
+ .build()
+ for config in (None, ttl):
+ with self.subTest(ttl=config):
+ backend, handler, remote = self._create_backend()
+ remote.get_raw.side_effect = [
+ (bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b'')]
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('map', VarIntCoder(),
VarIntCoder(), config)
+ self.assertFalse(state.is_empty())
+ self.assertFalse(state.is_empty())
+ remote.get_raw.assert_called_once()
+ self.assertEqual(remote.get_raw.call_args.args[1],
+
bytes([CachingMapStateHandler.CHECK_EMPTY_FLAG]))
+
+ def test_broadcast_map_state_keeps_default_cache_behavior(self):
+ _, handler, remote = self._create_backend()
+ backend = RemoteOperatorStateBackend(handler, 16, 16, 16)
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_broadcast_state(
+ MapStateDescriptor('broadcast', Types.INT(), Types.INT()))
+ state.put(1, 100)
+ backend.commit()
+ self.assertFalse(state._internal_map_state.is_empty())
+ self.assertEqual(state.get(1), 100)
+ remote.get_raw.assert_not_called()
+
+ def test_map_ttl_iteration_and_unfinished_iterator_cleanup(self):
+ ttl = next(self._ttl_configs())
+ backend, handler, remote = self._create_backend()
+ remote.get_raw.side_effect = [
+ (self._map_entry(1, 100), b''), (b'', b''),
+ (self._map_entry(1, 100), b'next-page'),
+ (self._map_entry(2, 200), b''),
+ (self._map_entry(1, 100), b'unfinished')]
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(),
ttl)
+ self.assertEqual(list(state.items()), [(1, 100)])
+ self.assertEqual(list(state.items()), [])
+ self.assertEqual(list(state.items()), [(1, 100), (2, 200)])
+ self.assertIn(b'next-page',
remote.get_raw.call_args_list[3].args[1])
+ self.assertEqual(next(state.items()), (1, 100))
+
self.assertEqual(backend._map_state_handler.get_cached_iterators_num(), 0)
+ backend.clear_cached_iterators()
+ remote.clear.assert_called_once()
+ clear_key = remote.clear.call_args.args[0].multimap_side_input
+ self.assertEqual(clear_key.transform_id, 'clear_iterators')
+ self.assertEqual(clear_key.key, b'key')
+
+ def test_ttl_handler_is_preserved_when_reusing_key_and_namespace(self):
+ ttl = next(self._ttl_configs())
+ backend, handler, remote = self._create_backend()
+ remote.get_raw.side_effect = [
+ (self._encode(100), b''), (self._encode(200), b''), (b'', b'')]
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_value_state('ttl', VarIntCoder(), ttl)
+ state.set_current_namespace('first')
+ self.assertEqual(state.value(), 100)
+ first_internal_state = state.get_internal_state()
+ backend.set_current_key('other')
+ state.set_current_namespace('second')
+ self.assertEqual(state.value(), 200)
+ backend.set_current_key('key')
+ state.set_current_namespace('first')
+ self.assertIs(state.get_internal_state(), first_internal_state)
+ self.assertIsNone(state.value())
+ first_key, second_key, reused_key = [
+ call.args[0] for call in remote.get_raw.call_args_list]
+ self.assertEqual(first_key, reused_key)
+ self.assertNotEqual(first_key, second_key)
+
+ def test_repeated_state_registration_with_disabled_cache(self):
+ ttl = next(self._ttl_configs())
+ backend, handler, remote = self._create_backend(cache_size=0)
+ remote.get_raw.side_effect = [(self._encode(100), b''), (b'', b'')]
+ with handler.process_instruction_id('bundle', []):
+ state = backend.get_value_state('ttl', VarIntCoder(), ttl)
+ self.assertEqual(state.value(), 100)
+ self.assertIs(backend.get_value_state('ttl', VarIntCoder(), ttl),
state)
+ self.assertIsNone(state.value())
+ self.assertEqual(remote.get_raw.call_count, 2)
+ self.assertFalse(handler._state_cache.is_cache_enabled())
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamKeyedStateStore.java
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamKeyedStateStore.java
index 83981df889b..b50242f3367 100644
---
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamKeyedStateStore.java
+++
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamKeyedStateStore.java
@@ -152,7 +152,12 @@ public class BeamKeyedStateStore implements BeamStateStore
{
byte[] keyBytes = mapUserState.getKey().toByteArray();
bais.setBuffer(keyBytes, 0, keyBytes.length);
Object key = keySerializer.deserialize(baisWrapper);
- setCurrentKey(key);
+ boolean readOnly =
+ request.hasGet()
+ && !request.getGet().getContinuationToken().isEmpty()
+ && request.getGet().getContinuationToken().byteAt(0)
+ ==
BeamMapStateHandler.CHECK_EMPTY_EXCLUDING_KEYS_FLAG;
+ setCurrentKey(key, readOnly);
// Retrieve state descriptor
byte[] data =
Base64.getDecoder().decode(mapUserState.getSideInputId());
@@ -198,13 +203,19 @@ public class BeamKeyedStateStore implements
BeamStateStore {
}
private void setCurrentKey(Object key) {
+ setCurrentKey(key, false);
+ }
+
+ private void setCurrentKey(Object key, boolean forStateAccess) {
if (keyedStateBackend.getKeySerializer() instanceof RowDataSerializer)
{
setCurrentKeyForStreaming(
(KeyedStateBackend<? super BinaryRowData>)
keyedStateBackend,
((RowDataSerializer) keyedStateBackend.getKeySerializer())
- .toBinaryRow((RowData) key));
+ .toBinaryRow((RowData) key),
+ forStateAccess);
} else {
- setCurrentKeyForStreaming((KeyedStateBackend<Object>)
keyedStateBackend, key);
+ setCurrentKeyForStreaming(
+ (KeyedStateBackend<Object>) keyedStateBackend, key,
forStateAccess);
}
}
}
diff --git
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandler.java
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandler.java
index 6bc2f7ad8aa..e5075ce67cd 100644
---
a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandler.java
+++
b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandler.java
@@ -25,16 +25,21 @@ import
org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.python.PythonOptions;
+import org.apache.flink.runtime.state.internal.InternalMapState;
import org.apache.flink.streaming.api.utils.ByteArrayWrapper;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
+import java.util.Set;
import java.util.UUID;
+import static org.apache.flink.util.Preconditions.checkArgument;
+
/** BeamMapStateHandler handles operations on a {@link MapState}. */
public class BeamMapStateHandler
extends AbstractBeamStateHandler<MapState<ByteArrayWrapper, byte[]>> {
@@ -46,6 +51,7 @@ public class BeamMapStateHandler
private static final byte ITERATE_FLAG = 1;
private static final byte CHECK_EMPTY_FLAG = 2;
+ static final byte CHECK_EMPTY_EXCLUDING_KEYS_FLAG = 3;
/** Map state GET response flags. */
private static final byte EXIST_FLAG = 0;
@@ -118,6 +124,8 @@ public class BeamMapStateHandler
// [flag (1 byte)][serialized map key]
// The continuation token structure of CHECK_EMPTY request is:
// [flag (1 byte)]
+ // The continuation token structure of CHECK_EMPTY_EXCLUDING_KEYS
request is:
+ // [flag (1 byte)][key count (int32)][key length (int32)][serialized
key]...
// The continuation token structure of ITERATE request is:
// [flag (1 byte)][iterate type (1 byte)][iterator token length
(int32)][iterator token]
byte[] getRequest =
request.getGet().getContinuationToken().toByteArray();
@@ -133,6 +141,22 @@ public class BeamMapStateHandler
case CHECK_EMPTY_FLAG:
response = handleCheckEmpty(mapState);
break;
+ case CHECK_EMPTY_EXCLUDING_KEYS_FLAG:
+ bais.setBuffer(getRequest, 1, getRequest.length - 1);
+ final int keyCount = baisWrapper.readInt();
+ checkArgument(keyCount >= 0, "The excluded key count must not
be negative.");
+ final Set<ByteArrayWrapper> excludedKeys = new HashSet<>();
+ for (int i = 0; i < keyCount; i++) {
+ final int keyLength = baisWrapper.readInt();
+ checkArgument(
+ keyLength >= 0 && keyLength <= bais.available(),
+ "Invalid excluded key length.");
+ final byte[] key = new byte[keyLength];
+ baisWrapper.readFully(key);
+ excludedKeys.add(new ByteArrayWrapper(key));
+ }
+ response = handleCheckEmptyExcludingKeys(mapState,
excludedKeys);
+ break;
case ITERATE_FLAG:
bais.setBuffer(getRequest, 1, getRequest.length - 1);
IterateType iterateType =
IterateType.fromOrd(baisWrapper.readByte());
@@ -245,6 +269,26 @@ public class BeamMapStateHandler
}
}
+ private BeamFnApi.StateGetResponse.Builder handleCheckEmptyExcludingKeys(
+ MapState<ByteArrayWrapper, byte[]> mapState, Set<ByteArrayWrapper>
excludedKeys)
+ throws Exception {
+ if (mapState instanceof InternalMapState<?, ?, ?, ?>) {
+ return ((InternalMapState<?, ?, ?, ?>)
mapState).isEmptyExcluding(excludedKeys)
+ ? IS_EMPTY_RESPONSE
+ : NOT_EMPTY_RESPONSE;
+ }
+ // Operator map state has no TTL decorator.
+ final Iterable<ByteArrayWrapper> keys = mapState.keys();
+ if (keys != null) {
+ for (ByteArrayWrapper key : keys) {
+ if (!excludedKeys.contains(key)) {
+ return NOT_EMPTY_RESPONSE;
+ }
+ }
+ }
+ return IS_EMPTY_RESPONSE;
+ }
+
private BeamFnApi.StateGetResponse.Builder handleIterate(
MapState<ByteArrayWrapper, byte[]> mapState,
IterateType iterateType,
diff --git
a/flink-python/src/main/java/org/apache/flink/streaming/api/utils/PythonOperatorUtils.java
b/flink-python/src/main/java/org/apache/flink/streaming/api/utils/PythonOperatorUtils.java
index dc26f452ab5..7d657d2969f 100644
---
a/flink-python/src/main/java/org/apache/flink/streaming/api/utils/PythonOperatorUtils.java
+++
b/flink-python/src/main/java/org/apache/flink/streaming/api/utils/PythonOperatorUtils.java
@@ -31,8 +31,18 @@ public class PythonOperatorUtils {
/** Set the current key for streaming operator. */
public static <K> void setCurrentKeyForStreaming(
KeyedStateBackend<K> stateBackend, K currentKey) {
+ setCurrentKeyForStreaming(stateBackend, currentKey, false);
+ }
+
+ /** Selects a state-request key without treating a read-only query as
another input record. */
+ public static <K> void setCurrentKeyForStreaming(
+ KeyedStateBackend<K> stateBackend, K currentKey, boolean
forStateAccess) {
if (!inBatchExecutionMode(stateBackend)) {
- stateBackend.setCurrentKey(currentKey);
+ if (forStateAccess) {
+ stateBackend.setCurrentKeyForStateAccess(currentKey);
+ } else {
+ stateBackend.setCurrentKey(currentKey);
+ }
}
}
diff --git
a/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandlerTest.java
b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandlerTest.java
new file mode 100644
index 00000000000..b75260c2d7e
--- /dev/null
+++
b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandlerTest.java
@@ -0,0 +1,425 @@
+/*
+ * 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 org.apache.flink.streaming.api.runners.python.beam.state;
+
+import org.apache.flink.api.common.state.MapState;
+import org.apache.flink.api.common.state.MapStateDescriptor;
+import org.apache.flink.api.common.state.StateTtlConfig;
+import org.apache.flink.api.common.typeinfo.Types;
+import org.apache.flink.api.common.typeutils.base.IntSerializer;
+import
org.apache.flink.api.common.typeutils.base.array.BytePrimitiveArraySerializer;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.StateLatencyTrackOptions;
+import org.apache.flink.fnexecution.v1.FlinkFnApi.StateDescriptor;
+import
org.apache.flink.fnexecution.v1.FlinkFnApi.StateDescriptor.StateTTLConfig;
+import
org.apache.flink.fnexecution.v1.FlinkFnApi.StateDescriptor.StateTTLConfig.CleanupStrategies;
+import org.apache.flink.python.PythonOptions;
+import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
+import org.apache.flink.runtime.state.VoidNamespace;
+import org.apache.flink.runtime.state.VoidNamespaceSerializer;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.api.operators.StreamMap;
+import org.apache.flink.streaming.api.utils.ByteArrayWrapper;
+import org.apache.flink.streaming.api.utils.ByteArrayWrapperSerializer;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests the map-state protocol against a real heap backend and a
controllable TTL clock. */
+class BeamMapStateHandlerTest {
+
+ @ParameterizedTest
+ @CsvSource({
+ "ITEMS, false, ReturnExpiredIfNotCleanedUp",
+ "KEYS, false, ReturnExpiredIfNotCleanedUp",
+ "VALUES, false, ReturnExpiredIfNotCleanedUp",
+ "ITEMS, true, ReturnExpiredIfNotCleanedUp",
+ "KEYS, true, ReturnExpiredIfNotCleanedUp",
+ "VALUES, true, ReturnExpiredIfNotCleanedUp",
+ "ITEMS, false, NeverReturnExpired",
+ "KEYS, false, NeverReturnExpired",
+ "VALUES, false, NeverReturnExpired",
+ "ITEMS, true, NeverReturnExpired",
+ "KEYS, true, NeverReturnExpired",
+ "VALUES, true, NeverReturnExpired"
+ })
+ void testMixedExpiryDoesNotInvalidateActiveIterator(
+ IterateType type, boolean metrics, StateTtlConfig.StateVisibility
visibility)
+ throws Exception {
+ try (Fixture fixture = new Fixture(metrics, visibility)) {
+ for (int key = 0; key < 64; key++) {
+ fixture.state.put(key(key), new byte[] {(byte) key});
+ }
+ List<Integer> order = new ArrayList<>();
+ for (ByteArrayWrapper key : fixture.state.keys()) {
+ order.add(Byte.toUnsignedInt(key.get(0)));
+ }
+ fixture.harness.setStateTtlProcessingTime(101);
+ for (int key : order.subList(0, 8)) {
+ fixture.state.put(key(key), new byte[] {(byte) key});
+ }
+
+ BeamFnApi.StateGetResponse page = iterate(fixture, type,
ByteString.EMPTY);
+ List<Integer> firstPage = decodeKeys(page, type);
+ assertThat(firstPage).containsExactlyElementsOf(order.subList(0,
2));
+
assertThat(page.getContinuationToken()).isNotEqualTo(ByteString.EMPTY);
+ assertThat(checkEmptyExcluding(fixture,
Collections.singletonList(firstPage.get(0))))
+ .isFalse();
+
+ List<Integer> visited = new ArrayList<>(firstPage);
+ while (!page.getContinuationToken().isEmpty()) {
+ page = iterate(fixture, type, page.getContinuationToken());
+ visited.addAll(decodeKeys(page, type));
+ }
+ // ReturnExpiredIfNotCleanedUp permits expired entries to be
returned during cleanup.
+ List<Integer> expected =
+ visibility ==
StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp
+ ? order
+ : order.subList(0, 8);
+ assertThat(visited).containsExactlyElementsOf(expected);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testExcludedKeysAndNullValues(boolean metrics) throws Exception {
+ try (Fixture fixture =
+ new Fixture(metrics,
StateTtlConfig.StateVisibility.NeverReturnExpired)) {
+ fixture.state.put(key(1), null);
+ fixture.state.put(key(2), new byte[] {2});
+ assertThat(checkEmptyExcluding(fixture,
Collections.emptyList())).isFalse();
+ assertThat(checkEmptyExcluding(fixture,
Collections.singletonList(2))).isFalse();
+ assertThat(checkEmptyExcluding(fixture, List.of(1, 2))).isTrue();
+ fixture.state.clear();
+ assertThat(checkEmptyExcluding(fixture,
Collections.emptyList())).isTrue();
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testCheckingEmptyDoesNotRefreshTtl(boolean metrics) throws Exception {
+ try (Fixture fixture =
+ new Fixture(metrics,
StateTtlConfig.StateVisibility.NeverReturnExpired)) {
+ fixture.state.put(key(1), new byte[] {1});
+ fixture.harness.setStateTtlProcessingTime(50);
+ assertThat(checkEmptyExcluding(fixture,
Collections.singletonList(2))).isFalse();
+ fixture.harness.setStateTtlProcessingTime(101);
+ assertThat(fixture.state.get(key(1))).isNull();
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testEmptyExclusionsDoNotRunIncrementalCleanup(boolean metrics) throws
Exception {
+ try (Fixture fixture =
+ new Fixture(metrics,
StateTtlConfig.StateVisibility.NeverReturnExpired, true)) {
+ fixture.state.put(key(1), new byte[] {1});
+ fixture.harness.setStateTtlProcessingTime(101);
+ assertThat(checkEmptyExcluding(fixture,
Collections.emptyList())).isTrue();
+ // A read-only query reports expiration while retaining the
physical state entry.
+ assertThat(fixture.harness.numKeyedStateEntries()).isEqualTo(1);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testReadOnlyRpcDoesNotTriggerPerRecordCleanup(boolean metrics) throws
Exception {
+ try (Fixture fixture =
+ new Fixture(
+ metrics,
StateTtlConfig.StateVisibility.NeverReturnExpired, false, 16)) {
+ BeamStateRequestHandler requests =
+ BeamStateRequestHandler.of(
+ fixture.backend, null, IntSerializer.INSTANCE,
null, fixture.config);
+ AtomicInteger selections = new AtomicInteger();
+ fixture.backend.registerKeySelectionListener(ignored ->
selections.incrementAndGet());
+ BeamFnApi.StateKey firstKey = rpcStateKey(0);
+ BeamFnApi.StateKey secondKey = rpcStateKey(1);
+ appendThroughStore(requests, firstKey);
+ appendThroughStore(requests, secondKey);
+ BeamFnApi.StateGetResponse page =
+ getThroughStore(
+ requests,
+ firstKey,
+ iterateRequest(IterateType.ITEMS,
ByteString.EMPTY));
+
assertThat(page.getContinuationToken()).isNotEqualTo(ByteString.EMPTY);
+ int beforeQuery = selections.get();
+ assertThat(beforeQuery).isEqualTo(3);
+ fixture.harness.setStateTtlProcessingTime(101);
+
+ for (int key : new int[] {0, 1, 0}) {
+ BeamFnApi.StateGetResponse response =
+ getThroughStore(
+ requests,
+ rpcStateKey(key),
+
checkEmptyExcludingRequest(Collections.emptyList()));
+
assertThat(response.getData()).isEqualTo(ByteString.copyFrom(new byte[] {3}));
+ assertThat(selections.get()).isEqualTo(beforeQuery);
+ assertThat(fixture.backend.getCurrentKey()).isEqualTo(key);
+ assertThat(fixture.backend.getCurrentKeyGroupIndex())
+
.isEqualTo(KeyGroupRangeAssignment.assignToKeyGroup(key, 16));
+
assertThat(fixture.harness.numKeyedStateEntries()).isEqualTo(2);
+ }
+
+ int ordinaryRequests = 0;
+ while (!page.getContinuationToken().isEmpty()) {
+ page =
+ getThroughStore(
+ requests,
+ firstKey,
+ iterateRequest(IterateType.ITEMS,
page.getContinuationToken()));
+ ordinaryRequests++;
+ }
+ assertThat(ordinaryRequests).isPositive();
+ assertThat(selections.get()).isEqualTo(beforeQuery +
ordinaryRequests);
+ }
+ }
+
+ private static BeamFnApi.StateKey rpcStateKey(int key) throws Exception {
+ CleanupStrategies cleanup =
+ CleanupStrategies.newBuilder()
+ .setIsCleanupInBackground(false)
+ .addStrategies(
+
CleanupStrategies.MapStrategiesEntry.newBuilder()
+ .setStrategy(
+
CleanupStrategies.Strategies.INCREMENTAL_CLEANUP)
+ .setIncrementalCleanupStrategy(
+
CleanupStrategies.IncrementalCleanupStrategy
+ .newBuilder()
+ .setCleanupSize(1)
+
.setRunCleanupForEveryRecord(true)))
+ .build();
+ StateDescriptor descriptor =
+ StateDescriptor.newBuilder()
+ .setStateName("rpc-map")
+ .setStateTtlConfig(
+ StateTTLConfig.newBuilder()
+ .setTtl(100)
+
.setUpdateType(StateTTLConfig.UpdateType.OnReadAndWrite)
+ .setStateVisibility(
+
StateTTLConfig.StateVisibility.NeverReturnExpired)
+ .setCleanupStrategies(cleanup))
+ .build();
+ ByteArrayOutputStream keyBytes = new ByteArrayOutputStream();
+ new DataOutputStream(keyBytes).writeInt(key);
+ return BeamFnApi.StateKey.newBuilder()
+ .setMultimapSideInput(
+ BeamFnApi.StateKey.MultimapSideInput.newBuilder()
+
.setKey(ByteString.copyFrom(keyBytes.toByteArray()))
+ .setSideInputId(
+ Base64.getEncoder()
+
.encodeToString(descriptor.toByteArray())))
+ .build();
+ }
+
+ private static void appendThroughStore(BeamStateRequestHandler handler,
BeamFnApi.StateKey key)
+ throws Exception {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ DataOutputStream data = new DataOutputStream(bytes);
+ data.writeInt(4);
+ for (int mapKey = 0; mapKey < 4; mapKey++) {
+ data.writeByte(2);
+ data.writeInt(1);
+ data.writeByte(mapKey);
+ data.writeInt(1);
+ data.writeByte(mapKey);
+ }
+ handler.handle(
+ BeamFnApi.StateRequest.newBuilder()
+ .setId("append")
+ .setStateKey(key)
+ .setAppend(
+
BeamFnApi.StateAppendRequest.newBuilder()
+
.setData(ByteString.copyFrom(bytes.toByteArray())))
+ .build())
+ .toCompletableFuture()
+ .get();
+ }
+
+ private static BeamFnApi.StateGetResponse getThroughStore(
+ BeamStateRequestHandler handler, BeamFnApi.StateKey key, byte[]
request)
+ throws Exception {
+ return handler.handle(
+ BeamFnApi.StateRequest.newBuilder()
+ .setId("get")
+ .setStateKey(key)
+ .setGet(
+ BeamFnApi.StateGetRequest.newBuilder()
+
.setContinuationToken(ByteString.copyFrom(request)))
+ .build())
+ .toCompletableFuture()
+ .get()
+ .build()
+ .getGet();
+ }
+
+ private static boolean checkEmptyExcluding(Fixture fixture, List<Integer>
excluded)
+ throws Exception {
+ BeamFnApi.StateGetResponse response = get(fixture,
checkEmptyExcludingRequest(excluded));
+
assertThat(response.getContinuationToken()).isEqualTo(ByteString.EMPTY);
+ assertThat(response.getData().size()).isEqualTo(1);
+ assertThat(response.getData().byteAt(0)).isIn((byte) 3, (byte) 4);
+ return response.getData().byteAt(0) == 3;
+ }
+
+ private static byte[] checkEmptyExcludingRequest(List<Integer> excluded)
throws Exception {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ DataOutputStream request = new DataOutputStream(bytes);
+ request.writeByte(3);
+ request.writeInt(excluded.size());
+ for (int key : excluded) {
+ request.writeInt(1);
+ request.writeByte(key);
+ }
+ return bytes.toByteArray();
+ }
+
+ private static BeamFnApi.StateGetResponse iterate(
+ Fixture fixture, IterateType type, ByteString token) throws
Exception {
+ return get(fixture, iterateRequest(type, token));
+ }
+
+ private static byte[] iterateRequest(IterateType type, ByteString token)
throws Exception {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ DataOutputStream request = new DataOutputStream(bytes);
+ request.writeByte(1);
+ request.writeByte(type.getOrd());
+ request.writeInt(token.size());
+ token.writeTo(request);
+ return bytes.toByteArray();
+ }
+
+ private static BeamFnApi.StateGetResponse get(Fixture fixture, byte[]
request)
+ throws Exception {
+ return fixture.handler
+ .handleGet(
+ BeamFnApi.StateRequest.newBuilder()
+ .setId("request")
+ .setGet(
+ BeamFnApi.StateGetRequest.newBuilder()
+
.setContinuationToken(ByteString.copyFrom(request)))
+ .build(),
+ fixture.state)
+ .build()
+ .getGet();
+ }
+
+ private static List<Integer> decodeKeys(BeamFnApi.StateGetResponse
response, IterateType type)
+ throws Exception {
+ DataInputStream input =
+ new DataInputStream(new
ByteArrayInputStream(response.getData().toByteArray()));
+ List<Integer> keys = new ArrayList<>();
+ while (input.available() > 0) {
+ keys.add(input.readUnsignedByte());
+ if (type != IterateType.KEYS && input.readBoolean()) {
+ input.readByte();
+ }
+ }
+ return keys;
+ }
+
+ private static ByteArrayWrapper key(int key) {
+ return new ByteArrayWrapper(new byte[] {(byte) key});
+ }
+
+ private static class Fixture implements AutoCloseable {
+ private final KeyedOneInputStreamOperatorTestHarness<Integer, Integer,
Integer> harness;
+ private final MapState<ByteArrayWrapper, byte[]> state;
+ private final BeamMapStateHandler handler;
+ private final AbstractKeyedStateBackend<Integer> backend;
+ private final Configuration config;
+
+ private Fixture(boolean metrics, StateTtlConfig.StateVisibility
visibility)
+ throws Exception {
+ this(metrics, visibility, false);
+ }
+
+ private Fixture(
+ boolean metrics, StateTtlConfig.StateVisibility visibility,
boolean incremental)
+ throws Exception {
+ this(metrics, visibility, incremental, 1);
+ }
+
+ private Fixture(
+ boolean metrics,
+ StateTtlConfig.StateVisibility visibility,
+ boolean incremental,
+ int maxParallelism)
+ throws Exception {
+ config = new Configuration();
+ config.set(PythonOptions.MAP_STATE_ITERATE_RESPONSE_BATCH_SIZE, 2);
+ config.set(StateLatencyTrackOptions.LATENCY_TRACK_ENABLED,
metrics);
+ config.set(StateLatencyTrackOptions.LATENCY_TRACK_SAMPLE_INTERVAL,
1);
+ StreamMap<Integer, Integer> operator = new StreamMap<>(value ->
value);
+ harness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ operator, value -> 0, Types.INT, maxParallelism,
1, 0);
+ harness.setStateBackend(
+ new HashMapStateBackend()
+ .configure(config,
Thread.currentThread().getContextClassLoader()));
+ harness.open();
+ harness.setStateTtlProcessingTime(0);
+ backend = (AbstractKeyedStateBackend<Integer>)
operator.<Integer>getKeyedStateBackend();
+ backend.setCurrentKey(0);
+ MapStateDescriptor<ByteArrayWrapper, byte[]> descriptor =
+ new MapStateDescriptor<>(
+ "map",
+ ByteArrayWrapperSerializer.INSTANCE,
+ BytePrimitiveArraySerializer.INSTANCE);
+ StateTtlConfig.Builder ttl =
+ StateTtlConfig.newBuilder(Duration.ofMillis(100))
+
.setUpdateType(StateTtlConfig.UpdateType.OnReadAndWrite)
+ .setStateVisibility(visibility)
+ .disableCleanupInBackground();
+ if (incremental) {
+ ttl.cleanupIncrementally(1, false);
+ }
+ descriptor.enableTimeToLive(ttl.build());
+ state =
+ backend.getPartitionedState(
+ VoidNamespace.INSTANCE,
VoidNamespaceSerializer.INSTANCE, descriptor);
+ if (metrics) {
+
assertThat(state.getClass().getSimpleName()).isEqualTo("MetricsTrackingMapState");
+ }
+ handler = new BeamMapStateHandler(config);
+ }
+
+ @Override
+ public void close() throws Exception {
+ harness.close();
+ }
+ }
+}
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java
index 1bec7b2218d..f34378a250a 100644
---
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java
@@ -262,6 +262,13 @@ public abstract class AbstractKeyedStateBackend<K>
KeyGroupRangeAssignment.assignToKeyGroup(newKey,
numberOfKeyGroups));
}
+ @Override
+ public void setCurrentKeyForStateAccess(K newKey) {
+ this.keyContext.setCurrentKey(newKey);
+ this.keyContext.setCurrentKeyGroupIndex(
+ KeyGroupRangeAssignment.assignToKeyGroup(newKey,
numberOfKeyGroups));
+ }
+
@Override
public void setCurrentKeyAndKeyGroup(K newKey, int newKeyGroupIndex) {
notifyKeySelected(newKey);
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
index e11ca590ab3..807bf83bc95 100644
---
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
@@ -43,6 +43,15 @@ public interface KeyedStateBackend<K>
*/
void setCurrentKey(K newKey);
+ /**
+ * Restores the key context for a state request rather than selecting a
key for another input
+ * record. Backends with record-triggered maintenance should override this
method to avoid
+ * running that maintenance for the request.
+ */
+ default void setCurrentKeyForStateAccess(K newKey) {
+ setCurrentKey(newKey);
+ }
+
/**
* @return Current key.
*/
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/internal/InternalMapState.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/internal/InternalMapState.java
index 92ae9c0c4c0..84572179c8a 100644
---
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/internal/InternalMapState.java
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/internal/InternalMapState.java
@@ -21,6 +21,7 @@ package org.apache.flink.runtime.state.internal;
import org.apache.flink.api.common.state.MapState;
import java.util.Map;
+import java.util.Set;
/**
* The peer to the {@link MapState} in the internal state type hierarchy.
@@ -33,4 +34,25 @@ import java.util.Map;
* @param <UV> Type of the value in the state
*/
public interface InternalMapState<K, N, UK, UV>
- extends InternalKvState<K, N, Map<UK, UV>>, MapState<UK, UV> {}
+ extends InternalKvState<K, N, Map<UK, UV>>, MapState<UK, UV> {
+
+ /**
+ * Returns whether this state has no visible entries other than the
excluded keys, without
+ * updating TTL timestamps or cleaning up expired entries.
+ *
+ * <p>This allows clients with buffered deletions to check emptiness
without applying those
+ * deletions or invalidating active iterators. TTL decorators must
override this method instead
+ * of using their ordinary, potentially mutating key iteration.
+ */
+ default boolean isEmptyExcluding(Set<?> excludedKeys) throws Exception {
+ final Iterable<UK> keys = keys();
+ if (keys != null) {
+ for (UK key : keys) {
+ if (!excludedKeys.contains(key)) {
+ return false;
+ }
+ }
+ }
+ return true;
+ }
+}
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/metrics/MetricsTrackingMapState.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/metrics/MetricsTrackingMapState.java
index 872b12ea691..f693ffdd97e 100644
---
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/metrics/MetricsTrackingMapState.java
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/metrics/MetricsTrackingMapState.java
@@ -29,6 +29,7 @@ import org.apache.flink.util.function.ThrowingConsumer;
import java.io.IOException;
import java.util.Iterator;
import java.util.Map;
+import java.util.Set;
/**
* This class wraps map state with latency tracking logic.
@@ -261,6 +262,22 @@ class MetricsTrackingMapState<K, N, UK, UV>
}
}
+ @Override
+ public boolean isEmptyExcluding(Set<?> excludedKeys) throws Exception {
+ if (sizeTrackingStateMetric != null &&
sizeTrackingStateMetric.trackMetricsOnIsEmpty()) {
+ sizeTrackingStateMetric.updateMetrics(
+ MapStateMetrics.MAP_STATE_IS_EMPTY_KEY_SIZE,
super.sizeOfKey());
+ }
+ if (latencyTrackingStateMetric != null
+ && latencyTrackingStateMetric.trackMetricsOnIsEmpty()) {
+ return trackLatencyWithException(
+ () -> original.isEmptyExcluding(excludedKeys),
+ MapStateMetrics.MAP_STATE_IS_EMPTY_LATENCY);
+ } else {
+ return original.isEmptyExcluding(excludedKeys);
+ }
+ }
+
protected long sizeOfKeyAndUserKey(UK userKey) throws IOException {
if (userKeySerializer == null || userKey == null) {
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlMapState.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlMapState.java
index 2dc2df3a031..242fa43c23d 100644
---
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlMapState.java
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlMapState.java
@@ -33,6 +33,7 @@ import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
+import java.util.Set;
import java.util.function.Function;
/**
@@ -132,6 +133,24 @@ class TtlMapState<K, N, UK, UV>
return original.isEmpty();
}
+ @Override
+ public boolean isEmptyExcluding(Set<?> excludedKeys) throws Exception {
+ // Ordinary TTL reads can clean up entries and invalidate another
active iterator.
+ final Iterable<Map.Entry<UK, TtlValue<UV>>> entries =
original.entries();
+ final long now = timeProvider.currentTimestamp();
+ if (entries != null) {
+ for (Map.Entry<UK, TtlValue<UV>> entry : entries) {
+ final TtlValue<UV> value = entry.getValue();
+ if (!excludedKeys.contains(entry.getKey())
+ && value != null
+ && (returnExpired || !TtlUtils.expired(value, ttl,
now))) {
+ return false;
+ }
+ }
+ }
+ return true;
+ }
+
@Nullable
@Override
public Map<UK, TtlValue<UV>> getUnexpiredOrNull(@Nonnull Map<UK,
TtlValue<UV>> ttlValue) {
diff --git
a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
index 5b2a80a9549..32ab5396c43 100644
---
a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
+++
b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
@@ -3651,6 +3651,37 @@ public abstract class StateBackendTestBase<B extends
AbstractStateBackend> {
}
}
+ @TestTemplate
+ void testMapStateReadOnlyKeyContext() throws Exception {
+ CheckpointableKeyedStateBackend<Integer> backend =
+ createKeyedBackend(IntSerializer.INSTANCE);
+ try {
+ MapState<Integer, Long> state =
+ backend.getPartitionedState(
+ VoidNamespace.INSTANCE,
+ VoidNamespaceSerializer.INSTANCE,
+ new MapStateDescriptor<>("map", Integer.class,
Long.class));
+ backend.setCurrentKey(1);
+ state.put(1, 100L);
+ backend.setCurrentKey(2);
+ state.put(1, 200L);
+
+ int[] selections = {0};
+ backend.registerKeySelectionListener(ignored -> selections[0]++);
+ for (int key : new int[] {1, 3, 2, 1}) {
+ backend.setCurrentKeyForStateAccess(key);
+ assertThat(backend.getCurrentKey()).isEqualTo(key);
+ assertThat(state.get(1)).isEqualTo(key == 3 ? null : key *
100L);
+ }
+ assertThat(selections[0]).isZero();
+ backend.setCurrentKey(1);
+ assertThat(selections[0]).isEqualTo(1);
+ } finally {
+ IOUtils.closeQuietly(backend);
+ backend.dispose();
+ }
+ }
+
@TestTemplate
void testMapStateIsEmpty() throws Exception {
MapStateDescriptor<Integer, Long> kvId =
diff --git
a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java
b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java
index 621aefb9f8b..624ba448adb 100644
---
a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java
+++
b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java
@@ -294,6 +294,11 @@ public class ChangelogKeyedStateBackend<K>
keyedStateBackend.setCurrentKey(newKey);
}
+ @Override
+ public void setCurrentKeyForStateAccess(K newKey) {
+ keyedStateBackend.setCurrentKeyForStateAccess(newKey);
+ }
+
@Override
public K getCurrentKey() {
return keyedStateBackend.getCurrentKey();
diff --git
a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java
b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java
index ba02e93d221..1c668a9a797 100644
---
a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java
+++
b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java
@@ -202,6 +202,11 @@ public class ChangelogMigrationRestoreTarget<K> implements
ChangelogRestoreTarge
keyedStateBackend.setCurrentKey(newKey);
}
+ @Override
+ public void setCurrentKeyForStateAccess(K newKey) {
+ keyedStateBackend.setCurrentKeyForStateAccess(newKey);
+ }
+
@Override
public void setCurrentKeyAndKeyGroup(K newKey, int
newKeyGroupIndex) {
keyedStateBackend.setCurrentKeyAndKeyGroup(newKey,
newKeyGroupIndex);
diff --git
a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java
b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java
index 97ab1dfbc1b..48e4ec447d1 100644
---
a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java
+++
b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java
@@ -524,6 +524,12 @@ public class ForStSyncKeyedStateBackend<K> extends
AbstractKeyedStateBackend<K>
sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(),
getCurrentKeyGroupIndex());
}
+ @Override
+ public void setCurrentKeyForStateAccess(K newKey) {
+ super.setCurrentKeyForStateAccess(newKey);
+ sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(),
getCurrentKeyGroupIndex());
+ }
+
@Override
public void setCurrentKeyAndKeyGroup(K newKey, int newKeyGroupIndex) {
super.setCurrentKeyAndKeyGroup(newKey, newKeyGroupIndex);
diff --git
a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java
b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java
index c6512f4a356..d9958b7566e 100644
---
a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java
+++
b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java
@@ -553,6 +553,12 @@ public class RocksDBKeyedStateBackend<K> extends
AbstractKeyedStateBackend<K> {
sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(),
getCurrentKeyGroupIndex());
}
+ @Override
+ public void setCurrentKeyForStateAccess(K newKey) {
+ super.setCurrentKeyForStateAccess(newKey);
+ sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(),
getCurrentKeyGroupIndex());
+ }
+
@Override
public void setCurrentKeyAndKeyGroup(K newKey, int newKeyGroupIndex) {
super.setCurrentKeyAndKeyGroup(newKey, newKeyGroupIndex);
