Adds the Delta Lake CDC read transforms to the Managed I/O API

.github/trigger_files/beam_PostCommit_Java_Delta_IO_Dataflow.json

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

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

@@ -107,6 +107,8 @@ message ManagedTransforms {
       "beam:schematransform:org.apache.beam:sql_server_write:v1"];
     DELTA_LAKE_READ = 13 [(org.apache.beam.model.pipeline.v1.beam_urn) =
       "beam:schematransform:org.apache.beam:delta_lake_read:v1"];
+    DELTA_LAKE_CDC_READ = 14 [(org.apache.beam.model.pipeline.v1.beam_urn) =
+      "beam:schematransform:org.apache.beam:delta_lake_cdc_read:v1"];
   }
 }
 

sdks/java/io/delta/build.gradle

@@ -101,7 +101,7 @@ task dataflowIntegrationTest(type: Test) {
     def dockerJavaImageName = project.project(':runners:google-cloud-dataflow-java').ext.dockerJavaImageName
 
     def args = [
-            "--runner=DataflowRunner",
+            "--runner=TestDataflowRunner",
             "--region=us-central1",
             "--project=${gcpProject}",
             "--tempLocation=${gcpTempLocation}",

sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java

@@ -61,11 +61,14 @@
 @DoFn.BoundedPerElement
 class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> {
   @Nullable Map<String, String> hadoopConfig;
+  private final @Nullable List<String> metadataColumns;
   private transient @Nullable Engine engine;
   private transient @Nullable Configuration conf;
 
-  public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) {
+  public DeltaCDCSourceDoFn(
+      @Nullable Map<String, String> hadoopConfig, @Nullable List<String> metadataColumns) {
     this.hadoopConfig = hadoopConfig;
+    this.metadataColumns = metadataColumns;
   }
 
   private synchronized Configuration getConfiguration() {
@@ -117,7 +120,8 @@ public void processElement(
 
     SerializableRow originalScanStateRow = task.getScanStateRow();
     StructType logicalTableSchema = ScanStateRow.getLogicalSchema(originalScanStateRow);
-    Schema publicBeamSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
+    Schema baseSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
+    Schema publicBeamSchema = DeltaIO.buildPublicBeamSchema(baseSchema, metadataColumns);
     StructType physicalTableSchema = ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow);
 
     StructType scanStateSchema = originalScanStateRow.getSchema();

sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCdcReadSchemaTransformProvider.java

@@ -0,0 +1,178 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import static org.apache.beam.sdk.io.delta.DeltaCdcReadSchemaTransformProvider.Configuration;
+import static org.apache.beam.sdk.util.construction.BeamUrns.getUrn;
+
+import com.google.auto.service.AutoService;
+import com.google.auto.value.AutoValue;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.model.pipeline.v1.ExternalTransforms;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.Row;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * SchemaTransform implementation for {@link DeltaIO#readChanges}. Reads change records from Delta
+ * Lake and outputs a {@link org.apache.beam.sdk.values.PCollection} of Beam {@link
+ * org.apache.beam.sdk.values.Row}s.
+ */
+@AutoService(SchemaTransformProvider.class)
+public class DeltaCdcReadSchemaTransformProvider
+    extends TypedSchemaTransformProvider<Configuration> {
+  static final String OUTPUT_TAG = "output";
+
+  @Override
+  protected SchemaTransform from(Configuration configuration) {
+    return new DeltaCdcReadSchemaTransform(configuration);
+  }
+
+  @Override
+  public List<String> outputCollectionNames() {
+    return Collections.singletonList(OUTPUT_TAG);
+  }
+
+  @Override
+  public String identifier() {
+    return getUrn(ExternalTransforms.ManagedTransforms.Urns.DELTA_LAKE_CDC_READ);
+  }
+
+  static class DeltaCdcReadSchemaTransform extends SchemaTransform {
+    private final Configuration configuration;
+
+    DeltaCdcReadSchemaTransform(Configuration configuration) {
+      this.configuration =
+          java.util.Objects.requireNonNull(configuration, "configuration cannot be null");
+    }
+
+    Row getConfigurationRow() {
+      try {
+        return SchemaRegistry.createDefault()
+            .getToRowFunction(Configuration.class)
+            .apply(configuration)
+            .sorted()
+            .toSnakeCase();
+      } catch (NoSuchSchemaException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    @Override
+    public PCollectionRowTuple expand(PCollectionRowTuple input) {
+      DeltaIO.ReadChanges read = DeltaIO.readChanges().from(configuration.getTable());
+      Long startVersion = configuration.getStartVersion();
+      if (startVersion != null) {
+        read = read.withStartVersion(startVersion);
+      }
+      String startTimestamp = configuration.getStartTimestamp();
+      if (startTimestamp != null) {
+        read = read.withStartTimestamp(startTimestamp);
+      }
+      Long endVersion = configuration.getEndVersion();
+      if (endVersion != null) {
+        read = read.withEndVersion(endVersion);
+      }
+      String endTimestamp = configuration.getEndTimestamp();
+      if (endTimestamp != null) {
+        read = read.withEndTimestamp(endTimestamp);
+      }
+      Map<String, String> hadoopConfig = configuration.getHadoopConfig();
+      if (hadoopConfig != null) {
+        read = read.withConfig(hadoopConfig);
+      }
+      List<String> includeMetadataColumns = configuration.getIncludeMetadataColumns();
+      if (includeMetadataColumns != null && !includeMetadataColumns.isEmpty()) {
+        read = read.withMetadataColumns(includeMetadataColumns.toArray(new String[0]));
+      }
+
+      PCollection<Row> output = input.getPipeline().apply(read);
+
+      return PCollectionRowTuple.of(OUTPUT_TAG, output);
+    }
+  }
+
+  @DefaultSchema(AutoValueSchema.class)
+  @AutoValue
+  public abstract static class Configuration {
+    static Builder builder() {
+      return new AutoValue_DeltaCdcReadSchemaTransformProvider_Configuration.Builder();
+    }
+
+    @SchemaFieldDescription("Identifier of the Delta Lake table.")
+    abstract String getTable();
+
+    @SchemaFieldDescription(
+        "Start version of the Delta Lake table to read changes from. Either this or the start timestamp has to be provided.")
+    @Nullable
+    abstract Long getStartVersion();
+
+    @SchemaFieldDescription(
+        "Start timestamp of the Delta Lake table to read changes from. Should be specified in the ISO 8601 standard. Either this or the start version has to be provided.")
+    @Nullable
+    abstract String getStartTimestamp();
+
+    @SchemaFieldDescription("End version of the Delta Lake table to read changes up to.")
+    @Nullable
+    abstract Long getEndVersion();
+
+    @SchemaFieldDescription(
+        "End timestamp of the Delta Lake table to read changes up to. Should be specified in the ISO 8601 standard.")
+    @Nullable
+    abstract String getEndTimestamp();
+
+    @SchemaFieldDescription("Properties passed to the Hadoop Configuration.")
+    @Nullable
+    abstract Map<String, String> getHadoopConfig();
+
+    @SchemaFieldDescription(
+        "Metadata columns to include in the output rows. Supported columns are: _change_type, _commit_version, and _commit_timestamp.")
+    @Nullable
+    abstract List<String> getIncludeMetadataColumns();
+
+    @AutoValue.Builder
+    abstract static class Builder {
+      abstract Builder setTable(String table);
+
+      abstract Builder setStartVersion(Long startVersion);
+
+      abstract Builder setStartTimestamp(String startTimestamp);
+
+      abstract Builder setEndVersion(Long endVersion);
+
+      abstract Builder setEndTimestamp(String endTimestamp);
+
+      abstract Builder setHadoopConfig(Map<String, String> hadoopConfig);
+
+      abstract Builder setIncludeMetadataColumns(List<String> includeMetadataColumns);
+
+      abstract Configuration build();
+    }
+  }
+}

sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java

@@ -37,6 +37,8 @@
 import io.delta.kernel.types.StructField;
 import io.delta.kernel.types.StructType;
 import io.delta.kernel.types.TimestampType;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.schemas.Schema;
@@ -70,7 +72,7 @@
   /**
    * Reads rows from a Delta Lake table.
    *
    * <p>Normally, it is recommended to use {@link org.apache.beam.sdk.managed.Managed#read(String)}
    * with {@code Managed.DELTA_LAKE} instead of directly using this transform.
    */
   public static ReadRows readRows() {
@@ -206,6 +208,26 @@
     }
   }
 
+  static Schema buildPublicBeamSchema(Schema baseSchema, @Nullable List<String> metadataColumns) {
+    if (metadataColumns == null || metadataColumns.isEmpty()) {
+      return baseSchema;
+    }
+    Schema.Builder builder = Schema.builder();
+    for (Schema.Field field : baseSchema.getFields()) {
+      builder.addField(field);
+    }
+    for (String col : metadataColumns) {
+      if (col.equals(CHANGE_TYPE_COLUMN)) {
+        builder.addField(CHANGE_TYPE_COLUMN, Schema.FieldType.STRING);
+      } else if (col.equals(COMMIT_VERSION_COLUMN)) {
+        builder.addField(COMMIT_VERSION_COLUMN, Schema.FieldType.INT64);
+      } else if (col.equals(COMMIT_TIMESTAMP_COLUMN)) {
+        builder.addField(COMMIT_TIMESTAMP_COLUMN, Schema.FieldType.DATETIME);
+      }
+    }
+    return builder.build();
+  }
+
   @AutoValue
   public abstract static class ReadChanges extends PTransform<PBegin, PCollection<Row>> {
     public abstract @Nullable String getTablePath();
@@ -218,6 +240,8 @@
 
     public abstract @Nullable String getEndTimestamp();
 
+    public abstract @Nullable List<String> getMetadataColumns();
+
     public abstract @Nullable Map<String, String> getHadoopConfig();
 
     abstract Builder toBuilder();
@@ -234,6 +258,8 @@
 
       abstract Builder setEndTimestamp(@Nullable String endTimestamp);
 
+      abstract Builder setMetadataColumns(@Nullable List<String> metadataColumns);
+
       abstract Builder setHadoopConfig(@Nullable Map<String, String> hadoopConfig);
 
       abstract ReadChanges build();
@@ -259,6 +285,20 @@
       return toBuilder().setEndTimestamp(endTimestamp).build();
     }
 
+    public ReadChanges withMetadataColumns(String... metadataColumns) {
+      for (String col : metadataColumns) {
+        if (!col.equals(CHANGE_TYPE_COLUMN)
+            && !col.equals(COMMIT_VERSION_COLUMN)
+            && !col.equals(COMMIT_TIMESTAMP_COLUMN)) {
+          throw new IllegalArgumentException(
+              String.format(
+                  "Unsupported metadata column %s. Supported columns are: %s, %s, and %s.",
+                  col, CHANGE_TYPE_COLUMN, COMMIT_VERSION_COLUMN, COMMIT_TIMESTAMP_COLUMN));
+        }
+      }
+      return toBuilder().setMetadataColumns(Arrays.asList(metadataColumns)).build();
+    }
+
     public ReadChanges withConfig(Map<String, String> config) {
       return toBuilder().setHadoopConfig(config).build();
     }
@@ -310,7 +350,8 @@
       if (deltaSchema == null) {
         throw new IllegalStateException("Table schema is null.");
       }
-      Schema beamSchema = ReadRows.convertToBeamSchema(deltaSchema);
+      Schema baseSchema = ReadRows.convertToBeamSchema(deltaSchema);
+      Schema publicBeamSchema = buildPublicBeamSchema(baseSchema, getMetadataColumns());
 
       return input
           .apply("Create Path", Create.of(path))
@@ -323,8 +364,9 @@
                       getStartTimestamp(),
                       getEndVersion(),
                       getEndTimestamp())))
-          .apply("Read CDF Data", ParDo.of(new DeltaCDCSourceDoFn(hadoopConfig)))
-          .setRowSchema(beamSchema);
+          .apply(
+              "Read CDF Data", ParDo.of(new DeltaCDCSourceDoFn(hadoopConfig, getMetadataColumns())))
+          .setRowSchema(publicBeamSchema);
     }
   }
 }