Java Code Examples for org.apache.beam.sdk.util.UserCodeException#wrap()

private ParDoEvaluator(
    PushbackSideInputDoFnRunner<InputT, ?> fnRunner,
    AppliedPTransform<?, ?, ?> transform,
    BundleOutputManager outputManager,
    DirectStepContext stepContext) {
  this.fnRunner = fnRunner;
  this.transform = transform;
  this.outputManager = outputManager;
  this.stepContext = stepContext;
  this.unprocessedElements = ImmutableList.builder();

  try {
    fnRunner.startBundle();
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
}
 

@Override
public TransformResult<InputT> finishBundle() {
  try {
    fnRunner.finishBundle();
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
  StepTransformResult.Builder<InputT> resultBuilder;
  CopyOnAccessInMemoryStateInternals state = stepContext.commitState();
  if (state != null) {
    resultBuilder =
        StepTransformResult.<InputT>withHold(transform, state.getEarliestWatermarkHold())
            .withState(state);
  } else {
    resultBuilder = StepTransformResult.withoutHold(transform);
  }
  return resultBuilder
      .addOutput(outputManager.bundles.values())
      .withTimerUpdate(stepContext.getTimerUpdate())
      .addUnprocessedElements(unprocessedElements.build())
      .build();
}
 

@Override
public void processElement(WindowedValue<KV<K, Iterable<AccumT>>> element) throws Exception {
  checkState(
      element.getWindows().size() == 1,
      "Expected inputs to %s to be in exactly one window. Got %s",
      MergeAccumulatorsAndExtractOutputEvaluator.class.getSimpleName(),
      element.getWindows().size());
  Iterable<AccumT> inputAccumulators = element.getValue().getValue();
  try {
    AccumT first = combineFn.createAccumulator();
    AccumT merged =
        combineFn.mergeAccumulators(
            Iterables.concat(
                Collections.singleton(first),
                inputAccumulators,
                Collections.singleton(combineFn.createAccumulator())));
    OutputT extracted = combineFn.extractOutput(merged);
    output.add(element.withValue(KV.of(element.getValue().getKey(), extracted)));
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
}
 

/** Outputs the given element to the specified set of consumers wrapping any exceptions. */
private <T> void outputTo(
    Collection<FnDataReceiver<WindowedValue<T>>> consumers, WindowedValue<T> output) {
  if (currentWatermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
    ((TimestampObservingWatermarkEstimator) currentWatermarkEstimator)
        .observeTimestamp(output.getTimestamp());
  }
  try {
    for (FnDataReceiver<WindowedValue<T>> consumer : consumers) {
      consumer.accept(output);
    }
  } catch (Throwable t) {
    throw UserCodeException.wrap(t);
  }
}
 

public SerializableArrayViaCoder(Coder<T> elementCoder, T[] value) {
  this.coder = ListCoder.of(elementCoder);
  this.value = value;
  try {
    this.encodedValue = CoderUtils.encodeToByteArray(coder, Arrays.asList(value));
  } catch (CoderException exc) {
    throw UserCodeException.wrap(exc);
  }
}
 

private ValueT toValue(WindowedValue<InputT> input) {
  try {
    return toValue.call(input.getValue());
  } catch (Exception ex) {
    throw UserCodeException.wrap(ex);
  }
}
 

private ValueT toValue(WindowedValue<InputT> value) {
  try {
    return toValue.call(value.getValue());
  } catch (Exception ex) {
    throw UserCodeException.wrap(ex);
  }
}
 

@Override
public void processElement(WindowedValue<BoundedSourceShard<OutputT>> element)
    throws Exception {
  BoundedSource<OutputT> source = element.getValue().getSource();
  try (final BoundedReader<OutputT> reader = source.createReader(options)) {
    boolean contentsRemaining = reader.start();
    Future<BoundedSource<OutputT>> residualFuture = startDynamicSplitThread(source, reader);
    UncommittedBundle<OutputT> output = evaluationContext.createBundle(outputPCollection);
    while (contentsRemaining) {
      output.add(
          WindowedValue.timestampedValueInGlobalWindow(
              reader.getCurrent(), reader.getCurrentTimestamp()));
      contentsRemaining = reader.advance();
    }
    resultBuilder.addOutput(output);
    try {
      BoundedSource<OutputT> residual = residualFuture.get();
      if (residual != null) {
        resultBuilder.addUnprocessedElements(
            element.withValue(BoundedSourceShard.of(residual)));
      }
    } catch (ExecutionException exex) {
      // Un-and-rewrap the exception thrown by attempting to split
      throw UserCodeException.wrap(exex.getCause());
    }
  }
}
 

@Override
public void processElement(WindowedValue<InputT> element) {
  try {
    Iterable<WindowedValue<InputT>> unprocessed = fnRunner.processElementInReadyWindows(element);
    unprocessedElements.addAll(unprocessed);
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
}
 

public <KeyT> void onTimer(TimerData timer, KeyT key, BoundedWindow window) {
  try {
    fnRunner.onTimer(
        timer.getTimerId(),
        timer.getTimerFamilyId(),
        key,
        window,
        timer.getTimestamp(),
        timer.getOutputTimestamp(),
        timer.getDomain());
  } catch (Exception e) {
    throw UserCodeException.wrap(e);
  }
}
 

@Override
public void beforeElement(WindowedValue<T> element) {
  try {
    mutationElements.put(
        element, MutationDetectors.forValueWithCoder(element.getValue(), coder));
  } catch (CoderException e) {
    throw UserCodeException.wrap(e);
  }
}
 

@Override
public UncommittedBundle<T> add(WindowedValue<T> element) {
  try {
    // Use the cloned value to ensure that if the coder behaves poorly (e.g. a NoOpCoder that
    // does not expect to be used) that is reflected in the values given to downstream
    // transforms
    WindowedValue<T> clone = element.withValue(CoderUtils.clone(coder, element.getValue()));
    underlying.add(clone);
  } catch (CoderException e) {
    throw UserCodeException.wrap(e);
  }
  return this;
}
 

private void output(Instant scheduledTime) {
  if (outputTimestamp != null) {
    checkArgument(
        !outputTimestamp.isBefore(elementTimestampOrTimerHoldTimestamp),
        "output timestamp %s should be after input message timestamp or output timestamp of firing timers %s",
        outputTimestamp,
        elementTimestampOrTimerHoldTimestamp);
  }

  // Output timestamp is set to the delivery time if not initialized by an user.
  if (outputTimestamp == null && TimeDomain.EVENT_TIME.equals(timeDomain)) {
    outputTimestamp = scheduledTime;
  }

  // For processing timers
  if (outputTimestamp == null) {
    // For processing timers output timestamp will be:
    // 1) timestamp of input element
    // OR
    // 2) hold timestamp of firing timer.
    outputTimestamp = elementTimestampOrTimerHoldTimestamp;
  }

  Instant windowExpiry = LateDataUtils.garbageCollectionTime(currentWindow, allowedLateness);
  if (TimeDomain.EVENT_TIME.equals(timeDomain)) {
    checkArgument(
        !outputTimestamp.isAfter(scheduledTime),
        "Attempted to set an event-time timer with an output timestamp of %s that is"
            + " after the timer firing timestamp %s",
        outputTimestamp,
        scheduledTime);
    checkArgument(
        !scheduledTime.isAfter(windowExpiry),
        "Attempted to set an event-time timer with a firing timestamp of %s that is"
            + " after the expiration of window %s",
        scheduledTime,
        windowExpiry);
  } else {
    checkArgument(
        !outputTimestamp.isAfter(windowExpiry),
        "Attempted to set a processing-time timer with an output timestamp of %s that is"
            + " after the expiration of window %s",
        outputTimestamp,
        windowExpiry);
  }

  TimerHandler<K> consumer = (TimerHandler) timerHandlers.get(timerId);
  try {
    consumer.accept(
        Timer.of(
            userKey,
            dynamicTimerTag,
            Collections.singletonList(boundedWindow),
            scheduledTime,
            outputTimestamp,
            paneInfo));
  } catch (Throwable t) {
    throw UserCodeException.wrap(t);
  }
}
 

@Override
public PipelineResult run(Pipeline pipeline) {
  Throwable t = new IllegalStateException("user code exception");
  throw UserCodeException.wrap(t);
}
 

@Override
// TODO: [BEAM-4563] Pass Future back to consumer to check for async errors
@SuppressWarnings("FutureReturnValueIgnored")
public void start(DirectGraph graph, RootProviderRegistry rootProviderRegistry) {
  int numTargetSplits = Math.max(3, targetParallelism);
  ImmutableMap.Builder<AppliedPTransform<?, ?, ?>, ConcurrentLinkedQueue<CommittedBundle<?>>>
      pendingRootBundles = ImmutableMap.builder();
  for (AppliedPTransform<?, ?, ?> root : graph.getRootTransforms()) {
    ConcurrentLinkedQueue<CommittedBundle<?>> pending = new ConcurrentLinkedQueue<>();
    try {
      Collection<CommittedBundle<?>> initialInputs =
          rootProviderRegistry.getInitialInputs(root, numTargetSplits);
      pending.addAll(initialInputs);
    } catch (Exception e) {
      throw UserCodeException.wrap(e);
    }
    pendingRootBundles.put(root, pending);
  }
  evaluationContext.initialize(pendingRootBundles.build());
  final ExecutionDriver executionDriver =
      QuiescenceDriver.create(
          evaluationContext, graph, this, visibleUpdates, pendingRootBundles.build());
  executorService.submit(
      new Runnable() {
        @Override
        public void run() {
          DriverState drive = executionDriver.drive();
          if (drive.isTermainal()) {
            State newPipelineState = State.UNKNOWN;
            switch (drive) {
              case FAILED:
                newPipelineState = State.FAILED;
                break;
              case SHUTDOWN:
                newPipelineState = State.DONE;
                break;
              case CONTINUE:
                throw new IllegalStateException(
                    String.format("%s should not be a terminal state", DriverState.CONTINUE));
              default:
                throw new IllegalArgumentException(
                    String.format("Unknown %s %s", DriverState.class.getSimpleName(), drive));
            }
            shutdownIfNecessary(newPipelineState);
          } else {
            executorService.submit(this);
          }
        }
      });
}