chamikaramj commented on code in PR #39426: URL: https://github.com/apache/beam/pull/39426#discussion_r3671079446
########## sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java: ########## Review Comment: Done. ########## 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: It should retry the whole workitem (including the whole element) if there's a crash. Agree that we can just return here instead of using `ProcessContinuation`. Updated. ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java: ########## @@ -174,4 +180,124 @@ static Schema.FieldType convertToBeamFieldType(DataType deltaType) { } } } + + @AutoValue + public abstract static class ReadChanges extends PTransform<PBegin, PCollection<Row>> { + public abstract @Nullable String getTablePath(); + + public abstract @Nullable Long getStartVersion(); + + public abstract @Nullable String getStartTimestamp(); + + public abstract @Nullable Long getEndVersion(); + + public abstract @Nullable String getEndTimestamp(); + + public abstract @Nullable Map<String, String> getHadoopConfig(); + + abstract Builder toBuilder(); + + @AutoValue.Builder + abstract static class Builder { + abstract Builder setTablePath(String tablePath); + + abstract Builder setStartVersion(@Nullable Long startVersion); + + abstract Builder setStartTimestamp(@Nullable String startTimestamp); + + abstract Builder setEndVersion(@Nullable Long endVersion); + + abstract Builder setEndTimestamp(@Nullable String endTimestamp); + + abstract Builder setHadoopConfig(@Nullable Map<String, String> hadoopConfig); + + abstract ReadChanges build(); + } + + public ReadChanges from(String tablePath) { + return toBuilder().setTablePath(tablePath).build(); + } + + public ReadChanges withStartVersion(long startVersion) { + return toBuilder().setStartVersion(startVersion).build(); + } + + public ReadChanges withStartTimestamp(String startTimestamp) { + return toBuilder().setStartTimestamp(startTimestamp).build(); + } + + public ReadChanges withEndVersion(long endVersion) { + return toBuilder().setEndVersion(endVersion).build(); + } + + public ReadChanges withEndTimestamp(String endTimestamp) { + return toBuilder().setEndTimestamp(endTimestamp).build(); + } + + public ReadChanges withConfig(Map<String, String> config) { + return toBuilder().setHadoopConfig(config).build(); + } + + @Override + public PCollection<Row> expand(PBegin input) { + String path = getTablePath(); + if (path == null) { + throw new IllegalArgumentException("Table path must be set."); + } + if (getStartVersion() == null && getStartTimestamp() == null) { + throw new IllegalArgumentException("Either startVersion or startTimestamp must be set."); + } Review Comment: I think given that this is a bounded source, it's cleaner for users to provide a starting point ? Starting from current doesn't make sense since that would not result in anything (makes sense to extend this after we support unbounded reads). Added a TODO. ########## 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: Done. ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java: ########## @@ -174,4 +180,124 @@ static Schema.FieldType convertToBeamFieldType(DataType deltaType) { } } } + + @AutoValue + public abstract static class ReadChanges extends PTransform<PBegin, PCollection<Row>> { + public abstract @Nullable String getTablePath(); + + public abstract @Nullable Long getStartVersion(); + + public abstract @Nullable String getStartTimestamp(); + + public abstract @Nullable Long getEndVersion(); + + public abstract @Nullable String getEndTimestamp(); + + public abstract @Nullable Map<String, String> getHadoopConfig(); + + abstract Builder toBuilder(); + + @AutoValue.Builder + abstract static class Builder { + abstract Builder setTablePath(String tablePath); + + abstract Builder setStartVersion(@Nullable Long startVersion); + + abstract Builder setStartTimestamp(@Nullable String startTimestamp); + + abstract Builder setEndVersion(@Nullable Long endVersion); + + abstract Builder setEndTimestamp(@Nullable String endTimestamp); + + abstract Builder setHadoopConfig(@Nullable Map<String, String> hadoopConfig); + + abstract ReadChanges build(); + } + + public ReadChanges from(String tablePath) { + return toBuilder().setTablePath(tablePath).build(); + } + + public ReadChanges withStartVersion(long startVersion) { + return toBuilder().setStartVersion(startVersion).build(); + } + + public ReadChanges withStartTimestamp(String startTimestamp) { + return toBuilder().setStartTimestamp(startTimestamp).build(); + } + + public ReadChanges withEndVersion(long endVersion) { + return toBuilder().setEndVersion(endVersion).build(); + } + + public ReadChanges withEndTimestamp(String endTimestamp) { + return toBuilder().setEndTimestamp(endTimestamp).build(); + } + + public ReadChanges withConfig(Map<String, String> config) { + return toBuilder().setHadoopConfig(config).build(); + } + + @Override + public PCollection<Row> expand(PBegin input) { + String path = getTablePath(); + if (path == null) { + throw new IllegalArgumentException("Table path must be set."); + } + if (getStartVersion() == null && getStartTimestamp() == null) { + throw new IllegalArgumentException("Either startVersion or startTimestamp must be set."); + } + if (getStartVersion() != null && getStartTimestamp() != null) { + throw new IllegalArgumentException("Cannot set both startVersion and startTimestamp."); + } + if (getEndVersion() != null && getEndTimestamp() != null) { + throw new IllegalArgumentException("Cannot set both endVersion and endTimestamp."); + } Review Comment: Yeah, this is handled. We only fail if end is specified both as a version and a timestampl. ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java: ########## @@ -0,0 +1,293 @@ +/* + * 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 io.delta.kernel.CommitRange; +import io.delta.kernel.CommitRangeBuilder; +import io.delta.kernel.Scan; +import io.delta.kernel.Snapshot; +import io.delta.kernel.Table; +import io.delta.kernel.TableManager; +import io.delta.kernel.data.ColumnarBatch; +import io.delta.kernel.data.Row; +import io.delta.kernel.defaults.engine.DefaultEngine; +import io.delta.kernel.engine.Engine; +import io.delta.kernel.internal.DeltaLogActionUtils.DeltaAction; +import io.delta.kernel.internal.TableImpl; +import io.delta.kernel.internal.actions.AddCDCFile; +import io.delta.kernel.internal.actions.AddFile; +import io.delta.kernel.internal.util.VectorUtils; +import io.delta.kernel.utils.CloseableIterator; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.hadoop.conf.Configuration; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** A DoFn that reads the Delta log and plans read tasks for Change Data Feed. */ +class CreateCDCReadTasksDoFn extends DoFn<String, DeltaCDCReadTask> { + private static final long MAX_TASK_SIZE_BYTES = 1024L * 1024L * 1024L; // 1 GB + private final @Nullable Map<String, String> hadoopConfig; + private final @Nullable Long startVersion; + private final @Nullable String startTimestamp; + private final @Nullable Long endVersion; + private final @Nullable String endTimestamp; + + public CreateCDCReadTasksDoFn( + @Nullable Map<String, String> hadoopConfig, + @Nullable Long startVersion, + @Nullable String startTimestamp, + @Nullable Long endVersion, + @Nullable String endTimestamp) { + this.hadoopConfig = hadoopConfig; + this.startVersion = startVersion; + this.startTimestamp = startTimestamp; + this.endVersion = endVersion; + this.endTimestamp = endTimestamp; + } + + @ProcessElement + public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCReadTask> out) + throws Exception { + Configuration conf = new Configuration(); + if (hadoopConfig != null) { + for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) { + conf.set(entry.getKey(), entry.getValue()); + } + } + Engine engine = DefaultEngine.create(conf); + Table table = Table.forPath(engine, tablePath); + TableImpl tableImpl = (TableImpl) table; + + // 1. Resolve starting and ending versions + long resolvedStartVersion; + if (startVersion != null) { + resolvedStartVersion = startVersion; + } else if (startTimestamp != null) { + long startMillis = Instant.parse(startTimestamp).toEpochMilli(); + resolvedStartVersion = tableImpl.getVersionAtOrAfterTimestamp(engine, startMillis); + } else { + throw new IllegalArgumentException("Starting version or timestamp must be specified."); + } + + long resolvedEndVersion; + if (endVersion != null) { + resolvedEndVersion = endVersion; + } else if (endTimestamp != null) { + long endMillis = Instant.parse(endTimestamp).toEpochMilli(); + resolvedEndVersion = tableImpl.getVersionBeforeOrAtTimestamp(engine, endMillis); + } else { + resolvedEndVersion = table.getLatestSnapshot(engine).getVersion(); + } + + if (resolvedStartVersion > resolvedEndVersion) { + throw new IllegalArgumentException( + String.format( + "Resolved start version %d is greater than resolved end version %d", + resolvedStartVersion, resolvedEndVersion)); + } + + // 2. Load snapshot at resolvedEndVersion to get the scanStateRow + // We use endVersion's schema because it represents the latest schema in the + // read range + // which handles schema evolution (older files will just lack new columns). + Snapshot endSnapshot = table.getSnapshotAsOfVersion(engine, resolvedEndVersion); + Scan scan = endSnapshot.getScanBuilder().build(); + Row scanState = scan.getScanState(engine); + SerializableRow serializableScanState = new SerializableRow(scanState); + + // 3. Load snapshot at resolvedStartVersion to initialize the CommitRange + Snapshot startSnapshot = table.getSnapshotAsOfVersion(engine, resolvedStartVersion); + + CommitRangeBuilder rangeBuilder = + TableManager.loadCommitRange( + tablePath, CommitRangeBuilder.CommitBoundary.atVersion(resolvedStartVersion)); + rangeBuilder.withEndBoundary(CommitRangeBuilder.CommitBoundary.atVersion(resolvedEndVersion)); + CommitRange range = rangeBuilder.build(engine); + + // We need both CDC and ADD actions. + // If a commit version has CDC files, we only read CDC files. + // If a commit version has no CDC files, we read ADD files (inserts). + Set<DeltaAction> actionSet = new HashSet<>(); + actionSet.add(DeltaAction.CDC); + actionSet.add(DeltaAction.ADD); + + // 4. Iterate over commits in the range and group actions by version + try (CloseableIterator<ColumnarBatch> batchIter = + range.getActions(engine, startSnapshot, actionSet)) { + Map<Long, CommitActionsInfo> commitActionsMap = new HashMap<>(); + + while (batchIter.hasNext()) { + ColumnarBatch batch = batchIter.next(); + int versionIdx = batch.getSchema().indexOf("version"); + int timestampIdx = batch.getSchema().indexOf("timestamp"); + int cdcIdx = batch.getSchema().indexOf("cdc"); + int addIdx = batch.getSchema().indexOf("add"); + + for (int i = 0; i < batch.getSize(); i++) { + long version = batch.getColumnVector(versionIdx).getLong(i); + long timestamp = batch.getColumnVector(timestampIdx).getLong(i); + + CommitActionsInfo info = + commitActionsMap.computeIfAbsent( + version, k -> new CommitActionsInfo(version, timestamp)); + + if (cdcIdx >= 0 && !batch.getColumnVector(cdcIdx).isNullAt(i)) { + Row cdcRow = + (Row) + VectorUtils.getValueAsObject( + batch.getColumnVector(cdcIdx), + batch.getSchema().at(cdcIdx).getDataType(), + i); + info.cdcRows.add(cdcRow); + } + if (addIdx >= 0 && !batch.getColumnVector(addIdx).isNullAt(i)) { + Row addRow = + (Row) + VectorUtils.getValueAsObject( + batch.getColumnVector(addIdx), + batch.getSchema().at(addIdx).getDataType(), + i); + AddFile addFile = new AddFile(addRow); + // Only consider add files that change data (ignore OPTIMIZE etc.) + if (addFile.getDataChange()) { + info.addRows.add(addRow); + } + } + } + } + + // 5. Emit tasks for each version + List<DeltaCDCReadTask> currentGroup = new ArrayList<>(); + long currentGroupSize = 0L; + + // Sort versions to process them in order + List<Long> versions = new ArrayList<>(commitActionsMap.keySet()); + Collections.sort(versions); + + for (long version : versions) { + CommitActionsInfo info = commitActionsMap.get(version); + if (info == null) { + throw new IllegalStateException("CommitActionsInfo was not found for version " + version); + } + boolean hasCDC = !info.cdcRows.isEmpty(); + + List<Row> rowsToProcess = hasCDC ? info.cdcRows : info.addRows; + boolean isCDC = hasCDC; + + for (Row fileRow : rowsToProcess) { + String relPath; + long size; + Map<String, String> partitionValues; + + if (isCDC) { + relPath = fileRow.getString(AddCDCFile.FULL_SCHEMA.indexOf("path")); + size = fileRow.getLong(AddCDCFile.FULL_SCHEMA.indexOf("size")); + partitionValues = + VectorUtils.toJavaMap( + fileRow.getMap(AddCDCFile.FULL_SCHEMA.indexOf("partitionValues"))); + } else { + AddFile addFile = new AddFile(fileRow); + relPath = addFile.getPath(); + size = addFile.getSize(); + partitionValues = VectorUtils.toJavaMap(addFile.getPartitionValues()); + } + + String fullPath = new org.apache.hadoop.fs.Path(tablePath, relPath).toString(); + List<Long> rowGroupSizes = getRowGroupSizes(fullPath, conf); + + DeltaCDCReadTask task = + new DeltaCDCReadTask( + fullPath, + size, + partitionValues, + info.version, + info.timestamp, + isCDC, + rowGroupSizes, + serializableScanState); + + if (size >= MAX_TASK_SIZE_BYTES) { + if (!currentGroup.isEmpty()) { + emitGroup(currentGroup, out); + currentGroup = new ArrayList<>(); + currentGroupSize = 0L; + } + out.output(task); + } else { + if (currentGroupSize + size > MAX_TASK_SIZE_BYTES) { + emitGroup(currentGroup, out); + currentGroup = new ArrayList<>(); + currentGroup.add(task); + currentGroupSize = size; + } else { + currentGroup.add(task); + currentGroupSize += size; + } + } + } + } + + if (!currentGroup.isEmpty()) { + emitGroup(currentGroup, out); + } + } + } + + private void emitGroup(List<DeltaCDCReadTask> group, OutputReceiver<DeltaCDCReadTask> out) { + for (DeltaCDCReadTask task : group) { + out.output(task); + } + } + + private List<Long> getRowGroupSizes(String pathStr, Configuration conf) { + List<Long> sizes = new ArrayList<>(); + try { + org.apache.hadoop.fs.Path hadoopPath = new org.apache.hadoop.fs.Path(pathStr); + org.apache.parquet.hadoop.metadata.ParquetMetadata metadata = + org.apache.parquet.hadoop.ParquetFileReader.readFooter( + conf, + hadoopPath, + org.apache.parquet.format.converter.ParquetMetadataConverter.NO_FILTER); + for (org.apache.parquet.hadoop.metadata.BlockMetaData block : metadata.getBlocks()) { + sizes.add(block.getTotalByteSize()); + } + } catch (java.io.IOException e) { + throw new RuntimeException("Failed to read Parquet footer for " + pathStr, e); + } + return sizes; + } + + private static class CommitActionsInfo { + final long version; + final long timestamp; + final List<Row> cdcRows = new ArrayList<>(); + final List<Row> addRows = new ArrayList<>(); Review Comment: Yeah, renamed to better convey this. ########## 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 { Review Comment: Done. ########## 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> { Review Comment: Yeah, for streaming, we could consider batching files. This just supports batch reads for now. Added a TODO. ########## 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: Done. ########## 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: Done. ########## sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java: ########## @@ -0,0 +1,293 @@ +/* + * 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 io.delta.kernel.CommitRange; +import io.delta.kernel.CommitRangeBuilder; +import io.delta.kernel.Scan; +import io.delta.kernel.Snapshot; +import io.delta.kernel.Table; +import io.delta.kernel.TableManager; +import io.delta.kernel.data.ColumnarBatch; +import io.delta.kernel.data.Row; +import io.delta.kernel.defaults.engine.DefaultEngine; +import io.delta.kernel.engine.Engine; +import io.delta.kernel.internal.DeltaLogActionUtils.DeltaAction; +import io.delta.kernel.internal.TableImpl; +import io.delta.kernel.internal.actions.AddCDCFile; +import io.delta.kernel.internal.actions.AddFile; +import io.delta.kernel.internal.util.VectorUtils; +import io.delta.kernel.utils.CloseableIterator; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.hadoop.conf.Configuration; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** A DoFn that reads the Delta log and plans read tasks for Change Data Feed. */ +class CreateCDCReadTasksDoFn extends DoFn<String, DeltaCDCReadTask> { + private static final long MAX_TASK_SIZE_BYTES = 1024L * 1024L * 1024L; // 1 GB + private final @Nullable Map<String, String> hadoopConfig; + private final @Nullable Long startVersion; + private final @Nullable String startTimestamp; + private final @Nullable Long endVersion; + private final @Nullable String endTimestamp; + + public CreateCDCReadTasksDoFn( + @Nullable Map<String, String> hadoopConfig, + @Nullable Long startVersion, + @Nullable String startTimestamp, + @Nullable Long endVersion, + @Nullable String endTimestamp) { + this.hadoopConfig = hadoopConfig; + this.startVersion = startVersion; + this.startTimestamp = startTimestamp; + this.endVersion = endVersion; + this.endTimestamp = endTimestamp; + } + + @ProcessElement + public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCReadTask> out) + throws Exception { + Configuration conf = new Configuration(); + if (hadoopConfig != null) { + for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) { + conf.set(entry.getKey(), entry.getValue()); + } + } + Engine engine = DefaultEngine.create(conf); + Table table = Table.forPath(engine, tablePath); + TableImpl tableImpl = (TableImpl) table; + + // 1. Resolve starting and ending versions + long resolvedStartVersion; + if (startVersion != null) { + resolvedStartVersion = startVersion; + } else if (startTimestamp != null) { + long startMillis = Instant.parse(startTimestamp).toEpochMilli(); + resolvedStartVersion = tableImpl.getVersionAtOrAfterTimestamp(engine, startMillis); + } else { + throw new IllegalArgumentException("Starting version or timestamp must be specified."); + } + + long resolvedEndVersion; + if (endVersion != null) { + resolvedEndVersion = endVersion; + } else if (endTimestamp != null) { + long endMillis = Instant.parse(endTimestamp).toEpochMilli(); + resolvedEndVersion = tableImpl.getVersionBeforeOrAtTimestamp(engine, endMillis); + } else { + resolvedEndVersion = table.getLatestSnapshot(engine).getVersion(); + } + + if (resolvedStartVersion > resolvedEndVersion) { + throw new IllegalArgumentException( + String.format( + "Resolved start version %d is greater than resolved end version %d", + resolvedStartVersion, resolvedEndVersion)); + } + + // 2. Load snapshot at resolvedEndVersion to get the scanStateRow + // We use endVersion's schema because it represents the latest schema in the + // read range + // which handles schema evolution (older files will just lack new columns). + Snapshot endSnapshot = table.getSnapshotAsOfVersion(engine, resolvedEndVersion); + Scan scan = endSnapshot.getScanBuilder().build(); + Row scanState = scan.getScanState(engine); + SerializableRow serializableScanState = new SerializableRow(scanState); + + // 3. Load snapshot at resolvedStartVersion to initialize the CommitRange + Snapshot startSnapshot = table.getSnapshotAsOfVersion(engine, resolvedStartVersion); + + CommitRangeBuilder rangeBuilder = + TableManager.loadCommitRange( + tablePath, CommitRangeBuilder.CommitBoundary.atVersion(resolvedStartVersion)); + rangeBuilder.withEndBoundary(CommitRangeBuilder.CommitBoundary.atVersion(resolvedEndVersion)); + CommitRange range = rangeBuilder.build(engine); + + // We need both CDC and ADD actions. + // If a commit version has CDC files, we only read CDC files. + // If a commit version has no CDC files, we read ADD files (inserts). + Set<DeltaAction> actionSet = new HashSet<>(); + actionSet.add(DeltaAction.CDC); + actionSet.add(DeltaAction.ADD); + + // 4. Iterate over commits in the range and group actions by version + try (CloseableIterator<ColumnarBatch> batchIter = + range.getActions(engine, startSnapshot, actionSet)) { + Map<Long, CommitActionsInfo> commitActionsMap = new HashMap<>(); + + while (batchIter.hasNext()) { + ColumnarBatch batch = batchIter.next(); + int versionIdx = batch.getSchema().indexOf("version"); + int timestampIdx = batch.getSchema().indexOf("timestamp"); + int cdcIdx = batch.getSchema().indexOf("cdc"); + int addIdx = batch.getSchema().indexOf("add"); + + for (int i = 0; i < batch.getSize(); i++) { + long version = batch.getColumnVector(versionIdx).getLong(i); + long timestamp = batch.getColumnVector(timestampIdx).getLong(i); + + CommitActionsInfo info = + commitActionsMap.computeIfAbsent( + version, k -> new CommitActionsInfo(version, timestamp)); + + if (cdcIdx >= 0 && !batch.getColumnVector(cdcIdx).isNullAt(i)) { + Row cdcRow = + (Row) + VectorUtils.getValueAsObject( + batch.getColumnVector(cdcIdx), + batch.getSchema().at(cdcIdx).getDataType(), + i); + info.cdcRows.add(cdcRow); + } + if (addIdx >= 0 && !batch.getColumnVector(addIdx).isNullAt(i)) { + Row addRow = + (Row) + VectorUtils.getValueAsObject( + batch.getColumnVector(addIdx), + batch.getSchema().at(addIdx).getDataType(), + i); + AddFile addFile = new AddFile(addRow); + // Only consider add files that change data (ignore OPTIMIZE etc.) + if (addFile.getDataChange()) { + info.addRows.add(addRow); + } + } + } + } + + // 5. Emit tasks for each version + List<DeltaCDCReadTask> currentGroup = new ArrayList<>(); + long currentGroupSize = 0L; + + // Sort versions to process them in order + List<Long> versions = new ArrayList<>(commitActionsMap.keySet()); + Collections.sort(versions); + + for (long version : versions) { + CommitActionsInfo info = commitActionsMap.get(version); + if (info == null) { + throw new IllegalStateException("CommitActionsInfo was not found for version " + version); + } + boolean hasCDC = !info.cdcRows.isEmpty(); + + List<Row> rowsToProcess = hasCDC ? info.cdcRows : info.addRows; + boolean isCDC = hasCDC; Review Comment: Done. ########## 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: Done. ########## 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) + .add("_commit_version", LongType.LONG, false) + .add("_commit_timestamp", TimestampType.TIMESTAMP, false); Review Comment: Done. Seems like these map well to what's defined in ValueKind. What's the incompatibility you expect for Iceberg sink ? ########## 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)); Review Comment: Added a new test for this. ########## 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(); + } Review Comment: Done. -- 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]
