@Override public <T> Schema schemaFor(TypeDescriptor<T> typeDescriptor) { if (typeDescriptor.equals(TypeDescriptors.strings())) { return STRING_SCHEMA; } else if (typeDescriptor.equals(TypeDescriptors.integers())) { return INTEGER_SCHEMA; } else { return null; } }
/** * Basic test of {@link MapElements} with a lambda wrapped into a {@link SimpleFunction} to * remember its type. */ @Test public void testGoodTypeForMethodRef() throws Exception { SimpleFunction<Integer, String> fn = new SimpleFunction<Integer, String>(SimpleFunctionTest::toStringThisThing) {}; assertThat(fn.getInputTypeDescriptor(), equalTo(TypeDescriptors.integers())); assertThat(fn.getOutputTypeDescriptor(), equalTo(TypeDescriptors.strings())); }
/** * Basic test of {@link MapElements} with a lambda (which is instantiated as a {@link * SerializableFunction}). */ @Test public void testGoodTypeForLambda() throws Exception { SimpleFunction<Integer, String> fn = new SimpleFunction<Integer, String>(Object::toString) {}; assertThat(fn.getInputTypeDescriptor(), equalTo(TypeDescriptors.integers())); assertThat(fn.getOutputTypeDescriptor(), equalTo(TypeDescriptors.strings())); }
@Test public void testTypeDescriptorsKV() throws Exception { TypeDescriptor<KV<String, Integer>> descriptor = kvs(strings(), integers()); assertEquals(descriptor, new TypeDescriptor<KV<String, Integer>>() {}); }
@Override public PCollection<KV<String, Integer>> expand(PCollection<GameActionInfo> gameInfo) { return gameInfo .apply( MapElements.into( TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers())) .via((GameActionInfo gInfo) -> KV.of(gInfo.getKey(field), gInfo.getScore()))) .apply(Sum.integersPerKey()); } }
@Test public void testSerializableFunctionDisplayData() { SerializableFunction<Integer, Integer> serializableFn = input -> input; MapElements<?, ?> serializableMap = MapElements.into(integers()).via(serializableFn); assertThat( DisplayData.from(serializableMap), hasDisplayItem("class", serializableFn.getClass())); }
@Test public void testTypeDescriptorsTypeParameterOf() throws Exception { assertEquals(strings(), extractFooT(new Generic<String, Integer>() {})); assertEquals(integers(), extractBarT(new Generic<String, Integer>() {})); assertEquals(kvs(strings(), integers()), extractKV(new Generic<String, Integer>() {})); }
@Override public PCollection<KV<Integer, KV<KeyT, ValueT>>> expand(PCollection<KV<KeyT, ValueT>> input) { return input .apply( "AssignTask", ParDo.of(new AssignTaskFn<KeyT, ValueT>(configView)).withSideInputs(configView)) .setTypeDescriptor( TypeDescriptors.kvs(TypeDescriptors.integers(), input.getTypeDescriptor())) .apply("GroupByTaskId", GroupByKey.create()) .apply("FlattenGroupedTasks", ParDo.of(new FlattenGroupedTasks<>())); } }
@Test @Category({ValidatesRunner.class, UsesImpulse.class}) public void testImpulse() { PCollection<Integer> result = p.apply(Impulse.create()) .apply( FlatMapElements.into(TypeDescriptors.integers()) .via(impulse -> Arrays.asList(1, 2, 3))); PAssert.that(result).containsInAnyOrder(1, 2, 3); p.run().waitUntilFinish(); } }
/** * Basic test of {@link MapElements} with a {@link SerializableFunction}. This style is generally * discouraged in Java 7, in favor of {@link SimpleFunction}. */ @Test @Category(NeedsRunner.class) public void testMapBasicSerializableFunction() throws Exception { PCollection<Integer> output = pipeline.apply(Create.of(1, 2, 3)).apply(MapElements.into(integers()).via(input -> -input)); PAssert.that(output).containsInAnyOrder(-2, -1, -3); pipeline.run(); }
/** * Basic test of {@link MapElements} with a lambda (which is instantiated as a {@link * SerializableFunction}). */ @Test @Category(NeedsRunner.class) public void testMapLambda() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( MapElements // Note that the type annotation is required. .into(TypeDescriptors.integers()) .via((Integer i) -> i * 2)); PAssert.that(output).containsInAnyOrder(6, 2, 4); pipeline.run(); }
/** Basic test of {@link FlatMapElements} with a method reference. */ @Test @Category(NeedsRunner.class) public void testFlatMapMethodReference() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( FlatMapElements // Note that the input type annotation is required. .into(TypeDescriptors.integers()) .via(new Negater()::numAndNegation)); PAssert.that(output).containsInAnyOrder(1, 3, -1, -3, 2, -2); pipeline.run(); }
/** Basic test of {@link MapElements} with a method reference. */ @Test @Category(NeedsRunner.class) public void testMapMethodReference() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( MapElements // Note that the type annotation is required. .into(TypeDescriptors.integers()) .via(new Doubler()::doubleIt)); PAssert.that(output).containsInAnyOrder(6, 2, 4); pipeline.run(); }
/** * Basic test of {@link FlatMapElements} with a lambda (which is instantiated as a {@link * SerializableFunction}). */ @Test @Category(NeedsRunner.class) public void testFlatMapBasicWithLambda() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( FlatMapElements // Note that the input type annotation is required. .into(TypeDescriptors.integers()) .via((Integer i) -> ImmutableList.of(i, -i))); PAssert.that(output).containsInAnyOrder(1, 3, -1, -3, 2, -2); pipeline.run(); }
@Test public void testRegisterForType() throws NoSuchSchemaException { SchemaRegistry registry = SchemaRegistry.createDefault(); registry.registerSchemaForType( TypeDescriptors.strings(), STRING_SCHEMA, s -> Row.withSchema(STRING_SCHEMA).addValue(s).build(), r -> r.getString("string")); registry.registerSchemaForType( TypeDescriptors.integers(), INTEGER_SCHEMA, s -> Row.withSchema(INTEGER_SCHEMA).addValue(s).build(), r -> r.getInt32("integer")); tryGetters(registry); }
/** Basic test of {@link MapElements} with a {@link Fn} and a side input. */ @Test @Category(NeedsRunner.class) public void testMapBasicWithSideInput() throws Exception { final PCollectionView<Integer> view = pipeline.apply("Create base", Create.of(40)).apply(View.asSingleton()); PCollection<Integer> output = pipeline .apply(Create.of(0, 1, 2)) .apply( MapElements.into(integers()) .via( fn((element, c) -> element + c.sideInput(view), requiresSideInputs(view)))); PAssert.that(output).containsInAnyOrder(40, 41, 42); pipeline.run(); }
@Test @Category(NeedsRunner.class) public void testComposeBasicSerializableFunction() throws Exception { PCollection<Integer> output = pipeline .apply(Create.of(1, 2, 3)) .apply( PTransform.compose( (PCollection<Integer> numbers) -> { PCollection<Integer> inverted = numbers.apply(MapElements.into(integers()).via(input -> -input)); return PCollectionList.of(numbers) .and(inverted) .apply(Flatten.pCollections()); })); PAssert.that(output).containsInAnyOrder(-2, -1, -3, 2, 1, 3); pipeline.run(); } }
private void tryGetters(SchemaRegistry registry) throws NoSuchSchemaException { assertEquals(STRING_SCHEMA, registry.getSchema(String.class)); assertEquals(STRING_SCHEMA, registry.getSchema(TypeDescriptors.strings())); assertEquals( Row.withSchema(STRING_SCHEMA).addValue("foobar").build(), registry.getToRowFunction(String.class).apply("foobar")); assertEquals( Row.withSchema(STRING_SCHEMA).addValue("foobar").build(), registry.getToRowFunction(TypeDescriptors.strings()).apply("foobar")); assertEquals(INTEGER_SCHEMA, registry.getSchema(Integer.class)); assertEquals(INTEGER_SCHEMA, registry.getSchema(TypeDescriptors.integers())); assertEquals( Row.withSchema(INTEGER_SCHEMA).addValue(42).build(), registry.getToRowFunction(Integer.class).apply(42)); assertEquals( Row.withSchema(INTEGER_SCHEMA).addValue(42).build(), registry.getToRowFunction(TypeDescriptors.integers()).apply(42)); thrown.expect(NoSuchSchemaException.class); registry.getSchema(Double.class); }
/** Test that bad input data is dropped appropriately. */ @Test @Category(ValidatesRunner.class) public void testUserScoresBadInput() throws Exception { PCollection<String> input = p.apply(Create.of(GAME_EVENTS2).withCoder(StringUtf8Coder.of())); PCollection<KV<String, Integer>> extract = input .apply(ParDo.of(new ParseEventFn())) .apply( MapElements.into( TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers())) .via((GameActionInfo gInfo) -> KV.of(gInfo.getUser(), gInfo.getScore()))); PAssert.that(extract).empty(); p.run().waitUntilFinish(); } }
/** Test the filtering. */ @Test @Category(ValidatesRunner.class) public void testUserScoresFilter() throws Exception { final Instant startMinTimestamp = new Instant(1447965680000L); PCollection<String> input = p.apply(Create.of(GAME_EVENTS).withCoder(StringUtf8Coder.of())); PCollection<KV<String, Integer>> output = input .apply("ParseGameEvent", ParDo.of(new ParseEventFn())) .apply( "FilterStartTime", Filter.by( (GameActionInfo gInfo) -> gInfo.getTimestamp() > startMinTimestamp.getMillis())) // run a map to access the fields in the result. .apply( MapElements.into( TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers())) .via((GameActionInfo gInfo) -> KV.of(gInfo.getUser(), gInfo.getScore()))); PAssert.that(output).containsInAnyOrder(FILTERED_EVENTS); p.run().waitUntilFinish(); }