Managed Transform protos & translation; Iceberg SchemaTransforms & translation

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto

@@ -443,12 +443,6 @@ message StandardPTransforms {
   enum GroupIntoBatchesComponents {
     GROUP_INTO_BATCHES = 0 [(beam_urn) = "beam:transform:group_into_batches:v1"];
   }
-  // Payload for all of these: SchemaTransformPayload (in external_transforms.proto)
-  enum ManagedSchemaTransforms {
-    MANAGED = 0 [(beam_urn) = "beam:transform:managed:v1"];
-    ICEBERG_READ = 1 [(beam_urn) = "beam:schematransform:org.apache.beam:iceberg_read:v1"];
-    ICEBERG_WRITE = 2 [(beam_urn) = "beam:schematransform:org.apache.beam:iceberg_write:v1"];
-  }
 }
 
 message StandardSideInputTypes {

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

@@ -35,7 +35,6 @@
 import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms;
 import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms.CombineComponents;
 import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms.SplittableParDoComponents;
-import org.apache.beam.model.pipeline.v1.SchemaAwareTransforms.ManagedSchemaAwareTransforms;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.coders.RowCoder;
@@ -93,7 +92,6 @@ public class PTransformTranslation {
   public static final String MAP_WINDOWS_TRANSFORM_URN = "beam:transform:map_windows:v1";
   public static final String MERGE_WINDOWS_TRANSFORM_URN = "beam:transform:merge_windows:v1";
   public static final String TO_STRING_TRANSFORM_URN = "beam:transform:to_string:v1";
-  public static final String MANAGED_TRANSFORM_URN = "beam:transform:managed:v1";
 
   // Required runner implemented transforms. These transforms should never specify an environment.
   public static final ImmutableSet<String> RUNNER_IMPLEMENTED_TRANSFORMS =
@@ -184,7 +182,6 @@ public class PTransformTranslation {
         MAP_WINDOWS_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.MAP_WINDOWS)));
     checkState(
         MERGE_WINDOWS_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.MERGE_WINDOWS)));
-    checkState(MANAGED_TRANSFORM_URN.equals(getUrn(ManagedSchemaAwareTransforms.Managed.MANAGED)));
 
     // DeprecatedPrimitives
     checkState(READ_TRANSFORM_URN.equals(getUrn(StandardPTransforms.DeprecatedPrimitives.READ)));

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergSchemaTransformTranslation.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.io.iceberg;
 
+import static org.apache.beam.model.pipeline.v1.ExternalTransforms.ExpansionMethods.Enum.SCHEMA_TRANSFORM;
 import static org.apache.beam.sdk.io.iceberg.IcebergReadSchemaTransformProvider.IcebergReadSchemaTransform;
 import static org.apache.beam.sdk.io.iceberg.IcebergWriteSchemaTransformProvider.IcebergWriteSchemaTransform;
 
@@ -33,6 +34,7 @@
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.SchemaTranslation;
 import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.util.construction.BeamUrns;
 import org.apache.beam.sdk.util.construction.PTransformTranslation.TransformPayloadTranslator;
 import org.apache.beam.sdk.util.construction.SdkComponents;
 import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
@@ -51,38 +53,47 @@ static class IcebergReadSchemaTransformTranslator
 
     @Override
     public String getUrn() {
-      return READ_PROVIDER.identifier();
+      return BeamUrns.getUrn(SCHEMA_TRANSFORM);
     }
 
     @Override
     public @Nullable FunctionSpec translate(
         AppliedPTransform<?, ?, IcebergReadSchemaTransform> application, SdkComponents components)
         throws IOException {
-      SchemaApi.Schema expansionSchema = SchemaTranslation.schemaToProto(READ_SCHEMA, true);
+      Schema snakeCaseSchema = READ_SCHEMA.toSnakeCase();
+      SchemaApi.Schema expansionSchema = SchemaTranslation.schemaToProto(snakeCaseSchema, true);
       Row configRow = toConfigRow(application.getTransform());
       ByteArrayOutputStream os = new ByteArrayOutputStream();
-      RowCoder.of(READ_SCHEMA).encode(configRow, os);
+      RowCoder.of(snakeCaseSchema).encode(configRow, os);
 
       return FunctionSpec.newBuilder()
           .setUrn(getUrn())
           .setPayload(
               SchemaTransformPayload.newBuilder()
+                  .setIdentifier(READ_PROVIDER.identifier())
                   .setConfigurationSchema(expansionSchema)
                   .setConfigurationRow(ByteString.copyFrom(os.toByteArray()))
-                  .setIdentifier(getUrn())
                   .build()
                   .toByteString())
           .build();
     }
 
     @Override
     public Row toConfigRow(IcebergReadSchemaTransform transform) {
-      return transform.getConfigurationRow();
+      // Will retrieve a Row with snake_case naming convention.
+      // Transform expects camelCase convention, so convert back
+      // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+      // TypedSchemaTransformProvider starts generating with snake_case convention
+      return transform.getConfigurationRow().toSnakeCase();
     }
 
     @Override
     public IcebergReadSchemaTransform fromConfigRow(Row configRow, PipelineOptions options) {
-      return (IcebergReadSchemaTransform) READ_PROVIDER.from(configRow);
+      // Will retrieve a Row with snake_case naming convention.
+      // Transform expects camelCase convention, so convert back
+      // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+      // TypedSchemaTransformProvider starts generating with snake_case convention
+      return (IcebergReadSchemaTransform) READ_PROVIDER.from(configRow.toCamelCase());
     }
   }
 
@@ -109,38 +120,48 @@ static class IcebergWriteSchemaTransformTranslator
 
     @Override
     public String getUrn() {
-      return WRITE_PROVIDER.identifier();
+      return BeamUrns.getUrn(SCHEMA_TRANSFORM);
     }
 
     @Override
     public @Nullable FunctionSpec translate(
         AppliedPTransform<?, ?, IcebergWriteSchemaTransform> application, SdkComponents components)
         throws IOException {
-      SchemaApi.Schema expansionSchema = SchemaTranslation.schemaToProto(WRITE_SCHEMA, true);
+      // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+      // TypedSchemaTransformProvider starts generating with snake_case convention
+      Schema snakeCaseSchema = WRITE_SCHEMA.toSnakeCase();
+      SchemaApi.Schema expansionSchema = SchemaTranslation.schemaToProto(snakeCaseSchema, true);
       Row configRow = toConfigRow(application.getTransform());
       ByteArrayOutputStream os = new ByteArrayOutputStream();
-      RowCoder.of(WRITE_SCHEMA).encode(configRow, os);
+      RowCoder.of(snakeCaseSchema).encode(configRow, os);
 
       return FunctionSpec.newBuilder()
           .setUrn(getUrn())
           .setPayload(
               SchemaTransformPayload.newBuilder()
+                  .setIdentifier(WRITE_PROVIDER.identifier())
                   .setConfigurationSchema(expansionSchema)
                   .setConfigurationRow(ByteString.copyFrom(os.toByteArray()))
-                  .setIdentifier(getUrn())
                   .build()
                   .toByteString())
           .build();
     }
 
     @Override
     public Row toConfigRow(IcebergWriteSchemaTransform transform) {
-      return transform.getConfigurationRow();
+      // Return with snake_case naming convention!
+      // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+      // TypedSchemaTransformProvider starts generating with snake_case convention
+      return transform.getConfigurationRow().toSnakeCase();
     }
 
     @Override
     public IcebergWriteSchemaTransform fromConfigRow(Row configRow, PipelineOptions options) {
-      return (IcebergWriteSchemaTransform) WRITE_PROVIDER.from(configRow);
+      // Will retrieve a Row with snake_case naming convention.
+      // Transform expects camelCase convention, so convert back
+      // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+      // TypedSchemaTransformProvider starts generating with snake_case convention
+      return (IcebergWriteSchemaTransform) WRITE_PROVIDER.from(configRow.toCamelCase());
     }
   }
 

sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergSchemaTransformTranslationTest.java

@@ -17,13 +17,15 @@
  */
 package org.apache.beam.sdk.io.iceberg;
 
+import static org.apache.beam.model.pipeline.v1.ExternalTransforms.ExpansionMethods.Enum.SCHEMA_TRANSFORM;
 import static org.apache.beam.sdk.io.iceberg.IcebergReadSchemaTransformProvider.IcebergReadSchemaTransform;
 import static org.apache.beam.sdk.io.iceberg.IcebergSchemaTransformTranslation.IcebergReadSchemaTransformTranslator.READ_SCHEMA;
 import static org.apache.beam.sdk.io.iceberg.IcebergSchemaTransformTranslation.IcebergWriteSchemaTransformTranslator.WRITE_SCHEMA;
 import static org.apache.beam.sdk.io.iceberg.IcebergWriteSchemaTransformProvider.INPUT_TAG;
 import static org.apache.beam.sdk.io.iceberg.IcebergWriteSchemaTransformProvider.IcebergWriteSchemaTransform;
 import static org.junit.Assert.assertEquals;
 
+import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
 import java.util.UUID;
@@ -32,16 +34,16 @@
 import org.apache.beam.model.pipeline.v1.RunnerApi;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.RowCoder;
-import org.apache.beam.sdk.managed.ManagedTransformConstants;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.schemas.NoSuchSchemaException;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.SchemaTranslation;
 import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.util.construction.BeamUrns;
 import org.apache.beam.sdk.util.construction.PipelineTranslation;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionRowTuple;
 import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.InvalidProtocolBufferException;
 import org.apache.iceberg.CatalogUtil;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.junit.ClassRule;
@@ -64,7 +66,7 @@ public class IcebergSchemaTransformTranslationTest {
       new IcebergReadSchemaTransformProvider();
 
   @Test
-  public void testReCreateWriteTransformFromRow() throws NoSuchSchemaException {
+  public void testReCreateWriteTransformFromRow() {
     Row catalogConfigRow =
         Row.withSchema(IcebergSchemaTransformCatalogConfig.SCHEMA)
             .withFieldValue("catalogName", "test_name")
@@ -91,7 +93,8 @@ public void testReCreateWriteTransformFromRow() throws NoSuchSchemaException {
   }
 
   @Test
-  public void testWriteTransformProtoTranslation() throws Exception {
+  public void testWriteTransformProtoTranslation()
+      throws InvalidProtocolBufferException, IOException {
     // First build a pipeline
     Pipeline p = Pipeline.create();
     Schema inputSchema = Schema.builder().addStringField("str").build();
@@ -122,18 +125,31 @@ public void testWriteTransformProtoTranslation() throws Exception {
     RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
     List<RunnerApi.PTransform> writeTransformProto =
         pipelineProto.getComponents().getTransformsMap().values().stream()
-            .filter(tr -> tr.getSpec().getUrn().equals(ManagedTransformConstants.ICEBERG_WRITE))
+            .filter(
+                tr -> {
+                  RunnerApi.FunctionSpec spec = tr.getSpec();
+                  try {
+                    return spec.getUrn().equals(BeamUrns.getUrn(SCHEMA_TRANSFORM))
+                        && SchemaTransformPayload.parseFrom(spec.getPayload())
+                            .getIdentifier()
+                            .equals(WRITE_PROVIDER.identifier());
+                  } catch (InvalidProtocolBufferException e) {
+                    throw new RuntimeException(e);
+                  }
+                })
             .collect(Collectors.toList());
     assertEquals(1, writeTransformProto.size());
     RunnerApi.FunctionSpec spec = writeTransformProto.get(0).getSpec();
 
     // Check that the proto contains correct values
     SchemaTransformPayload payload = SchemaTransformPayload.parseFrom(spec.getPayload());
     Schema schemaFromSpec = SchemaTranslation.schemaFromProto(payload.getConfigurationSchema());
-    assertEquals(WRITE_SCHEMA, schemaFromSpec);
+    // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+    // TypedSchemaTransformProvider starts generating with snake_case convention
+    assertEquals(WRITE_SCHEMA.toSnakeCase(), schemaFromSpec);
     Row rowFromSpec = RowCoder.of(schemaFromSpec).decode(payload.getConfigurationRow().newInput());
 
-    assertEquals(transformConfigRow, rowFromSpec);
+    assertEquals(transformConfigRow, rowFromSpec.toCamelCase());
 
     // Use the information in the proto to recreate the IcebergWriteSchemaTransform
     IcebergSchemaTransformTranslation.IcebergWriteSchemaTransformTranslator translator =
@@ -145,7 +161,7 @@ public void testWriteTransformProtoTranslation() throws Exception {
   }
 
   @Test
-  public void testReCreateReadTransformFromRow() throws NoSuchSchemaException {
+  public void testReCreateReadTransformFromRow() {
     // setting a subset of fields here.
     Row catalogConfigRow =
         Row.withSchema(IcebergSchemaTransformCatalogConfig.SCHEMA)
@@ -174,7 +190,8 @@ public void testReCreateReadTransformFromRow() throws NoSuchSchemaException {
   }
 
   @Test
-  public void testReadTransformProtoTranslation() throws Exception {
+  public void testReadTransformProtoTranslation()
+      throws InvalidProtocolBufferException, IOException {
     // First build a pipeline
     Pipeline p = Pipeline.create();
     Row catalogConfigRow =
@@ -201,17 +218,30 @@ public void testReadTransformProtoTranslation() throws Exception {
     RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
     List<RunnerApi.PTransform> readTransformProto =
         pipelineProto.getComponents().getTransformsMap().values().stream()
-            .filter(tr -> tr.getSpec().getUrn().equals(ManagedTransformConstants.ICEBERG_READ))
+            .filter(
+                tr -> {
+                  RunnerApi.FunctionSpec spec = tr.getSpec();
+                  try {
+                    return spec.getUrn().equals(BeamUrns.getUrn(SCHEMA_TRANSFORM))
+                        && SchemaTransformPayload.parseFrom(spec.getPayload())
+                            .getIdentifier()
+                            .equals(READ_PROVIDER.identifier());
+                  } catch (InvalidProtocolBufferException e) {
+                    throw new RuntimeException(e);
+                  }
+                })
             .collect(Collectors.toList());
     assertEquals(1, readTransformProto.size());
     RunnerApi.FunctionSpec spec = readTransformProto.get(0).getSpec();
 
     // Check that the proto contains correct values
     SchemaTransformPayload payload = SchemaTransformPayload.parseFrom(spec.getPayload());
     Schema schemaFromSpec = SchemaTranslation.schemaFromProto(payload.getConfigurationSchema());
-    assertEquals(READ_SCHEMA, schemaFromSpec);
+    // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+    // TypedSchemaTransformProvider starts generating with snake_case convention
+    assertEquals(READ_SCHEMA.toSnakeCase(), schemaFromSpec);
     Row rowFromSpec = RowCoder.of(schemaFromSpec).decode(payload.getConfigurationRow().newInput());
-    assertEquals(transformConfigRow, rowFromSpec);
+    assertEquals(transformConfigRow, rowFromSpec.toCamelCase());
 
     // Use the information in the proto to recreate the IcebergReadSchemaTransform
     IcebergSchemaTransformTranslation.IcebergReadSchemaTransformTranslator translator =

sdks/java/managed/src/main/java/org/apache/beam/sdk/managed/ManagedSchemaTransformTranslation.java

@@ -17,8 +17,8 @@
  */
 package org.apache.beam.sdk.managed;
 
+import static org.apache.beam.model.pipeline.v1.ExternalTransforms.ExpansionMethods.Enum.SCHEMA_TRANSFORM;
 import static org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
-import static org.apache.beam.sdk.managed.ManagedSchemaTransformProvider.ManagedConfig;
 import static org.apache.beam.sdk.managed.ManagedSchemaTransformProvider.ManagedSchemaTransform;
 import static org.apache.beam.sdk.util.construction.PTransformTranslation.TransformPayloadTranslator;
 
@@ -31,11 +31,10 @@
 import org.apache.beam.sdk.coders.RowCoder;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.runners.AppliedPTransform;
-import org.apache.beam.sdk.schemas.NoSuchSchemaException;
 import org.apache.beam.sdk.schemas.Schema;
-import org.apache.beam.sdk.schemas.SchemaRegistry;
 import org.apache.beam.sdk.schemas.SchemaTranslation;
 import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.util.construction.BeamUrns;
 import org.apache.beam.sdk.util.construction.SdkComponents;
 import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
 import org.apache.beam.sdk.values.Row;
@@ -51,25 +50,27 @@ static class ManagedSchemaTransformTranslator
 
     @Override
     public String getUrn() {
-      return PROVIDER.identifier();
+      return BeamUrns.getUrn(SCHEMA_TRANSFORM);
     }
 
     @Override
     @SuppressWarnings("argument")
     public @Nullable FunctionSpec translate(
         AppliedPTransform<?, ?, ManagedSchemaTransform> application, SdkComponents components)
         throws IOException {
-      SchemaApi.Schema expansionSchema = SchemaTranslation.schemaToProto(SCHEMA, true);
-      ManagedConfig managedConfig = application.getTransform().getManagedConfig();
+      // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+      // TypedSchemaTransformProvider starts generating with snake_case convention
+      Schema snakeCaseSchema = SCHEMA.toSnakeCase();
+      SchemaApi.Schema expansionSchema = SchemaTranslation.schemaToProto(snakeCaseSchema, true);
       Row configRow = toConfigRow(application.getTransform());
       ByteArrayOutputStream os = new ByteArrayOutputStream();
-      RowCoder.of(SCHEMA).encode(configRow, os);
+      RowCoder.of(snakeCaseSchema).encode(configRow, os);
 
       return FunctionSpec.newBuilder()
           .setUrn(getUrn())
           .setPayload(
               SchemaTransformPayload.newBuilder()
-                  .setIdentifier(managedConfig.getTransformIdentifier())
+                  .setIdentifier(PROVIDER.identifier())
                   .setConfigurationSchema(expansionSchema)
                   .setConfigurationRow(ByteString.copyFrom(os.toByteArray()))
                   .build()
@@ -79,25 +80,19 @@ public String getUrn() {
 
     @Override
     public Row toConfigRow(ManagedSchemaTransform transform) {
-      ManagedConfig managedConfig = transform.getManagedConfig();
-      Row configRow;
-      try {
-        configRow =
-            SchemaRegistry.createDefault()
-                .getToRowFunction(ManagedConfig.class)
-                .apply(managedConfig);
-      } catch (NoSuchSchemaException e) {
-        throw new RuntimeException(e);
-      }
-      // Sanity check: sort fields according to the configuration schema
-      return SCHEMA.getFields().stream()
-          .map(field -> configRow.getValue(field.getName()))
-          .collect(Row.toRow(SCHEMA));
+      // Return with snake_case naming convention!
+      // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+      // TypedSchemaTransformProvider starts generating with snake_case convention
+      return transform.getConfigurationRow().toSnakeCase();
     }
 
     @Override
     public ManagedSchemaTransform fromConfigRow(Row configRow, PipelineOptions options) {
-      return (ManagedSchemaTransform) PROVIDER.from(configRow);
+      // Will retrieve a Row with snake_case naming convention.
+      // Transform expects camelCase convention, so convert back
+      // TODO(https://github.com/apache/beam/issues/31061): Remove conversion when
+      // TypedSchemaTransformProvider starts generating with snake_case convention
+      return (ManagedSchemaTransform) PROVIDER.from(configRow.toCamelCase());
     }
   }