FlatMapElements extends PTransform<PCollection<? extends InputT>, PCollection<OutputT>> { public static <InputT, OutputT> FlatMapElements<InputT, OutputT> via( InferableFunction<? super InputT, ? extends Iterable<OutputT>> fn) { Contextful<Fn<InputT, Iterable<OutputT>>> wrapped = (Contextful) Contextful.fn(fn); TypeDescriptor<OutputT> outputType = TypeDescriptors.extractFromTypeParameters( (TypeDescriptor<Iterable<OutputT>>) fn.getOutputTypeDescriptor(), Iterable.class, new TypeDescriptors.TypeVariableExtractor<Iterable<OutputT>, OutputT>() {}); TypeDescriptor<InputT> inputType = (TypeDescriptor<InputT>) fn.getInputTypeDescriptor(); return new FlatMapElements<>(wrapped, fn, inputType, outputType); } private  FlatMapElements(
      @Nullable Contextful<Fn<InputT, Iterable<OutputT>>> fn,
      @Nullable Object originalFnForDisplayData,
      @Nullable TypeDescriptor<InputT> inputType,
      TypeDescriptor<OutputT> outputType); static FlatMapElements<InputT, OutputT> via(
      InferableFunction<? super InputT, ? extends Iterable<OutputT>> fn); static FlatMapElements<InputT, OutputT> via(
      SimpleFunction<? super InputT, ? extends Iterable<OutputT>> fn); static FlatMapElements<?, OutputT> into(
      final TypeDescriptor<OutputT> outputType); FlatMapElements<NewInputT, OutputT> via(
      ProcessFunction<NewInputT, ? extends Iterable<OutputT>> fn); FlatMapElements<NewInputT, OutputT> via(
      SerializableFunction<NewInputT, ? extends Iterable<OutputT>> fn); @Experimental(Kind.CONTEXTFUL) FlatMapElements<NewInputT, OutputT> via(
      Contextful<Fn<NewInputT, Iterable<OutputT>>> fn); @Override PCollection<OutputT> expand(PCollection<? extends InputT> input); @Override void populateDisplayData(DisplayData.Builder builder); @Experimental(Kind.WITH_EXCEPTIONS) FlatMapWithFailures<InputT, OutputT, NewFailureT> exceptionsInto(
      TypeDescriptor<NewFailureT> failureTypeDescriptor); @Experimental(Kind.WITH_EXCEPTIONS) FlatMapWithFailures<InputT, OutputT, FailureT> exceptionsVia(
      InferableFunction<ExceptionElement<InputT>, FailureT> exceptionHandler);  }
@Test @Category(NeedsRunner.class) public void testFlatMapSimpleFunction() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( FlatMapElements.via( new SimpleFunction<Integer, List<Integer>>() { @Override public List<Integer> apply(Integer input) { return ImmutableList.of(-input, input); } })); PAssert.that(output).containsInAnyOrder(1, -2, -1, -3, 2, 3); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testFlatMapInferableFunction() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( FlatMapElements.via( new InferableFunction<Integer, List<Integer>>() { @Override public List<Integer> apply(Integer input) throws Exception { return ImmutableList.of(-input, input); } })); PAssert.that(output).containsInAnyOrder(1, -2, -1, -3, 2, 3); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testFlatMapFnOutputTypeDescriptor() throws Exception { PCollection<String> output = pipeline .apply(Create.of("hello")) .apply( FlatMapElements.via( new SimpleFunction<String, Set<String>>() { @Override public Set<String> apply(String input) { return ImmutableSet.copyOf(input.split("")); } })); assertThat( output.getTypeDescriptor(), equalTo((TypeDescriptor<String>) new TypeDescriptor<String>() {})); assertThat( pipeline.getCoderRegistry().getCoder(output.getTypeDescriptor()), equalTo(pipeline.getCoderRegistry().getCoder(new TypeDescriptor<String>() {}))); pipeline.run(); }
@Test public void testPolymorphicSimpleFunction() throws Exception { pipeline.enableAbandonedNodeEnforcement(false); pipeline .apply(Create.of(1, 2, 3)) .apply("Polymorphic Identity", MapElements.via(new PolymorphicSimpleFunction<>())) .apply( "Test Consumer", MapElements.via( new SimpleFunction<Iterable<Integer>, Integer>() { @Override public Integer apply(Iterable<Integer> input) { return 42; } })); }
@Test public void testSimpleFunctionClassDisplayData() { SimpleFunction<Integer, List<Integer>> simpleFn = new SimpleFunction<Integer, List<Integer>>() { @Override public List<Integer> apply(Integer input) { return Collections.emptyList(); } }; FlatMapElements<?, ?> simpleMap = FlatMapElements.via(simpleFn); assertThat(DisplayData.from(simpleMap), hasDisplayItem("class", simpleFn.getClass())); }
@Test public void testInferableFunctionClassDisplayData() { InferableFunction<Integer, List<Integer>> inferableFn = new InferableFunction<Integer, List<Integer>>() { @Override public List<Integer> apply(Integer input) { return Collections.emptyList(); } }; FlatMapElements<?, ?> inferableMap = FlatMapElements.via(inferableFn); assertThat(DisplayData.from(inferableMap), hasDisplayItem("class", inferableFn.getClass())); }
@Test public void testSimpleFunctionDisplayData() { SimpleFunction<Integer, List<Integer>> simpleFn = new SimpleFunction<Integer, List<Integer>>() { @Override public List<Integer> apply(Integer input) { return Collections.emptyList(); } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo", "baz")); } }; FlatMapElements<?, ?> simpleFlatMap = FlatMapElements.via(simpleFn); assertThat(DisplayData.from(simpleFlatMap), hasDisplayItem("class", simpleFn.getClass())); assertThat(DisplayData.from(simpleFlatMap), hasDisplayItem("foo", "baz")); }
@Test public void testInferableFunctionDisplayData() { InferableFunction<Integer, List<Integer>> inferableFn = new InferableFunction<Integer, List<Integer>>() { @Override public List<Integer> apply(Integer input) { return Collections.emptyList(); } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo", "baz")); } }; FlatMapElements<?, ?> inferableFlatMap = FlatMapElements.via(inferableFn); assertThat(DisplayData.from(inferableFlatMap), hasDisplayItem("class", inferableFn.getClass())); assertThat(DisplayData.from(inferableFlatMap), hasDisplayItem("foo", "baz")); }