Examples with KeyGroupStatePartitionStreamProvider org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider used on opensource projects

Example 1 with KeyGroupStatePartitionStreamProvider

use of org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider in project flink by apache.

the class StateInitializationContextImplTest method close.

@Test
public void close() throws Exception {
    int count = 0;
    int stopCount = NUM_HANDLES / 2;
    boolean isClosed = false;
    try {
        for (KeyGroupStatePartitionStreamProvider stateStreamProvider : initializationContext.getRawKeyedStateInputs()) {
            Assert.assertNotNull(stateStreamProvider);
            if (count == stopCount) {
                initializationContext.close();
                isClosed = true;
            }
            try (InputStream is = stateStreamProvider.getStream()) {
                DataInputView div = new DataInputViewStreamWrapper(is);
                try {
                    int val = div.readInt();
                    Assert.assertEquals(stateStreamProvider.getKeyGroupId(), val);
                    if (isClosed) {
                        Assert.fail("Close was ignored: stream");
                    }
                    ++count;
                } catch (IOException ioex) {
                    if (!isClosed) {
                        throw ioex;
                    }
                }
            }
        }
        Assert.fail("Close was ignored: registry");
    } catch (IOException iex) {
        Assert.assertTrue(isClosed);
        Assert.assertEquals(stopCount, count);
    }
}

Example 2 with KeyGroupStatePartitionStreamProvider

use of org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider in project beam by apache.

the class DoFnOperator method initializeState.

@Override
public void initializeState(StateInitializationContext context) throws Exception {
    if (getKeyedStateBackend() != null) {
        int totalKeyGroups = getKeyedStateBackend().getNumberOfKeyGroups();
        KeyGroupsList localKeyGroupRange = getKeyedStateBackend().getKeyGroupRange();
        for (KeyGroupStatePartitionStreamProvider streamProvider : context.getRawKeyedStateInputs()) {
            DataInputViewStreamWrapper div = new DataInputViewStreamWrapper(streamProvider.getStream());
            int keyGroupIdx = streamProvider.getKeyGroupId();
            checkArgument(localKeyGroupRange.contains(keyGroupIdx), "Key Group " + keyGroupIdx + " does not belong to the local range.");
            // if (this instanceof KeyGroupRestoringOperator)
            restoreKeyGroupState(keyGroupIdx, div);
            // We just initialize our timerService
            if (keyCoder != null) {
                if (timerService == null) {
                    timerService = new HeapInternalTimerService<>(totalKeyGroups, localKeyGroupRange, this, getRuntimeContext().getProcessingTimeService());
                }
                timerService.restoreTimersForKeyGroup(div, keyGroupIdx, getUserCodeClassloader());
            }
        }
    }
}

Example 3 with KeyGroupStatePartitionStreamProvider

use of org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider in project flink by apache.

the class AbstractStreamOperator method initializeState.

/**
	 * Stream operators with state which can be restored need to override this hook method.
	 *
	 * @param context context that allows to register different states.
	 */
public void initializeState(StateInitializationContext context) throws Exception {
    if (getKeyedStateBackend() != null) {
        KeyGroupsList localKeyGroupRange = getKeyedStateBackend().getKeyGroupRange();
        // and then initialize the timer services
        for (KeyGroupStatePartitionStreamProvider streamProvider : context.getRawKeyedStateInputs()) {
            int keyGroupIdx = streamProvider.getKeyGroupId();
            checkArgument(localKeyGroupRange.contains(keyGroupIdx), "Key Group " + keyGroupIdx + " does not belong to the local range.");
            timeServiceManager.restoreStateForKeyGroup(new DataInputViewStreamWrapper(streamProvider.getStream()), keyGroupIdx, getUserCodeClassloader());
        }
    }
}

Example 4 with KeyGroupStatePartitionStreamProvider

use of org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider in project flink by apache.

the class StateInitializationContextImplTest method getKeyedStateStreams.

@Test
public void getKeyedStateStreams() throws Exception {
    int readKeyGroupCount = 0;
    for (KeyGroupStatePartitionStreamProvider stateStreamProvider : initializationContext.getRawKeyedStateInputs()) {
        Assert.assertNotNull(stateStreamProvider);
        try (InputStream is = stateStreamProvider.getStream()) {
            DataInputView div = new DataInputViewStreamWrapper(is);
            int val = div.readInt();
            ++readKeyGroupCount;
            Assert.assertEquals(stateStreamProvider.getKeyGroupId(), val);
        }
    }
    Assert.assertEquals(writtenKeyGroups, readKeyGroupCount);
}

Example 5 with KeyGroupStatePartitionStreamProvider

use of org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider in project beam by apache.

the class DedupingOperator method initializeState.

@Override
public void initializeState(StateInitializationContext context) throws Exception {
    if (getKeyedStateBackend() != null) {
        KeyGroupsList localKeyGroupRange = getKeyedStateBackend().getKeyGroupRange();
        for (KeyGroupStatePartitionStreamProvider streamProvider : context.getRawKeyedStateInputs()) {
            DataInputViewStreamWrapper div = new DataInputViewStreamWrapper(streamProvider.getStream());
            int keyGroupIdx = streamProvider.getKeyGroupId();
            checkArgument(localKeyGroupRange.contains(keyGroupIdx), "Key Group " + keyGroupIdx + " does not belong to the local range.");
            // if (this instanceof KeyGroupRestoringOperator)
            restoreKeyGroupState(keyGroupIdx, div);
        }
    }
}