Managed Transform protos & translation; Iceberg SchemaTransforms & translation

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json

@@ -1,3 +1,4 @@
 {
-  "comment": "Modify this file in a trivial way to cause this test suite to run"
+  "comment": "Modify this file in a trivial way to cause this test suite to run",
+  "modification": 1
 }

.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json

@@ -1,3 +1,4 @@
 {
-  "comment": "Modify this file in a trivial way to cause this test suite to run"
+  "comment": "Modify this file in a trivial way to cause this test suite to run",
+  "modification": 1
 }

.github/workflows/beam_PreCommit_Xlang_Generated_Transforms.yml

@@ -111,4 +111,4 @@ jobs:
       - name: run Cross-Language Wrapper Validation script
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:
-          gradle-command: :sdks:python:test-suites:direct:crossLanguageWrapperValidationPreCommit
+          gradle-command: :sdks:python:test-suites:direct:crossLanguageWrapperValidationPreCommit --info

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

@@ -111,6 +111,10 @@ message BuilderMethod {
   bytes payload = 3;
 }
 
+// Payload for a Schema-aware PTransform.
+// This is a transform that is aware of its input and output PCollection schemas
+// and is configured using Beam Schema-compatible parameters.
+// The information available in the payload can be used to instantiate the schema-aware transform.
 message SchemaTransformPayload {
   // The identifier of the SchemaTransform (typically a URN).
   string identifier = 1;

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/Schema.java

@@ -37,6 +37,7 @@
 import java.util.stream.Collectors;
 import javax.annotation.concurrent.Immutable;
 import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.CaseFormat;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 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.HashBiMap;
@@ -326,7 +327,10 @@ public static Schema of(Field... fields) {
     return Schema.builder().addFields(fields).build();
   }
 
-  /** Returns an identical Schema with sorted fields. */
+  /**
+   * Returns an identical Schema with lexicographically sorted fields. Recursively sorts nested
+   * fields.
+   */
   public Schema sorted() {
     // Create a new schema and copy over the appropriate Schema object attributes:
     // {fields, uuid, options}
@@ -336,6 +340,16 @@ public Schema sorted() {
     Schema sortedSchema =
         this.fields.stream()
             .sorted(Comparator.comparing(Field::getName))
+            .map(
+                field -> {
+                  FieldType innerType = field.getType();
+                  if (innerType.getRowSchema() != null) {
+                    Schema innerSortedSchema = innerType.getRowSchema().sorted();
+                    innerType = innerType.toBuilder().setRowSchema(innerSortedSchema).build();
+                    return field.toBuilder().setType(innerType).build();
+                  }
+                  return field;
+                })
             .collect(Schema.toSchema())
             .withOptions(getOptions());
     sortedSchema.setUUID(getUUID());
@@ -1451,4 +1465,42 @@ public int getFieldCount() {
   public Options getOptions() {
     return this.options;
   }
+
+  /** Recursively converts all field names to `snake_case`. */
+  public Schema toSnakeCase() {
+    return this.getFields().stream()
+        .map(
+            field -> {
+              FieldType innerType = field.getType();
+              if (innerType.getRowSchema() != null) {
+                Schema innerSnakeCaseSchema = innerType.getRowSchema().toSnakeCase();
+                innerType = innerType.toBuilder().setRowSchema(innerSnakeCaseSchema).build();
+                field = field.toBuilder().setType(innerType).build();
+              }
+              return field
+                  .toBuilder()
+                  .setName(CaseFormat.LOWER_CAMEL.to(CaseFormat.LOWER_UNDERSCORE, field.getName()))
+                  .build();
+            })
+        .collect(toSchema());
+  }
+
+  /** Recursively converts all field names to `lowerCamelCase`. */
+  public Schema toCamelCase() {
+    return this.getFields().stream()
+        .map(
+            field -> {
+              FieldType innerType = field.getType();
+              if (innerType.getRowSchema() != null) {
+                Schema innerCamelCaseSchema = innerType.getRowSchema().toCamelCase();
+                innerType = innerType.toBuilder().setRowSchema(innerCamelCaseSchema).build();
+                field = field.toBuilder().setType(innerType).build();
+              }
+              return field
+                  .toBuilder()
+                  .setName(CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.LOWER_CAMEL, field.getName()))
+                  .build();
+            })
+        .collect(toSchema());
+  }
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaRegistry.java

@@ -302,6 +302,29 @@ public <T> SchemaCoder<T> getSchemaCoder(TypeDescriptor<T> typeDescriptor)
         getFromRowFunction(typeDescriptor));
   }
 
+  /**
+   * Retrieve a registered {@link SchemaProvider} for a given {@link TypeDescriptor}. If no schema
+   * exists, throws {@link * NoSuchSchemaException}.
+   */
+  public <T> SchemaProvider getSchemaProvider(TypeDescriptor<T> typeDescriptor)
+      throws NoSuchSchemaException {
+    for (SchemaProvider provider : providers) {
+      Schema schema = provider.schemaFor(typeDescriptor);
+      if (schema != null) {
+        return provider;
+      }
+    }
+    throw new NoSuchSchemaException();
+  }
+
+  /**
+   * Retrieve a registered {@link SchemaProvider} for a given {@link Class}. If no schema exists,
+   * throws {@link * NoSuchSchemaException}.
+   */
+  public <T> SchemaProvider getSchemaProvider(Class<T> clazz) throws NoSuchSchemaException {
+    return getSchemaProvider(TypeDescriptor.of(clazz));
+  }
+
   private <ReturnT> ReturnT getProviderResult(Function<SchemaProvider, ReturnT> f)
       throws NoSuchSchemaException {
     for (SchemaProvider provider : providers) {

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/annotations/DefaultSchema.java

@@ -123,6 +123,24 @@ public ProviderAndDescriptor(
           });
     }
 
+    /**
+     * Retrieves the underlying {@link SchemaProvider} for the given {@link TypeDescriptor}. If no
+     * provider is found, returns null.
+     */
+    public @Nullable <T> SchemaProvider getUnderlyingSchemaProvider(
+        TypeDescriptor<T> typeDescriptor) {
+      ProviderAndDescriptor providerAndDescriptor = getSchemaProvider(typeDescriptor);
+      return providerAndDescriptor != null ? providerAndDescriptor.schemaProvider : null;
+    }
+
+    /**
+     * Retrieves the underlying {@link SchemaProvider} for the given {@link Class}. If no provider
+     * is found, returns null.
+     */
+    public @Nullable <T> SchemaProvider getUnderlyingSchemaProvider(Class<T> clazz) {
+      return getUnderlyingSchemaProvider(TypeDescriptor.of(clazz));
+    }
+
     @Override
     public <T> Schema schemaFor(TypeDescriptor<T> typeDescriptor) {
       ProviderAndDescriptor providerAndDescriptor = getSchemaProvider(typeDescriptor);

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/transforms/TypedSchemaTransformProvider.java

@@ -17,17 +17,21 @@
  */
 package org.apache.beam.sdk.schemas.transforms;
 
+import static org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider;
 import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 
 import java.lang.reflect.ParameterizedType;
 import java.util.List;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.options.PipelineOptions;
+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.SchemaProvider;
 import org.apache.beam.sdk.schemas.SchemaRegistry;
 import org.apache.beam.sdk.values.Row;
 
@@ -77,7 +81,8 @@ Optional<List<String>> dependencies(ConfigT configuration, PipelineOptions optio
   public final Schema configurationSchema() {
     try {
       // Sort the fields by name to ensure a consistent schema is produced
-      return SchemaRegistry.createDefault().getSchema(configurationClass()).sorted();
+      // We also establish a `snake_case` convention for all SchemaTransform configurations
+      return SchemaRegistry.createDefault().getSchema(configurationClass()).sorted().toSnakeCase();
     } catch (NoSuchSchemaException e) {
       throw new RuntimeException(
           "Unable to find schema for "
@@ -86,6 +91,10 @@ public final Schema configurationSchema() {
     }
   }
 
+  /**
+   * Produces a {@link SchemaTransform} from a Row configuration. Row fields are expected to have
+   * `snake_case` naming convention.
+   */
   @Override
   public final SchemaTransform from(Row configuration) {
     return from(configFromRow(configuration));
@@ -98,9 +107,20 @@ public final Optional<List<String>> dependencies(Row configuration, PipelineOpti
 
   private ConfigT configFromRow(Row configuration) {
     try {
-      return SchemaRegistry.createDefault()
-          .getFromRowFunction(configurationClass())
-          .apply(configuration);
+      SchemaRegistry registry = SchemaRegistry.createDefault();
+
+      // Configuration objects handled by the AutoValueSchema provider will expect Row fields with
+      // camelCase naming convention
+      SchemaProvider schemaProvider = registry.getSchemaProvider(configurationClass());
+      if (schemaProvider.getClass().equals(DefaultSchemaProvider.class)
+          && checkNotNull(
+                  ((DefaultSchemaProvider) schemaProvider)
+                      .getUnderlyingSchemaProvider(configurationClass()))
+              .getClass()
+              .equals(AutoValueSchema.class)) {
+        configuration = configuration.toCamelCase();
+      }
+      return registry.getFromRowFunction(configurationClass()).apply(configuration);
     } catch (NoSuchSchemaException e) {
       throw new RuntimeException(
           "Unable to find schema for " + identifier() + "SchemaTransformProvider's config");

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/YamlUtils.java

@@ -147,14 +147,27 @@ public static Row toBeamRow(
   }
 
   @SuppressWarnings("nullness")
-  public static Row toBeamRow(Map<String, Object> yamlMap, Schema rowSchema, boolean toCamelCase) {
+  public static Row toBeamRow(
+      @Nullable Map<String, Object> map, Schema rowSchema, boolean toCamelCase) {
+    if (map == null || map.isEmpty()) {
+      List<Field> requiredFields =
+          rowSchema.getFields().stream()
+              .filter(field -> !field.getType().getNullable())
+              .collect(Collectors.toList());
+      if (requiredFields.isEmpty()) {
+        return Row.nullRow(rowSchema);
+      } else {
+        throw new IllegalArgumentException(
+            String.format(
+                "Received an empty Map, but output schema contains required fields: %s",
+                requiredFields));
+      }
+    }
     return rowSchema.getFields().stream()
         .map(
             field ->
                 toBeamValue(
-                    field,
-                    yamlMap.get(maybeGetSnakeCase(field.getName(), toCamelCase)),
-                    toCamelCase))
+                    field, map.get(maybeGetSnakeCase(field.getName(), toCamelCase)), toCamelCase))
         .collect(toRow(rowSchema));
   }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/values/Row.java

@@ -893,4 +893,53 @@ public static Row nullRow(Schema schema) {
         .addValues(Collections.nCopies(schema.getFieldCount(), null))
         .build();
   }
+
+  /** Returns an equivalent {@link Row} with fields lexicographically sorted by their name. */
+  public Row sorted() {
+    Schema sortedSchema = getSchema().sorted();
+    return sortedSchema.getFields().stream()
+        .map(
+            field -> {
+              if (field.getType().getRowSchema() != null) {
+                Row innerRow = getValue(field.getName());
+                if (innerRow != null) {
+                  return innerRow.sorted();
+                }
+              }
+              return (Object) getValue(field.getName());
+            })
+        .collect(Row.toRow(sortedSchema));
+  }
+
+  /** Returns an equivalent {@link Row} with `snake_case` field names. */
+  public Row toSnakeCase() {
+    return getSchema().getFields().stream()
+        .map(
+            field -> {
+              if (field.getType().getRowSchema() != null) {
+                Row innerRow = getValue(field.getName());
+                if (innerRow != null) {
+                  return innerRow.toSnakeCase();
+                }
+              }
+              return (Object) getValue(field.getName());
+            })
+        .collect(toRow(getSchema().toSnakeCase()));
+  }
+
+  /** Returns an equivalent {@link Row} with `lowerCamelCase` field names. */
+  public Row toCamelCase() {
+    return getSchema().getFields().stream()
+        .map(
+            field -> {
+              if (field.getType().getRowSchema() != null) {
+                Row innerRow = getValue(field.getName());
+                if (innerRow != null) {
+                  return innerRow.toCamelCase();
+                }
+              }
+              return (Object) getValue(field.getName());
+            })
+        .collect(toRow(getSchema().toCamelCase()));
+  }
 }

sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/SchemaRegistryTest.java

@@ -17,12 +17,14 @@
  */
 package org.apache.beam.sdk.schemas;
 
+import static org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider;
 import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.SIMPLE_BEAN_SCHEMA;
 import static org.apache.beam.sdk.schemas.utils.TestPOJOs.SIMPLE_POJO_SCHEMA;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 import com.google.auto.service.AutoService;
+import com.google.auto.value.AutoValue;
 import java.util.List;
 import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
 import org.apache.beam.sdk.schemas.utils.TestJavaBeans.SimpleBean;
@@ -248,4 +250,51 @@ public void testGetSchemaCoder() throws NoSuchSchemaException {
     thrown.expect(NoSuchSchemaException.class);
     registry.getSchemaCoder(Double.class);
   }
+
+  @Test
+  public void testGetSchemaProvider() throws NoSuchSchemaException {
+    SchemaRegistry registry = SchemaRegistry.createDefault();
+
+    SchemaProvider testDefaultSchemaProvider =
+        registry.getSchemaProvider(TestDefaultSchemaClass.class);
+    assertEquals(DefaultSchemaProvider.class, testDefaultSchemaProvider.getClass());
+    assertEquals(
+        TestDefaultSchemaProvider.class,
+        ((DefaultSchemaProvider) testDefaultSchemaProvider)
+            .getUnderlyingSchemaProvider(TestDefaultSchemaClass.class)
+            .getClass());
+
+    SchemaProvider autoValueSchemaProvider = registry.getSchemaProvider(TestAutoValue.class);
+    assertEquals(DefaultSchemaProvider.class, autoValueSchemaProvider.getClass());
+    assertEquals(
+        AutoValueSchema.class,
+        ((DefaultSchemaProvider) autoValueSchemaProvider)
+            .getUnderlyingSchemaProvider(TestAutoValue.class)
+            .getClass());
+
+    SchemaProvider simpleBeanSchemaProvider = registry.getSchemaProvider(SimpleBean.class);
+    assertEquals(DefaultSchemaProvider.class, simpleBeanSchemaProvider.getClass());
+    assertEquals(
+        JavaBeanSchema.class,
+        ((DefaultSchemaProvider) simpleBeanSchemaProvider)
+            .getUnderlyingSchemaProvider(SimpleBean.class)
+            .getClass());
+  }
+
+  @DefaultSchema(AutoValueSchema.class)
+  @AutoValue
+  public abstract static class TestAutoValue {
+    public static Builder builder() {
+      return new AutoValue_SchemaRegistryTest_TestAutoValue.Builder();
+    }
+
+    public abstract String getStr();
+
+    @AutoValue.Builder
+    public abstract static class Builder {
+      public abstract Builder setStr(String str);
+
+      public abstract TestAutoValue build();
+    }
+  }
 }