MapElements extends PTransform<PCollection<? extends InputT>, PCollection<OutputT>> { public static <InputT, OutputT> MapElements<InputT, OutputT> via( final InferableFunction<InputT, OutputT> fn) { return new MapElements<>( Contextful.fn(fn), fn, fn.getInputTypeDescriptor(), fn.getOutputTypeDescriptor()); } private  MapElements(
      @Nullable Contextful<Fn<InputT, OutputT>> fn,
      @Nullable Object originalFnForDisplayData,
      @Nullable TypeDescriptor<InputT> inputType,
      TypeDescriptor<OutputT> outputType); static MapElements<InputT, OutputT> via(
      final InferableFunction<InputT, OutputT> fn); static MapElements<InputT, OutputT> via(
      final SimpleFunction<InputT, OutputT> fn); static MapElements<?, OutputT> into(final TypeDescriptor<OutputT> outputType); MapElements<NewInputT, OutputT> via(ProcessFunction<NewInputT, OutputT> fn); MapElements<NewInputT, OutputT> via(
      SerializableFunction<NewInputT, OutputT> fn); @Experimental(Kind.CONTEXTFUL) MapElements<NewInputT, OutputT> via(Contextful<Fn<NewInputT, OutputT>> fn); @Override PCollection<OutputT> expand(PCollection<? extends InputT> input); @Override void populateDisplayData(DisplayData.Builder builder); @Experimental(Kind.WITH_EXCEPTIONS) MapWithFailures<InputT, OutputT, NewFailureT> exceptionsInto(
      TypeDescriptor<NewFailureT> failureTypeDescriptor); @Experimental(Kind.WITH_EXCEPTIONS) MapWithFailures<InputT, OutputT, FailureT> exceptionsVia(
      InferableFunction<ExceptionElement<InputT>, FailureT> exceptionHandler);  }
@Test @Category(NeedsRunner.class) public void testMapSimpleFunction() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( MapElements.via( new SimpleFunction<Integer, Integer>() { @Override public Integer apply(Integer input) { return -input; } })); PAssert.that(output).containsInAnyOrder(-2, -1, -3); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testMapInferableFunction() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( MapElements.via( new InferableFunction<Integer, Integer>() { @Override public Integer apply(Integer input) throws Exception { return -input; } })); PAssert.that(output).containsInAnyOrder(-2, -1, -3); 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<Integer, Integer>() { @Override public Integer apply(Integer input) { return input; } })); }
@Test public void testPolymorphicInferableFunction() throws Exception { pipeline.enableAbandonedNodeEnforcement(false); pipeline .apply(Create.of(1, 2, 3)) .apply("Polymorphic Identity", MapElements.via(new PolymorphicInferableFunction<>())) .apply( "Test Consumer", MapElements.via( new InferableFunction<Integer, Integer>() { @Override public Integer apply(Integer input) throws Exception { return input; } })); }
@Test public void testNestedPolymorphicSimpleFunction() throws Exception { pipeline.enableAbandonedNodeEnforcement(false); pipeline .apply(Create.of(1, 2, 3)) .apply("Polymorphic Identity", MapElements.via(new NestedPolymorphicSimpleFunction<>())) .apply( "Test Consumer", MapElements.via( new SimpleFunction<KV<Integer, String>, Integer>() { @Override public Integer apply(KV<Integer, String> input) { return 42; } })); }
@Test public void testNestedPolymorphicInferableFunction() throws Exception { pipeline.enableAbandonedNodeEnforcement(false); pipeline .apply(Create.of(1, 2, 3)) .apply("Polymorphic Identity", MapElements.via(new NestedPolymorphicInferableFunction<>())) .apply( "Test Consumer", MapElements.via( new InferableFunction<KV<Integer, String>, Integer>() { @Override public Integer apply(KV<Integer, String> input) throws Exception { return 42; } })); }
@Test @Category(NeedsRunner.class) public void testSimpleFunctionOutputTypeDescriptor() throws Exception { PCollection<String> output = pipeline .apply(Create.of("hello")) .apply( MapElements.via( new SimpleFunction<String, String>() { @Override public String apply(String input) { return input; } })); 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 @Category(NeedsRunner.class) public void testInferableFunctionOutputTypeDescriptor() throws Exception { PCollection<String> output = pipeline .apply(Create.of("hello")) .apply( MapElements.via( new InferableFunction<String, String>() { @Override public String apply(String input) throws Exception { return input; } })); 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 testSimpleFunctionClassDisplayData() { SimpleFunction<?, ?> simpleFn = new SimpleFunction<Integer, Integer>() { @Override public Integer apply(Integer input) { return input; } }; MapElements<?, ?> simpleMap = MapElements.via(simpleFn); assertThat(DisplayData.from(simpleMap), hasDisplayItem("class", simpleFn.getClass())); }
@Test public void testInferableFunctionClassDisplayData() { InferableFunction<?, ?> inferableFn = new InferableFunction<Integer, Integer>() { @Override public Integer apply(Integer input) throws Exception { return input; } }; MapElements<?, ?> inferableMap = MapElements.via(inferableFn); assertThat(DisplayData.from(inferableMap), hasDisplayItem("class", inferableFn.getClass())); }
@Test public void testSimpleFunctionDisplayData() { SimpleFunction<Integer, ?> simpleFn = new SimpleFunction<Integer, Integer>() { @Override public Integer apply(Integer input) { return input; } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo", "baz")); } }; MapElements<?, ?> simpleMap = MapElements.via(simpleFn); assertThat(DisplayData.from(simpleMap), hasDisplayItem("class", simpleFn.getClass())); assertThat(DisplayData.from(simpleMap), hasDisplayItem("foo", "baz")); }
@Test public void testInferableFunctionDisplayData() { InferableFunction<Integer, ?> inferableFn = new InferableFunction<Integer, Integer>() { @Override public Integer apply(Integer input) { return input; } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo", "baz")); } }; MapElements<?, ?> inferableMap = MapElements.via(inferableFn); assertThat(DisplayData.from(inferableMap), hasDisplayItem("class", inferableFn.getClass())); assertThat(DisplayData.from(inferableMap), hasDisplayItem("foo", "baz")); }
@Test @Category(ValidatesRunner.class) public void testPrimitiveDisplayData() { SimpleFunction<Integer, ?> mapFn = new SimpleFunction<Integer, Integer>() { @Override public Integer apply(Integer input) { return input; } }; MapElements<Integer, ?> map = MapElements.via(mapFn); DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(map); assertThat( "MapElements should include the mapFn in its primitive display data", displayData, hasItem(hasDisplayItem("class", mapFn.getClass()))); }
@Test @Category(NeedsRunner.class) public void testMapWrappedLambda() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply(MapElements.via(new SimpleFunction<Integer, Integer>((Integer i) -> i * 2) {})); PAssert.that(output).containsInAnyOrder(6, 2, 4); pipeline.run(); }