Adds a new coder translator for Java SchemaCoder. (#37631)

* Adds a new coder translator for Java SchemaCoder. Adds PipelineOptions to translation context so we can disable the new translator based on pipeline compatibility version.
* Refactors ModelCoderRegistrars so decisions about known coders are made there instead of in the individual CoderTranslators. Removed SdkComponents from getCoderTranslator and getCoderForUrn since it wasn't needed yet.
* Adds option to pass environment to SdkComponents. Otherwise, we register a default environment in the constructor, which was preventing DataflowRunner from registering its own custom environment.
* Fixup DataflowPipelineTranslatorTest to also use the new SdkComponent constructor.

Author: acrites

CHANGES.md

@@ -82,6 +82,9 @@
 
 ## Breaking Changes
 
+* 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.73")](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)).
 * (Python) Made Beartype the default fallback type checking tool. This can be disabled with the `--disable_beartype` pipeline option. ([#38275](https://github.com/apache/beam/issues/38275))
 * X behavior was changed ([#X](https://github.com/apache/beam/issues/X)).
 
@@ -2435,4 +2438,4 @@ Schema Options, it will be removed in version `2.23.0`. ([BEAM-9704](https://iss
 
 ## Highlights
 
-- For versions 2.19.0 and older release notes are available on [Apache Beam Blog](https://beam.apache.org/blog/).
+- For versions 2.19.0 and older release notes are available on [Apache Beam Blog](https://beam.apache.org/blog/).

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java

@@ -221,7 +221,7 @@ public Boolean visit(OrFinallyTrigger trigger) {
   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));

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java

@@ -1333,19 +1333,20 @@ public DataflowPipelineJob run(Pipeline pipeline) {
     // 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);
@@ -1374,28 +1375,30 @@ public DataflowPipelineJob run(Pipeline pipeline) {
     options.as(SdkHarnessOptions.class).setPipelineProtoHash(pipelineProtoHash);
 
     if (useUnifiedWorker(options)) {
-      LOG.info("Skipping v1 transform replacements since job will run on v2.");
+      LOG.info(
+          "Skipping non-portable transform replacements since job will run on portable worker.");
     } else {
-      // Now rewrite things to be as needed for v1 (mutates the pipeline)
-      // This way the job submitted is valid for v1 and v2, simultaneously
+      // Now rewrite things to be as needed for non-portable (mutates the pipeline).
+      // This way the job submitted is valid for portable and non-portable, 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());
-    // No need to perform transform upgrading for the Runner v1 proto.
-    RunnerApi.Pipeline dataflowV1PipelineProto =
-        PipelineTranslation.toProto(pipeline, dataflowV1Components, true, false);
+    // 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 non-portable runner proto.
+    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.
@@ -1415,7 +1418,11 @@ public DataflowPipelineJob run(Pipeline pipeline) {
 
     JobSpecification jobSpecification =
         translator.translate(
-            pipeline, dataflowV1PipelineProto, dataflowV1Components, this, packages);
+            pipeline,
+            dataflowNonPortablePipelineProto,
+            dataflowNonPortableComponents,
+            this,
+            packages);
 
     if (!isNullOrEmpty(dataflowOptions.getDataflowWorkerJar()) && !useUnifiedWorker(options)) {
       List<String> experiments =

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.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.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 String apply(byte[] input) {
     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,53 @@ public OffsetRange getInitialRange(@SuppressWarnings("unused") @Element String e
       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))));
+    {
+      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());
+    }
+
+    // Prior to version 2.74, SchemaCoders are translated as custom java coders.
+    {
+      options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.73");
+      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());
+    }
+  }
 }

runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptorsTest.java

@@ -78,7 +78,8 @@ public void testLengthPrefixingOfKeyCoderInStatefulExecutableStage() throws Exce
     // 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 void onTimer() {}
   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(

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslation.java

@@ -25,11 +25,13 @@
 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 @@ private static class DefaultTranslationContext implements TranslationContext {}
 
   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 @@ static BiMap<Class<? extends Coder>, String> getKnownCoderUrns() {
     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 static RunnerApi.MessageWithComponents toProto(Coder<?> coder) throws IOE
 
   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 @@ private static RunnerApi.Coder toUnknownCoderWrapper(UnknownCoderWrapper coder)
 
   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 @@ private static Coder<?> fromKnownCoder(
                   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);

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslator.java

@@ -28,6 +28,7 @@
  * additional payload, which is not currently supported. This exists as a temporary measure.
  */
 public interface CoderTranslator<T extends Coder<?>> {
+
   /** Extract all component {@link Coder coders} within a coder. */
   List<? extends Coder<?>> getComponents(T from);
 

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/CoderTranslatorRegistrar.java

@@ -19,6 +19,8 @@
 
 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);
 }