Support ordered list states in python sdk and fnapi runner (#32326)

* Support ordered list state in python sdk and fnapi runner.

* Add test to verify integrity of multiple iterators

* Add fuzz tests and fix two edge cases.

* Add sortedcontainer in package dependency

* Code refactoring and add a check for the supported maximum key

* regenerate requirments for python images.

* Refactor portable runner code for ordered list state

* Return continuation tokens in portable runnner for ordered list state

* Fix some lints

* Apply yapf

* Fix lints

* Sync base image requirements with master.

* Add typing for ordered list state apis.

* Add typing to orderedliststate user state.

* Fix a typo.

* Refactor some code based on the feedback.

* Fix lints

* Remove the support of int argument type in ordered list state apis.

* Fix formats and lints

* More lints

* Refactor the code to use the continuation token logic.

* Fix lints

sdks/python/apache_beam/runners/portability/fn_api_runner/worker_handlers.py

@@ -48,7 +48,9 @@
 from typing import overload
 
 import grpc
+from sortedcontainers import SortedSet
 
+from apache_beam import coders
 from apache_beam.io import filesystems
 from apache_beam.io.filesystems import CompressionTypes
 from apache_beam.portability import common_urns
@@ -959,7 +961,8 @@ class StateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer,
       'multimap_keys_values_side_input',
       'iterable_side_input',
       'bag_user_state',
-      'multimap_user_state'
+      'multimap_user_state',
+      'ordered_list_user_state'
   ])
 
   class CopyOnWriteState(object):
@@ -1021,6 +1024,8 @@ def __init__(self):
     self._checkpoint = None  # type: Optional[StateServicer.StateType]
     self._use_continuation_tokens = False
     self._continuations = {}  # type: Dict[bytes, Tuple[bytes, ...]]
+    self._ordered_list_keys = collections.defaultdict(
+        SortedSet)  # type: DefaultDict[bytes, SortedSet]
 
   def checkpoint(self):
     # type: () -> None
@@ -1050,6 +1055,14 @@ def process_instruction_id(self, unused_instruction_id):
     # type: (Any) -> Iterator
     yield
 
+  def _get_one_interval_key(self, state_key, start):
+    # type: (beam_fn_api_pb2.StateKey, int) -> bytes
+    state_key_copy = beam_fn_api_pb2.StateKey()
+    state_key_copy.CopyFrom(state_key)
+    state_key_copy.ordered_list_user_state.range.start = start
+    state_key_copy.ordered_list_user_state.range.end = start + 1
+    return self._to_key(state_key_copy)
+
   def get_raw(self,
       state_key,  # type: beam_fn_api_pb2.StateKey
       continuation_token=None  # type: Optional[bytes]
@@ -1061,7 +1074,30 @@ def get_raw(self,
           'Unknown state type: ' + state_key.WhichOneof('type'))
 
     with self._lock:
-      full_state = self._state[self._to_key(state_key)]
+      if not continuation_token:
+        # Compute full_state only when no continuation token is provided.
+        # If there is continuation token, full_state is already in
+        # continuation cache. No need to recompute.
+        full_state = []  # type: List[bytes]
+        if state_key.WhichOneof('type') == 'ordered_list_user_state':
+          maybe_start = state_key.ordered_list_user_state.range.start
+          maybe_end = state_key.ordered_list_user_state.range.end
+          persistent_state_key = beam_fn_api_pb2.StateKey()
+          persistent_state_key.CopyFrom(state_key)
+          persistent_state_key.ordered_list_user_state.ClearField("range")
+
+          available_keys = self._ordered_list_keys[self._to_key(
+              persistent_state_key)]
+
+          for i in available_keys.irange(maybe_start,
+                                         maybe_end,
+                                         inclusive=(True, False)):
+            entries = self._state[self._get_one_interval_key(
+                persistent_state_key, i)]
+            full_state.extend(entries)
+        else:
+          full_state.extend(self._state[self._to_key(state_key)])
+
       if self._use_continuation_tokens:
         # The token is "nonce:index".
         if not continuation_token:
@@ -1087,14 +1123,40 @@ def append_raw(
   ):
     # type: (...) -> _Future
     with self._lock:
-      self._state[self._to_key(state_key)].append(data)
+      if state_key.WhichOneof('type') == 'ordered_list_user_state':
+        coder = coders.TupleCoder([
+            coders.VarIntCoder(),
+            coders.coders.LengthPrefixCoder(coders.BytesCoder())
+        ]).get_impl()
+
+        for key, value in coder.decode_all(data):
+          self._state[self._get_one_interval_key(state_key, key)].append(
+              coder.encode((key, value)))
+          self._ordered_list_keys[self._to_key(state_key)].add(key)
+      else:
+        self._state[self._to_key(state_key)].append(data)
     return _Future.done()
 
   def clear(self, state_key):
     # type: (beam_fn_api_pb2.StateKey) -> _Future
     with self._lock:
       try:
-        del self._state[self._to_key(state_key)]
+        if state_key.WhichOneof('type') == 'ordered_list_user_state':
+          start = state_key.ordered_list_user_state.range.start
+          end = state_key.ordered_list_user_state.range.end
+          persistent_state_key = beam_fn_api_pb2.StateKey()
+          persistent_state_key.CopyFrom(state_key)
+          persistent_state_key.ordered_list_user_state.ClearField("range")
+          available_keys = self._ordered_list_keys[self._to_key(
+              persistent_state_key)]
+
+          for i in list(available_keys.irange(start,
+                                              end,
+                                              inclusive=(True, False))):
+            del self._state[self._get_one_interval_key(persistent_state_key, i)]
+            available_keys.remove(i)
+        else:
+          del self._state[self._to_key(state_key)]
       except KeyError:
         # This may happen with the caching layer across bundles. Caching may
         # skip this storage layer for a blocking_get(key) request. Without

sdks/python/apache_beam/runners/worker/bundle_processor.py

@@ -19,16 +19,21 @@
 
 # pytype: skip-file
 
+from __future__ import annotations
+
 import base64
 import bisect
 import collections
 import copy
+import heapq
+import itertools
 import json
 import logging
 import random
 import threading
 from dataclasses import dataclass
 from dataclasses import field
+from itertools import chain
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
@@ -50,6 +55,8 @@
 
 from google.protobuf import duration_pb2
 from google.protobuf import timestamp_pb2
+from sortedcontainers import SortedDict
+from sortedcontainers import SortedList
 
 import apache_beam as beam
 from apache_beam import coders
@@ -104,7 +111,8 @@
 FnApiUserRuntimeStateTypes = Union['ReadModifyWriteRuntimeState',
                                    'CombiningValueRuntimeState',
                                    'SynchronousSetRuntimeState',
-                                   'SynchronousBagRuntimeState']
+                                   'SynchronousBagRuntimeState',
+                                   'SynchronousOrderedListRuntimeState']
 
 DATA_INPUT_URN = 'beam:runner:source:v1'
 DATA_OUTPUT_URN = 'beam:runner:sink:v1'
@@ -704,6 +712,180 @@ def commit(self):
       to_await.get()
 
 
+class RangeSet:
+  """For Internal Use only. A simple range set for ranges of [x,y)."""
+  def __init__(self) -> None:
+    # The start points and end points are stored separately in order.
+    self._sorted_starts = SortedList()
+    self._sorted_ends = SortedList()
+
+  def add(self, start: int, end: int) -> None:
+    if start >= end:
+      return
+
+    # ranges[:min_idx] and ranges[max_idx:] is unaffected by this insertion
+    # the first range whose end point >= the start of the new range
+    min_idx = self._sorted_ends.bisect_left(start)
+    # the first range whose start point > the end point of the new range
+    max_idx = self._sorted_starts.bisect_right(end)
+
+    if min_idx >= len(self._sorted_starts) or max_idx <= 0:
+      # the new range is beyond any current ranges
+      new_start = start
+      new_end = end
+    else:
+      # the new range overlaps with ranges[min_idx:max_idx]
+      new_start = min(start, self._sorted_starts[min_idx])
+      new_end = max(end, self._sorted_ends[max_idx - 1])
+
+      del self._sorted_starts[min_idx:max_idx]
+      del self._sorted_ends[min_idx:max_idx]
+
+    self._sorted_starts.add(new_start)
+    self._sorted_ends.add(new_end)
+
+  def __contains__(self, key: int) -> bool:
+    idx = self._sorted_starts.bisect_left(key)
+    return (idx < len(self._sorted_starts) and self._sorted_starts[idx] == key
+            ) or (idx > 0 and self._sorted_ends[idx - 1] > key)
+
+  def __len__(self) -> int:
+    assert len(self._sorted_starts) == len(self._sorted_ends)
+    return len(self._sorted_starts)
+
+  def __iter__(self) -> Iterator[Tuple[int, int]]:
+    return zip(self._sorted_starts, self._sorted_ends)
+
+  def __str__(self) -> str:
+    return str(list(zip(self._sorted_starts, self._sorted_ends)))
+
+
+class SynchronousOrderedListRuntimeState(userstate.OrderedListRuntimeState):
+  RANGE_MIN = -(1 << 63)
+  RANGE_MAX = (1 << 63) - 1
+  TIMESTAMP_RANGE_MIN = timestamp.Timestamp(micros=RANGE_MIN)
+  TIMESTAMP_RANGE_MAX = timestamp.Timestamp(micros=RANGE_MAX)
+
+  def __init__(
+      self,
+      state_handler: sdk_worker.CachingStateHandler,
+      state_key: beam_fn_api_pb2.StateKey,
+      value_coder: coders.Coder) -> None:
+    self._state_handler = state_handler
+    self._state_key = state_key
+    self._elem_coder = beam.coders.TupleCoder(
+        [coders.VarIntCoder(), coders.coders.LengthPrefixCoder(value_coder)])
+    self._cleared = False
+    self._pending_adds = SortedDict()
+    self._pending_removes = RangeSet()
+
+  def add(self, elem: Tuple[timestamp.Timestamp, Any]) -> None:
+    assert len(elem) == 2
+    key_ts, value = elem
+    key = key_ts.micros
+
+    if key >= self.RANGE_MAX or key < self.RANGE_MIN:
+      raise ValueError("key value %d is out of range" % key)
+    self._pending_adds.setdefault(key, []).append(value)
+
+  def read(self) -> Iterable[Tuple[timestamp.Timestamp, Any]]:
+    return self.read_range(self.TIMESTAMP_RANGE_MIN, self.TIMESTAMP_RANGE_MAX)
+
+  def read_range(
+      self,
+      min_timestamp: timestamp.Timestamp,
+      limit_timestamp: timestamp.Timestamp
+  ) -> Iterable[Tuple[timestamp.Timestamp, Any]]:
+    # convert timestamp to int, as sort keys are stored as int internally.
+    min_key = min_timestamp.micros
+    limit_key = limit_timestamp.micros
+
+    keys_to_add = self._pending_adds.irange(
+        min_key, limit_key, inclusive=(True, False))
+
+    # use list interpretation here to construct the actual list
+    # of iterators of the selected range.
+    local_items = chain.from_iterable([
+        itertools.islice(
+            zip(itertools.cycle([
+                k,
+            ]), self._pending_adds[k]),
+            len(self._pending_adds[k])) for k in keys_to_add
+    ])
+
+    if not self._cleared:
+      range_query_state_key = beam_fn_api_pb2.StateKey()
+      range_query_state_key.CopyFrom(self._state_key)
+      range_query_state_key.ordered_list_user_state.range.start = min_key
+      range_query_state_key.ordered_list_user_state.range.end = limit_key
+
+      # make a deep copy here because there could be other operations occur in
+      # the middle of an iteration and change pending_removes
+      pending_removes_snapshot = copy.deepcopy(self._pending_removes)
+      persistent_items = filter(
+          lambda kv: kv[0] not in pending_removes_snapshot,
+          _StateBackedIterable(
+              self._state_handler, range_query_state_key, self._elem_coder))
+
+      return map(
+          lambda x: (timestamp.Timestamp(micros=x[0]), x[1]),
+          heapq.merge(persistent_items, local_items))
+
+    return map(lambda x: (timestamp.Timestamp(micros=x[0]), x[1]), local_items)
+
+  def clear(self) -> None:
+    self._cleared = True
+    self._pending_adds = SortedDict()
+    self._pending_removes = RangeSet()
+    self._pending_removes.add(self.RANGE_MIN, self.RANGE_MAX)
+
+  def clear_range(
+      self,
+      min_timestamp: timestamp.Timestamp,
+      limit_timestamp: timestamp.Timestamp) -> None:
+    min_key = min_timestamp.micros
+    limit_key = limit_timestamp.micros
+
+    # materialize the keys to remove before the actual removal
+    keys_to_remove = list(
+        self._pending_adds.irange(min_key, limit_key, inclusive=(True, False)))
+    for k in keys_to_remove:
+      del self._pending_adds[k]
+
+    if not self._cleared:
+      self._pending_removes.add(min_key, limit_key)
+
+  def commit(self) -> None:
+    futures = []
+    if self._pending_removes:
+      for start, end in self._pending_removes:
+        range_query_state_key = beam_fn_api_pb2.StateKey()
+        range_query_state_key.CopyFrom(self._state_key)
+        range_query_state_key.ordered_list_user_state.range.start = start
+        range_query_state_key.ordered_list_user_state.range.end = end
+        futures.append(self._state_handler.clear(range_query_state_key))
+
+      self._pending_removes = RangeSet()
+
+    if self._pending_adds:
+      items_to_add = []
+      for k in self._pending_adds:
+        items_to_add.extend(zip(itertools.cycle([
+            k,
+        ]), self._pending_adds[k]))
+      futures.append(
+          self._state_handler.extend(
+              self._state_key, self._elem_coder.get_impl(), items_to_add))
+      self._pending_adds = SortedDict()
+
+    if len(futures):
+      # To commit, we need to wait on every state request futures to complete.
+      for to_await in futures:
+        to_await.get()
+
+    self._cleared = False
+
+
 class OutputTimer(userstate.BaseTimer):
   def __init__(self,
                key,
@@ -850,6 +1032,17 @@ def _create_state(self,
                   # State keys are expected in nested encoding format
                   key=self._key_coder.encode_nested(key))),
           value_coder=state_spec.coder)
+    elif isinstance(state_spec, userstate.OrderedListStateSpec):
+      return SynchronousOrderedListRuntimeState(
+          self._state_handler,
+          state_key=beam_fn_api_pb2.StateKey(
+              ordered_list_user_state=beam_fn_api_pb2.StateKey.
+              OrderedListUserState(
+                  transform_id=self._transform_id,
+                  user_state_id=state_spec.name,
+                  window=self._window_coder.encode(window),
+                  key=self._key_coder.encode_nested(key))),
+          value_coder=state_spec.coder)
     else:
       raise NotImplementedError(state_spec)