@Override
    public KeyGroupRangeInputSplit[] createInputSplits(int minNumSplits) throws IOException {
        final int maxParallelism = operatorState.getMaxParallelism();

        final List<KeyGroupRange> keyGroups = sortedKeyGroupRanges(minNumSplits, maxParallelism);

        return CollectionUtil.mapWithIndex(
                        keyGroups,
                        (keyGroupRange, index) ->
                                createKeyGroupRangeInputSplit(
                                        operatorState, maxParallelism, keyGroupRange, index))
                .toArray(KeyGroupRangeInputSplit[]::new);
    }
@Test
    public void testCreatePartitionedInputSplits() 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),
                        new ExecutionConfig());
        KeyGroupRangeInputSplit[] splits = format.createInputSplits(4);
        Assert.assertEquals(
                "Failed to properly partition operator state into input splits", 4, splits.length);
    }