Managed Transform protos & translation; Iceberg SchemaTransforms & translation

.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 `lower_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/utils/YamlUtils.java

@@ -53,7 +53,7 @@ public class YamlUtils {
           .put(Schema.TypeName.BYTES, str -> BaseEncoding.base64().decode(str))
           .build();
 
-  public static Row toBeamRow(@Nullable String yamlString, Schema schema) {
+  public static @Nullable Row toBeamRow(@Nullable String yamlString, Schema schema) {
     return toBeamRow(yamlString, schema, false);
   }
 
@@ -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,44 @@ public static Row nullRow(Schema schema) {
         .addValues(Collections.nCopies(schema.getFieldCount(), null))
         .build();
   }
+
+  /** Returns an identical {@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) {
+                return ((Row) getValue(field.getName())).sorted();
+              }
+              return (Object) getValue(field.getName());
+            })
+        .collect(Row.toRow(sortedSchema));
+  }
+
+  /** Returns an equivalent {@link Row} with `lower_snake_case` field names. */
+  public Row toSnakeCase() {
+    return getSchema().getFields().stream()
+        .map(
+            field -> {
+              if (field.getType().getRowSchema() != null) {
+                return ((Row) getValue(field.getName())).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) {
+                return ((Row) getValue(field.getName())).toCamelCase();
+              }
+              return (Object) getValue(field.getName());
+            })
+        .collect(toRow(getSchema().toCamelCase()));
+  }
 }

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

@@ -194,6 +194,71 @@ public void testCollector() {
     assertEquals(FieldType.STRING, schema.getField(1).getType());
   }
 
+  @Test
+  public void testToSnakeCase() {
+    Schema innerSchema =
+        Schema.builder()
+            .addStringField("myFirstNestedStringField")
+            .addStringField("mySecondNestedStringField")
+            .build();
+    Schema schema =
+        Schema.builder()
+            .addStringField("myFirstStringField")
+            .addStringField("mySecondStringField")
+            .addRowField("myRowField", innerSchema)
+            .build();
+
+    Schema expectedInnerSnakeCaseSchema =
+        Schema.builder()
+            .addStringField("my_first_nested_string_field")
+            .addStringField("my_second_nested_string_field")
+            .build();
+    Schema expectedSnakeCaseSchema =
+        Schema.builder()
+            .addStringField("my_first_string_field")
+            .addStringField("my_second_string_field")
+            .addRowField("my_row_field", expectedInnerSnakeCaseSchema)
+            .build();
+
+    assertEquals(
+        expectedInnerSnakeCaseSchema,
+        schema.toSnakeCase().getField("my_row_field").getType().getRowSchema());
+    assertEquals(expectedSnakeCaseSchema, schema.toSnakeCase());
+  }
+
+  @Test
+  public void testToCamelCase() {
+    Schema innerSchema =
+        Schema.builder()
+            .addStringField("my_first_nested_string_field")
+            .addStringField("my_second_nested_string_field")
+            .build();
+    Schema schema =
+        Schema.builder()
+            .addStringField("my_first_string_field")
+            .addStringField("my_second_string_field")
+            .addRowField("my_row_field", innerSchema)
+            .build();
+
+    Schema expectedInnerCamelCaseSchema =
+        Schema.builder()
+            .addStringField("myFirstNestedStringField")
+            .addStringField("mySecondNestedStringField")
+            .build();
+    Schema expectedCamelCaseSchema =
+        Schema.builder()
+            .addStringField("myFirstStringField")
+            .addStringField("mySecondStringField")
+            .addRowField("myRowField", expectedInnerCamelCaseSchema)
+            .build();
+
+    assertTrue(schema.toCamelCase().hasField("myRowField"));
+    assertEquals(
+        expectedInnerCamelCaseSchema,
+        schema.toCamelCase().getField("myRowField").getType().getRowSchema());
+    assertEquals(expectedCamelCaseSchema, schema.toCamelCase());
+  }
+
   @Test
   public void testSorted() {
     Options testOptions =
@@ -232,6 +297,41 @@ public void testSorted() {
                 sortedSchema.getEncodingPositions()));
   }
 
+  @Test
+  public void testNestedSorted() {
+    Schema unsortedNestedSchema =
+        Schema.builder().addStringField("bb").addInt32Field("aa").addStringField("cc").build();
+    Schema unsortedSchema =
+        Schema.builder()
+            .addStringField("d")
+            .addInt32Field("c")
+            .addRowField("e", unsortedNestedSchema)
+            .addStringField("b")
+            .addByteField("a")
+            .build();
+
+    Schema sortedSchema = unsortedSchema.sorted();
+
+    Schema expectedInnerSortedSchema =
+        Schema.builder().addInt32Field("aa").addStringField("bb").addStringField("cc").build();
+    Schema expectedSortedSchema =
+        Schema.builder()
+            .addByteField("a")
+            .addStringField("b")
+            .addInt32Field("c")
+            .addStringField("d")
+            .addRowField("e", expectedInnerSortedSchema)
+            .build();
+
+    assertTrue(unsortedSchema.equivalent(sortedSchema));
+    assertEquals(expectedSortedSchema.getFields(), sortedSchema.getFields());
+    assertEquals(expectedSortedSchema.getEncodingPositions(), sortedSchema.getEncodingPositions());
+    assertEquals(expectedInnerSortedSchema, sortedSchema.getField("e").getType().getRowSchema());
+    assertEquals(
+        expectedInnerSortedSchema.getEncodingPositions(),
+        sortedSchema.getField("e").getType().getRowSchema().getEncodingPositions());
+  }
+
   @Test
   public void testSortedMethodIncludesAllSchemaFields() {
     // This test is most likely to break when new Schema object attributes are added. It is designed

sdks/java/core/src/test/java/org/apache/beam/sdk/util/YamlUtilsTest.java

@@ -21,11 +21,13 @@
 
 import java.math.BigDecimal;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.utils.YamlUtils;
 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.io.BaseEncoding;
 import org.junit.Rule;
 import org.junit.Test;
@@ -225,4 +227,35 @@ public void testNestedArray() {
 
     assertEquals(expectedRow, YamlUtils.toBeamRow(yamlString, schema));
   }
+
+  private static final Schema FLAT_SCHEMA_CAMEL_CASE =
+      Schema.builder()
+          .addFields(
+              FLAT_SCHEMA.getFields().stream()
+                  .map(
+                      field ->
+                          field.withName(
+                              CaseFormat.LOWER_UNDERSCORE.to(
+                                  CaseFormat.LOWER_CAMEL, field.getName())))
+                  .collect(Collectors.toList()))
+          .build();
+
+  private static final Map<String, Object> FLAT_MAP =
+      FLAT_SCHEMA.getFields().stream()
+          .collect(
+              Collectors.toMap(
+                  Schema.Field::getName,
+                  field -> Preconditions.checkArgumentNotNull(FLAT_ROW.getValue(field.getName()))));
+
+  @Test
+  public void testSnakeCaseMapToCamelCaseRow() {
+    Row expectedRow =
+        FLAT_SCHEMA.getFields().stream()
+            .map(field -> Preconditions.checkStateNotNull(FLAT_ROW.getValue(field.getName())))
+            .collect(Row.toRow(FLAT_SCHEMA_CAMEL_CASE));
+
+    Row convertedRow = YamlUtils.toBeamRow(FLAT_MAP, FLAT_SCHEMA_CAMEL_CASE, true);
+
+    assertEquals(expectedRow, convertedRow);
+  }
 }