Search in sources :

Example 1 with GeneratorConfig

use of org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig in project beam by apache.

the class GeneratorTest method splitPreservesOverallEventCount.

@Test
public void splitPreservesOverallEventCount() {
    long n = 51237L;
    GeneratorConfig initialConfig = makeConfig(n);
    long expected = initialConfig.getStopEventId() - initialConfig.getStartEventId();
    List<Generator> generators = new ArrayList<>();
    for (GeneratorConfig subConfig : initialConfig.split(20)) {
        generators.add(new Generator(subConfig));
    }
    long actual = 0;
    for (Generator generator : generators) {
        actual += consume(generator);
    }
    assertEquals(expected, actual);
}
Also used : ArrayList(java.util.ArrayList) GeneratorConfig(org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig) Generator(org.apache.beam.sdk.nexmark.sources.generator.Generator) Test(org.junit.Test)

Example 2 with GeneratorConfig

use of org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig in project beam by apache.

the class BoundedEventSource method split.

@Override
public List<BoundedEventSource> split(long desiredBundleSizeBytes, PipelineOptions options) {
    NexmarkUtils.info("slitting bounded source %s into %d sub-sources", config, numEventGenerators);
    List<BoundedEventSource> results = new ArrayList<>();
    // Ignore desiredBundleSizeBytes and use numEventGenerators instead.
    for (GeneratorConfig subConfig : config.split(numEventGenerators)) {
        results.add(new BoundedEventSource(subConfig, 1));
    }
    return results;
}
Also used : ArrayList(java.util.ArrayList) GeneratorConfig(org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig)

Example 3 with GeneratorConfig

use of org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig in project beam by apache.

the class UnboundedEventSourceTest method resumeFromCheckpoint.

/**
 * Check aggressively checkpointing and resuming a reader gives us exactly the same event stream
 * as reading directly.
 */
@Ignore("TODO(BEAM-5070) Test is flaky. Fix before reenabling.")
@Test
public void resumeFromCheckpoint() throws IOException {
    Random random = new Random(297);
    int n = 47293;
    GeneratorConfig config = makeConfig(n);
    Generator modelGenerator = new Generator(config);
    EventIdChecker checker = new EventIdChecker();
    PipelineOptions options = TestPipeline.testingPipelineOptions();
    UnboundedEventSource source = new UnboundedEventSource(config, 1, 0, false);
    UnboundedReader<Event> reader = source.createReader(options, null);
    while (n > 0) {
        int m = Math.min(459 + random.nextInt(455), n);
        System.out.printf("reading %d...%n", m);
        checker.add(m, reader, modelGenerator);
        n -= m;
        System.out.printf("splitting with %d remaining...%n", n);
        CheckpointMark checkpointMark = reader.getCheckpointMark();
        reader = source.createReader(options, (GeneratorCheckpoint) checkpointMark);
    }
    assertFalse(reader.advance());
}
Also used : Random(java.util.Random) GeneratorCheckpoint(org.apache.beam.sdk.nexmark.sources.generator.GeneratorCheckpoint) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) CheckpointMark(org.apache.beam.sdk.io.UnboundedSource.CheckpointMark) Event(org.apache.beam.sdk.nexmark.model.Event) GeneratorConfig(org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig) GeneratorCheckpoint(org.apache.beam.sdk.nexmark.sources.generator.GeneratorCheckpoint) Generator(org.apache.beam.sdk.nexmark.sources.generator.Generator) Ignore(org.junit.Ignore) Test(org.junit.Test)

Example 4 with GeneratorConfig

use of org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig in project beam by apache.

the class GeneratorTest method splitAtFractionPreservesOverallEventCount.

@Test
public void splitAtFractionPreservesOverallEventCount() {
    long n = 55729L;
    GeneratorConfig initialConfig = makeConfig(n);
    long expected = initialConfig.getStopEventId() - initialConfig.getStartEventId();
    long actual = 0;
    Generator initialGenerator = new Generator(initialConfig);
    // Consume some events.
    actual += consume(5000, initialGenerator);
    // Split once.
    GeneratorConfig remainConfig1 = initialGenerator.splitAtEventId(9000L);
    Generator remainGenerator1 = new Generator(remainConfig1);
    // Consume some more events.
    actual += consume(2000, initialGenerator);
    actual += consume(3000, remainGenerator1);
    // Split again.
    GeneratorConfig remainConfig2 = remainGenerator1.splitAtEventId(30000L);
    Generator remainGenerator2 = new Generator(remainConfig2);
    // Run to completion.
    actual += consume(initialGenerator);
    actual += consume(remainGenerator1);
    actual += consume(remainGenerator2);
    assertEquals(expected, actual);
}
Also used : GeneratorConfig(org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig) Generator(org.apache.beam.sdk.nexmark.sources.generator.Generator) Test(org.junit.Test)

Example 5 with GeneratorConfig

use of org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig in project beam by apache.

the class UnboundedEventSource method split.

@Override
public List<UnboundedEventSource> split(int desiredNumSplits, PipelineOptions options) {
    LOG.trace("splitting unbounded source into {} sub-sources", numEventGenerators);
    List<UnboundedEventSource> results = new ArrayList<>();
    // Ignore desiredNumSplits and use numEventGenerators instead.
    for (GeneratorConfig subConfig : config.split(numEventGenerators)) {
        results.add(new UnboundedEventSource(subConfig, 1, watermarkHoldbackSec, isRateLimited));
    }
    return results;
}
Also used : ArrayList(java.util.ArrayList) GeneratorConfig(org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig)

Aggregations

GeneratorConfig (org.apache.beam.sdk.nexmark.sources.generator.GeneratorConfig)5 ArrayList (java.util.ArrayList)3 Generator (org.apache.beam.sdk.nexmark.sources.generator.Generator)3 Test (org.junit.Test)3 Random (java.util.Random)1 CheckpointMark (org.apache.beam.sdk.io.UnboundedSource.CheckpointMark)1 Event (org.apache.beam.sdk.nexmark.model.Event)1 GeneratorCheckpoint (org.apache.beam.sdk.nexmark.sources.generator.GeneratorCheckpoint)1 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)1 Ignore (org.junit.Ignore)1