public <InputT, OutputT, CollectionT extends PCollection<? extends InputT>>
      DataStream<OutputT> applyBeamPTransform(
          DataStream<InputT> input, PTransform<CollectionT, PCollection<OutputT>> transform) {
    return (DataStream)
        getNonNull(
            applyBeamPTransformInternal(
                ImmutableMap.of("input", input),
                (pipeline, map) -> (CollectionT) getNonNull(map, "input"),
                (output) -> ImmutableMap.of("output", output),
                transform,
                input.getExecutionEnvironment()),
            "output");
  }
@Test
  public void testApplyPreservesInputTimestamps() throws Exception {
    StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironment();
    env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);

    DataStream<Long> input =
        env.fromCollection(ImmutableList.of(1L, 2L, 12L))
            .assignTimestampsAndWatermarks(
                WatermarkStrategy.<Long>forBoundedOutOfOrderness(java.time.Duration.ofMillis(100))
                    .withTimestampAssigner(
                        (SerializableTimestampAssigner<Long>)
                            (element, recordTimestamp) -> element));
    DataStream<Long> result =
        new BeamFlinkDataStreamAdapter()
            .applyBeamPTransform(
                input,
                new PTransform<PCollection<Long>, PCollection<Long>>() {
                  @Override
                  public PCollection<Long> expand(PCollection<Long> input) {
                    return input
                        .apply(Window.into(FixedWindows.of(Duration.millis(10))))
                        .apply(Sum.longsGlobally().withoutDefaults());
                  }
                });

    assertThat(ImmutableList.copyOf(result.executeAndCollect()), containsInAnyOrder(3L, 12L));
  }