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

Example 6 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 7 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 8 with CoderRegistry

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

the class ReduceFnTester method combining.

/**
   * Creates a {@link ReduceFnTester} for the given {@link WindowingStrategy} and
   * {@link CombineFn}, creating a {@link TriggerStateMachine} from the
   * {@link Trigger} in the {@link WindowingStrategy}.
   */
public static <W extends BoundedWindow, AccumT, OutputT> ReduceFnTester<Integer, OutputT, W> combining(WindowingStrategy<?, W> strategy, CombineFn<Integer, AccumT, OutputT> combineFn, Coder<OutputT> outputCoder) throws Exception {
    CoderRegistry registry = CoderRegistry.createDefault();
    AppliedCombineFn<String, Integer, AccumT, OutputT> fn = AppliedCombineFn.<String, Integer, AccumT, OutputT>withInputCoder(combineFn, registry, KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
    return combining(strategy, TriggerStateMachines.stateMachineForTrigger(TriggerTranslation.toProto(strategy.getTrigger())), combineFn, outputCoder);
}

Example 9 with CoderRegistry

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

the class ReduceFnTester method combining.

/**
   * Creates a {@link ReduceFnTester} for the given {@link WindowingStrategy},
   * {@link CombineFn}, and {@link TriggerStateMachine}, for mocking the interaction
   * between {@link ReduceFnRunner} and the {@link TriggerStateMachine}.
   * Ignores the {@link Trigger} in the {@link WindowingStrategy}.
   */
public static <W extends BoundedWindow, AccumT, OutputT> ReduceFnTester<Integer, OutputT, W> combining(WindowingStrategy<?, W> strategy, TriggerStateMachine triggerStateMachine, CombineFn<Integer, AccumT, OutputT> combineFn, Coder<OutputT> outputCoder) throws Exception {
    CoderRegistry registry = CoderRegistry.createDefault();
    AppliedCombineFn<String, Integer, AccumT, OutputT> fn = AppliedCombineFn.<String, Integer, AccumT, OutputT>withInputCoder(combineFn, registry, KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
    return new ReduceFnTester<Integer, OutputT, W>(strategy, triggerStateMachine, SystemReduceFn.<String, Integer, AccumT, OutputT, W>combining(StringUtf8Coder.of(), fn), outputCoder, PipelineOptionsFactory.create(), NullSideInputReader.empty());
}

Example 10 with CoderRegistry

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

the class CopyOnAccessInMemoryStateInternalsTest method testAccumulatorCombiningStateWithUnderlying.

@Test
public void testAccumulatorCombiningStateWithUnderlying() throws CannotProvideCoderException {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    CombineFn<Long, long[], Long> sumLongFn = Sum.ofLongs();
    StateNamespace namespace = new StateNamespaceForTest("foo");
    CoderRegistry reg = pipeline.getCoderRegistry();
    StateTag<CombiningState<Long, long[], Long>> stateTag = StateTags.combiningValue("summer", sumLongFn.getAccumulatorCoder(reg, reg.getCoder(Long.class)), sumLongFn);
    GroupingState<Long, Long> underlyingValue = underlying.state(namespace, stateTag);
    assertThat(underlyingValue.read(), equalTo(0L));
    underlyingValue.add(1L);
    assertThat(underlyingValue.read(), equalTo(1L));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    GroupingState<Long, Long> copyOnAccessState = internals.state(namespace, stateTag);
    assertThat(copyOnAccessState.read(), equalTo(1L));
    copyOnAccessState.add(4L);
    assertThat(copyOnAccessState.read(), equalTo(5L));
    assertThat(underlyingValue.read(), equalTo(1L));
    GroupingState<Long, Long> reReadUnderlyingValue = underlying.state(namespace, stateTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}