ahmedabu98 commented on code in PR #39599: URL: https://github.com/apache/beam/pull/39599#discussion_r3708669874
########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCdcReadSchemaTransformProvider.java: ########## @@ -0,0 +1,174 @@ +/* + * 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.") + @Nullable + abstract Long getStartVersion(); + + @SchemaFieldDescription("Start timestamp of the Delta Lake table to read changes from.") + @Nullable + abstract String getStartTimestamp(); Review Comment: Mention ISO timestamp format is expected? same with end timestamp below ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCdcReadSchemaTransformProvider.java: ########## @@ -0,0 +1,174 @@ +/* + * 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.") + @Nullable + abstract Long getStartVersion(); Review Comment: Mention it's required to provide either this or start timestamp ########## website/www/site/content/en/documentation/io/managed-io.md: ########## @@ -70,6 +70,21 @@ and Beam SQL is invoked via the Managed API under the hood. Unavailable </td> </tr> + <tr> Review Comment: This markdown is automatically generated at release time right? ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCdcReadSchemaTransformProvider.java: ########## @@ -0,0 +1,174 @@ +/* + * 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.") + @Nullable + abstract Long getStartVersion(); + + @SchemaFieldDescription("Start timestamp of the Delta Lake table to read changes from.") + @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.") + @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.") + @Nullable + abstract List<String> getIncludeMetadataColumns(); Review Comment: Mention the full list of possible column names ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java: ########## @@ -259,6 +285,17 @@ public ReadChanges withEndTimestamp(String endTimestamp) { 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("Unsupported metadata column: " + col); Review Comment: Mention what the supported columns are ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCdcReadSchemaTransformProvider.java: ########## @@ -0,0 +1,174 @@ +/* + * 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.") + @Nullable + abstract Long getStartVersion(); + + @SchemaFieldDescription("Start timestamp of the Delta Lake table to read changes from.") + @Nullable + abstract String getStartTimestamp(); Review Comment: Also maybe this is a late suggestion (since DeltaIO is already implemented), but I wonder if a micros (long) would be better here ########## sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOIT.java: ########## @@ -261,4 +288,329 @@ public void testReadDeltaLakeTable() { PAssert.that(output).containsInAnyOrder(TEST_ROWS); readPipeline.run().waitUntilFinish(); } + + @Test + public void testReadChangesDeltaLake() throws Exception { + Map<String, String> hadoopConfig = new HashMap<>(); + hadoopConfig.put("fs.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem"); + hadoopConfig.put( + "fs.AbstractFileSystem.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFS"); + hadoopConfig.put("fs.gs.auth.type", "APPLICATION_DEFAULT"); + String project = + readPipeline + .getOptions() + .as(org.apache.beam.sdk.extensions.gcp.options.GcpOptions.class) + .getProject(); + if (project != null) { + hadoopConfig.put("fs.gs.project.id", project); + } + + org.apache.hadoop.conf.Configuration conf = new org.apache.hadoop.conf.Configuration(); + for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) { + conf.set(entry.getKey(), entry.getValue()); + } + Engine engine = DefaultEngine.create(conf); + + StructType deltaSchema = + new StructType().add("id", IntegerType.INTEGER).add("name", StringType.STRING); + + // 1. Write version 1 containing cdc actions for testing updates and deletes + Schema cdcWriteSchema = + Schema.builder() + .addField("id", Schema.FieldType.INT32) + .addField("name", Schema.FieldType.STRING) + .addField(DeltaIO.CHANGE_TYPE_COLUMN, Schema.FieldType.STRING) + .addField(DeltaIO.COMMIT_VERSION_COLUMN, Schema.FieldType.INT64) + .addField(DeltaIO.COMMIT_TIMESTAMP_COLUMN, Schema.FieldType.DATETIME) + .build(); + StructType cdcWriteDeltaSchema = + new StructType() + .add("id", IntegerType.INTEGER) + .add("name", StringType.STRING) + .add(DeltaIO.CHANGE_TYPE_COLUMN, StringType.STRING) + .add(DeltaIO.COMMIT_VERSION_COLUMN, LongType.LONG) + .add(DeltaIO.COMMIT_TIMESTAMP_COLUMN, TimestampType.TIMESTAMP); + + Row cdcRow1 = + Row.withSchema(cdcWriteSchema) + .addValues(0, "name_0", "delete", 1L, new Instant(123456789000L)) + .build(); + Row cdcRow2 = + Row.withSchema(cdcWriteSchema) + .addValues(1, "name_1", "update_preimage", 1L, new Instant(123456789000L)) + .build(); + Row cdcRow3 = + Row.withSchema(cdcWriteSchema) + .addValues(1, "name_1_updated", "update_postimage", 1L, new Instant(123456789000L)) + .build(); + + writeCdcCommit( + engine, + repoPath, + 1L, + deltaSchema, + null, + version0FilePath, + java.util.Arrays.asList(cdcRow1, cdcRow2, cdcRow3), + cdcWriteDeltaSchema); + + // 2. Read CDF data from table using Managed.read(Managed.DELTA_LAKE_CDC) + Map<String, Object> readConfig = new HashMap<>(); + readConfig.put("table", repoPath); + readConfig.put("start_version", 0L); + readConfig.put("hadoop_config", hadoopConfig); + readConfig.put( + "include_metadata_columns", + java.util.Arrays.asList( + DeltaIO.CHANGE_TYPE_COLUMN, + DeltaIO.COMMIT_VERSION_COLUMN, + DeltaIO.COMMIT_TIMESTAMP_COLUMN)); + + PCollection<Row> output = + readPipeline + .apply(Managed.read(Managed.DELTA_LAKE_CDC).withConfig(readConfig)) + .getSinglePCollection(); + + PCollection<String> formattedOutput = + output.apply("Format Row with Metadata", ParDo.of(new FormatITRowWithMetadata())); + + // Generate expected outputs for version 0 (inserts of id 0-99) + List<String> expectedOutputs = new ArrayList<>(); + for (int i = 0; i < 100; i++) { + expectedOutputs.add(String.format("%d:name_%d:insert:v0", i, i)); + } + // Expected outputs for version 1 + expectedOutputs.add("0:name_0:delete:v1"); + expectedOutputs.add("1:name_1:update_preimage:v1"); + expectedOutputs.add("1:name_1_updated:update_postimage:v1"); + + PAssert.that(formattedOutput).containsInAnyOrder(expectedOutputs); + + readPipeline.run().waitUntilFinish(); + } + + private static final class FormatITRowWithMetadata extends DoFn<Row, String> { + @ProcessElement + public void process(@Element Row row, OutputReceiver<String> out) { + out.output( + String.format( + "%d:%s:%s:v%d", + row.getInt32("id"), + row.getString("name"), + row.getString(DeltaIO.CHANGE_TYPE_COLUMN), + row.getInt64(DeltaIO.COMMIT_VERSION_COLUMN))); + } + } + + private static final StructType CDC_ACTION_SCHEMA = + new StructType() + .add("path", StringType.STRING, false) + .add("partitionValues", new MapType(StringType.STRING, StringType.STRING, false), false) + .add("size", LongType.LONG, false) + .add("dataChange", BooleanType.BOOLEAN, false); + + private static StructType getCustomSingleActionSchema() { + StructType originalSchema = io.delta.kernel.internal.actions.SingleAction.FULL_SCHEMA; + List<StructField> fields = new ArrayList<>(); + for (StructField field : originalSchema.fields()) { + if (field.getName().equals("cdc")) { + fields.add(new StructField("cdc", CDC_ACTION_SCHEMA, true)); + } else { + fields.add(field); + } + } + return new StructType(fields); + } + + private static io.delta.kernel.data.Row createSingleAction( + StructType customSingleActionSchema, String actionName, io.delta.kernel.data.Row actionRow) { + Map<Integer, Object> values = new HashMap<>(); + values.put(customSingleActionSchema.indexOf(actionName), actionRow); + return new GenericRow(customSingleActionSchema, values); + } + + private static io.delta.kernel.data.Row createRemoveAction( + StructType removeSchema, String path, long deletionTimestamp) { + Map<Integer, Object> values = new HashMap<>(); + values.put(removeSchema.indexOf("path"), path); + values.put(removeSchema.indexOf("deletionTimestamp"), deletionTimestamp); + values.put(removeSchema.indexOf("dataChange"), true); + values.put(removeSchema.indexOf("size"), 100L); + return new GenericRow(removeSchema, values); + } + + private static io.delta.kernel.data.Row createCdcAction( + StructType cdcSchema, String path, long size) { + Map<Integer, Object> values = new HashMap<>(); + values.put(cdcSchema.indexOf("path"), path); + values.put( + cdcSchema.indexOf("partitionValues"), + io.delta.kernel.internal.util.VectorUtils.stringStringMapValue(Collections.emptyMap())); + values.put(cdcSchema.indexOf("size"), size); + values.put(cdcSchema.indexOf("dataChange"), true); + return new GenericRow(cdcSchema, values); + } + + private static ColumnVector createColumnVector( + List<Row> rows, int fieldIndex, DataType dataType) { + return new ColumnVector() { + @Override + public DataType getDataType() { + return dataType; + } + + @Override + public int getSize() { + return rows.size(); + } + + @Override + public void close() {} + + @Override + public boolean isNullAt(int rowId) { + return rows.get(rowId).getValue(fieldIndex) == null; + } + + @Override + public boolean getBoolean(int rowId) { + return rows.get(rowId).getBoolean(fieldIndex); + } + + @Override + public int getInt(int rowId) { + return rows.get(rowId).getInt32(fieldIndex); + } + + @Override + public long getLong(int rowId) { + if (dataType instanceof TimestampType) { + Instant instant = rows.get(rowId).getDateTime(fieldIndex).toInstant(); + return instant.getMillis() * 1000L; + } + return rows.get(rowId).getInt64(fieldIndex); + } + + @Override + public String getString(int rowId) { + return rows.get(rowId).getString(fieldIndex); + } + }; + } + + /** + * Writes a Delta commit containing CDC actions (simulating updates/deletes). + * + * <p>Note on why this is manual: In a standard Spark or Flink writer, setting the table property + * {@code "delta.enableChangeDataFeed" = "true"} automatically instructs the engine to compute and + * write the change data files to {@code _change_data/} and append the {@code cdc} actions to the + * commit log whenever DML statements (like UPDATE/DELETE) are executed. + * + * <p>However, we are using the Delta Lake Kernel API which does not contain an SQL execution + * engine or a DML parser. Thus, it cannot automatically compute which rows were deleted or + * updated. To generate a realistic integration test dataset, we must manually construct these + * change records, write them into the GCS {@code _change_data/} directory using the low-level + * parquet handler, and manually register them as {@code cdc} actions in the committed + * transaction. + */ + private void writeCdcCommit( + Engine engine, + String tablePath, + long expectedVersion, + StructType deltaSchema, + @Nullable List<Row> addBeamRows, Review Comment: Looks like we never pass `addBeamRows` into here? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
