SparkExecutableStageFunction implements FlatMapFunction<Iterator<WindowedValue<InputT>>, RawUnionValue> { @Override public Iterator<RawUnionValue> call(Iterator<WindowedValue<InputT>> inputs) throws Exception { if (!inputs.hasNext()) { return Collections.emptyIterator(); } try (ExecutableStageContext stageContext = contextFactory.get(jobInfo)) { ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); try (StageBundleFactory stageBundleFactory = stageContext.getStageBundleFactory(executableStage)) { ConcurrentLinkedQueue<RawUnionValue> collector = new ConcurrentLinkedQueue<>(); StateRequestHandler stateRequestHandler = getStateRequestHandler( executableStage, stageBundleFactory.getProcessBundleDescriptor()); if (executableStage.getTimers().size() == 0) { ReceiverFactory receiverFactory = new ReceiverFactory(collector, outputMap); processElements( executableStage, stateRequestHandler, receiverFactory, null, stageBundleFactory, inputs); return collector.iterator(); } final InMemoryTimerInternals timerInternals = new InMemoryTimerInternals(); timerInternals.advanceProcessingTime(Instant.now()); timerInternals.advanceSynchronizedProcessingTime(Instant.now()); ReceiverFactory receiverFactory = new ReceiverFactory(collector, outputMap); TimerReceiverFactory timerReceiverFactory = new TimerReceiverFactory( stageBundleFactory, (Timer<?> timer, TimerInternals.TimerData timerData) -> { currentTimerKey = timer.getUserKey(); if (timer.getClearBit()) { timerInternals.deleteTimer(timerData); } else { timerInternals.setTimer(timerData); } }, windowCoder); processElements( executableStage, stateRequestHandler, receiverFactory, timerReceiverFactory, stageBundleFactory, inputs); timerInternals.advanceInputWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE); timerInternals.advanceProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE); timerInternals.advanceSynchronizedProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE); while (timerInternals.hasPendingTimers()) { try (RemoteBundle bundle = stageBundleFactory.getBundle( receiverFactory, timerReceiverFactory, stateRequestHandler, getBundleProgressHandler())) { PipelineTranslatorUtils.fireEligibleTimers( timerInternals, bundle.getTimerReceivers(), currentTimerKey); } } return collector.iterator(); } } } SparkExecutableStageFunction(
      RunnerApi.ExecutableStagePayload stagePayload,
      JobInfo jobInfo,
      Map<String, Integer> outputMap,
      SparkExecutableStageContextFactory contextFactory,
      Map<String, Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<SideInputT>>> sideInputs,
      MetricsContainerStepMapAccumulator metricsAccumulator,
      Coder windowCoder); @Override Iterator<RawUnionValue> call(Iterator<WindowedValue<InputT>> inputs);  }
@Test(expected = Exception.class) public void sdkErrorsSurfaceOnClose() throws Exception { SparkExecutableStageFunction<Integer, ?> function = getFunction(Collections.emptyMap()); doThrow(new Exception()).when(remoteBundle).close(); List<WindowedValue<Integer>> inputs = new ArrayList<>(); inputs.add(WindowedValue.valueInGlobalWindow(0)); function.call(inputs.iterator()); }
@Test public void expectedInputsAreSent() throws Exception { SparkExecutableStageFunction<Integer, ?> function = getFunction(Collections.emptyMap()); RemoteBundle bundle = Mockito.mock(RemoteBundle.class); when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle); @SuppressWarnings("unchecked") FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class); when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of(inputId, receiver)); WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1); WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2); WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3); function.call(Arrays.asList(one, two, three).iterator()); verify(receiver).accept(one); verify(receiver).accept(two); verify(receiver).accept(three); verifyNoMoreInteractions(receiver); }
@Test public void outputsAreTaggedCorrectly() throws Exception { WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3); WindowedValue<Integer> four = WindowedValue.valueInGlobalWindow(4); WindowedValue<Integer> five = WindowedValue.valueInGlobalWindow(5); Map<String, Integer> outputTagMap = ImmutableMap.of( "one", 1, "two", 2, "three", 3); StageBundleFactory stageBundleFactory = new StageBundleFactory() { private boolean once; @Override public RemoteBundle getBundle( OutputReceiverFactory receiverFactory, TimerReceiverFactory timerReceiverFactory, StateRequestHandler stateRequestHandler, BundleProgressHandler progressHandler) { return new RemoteBundle() { @Override public String getId() { return "bundle-id"; } @Override public Map<String, FnDataReceiver> getInputReceivers() { return ImmutableMap.of( "input", input -> { }); } @Override public Map<KV<String, String>, FnDataReceiver<Timer>> getTimerReceivers() { return Collections.emptyMap(); } @Override public void requestProgress() { throw new UnsupportedOperationException(); } @Override public void split(double fractionOfRemainder) { throw new UnsupportedOperationException(); } @Override public void close() throws Exception { if (once) { return; } receiverFactory.create("one").accept(three); receiverFactory.create("two").accept(four); receiverFactory.create("three").accept(five); once = true; } }; } @Override public ProcessBundleDescriptors.ExecutableProcessBundleDescriptor getProcessBundleDescriptor() { return Mockito.mock(ProcessBundleDescriptors.ExecutableProcessBundleDescriptor.class); } @Override public void close() {} }; when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory); SparkExecutableStageFunction<Integer, ?> function = getFunction(outputTagMap); List<WindowedValue<Integer>> inputs = new ArrayList<>(); inputs.add(WindowedValue.valueInGlobalWindow(0)); Iterator<RawUnionValue> iterator = function.call(inputs.iterator()); Iterable<RawUnionValue> iterable = () -> iterator; assertThat( iterable, contains( new RawUnionValue(1, three), new RawUnionValue(2, four), new RawUnionValue(3, five))); }
@Test public void testStageBundleClosed() throws Exception { SparkExecutableStageFunction<Integer, ?> function = getFunction(Collections.emptyMap()); List<WindowedValue<Integer>> inputs = new ArrayList<>(); inputs.add(WindowedValue.valueInGlobalWindow(0)); function.call(inputs.iterator()); verify(stageBundleFactory).getBundle(any(), any(), any(), any()); verify(stageBundleFactory).getProcessBundleDescriptor(); verify(stageBundleFactory).close(); verifyNoMoreInteractions(stageBundleFactory); }
@Test public void testNoCallOnEmptyInputIterator() throws Exception { SparkExecutableStageFunction<Integer, ?> function = getFunction(Collections.emptyMap()); function.call(Collections.emptyIterator()); verifyZeroInteractions(stageBundleFactory); }