You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ke...@apache.org on 2025/01/14 17:24:10 UTC
(beam) branch master updated: Improve memory efficiency in FnApiDoFnRunner (#33522)
This is an automated email from the ASF dual-hosted git repository.
kenn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 93476eb85bf Improve memory efficiency in FnApiDoFnRunner (#33522)
93476eb85bf is described below
commit 93476eb85bff504b4b2253048c39cfb78ef3d8d1
Author: Radosław Stankiewicz <ra...@google.com>
AuthorDate: Tue Jan 14 18:24:01 2025 +0100
Improve memory efficiency in FnApiDoFnRunner (#33522)
---
.../beam/fn/harness/BeamFnDataReadRunner.java | 12 +-
.../beam/fn/harness/BeamFnDataWriteRunner.java | 7 +-
.../org/apache/beam/fn/harness/CombineRunners.java | 11 +-
.../apache/beam/fn/harness/FnApiDoFnRunner.java | 24 +--
.../beam/fn/harness/PTransformRunnerFactory.java | 16 +-
.../fn/harness/control/ProcessBundleHandler.java | 27 ++-
.../beam/fn/harness/AssignWindowsRunnerTest.java | 7 +-
.../beam/fn/harness/BeamFnDataReadRunnerTest.java | 51 +++--
.../beam/fn/harness/BeamFnDataWriteRunnerTest.java | 15 +-
.../apache/beam/fn/harness/CombineRunnersTest.java | 37 +++-
.../apache/beam/fn/harness/FlattenRunnerTest.java | 14 +-
.../beam/fn/harness/FnApiDoFnRunnerTest.java | 225 ++++++++++++++-------
.../apache/beam/fn/harness/MapFnRunnersTest.java | 21 +-
.../PTransformRunnerFactoryTestContext.java | 18 +-
14 files changed, 305 insertions(+), 180 deletions(-)
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
index 26d07b3d8b2..a34eb2c4ab0 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
@@ -38,7 +38,6 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitResponse
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
import org.apache.beam.runners.core.metrics.MonitoringInfoConstants.Urns;
import org.apache.beam.runners.core.metrics.MonitoringInfoEncodings;
@@ -95,7 +94,7 @@ public class BeamFnDataReadRunner<OutputT> {
context.getPTransformId(),
context.getPTransform(),
context.getProcessBundleInstructionIdSupplier(),
- context.getCoders(),
+ context.getComponents(),
context.getBeamFnStateClient(),
context::addBundleProgressReporter,
consumer);
@@ -127,7 +126,7 @@ public class BeamFnDataReadRunner<OutputT> {
String pTransformId,
RunnerApi.PTransform grpcReadNode,
Supplier<String> processBundleInstructionIdSupplier,
- Map<String, RunnerApi.Coder> coders,
+ RunnerApi.Components components,
BeamFnStateClient beamFnStateClient,
Consumer<BundleProgressReporter> addBundleProgressReporter,
FnDataReceiver<WindowedValue<OutputT>> consumer)
@@ -138,13 +137,12 @@ public class BeamFnDataReadRunner<OutputT> {
this.processBundleInstructionIdSupplier = processBundleInstructionIdSupplier;
this.consumer = consumer;
- RehydratedComponents components =
- RehydratedComponents.forComponents(Components.newBuilder().putAllCoders(coders).build());
+ RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components);
this.coder =
(Coder<WindowedValue<OutputT>>)
CoderTranslation.fromProto(
- coders.get(port.getCoderId()),
- components,
+ components.getCodersMap().get(port.getCoderId()),
+ rehydratedComponents,
new StateBackedIterableTranslationContext() {
@Override
public Supplier<Cache<?, ?>> getCache() {
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataWriteRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataWriteRunner.java
index ade5c32a8c1..e3474fb1e11 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataWriteRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataWriteRunner.java
@@ -27,7 +27,6 @@ import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.fn.harness.state.StateBackedIterable.StateBackedIterableTranslationContext;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort;
-import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.data.RemoteGrpcPortWrite;
import org.apache.beam.sdk.util.WindowedValue;
@@ -64,13 +63,11 @@ public class BeamFnDataWriteRunner<InputT> {
public BeamFnDataWriteRunner createRunnerForPTransform(Context context) throws IOException {
RemoteGrpcPort port = RemoteGrpcPortWrite.fromPTransform(context.getPTransform()).getPort();
- RehydratedComponents components =
- RehydratedComponents.forComponents(
- Components.newBuilder().putAllCoders(context.getCoders()).build());
+ RehydratedComponents components = RehydratedComponents.forComponents(context.getComponents());
Coder<WindowedValue<InputT>> coder =
(Coder<WindowedValue<InputT>>)
CoderTranslation.fromProto(
- context.getCoders().get(port.getCoderId()),
+ context.getComponents().getCodersMap().get(port.getCoderId()),
components,
new StateBackedIterableTranslationContext() {
@Override
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
index 984aa30e7fb..e0fefafbd1b 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
@@ -137,15 +137,14 @@ public class CombineRunners {
throws IOException {
// Get objects needed to create the runner.
RehydratedComponents rehydratedComponents =
- RehydratedComponents.forComponents(
- RunnerApi.Components.newBuilder()
- .putAllCoders(context.getCoders())
- .putAllWindowingStrategies(context.getWindowingStrategies())
- .build());
+ RehydratedComponents.forComponents(context.getComponents());
String mainInputTag =
Iterables.getOnlyElement(context.getPTransform().getInputsMap().keySet());
RunnerApi.PCollection mainInput =
- context.getPCollections().get(context.getPTransform().getInputsOrThrow(mainInputTag));
+ context
+ .getComponents()
+ .getPcollectionsMap()
+ .get(context.getPTransform().getInputsOrThrow(mainInputTag));
// Input coder may sometimes be WindowedValueCoder depending on runner, instead of the
// expected KvCoder.
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
index c39722c90d8..b11dbd84d90 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
@@ -180,9 +180,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
context.getCacheTokensSupplier(),
context.getBundleCacheSupplier(),
context.getProcessWideCache(),
- context.getPCollections(),
- context.getCoders(),
- context.getWindowingStrategies(),
+ context.getComponents(),
context::addStartBundleFunction,
context::addFinishBundleFunction,
context::addResetFunction,
@@ -354,9 +352,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
Supplier<List<BeamFnApi.ProcessBundleRequest.CacheToken>> cacheTokens,
Supplier<Cache<?, ?>> bundleCache,
Cache<?, ?> processWideCache,
- Map<String, PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
+ RunnerApi.Components components,
Consumer<ThrowingRunnable> addStartFunction,
Consumer<ThrowingRunnable> addFinishFunction,
Consumer<ThrowingRunnable> addResetFunction,
@@ -375,13 +371,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
ImmutableMap.builder();
try {
rehydratedComponents =
- RehydratedComponents.forComponents(
- RunnerApi.Components.newBuilder()
- .putAllCoders(coders)
- .putAllPcollections(pCollections)
- .putAllWindowingStrategies(windowingStrategies)
- .build())
- .withPipeline(Pipeline.create());
+ RehydratedComponents.forComponents(components).withPipeline(Pipeline.create());
parDoPayload = ParDoPayload.parseFrom(pTransform.getSpec().getPayload());
doFn = (DoFn) ParDoTranslation.getDoFn(parDoPayload);
doFnSignature = DoFnSignatures.signatureForDoFn(doFn);
@@ -404,7 +394,8 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
Iterables.getOnlyElement(
Sets.difference(
pTransform.getInputsMap().keySet(), parDoPayload.getSideInputsMap().keySet()));
- PCollection mainInput = pCollections.get(pTransform.getInputsOrThrow(mainInputTag));
+ PCollection mainInput =
+ components.getPcollectionsMap().get(pTransform.getInputsOrThrow(mainInputTag));
Coder<?> maybeWindowedValueInputCoder = rehydratedComponents.getCoder(mainInput.getCoderId());
// TODO: Stop passing windowed value coders within PCollections.
if (maybeWindowedValueInputCoder instanceof WindowedValue.WindowedValueCoder) {
@@ -426,7 +417,8 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
outputCoders = Maps.newHashMap();
for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) {
TupleTag<?> outputTag = new TupleTag<>(entry.getKey());
- RunnerApi.PCollection outputPCollection = pCollections.get(entry.getValue());
+ RunnerApi.PCollection outputPCollection =
+ components.getPcollectionsMap().get(entry.getValue());
Coder<?> outputCoder = rehydratedComponents.getCoder(outputPCollection.getCoderId());
if (outputCoder instanceof WindowedValueCoder) {
outputCoder = ((WindowedValueCoder) outputCoder).getValueCoder();
@@ -443,7 +435,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
String sideInputTag = entry.getKey();
RunnerApi.SideInput sideInput = entry.getValue();
PCollection sideInputPCollection =
- pCollections.get(pTransform.getInputsOrThrow(sideInputTag));
+ components.getPcollectionsMap().get(pTransform.getInputsOrThrow(sideInputTag));
WindowingStrategy sideInputWindowingStrategy =
rehydratedComponents.getWindowingStrategy(
sideInputPCollection.getWindowingStrategyId());
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java
index d75d7eea39c..cc0eadebf41 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java
@@ -62,6 +62,13 @@ public interface PTransformRunnerFactory<T> {
/** The id of the PTransform. */
String getPTransformId();
+ /**
+ * An immutable component with mapping from coder id to coder definition, mapping from windowing
+ * strategy id to windowing strategy definition and mapping from PCollection id to PCollection
+ * definition.
+ */
+ RunnerApi.Components getComponents();
+
/** The PTransform definition. */
RunnerApi.PTransform getPTransform();
@@ -77,15 +84,6 @@ public interface PTransformRunnerFactory<T> {
/** A cache that is process wide and persists across bundle boundaries. */
Cache<?, ?> getProcessWideCache();
- /** An immutable mapping from PCollection id to PCollection definition. */
- Map<String, RunnerApi.PCollection> getPCollections();
-
- /** An immutable mapping from coder id to coder definition. */
- Map<String, RunnerApi.Coder> getCoders();
-
- /** An immutable mapping from windowing strategy id to windowing strategy definition. */
- Map<String, RunnerApi.WindowingStrategy> getWindowingStrategies();
-
/** An immutable set of runner capability urns. */
Set<String> getRunnerCapabilities();
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
index 300796ac6f1..e6f87d58bae 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
@@ -66,11 +66,8 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.model.pipeline.v1.RunnerApi.Coder;
-import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.model.pipeline.v1.RunnerApi.StandardRunnerProtocols;
-import org.apache.beam.model.pipeline.v1.RunnerApi.WindowingStrategy;
import org.apache.beam.runners.core.metrics.MonitoringInfoConstants.Urns;
import org.apache.beam.runners.core.metrics.ShortIdMap;
import org.apache.beam.sdk.fn.data.BeamFnDataInboundObserver;
@@ -236,6 +233,7 @@ public class ProcessBundleHandler {
Supplier<List<CacheToken>> cacheTokens,
Supplier<Cache<?, ?>> bundleCache,
ProcessBundleDescriptor processBundleDescriptor,
+ RunnerApi.Components components,
SetMultimap<String, String> pCollectionIdsToConsumingPTransforms,
PCollectionConsumerRegistry pCollectionConsumerRegistry,
Set<String> processedPTransformIds,
@@ -268,6 +266,7 @@ public class ProcessBundleHandler {
cacheTokens,
bundleCache,
processBundleDescriptor,
+ components,
pCollectionIdsToConsumingPTransforms,
pCollectionConsumerRegistry,
processedPTransformIds,
@@ -358,18 +357,8 @@ public class ProcessBundleHandler {
}
@Override
- public Map<String, PCollection> getPCollections() {
- return processBundleDescriptor.getPcollectionsMap();
- }
-
- @Override
- public Map<String, Coder> getCoders() {
- return processBundleDescriptor.getCodersMap();
- }
-
- @Override
- public Map<String, WindowingStrategy> getWindowingStrategies() {
- return processBundleDescriptor.getWindowingStrategiesMap();
+ public RunnerApi.Components getComponents() {
+ return components;
}
@Override
@@ -867,6 +856,13 @@ public class ProcessBundleHandler {
continue;
}
+ RunnerApi.Components components =
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(bundleDescriptor.getCodersMap())
+ .putAllPcollections(bundleDescriptor.getPcollectionsMap())
+ .putAllWindowingStrategies(bundleDescriptor.getWindowingStrategiesMap())
+ .build();
+
createRunnerAndConsumersForPTransformRecursively(
beamFnStateClient,
beamFnDataClient,
@@ -876,6 +872,7 @@ public class ProcessBundleHandler {
bundleProcessor::getCacheTokens,
bundleProcessor::getBundleCache,
bundleDescriptor,
+ components,
pCollectionIdsToConsumingPTransforms,
pCollectionConsumerRegistry,
processedPTransformIds,
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
index 3d81564fee3..6c0e42de39d 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
@@ -192,8 +192,11 @@ public class AssignWindowsRunnerTest implements Serializable {
.build()
.toByteString()))
.build())
- .pCollections(Collections.singletonMap("input", pCollection))
- .coders(Collections.singletonMap("coder-id", coder))
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(Collections.singletonMap("input", pCollection))
+ .putAllCoders(Collections.singletonMap("coder-id", coder))
+ .build())
.build();
Collection<WindowedValue<?>> outputs = new ArrayList<>();
context.addPCollectionConsumer("output", outputs::add);
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
index 55342f9362b..adc1bd0eae1 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
@@ -149,12 +149,17 @@ public class BeamFnDataReadRunnerTest {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(INPUT_TRANSFORM_ID, pTransform)
.processBundleInstructionId(DEFAULT_BUNDLE_ID)
- .pCollections(
- ImmutableMap.of(
- localOutputId,
- RunnerApi.PCollection.newBuilder().setCoderId(ELEMENT_CODER_SPEC_ID).build()))
- .coders(COMPONENTS.getCodersMap())
- .windowingStrategies(COMPONENTS.getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(
+ ImmutableMap.of(
+ localOutputId,
+ RunnerApi.PCollection.newBuilder()
+ .setCoderId(ELEMENT_CODER_SPEC_ID)
+ .build()))
+ .putAllCoders(COMPONENTS.getCodersMap())
+ .putAllWindowingStrategies(COMPONENTS.getWindowingStrategiesMap())
+ .build())
.build();
context.<String>addPCollectionConsumer(localOutputId, outputValues::add);
@@ -187,12 +192,17 @@ public class BeamFnDataReadRunnerTest {
INPUT_TRANSFORM_ID,
RemoteGrpcPortRead.readFromPort(PORT_SPEC, localOutputId).toPTransform())
.processBundleInstructionIdSupplier(bundleId::get)
- .pCollections(
- ImmutableMap.of(
- localOutputId,
- RunnerApi.PCollection.newBuilder().setCoderId(ELEMENT_CODER_SPEC_ID).build()))
- .coders(COMPONENTS.getCodersMap())
- .windowingStrategies(COMPONENTS.getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(
+ ImmutableMap.of(
+ localOutputId,
+ RunnerApi.PCollection.newBuilder()
+ .setCoderId(ELEMENT_CODER_SPEC_ID)
+ .build()))
+ .putAllCoders(COMPONENTS.getCodersMap())
+ .putAllWindowingStrategies(COMPONENTS.getWindowingStrategiesMap())
+ .build())
.build();
context.<String>addPCollectionConsumer(localOutputId, outputValues::add);
@@ -659,12 +669,17 @@ public class BeamFnDataReadRunnerTest {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(pTransformId, pTransform)
.processBundleInstructionId(DEFAULT_BUNDLE_ID)
- .pCollections(
- ImmutableMap.of(
- localOutputId,
- RunnerApi.PCollection.newBuilder().setCoderId(ELEMENT_CODER_SPEC_ID).build()))
- .coders(COMPONENTS.getCodersMap())
- .windowingStrategies(COMPONENTS.getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(
+ ImmutableMap.of(
+ localOutputId,
+ RunnerApi.PCollection.newBuilder()
+ .setCoderId(ELEMENT_CODER_SPEC_ID)
+ .build()))
+ .putAllCoders(COMPONENTS.getCodersMap())
+ .putAllWindowingStrategies(COMPONENTS.getWindowingStrategiesMap())
+ .build())
.build();
context.addPCollectionConsumer(localOutputId, consumer);
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java
index 071ac037883..7835516f1b5 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataWriteRunnerTest.java
@@ -153,12 +153,15 @@ public class BeamFnDataWriteRunnerTest {
.beamFnDataClient(mockBeamFnDataClient)
.processBundleInstructionIdSupplier(bundleId::get)
.outboundAggregators(aggregators)
- .pCollections(
- ImmutableMap.of(
- localInputId,
- RunnerApi.PCollection.newBuilder().setCoderId(ELEM_CODER_ID).build()))
- .coders(COMPONENTS.getCodersMap())
- .windowingStrategies(COMPONENTS.getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(
+ ImmutableMap.of(
+ localInputId,
+ RunnerApi.PCollection.newBuilder().setCoderId(ELEM_CODER_ID).build()))
+ .putAllCoders(COMPONENTS.getCodersMap())
+ .putAllWindowingStrategies(COMPONENTS.getWindowingStrategiesMap())
+ .build())
.build();
new BeamFnDataWriteRunner.Factory<String>().createRunnerForPTransform(context);
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/CombineRunnersTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/CombineRunnersTest.java
index 1bdc3d89971..d441045261b 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/CombineRunnersTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/CombineRunnersTest.java
@@ -120,9 +120,12 @@ public class CombineRunnersTest {
public void testPrecombine() throws Exception {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_COMBINE_ID, pTransform)
- .pCollections(pProto.getComponents().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponents().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build())
.build();
// Add a consumer and output target to check output values.
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new ArrayDeque<>();
@@ -201,8 +204,11 @@ public class CombineRunnersTest {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_COMBINE_ID, pTransform)
- .pCollections(pCollectionMap)
- .coders(coderMap)
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pCollectionMap)
+ .putAllCoders(coderMap)
+ .build())
.build();
// Add a consumer and output target to check output values.
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new ArrayDeque<>();
@@ -262,8 +268,11 @@ public class CombineRunnersTest {
.build());
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_COMBINE_ID, pTransform)
- .pCollections(pCollectionMap)
- .coders(coderMap)
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pCollectionMap)
+ .putAllCoders(coderMap)
+ .build())
.build();
// Add a consumer and output target to check output values.
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new ArrayDeque<>();
@@ -306,8 +315,11 @@ public class CombineRunnersTest {
public void testConvertToAccumulators() throws Exception {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_COMBINE_ID, pTransform)
- .pCollections(pProto.getComponents().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponents().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .build())
.build();
// Add a consumer and output target to check output values.
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new ArrayDeque<>();
@@ -373,8 +385,11 @@ public class CombineRunnersTest {
.build());
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_COMBINE_ID, pTransform)
- .pCollections(pCollectionMap)
- .coders(coderMap)
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pCollectionMap)
+ .putAllCoders(coderMap)
+ .build())
.build();
// Add a consumer and output target to check output values.
Deque<WindowedValue<KV<String, Integer>>> mainOutputValues = new ArrayDeque<>();
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FlattenRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FlattenRunnerTest.java
index 828d146f0b9..bbb86cab959 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FlattenRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FlattenRunnerTest.java
@@ -90,8 +90,11 @@ public class FlattenRunnerTest {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(pTransformId, pTransform)
.processBundleInstructionId("57")
- .pCollections(pCollectionMap)
- .coders(Collections.singletonMap("coder-id", coder))
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pCollectionMap)
+ .putAllCoders(Collections.singletonMap("coder-id", coder))
+ .build())
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
context.addPCollectionConsumer(
@@ -150,8 +153,11 @@ public class FlattenRunnerTest {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(pTransformId, pTransform)
.processBundleInstructionId("57")
- .pCollections(Collections.singletonMap("inputATarget", pCollection))
- .coders(Collections.singletonMap("coder-id", coder))
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(Collections.singletonMap("inputATarget", pCollection))
+ .putAllCoders(Collections.singletonMap("coder-id", coder))
+ .build())
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
context.addPCollectionConsumer(
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
index f4d555dabcc..8f4cfa1acf6 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
@@ -268,9 +268,12 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponents().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
context.addPCollectionConsumer(
@@ -428,9 +431,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponents().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
List<WindowedValue<String>> additionalOutputValues = new ArrayList<>();
@@ -531,9 +538,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponents().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
List<WindowedValue<String>> additionalOutputValues = new ArrayList<>();
@@ -670,9 +681,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponents().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<Iterable<String>>> mainOutputValues = new ArrayList<>();
context.addPCollectionConsumer(
@@ -772,9 +787,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponents().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<Iterable<String>>> mainOutputValues = new ArrayList<>();
context.addPCollectionConsumer(
@@ -942,9 +961,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeStateClient)
.processBundleInstructionId("57L")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.outboundAggregators(
ImmutableMap.of(ApiServiceDescriptor.getDefaultInstance(), aggregator))
.timerApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance())
@@ -1678,9 +1701,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.splitListener(splitListener)
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
@@ -1968,9 +1995,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.splitListener(splitListener)
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
@@ -2178,9 +2209,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.splitListener(splitListener)
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
@@ -2372,9 +2407,12 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build())
.splitListener(splitListener)
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
@@ -2790,9 +2828,12 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build())
.build();
List<WindowedValue<KV<String, OffsetRange>>> mainOutputValues = new ArrayList<>();
context.addPCollectionConsumer(outputPCollectionId, ((List) mainOutputValues)::add);
@@ -2875,9 +2916,12 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build())
.build();
List<WindowedValue<KV<String, OffsetRange>>> mainOutputValues = new ArrayList<>();
context.addPCollectionConsumer(outputPCollectionId, ((List) mainOutputValues)::add);
@@ -2986,9 +3030,12 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build())
.build();
List<WindowedValue<KV<String, OffsetRange>>> mainOutputValues = new ArrayList<>();
context.addPCollectionConsumer(outputPCollectionId, ((List) mainOutputValues)::add);
@@ -3081,9 +3128,12 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build())
.build();
List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
Coder coder =
@@ -3180,9 +3230,12 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build())
.build();
List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
Coder coder =
@@ -3321,9 +3374,12 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build())
.build();
List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
Coder coder =
@@ -3516,9 +3572,16 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ // .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ // .coders(pProto.getComponents().getCodersMap())
+ // .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
.build();
List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
Coder coder =
@@ -3666,9 +3729,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
Coder coder =
@@ -3722,9 +3789,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
Coder coder =
@@ -3820,9 +3891,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
Coder coder =
@@ -3938,9 +4013,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
Coder coder =
@@ -4066,9 +4145,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
Coder coder = StringUtf8Coder.of();
@@ -4125,9 +4208,13 @@ public class FnApiDoFnRunnerTest implements Serializable {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
.processBundleInstructionId("57")
- .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .coders(pProto.getComponents().getCodersMap())
- .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
.build();
Coder coder = StringUtf8Coder.of();
context.addPCollectionConsumer(
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/MapFnRunnersTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/MapFnRunnersTest.java
index cf750bfd7d6..18b15e7785b 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/MapFnRunnersTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/MapFnRunnersTest.java
@@ -71,8 +71,11 @@ public class MapFnRunnersTest {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(EXPECTED_ID, EXPECTED_PTRANSFORM)
.processBundleInstructionId("57")
- .pCollections(Collections.singletonMap("inputPC", INPUT_PCOLLECTION))
- .coders(Collections.singletonMap("coder-id", valueCoder))
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(Collections.singletonMap("inputPC", INPUT_PCOLLECTION))
+ .putAllCoders(Collections.singletonMap("coder-id", valueCoder))
+ .build())
.build();
List<WindowedValue<?>> outputConsumer = new ArrayList<>();
context.addPCollectionConsumer("outputPC", outputConsumer::add);
@@ -97,8 +100,11 @@ public class MapFnRunnersTest {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(EXPECTED_ID, EXPECTED_PTRANSFORM)
.processBundleInstructionId("57")
- .pCollections(Collections.singletonMap("inputPC", INPUT_PCOLLECTION))
- .coders(Collections.singletonMap("coder-id", valueCoder))
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(Collections.singletonMap("inputPC", INPUT_PCOLLECTION))
+ .putAllCoders(Collections.singletonMap("coder-id", valueCoder))
+ .build())
.build();
List<WindowedValue<?>> outputConsumer = new ArrayList<>();
context.addPCollectionConsumer("outputPC", outputConsumer::add);
@@ -122,8 +128,11 @@ public class MapFnRunnersTest {
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(EXPECTED_ID, EXPECTED_PTRANSFORM)
.processBundleInstructionId("57")
- .pCollections(Collections.singletonMap("inputPC", INPUT_PCOLLECTION))
- .coders(Collections.singletonMap("coder-id", valueCoder))
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllPcollections(Collections.singletonMap("inputPC", INPUT_PCOLLECTION))
+ .putAllCoders(Collections.singletonMap("coder-id", valueCoder))
+ .build())
.build();
List<WindowedValue<?>> outputConsumer = new ArrayList<>();
context.addPCollectionConsumer("outputPC", outputConsumer::add);
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java
index acfd3bb7020..bd7b7726e69 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java
@@ -114,9 +114,13 @@ public abstract class PTransformRunnerFactoryTestContext
.cacheTokensSupplier(() -> Collections.emptyList())
.bundleCacheSupplier(() -> Caches.noop())
.processWideCache(Caches.noop())
- .pCollections(Collections.emptyMap()) // expected to be immutable
- .coders(Collections.emptyMap()) // expected to be immutable
- .windowingStrategies(Collections.emptyMap()) // expected to be immutable
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(Collections.emptyMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(Collections.emptyMap())
+ .putAllPcollections(Collections.emptyMap())
+ .build())
.pCollectionConsumers(new HashMap<>())
.startBundleFunctions(new ArrayList<>())
.finishBundleFunctions(new ArrayList<>())
@@ -175,11 +179,13 @@ public abstract class PTransformRunnerFactoryTestContext
return processBundleInstructionIdSupplier(() -> value);
}
- Builder pCollections(Map<String, RunnerApi.PCollection> value);
+ // Builder pCollections(Map<String, RunnerApi.PCollection> value);
- Builder coders(Map<String, RunnerApi.Coder> value);
+ Builder components(RunnerApi.Components value);
- Builder windowingStrategies(Map<String, RunnerApi.WindowingStrategy> value);
+ // Builder coders(Map<String, RunnerApi.Coder> value);
+
+ // Builder windowingStrategies(Map<String, RunnerApi.WindowingStrategy> value);
Builder runnerCapabilities(Set<String> value);