This is an automated email from the ASF dual-hosted git repository.
scwhittle 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 66ecc31a22c Adds a new CoderTranslator for Java SchemaCoders. (#39594)
66ecc31a22c is described below
commit 66ecc31a22c9fa39dbf70965f17d10b114022185
Author: Andrew Crites <[email protected]>
AuthorDate: Tue Aug 25 00:42:07 2026 -0700
Adds a new CoderTranslator for Java SchemaCoders. (#39594)
The updateCompatabilityVersion option can be used to select the previous
encoding.
---
CHANGES.md | 4 +-
.../dataflow/DataflowPipelineTranslator.java | 2 +-
.../beam/runners/dataflow/DataflowRunner.java | 54 ++++++++-------
.../dataflow/DataflowPipelineTranslatorTest.java | 78 +++++++++++++++++----
.../control/ProcessBundleDescriptorsTest.java | 6 +-
.../sdk/util/construction/CoderTranslation.java | 62 +++++++++++++++--
.../construction/CoderTranslatorRegistrar.java | 16 +++++
.../sdk/util/construction/CoderTranslators.java | 79 ++++++++++++++++++++++
.../sdk/util/construction/ModelCoderRegistrar.java | 28 ++++++--
.../beam/sdk/util/construction/ModelCoders.java | 2 +
.../util/construction/RehydratedComponents.java | 3 +-
.../beam/sdk/util/construction/SdkComponents.java | 39 ++++++-----
.../util/construction/CoderTranslationTest.java | 38 ++++++++++-
.../extensions/avro/AvroGenericCoderRegistrar.java | 18 +++++
.../beam/fn/harness/state/StateBackedIterable.java | 22 ++++++
.../runners/portability/fn_api_runner/execution.py | 4 +-
.../portability/fn_api_runner/translations.py | 4 +-
17 files changed, 386 insertions(+), 73 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 3994aea4d93..3f097d96dab 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -74,7 +74,9 @@
## Breaking Changes
-* X behavior was changed ([#X](https://github.com/apache/beam/issues/X)).
+* Portable Java SDK now encodes SchemaCoders in a portable way
([#34672](https://github.com/apache/beam/issues/34672)).
+ - Original custom Java coder encoding can still be obtained using
[StreamingOptions.setUpdateCompatibilityVersion("2.76")](https://github.com/apache/beam/blob/2cf0930e7ae1aa389c26ce6639b584877a3e31d9/sdks/java/core/src/main/java/org/apache/beam/sdk/options/StreamingOptions.java#L47)
([#34672](https://github.com/apache/beam/issues/34672)).
+ - Fixes ([#36496](https://github.com/apache/beam/issues/36496)),
([#30276](https://github.com/apache/beam/issues/30276)),
([#29245](https://github.com/apache/beam/issues/29245)).
## Deprecations
diff --git
a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
index 4016f31a547..1609cf6ea23 100644
---
a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
+++
b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
@@ -221,7 +221,7 @@ public class DataflowPipelineTranslator {
private static byte[] serializeWindowingStrategy(
WindowingStrategy<?, ?> windowingStrategy, PipelineOptions options) {
try {
- SdkComponents sdkComponents = SdkComponents.create();
+ SdkComponents sdkComponents = SdkComponents.create(options);
String workerHarnessContainerImageURL =
DataflowRunner.getContainerImageForJob(options.as(DataflowPipelineOptions.class));
diff --git
a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
index 1b03cc10351..ee749901ec6 100644
---
a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
+++
b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
@@ -1355,19 +1355,20 @@ public class DataflowRunner extends
PipelineRunner<DataflowPipelineJob> {
// with the SDK harness image (which implements Fn API).
//
// The same Environment is used in different and contradictory ways,
depending on whether
- // it is a v1 or v2 job submission.
+ // it is a portable or non-portable job submission.
RunnerApi.Environment defaultEnvironmentForDataflow =
Environments.createDockerEnvironment(workerHarnessContainerImageURL);
- // The SdkComponents for portable an non-portable job submission must be
kept distinct. Both
+ // The SdkComponents for portable and non-portable job submission must be
kept distinct. Both
// need the default environment.
- SdkComponents portableComponents = SdkComponents.create();
- portableComponents.registerEnvironment(
- defaultEnvironmentForDataflow
- .toBuilder()
- .addAllDependencies(getDefaultArtifacts())
- .addAllCapabilities(Environments.getJavaCapabilities())
- .build());
+ SdkComponents portableComponents =
+ SdkComponents.create(
+ options,
+ defaultEnvironmentForDataflow
+ .toBuilder()
+ .addAllDependencies(getDefaultArtifacts())
+ .addAllCapabilities(Environments.getJavaCapabilities())
+ .build());
RunnerApi.Pipeline portablePipelineProto =
PipelineTranslation.toProto(pipeline, portableComponents, false);
@@ -1400,27 +1401,28 @@ public class DataflowRunner extends
PipelineRunner<DataflowPipelineJob> {
"Skipping Dataflow Streaming Java Runner transform replacements
since job will run on Dataflow Portable Runner.");
} else {
// Now rewrite things to be as needed for Dataflow Streaming Java Runner
(mutates the
- // pipeline)
+ // pipeline).
// This way the job submitted is valid for Dataflow Streaming Java
Runner and Dataflow
- // Portable Runner, simultaneously
+ // Portable Runner, simultaneously.
replaceV1Transforms(pipeline);
}
- // Capture the SdkComponents for look up during step translations
- SdkComponents dataflowV1Components = SdkComponents.create();
- dataflowV1Components.registerEnvironment(
- defaultEnvironmentForDataflow
- .toBuilder()
- .addAllDependencies(getDefaultArtifacts())
- .addAllCapabilities(Environments.getJavaCapabilities())
- .build());
+ // Capture the SdkComponents for look up during step translations.
+ SdkComponents dataflowNonPortableComponents =
+ SdkComponents.create(
+ options,
+ defaultEnvironmentForDataflow
+ .toBuilder()
+ .addAllDependencies(getDefaultArtifacts())
+ .addAllCapabilities(Environments.getJavaCapabilities())
+ .build());
// No need to perform transform upgrading for the Dataflow Streaming Java
Runner proto.
- RunnerApi.Pipeline dataflowV1PipelineProto =
- PipelineTranslation.toProto(pipeline, dataflowV1Components, true,
false);
+ RunnerApi.Pipeline dataflowNonPortablePipelineProto =
+ PipelineTranslation.toProto(pipeline, dataflowNonPortableComponents,
true, false);
if (LOG.isDebugEnabled()) {
LOG.debug(
- "Dataflow v1 pipeline proto:\n{}",
- TextFormat.printer().printToString(dataflowV1PipelineProto));
+ "Dataflow non-portable worker pipeline proto:\n{}",
+
TextFormat.printer().printToString(dataflowNonPortablePipelineProto));
}
// Set a unique client_request_id in the CreateJob request.
@@ -1440,7 +1442,11 @@ public class DataflowRunner extends
PipelineRunner<DataflowPipelineJob> {
JobSpecification jobSpecification =
translator.translate(
- pipeline, dataflowV1PipelineProto, dataflowV1Components, this,
packages);
+ pipeline,
+ dataflowNonPortablePipelineProto,
+ dataflowNonPortableComponents,
+ this,
+ packages);
if (!isNullOrEmpty(dataflowOptions.getDataflowWorkerJar()) &&
!useUnifiedWorker(options)) {
List<String> experiments =
diff --git
a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java
b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java
index f8818931a68..127402f6f54 100644
---
a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java
+++
b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java
@@ -47,6 +47,7 @@ import com.google.api.services.dataflow.Dataflow;
import com.google.api.services.dataflow.model.Job;
import com.google.api.services.dataflow.model.Step;
import com.google.api.services.dataflow.model.WorkerPool;
+import com.google.auto.value.AutoValue;
import java.io.File;
import java.io.IOException;
import java.io.Serializable;
@@ -92,6 +93,8 @@ import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.StreamingOptions;
import org.apache.beam.sdk.options.ValueProvider;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.StateSpecs;
import org.apache.beam.sdk.state.ValueState;
@@ -166,15 +169,11 @@ public class DataflowPipelineTranslatorTest implements
Serializable {
@Rule public transient ExpectedException thrown = ExpectedException.none();
private SdkComponents createSdkComponents(PipelineOptions options) {
- SdkComponents sdkComponents = SdkComponents.create();
-
String containerImageURL =
DataflowRunner.getContainerImageForJob(options.as(DataflowPipelineOptions.class));
RunnerApi.Environment defaultEnvironmentForDataflow =
Environments.createDockerEnvironment(containerImageURL);
-
- sdkComponents.registerEnvironment(defaultEnvironmentForDataflow);
- return sdkComponents;
+ return SdkComponents.create(options, defaultEnvironmentForDataflow);
}
// A Custom Mockito matcher for an initial Job that checks that all
@@ -1294,15 +1293,16 @@ public class DataflowPipelineTranslatorTest implements
Serializable {
file1.deleteOnExit();
File file2 = File.createTempFile("file2-", ".txt");
file2.deleteOnExit();
- SdkComponents sdkComponents = SdkComponents.create();
- sdkComponents.registerEnvironment(
-
Environments.createDockerEnvironment(DataflowRunner.getContainerImageForJob(options))
- .toBuilder()
- .addAllDependencies(
- Environments.getArtifacts(
- ImmutableList.of("file1.txt=" + file1, "file2.txt=" +
file2)))
- .addAllCapabilities(Environments.getJavaCapabilities())
- .build());
+ SdkComponents sdkComponents =
+ SdkComponents.create(
+ options,
+
Environments.createDockerEnvironment(DataflowRunner.getContainerImageForJob(options))
+ .toBuilder()
+ .addAllDependencies(
+ Environments.getArtifacts(
+ ImmutableList.of("file1.txt=" + file1, "file2.txt=" +
file2)))
+ .addAllCapabilities(Environments.getJavaCapabilities())
+ .build());
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline,
sdkComponents, true);
@@ -1870,4 +1870,54 @@ public class DataflowPipelineTranslatorTest implements
Serializable {
return null;
}
}
+
+ @AutoValue
+ @DefaultSchema(AutoValueSchema.class)
+ public abstract static class SimpleAutoValue {
+ public abstract String getString();
+
+ public abstract int getInt32();
+
+ public abstract long getInt64();
+
+ public static DataflowPipelineTranslatorTest.SimpleAutoValue of(
+ String string, int int32, long int64) {
+ return new
AutoValue_DataflowPipelineTranslatorTest_SimpleAutoValue(string, int32, int64);
+ }
+ }
+
+ @Test
+ public void testSchemaCoderTranslation() throws Exception {
+ DataflowPipelineOptions options = buildPipelineOptions();
+ Pipeline pipeline = Pipeline.create(options);
+ pipeline
+ .apply(Impulse.create())
+ .apply(
+ MapElements.via(
+ new SimpleFunction<byte[], SimpleAutoValue>() {
+ @Override
+ public SimpleAutoValue apply(byte[] input) {
+ return SimpleAutoValue.of("foo", 5, 10L);
+ }
+ }))
+ .apply(Window.into(FixedWindows.of(Duration.standardMinutes(1))));
+ {
+ // By default, new pipelines get the new schema coder encoding.
+ SdkComponents sdkComponents = createSdkComponents(options);
+ RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline,
sdkComponents, true);
+ Map<String, RunnerApi.Coder> coders =
pipelineProto.getComponents().getCodersMap();
+ assertTrue(coders.containsKey("SchemaCoder"));
+ assertEquals("beam:coder:schema:v1",
coders.get("SchemaCoder").getSpec().getUrn());
+ }
+
+ {
+ // Check that we get the java coder encoding at version 2.76 (and
earlier).
+ options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.76");
+ SdkComponents sdkComponents = createSdkComponents(options);
+ RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline,
sdkComponents, true);
+ Map<String, RunnerApi.Coder> coders =
pipelineProto.getComponents().getCodersMap();
+ assertTrue(coders.containsKey("SchemaCoder"));
+ assertEquals("beam:coders:javasdk:0.1",
coders.get("SchemaCoder").getSpec().getUrn());
+ }
+ }
}
diff --git
a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptorsTest.java
b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptorsTest.java
index 21d7550c38b..9ea7404053d 100644
---
a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptorsTest.java
+++
b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptorsTest.java
@@ -78,7 +78,8 @@ public class ProcessBundleDescriptorsTest implements
Serializable {
// Add another stateful stage with a non-standard key coder
Pipeline p = Pipeline.create();
Coder<Void> keycoder = VoidCoder.of();
- assertThat(ModelCoderRegistrar.isKnownCoder(keycoder), is(false));
+ ModelCoderRegistrar coderRegistrar = new ModelCoderRegistrar();
+ assertThat(coderRegistrar.isKnownCoder(keycoder, p.getOptions()),
is(false));
p.apply("impulse", Impulse.create())
.apply(
"create",
@@ -165,7 +166,8 @@ public class ProcessBundleDescriptorsTest implements
Serializable {
public void testLengthPrefixingOfInputCoderExecutableStage() throws
Exception {
Pipeline p = Pipeline.create();
Coder<Void> voidCoder = VoidCoder.of();
- assertThat(ModelCoderRegistrar.isKnownCoder(voidCoder), is(false));
+ ModelCoderRegistrar coderRegistrar = new ModelCoderRegistrar();
+ assertThat(coderRegistrar.isKnownCoder(voidCoder, p.getOptions()),
is(false));
p.apply("impulse", Impulse.create())
.apply(
ParDo.of(
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslation.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslation.java
index 22859dc68b9..2cc4bf0c6a0 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslation.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslation.java
@@ -25,11 +25,13 @@ import java.util.ServiceLoader;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.BiMap;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableBiMap;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
import org.checkerframework.dataflow.qual.Deterministic;
@@ -62,6 +64,8 @@ public class CoderTranslation {
private static @MonotonicNonNull BiMap<Class<? extends Coder>, String>
knownCoderUrns;
+ private static @MonotonicNonNull List<CoderTranslatorRegistrar>
coderTranslatorRegistrars;
+
private static @MonotonicNonNull Map<Class<? extends Coder>,
CoderTranslator<? extends Coder>>
knownTranslators;
@@ -80,6 +84,53 @@ public class CoderTranslation {
return knownCoderUrns;
}
+ private static void initializeCoderTranslatorRegistrars() {
+ ImmutableList.Builder<CoderTranslatorRegistrar> registrars =
ImmutableList.builder();
+ for (CoderTranslatorRegistrar coderTranslatorRegistrar :
+ ServiceLoader.load(CoderTranslatorRegistrar.class)) {
+ registrars.add(coderTranslatorRegistrar);
+ }
+ coderTranslatorRegistrars = registrars.build();
+ }
+
+ static boolean isKnownCoder(Coder<?> coder, PipelineOptions options) {
+ if (coderTranslatorRegistrars == null) {
+ initializeCoderTranslatorRegistrars();
+ }
+ for (CoderTranslatorRegistrar registrar : coderTranslatorRegistrars) {
+ if (registrar.isKnownCoder(coder, options)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ static CoderTranslator<? extends Coder> getCoderTranslator(Class<? extends
Coder> coderClass) {
+ if (coderTranslatorRegistrars == null) {
+ initializeCoderTranslatorRegistrars();
+ }
+ for (CoderTranslatorRegistrar registrar : coderTranslatorRegistrars) {
+ CoderTranslator translator = registrar.getCoderTranslator(coderClass);
+ if (translator != null) {
+ return translator;
+ }
+ }
+ return null;
+ }
+
+ static Class<? extends Coder> getCoderForUrn(String coderUrn) {
+ if (coderTranslatorRegistrars == null) {
+ initializeCoderTranslatorRegistrars();
+ }
+ for (CoderTranslatorRegistrar registrar : coderTranslatorRegistrars) {
+ Class<? extends Coder> coder = registrar.getCoderForUrn(coderUrn);
+ if (coder != null) {
+ return coder;
+ }
+ }
+ return null;
+ }
+
@VisibleForTesting
@Deterministic
static Map<Class<? extends Coder>, CoderTranslator<? extends Coder>>
getKnownTranslators() {
@@ -107,7 +158,7 @@ public class CoderTranslation {
public static RunnerApi.Coder toProto(Coder<?> coder, SdkComponents
components)
throws IOException {
- if (getKnownCoderUrns().containsKey(coder.getClass())) {
+ if (isKnownCoder(coder, components.getPipelineOptions())) {
return toKnownCoder(coder, components);
}
@@ -129,7 +180,10 @@ public class CoderTranslation {
private static RunnerApi.Coder toKnownCoder(Coder<?> coder, SdkComponents
components)
throws IOException {
- CoderTranslator translator = getKnownTranslators().get(coder.getClass());
+ CoderTranslator translator = getCoderTranslator(coder.getClass());
+ if (translator == null) {
+ throw new IOException("Unable to find CoderTranslator for known Coder");
+ }
List<String> componentIds = registerComponents(coder, translator,
components);
return RunnerApi.Coder.newBuilder()
.addAllComponentCoderIds(componentIds)
@@ -186,8 +240,8 @@ public class CoderTranslation {
components.getComponents().getCodersOrThrow(componentId),
components, context);
coderComponents.add(innerCoder);
}
- Class<? extends Coder> coderType =
getKnownCoderUrns().inverse().get(coderUrn);
- CoderTranslator<?> translator = getKnownTranslators().get(coderType);
+ Class<? extends Coder> coderType = getCoderForUrn(coderUrn);
+ CoderTranslator<?> translator = getCoderTranslator(coderType);
if (translator != null) {
return translator.fromComponents(
coderComponents, coder.getSpec().getPayload().toByteArray(),
context);
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslatorRegistrar.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslatorRegistrar.java
index b69d0290de5..44e8c2956ae 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslatorRegistrar.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslatorRegistrar.java
@@ -19,6 +19,8 @@ package org.apache.beam.sdk.util.construction;
import java.util.Map;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** A registrar of {@link Coder} URNs to the associated {@link
CoderTranslator}. */
@SuppressWarnings({
@@ -34,4 +36,18 @@ public interface CoderTranslatorRegistrar {
/** Returns a mapping of URN to {@link CoderTranslator}. */
Map<Class<? extends Coder>, CoderTranslator<? extends Coder>>
getCoderTranslators();
+
+ /**
+ * Returns whether the given Coder is known to this
CoderTranslatorRegistrar. If the Coder is
+ * known, then getCoderTranslator() will return a non-null CoderTranslator.
+ */
+ boolean isKnownCoder(Coder<?> coder, PipelineOptions options);
+
+ /** Returns the CoderTranslator to use for this Coder, or null if the Coder
is not known. */
+ @Nullable
+ CoderTranslator<? extends Coder> getCoderTranslator(Class<? extends Coder>
coderClass);
+
+ /** Returns the Coder to use for the given Urn, or null if the Urn is for an
unknown Coder. */
+ @Nullable
+ Class<? extends Coder> getCoderForUrn(String coderUrn);
}
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslators.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslators.java
index 84a90721a98..a847bf780df 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslators.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslators.java
@@ -19,6 +19,7 @@ package org.apache.beam.sdk.util.construction;
import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+import java.io.IOException;
import java.util.Collections;
import java.util.List;
import org.apache.beam.model.pipeline.v1.SchemaApi;
@@ -30,12 +31,19 @@ import org.apache.beam.sdk.coders.NullableCoder;
import org.apache.beam.sdk.coders.RowCoder;
import org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
import org.apache.beam.sdk.schemas.SchemaTranslation;
+import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.InstanceBuilder;
+import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.ShardedKey;
+import
org.apache.beam.sdk.util.construction.CoderTranslation.TranslationContext;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import
org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
@@ -177,6 +185,77 @@ class CoderTranslators {
};
}
+ static <T> CoderTranslator<SchemaCoder<T>> schema() {
+ return new CoderTranslator<SchemaCoder<T>>() {
+ private static final String TO_ROW_FUNCTION_URN =
"beam:torowfn:javasdk:v1";
+ private static final String FROM_ROW_FUNCTION_URN =
"beam:fromrowfn:javasdk:v1";
+ private static final String TYPE_DESCRIPTOR_URN =
"beam:typedescriptor:javasdk:v1";
+
+ @Override
+ public ImmutableList<? extends Coder<?>> getComponents(SchemaCoder<T>
from) {
+ return ImmutableList.of();
+ }
+
+ @Override
+ public byte[] getPayload(SchemaCoder<T> from) {
+ SchemaApi.SchemaCoderPayload.Builder payload =
SchemaApi.SchemaCoderPayload.newBuilder();
+ payload.setSchema(SchemaTranslation.schemaToProto(from.getSchema(),
true));
+ payload
+ .getToRowFnBuilder()
+ .setUrn(TO_ROW_FUNCTION_URN)
+ .setPayload(
+ ByteString.copyFrom(
+
SerializableUtils.serializeToByteArray(from.getToRowFunction())));
+ payload
+ .getFromRowFnBuilder()
+ .setUrn(FROM_ROW_FUNCTION_URN)
+ .setPayload(
+ ByteString.copyFrom(
+
SerializableUtils.serializeToByteArray(from.getFromRowFunction())));
+ payload
+ .addAdditionalCoderInfosBuilder()
+ .setUrn(TYPE_DESCRIPTOR_URN)
+ .setPayload(
+ ByteString.copyFrom(
+
SerializableUtils.serializeToByteArray(from.getEncodedTypeDescriptor())));
+ return payload.build().toByteArray();
+ }
+
+ @Override
+ public SchemaCoder<T> fromComponents(
+ List<Coder<?>> components, byte[] payload, TranslationContext
context) {
+ checkArgument(
+ components.isEmpty(), "Expected empty component list, but
received: %s", components);
+ try {
+ SchemaApi.SchemaCoderPayload schemaCoderPayload =
+ SchemaApi.SchemaCoderPayload.parseFrom(payload);
+ if (schemaCoderPayload.getAdditionalCoderInfosCount() == 0) {
+ throw new IllegalArgumentException("Missing serialized
typeDescriptor");
+ }
+ TypeDescriptor<T> typeDescriptor =
+ (TypeDescriptor<T>)
+ SerializableUtils.deserializeFromByteArray(
+
schemaCoderPayload.getAdditionalCoderInfos(0).getPayload().toByteArray(),
+ "typeDescriptor");
+ SerializableFunction<T, Row> toRowFunction =
+ (SerializableFunction<T, Row>)
+ SerializableUtils.deserializeFromByteArray(
+
schemaCoderPayload.getToRowFn().getPayload().toByteArray(), "toRowFunction");
+ SerializableFunction<Row, T> fromRowFunction =
+ (SerializableFunction<Row, T>)
+ SerializableUtils.deserializeFromByteArray(
+
schemaCoderPayload.getFromRowFn().getPayload().toByteArray(),
+ "fromRowFunction");
+
+ Schema schema =
SchemaTranslation.schemaFromProto(schemaCoderPayload.getSchema());
+ return SchemaCoder.of(schema, typeDescriptor, toRowFunction,
fromRowFunction);
+ } catch (IOException | IllegalArgumentException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ };
+ }
+
static CoderTranslator<ShardedKey.Coder<?>> shardedKey() {
return new SimpleStructuredCoderTranslator<ShardedKey.Coder<?>>() {
@Override
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/ModelCoderRegistrar.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/ModelCoderRegistrar.java
index 5b0d5aedd61..5728aacd805 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/ModelCoderRegistrar.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/ModelCoderRegistrar.java
@@ -34,6 +34,9 @@ import org.apache.beam.sdk.coders.RowCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder;
import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.StreamingOptions;
+import org.apache.beam.sdk.schemas.SchemaCoder;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import
org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder;
import org.apache.beam.sdk.util.ShardedKey;
@@ -71,6 +74,7 @@ public class ModelCoderRegistrar implements
CoderTranslatorRegistrar {
ModelCoders.PARAM_WINDOWED_VALUE_CODER_URN)
.put(DoubleCoder.class, ModelCoders.DOUBLE_CODER_URN)
.put(RowCoder.class, ModelCoders.ROW_CODER_URN)
+ .put(SchemaCoder.class, ModelCoders.SCHEMA_CODER_URN)
.put(ShardedKey.Coder.class, ModelCoders.SHARDED_KEY_CODER_URN)
.put(TimestampPrefixingWindowCoder.class,
ModelCoders.CUSTOM_WINDOW_CODER_URN)
.put(NullableCoder.class, ModelCoders.NULLABLE_CODER_URN)
@@ -96,6 +100,7 @@ public class ModelCoderRegistrar implements
CoderTranslatorRegistrar {
CoderTranslators.paramWindowedValue())
.put(DoubleCoder.class,
CoderTranslators.atomic(DoubleCoder.class))
.put(RowCoder.class, CoderTranslators.row())
+ .put(SchemaCoder.class, CoderTranslators.schema())
.put(ShardedKey.Coder.class, CoderTranslators.shardedKey())
.put(TimestampPrefixingWindowCoder.class,
CoderTranslators.timestampPrefixingWindow())
.put(NullableCoder.class, CoderTranslators.nullable())
@@ -123,10 +128,6 @@ public class ModelCoderRegistrar implements
CoderTranslatorRegistrar {
Coder.class.getSimpleName());
}
- public static boolean isKnownCoder(Coder<?> coder) {
- return BEAM_MODEL_CODER_URNS.containsKey(coder.getClass());
- }
-
@Override
public Map<Class<? extends Coder>, String> getCoderURNs() {
return BEAM_MODEL_CODER_URNS;
@@ -136,4 +137,23 @@ public class ModelCoderRegistrar implements
CoderTranslatorRegistrar {
public Map<Class<? extends Coder>, CoderTranslator<? extends Coder>>
getCoderTranslators() {
return BEAM_MODEL_CODERS;
}
+
+ @Override
+ public boolean isKnownCoder(Coder<?> coder, PipelineOptions options) {
+ if (coder.getClass() == SchemaCoder.class
+ && StreamingOptions.updateCompatibilityVersionLessThan(options,
"2.77")) {
+ return false;
+ }
+ return BEAM_MODEL_CODER_URNS.containsKey(coder.getClass());
+ }
+
+ @Override
+ public CoderTranslator<? extends Coder> getCoderTranslator(Class<? extends
Coder> coderClass) {
+ return BEAM_MODEL_CODERS.getOrDefault(coderClass, null);
+ }
+
+ @Override
+ public Class<? extends Coder> getCoderForUrn(String coderUrn) {
+ return BEAM_MODEL_CODER_URNS.inverse().getOrDefault(coderUrn, null);
+ }
}
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/ModelCoders.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/ModelCoders.java
index 7b7546aceb6..5059cc1c6b8 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/ModelCoders.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/ModelCoders.java
@@ -61,6 +61,7 @@ public class ModelCoders {
getUrn(StandardCoders.Enum.PARAM_WINDOWED_VALUE);
public static final String ROW_CODER_URN = getUrn(StandardCoders.Enum.ROW);
+ public static final String SCHEMA_CODER_URN =
getUrn(StandardCoders.Enum.SCHEMA);
public static final String STATE_BACKED_ITERABLE_CODER_URN =
"beam:coder:state_backed_iterable:v1";
@@ -90,6 +91,7 @@ public class ModelCoders {
WINDOWED_VALUE_CODER_URN,
DOUBLE_CODER_URN,
ROW_CODER_URN,
+ SCHEMA_CODER_URN,
PARAM_WINDOWED_VALUE_CODER_URN,
STATE_BACKED_ITERABLE_CODER_URN,
SHARDED_KEY_CODER_URN,
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/RehydratedComponents.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/RehydratedComponents.java
index f7969621436..64c7898a37b 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/RehydratedComponents.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/RehydratedComponents.java
@@ -189,6 +189,7 @@ public class RehydratedComponents {
windowingStrategies.asMap(),
coders.asMap(),
Collections.emptyMap(),
- requirements);
+ requirements,
+ pipeline.getOptions());
}
}
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SdkComponents.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SdkComponents.java
index 446697f24a8..6288649aba3 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SdkComponents.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SdkComponents.java
@@ -63,6 +63,7 @@ public class SdkComponents {
private final BiMap<Environment, String> environmentIds = HashBiMap.create();
private final BiMap<RunnerApi.Coder, String> coderProtoToId =
HashBiMap.create();
private final Set<String> requirements;
+ private final PipelineOptions pipelineOptions;
private final Set<String> reservedIds = new HashSet<>();
@@ -71,17 +72,7 @@ public class SdkComponents {
/** Create a new {@link SdkComponents} with no components. */
public static SdkComponents create() {
- return new SdkComponents(RunnerApi.Components.getDefaultInstance(), null,
"");
- }
-
- /**
- * Create new {@link SdkComponents} importing all items from provided {@link
Components} object.
- *
- * <p>WARNING: This action might cause some of duplicate items created.
- */
- public static SdkComponents create(
- RunnerApi.Components components, Collection<String> requirements) {
- return new SdkComponents(components, requirements, "");
+ return new SdkComponents(RunnerApi.Components.getDefaultInstance(), null,
"", null);
}
/*package*/ static SdkComponents create(
@@ -91,8 +82,9 @@ public class SdkComponents {
Map<String, WindowingStrategy<?, ?>> windowingStrategies,
Map<String, Coder<?>> coders,
Map<String, Environment> environments,
- Collection<String> requirements) {
- SdkComponents sdkComponents = SdkComponents.create(components,
requirements);
+ Collection<String> requirements,
+ PipelineOptions pipelineOptions) {
+ SdkComponents sdkComponents = new SdkComponents(components, requirements,
"", pipelineOptions);
sdkComponents.transformIds.inverse().putAll(transforms);
sdkComponents.pCollectionIds.inverse().putAll(pCollections);
sdkComponents.windowingStrategyIds.inverse().putAll(windowingStrategies);
@@ -103,19 +95,28 @@ public class SdkComponents {
public static SdkComponents create(PipelineOptions options) {
SdkComponents sdkComponents =
- new SdkComponents(RunnerApi.Components.getDefaultInstance(), null, "");
+ new SdkComponents(RunnerApi.Components.getDefaultInstance(), null, "",
options);
PortablePipelineOptions portablePipelineOptions =
options.as(PortablePipelineOptions.class);
sdkComponents.registerEnvironment(
Environments.createOrGetDefaultEnvironment(portablePipelineOptions));
return sdkComponents;
}
+ public static SdkComponents create(PipelineOptions options, Environment
environment) {
+ SdkComponents sdkComponents =
+ new SdkComponents(RunnerApi.Components.getDefaultInstance(), null, "",
options);
+ sdkComponents.registerEnvironment(environment);
+ return sdkComponents;
+ }
+
private SdkComponents(
@Nullable Components components,
@Nullable Collection<String> requirements,
- String newIdPrefix) {
+ String newIdPrefix,
+ @Nullable PipelineOptions pipelineOptions) {
this.newIdPrefix = newIdPrefix;
this.requirements = new HashSet<>();
+ this.pipelineOptions = pipelineOptions;
if (components == null) {
if (requirements != null) {
@@ -153,7 +154,7 @@ public class SdkComponents {
*/
public SdkComponents withNewIdPrefix(String newIdPrefix) {
SdkComponents sdkComponents =
- new SdkComponents(componentsBuilder.build(), requirements,
newIdPrefix);
+ new SdkComponents(componentsBuilder.build(), requirements,
newIdPrefix, pipelineOptions);
sdkComponents.transformIds.putAll(transformIds);
sdkComponents.pCollectionIds.putAll(pCollectionIds);
sdkComponents.windowingStrategyIds.putAll(windowingStrategyIds);
@@ -174,7 +175,7 @@ public class SdkComponents {
throws IOException {
String name = getApplicationName(appliedPTransform);
// If this transform is present in the components, nothing to do. return
the existing name.
- // Otherwise the transform must be translated and added to the components.
+ // Otherwise, the transform must be translated and added to the components.
if (componentsBuilder.getTransformsOrDefault(name, null) != null) {
return name;
}
@@ -375,4 +376,8 @@ public class SdkComponents {
public Collection<String> requirements() {
return ImmutableSet.copyOf(requirements);
}
+
+ public PipelineOptions getPipelineOptions() {
+ return pipelineOptions;
+ }
}
diff --git
a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/CoderTranslationTest.java
b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/CoderTranslationTest.java
index b8f92ff0053..1ec0a74f5be 100644
---
a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/CoderTranslationTest.java
+++
b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/CoderTranslationTest.java
@@ -22,6 +22,7 @@ import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.hasItems;
import static org.hamcrest.Matchers.not;
+import com.google.auto.value.AutoValue;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
@@ -45,14 +46,20 @@ import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder;
import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.Field;
import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
import org.apache.beam.sdk.schemas.logicaltypes.FixedBytes;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import
org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder;
import org.apache.beam.sdk.util.ShardedKey;
import
org.apache.beam.sdk.util.construction.CoderTranslation.TranslationContext;
+import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
@@ -70,6 +77,34 @@ import org.junit.runners.Parameterized.Parameters;
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
})
public class CoderTranslationTest {
+ @AutoValue
+ @DefaultSchema(AutoValueSchema.class)
+ public abstract static class SimpleAutoValue {
+ public abstract String getString();
+
+ public abstract int getInt32();
+
+ public abstract long getInt64();
+
+ public static SimpleAutoValue of(String string, Integer int32, Long int64)
{
+ return new AutoValue_CoderTranslationTest_SimpleAutoValue(string, int32,
int64);
+ }
+ }
+
+ private static final SchemaRegistry REGISTRY =
SchemaRegistry.createDefault();
+
+ private static SchemaCoder schemaCoderFrom(TypeDescriptor typeDescriptor) {
+ try {
+ return SchemaCoder.of(
+ REGISTRY.getSchema(typeDescriptor),
+ typeDescriptor,
+ REGISTRY.getToRowFunction(typeDescriptor),
+ REGISTRY.getFromRowFunction(typeDescriptor));
+ } catch (NoSuchSchemaException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
private static final Set<Coder<?>> KNOWN_CODERS =
ImmutableSet.<Coder<?>>builder()
.add(ByteArrayCoder.of())
@@ -94,6 +129,7 @@ public class CoderTranslationTest {
Field.of("array", FieldType.array(FieldType.STRING)),
Field.of("map", FieldType.map(FieldType.STRING,
FieldType.INT32)),
Field.of("bar",
FieldType.logicalType(FixedBytes.of(123))))))
+ .add(schemaCoderFrom(TypeDescriptor.of(SimpleAutoValue.class)))
.add(ShardedKey.Coder.of(StringUtf8Coder.of()))
.add(TimestampPrefixingWindowCoder.of(IntervalWindowCoder.of()))
.add(NullableCoder.of(ByteArrayCoder.of()))
@@ -127,7 +163,7 @@ public class CoderTranslationTest {
}
@Test
- public void validateCoderTranslators() {
+ public void validateModelCoderTranslators() {
assertThat(
"Every Model Coder must have a Translator",
new ModelCoderRegistrar().getCoderURNs().keySet(),
diff --git
a/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/AvroGenericCoderRegistrar.java
b/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/AvroGenericCoderRegistrar.java
index 14ab48f6669..8bd18fd8e25 100644
---
a/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/AvroGenericCoderRegistrar.java
+++
b/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/AvroGenericCoderRegistrar.java
@@ -21,9 +21,11 @@ import com.google.auto.service.AutoService;
import java.util.Map;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.extensions.avro.coders.AvroGenericCoder;
+import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.construction.CoderTranslator;
import org.apache.beam.sdk.util.construction.CoderTranslatorRegistrar;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** Coder registrar for AvroGenericCoder. */
@AutoService(CoderTranslatorRegistrar.class)
@@ -42,4 +44,20 @@ public class AvroGenericCoderRegistrar implements
CoderTranslatorRegistrar {
public Map<Class<? extends Coder>, CoderTranslator<? extends Coder>>
getCoderTranslators() {
return ImmutableMap.of(AvroGenericCoder.class, new
AvroGenericCoderTranslator());
}
+
+ @Override
+ public boolean isKnownCoder(Coder<?> coder, PipelineOptions options) {
+ return coder.getClass() == AvroGenericCoder.class;
+ }
+
+ @Override
+ public @Nullable CoderTranslator<? extends Coder> getCoderTranslator(
+ Class<? extends Coder> coderClass) {
+ return coderClass == AvroGenericCoder.class ? new
AvroGenericCoderTranslator() : null;
+ }
+
+ @Override
+ public @Nullable Class<? extends Coder> getCoderForUrn(String coderUrn) {
+ return AVRO_CODER_URN.equals(coderUrn) ? AvroGenericCoder.class : null;
+ }
}
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateBackedIterable.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateBackedIterable.java
index ef8d69bc1ec..42a6f8d11c2 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateBackedIterable.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateBackedIterable.java
@@ -38,6 +38,7 @@ import
org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
import org.apache.beam.sdk.coders.IterableLikeCoder;
import org.apache.beam.sdk.fn.stream.PrefetchableIterable;
import org.apache.beam.sdk.fn.stream.PrefetchableIterators;
+import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.BufferedElementCountingOutputStream;
import org.apache.beam.sdk.util.VarInt;
import org.apache.beam.sdk.util.common.ElementByteSizeObservableIterable;
@@ -52,6 +53,7 @@ import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Immuta
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams;
+import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -300,6 +302,26 @@ public class StateBackedIterable<T>
getCoderTranslators() {
return ImmutableMap.of(StateBackedIterable.Coder.class, new
Translator());
}
+
+ @Override
+ public boolean isKnownCoder(
+ org.apache.beam.sdk.coders.Coder<?> coder, PipelineOptions options) {
+ return coder.getClass() == StateBackedIterable.Coder.class;
+ }
+
+ @Override
+ public @Nullable CoderTranslator<? extends
org.apache.beam.sdk.coders.Coder> getCoderTranslator(
+ Class<? extends org.apache.beam.sdk.coders.Coder> coderClass) {
+ return coderClass == StateBackedIterable.Coder.class ? new Translator()
: null;
+ }
+
+ @Override
+ public @Nullable Class<? extends org.apache.beam.sdk.coders.Coder>
getCoderForUrn(
+ String coderUrn) {
+ return STATE_BACKED_ITERABLE_CODER_URN.equals(coderUrn)
+ ? StateBackedIterable.Coder.class
+ : null;
+ }
}
/**
diff --git
a/sdks/python/apache_beam/runners/portability/fn_api_runner/execution.py
b/sdks/python/apache_beam/runners/portability/fn_api_runner/execution.py
index d2f0f97a8a3..fe2421f8ee8 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/execution.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/execution.py
@@ -780,9 +780,9 @@ class FnApiRunnerExecutionContext(object):
if transform.spec.urn == bundle_processor.DATA_INPUT_URN:
coder_id = self.data_channel_coders[only_element(
transform.outputs.values())]
- coder = self.pipeline_context.coders[self.safe_coders.get(
- coder_id, coder_id)]
if transform.spec.payload == translations.IMPULSE_BUFFER:
+ coder = self.pipeline_context.coders[self.safe_coders.get(
+ coder_id, coder_id)]
data_input[transform.unique_name] = ListBuffer(coder.get_impl())
data_input[transform.unique_name].append(ENCODED_IMPULSE_VALUE)
else:
diff --git
a/sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py
b/sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py
index 7261a815f7e..c77b10cd855 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py
@@ -387,8 +387,8 @@ class TransformContext(object):
_COMMON_CODER_URNS = set(
value.urn for (key, value) in common_urns.coders.__dict__.items()
if not key.startswith('_')
- # Length prefix Rows rather than re-coding them.
- ) - set([common_urns.coders.ROW.urn])
+ # Length prefix Rows and Schemas rather than re-coding them.
+ ) - set([common_urns.coders.ROW.urn, common_urns.coders.SCHEMA.urn])
_REQUIRED_CODER_URNS = set([
common_urns.coders.WINDOWED_VALUE.urn,
