use of org.apache.beam.sdk.values.PCollectionView in project beam by apache.
the class FetchAndFilterStreamingSideInputsOperation method buildPCollectionViewsWithSdkSupportedWindowMappingFn.
private Iterable<PCollectionView<?>> buildPCollectionViewsWithSdkSupportedWindowMappingFn(IdGenerator idGenerator, InstructionRequestHandler instructionRequestHandler, FnDataService beamFnDataService, ApiServiceDescriptor dataServiceApiServiceDescriptor, Coder<BoundedWindow> mainInputWindowCoder, Map<PCollectionView<?>, RunnerApi.FunctionSpec> pCollectionViewsToWindowMappingFns) {
ImmutableList.Builder<PCollectionView<?>> wrappedViews = ImmutableList.builder();
for (Map.Entry<PCollectionView<?>, RunnerApi.FunctionSpec> entry : pCollectionViewsToWindowMappingFns.entrySet()) {
WindowMappingFn windowMappingFn = new FnApiWindowMappingFn(idGenerator, instructionRequestHandler, dataServiceApiServiceDescriptor, beamFnDataService, entry.getValue(), mainInputWindowCoder, entry.getKey().getWindowingStrategyInternal().getWindowFn().windowCoder());
wrappedViews.add(new ForwardingPCollectionView<Materializations.MultimapView>((PCollectionView) entry.getKey()) {
@Override
public WindowMappingFn<?> getWindowMappingFn() {
return windowMappingFn;
}
});
}
return wrappedViews.build();
}
use of org.apache.beam.sdk.values.PCollectionView in project beam by apache.
the class PartialGroupByKeyParDoFnsTest method testCreateWithCombinerAndBatchSideInputs.
@Test
public void testCreateWithCombinerAndBatchSideInputs() throws Exception {
PipelineOptions options = PipelineOptionsFactory.create();
Coder keyCoder = StringUtf8Coder.of();
Coder valueCoder = BigEndianIntegerCoder.of();
KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder);
TestOutputReceiver receiver = new TestOutputReceiver(new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)), counterSet, NameContextsForTests.nameContextForTest());
StepContext stepContext = BatchModeExecutionContext.forTesting(options, "testStage").getStepContext(TestOperationContext.create(counterSet));
when(mockSideInputReader.isEmpty()).thenReturn(false);
ParDoFn pgbk = PartialGroupByKeyParDoFns.create(options, kvCoder, AppliedCombineFn.withInputCoder(Sum.ofIntegers(), CoderRegistry.createDefault(), kvCoder, ImmutableList.<PCollectionView<?>>of(), WindowingStrategy.globalDefault()), mockSideInputReader, receiver, stepContext);
assertTrue(pgbk instanceof BatchSideInputPGBKParDoFn);
}
use of org.apache.beam.sdk.values.PCollectionView in project beam by apache.
the class StreamingSideInputDoFnRunnerTest method testMultipleSideInputs.
@Test
public void testMultipleSideInputs() throws Exception {
PCollectionView<String> view1 = createView();
PCollectionView<String> view2 = createView();
IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10));
Windmill.GlobalDataId id = Windmill.GlobalDataId.newBuilder().setTag(view1.getTagInternal().getId()).setVersion(ByteString.copyFrom(CoderUtils.encodeToByteArray(IntervalWindow.getCoder(), window))).build();
Set<Windmill.GlobalDataRequest> requestSet = new HashSet<>();
requestSet.add(Windmill.GlobalDataRequest.newBuilder().setDataId(id).build());
Map<IntervalWindow, Set<Windmill.GlobalDataRequest>> blockedMap = new HashMap<>();
blockedMap.put(window, requestSet);
ValueState<Map<IntervalWindow, Set<GlobalDataRequest>>> blockedMapState = state.state(StateNamespaces.global(), StreamingSideInputFetcher.blockedMapAddr(WINDOW_FN.windowCoder()));
blockedMapState.write(blockedMap);
when(stepContext.getSideInputNotifications()).thenReturn(Arrays.asList(id));
when(stepContext.issueSideInputFetch(any(PCollectionView.class), any(BoundedWindow.class), any(SideInputState.class))).thenReturn(true);
when(execContext.getSideInputReaderForViews(Mockito.<Iterable<? extends PCollectionView<?>>>any())).thenReturn(mockSideInputReader);
when(mockSideInputReader.contains(eq(view1))).thenReturn(true);
when(mockSideInputReader.contains(eq(view2))).thenReturn(true);
when(mockSideInputReader.get(eq(view1), any(BoundedWindow.class))).thenReturn("data1");
when(mockSideInputReader.get(eq(view2), any(BoundedWindow.class))).thenReturn("data2");
ListOutputManager outputManager = new ListOutputManager();
List<PCollectionView<String>> views = Arrays.asList(view1, view2);
StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher = createFetcher(views);
StreamingSideInputDoFnRunner<String, String, IntervalWindow> runner = createRunner(outputManager, views, sideInputFetcher);
sideInputFetcher.watermarkHold(createWindow(0)).add(new Instant(0));
sideInputFetcher.elementBag(createWindow(0)).add(createDatum("e1", 0));
runner.startBundle();
runner.processElement(createDatum("e2", 2));
runner.finishBundle();
assertThat(outputManager.getOutput(mainOutputTag), contains(createDatum("e1:data1:data2", 0), createDatum("e2:data1:data2", 2)));
assertThat(blockedMapState.read(), Matchers.nullValue());
assertThat(sideInputFetcher.watermarkHold(createWindow(0)).read(), Matchers.nullValue());
assertThat(sideInputFetcher.elementBag(createWindow(0)).read(), Matchers.emptyIterable());
}
use of org.apache.beam.sdk.values.PCollectionView in project beam by apache.
the class StreamingSideInputDoFnRunnerTest method testSideInputReady.
@Test
public void testSideInputReady() throws Exception {
PCollectionView<String> view = createView();
when(stepContext.getSideInputNotifications()).thenReturn(Arrays.<Windmill.GlobalDataId>asList());
when(stepContext.issueSideInputFetch(eq(view), any(BoundedWindow.class), eq(SideInputState.UNKNOWN))).thenReturn(true);
when(execContext.getSideInputReaderForViews(Mockito.<Iterable<? extends PCollectionView<?>>>any())).thenReturn(mockSideInputReader);
when(mockSideInputReader.contains(eq(view))).thenReturn(true);
when(mockSideInputReader.get(eq(view), any(BoundedWindow.class))).thenReturn("data");
ListOutputManager outputManager = new ListOutputManager();
List<PCollectionView<String>> views = Arrays.asList(view);
StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher = createFetcher(views);
StreamingSideInputDoFnRunner<String, String, IntervalWindow> runner = createRunner(outputManager, views, sideInputFetcher);
runner.startBundle();
runner.processElement(createDatum("e", 0));
runner.finishBundle();
assertThat(outputManager.getOutput(mainOutputTag), contains(createDatum("e:data", 0)));
}
use of org.apache.beam.sdk.values.PCollectionView in project beam by apache.
the class StreamingSideInputDoFnRunnerTest method testSideInputNotification.
@Test
public void testSideInputNotification() throws Exception {
PCollectionView<String> view = createView();
IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10));
Windmill.GlobalDataId id = Windmill.GlobalDataId.newBuilder().setTag(view.getTagInternal().getId()).setVersion(ByteString.copyFrom(CoderUtils.encodeToByteArray(IntervalWindow.getCoder(), window))).build();
Set<Windmill.GlobalDataRequest> requestSet = new HashSet<>();
requestSet.add(Windmill.GlobalDataRequest.newBuilder().setDataId(id).build());
Map<IntervalWindow, Set<Windmill.GlobalDataRequest>> blockedMap = new HashMap<>();
blockedMap.put(window, requestSet);
ValueState<Map<IntervalWindow, Set<GlobalDataRequest>>> blockedMapState = state.state(StateNamespaces.global(), StreamingSideInputFetcher.blockedMapAddr(WINDOW_FN.windowCoder()));
blockedMapState.write(blockedMap);
ListOutputManager outputManager = new ListOutputManager();
List<PCollectionView<String>> views = Arrays.asList(view);
StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher = createFetcher(views);
StreamingSideInputDoFnRunner<String, String, IntervalWindow> runner = createRunner(outputManager, views, sideInputFetcher);
sideInputFetcher.watermarkHold(createWindow(0)).add(new Instant(0));
sideInputFetcher.elementBag(createWindow(0)).add(createDatum("e", 0));
when(stepContext.getSideInputNotifications()).thenReturn(Arrays.asList(id));
when(stepContext.issueSideInputFetch(eq(view), any(BoundedWindow.class), eq(SideInputState.UNKNOWN))).thenReturn(false);
when(stepContext.issueSideInputFetch(eq(view), any(BoundedWindow.class), eq(SideInputState.KNOWN_READY))).thenReturn(true);
when(execContext.getSideInputReaderForViews(Mockito.<Iterable<? extends PCollectionView<?>>>any())).thenReturn(mockSideInputReader);
when(mockSideInputReader.contains(eq(view))).thenReturn(true);
when(mockSideInputReader.get(eq(view), any(BoundedWindow.class))).thenReturn("data");
runner.startBundle();
runner.finishBundle();
assertThat(outputManager.getOutput(mainOutputTag), contains(createDatum("e:data", 0)));
assertThat(blockedMapState.read(), Matchers.nullValue());
assertThat(sideInputFetcher.watermarkHold(createWindow(0)).read(), Matchers.nullValue());
assertThat(sideInputFetcher.elementBag(createWindow(0)).read(), Matchers.emptyIterable());
}
Aggregations