@Override
    public void applyToAllPartitions(ApplyPartitionFunction<OUT> applyPartitionFunction)
            throws Exception {
        if (isKeyed) {
            for (Object key : keySet) {
                partitionedContext
                        .getStateManager()
                        .executeInKeyContext(
                                () -> {
                                    try {
                                        applyPartitionFunction.apply(collector, partitionedContext);
                                    } catch (Exception e) {
                                        throw new RuntimeException(e);
                                    }
                                },
                                key);
            }
        } else {
            // non-keyed operator has only one partition.
            applyPartitionFunction.apply(collector, partitionedContext);
        }
    }
@Test
    void testApplyToAllPartitions() throws Exception {
        AtomicInteger counter = new AtomicInteger(0);
        List<Integer> collectedData = new ArrayList<>();

        TestingTimestampCollector<Integer> collector =
                TestingTimestampCollector.<Integer>builder()
                        .setCollectConsumer(collectedData::add)
                        .build();
        CompletableFuture<Void> cf = new CompletableFuture<>();
        StreamingRuntimeContext operatorRuntimeContext =
                ContextTestUtils.createStreamingRuntimeContext();
        DefaultRuntimeContext runtimeContext =
                new DefaultRuntimeContext(
                        operatorRuntimeContext.getJobInfo().getJobName(),
                        operatorRuntimeContext.getJobType(),
                        1,
                        2,
                        "mock-task",
                        operatorRuntimeContext.getMetricGroup());

        DefaultNonPartitionedContext<Integer> nonPartitionedContext =
                new DefaultNonPartitionedContext<>(
                        runtimeContext,
                        new DefaultPartitionedContext(
                                runtimeContext,
                                Optional::empty,
                                (key) -> cf.complete(null),
                                UnsupportedProcessingTimeManager.INSTANCE,
                                ContextTestUtils.createStreamingRuntimeContext(),
                                new MockOperatorStateStore()),
                        collector,
                        false,
                        null);
        nonPartitionedContext.applyToAllPartitions(
                (out, ctx) -> {
                    counter.incrementAndGet();
                    out.collect(10);
                });
        assertThat(counter.get()).isEqualTo(1);
        assertThat(cf).isNotCompleted();
        assertThat(collectedData).containsExactly(10);
    }