BoundedSourceSystem { private static <T> List<BoundedSource<T>> split( BoundedSource<T> source, SamzaPipelineOptions pipelineOptions) throws Exception { final int numSplits = pipelineOptions.getMaxSourceParallelism(); if (numSplits > 1) { final long estimatedSize = source.getEstimatedSizeBytes(pipelineOptions); final long bundleSize = (estimatedSize + numSplits - 1) / numSplits; @SuppressWarnings("unchecked") final List<BoundedSource<T>> splits = (List<BoundedSource<T>>) source.split(bundleSize, pipelineOptions); if (!splits.isEmpty()) { return splits; } } return Collections.singletonList(source); }    }
@Test public void testSplit() throws IOException, InterruptedException { final TestBoundedSource.SplittableBuilder<String> builder = TestBoundedSource.<String>createSplits(3); builder.forSplit(0).addElements("split-0"); builder.forSplit(1).addElements("split-1"); builder.forSplit(2).addElements("split-2"); final TestBoundedSource<String> source = builder.build(); final BoundedSourceSystem.Consumer<String> consumer = createConsumer(source, 3); consumer.register(ssp(0), NULL_STRING); consumer.register(ssp(1), NULL_STRING); consumer.register(ssp(2), NULL_STRING); consumer.start(); final Set<String> offsets = new HashSet<>(); List<IncomingMessageEnvelope> envelopes = consumeUntilTimeoutOrEos(consumer, ssp(0), DEFAULT_TIMEOUT_MILLIS); assertEquals( Arrays.asList( createElementMessage( ssp(0), envelopes.get(0).getOffset(), "split-0", BoundedWindow.TIMESTAMP_MIN_VALUE), createWatermarkMessage(ssp(0), BoundedWindow.TIMESTAMP_MAX_VALUE), createEndOfStreamMessage(ssp(0))), envelopes); offsets.add(envelopes.get(0).getOffset()); envelopes = consumeUntilTimeoutOrEos(consumer, ssp(1), DEFAULT_TIMEOUT_MILLIS); assertEquals( Arrays.asList( createElementMessage( ssp(1), envelopes.get(0).getOffset(), "split-1", BoundedWindow.TIMESTAMP_MIN_VALUE), createWatermarkMessage(ssp(1), BoundedWindow.TIMESTAMP_MAX_VALUE), createEndOfStreamMessage(ssp(1))), envelopes); offsets.add(envelopes.get(0).getOffset()); envelopes = consumeUntilTimeoutOrEos(consumer, ssp(2), DEFAULT_TIMEOUT_MILLIS); assertEquals( Arrays.asList( createElementMessage( ssp(2), envelopes.get(0).getOffset(), "split-2", BoundedWindow.TIMESTAMP_MIN_VALUE), createWatermarkMessage(ssp(2), BoundedWindow.TIMESTAMP_MAX_VALUE), createEndOfStreamMessage(ssp(2))), envelopes); offsets.add(envelopes.get(0).getOffset()); assertEquals(Sets.newHashSet("0", "1", "2"), offsets); consumer.stop(); }