Examples with BoundedSourceShard org.apache.beam.runners.direct.BoundedReadEvaluatorFactory.BoundedSourceShard used on opensource projects

Example 1 with BoundedSourceShard

use of org.apache.beam.runners.direct.BoundedReadEvaluatorFactory.BoundedSourceShard in project beam by apache.

the class BoundedReadEvaluatorFactoryTest method boundedSourceInMemoryTransformEvaluatorShardsOfSource.

@Test
public void boundedSourceInMemoryTransformEvaluatorShardsOfSource() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    List<? extends BoundedSource<Long>> splits = source.split(source.getEstimatedSizeBytes(options) / 2, options);
    UncommittedBundle<BoundedSourceShard<Long>> rootBundle = bundleFactory.createRootBundle();
    for (BoundedSource<Long> split : splits) {
        BoundedSourceShard<Long> shard = BoundedSourceShard.of(split);
        rootBundle.add(WindowedValue.valueInGlobalWindow(shard));
    }
    CommittedBundle<BoundedSourceShard<Long>> shards = rootBundle.commit(Instant.now());
    TransformEvaluator<BoundedSourceShard<Long>> evaluator = factory.forApplication(longsProducer, shards);
    for (WindowedValue<BoundedSourceShard<Long>> shard : shards.getElements()) {
        UncommittedBundle<Long> outputBundle = bundleFactory.createBundle(longs);
        when(context.createBundle(longs)).thenReturn(outputBundle);
        evaluator.processElement(shard);
    }
    TransformResult<?> result = evaluator.finishBundle();
    assertThat(Iterables.size(result.getOutputBundles()), equalTo(splits.size()));
    List<WindowedValue<?>> outputElems = new ArrayList<>();
    for (UncommittedBundle<?> outputBundle : result.getOutputBundles()) {
        CommittedBundle<?> outputs = outputBundle.commit(Instant.now());
        for (WindowedValue<?> outputElem : outputs.getElements()) {
            outputElems.add(outputElem);
        }
    }
    assertThat(outputElems, Matchers.<WindowedValue<?>>containsInAnyOrder(gw(1L), gw(2L), gw(4L), gw(8L), gw(9L), gw(7L), gw(6L), gw(5L), gw(3L), gw(0L)));
}