Apache Beam UnboundedSource that repeats one item.Inspired by: https://github.com/eljefe6a/beamexample/blob/master/BeamTutorial/src/main/java/org/apache/beam/examples/tutorial/game/injector/InjectorUnboundedSource.java

RepeatedSource.java

package org.apache.beam.examples;

import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.PCollection;
import org.joda.time.Instant;

import javax.annotation.Nullable;
import java.io.IOException;
import java.io.Serializable;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Random;
import java.util.stream.Collectors;
import java.util.stream.Stream;


/**
 * Created by logicalguess on 2/20/17.
 */
public class RepeatedSource<E> extends UnboundedSource<E, RepeatedSource.Checkpoint> {

    public RepeatedSource(E item, Coder<E> coder) {
        this.item = item;
        this.coder = coder;
    }

    public static class Checkpoint implements UnboundedSource.CheckpointMark, Serializable {

        private final Random random;

        public Checkpoint(Random random) {
            this.random = SerializableUtils.clone(random);
        }

        @Override
        public void finalizeCheckpoint() throws IOException {
            // Nothing is necessary for checkpoints.
        }

        public Random getRandom() {
            return random;
        }
    }

    private static class RepeteadReader<E> extends UnboundedReader<E> {
        private final RepeatedSource<E> source;
        private final Random random;

        private RepeteadReader(RepeatedSource<E> source, @Nullable Checkpoint initialCheckpoint) {
            this.source = source;
            random = initialCheckpoint == null ? new Random() : initialCheckpoint.getRandom();
        }


        @Override
        public boolean start() throws IOException {
            return advance();
        }

        @Override
        public boolean advance() throws IOException {
            return true;
        }

        @Override
        public E getCurrent() throws NoSuchElementException {
            return source.item;
        }

        @Override
        public Instant getCurrentTimestamp() throws NoSuchElementException {
            return Instant.now(); //TODO
        }

        @Override
        public void close() throws IOException {

        }

        @Override
        public Instant getWatermark() {
            return Instant.now();
        }

        @Override
        public CheckpointMark getCheckpointMark() {
            return new Checkpoint(random);
        }

        @Override
        public UnboundedSource<E, ?> getCurrentSource() {
            return source;
        }
    }

    private final E item;
    private final Coder<E> coder;

    @Override
    public List<UnboundedSource<E, Checkpoint>> generateInitialSplits(int i, PipelineOptions pipelineOptions)
            throws Exception {
        return Stream.generate(() -> new RepeatedSource<E>(item, coder)).limit(i).collect(Collectors.toList());
    }

    @Override
    public UnboundedReader<E> createReader(PipelineOptions pipelineOptions, @Nullable Checkpoint checkpoint)
            throws IOException {
        return new RepeteadReader<E>(this, checkpoint);
    }

    @Nullable
    @Override
    public Coder<Checkpoint> getCheckpointMarkCoder() {
        return SerializableCoder.of(Checkpoint.class);
    }

    @Override
    public void validate() {
        //nothing to validate
    }

    @Override
    public Coder<E> getDefaultOutputCoder() {
        return coder;
    }

    private static class DebugFunction<E> extends SimpleFunction<E, E> {
        @Override
        public E apply(E in) {
            System.out.println(in);
            return in;
        }
    }

    public static void main(String[] args) {
        TestPipeline p = TestPipeline.create();
        PCollection<String> input = p.apply(Read.from(new RepeatedSource<String>("abc", StringUtf8Coder.of())));
        PCollection<String> output = input.apply(MapElements.via(new DebugFunction<String>()));
        p.run().waitUntilFinish();
    }
}