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

Example 21 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()));
}

Example 22 with CoderRegistry

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

the class CombineValuesFnFactoryTest method testCombineValuesFnCoders.

@Test
public void testCombineValuesFnCoders() throws Exception {
    CoderRegistry registry = CoderRegistry.createDefault();
    MeanInts meanInts = new MeanInts();
    CountSum countSum = new CountSum(6, 27);
    Coder<CountSum> coder = meanInts.getAccumulatorCoder(registry, registry.getCoder(TypeDescriptor.of(Integer.class)));
    assertEquals(countSum, CoderUtils.decodeFromByteArray(coder, CoderUtils.encodeToByteArray(coder, countSum)));
}

Example 23 with CoderRegistry

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

the class StreamingKeyedWorkItemSideInputDoFnRunnerTest method createRunner.

@SuppressWarnings("unchecked")
private StreamingKeyedWorkItemSideInputDoFnRunner<String, Integer, KV<String, Integer>, IntervalWindow> createRunner(DoFnRunners.OutputManager outputManager) throws Exception {
    CoderRegistry registry = CoderRegistry.createDefault();
    Coder<String> keyCoder = StringUtf8Coder.of();
    Coder<Integer> inputCoder = BigEndianIntegerCoder.of();
    AppliedCombineFn<String, Integer, ?, Integer> combineFn = AppliedCombineFn.withInputCoder(Sum.ofIntegers(), registry, KvCoder.of(keyCoder, inputCoder));
    WindowingStrategy<Object, IntervalWindow> windowingStrategy = WindowingStrategy.of(WINDOW_FN);
    @SuppressWarnings("rawtypes") StreamingGroupAlsoByWindowViaWindowSetFn doFn = (StreamingGroupAlsoByWindowViaWindowSetFn) StreamingGroupAlsoByWindowsDoFns.create(windowingStrategy, key -> state, combineFn, keyCoder);
    DoFnRunner<KeyedWorkItem<String, Integer>, KV<String, Integer>> simpleDoFnRunner = new GroupAlsoByWindowFnRunner<>(PipelineOptionsFactory.create(), doFn.asDoFn(), mockSideInputReader, outputManager, mainOutputTag, stepContext);
    return new StreamingKeyedWorkItemSideInputDoFnRunner<String, Integer, KV<String, Integer>, IntervalWindow>(simpleDoFnRunner, keyCoder, sideInputFetcher, stepContext);
}

Example 24 with CoderRegistry

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

the class DoFnInvokersTest method testSplittableDoFnWithHasDefaultMethods.

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

        @ProcessElement
        public void processElement(ProcessContext c, RestrictionTracker<RestrictionWithBoundedDefaultTracker, Void> tracker, WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator> watermarkEstimator) {
        }

        @GetInitialRestriction
        public RestrictionWithBoundedDefaultTracker getInitialRestriction(@Element String element) {
            return null;
        }

        @GetInitialWatermarkEstimatorState
        public WatermarkEstimatorStateWithDefaultWatermarkEstimator getInitialWatermarkEstimatorState() {
            return null;
        }
    }
    MockFn fn = mock(MockFn.class);
    DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
    CoderRegistry coderRegistry = CoderRegistry.createDefault();
    coderRegistry.registerCoderProvider(CoderProviders.fromStaticMethods(RestrictionWithBoundedDefaultTracker.class, CoderForDefaultTracker.class));
    coderRegistry.registerCoderForClass(WatermarkEstimatorStateWithDefaultWatermarkEstimator.class, new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator());
    assertThat(invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class));
    assertThat(invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry), instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class));
    invoker.invokeSplitRestriction(new FakeArgumentProvider<String, String>() {

        @Override
        public String element(DoFn<String, String> doFn) {
            return "blah";
        }

        @Override
        public Object restriction() {
            return "foo";
        }

        @Override
        public OutputReceiver<String> outputReceiver(DoFn<String, String> doFn) {
            return new DoFn.OutputReceiver<String>() {

                private boolean invoked;

                @Override
                public void output(String output) {
                    assertFalse(invoked);
                    invoked = true;
                    assertEquals("foo", output);
                }

                @Override
                public void outputWithTimestamp(String output, Instant instant) {
                    assertFalse(invoked);
                    invoked = true;
                    assertEquals("foo", output);
                }
            };
        }
    });
    assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider));
    assertThat(invoker.invokeNewTracker(new FakeArgumentProvider<String, String>() {

        @Override
        public Object restriction() {
            return new RestrictionWithBoundedDefaultTracker();
        }
    }), instanceOf(BoundedDefaultTracker.class));
    assertThat(invoker.invokeNewWatermarkEstimator(new FakeArgumentProvider<String, String>() {

        @Override
        public Object watermarkEstimatorState() {
            return new WatermarkEstimatorStateWithDefaultWatermarkEstimator();
        }
    }), instanceOf(DefaultWatermarkEstimator.class));
}