public <InputT, OutputT, CollectionT extends PCollection<? extends InputT>>
      DataSet<OutputT> applyBeamPTransform(
          DataSet<InputT> input, PTransform<CollectionT, PCollection<OutputT>> transform) {
    return (DataSet)
        getNonNull(
            applyBeamPTransformInternal(
                ImmutableMap.of("input", input),
                (pipeline, map) -> (CollectionT) getNonNull(map, "input"),
                (output) -> ImmutableMap.of("output", output),
                transform,
                input.getExecutionEnvironment()),
            "output");
  }
@Test
  public void testApplyGroupingTransform() throws Exception {
    ExecutionEnvironment env = ExecutionEnvironment.createCollectionsEnvironment();

    DataSet<String> input = env.fromCollection(ImmutableList.of("a", "a", "b"));
    DataSet<KV<String, Long>> result =
        new BeamFlinkDataSetAdapter().applyBeamPTransform(input, Count.perElement());

    assertThat(result.collect(), containsInAnyOrder(KV.of("a", 2L), KV.of("b", 1L)));
  }