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

Example 1 with DesiredSplit

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

the class BeamFnDataReadRunner method trySplit.

public void trySplit(ProcessBundleSplitRequest request, ProcessBundleSplitResponse.Builder response) {
    DesiredSplit desiredSplit = request.getDesiredSplitsMap().get(pTransformId);
    if (desiredSplit == null) {
        return;
    }
    long totalBufferSize = desiredSplit.getEstimatedInputElements();
    List<Long> allowedSplitPoints = new ArrayList<>(desiredSplit.getAllowedSplitPointsList());
    HandlesSplits splittingConsumer = null;
    if (consumer instanceof HandlesSplits) {
        splittingConsumer = ((HandlesSplits) consumer);
    }
    synchronized (splittingLock) {
        // provide.
        if (index == stopIndex) {
            return;
        }
        // being released.
        if (!request.getInstructionId().equals(processBundleInstructionIdSupplier.get())) {
            return;
        }
        // split request is bounded incorrectly, use the stop index as the upper bound.
        if (totalBufferSize < index + 1) {
            totalBufferSize = index + 1;
        } else if (totalBufferSize > stopIndex) {
            totalBufferSize = stopIndex;
        }
        // In the case where we have yet to process an element, set the current element progress to 1.
        double currentElementProgress = 1;
        // progress defaulting to 0.5 if no progress was able to get fetched.
        if (index >= 0) {
            if (splittingConsumer != null) {
                currentElementProgress = splittingConsumer.getProgress();
            } else {
                currentElementProgress = 0.5;
            }
        }
        // Now figure out where to split.
        // 
        // The units here (except for keepOfElementRemainder) are all in terms of number or
        // (possibly fractional) elements.
        // Compute the amount of "remaining" work that we know of.
        double remainder = totalBufferSize - index - currentElementProgress;
        // Compute the number of elements (including fractional elements) that we should "keep".
        double keep = remainder * desiredSplit.getFractionOfRemainder();
        // splittable.
        if (currentElementProgress < 1) {
            // See if the amount we need to keep falls within the current element's remainder and if
            // so, attempt to split it.
            double keepOfElementRemainder = keep / (1 - currentElementProgress);
            // If both index and index are allowed split point, we can split at index.
            if (keepOfElementRemainder < 1 && isValidSplitPoint(allowedSplitPoints, index) && isValidSplitPoint(allowedSplitPoints, index + 1)) {
                SplitResult splitResult = splittingConsumer != null ? splittingConsumer.trySplit(keepOfElementRemainder) : null;
                if (splitResult != null) {
                    stopIndex = index + 1;
                    response.addAllPrimaryRoots(splitResult.getPrimaryRoots()).addAllResidualRoots(splitResult.getResidualRoots()).addChannelSplitsBuilder().setLastPrimaryElement(index - 1).setFirstResidualElement(stopIndex);
                    return;
                }
            }
        }
        // Otherwise, split at the closest allowed element boundary.
        long newStopIndex = index + Math.max(1, Math.round(currentElementProgress + keep));
        if (!isValidSplitPoint(allowedSplitPoints, newStopIndex)) {
            // Choose the closest allowed split point.
            Collections.sort(allowedSplitPoints);
            int closestSplitPointIndex = -(Collections.binarySearch(allowedSplitPoints, newStopIndex) + 1);
            if (closestSplitPointIndex == 0) {
                newStopIndex = allowedSplitPoints.get(0);
            } else if (closestSplitPointIndex == allowedSplitPoints.size()) {
                newStopIndex = allowedSplitPoints.get(closestSplitPointIndex - 1);
            } else {
                long prevPoint = allowedSplitPoints.get(closestSplitPointIndex - 1);
                long nextPoint = allowedSplitPoints.get(closestSplitPointIndex);
                if (index < prevPoint && newStopIndex - prevPoint < nextPoint - newStopIndex) {
                    newStopIndex = prevPoint;
                } else {
                    newStopIndex = nextPoint;
                }
            }
        }
        if (newStopIndex < stopIndex && newStopIndex > index) {
            stopIndex = newStopIndex;
            response.addChannelSplitsBuilder().setLastPrimaryElement(stopIndex - 1).setFirstResidualElement(stopIndex);
            return;
        }
    }
}