Examples with ThrowingConsumer org.apache.beam.fn.harness.fn.ThrowingConsumer used on opensource projects

Example 1 with ThrowingConsumer

use of org.apache.beam.fn.harness.fn.ThrowingConsumer in project beam by apache.

the class ProcessBundleHandlerTest method testCreatingAndProcessingBeamFnDataReadRunner.

@Test
public void testCreatingAndProcessingBeamFnDataReadRunner() throws Exception {
    Map<String, Message> fnApiRegistry = ImmutableMap.of(STRING_CODER_SPEC_ID, STRING_CODER_SPEC);
    String bundleId = "57L";
    String primitiveTransformId = "100L";
    long outputId = 101L;
    List<WindowedValue<String>> outputValues = new ArrayList<>();
    BeamFnApi.Target outputTarget = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference(primitiveTransformId).setName(Long.toString(outputId)).build();
    Multimap<BeamFnApi.Target, ThrowingConsumer<WindowedValue<String>>> existingConsumers = ImmutableMultimap.of(outputTarget, outputValues::add);
    Multimap<BeamFnApi.Target, ThrowingConsumer<WindowedValue<String>>> newConsumers = HashMultimap.create();
    List<ThrowingRunnable> startFunctions = new ArrayList<>();
    List<ThrowingRunnable> finishFunctions = new ArrayList<>();
    BeamFnApi.FunctionSpec functionSpec = BeamFnApi.FunctionSpec.newBuilder().setId("1L").setUrn(DATA_INPUT_URN).setData(Any.pack(REMOTE_PORT)).build();
    BeamFnApi.PrimitiveTransform primitiveTransform = BeamFnApi.PrimitiveTransform.newBuilder().setId(primitiveTransformId).setFunctionSpec(functionSpec).putInputs("input", BeamFnApi.Target.List.getDefaultInstance()).putOutputs(Long.toString(outputId), BeamFnApi.PCollection.newBuilder().setCoderReference(STRING_CODER_SPEC_ID).build()).build();
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient);
    handler.createConsumersForPrimitiveTransform(primitiveTransform, Suppliers.ofInstance(bundleId)::get, existingConsumers::get, newConsumers::put, startFunctions::add, finishFunctions::add);
    verifyZeroInteractions(beamFnDataClient);
    CompletableFuture<Void> completionFuture = new CompletableFuture<>();
    when(beamFnDataClient.forInboundConsumer(any(), any(), any(), any())).thenReturn(completionFuture);
    Iterables.getOnlyElement(startFunctions).run();
    verify(beamFnDataClient).forInboundConsumer(eq(REMOTE_PORT.getApiServiceDescriptor()), eq(KV.of(bundleId, BeamFnApi.Target.newBuilder().setPrimitiveTransformReference(primitiveTransformId).setName("input").build())), eq(STRING_CODER), consumerCaptor.capture());
    consumerCaptor.getValue().accept(valueInGlobalWindow("TestValue"));
    assertThat(outputValues, contains(valueInGlobalWindow("TestValue")));
    outputValues.clear();
    assertThat(newConsumers.keySet(), empty());
    completionFuture.complete(null);
    Iterables.getOnlyElement(finishFunctions).run();
    verifyNoMoreInteractions(beamFnDataClient);
}

Example 2 with ThrowingConsumer

use of org.apache.beam.fn.harness.fn.ThrowingConsumer in project beam by apache.

the class ProcessBundleHandlerTest method testCreatingAndProcessingSource.

@Test
public void testCreatingAndProcessingSource() throws Exception {
    Map<String, Message> fnApiRegistry = ImmutableMap.of(LONG_CODER_SPEC_ID, LONG_CODER_SPEC);
    String primitiveTransformId = "100L";
    long outputId = 101L;
    BeamFnApi.Target inputTarget = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference("1000L").setName("inputTarget").build();
    List<WindowedValue<String>> outputValues = new ArrayList<>();
    BeamFnApi.Target outputTarget = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference(primitiveTransformId).setName(Long.toString(outputId)).build();
    Multimap<BeamFnApi.Target, ThrowingConsumer<WindowedValue<String>>> existingConsumers = ImmutableMultimap.of(outputTarget, outputValues::add);
    Multimap<BeamFnApi.Target, ThrowingConsumer<WindowedValue<BoundedSource<Long>>>> newConsumers = HashMultimap.create();
    List<ThrowingRunnable> startFunctions = new ArrayList<>();
    List<ThrowingRunnable> finishFunctions = new ArrayList<>();
    BeamFnApi.FunctionSpec functionSpec = BeamFnApi.FunctionSpec.newBuilder().setId("1L").setUrn(JAVA_SOURCE_URN).setData(Any.pack(BytesValue.newBuilder().setValue(ByteString.copyFrom(SerializableUtils.serializeToByteArray(CountingSource.upTo(3)))).build())).build();
    BeamFnApi.PrimitiveTransform primitiveTransform = BeamFnApi.PrimitiveTransform.newBuilder().setId(primitiveTransformId).setFunctionSpec(functionSpec).putInputs("input", BeamFnApi.Target.List.newBuilder().addTarget(inputTarget).build()).putOutputs(Long.toString(outputId), BeamFnApi.PCollection.newBuilder().setCoderReference(LONG_CODER_SPEC_ID).build()).build();
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient);
    handler.createConsumersForPrimitiveTransform(primitiveTransform, Suppliers.ofInstance("57L")::get, existingConsumers::get, newConsumers::put, startFunctions::add, finishFunctions::add);
    // This is testing a deprecated way of running sources and should be removed
    // once all source definitions are instead propagated along the input edge.
    Iterables.getOnlyElement(startFunctions).run();
    assertThat(outputValues, contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(2L)));
    outputValues.clear();
    // Check that when passing a source along as an input, the source is processed.
    assertEquals(newConsumers.keySet(), ImmutableSet.of(inputTarget));
    Iterables.getOnlyElement(newConsumers.get(inputTarget)).accept(valueInGlobalWindow(CountingSource.upTo(2)));
    assertThat(outputValues, contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L)));
    assertThat(finishFunctions, empty());
}

Example 3 with ThrowingConsumer

use of org.apache.beam.fn.harness.fn.ThrowingConsumer in project beam by apache.

the class ProcessBundleHandlerTest method testCreatingAndProcessingBeamFnDataWriteRunner.

@Test
public void testCreatingAndProcessingBeamFnDataWriteRunner() throws Exception {
    Map<String, Message> fnApiRegistry = ImmutableMap.of(STRING_CODER_SPEC_ID, STRING_CODER_SPEC);
    String bundleId = "57L";
    String primitiveTransformId = "100L";
    long outputId = 101L;
    BeamFnApi.Target inputTarget = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference("1000L").setName("inputTarget").build();
    Multimap<BeamFnApi.Target, ThrowingConsumer<WindowedValue<String>>> existingConsumers = ImmutableMultimap.of();
    Multimap<BeamFnApi.Target, ThrowingConsumer<WindowedValue<String>>> newConsumers = HashMultimap.create();
    List<ThrowingRunnable> startFunctions = new ArrayList<>();
    List<ThrowingRunnable> finishFunctions = new ArrayList<>();
    BeamFnApi.FunctionSpec functionSpec = BeamFnApi.FunctionSpec.newBuilder().setId("1L").setUrn(DATA_OUTPUT_URN).setData(Any.pack(REMOTE_PORT)).build();
    BeamFnApi.PrimitiveTransform primitiveTransform = BeamFnApi.PrimitiveTransform.newBuilder().setId(primitiveTransformId).setFunctionSpec(functionSpec).putInputs("input", BeamFnApi.Target.List.newBuilder().addTarget(inputTarget).build()).putOutputs(Long.toString(outputId), BeamFnApi.PCollection.newBuilder().setCoderReference(STRING_CODER_SPEC_ID).build()).build();
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient);
    handler.createConsumersForPrimitiveTransform(primitiveTransform, Suppliers.ofInstance(bundleId)::get, existingConsumers::get, newConsumers::put, startFunctions::add, finishFunctions::add);
    verifyZeroInteractions(beamFnDataClient);
    List<WindowedValue<String>> outputValues = new ArrayList<>();
    AtomicBoolean wasCloseCalled = new AtomicBoolean();
    CloseableThrowingConsumer<WindowedValue<String>> outputConsumer = new CloseableThrowingConsumer<WindowedValue<String>>() {

        @Override
        public void close() throws Exception {
            wasCloseCalled.set(true);
        }

        @Override
        public void accept(WindowedValue<String> t) throws Exception {
            outputValues.add(t);
        }
    };
    when(beamFnDataClient.forOutboundConsumer(any(), any(), Matchers.<Coder<WindowedValue<String>>>any())).thenReturn(outputConsumer);
    Iterables.getOnlyElement(startFunctions).run();
    verify(beamFnDataClient).forOutboundConsumer(eq(REMOTE_PORT.getApiServiceDescriptor()), eq(KV.of(bundleId, BeamFnApi.Target.newBuilder().setPrimitiveTransformReference(primitiveTransformId).setName(Long.toString(outputId)).build())), eq(STRING_CODER));
    assertEquals(newConsumers.keySet(), ImmutableSet.of(inputTarget));
    Iterables.getOnlyElement(newConsumers.get(inputTarget)).accept(valueInGlobalWindow("TestValue"));
    assertThat(outputValues, contains(valueInGlobalWindow("TestValue")));
    outputValues.clear();
    assertFalse(wasCloseCalled.get());
    Iterables.getOnlyElement(finishFunctions).run();
    assertTrue(wasCloseCalled.get());
    verifyNoMoreInteractions(beamFnDataClient);
}

Example 4 with ThrowingConsumer

use of org.apache.beam.fn.harness.fn.ThrowingConsumer in project beam by apache.

the class ProcessBundleHandlerTest method testCreatingAndProcessingDoFn.

/**
   * Create a DoFn that has 3 inputs (inputATarget1, inputATarget2, inputBTarget) and 2 outputs
   * (mainOutput, output). Validate that inputs are fed to the {@link DoFn} and that outputs
   * are directed to the correct consumers.
   */
@Test
public void testCreatingAndProcessingDoFn() throws Exception {
    Map<String, Message> fnApiRegistry = ImmutableMap.of(STRING_CODER_SPEC_ID, STRING_CODER_SPEC);
    String primitiveTransformId = "100L";
    long mainOutputId = 101L;
    long additionalOutputId = 102L;
    DoFnInfo<?, ?> doFnInfo = DoFnInfo.forFn(new TestDoFn(), WindowingStrategy.globalDefault(), ImmutableList.of(), StringUtf8Coder.of(), mainOutputId, ImmutableMap.of(mainOutputId, TestDoFn.mainOutput, additionalOutputId, TestDoFn.additionalOutput));
    BeamFnApi.FunctionSpec functionSpec = BeamFnApi.FunctionSpec.newBuilder().setId("1L").setUrn(JAVA_DO_FN_URN).setData(Any.pack(BytesValue.newBuilder().setValue(ByteString.copyFrom(SerializableUtils.serializeToByteArray(doFnInfo))).build())).build();
    BeamFnApi.Target inputATarget1 = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference("1000L").setName("inputATarget1").build();
    BeamFnApi.Target inputATarget2 = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference("1001L").setName("inputATarget1").build();
    BeamFnApi.Target inputBTarget = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference("1002L").setName("inputBTarget").build();
    BeamFnApi.PrimitiveTransform primitiveTransform = BeamFnApi.PrimitiveTransform.newBuilder().setId(primitiveTransformId).setFunctionSpec(functionSpec).putInputs("inputA", BeamFnApi.Target.List.newBuilder().addTarget(inputATarget1).addTarget(inputATarget2).build()).putInputs("inputB", BeamFnApi.Target.List.newBuilder().addTarget(inputBTarget).build()).putOutputs(Long.toString(mainOutputId), BeamFnApi.PCollection.newBuilder().setCoderReference(STRING_CODER_SPEC_ID).build()).putOutputs(Long.toString(additionalOutputId), BeamFnApi.PCollection.newBuilder().setCoderReference(STRING_CODER_SPEC_ID).build()).build();
    List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
    List<WindowedValue<String>> additionalOutputValues = new ArrayList<>();
    BeamFnApi.Target mainOutputTarget = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference(primitiveTransformId).setName(Long.toString(mainOutputId)).build();
    BeamFnApi.Target additionalOutputTarget = BeamFnApi.Target.newBuilder().setPrimitiveTransformReference(primitiveTransformId).setName(Long.toString(additionalOutputId)).build();
    Multimap<BeamFnApi.Target, ThrowingConsumer<WindowedValue<String>>> existingConsumers = ImmutableMultimap.of(mainOutputTarget, mainOutputValues::add, additionalOutputTarget, additionalOutputValues::add);
    Multimap<BeamFnApi.Target, ThrowingConsumer<WindowedValue<String>>> newConsumers = HashMultimap.create();
    List<ThrowingRunnable> startFunctions = new ArrayList<>();
    List<ThrowingRunnable> finishFunctions = new ArrayList<>();
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), fnApiRegistry::get, beamFnDataClient);
    handler.createConsumersForPrimitiveTransform(primitiveTransform, Suppliers.ofInstance("57L")::get, existingConsumers::get, newConsumers::put, startFunctions::add, finishFunctions::add);
    Iterables.getOnlyElement(startFunctions).run();
    mainOutputValues.clear();
    assertEquals(newConsumers.keySet(), ImmutableSet.of(inputATarget1, inputATarget2, inputBTarget));
    Iterables.getOnlyElement(newConsumers.get(inputATarget1)).accept(valueInGlobalWindow("A1"));
    Iterables.getOnlyElement(newConsumers.get(inputATarget1)).accept(valueInGlobalWindow("A2"));
    Iterables.getOnlyElement(newConsumers.get(inputATarget1)).accept(valueInGlobalWindow("B"));
    assertThat(mainOutputValues, contains(valueInGlobalWindow("MainOutputA1"), valueInGlobalWindow("MainOutputA2"), valueInGlobalWindow("MainOutputB")));
    assertThat(additionalOutputValues, contains(valueInGlobalWindow("AdditionalOutputA1"), valueInGlobalWindow("AdditionalOutputA2"), valueInGlobalWindow("AdditionalOutputB")));
    mainOutputValues.clear();
    additionalOutputValues.clear();
    Iterables.getOnlyElement(finishFunctions).run();
    assertThat(mainOutputValues, contains(timestampedValueInGlobalWindow("FinishBundle", GlobalWindow.INSTANCE.maxTimestamp())));
    mainOutputValues.clear();
}

Example 5 with ThrowingConsumer

use of org.apache.beam.fn.harness.fn.ThrowingConsumer in project beam by apache.

the class ProcessBundleHandler method createDoFnRunner.

/**
   * Converts a {@link org.apache.beam.fn.v1.BeamFnApi.FunctionSpec} into a {@link DoFnRunner}.
   */
private <InputT, OutputT> DoFnRunner<InputT, OutputT> createDoFnRunner(BeamFnApi.FunctionSpec functionSpec, Map<String, Collection<ThrowingConsumer<WindowedValue<OutputT>>>> outputMap) {
    ByteString serializedFn;
    try {
        serializedFn = functionSpec.getData().unpack(BytesValue.class).getValue();
    } catch (InvalidProtocolBufferException e) {
        throw new IllegalArgumentException(String.format("Unable to unwrap DoFn %s", functionSpec), e);
    }
    DoFnInfo<?, ?> doFnInfo = (DoFnInfo<?, ?>) SerializableUtils.deserializeFromByteArray(serializedFn.toByteArray(), "DoFnInfo");
    checkArgument(Objects.equals(new HashSet<>(Collections2.transform(outputMap.keySet(), Long::parseLong)), doFnInfo.getOutputMap().keySet()), "Unexpected mismatch between transform output map %s and DoFnInfo output map %s.", outputMap.keySet(), doFnInfo.getOutputMap());
    ImmutableMultimap.Builder<TupleTag<?>, ThrowingConsumer<WindowedValue<OutputT>>> tagToOutput = ImmutableMultimap.builder();
    for (Map.Entry<Long, TupleTag<?>> entry : doFnInfo.getOutputMap().entrySet()) {
        tagToOutput.putAll(entry.getValue(), outputMap.get(Long.toString(entry.getKey())));
    }
    @SuppressWarnings({ "unchecked", "rawtypes" }) final Map<TupleTag<?>, Collection<ThrowingConsumer<WindowedValue<?>>>> tagBasedOutputMap = (Map) tagToOutput.build().asMap();
    OutputManager outputManager = new OutputManager() {

        Map<TupleTag<?>, Collection<ThrowingConsumer<WindowedValue<?>>>> tupleTagToOutput = tagBasedOutputMap;

        @Override
        public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
            try {
                Collection<ThrowingConsumer<WindowedValue<?>>> consumers = tupleTagToOutput.get(tag);
                if (consumers == null) {
                    /* This is a normal case, e.g., if a DoFn has output but that output is not
                 * consumed. Drop the output. */
                    return;
                }
                for (ThrowingConsumer<WindowedValue<?>> consumer : consumers) {
                    consumer.accept(output);
                }
            } catch (Throwable t) {
                throw new RuntimeException(t);
            }
        }
    };
    @SuppressWarnings({ "unchecked", "rawtypes", "deprecation" }) DoFnRunner<InputT, OutputT> runner = DoFnRunners.simpleRunner(PipelineOptionsFactory.create(), /* TODO */
    (DoFn) doFnInfo.getDoFn(), NullSideInputReader.empty(), /* TODO */
    outputManager, (TupleTag) doFnInfo.getOutputMap().get(doFnInfo.getMainOutput()), new ArrayList<>(doFnInfo.getOutputMap().values()), new FakeStepContext(), (WindowingStrategy) doFnInfo.getWindowingStrategy());
    return runner;
}