Examples with BundleApplication org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication used on opensource projects

Example 1 with BundleApplication

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication in project beam by apache.

the class FnApiDoFnRunner method constructSplitResult.

@VisibleForTesting
static <WatermarkEstimatorStateT> HandlesSplits.SplitResult constructSplitResult(WindowedSplitResult windowedSplitResult, HandlesSplits.SplitResult downstreamElementSplit, Coder fullInputCoder, Instant initialWatermark, KV<Instant, WatermarkEstimatorStateT> watermarkAndState, String pTransformId, String mainInputId, Collection<String> outputIds, Duration resumeDelay) {
    // The element split cannot from both windowedSplitResult and downstreamElementSplit.
    checkArgument((windowedSplitResult == null || windowedSplitResult.getResidualSplitRoot() == null) || downstreamElementSplit == null);
    List<BundleApplication> primaryRoots = new ArrayList<>();
    List<DelayedBundleApplication> residualRoots = new ArrayList<>();
    // Encode window splits.
    if (windowedSplitResult != null && windowedSplitResult.getPrimaryInFullyProcessedWindowsRoot() != null) {
        ByteString.Output primaryInOtherWindowsBytes = ByteString.newOutput();
        try {
            fullInputCoder.encode(windowedSplitResult.getPrimaryInFullyProcessedWindowsRoot(), primaryInOtherWindowsBytes);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        BundleApplication.Builder primaryApplicationInOtherWindows = BundleApplication.newBuilder().setTransformId(pTransformId).setInputId(mainInputId).setElement(primaryInOtherWindowsBytes.toByteString());
        primaryRoots.add(primaryApplicationInOtherWindows.build());
    }
    if (windowedSplitResult != null && windowedSplitResult.getResidualInUnprocessedWindowsRoot() != null) {
        ByteString.Output bytesOut = ByteString.newOutput();
        try {
            fullInputCoder.encode(windowedSplitResult.getResidualInUnprocessedWindowsRoot(), bytesOut);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        BundleApplication.Builder residualInUnprocessedWindowsRoot = BundleApplication.newBuilder().setTransformId(pTransformId).setInputId(mainInputId).setElement(bytesOut.toByteString());
        // We don't want to change the output watermarks or set the checkpoint resume time since
        // that applies to the current window.
        Map<String, org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Timestamp> outputWatermarkMapForUnprocessedWindows = new HashMap<>();
        if (!initialWatermark.equals(GlobalWindow.TIMESTAMP_MIN_VALUE)) {
            org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Timestamp outputWatermark = org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Timestamp.newBuilder().setSeconds(initialWatermark.getMillis() / 1000).setNanos((int) (initialWatermark.getMillis() % 1000) * 1000000).build();
            for (String outputId : outputIds) {
                outputWatermarkMapForUnprocessedWindows.put(outputId, outputWatermark);
            }
        }
        residualInUnprocessedWindowsRoot.putAllOutputWatermarks(outputWatermarkMapForUnprocessedWindows);
        residualRoots.add(DelayedBundleApplication.newBuilder().setApplication(residualInUnprocessedWindowsRoot).build());
    }
    ByteString.Output primaryBytes = ByteString.newOutput();
    ByteString.Output residualBytes = ByteString.newOutput();
    // that there is no element split.
    if (windowedSplitResult != null && windowedSplitResult.getResidualSplitRoot() != null) {
        // When there is element split in windowedSplitResult, the resumeDelay should not be null.
        checkNotNull(resumeDelay);
        try {
            fullInputCoder.encode(windowedSplitResult.getPrimarySplitRoot(), primaryBytes);
            fullInputCoder.encode(windowedSplitResult.getResidualSplitRoot(), residualBytes);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        primaryRoots.add(BundleApplication.newBuilder().setTransformId(pTransformId).setInputId(mainInputId).setElement(primaryBytes.toByteString()).build());
        BundleApplication.Builder residualApplication = BundleApplication.newBuilder().setTransformId(pTransformId).setInputId(mainInputId).setElement(residualBytes.toByteString());
        Map<String, org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Timestamp> outputWatermarkMap = new HashMap<>();
        if (!watermarkAndState.getKey().equals(GlobalWindow.TIMESTAMP_MIN_VALUE)) {
            org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Timestamp outputWatermark = org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Timestamp.newBuilder().setSeconds(watermarkAndState.getKey().getMillis() / 1000).setNanos((int) (watermarkAndState.getKey().getMillis() % 1000) * 1000000).build();
            for (String outputId : outputIds) {
                outputWatermarkMap.put(outputId, outputWatermark);
            }
        }
        residualApplication.putAllOutputWatermarks(outputWatermarkMap);
        residualRoots.add(DelayedBundleApplication.newBuilder().setApplication(residualApplication).setRequestedTimeDelay(Durations.fromMillis(resumeDelay.getMillis())).build());
    } else if (downstreamElementSplit != null) {
        primaryRoots.add(Iterables.getOnlyElement(downstreamElementSplit.getPrimaryRoots()));
        residualRoots.add(Iterables.getOnlyElement(downstreamElementSplit.getResidualRoots()));
    }
    return HandlesSplits.SplitResult.of(primaryRoots, residualRoots);
}