Abacn commented on code in PR #39426: URL: https://github.com/apache/beam/pull/39426#discussion_r3660652412
########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCReadTask.java: ########## @@ -0,0 +1,140 @@ +/* + * 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 java.io.Serializable; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A serializable task containing the necessary metadata to read a CDF (Change Data Feed) file. This + * can be either a CDC parquet file or a regular data file (representing inserts) from a commit. + */ +public class DeltaCDCReadTask implements Serializable { + private static final long serialVersionUID = 1L; + + private final String path; + private final long size; + private final Map<String, String> partitionValues; + private final long version; + private final long timestamp; + private final boolean isCDC; + private final List<Long> rowGroupSizes; + private final SerializableRow scanStateRow; + + public DeltaCDCReadTask( + String path, + long size, + Map<String, String> partitionValues, + long version, + long timestamp, + boolean isCDC, + List<Long> rowGroupSizes, + SerializableRow scanStateRow) { + this.path = path; + this.size = size; + this.partitionValues = partitionValues; + this.version = version; + this.timestamp = timestamp; + this.isCDC = isCDC; + this.rowGroupSizes = rowGroupSizes; + this.scanStateRow = scanStateRow; + } + + public String getPath() { + return path; + } + + public long getSize() { + return size; + } + + public Map<String, String> getPartitionValues() { + return partitionValues; + } + + public long getVersion() { + return version; + } + + public long getTimestamp() { + return timestamp; + } + + public boolean isCDC() { + return isCDC; + } + + public List<Long> getRowGroupSizes() { + return rowGroupSizes; + } + + public SerializableRow getScanStateRow() { + return scanStateRow; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof DeltaCDCReadTask)) { + return false; + } + DeltaCDCReadTask that = (DeltaCDCReadTask) o; + return size == that.size + && version == that.version + && timestamp == that.timestamp + && isCDC == that.isCDC + && Objects.equals(path, that.path) + && Objects.equals(partitionValues, that.partitionValues) + && Objects.equals(rowGroupSizes, that.rowGroupSizes) + && Objects.equals(scanStateRow, that.scanStateRow); + } + + @Override + public int hashCode() { + return Objects.hash( + path, size, partitionValues, version, timestamp, isCDC, rowGroupSizes, scanStateRow); + } + + @Override + public String toString() { + return "DeltaCDCReadTask{" Review Comment: nit: prefer String.format ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java: ########## @@ -0,0 +1,353 @@ +/* + * 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 io.delta.kernel.internal.DeltaErrors.wrapEngineException; + +import io.delta.kernel.Scan; +import io.delta.kernel.data.ColumnVector; +import io.delta.kernel.data.ColumnarBatch; +import io.delta.kernel.data.FilteredColumnarBatch; +import io.delta.kernel.data.MapValue; +import io.delta.kernel.defaults.engine.DefaultEngine; +import io.delta.kernel.engine.Engine; +import io.delta.kernel.engine.FileReadResult; +import io.delta.kernel.expressions.ExpressionEvaluator; +import io.delta.kernel.expressions.Literal; +import io.delta.kernel.internal.InternalScanFileUtils; +import io.delta.kernel.internal.data.GenericRow; +import io.delta.kernel.internal.data.ScanStateRow; +import io.delta.kernel.internal.util.Utils; +import io.delta.kernel.internal.util.VectorUtils; +import io.delta.kernel.types.LongType; +import io.delta.kernel.types.StringType; +import io.delta.kernel.types.StructField; +import io.delta.kernel.types.StructType; +import io.delta.kernel.types.TimestampType; +import io.delta.kernel.utils.CloseableIterator; +import io.delta.kernel.utils.FileStatus; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import org.apache.beam.sdk.io.range.OffsetRange; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.hadoop.conf.Configuration; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and reads Change Data Feed + * files, converting rows to Beam Rows. + */ [email protected] +class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> { + @Nullable Map<String, String> hadoopConfig; + private transient @Nullable Engine engine; + private transient @Nullable Configuration conf; + + public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) { + this.hadoopConfig = hadoopConfig; + } + + private synchronized Configuration getConfiguration() { + Configuration localConf = conf; + if (localConf == null) { + localConf = new Configuration(); + if (hadoopConfig != null) { + for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) { + localConf.set(entry.getKey(), entry.getValue()); + } + } + conf = localConf; + } + return localConf; + } + + private List<Long> getRowGroupSizes(DeltaCDCReadTask task) { + return task.getRowGroupSizes(); + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) { + List<Long> rowGroupSizes = getRowGroupSizes(task); + return new OffsetRange(0L, rowGroupSizes.size()); + } + + @NewTracker + public DeltaReadTaskTracker newTracker( + @Restriction OffsetRange restriction, @Element DeltaCDCReadTask task) { + return new DeltaReadTaskTracker(restriction, getRowGroupSizes(task)); + } + + @Setup + public void setUp() { + engine = DefaultEngine.create(getConfiguration()); + } + + @ProcessElement + public ProcessContinuation processElement( + @Element DeltaCDCReadTask task, + RestrictionTracker<OffsetRange, Long> tracker, + OutputReceiver<Row> out) + throws Exception { + + Engine currentEngine = engine; + if (currentEngine == null) { + throw new IllegalArgumentException("Expected the engine to not be null"); + } + + SerializableRow originalScanStateRow = task.getScanStateRow(); + StructType logicalTableSchema = ScanStateRow.getLogicalSchema(originalScanStateRow); + Schema publicBeamSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema); + StructType physicalTableSchema = ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow); + + StructType scanStateSchema = originalScanStateRow.getSchema(); + + // 1. Build modified scanState and scanFile rows depending on whether we read a CDC file or ADD + // file. + io.delta.kernel.data.Row scanStateRow; + StructType readPhysicalSchema; + StructType readLogicalSchema; + Schema beamSchema; + + if (task.isCDC()) { + readLogicalSchema = appendCDFColumns(logicalTableSchema); + readPhysicalSchema = appendCDFColumns(physicalTableSchema); + beamSchema = DeltaIO.ReadRows.convertToBeamSchema(readLogicalSchema); + + HashMap<Integer, Object> valueMap = new HashMap<>(); + + // Tracking row level lineage is not needed. + Map<String, String> config = + new HashMap<>(ScanStateRow.getConfiguration(originalScanStateRow)); + config.put("delta.enableRowTracking", "false"); + + valueMap.put( + scanStateSchema.indexOf("configuration"), VectorUtils.stringStringMapValue(config)); + valueMap.put(scanStateSchema.indexOf("logicalSchemaJson"), readLogicalSchema.toJson()); + valueMap.put(scanStateSchema.indexOf("physicalSchemaJson"), readPhysicalSchema.toJson()); + valueMap.put( + scanStateSchema.indexOf("partitionColumns"), + originalScanStateRow.getArray(scanStateSchema.indexOf("partitionColumns"))); + valueMap.put( + scanStateSchema.indexOf("minReaderVersion"), + originalScanStateRow.getInt(scanStateSchema.indexOf("minReaderVersion"))); + valueMap.put( + scanStateSchema.indexOf("minWriterVersion"), + originalScanStateRow.getInt(scanStateSchema.indexOf("minWriterVersion"))); + valueMap.put( + scanStateSchema.indexOf("tablePath"), + originalScanStateRow.getString(scanStateSchema.indexOf("tablePath"))); + + scanStateRow = new ScanStateRow(valueMap); + } else { + // For ADD files, we read the table schema and append the CDF columns manually afterwards. + // readLogicalSchema = logicalTableSchema; + readPhysicalSchema = physicalTableSchema; + beamSchema = DeltaIO.ReadRows.convertToBeamSchema(appendCDFColumns(logicalTableSchema)); + scanStateRow = originalScanStateRow; + } + + io.delta.kernel.data.Row scanFileRow = + generateScanFileRow(task.getPath(), task.getPartitionValues()); + FileStatus fileStatus = FileStatus.of(task.getPath(), task.getSize(), task.getTimestamp()); + + BeamParquetHandler parquetHandler = + new BeamParquetHandler(getConfiguration(), currentEngine.getParquetHandler(), tracker); + BeamEngine beamEngine = new BeamEngine(currentEngine, parquetHandler); + + long currentStartRgIndex = 0L; + + try (CloseableIterator<FileReadResult> fileReadResults = + parquetHandler.readParquetFiles( + Utils.singletonCloseableIterator(fileStatus), + readPhysicalSchema, + Optional.empty(), + currentStartRgIndex)) { + + CloseableIterator<ColumnarBatch> physicalData = + new CloseableIterator<ColumnarBatch>() { + @Override + public void close() throws java.io.IOException {} + + @Override + public boolean hasNext() { + return fileReadResults.hasNext(); + } + + @Override + public ColumnarBatch next() { + return fileReadResults.next().getData(); + } + }; + + try (CloseableIterator<FilteredColumnarBatch> logicalBatches = + Scan.transformPhysicalData(beamEngine, scanStateRow, scanFileRow, physicalData)) { + + while (logicalBatches.hasNext()) { + FilteredColumnarBatch batch = logicalBatches.next(); + ColumnarBatch logicalBatch = batch.getData(); + + if (!task.isCDC()) { + // For ADD files, we need to append the constant CDF columns: + // _change_type = "insert", _commit_version = task.version, _commit_timestamp = + // task.timestamp + logicalBatch = + appendConstantCDFColumns( + currentEngine, logicalBatch, task.getVersion(), task.getTimestamp()); + } + + try (CloseableIterator<io.delta.kernel.data.Row> logicalRows = logicalBatch.getRows()) { + while (logicalRows.hasNext()) { + io.delta.kernel.data.Row deltaRow = logicalRows.next(); + Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema); + String changeType = beamRow.getString("_change_type"); + if (changeType == null) { + throw new IllegalStateException("Field _change_type must not be null."); + } + ValueKind kind = getValueKind(changeType); + Row publicRow = projectRow(beamRow, publicBeamSchema); + out.builder(publicRow).setValueKind(kind).output(); + } + } + } + } + } + + return ProcessContinuation.stop(); + } + + private static Row projectRow(Row row, Schema targetSchema) { + Row.Builder builder = Row.withSchema(targetSchema); + for (Schema.Field field : targetSchema.getFields()) { + builder.addValue(row.getValue(field.getName())); + } + return builder.build(); + } + + private static ValueKind getValueKind(String changeType) { + // Maps Delta CDC change types to Beam's ValueKind enum. + // https://docs.delta.io/delta-change-data-feed/#what-is-the-schema-for-the-change-data-feed + switch (changeType) { + case "insert": + return ValueKind.INSERT; + case "delete": + return ValueKind.DELETE; + case "update_preimage": + return ValueKind.UPDATE_BEFORE; + case "update_postimage": + return ValueKind.UPDATE_AFTER; + default: + throw new IllegalArgumentException("Unsupported change type: " + changeType); + } + } + + private static StructType appendCDFColumns(StructType schema) { + return schema + .add("_change_type", StringType.STRING, false) Review Comment: nit: consider delcare these field name as constants. There are used in multiple places ########## sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java: ########## @@ -810,4 +815,263 @@ public boolean tryClaim(Long i) { } } } + + @Test + public void testReadChanges() throws Exception { + File tableDir = tempFolder.newFolder("delta-table-changes"); + File logDir = new File(tableDir, "_delta_log"); + logDir.mkdirs(); + + // 1. Write parquet files for Version 0 (insert-only commit) + Schema tableSchema = Schema.builder().addField("name", Schema.FieldType.STRING).build(); + Row tableRow1 = Row.withSchema(tableSchema).addValues("row-1").build(); + Row tableRow2 = Row.withSchema(tableSchema).addValues("row-2").build(); + + File partFile = new File(tableDir, "part-00000.parquet"); + byte[] partBytes = + writeParquetFile(partFile, tableSchema, java.util.Arrays.asList(tableRow1, tableRow2)); + + writeCommit( + logDir, 0L, 100000000000L, "part-00000.parquet", partBytes.length, null, null, 0L, true); + + // 2. Write cdc parquet file for Version 1 (commit with cdc actions) + Schema cdcWriteSchema = + Schema.builder() + .addField("name", Schema.FieldType.STRING) + .addField("_change_type", Schema.FieldType.STRING) + .addField("_commit_version", Schema.FieldType.INT64) + .addField("_commit_timestamp", Schema.FieldType.DATETIME) + .build(); + + Row cdcRow1 = + Row.withSchema(cdcWriteSchema) + .addValues("row-1", "update_preimage", 1L, new Instant(123456789000L)) + .build(); + Row cdcRow2 = + Row.withSchema(cdcWriteSchema) + .addValues("row-1-updated", "update_postimage", 1L, new Instant(123456789000L)) + .build(); + Row cdcRow3 = + Row.withSchema(cdcWriteSchema) + .addValues("row-2", "delete", 1L, new Instant(123456789000L)) + .build(); + + File changeFile = new File(tableDir, "change-00000.parquet"); + byte[] changeBytes = + writeParquetFile( + changeFile, cdcWriteSchema, java.util.Arrays.asList(cdcRow1, cdcRow2, cdcRow3)); + + writeCommit( + logDir, + 1L, + 200000000000L, + null, + 0L, + null, + "change-00000.parquet", + changeBytes.length, + false); + + // 3. Read CDF data from table using ReadChanges + PCollection<Row> output = + readPipeline.apply( + DeltaIO.readChanges().from(tableDir.getAbsolutePath()).withStartVersion(0L)); + + PCollection<String> formattedOutput = + output.apply("Format ValueKind and Row", ParDo.of(new FormatValueKindAndRow())); + + PAssert.that(formattedOutput) + .containsInAnyOrder( + "INSERT:row-1", + "INSERT:row-2", + "UPDATE_BEFORE:row-1", + "UPDATE_AFTER:row-1-updated", + "DELETE:row-2"); + + readPipeline.run().waitUntilFinish(); + } + + @Test + public void testReadChangesRanges() throws Exception { + File tableDir = tempFolder.newFolder("delta-table-changes-ranges"); + File logDir = new File(tableDir, "_delta_log"); + logDir.mkdirs(); + + Schema tableSchema = Schema.builder().addField("name", Schema.FieldType.STRING).build(); + + // 1. Write parquet files for Version 0 (insert-only commit) + Row tableRow1 = Row.withSchema(tableSchema).addValues("row-1").build(); + Row tableRow2 = Row.withSchema(tableSchema).addValues("row-2").build(); + File partFile0 = new File(tableDir, "part-00000.parquet"); + byte[] partBytes0 = + writeParquetFile(partFile0, tableSchema, java.util.Arrays.asList(tableRow1, tableRow2)); + writeCommit( + logDir, 0L, 100000000000L, "part-00000.parquet", partBytes0.length, null, null, 0L, true); + + // 2. Write parquet files for Version 1 (commit with updates and deletes) + Schema cdcWriteSchema = + Schema.builder() + .addField("name", Schema.FieldType.STRING) + .addField("_change_type", Schema.FieldType.STRING) + .addField("_commit_version", Schema.FieldType.INT64) + .addField("_commit_timestamp", Schema.FieldType.DATETIME) + .build(); + Row cdcRow1 = + Row.withSchema(cdcWriteSchema) + .addValues("row-1", "update_preimage", 1L, new Instant(200000000000L)) + .build(); + Row cdcRow2 = + Row.withSchema(cdcWriteSchema) + .addValues("row-1-updated", "update_postimage", 1L, new Instant(200000000000L)) + .build(); + Row cdcRow3 = + Row.withSchema(cdcWriteSchema) + .addValues("row-2", "delete", 1L, new Instant(200000000000L)) + .build(); + File changeFile0 = new File(tableDir, "change-00000.parquet"); + byte[] changeBytes0 = + writeParquetFile( + changeFile0, cdcWriteSchema, java.util.Arrays.asList(cdcRow1, cdcRow2, cdcRow3)); + + Row tableRow1Updated = Row.withSchema(tableSchema).addValues("row-1-updated").build(); + File partFile1 = new File(tableDir, "part-00001.parquet"); + byte[] partBytes1 = + writeParquetFile(partFile1, tableSchema, java.util.Arrays.asList(tableRow1Updated)); + + writeCommit( + logDir, + 1L, + 200000000000L, + "part-00001.parquet", + partBytes1.length, + "part-00000.parquet", + "change-00000.parquet", + changeBytes0.length, + false); + + // 3. Write parquet files for Version 2 (insert-only commit) + Row tableRow3 = Row.withSchema(tableSchema).addValues("row-3").build(); + File partFile2 = new File(tableDir, "part-00002.parquet"); + byte[] partBytes2 = + writeParquetFile(partFile2, tableSchema, java.util.Arrays.asList(tableRow3)); + writeCommit( + logDir, 2L, 300000000000L, "part-00002.parquet", partBytes2.length, null, null, 0L, false); + + // Test 1: Read changes between start version 0 and end version 2 + PCollection<Row> outputVersions = + readPipeline.apply( + "Read Changes Version Range", + DeltaIO.readChanges() + .from(tableDir.getAbsolutePath()) + .withStartVersion(0L) + .withEndVersion(2L)); + + PCollection<String> formattedVersions = + outputVersions.apply("Format Version Output", ParDo.of(new FormatValueKindAndRow())); + + PAssert.that(formattedVersions) + .containsInAnyOrder( + "INSERT:row-1", + "INSERT:row-2", + "UPDATE_BEFORE:row-1", + "UPDATE_AFTER:row-1-updated", + "DELETE:row-2", + "INSERT:row-3"); + + // Test 2: Read changes between start timestamp (after version 0) and end timestamp (after + // version 2) + String startTimestamp = java.time.Instant.ofEpochMilli(150000000000L).toString(); + String endTimestamp = java.time.Instant.ofEpochMilli(350000000000L).toString(); + + PCollection<Row> outputTimestamps = + filteringPipeline.apply( + "Read Changes Timestamp Range", + DeltaIO.readChanges() + .from(tableDir.getAbsolutePath()) + .withStartTimestamp(startTimestamp) + .withEndTimestamp(endTimestamp)); + + PCollection<String> formattedTimestamps = + outputTimestamps.apply("Format Timestamp Output", ParDo.of(new FormatValueKindAndRow())); + + PAssert.that(formattedTimestamps) + .containsInAnyOrder( + "UPDATE_BEFORE:row-1", "UPDATE_AFTER:row-1-updated", "DELETE:row-2", "INSERT:row-3"); + + readPipeline.run().waitUntilFinish(); + filteringPipeline.run().waitUntilFinish(); + } + + private void writeCommit( + File logDir, + long version, + long timestamp, + @Nullable String addPath, + long addSize, + @Nullable String removePath, + @Nullable String cdcPath, + long cdcSize, + boolean writeMetadata) + throws IOException { + File commitFile = new File(logDir, String.format("%020d.json", version)); + StringBuilder content = new StringBuilder(); + if (version == 0 || writeMetadata) { + content.append("{\"protocol\":{\"minReaderVersion\":1,\"minWriterVersion\":2}}\n"); + content.append( + "{\"metaData\":{\"id\":\"test-id\",\"format\":{\"provider\":\"parquet\",\"options\":{}},\"schemaString\":\"{\\\"type\\\":\\\"struct\\\",\\\"fields\\\":[{\\\"name\\\":\\\"name\\\",\\\"type\\\":\\\"string\\\",\\\"nullable\\\":true,\\\"metadata\\\":{}}]}\",\"partitionColumns\":[],\"configuration\":{\"delta.enableChangeDataFeed\":\"true\"},\"createdAt\":123456789}}\n"); + } + if (addPath != null) { + content.append( + String.format( + "{\"add\":{\"path\":\"%s\",\"partitionValues\":{},\"size\":%d,\"modificationTime\":%d,\"dataChange\":true}}\n", + addPath, addSize, timestamp)); + } + if (removePath != null) { + content.append( + String.format( + "{\"remove\":{\"path\":\"%s\",\"deletionTimestamp\":%d,\"dataChange\":true}}\n", + removePath, timestamp)); + } + if (cdcPath != null) { + content.append( + String.format( + "{\"cdc\":{\"path\":\"%s\",\"partitionValues\":{},\"size\":%d,\"dataChange\":true}}\n", + cdcPath, cdcSize)); + } + Files.write(commitFile.toPath(), content.toString().getBytes(StandardCharsets.UTF_8)); + commitFile.setLastModified(timestamp); + } + + private static final class FormatValueKindAndRow extends DoFn<Row, String> { + @ProcessElement + public void process( + @Element Row row, ValueKind valueKind, OutputReceiver<String> outputReceiver) { + outputReceiver.output(valueKind.name() + ":" + row.getString("name")); + } + } + + private byte[] writeParquetFile(File file, Schema schema, java.util.List<Row> rows) Review Comment: There is already a `writeParquetFile` https://github.com/chamikaramj/beam/blob/82ab11edcafe9251d75af1185daf33f247ca261c/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java#L379 ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java: ########## @@ -0,0 +1,353 @@ +/* + * 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 io.delta.kernel.internal.DeltaErrors.wrapEngineException; + +import io.delta.kernel.Scan; +import io.delta.kernel.data.ColumnVector; +import io.delta.kernel.data.ColumnarBatch; +import io.delta.kernel.data.FilteredColumnarBatch; +import io.delta.kernel.data.MapValue; +import io.delta.kernel.defaults.engine.DefaultEngine; +import io.delta.kernel.engine.Engine; +import io.delta.kernel.engine.FileReadResult; +import io.delta.kernel.expressions.ExpressionEvaluator; +import io.delta.kernel.expressions.Literal; +import io.delta.kernel.internal.InternalScanFileUtils; +import io.delta.kernel.internal.data.GenericRow; +import io.delta.kernel.internal.data.ScanStateRow; +import io.delta.kernel.internal.util.Utils; +import io.delta.kernel.internal.util.VectorUtils; +import io.delta.kernel.types.LongType; +import io.delta.kernel.types.StringType; +import io.delta.kernel.types.StructField; +import io.delta.kernel.types.StructType; +import io.delta.kernel.types.TimestampType; +import io.delta.kernel.utils.CloseableIterator; +import io.delta.kernel.utils.FileStatus; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import org.apache.beam.sdk.io.range.OffsetRange; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.hadoop.conf.Configuration; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and reads Change Data Feed + * files, converting rows to Beam Rows. + */ [email protected] +class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> { + @Nullable Map<String, String> hadoopConfig; + private transient @Nullable Engine engine; + private transient @Nullable Configuration conf; + + public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) { + this.hadoopConfig = hadoopConfig; + } + + private synchronized Configuration getConfiguration() { + Configuration localConf = conf; + if (localConf == null) { + localConf = new Configuration(); + if (hadoopConfig != null) { + for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) { + localConf.set(entry.getKey(), entry.getValue()); + } + } + conf = localConf; + } + return localConf; + } + + private List<Long> getRowGroupSizes(DeltaCDCReadTask task) { + return task.getRowGroupSizes(); + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) { + List<Long> rowGroupSizes = getRowGroupSizes(task); + return new OffsetRange(0L, rowGroupSizes.size()); + } + + @NewTracker + public DeltaReadTaskTracker newTracker( + @Restriction OffsetRange restriction, @Element DeltaCDCReadTask task) { + return new DeltaReadTaskTracker(restriction, getRowGroupSizes(task)); + } + + @Setup + public void setUp() { + engine = DefaultEngine.create(getConfiguration()); + } + + @ProcessElement + public ProcessContinuation processElement( + @Element DeltaCDCReadTask task, + RestrictionTracker<OffsetRange, Long> tracker, + OutputReceiver<Row> out) + throws Exception { + + Engine currentEngine = engine; + if (currentEngine == null) { + throw new IllegalArgumentException("Expected the engine to not be null"); + } + + SerializableRow originalScanStateRow = task.getScanStateRow(); + StructType logicalTableSchema = ScanStateRow.getLogicalSchema(originalScanStateRow); + Schema publicBeamSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema); + StructType physicalTableSchema = ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow); + + StructType scanStateSchema = originalScanStateRow.getSchema(); + + // 1. Build modified scanState and scanFile rows depending on whether we read a CDC file or ADD + // file. + io.delta.kernel.data.Row scanStateRow; + StructType readPhysicalSchema; + StructType readLogicalSchema; + Schema beamSchema; + + if (task.isCDC()) { + readLogicalSchema = appendCDFColumns(logicalTableSchema); + readPhysicalSchema = appendCDFColumns(physicalTableSchema); + beamSchema = DeltaIO.ReadRows.convertToBeamSchema(readLogicalSchema); + + HashMap<Integer, Object> valueMap = new HashMap<>(); + + // Tracking row level lineage is not needed. + Map<String, String> config = + new HashMap<>(ScanStateRow.getConfiguration(originalScanStateRow)); + config.put("delta.enableRowTracking", "false"); + + valueMap.put( + scanStateSchema.indexOf("configuration"), VectorUtils.stringStringMapValue(config)); + valueMap.put(scanStateSchema.indexOf("logicalSchemaJson"), readLogicalSchema.toJson()); + valueMap.put(scanStateSchema.indexOf("physicalSchemaJson"), readPhysicalSchema.toJson()); + valueMap.put( + scanStateSchema.indexOf("partitionColumns"), + originalScanStateRow.getArray(scanStateSchema.indexOf("partitionColumns"))); + valueMap.put( + scanStateSchema.indexOf("minReaderVersion"), + originalScanStateRow.getInt(scanStateSchema.indexOf("minReaderVersion"))); + valueMap.put( + scanStateSchema.indexOf("minWriterVersion"), + originalScanStateRow.getInt(scanStateSchema.indexOf("minWriterVersion"))); + valueMap.put( + scanStateSchema.indexOf("tablePath"), + originalScanStateRow.getString(scanStateSchema.indexOf("tablePath"))); + + scanStateRow = new ScanStateRow(valueMap); + } else { + // For ADD files, we read the table schema and append the CDF columns manually afterwards. + // readLogicalSchema = logicalTableSchema; + readPhysicalSchema = physicalTableSchema; + beamSchema = DeltaIO.ReadRows.convertToBeamSchema(appendCDFColumns(logicalTableSchema)); + scanStateRow = originalScanStateRow; + } + + io.delta.kernel.data.Row scanFileRow = + generateScanFileRow(task.getPath(), task.getPartitionValues()); + FileStatus fileStatus = FileStatus.of(task.getPath(), task.getSize(), task.getTimestamp()); + + BeamParquetHandler parquetHandler = + new BeamParquetHandler(getConfiguration(), currentEngine.getParquetHandler(), tracker); + BeamEngine beamEngine = new BeamEngine(currentEngine, parquetHandler); + + long currentStartRgIndex = 0L; + + try (CloseableIterator<FileReadResult> fileReadResults = + parquetHandler.readParquetFiles( + Utils.singletonCloseableIterator(fileStatus), + readPhysicalSchema, + Optional.empty(), + currentStartRgIndex)) { + + CloseableIterator<ColumnarBatch> physicalData = + new CloseableIterator<ColumnarBatch>() { + @Override + public void close() throws java.io.IOException {} + + @Override + public boolean hasNext() { + return fileReadResults.hasNext(); + } + + @Override + public ColumnarBatch next() { + return fileReadResults.next().getData(); + } + }; + + try (CloseableIterator<FilteredColumnarBatch> logicalBatches = + Scan.transformPhysicalData(beamEngine, scanStateRow, scanFileRow, physicalData)) { + + while (logicalBatches.hasNext()) { + FilteredColumnarBatch batch = logicalBatches.next(); + ColumnarBatch logicalBatch = batch.getData(); Review Comment: (based on AI comment): `batch.getData()` return raw data, ignoring `selectionVector` and causing filtered-out or deleted rows to be iterated over and emitted into the output PCollection. Suggestion, call `batch. getRows` which handles `selectionVector` ```java while (logicalBatches.hasNext()) { FilteredColumnarBatch batch = logicalBatches.next(); if (!task.isCDC()) { ColumnarBatch logicalBatch = appendConstantCDFColumns( currentEngine, batch.getData(), task.getVersion(), task.getTimestamp()); batch = new FilteredColumnarBatch(logicalBatch, batch.getSelectionVector()); } try (CloseableIterator<io.delta.kernel.data.Row> logicalRows = batch.getRows()) { // <-- Uses FilteredColumnarBatch.getRows() while (logicalRows.hasNext()) { io.delta.kernel.data.Row deltaRow = logicalRows.next(); ... } } } ``` ########## sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java: ########## Review Comment: Not related to the change. But running the test realize there are `System.out.println` and `System.err.println` debugging leftovers. Consider removing them or using slf4j LOGGER. ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java: ########## @@ -0,0 +1,353 @@ +/* + * 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 io.delta.kernel.internal.DeltaErrors.wrapEngineException; + +import io.delta.kernel.Scan; +import io.delta.kernel.data.ColumnVector; +import io.delta.kernel.data.ColumnarBatch; +import io.delta.kernel.data.FilteredColumnarBatch; +import io.delta.kernel.data.MapValue; +import io.delta.kernel.defaults.engine.DefaultEngine; +import io.delta.kernel.engine.Engine; +import io.delta.kernel.engine.FileReadResult; +import io.delta.kernel.expressions.ExpressionEvaluator; +import io.delta.kernel.expressions.Literal; +import io.delta.kernel.internal.InternalScanFileUtils; +import io.delta.kernel.internal.data.GenericRow; +import io.delta.kernel.internal.data.ScanStateRow; +import io.delta.kernel.internal.util.Utils; +import io.delta.kernel.internal.util.VectorUtils; +import io.delta.kernel.types.LongType; +import io.delta.kernel.types.StringType; +import io.delta.kernel.types.StructField; +import io.delta.kernel.types.StructType; +import io.delta.kernel.types.TimestampType; +import io.delta.kernel.utils.CloseableIterator; +import io.delta.kernel.utils.FileStatus; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import org.apache.beam.sdk.io.range.OffsetRange; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.hadoop.conf.Configuration; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and reads Change Data Feed + * files, converting rows to Beam Rows. + */ [email protected] +class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> { + @Nullable Map<String, String> hadoopConfig; + private transient @Nullable Engine engine; + private transient @Nullable Configuration conf; + + public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) { + this.hadoopConfig = hadoopConfig; + } + + private synchronized Configuration getConfiguration() { + Configuration localConf = conf; + if (localConf == null) { + localConf = new Configuration(); + if (hadoopConfig != null) { + for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) { + localConf.set(entry.getKey(), entry.getValue()); + } + } + conf = localConf; + } + return localConf; + } + + private List<Long> getRowGroupSizes(DeltaCDCReadTask task) { + return task.getRowGroupSizes(); + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) { + List<Long> rowGroupSizes = getRowGroupSizes(task); + return new OffsetRange(0L, rowGroupSizes.size()); + } + + @NewTracker + public DeltaReadTaskTracker newTracker( + @Restriction OffsetRange restriction, @Element DeltaCDCReadTask task) { + return new DeltaReadTaskTracker(restriction, getRowGroupSizes(task)); + } + + @Setup + public void setUp() { + engine = DefaultEngine.create(getConfiguration()); + } + + @ProcessElement + public ProcessContinuation processElement( + @Element DeltaCDCReadTask task, + RestrictionTracker<OffsetRange, Long> tracker, + OutputReceiver<Row> out) + throws Exception { + + Engine currentEngine = engine; + if (currentEngine == null) { + throw new IllegalArgumentException("Expected the engine to not be null"); + } + + SerializableRow originalScanStateRow = task.getScanStateRow(); + StructType logicalTableSchema = ScanStateRow.getLogicalSchema(originalScanStateRow); + Schema publicBeamSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema); + StructType physicalTableSchema = ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow); + + StructType scanStateSchema = originalScanStateRow.getSchema(); + + // 1. Build modified scanState and scanFile rows depending on whether we read a CDC file or ADD + // file. + io.delta.kernel.data.Row scanStateRow; + StructType readPhysicalSchema; + StructType readLogicalSchema; + Schema beamSchema; + + if (task.isCDC()) { + readLogicalSchema = appendCDFColumns(logicalTableSchema); + readPhysicalSchema = appendCDFColumns(physicalTableSchema); + beamSchema = DeltaIO.ReadRows.convertToBeamSchema(readLogicalSchema); + + HashMap<Integer, Object> valueMap = new HashMap<>(); + + // Tracking row level lineage is not needed. + Map<String, String> config = + new HashMap<>(ScanStateRow.getConfiguration(originalScanStateRow)); + config.put("delta.enableRowTracking", "false"); + + valueMap.put( + scanStateSchema.indexOf("configuration"), VectorUtils.stringStringMapValue(config)); + valueMap.put(scanStateSchema.indexOf("logicalSchemaJson"), readLogicalSchema.toJson()); + valueMap.put(scanStateSchema.indexOf("physicalSchemaJson"), readPhysicalSchema.toJson()); + valueMap.put( + scanStateSchema.indexOf("partitionColumns"), + originalScanStateRow.getArray(scanStateSchema.indexOf("partitionColumns"))); + valueMap.put( + scanStateSchema.indexOf("minReaderVersion"), + originalScanStateRow.getInt(scanStateSchema.indexOf("minReaderVersion"))); + valueMap.put( + scanStateSchema.indexOf("minWriterVersion"), + originalScanStateRow.getInt(scanStateSchema.indexOf("minWriterVersion"))); + valueMap.put( + scanStateSchema.indexOf("tablePath"), + originalScanStateRow.getString(scanStateSchema.indexOf("tablePath"))); + + scanStateRow = new ScanStateRow(valueMap); + } else { + // For ADD files, we read the table schema and append the CDF columns manually afterwards. + // readLogicalSchema = logicalTableSchema; + readPhysicalSchema = physicalTableSchema; + beamSchema = DeltaIO.ReadRows.convertToBeamSchema(appendCDFColumns(logicalTableSchema)); + scanStateRow = originalScanStateRow; + } + + io.delta.kernel.data.Row scanFileRow = + generateScanFileRow(task.getPath(), task.getPartitionValues()); + FileStatus fileStatus = FileStatus.of(task.getPath(), task.getSize(), task.getTimestamp()); + + BeamParquetHandler parquetHandler = + new BeamParquetHandler(getConfiguration(), currentEngine.getParquetHandler(), tracker); + BeamEngine beamEngine = new BeamEngine(currentEngine, parquetHandler); + + long currentStartRgIndex = 0L; + + try (CloseableIterator<FileReadResult> fileReadResults = + parquetHandler.readParquetFiles( + Utils.singletonCloseableIterator(fileStatus), + readPhysicalSchema, + Optional.empty(), + currentStartRgIndex)) { + + CloseableIterator<ColumnarBatch> physicalData = + new CloseableIterator<ColumnarBatch>() { + @Override + public void close() throws java.io.IOException {} + + @Override + public boolean hasNext() { + return fileReadResults.hasNext(); + } + + @Override + public ColumnarBatch next() { + return fileReadResults.next().getData(); + } + }; + + try (CloseableIterator<FilteredColumnarBatch> logicalBatches = + Scan.transformPhysicalData(beamEngine, scanStateRow, scanFileRow, physicalData)) { + + while (logicalBatches.hasNext()) { + FilteredColumnarBatch batch = logicalBatches.next(); + ColumnarBatch logicalBatch = batch.getData(); + + if (!task.isCDC()) { + // For ADD files, we need to append the constant CDF columns: + // _change_type = "insert", _commit_version = task.version, _commit_timestamp = + // task.timestamp + logicalBatch = + appendConstantCDFColumns( + currentEngine, logicalBatch, task.getVersion(), task.getTimestamp()); + } + + try (CloseableIterator<io.delta.kernel.data.Row> logicalRows = logicalBatch.getRows()) { + while (logicalRows.hasNext()) { + io.delta.kernel.data.Row deltaRow = logicalRows.next(); + Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema); + String changeType = beamRow.getString("_change_type"); + if (changeType == null) { + throw new IllegalStateException("Field _change_type must not be null."); + } + ValueKind kind = getValueKind(changeType); + Row publicRow = projectRow(beamRow, publicBeamSchema); + out.builder(publicRow).setValueKind(kind).output(); + } + } + } + } + } + + return ProcessContinuation.stop(); Review Comment: This is the only place where DoFn exits. if there is error throws, will it retry the whole element? In other words, would bounded SDF just resume from a claimed restriction even though previously crashed in the middle of a DoFn.process invocation? -- 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]
