Examples with KeyedStateReaderFunction org.apache.flink.state.api.functions.KeyedStateReaderFunction used on opensource projects

Example 1 with KeyedStateReaderFunction

use of org.apache.flink.state.api.functions.KeyedStateReaderFunction in project flink by apache.

the class ExistingSavepoint method readKeyedState.

/**
 * Read keyed state from an operator in a {@code Savepoint}.
 *
 * @param uid The uid of the operator.
 * @param function The {@link KeyedStateReaderFunction} that is called for each key in state.
 * @param <K> The type of the key in state.
 * @param <OUT> The output type of the transform function.
 * @return A {@code DataSet} of objects read from keyed state.
 * @throws IOException If the savepoint does not contain operator state with the given uid.
 */
public <K, OUT> DataSource<OUT> readKeyedState(String uid, KeyedStateReaderFunction<K, OUT> function) throws IOException {
    TypeInformation<K> keyTypeInfo;
    TypeInformation<OUT> outType;
    try {
        keyTypeInfo = TypeExtractor.createTypeInfo(KeyedStateReaderFunction.class, function.getClass(), 0, null, null);
    } catch (InvalidTypesException e) {
        throw new InvalidProgramException("The key type of the KeyedStateReaderFunction could not be automatically determined. Please use " + "Savepoint#readKeyedState(String, KeyedStateReaderFunction, TypeInformation, TypeInformation) instead.", e);
    }
    try {
        outType = TypeExtractor.getUnaryOperatorReturnType(function, KeyedStateReaderFunction.class, 0, 1, TypeExtractor.NO_INDEX, keyTypeInfo, Utils.getCallLocationName(), false);
    } catch (InvalidTypesException e) {
        throw new InvalidProgramException("The output type of the KeyedStateReaderFunction could not be automatically determined. Please use " + "Savepoint#readKeyedState(String, KeyedStateReaderFunction, TypeInformation, TypeInformation) instead.", e);
    }
    return readKeyedState(uid, function, keyTypeInfo, outType);
}

Example 2 with KeyedStateReaderFunction

use of org.apache.flink.state.api.functions.KeyedStateReaderFunction in project flink by apache.

the class SavepointReader method readKeyedState.

/**
 * Read keyed state from an operator in a {@code Savepoint}.
 *
 * @param uid The uid of the operator.
 * @param function The {@link KeyedStateReaderFunction} that is called for each key in state.
 * @param <K> The type of the key in state.
 * @param <OUT> The output type of the transform function.
 * @return A {@code DataStream} of objects read from keyed state.
 * @throws IOException If the savepoint does not contain operator state with the given uid.
 */
public <K, OUT> DataStream<OUT> readKeyedState(String uid, KeyedStateReaderFunction<K, OUT> function) throws IOException {
    TypeInformation<K> keyTypeInfo;
    TypeInformation<OUT> outType;
    try {
        keyTypeInfo = TypeExtractor.createTypeInfo(KeyedStateReaderFunction.class, function.getClass(), 0, null, null);
    } catch (InvalidTypesException e) {
        throw new InvalidProgramException("The key type of the KeyedStateReaderFunction could not be automatically determined. Please use " + "Savepoint#readKeyedState(String, KeyedStateReaderFunction, TypeInformation, TypeInformation) instead.", e);
    }
    try {
        outType = TypeExtractor.getUnaryOperatorReturnType(function, KeyedStateReaderFunction.class, 0, 1, TypeExtractor.NO_INDEX, keyTypeInfo, Utils.getCallLocationName(), false);
    } catch (InvalidTypesException e) {
        throw new InvalidProgramException("The output type of the KeyedStateReaderFunction could not be automatically determined. Please use " + "Savepoint#readKeyedState(String, KeyedStateReaderFunction, TypeInformation, TypeInformation) instead.", e);
    }
    return readKeyedState(uid, function, keyTypeInfo, outType);
}

Example 3 with KeyedStateReaderFunction

use of org.apache.flink.state.api.functions.KeyedStateReaderFunction in project flink by apache.

the class KeyedStateInputFormatTest method testReadState.

@Test
public void testReadState() throws Exception {
    OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
    OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
    OperatorState operatorState = new OperatorState(operatorID, 1, 128);
    operatorState.putState(0, state);
    KeyedStateInputFormat<?, ?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), new Configuration(), new KeyedStateReaderOperator<>(new ReaderFunction(), Types.INT));
    KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];
    KeyedStateReaderFunction<Integer, Integer> userFunction = new ReaderFunction();
    List<Integer> data = readInputSplit(split, userFunction);
    Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 2, 3), data);
}

Example 4 with KeyedStateReaderFunction

use of org.apache.flink.state.api.functions.KeyedStateReaderFunction in project flink by apache.

the class KeyedStateInputFormatTest method testInvalidProcessReaderFunctionFails.

@Test(expected = IOException.class)
public void testInvalidProcessReaderFunctionFails() throws Exception {
    OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
    OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
    OperatorState operatorState = new OperatorState(operatorID, 1, 128);
    operatorState.putState(0, state);
    KeyedStateInputFormat<?, ?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), new Configuration(), new KeyedStateReaderOperator<>(new ReaderFunction(), Types.INT));
    KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];
    KeyedStateReaderFunction<Integer, Integer> userFunction = new InvalidReaderFunction();
    readInputSplit(split, userFunction);
    Assert.fail("KeyedStateReaderFunction did not fail on invalid RuntimeContext use");
}

Example 5 with KeyedStateReaderFunction

use of org.apache.flink.state.api.functions.KeyedStateReaderFunction in project flink by apache.

the class KeyedStateInputFormatTest method readInputSplit.

@Nonnull
private List<Integer> readInputSplit(KeyGroupRangeInputSplit split, KeyedStateReaderFunction<Integer, Integer> userFunction) throws IOException {
    KeyedStateInputFormat<Integer, VoidNamespace, Integer> format = new KeyedStateInputFormat<>(new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4), new MemoryStateBackend(), new Configuration(), new KeyedStateReaderOperator<>(userFunction, Types.INT));
    List<Integer> data = new ArrayList<>();
    format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
    format.openInputFormat();
    format.open(split);
    while (!format.reachedEnd()) {
        data.add(format.nextRecord(0));
    }
    format.close();
    format.closeInputFormat();
    data.sort(Comparator.comparingInt(id -> id));
    return data;
}