Examples with CoderRegistry org.apache.beam.sdk.coders.CoderRegistry used on opensource projects

Example 1 with CoderRegistry

use of org.apache.beam.sdk.coders.CoderRegistry in project beam by apache.

the class KafkaIOTest method testInferKeyCoder.

@Test
public void testInferKeyCoder() {
    CoderRegistry registry = CoderRegistry.createDefault();
    assertTrue(KafkaIO.inferCoder(registry, LongDeserializer.class).getValueCoder() instanceof VarLongCoder);
    assertTrue(KafkaIO.inferCoder(registry, StringDeserializer.class).getValueCoder() instanceof StringUtf8Coder);
    assertTrue(KafkaIO.inferCoder(registry, InstantDeserializer.class).getValueCoder() instanceof InstantCoder);
    assertTrue(KafkaIO.inferCoder(registry, DeserializerWithInterfaces.class).getValueCoder() instanceof VarLongCoder);
}

Example 2 with CoderRegistry

use of org.apache.beam.sdk.coders.CoderRegistry in project beam by apache.

the class DoFnInvokersTest method testSplittableDoFnDefaultMethods.

@Test
public void testSplittableDoFnDefaultMethods() throws Exception {
    class MockFn extends DoFn<String, String> {

        @ProcessElement
        public void processElement(ProcessContext c, DefaultTracker tracker) {
        }

        @GetInitialRestriction
        public RestrictionWithDefaultTracker getInitialRestriction(String element) {
            return null;
        }
    }
    MockFn fn = mock(MockFn.class);
    DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
    CoderRegistry coderRegistry = CoderRegistry.createDefault();
    coderRegistry.registerCoderProvider(CoderProviders.fromStaticMethods(RestrictionWithDefaultTracker.class, CoderForDefaultTracker.class));
    assertThat(invoker.<RestrictionWithDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class));
    invoker.invokeSplitRestriction("blah", "foo", new DoFn.OutputReceiver<String>() {

        private boolean invoked;

        @Override
        public void output(String output) {
            assertFalse(invoked);
            invoked = true;
            assertEquals("foo", output);
        }
    });
    invoker.invokeProcessElement(mockArgumentProvider);
    assertThat(invoker.invokeNewTracker(new RestrictionWithDefaultTracker()), instanceOf(DefaultTracker.class));
}

Example 3 with CoderRegistry

use of org.apache.beam.sdk.coders.CoderRegistry in project beam by apache.

the class WithKeys method expand.

@Override
public PCollection<KV<K, V>> expand(PCollection<V> in) {
    PCollection<KV<K, V>> result = in.apply("AddKeys", MapElements.via(new SimpleFunction<V, KV<K, V>>() {

        @Override
        public KV<K, V> apply(V element) {
            return KV.of(fn.apply(element), element);
        }
    }));
    try {
        Coder<K> keyCoder;
        CoderRegistry coderRegistry = in.getPipeline().getCoderRegistry();
        if (keyType == null) {
            keyCoder = coderRegistry.getOutputCoder(fn, in.getCoder());
        } else {
            keyCoder = coderRegistry.getCoder(keyType);
        }
        // TODO: Remove when we can set the coder inference context.
        result.setCoder(KvCoder.of(keyCoder, in.getCoder()));
    } catch (CannotProvideCoderException exc) {
    // let lazy coder inference have a try
    }
    return result;
}

Example 4 with CoderRegistry

use of org.apache.beam.sdk.coders.CoderRegistry in project beam by apache.

the class ConfluentSchemaRegistryDeserializerProviderTest method testGetCoder.

@Test
public void testGetCoder() {
    String schemaRegistryUrl = "mock://my-scope-name";
    String subject = "mytopic";
    SchemaRegistryClient mockRegistryClient = mockSchemaRegistryClient(schemaRegistryUrl, subject);
    CoderRegistry coderRegistry = CoderRegistry.createDefault();
    AvroCoder coderV0 = (AvroCoder) mockDeserializerProvider(schemaRegistryUrl, subject, null).getCoder(coderRegistry);
    assertEquals(AVRO_SCHEMA, coderV0.getSchema());
    try {
        Integer version = mockRegistryClient.getVersion(subject, AVRO_SCHEMA_V1);
        AvroCoder coderV1 = (AvroCoder) mockDeserializerProvider(schemaRegistryUrl, subject, version).getCoder(coderRegistry);
        assertEquals(AVRO_SCHEMA_V1, coderV1.getSchema());
    } catch (IOException | RestClientException e) {
        throw new RuntimeException("Unable to register schema for subject: " + subject, e);
    }
}

Example 5 with CoderRegistry

use of org.apache.beam.sdk.coders.CoderRegistry in project beam by apache.

the class LocalDeserializerProviderTest method testInferKeyCoderFailure.

@Test
public void testInferKeyCoderFailure() throws Exception {
    cannotInferException.expect(RuntimeException.class);
    CoderRegistry registry = CoderRegistry.createDefault();
    LocalDeserializerProvider.of(NonInferableObjectDeserializer.class).getCoder(registry);
}