Examples with Source org.apache.beam.sdk.io.Source used on opensource projects

Example 1 with Source

use of org.apache.beam.sdk.io.Source in project beam by apache.

the class CustomSources method serializeToCloudSource.

public static com.google.api.services.dataflow.model.Source serializeToCloudSource(Source<?> source, PipelineOptions options) throws Exception {
    com.google.api.services.dataflow.model.Source cloudSource = new com.google.api.services.dataflow.model.Source();
    // We ourselves act as the SourceFormat.
    cloudSource.setSpec(CloudObject.forClass(CustomSources.class));
    addString(cloudSource.getSpec(), SERIALIZED_SOURCE, encodeBase64String(serializeToByteArray(source)));
    SourceMetadata metadata = new SourceMetadata();
    if (source instanceof BoundedSource) {
        BoundedSource<?> boundedSource = (BoundedSource<?>) source;
        // Size estimation is best effort so we continue even if it fails here.
        try {
            metadata.setEstimatedSizeBytes(boundedSource.getEstimatedSizeBytes(options));
        } catch (Exception e) {
            LOG.warn("Size estimation of the source failed: " + source, e);
        }
    } else if (source instanceof UnboundedSource) {
        UnboundedSource<?, ?> unboundedSource = (UnboundedSource<?, ?>) source;
        metadata.setInfinite(true);
        List<String> encodedSplits = new ArrayList<>();
        int desiredNumSplits = getDesiredNumUnboundedSourceSplits(options.as(DataflowPipelineOptions.class));
        for (UnboundedSource<?, ?> split : unboundedSource.split(desiredNumSplits, options)) {
            encodedSplits.add(encodeBase64String(serializeToByteArray(split)));
        }
        checkArgument(!encodedSplits.isEmpty(), "UnboundedSources must have at least one split");
        addStringList(cloudSource.getSpec(), SERIALIZED_SOURCE_SPLITS, encodedSplits);
    } else {
        throw new IllegalArgumentException("Unexpected source kind: " + source.getClass());
    }
    cloudSource.setMetadata(metadata);
    return cloudSource;
}

Example 2 with Source

use of org.apache.beam.sdk.io.Source in project beam by apache.

the class StateSpecFunctions method mapSourceFunction.

/**
   * A {@link org.apache.spark.streaming.StateSpec} function to support reading from
   * an {@link UnboundedSource}.
   *
   * <p>This StateSpec function expects the following:
   * <ul>
   * <li>Key: The (partitioned) Source to read from.</li>
   * <li>Value: An optional {@link UnboundedSource.CheckpointMark} to start from.</li>
   * <li>State: A byte representation of the (previously) persisted CheckpointMark.</li>
   * </ul>
   * And returns an iterator over all read values (for the micro-batch).
   *
   * <p>This stateful operation could be described as a flatMap over a single-element stream, which
   * outputs all the elements read from the {@link UnboundedSource} for this micro-batch.
   * Since micro-batches are bounded, the provided UnboundedSource is wrapped by a
   * {@link MicrobatchSource} that applies bounds in the form of duration and max records
   * (per micro-batch).
   *
   *
   * <p>In order to avoid using Spark Guava's classes which pollute the
   * classpath, we use the {@link StateSpec#function(scala.Function3)} signature which employs
   * scala's native {@link scala.Option}, instead of the
   * {@link StateSpec#function(org.apache.spark.api.java.function.Function3)} signature,
   * which employs Guava's {@link com.google.common.base.Optional}.
   *
   * <p>See also <a href="https://issues.apache.org/jira/browse/SPARK-4819">SPARK-4819</a>.</p>
   *
   * @param runtimeContext    A serializable {@link SparkRuntimeContext}.
   * @param <T>               The type of the input stream elements.
   * @param <CheckpointMarkT> The type of the {@link UnboundedSource.CheckpointMark}.
   * @return The appropriate {@link org.apache.spark.streaming.StateSpec} function.
   */
public static <T, CheckpointMarkT extends UnboundedSource.CheckpointMark> scala.Function3<Source<T>, scala.Option<CheckpointMarkT>, State<Tuple2<byte[], Instant>>, Tuple2<Iterable<byte[]>, Metadata>> mapSourceFunction(final SparkRuntimeContext runtimeContext, final String stepName) {
    return new SerializableFunction3<Source<T>, Option<CheckpointMarkT>, State<Tuple2<byte[], Instant>>, Tuple2<Iterable<byte[]>, Metadata>>() {

        @Override
        public Tuple2<Iterable<byte[]>, Metadata> apply(Source<T> source, scala.Option<CheckpointMarkT> startCheckpointMark, State<Tuple2<byte[], Instant>> state) {
            MetricsContainerStepMap metricsContainers = new MetricsContainerStepMap();
            MetricsContainer metricsContainer = metricsContainers.getContainer(stepName);
            // since they may report metrics.
            try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metricsContainer)) {
                // source as MicrobatchSource
                MicrobatchSource<T, CheckpointMarkT> microbatchSource = (MicrobatchSource<T, CheckpointMarkT>) source;
                // Initial high/low watermarks.
                Instant lowWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
                final Instant highWatermark;
                // if state exists, use it, otherwise it's first time so use the startCheckpointMark.
                // startCheckpointMark may be EmptyCheckpointMark (the Spark Java API tries to apply
                // Optional(null)), which is handled by the UnboundedSource implementation.
                Coder<CheckpointMarkT> checkpointCoder = microbatchSource.getCheckpointMarkCoder();
                CheckpointMarkT checkpointMark;
                if (state.exists()) {
                    // previous (output) watermark is now the low watermark.
                    lowWatermark = state.get()._2();
                    checkpointMark = CoderHelpers.fromByteArray(state.get()._1(), checkpointCoder);
                    LOG.info("Continue reading from an existing CheckpointMark.");
                } else if (startCheckpointMark.isDefined() && !startCheckpointMark.get().equals(EmptyCheckpointMark.get())) {
                    checkpointMark = startCheckpointMark.get();
                    LOG.info("Start reading from a provided CheckpointMark.");
                } else {
                    checkpointMark = null;
                    LOG.info("No CheckpointMark provided, start reading from default.");
                }
                // create reader.
                final MicrobatchSource.Reader /*<T>*/
                microbatchReader;
                final Stopwatch stopwatch = Stopwatch.createStarted();
                long readDurationMillis = 0;
                try {
                    microbatchReader = (MicrobatchSource.Reader) microbatchSource.getOrCreateReader(runtimeContext.getPipelineOptions(), checkpointMark);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
                // read microbatch as a serialized collection.
                final List<byte[]> readValues = new ArrayList<>();
                WindowedValue.FullWindowedValueCoder<T> coder = WindowedValue.FullWindowedValueCoder.of(source.getDefaultOutputCoder(), GlobalWindow.Coder.INSTANCE);
                try {
                    // measure how long a read takes per-partition.
                    boolean finished = !microbatchReader.start();
                    while (!finished) {
                        final WindowedValue<T> wv = WindowedValue.of((T) microbatchReader.getCurrent(), microbatchReader.getCurrentTimestamp(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING);
                        readValues.add(CoderHelpers.toByteArray(wv, coder));
                        finished = !microbatchReader.advance();
                    }
                    // end-of-read watermark is the high watermark, but don't allow decrease.
                    final Instant sourceWatermark = microbatchReader.getWatermark();
                    highWatermark = sourceWatermark.isAfter(lowWatermark) ? sourceWatermark : lowWatermark;
                    readDurationMillis = stopwatch.stop().elapsed(TimeUnit.MILLISECONDS);
                    LOG.info("Source id {} spent {} millis on reading.", microbatchSource.getId(), readDurationMillis);
                    // if the Source does not supply a CheckpointMark skip updating the state.
                    @SuppressWarnings("unchecked") final CheckpointMarkT finishedReadCheckpointMark = (CheckpointMarkT) microbatchReader.getCheckpointMark();
                    byte[] codedCheckpoint = new byte[0];
                    if (finishedReadCheckpointMark != null) {
                        codedCheckpoint = CoderHelpers.toByteArray(finishedReadCheckpointMark, checkpointCoder);
                    } else {
                        LOG.info("Skipping checkpoint marking because the reader failed to supply one.");
                    }
                    // persist the end-of-read (high) watermark for following read, where it will become
                    // the next low watermark.
                    state.update(new Tuple2<>(codedCheckpoint, highWatermark));
                } catch (IOException e) {
                    throw new RuntimeException("Failed to read from reader.", e);
                }
                final ArrayList<byte[]> payload = Lists.newArrayList(Iterators.unmodifiableIterator(readValues.iterator()));
                return new Tuple2<>((Iterable<byte[]>) payload, new Metadata(readValues.size(), lowWatermark, highWatermark, readDurationMillis, metricsContainers));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    };
}

Example 3 with Source

use of org.apache.beam.sdk.io.Source in project beam by apache.

the class SparkUnboundedSource method read.

public static <T, CheckpointMarkT extends CheckpointMark> UnboundedDataset<T> read(JavaStreamingContext jssc, SparkRuntimeContext rc, UnboundedSource<T, CheckpointMarkT> source, String stepName) {
    SparkPipelineOptions options = rc.getPipelineOptions().as(SparkPipelineOptions.class);
    Long maxRecordsPerBatch = options.getMaxRecordsPerBatch();
    SourceDStream<T, CheckpointMarkT> sourceDStream = new SourceDStream<>(jssc.ssc(), source, rc, maxRecordsPerBatch);
    JavaPairInputDStream<Source<T>, CheckpointMarkT> inputDStream = JavaPairInputDStream$.MODULE$.fromInputDStream(sourceDStream, JavaSparkContext$.MODULE$.<Source<T>>fakeClassTag(), JavaSparkContext$.MODULE$.<CheckpointMarkT>fakeClassTag());
    // call mapWithState to read from a checkpointable sources.
    JavaMapWithStateDStream<Source<T>, CheckpointMarkT, Tuple2<byte[], Instant>, Tuple2<Iterable<byte[]>, Metadata>> mapWithStateDStream = inputDStream.mapWithState(StateSpec.function(StateSpecFunctions.<T, CheckpointMarkT>mapSourceFunction(rc, stepName)).numPartitions(sourceDStream.getNumPartitions()));
    // set checkpoint duration for read stream, if set.
    checkpointStream(mapWithStateDStream, options);
    // report the number of input elements for this InputDStream to the InputInfoTracker.
    int id = inputDStream.inputDStream().id();
    JavaDStream<Metadata> metadataDStream = mapWithStateDStream.map(new Tuple2MetadataFunction());
    // register ReadReportDStream to report information related to this read.
    new ReadReportDStream(metadataDStream.dstream(), id, getSourceName(source, id), stepName).register();
    // output the actual (deserialized) stream.
    WindowedValue.FullWindowedValueCoder<T> coder = WindowedValue.FullWindowedValueCoder.of(source.getDefaultOutputCoder(), GlobalWindow.Coder.INSTANCE);
    JavaDStream<WindowedValue<T>> readUnboundedStream = mapWithStateDStream.flatMap(new Tuple2byteFlatMapFunction()).map(CoderHelpers.fromByteFunction(coder));
    return new UnboundedDataset<>(readUnboundedStream, Collections.singletonList(id));
}