public List<SplitT> getAndRemoveUncheckpointedAssignment(
            int subtaskId, long restoredCheckpointId) {
        final ArrayList<SplitT> splits = new ArrayList<>();

        for (final Map.Entry<Long, Map<Integer, LinkedHashSet<SplitT>>> entry :
                assignmentsByCheckpointId.entrySet()) {
            if (entry.getKey() > restoredCheckpointId) {
                removeFromAssignment(subtaskId, entry.getValue(), splits);
            }
        }

        removeFromAssignment(subtaskId, uncheckpointedAssignments, splits);
        return splits;
    }
@Test
    void testGetAndRemoveUncheckpointedAssignment() throws Exception {
        final long checkpointId1 = 100L;
        final long checkpointId2 = 101L;
        SplitAssignmentTracker<MockSourceSplit> tracker = new SplitAssignmentTracker<>();

        // Assign some splits and take snapshot 1.
        tracker.recordSplitAssignment(getSplitsAssignment(2, 0));
        tracker.onCheckpoint(checkpointId1);

        // Assign some more splits and take snapshot 2.
        tracker.recordSplitAssignment(getSplitsAssignment(2, 3));
        tracker.onCheckpoint(checkpointId2);

        // Now assume subtask 0 has failed.
        List<MockSourceSplit> splitsToPutBack =
                tracker.getAndRemoveUncheckpointedAssignment(0, checkpointId1 - 1);
        verifyAssignment(Arrays.asList("0", "3"), splitsToPutBack);
    }