@Override
    public void initializeState(StateInitializationContext context) throws Exception {
        if (isPartitionCommitTriggerEnabled()) {
            partitionCommitPredicate =
                    PartitionCommitPredicate.create(conf, getUserCodeClassloader(), partitionKeys);
        }

        currentNewPartitions = new HashSet<>();
        newPartitions = new TreeMap<>();
        committablePartitions = new HashSet<>();
        inProgressPartitions = new HashMap<>();
        super.initializeState(context);
    }
@Test
    void testFailover() throws Exception {
        OperatorSubtaskState state;
        try (OneInputStreamOperatorTestHarness<RowData, PartitionCommitInfo> harness = create()) {
            harness.setup();
            harness.initializeEmptyState();
            harness.open();
            harness.processElement(row("1"), 0);
            harness.processElement(row("2"), 0);
            harness.processElement(row("2"), 0);
            state = harness.snapshot(1, 1);
            harness.processElement(row("3"), 0);
            harness.processElement(row("4"), 0);
            harness.notifyOfCompletedCheckpoint(1);
            List<String> partitions = collect(harness);
            assertThat(partitions).containsExactly("1", "2");
        }

        // first retry, no partition {1, 2} records
        try (OneInputStreamOperatorTestHarness<RowData, PartitionCommitInfo> harness = create()) {
            harness.setup();
            harness.initializeState(state);
            harness.open();
            harness.processElement(row("3"), 0);
            harness.processElement(row("4"), 0);
            state = harness.snapshot(2, 2);
            harness.notifyOfCompletedCheckpoint(2);
            List<String> partitions = collect(harness);
            assertThat(partitions).containsExactly("1", "2", "3", "4");
        }

        // second retry, partition {4} repeat
        try (OneInputStreamOperatorTestHarness<RowData, PartitionCommitInfo> harness = create()) {
            harness.setup();
            harness.initializeState(state);
            harness.open();
            harness.processElement(row("4"), 0);
            harness.processElement(row("5"), 0);
            state = harness.snapshot(3, 3);
            harness.notifyOfCompletedCheckpoint(3);
            List<String> partitions = collect(harness);
            assertThat(partitions).containsExactly("3", "4", "5");
        }

        // third retry, multiple snapshots
        try (OneInputStreamOperatorTestHarness<RowData, PartitionCommitInfo> harness = create()) {
            harness.setup();
            harness.initializeState(state);
            harness.open();
            harness.processElement(row("6"), 0);
            harness.processElement(row("7"), 0);
            harness.snapshot(4, 4);
            harness.processElement(row("8"), 0);
            harness.snapshot(5, 5);
            harness.processElement(row("9"), 0);
            harness.snapshot(6, 6);
            harness.notifyOfCompletedCheckpoint(5);
            List<String> partitions = collect(harness);
            // should not contains partition {9}
            assertThat(partitions).containsExactly("4", "5", "6", "7", "8");
        }
    }