[ 
https://issues.apache.org/jira/browse/BEAM-2918?focusedWorklogId=157773&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-157773
 ]

ASF GitHub Bot logged work on BEAM-2918:
----------------------------------------

                Author: ASF GitHub Bot
            Created on: 23/Oct/18 19:11
            Start Date: 23/Oct/18 19:11
    Worklog Time Spent: 10m 
      Work Description: tweise commented on a change in pull request #6726: 
[BEAM-2918] Add state support for streaming in portable FlinkRunner
URL: https://github.com/apache/beam/pull/6726#discussion_r227527404
 
 

 ##########
 File path: 
runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
 ##########
 @@ -135,31 +149,131 @@ public void open() throws Exception {
     // ownership of the higher level "factories" explicit? Do we care?
     stageContext = contextFactory.get(jobInfo);
 
-    stateRequestHandler = getStateRequestHandler(executableStage);
     stageBundleFactory = stageContext.getStageBundleFactory(executableStage);
+    stateRequestHandler = getStateRequestHandler(executableStage);
     progressHandler = BundleProgressHandler.unsupported();
     outputQueue = new LinkedBlockingQueue<>();
   }
 
   private StateRequestHandler getStateRequestHandler(ExecutableStage 
executableStage) {
 
+    final StateRequestHandler sideInputStateHandler;
     if (executableStage.getSideInputs().size() > 0) {
       checkNotNull(super.sideInputHandler);
       StateRequestHandlers.SideInputHandlerFactory sideInputHandlerFactory =
           Preconditions.checkNotNull(
               FlinkStreamingSideInputHandlerFactory.forStage(
                   executableStage, sideInputIds, super.sideInputHandler));
       try {
-        return StateRequestHandlers.forSideInputHandlerFactory(
-            ProcessBundleDescriptors.getSideInputs(executableStage), 
sideInputHandlerFactory);
+        sideInputStateHandler =
+            StateRequestHandlers.forSideInputHandlerFactory(
+                ProcessBundleDescriptors.getSideInputs(executableStage), 
sideInputHandlerFactory);
       } catch (IOException e) {
-        throw new RuntimeException(e);
+        throw new RuntimeException("Failed to initialize SideInputHandler", e);
+      }
+    } else {
+      sideInputStateHandler = StateRequestHandler.unsupported();
+    }
+
+    final StateRequestHandler userStateRequestHandler;
+    if (executableStage.getUserStates().size() > 0) {
+      if (keyedStateInternals == null) {
+        throw new IllegalStateException("Input must be keyed when user state 
is used");
       }
+      userStateRequestHandler =
+          StateRequestHandlers.forBagUserStateHandlerFactory(
+              stageBundleFactory.getProcessBundleDescriptor(),
+              new BagUserStateFactory(keyedStateInternals, 
getKeyedStateBackend()));
     } else {
-      return StateRequestHandler.unsupported();
+      userStateRequestHandler = StateRequestHandler.unsupported();
+    }
+
+    EnumMap<TypeCase, StateRequestHandler> handlerMap = new 
EnumMap<>(TypeCase.class);
+    handlerMap.put(TypeCase.MULTIMAP_SIDE_INPUT, sideInputStateHandler);
+    handlerMap.put(TypeCase.BAG_USER_STATE, userStateRequestHandler);
+
+    return StateRequestHandlers.delegateBasedUponType(handlerMap);
+  }
+
+  private static class BagUserStateFactory
+      implements StateRequestHandlers.BagUserStateHandlerFactory {
+
+    private final StateInternals stateInternals;
+    private final KeyedStateBackend<ByteBuffer> keyedStateBackend;
+
+    private BagUserStateFactory(
+        StateInternals stateInternals, KeyedStateBackend<ByteBuffer> 
keyedStateBackend) {
+
+      this.stateInternals = stateInternals;
+      this.keyedStateBackend = keyedStateBackend;
+    }
+
+    @Override
+    public <K, V, W extends BoundedWindow>
+        StateRequestHandlers.BagUserStateHandler<K, V, W> forUserState(
+            String pTransformId,
+            String userStateId,
+            Coder<K> keyCoder,
+            Coder<V> valueCoder,
+            Coder<W> windowCoder) {
+      return new StateRequestHandlers.BagUserStateHandler<K, V, W>() {
+        @Override
+        public Iterable<V> get(K key, W window) {
+          prepareStateBackend(key, keyCoder);
+          StateNamespace namespace = StateNamespaces.window(windowCoder, 
window);
+          BagState<V> bagState =
+              stateInternals.state(namespace, StateTags.bag(userStateId, 
valueCoder));
+          return bagState.read();
+        }
+
+        @Override
+        public void append(K key, W window, Iterator<V> values) {
+          prepareStateBackend(key, keyCoder);
+          StateNamespace namespace = StateNamespaces.window(windowCoder, 
window);
+          BagState<V> bagState =
+              stateInternals.state(namespace, StateTags.bag(userStateId, 
valueCoder));
+          while (values.hasNext()) {
+            bagState.add(values.next());
+          }
+        }
+
+        @Override
+        public void clear(K key, W window) {
+          prepareStateBackend(key, keyCoder);
+          StateNamespace namespace = StateNamespaces.window(windowCoder, 
window);
+          BagState<V> bagState =
+              stateInternals.state(namespace, StateTags.bag(userStateId, 
valueCoder));
+          bagState.clear();
+        }
+
+        private void prepareStateBackend(K key, Coder<K> keyCoder) {
+          ByteArrayOutputStream baos = new ByteArrayOutputStream();
+          try {
+            keyCoder.encode(key, baos);
+          } catch (IOException e) {
+            throw new RuntimeException("Failed to encode key for Flink state 
backend", e);
+          }
+          keyedStateBackend.setCurrentKey(ByteBuffer.wrap(baos.toByteArray()));
+        }
+      };
     }
   }
 
+  @Override
+  public void setKeyContextElement1(StreamRecord record) throws Exception {
+    // Note: This is only relevant when we have a stateful DoFn.
+    // We want to control the key of the state backend ourselves and
+    // we must avoid any concurrent setting of the current active key.
+    // By overwriting this, we also prevent unnecessary serialization
+    // as the key has to be encoded as a byte array.
+  }
+
+  @Override
+  public void setCurrentKey(Object key) {
+    throw new IllegalStateException(
+        "This should never be called. Current key for StateBackend should be 
set manually.");
 
 Review comment:
   Maybe use UnsupportedOperationException and "Current key for state backend 
can only be set by SDK workers."

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
[email protected]


Issue Time Tracking
-------------------

    Worklog Id:     (was: 157773)
    Time Spent: 8h 40m  (was: 8.5h)

> Flink support for portable user state
> -------------------------------------
>
>                 Key: BEAM-2918
>                 URL: https://issues.apache.org/jira/browse/BEAM-2918
>             Project: Beam
>          Issue Type: Sub-task
>          Components: runner-flink
>            Reporter: Henning Rohde
>            Assignee: Maximilian Michels
>            Priority: Minor
>              Labels: portability
>          Time Spent: 8h 40m
>  Remaining Estimate: 0h
>




--
This message was sent by Atlassian JIRA
(v7.6.3#76005)

Reply via email to