From 6502458e0fdc9cbf7fdb8ef89ce2e912ce2d2854 Mon Sep 17 00:00:00 2001 From: Chamikara Jayalath Date: Tue, 21 Jul 2026 22:52:59 -0700 Subject: [PATCH 1/2] Updates the Delta Lake source to support reading bounded change data --- .../sdk/io/delta/CreateCDCReadTasksDoFn.java | 286 +++++++++++++++ .../beam/sdk/io/delta/DeltaCDCReadTask.java | 140 +++++++ .../beam/sdk/io/delta/DeltaCDCSourceDoFn.java | 341 ++++++++++++++++++ .../org/apache/beam/sdk/io/delta/DeltaIO.java | 131 +++++++ .../beam/sdk/io/delta/DeltaSourceDoFn.java | 2 +- .../apache/beam/sdk/io/delta/DeltaIOTest.java | 152 ++++++++ 6 files changed, 1051 insertions(+), 1 deletion(-) create mode 100644 sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java create mode 100644 sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCReadTask.java create mode 100644 sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java new file mode 100644 index 000000000000..16843b91906f --- /dev/null +++ b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java @@ -0,0 +1,286 @@ +/* + * 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 { + private static final long MAX_TASK_SIZE_BYTES = 1024L * 1024L * 1024L; // 1 GB + private final @Nullable Map 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 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 out) + throws Exception { + Configuration conf = new Configuration(); + if (hadoopConfig != null) { + for (Map.Entry 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 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 batchIter = range.getActions(engine, startSnapshot, actionSet)) { + Map 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 currentGroup = new ArrayList<>(); + long currentGroupSize = 0L; + + // Sort versions to process them in order + List 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 rowsToProcess = hasCDC ? info.cdcRows : info.addRows; + boolean isCDC = hasCDC; + + for (Row fileRow : rowsToProcess) { + String relPath; + long size; + Map 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 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 group, OutputReceiver out) { + for (DeltaCDCReadTask task : group) { + out.output(task); + } + } + + private List getRowGroupSizes(String pathStr, Configuration conf) { + List 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 cdcRows = new ArrayList<>(); + final List addRows = new ArrayList<>(); + + CommitActionsInfo(long version, long timestamp) { + this.version = version; + this.timestamp = timestamp; + } + } +} diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCReadTask.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCReadTask.java new file mode 100644 index 000000000000..f2adf354dc09 --- /dev/null +++ b/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 partitionValues; + private final long version; + private final long timestamp; + private final boolean isCDC; + private final List rowGroupSizes; + private final SerializableRow scanStateRow; + + public DeltaCDCReadTask( + String path, + long size, + Map partitionValues, + long version, + long timestamp, + boolean isCDC, + List 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 getPartitionValues() { + return partitionValues; + } + + public long getVersion() { + return version; + } + + public long getTimestamp() { + return timestamp; + } + + public boolean isCDC() { + return isCDC; + } + + public List 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{" + + "path='" + + path + + '\'' + + ", size=" + + size + + ", partitionValues=" + + partitionValues + + ", version=" + + version + + ", timestamp=" + + timestamp + + ", isCDC=" + + isCDC + + ", rowGroupSizes=" + + rowGroupSizes + + ", scanStateRow=" + + scanStateRow + + '}'; + } +} diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java new file mode 100644 index 000000000000..41cfd5d2e716 --- /dev/null +++ b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java @@ -0,0 +1,341 @@ +/* + * 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. + */ +@DoFn.BoundedPerElement +class DeltaCDCSourceDoFn extends DoFn { + @Nullable Map hadoopConfig; + private transient @Nullable Engine engine; + private transient @Nullable Configuration conf; + + public DeltaCDCSourceDoFn(@Nullable Map hadoopConfig) { + this.hadoopConfig = hadoopConfig; + } + + private synchronized Configuration getConfiguration() { + Configuration localConf = conf; + if (localConf == null) { + localConf = new Configuration(); + if (hadoopConfig != null) { + for (Map.Entry entry : hadoopConfig.entrySet()) { + localConf.set(entry.getKey(), entry.getValue()); + } + } + conf = localConf; + } + return localConf; + } + + private List getRowGroupSizes(DeltaCDCReadTask task) { + return task.getRowGroupSizes(); + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) { + List 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 tracker, + OutputReceiver 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); + 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 valueMap = new HashMap<>(); + + // Tracking row level lineage is not needed. + Map 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 fileReadResults = + parquetHandler.readParquetFiles( + Utils.singletonCloseableIterator(fileStatus), + readPhysicalSchema, + Optional.empty(), + currentStartRgIndex)) { + + CloseableIterator physicalData = + new CloseableIterator() { + @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 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 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); + out.builder(beamRow).setValueKind(kind).output(); + } + } + } + } + } + + return ProcessContinuation.stop(); + } + + private static ValueKind getValueKind(String changeType) { + 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); + } + + private ColumnarBatch appendConstantCDFColumns( + Engine engine, ColumnarBatch batch, long version, long timestamp) { + StructType schemaForEval = batch.getSchema(); + + ExpressionEvaluator changeTypeGenerator = + wrapEngineException( + () -> + engine + .getExpressionHandler() + .getEvaluator(schemaForEval, Literal.ofString("insert"), StringType.STRING), + "Get the expression evaluator for change type"); + + ExpressionEvaluator commitVersionGenerator = + wrapEngineException( + () -> + engine + .getExpressionHandler() + .getEvaluator(schemaForEval, Literal.ofLong(version), LongType.LONG), + "Get the expression evaluator for commit version"); + + ExpressionEvaluator commitTimestampGenerator = + wrapEngineException( + () -> + engine + .getExpressionHandler() + // Microseconds since epoch is expected for TimestampType + .getEvaluator( + schemaForEval, + Literal.ofTimestamp(timestamp * 1000L), + TimestampType.TIMESTAMP), + "Get the expression evaluator for commit timestamp"); + + ColumnVector changeTypeVector = + wrapEngineException( + () -> changeTypeGenerator.eval(batch), "Evaluating change type expression"); + + ColumnVector commitVersionVector = + wrapEngineException( + () -> commitVersionGenerator.eval(batch), "Evaluating commit version expression"); + + ColumnVector commitTimestampVector = + wrapEngineException( + () -> commitTimestampGenerator.eval(batch), "Evaluating commit timestamp expression"); + + int numCols = batch.getSchema().length(); + return batch + .withNewColumn( + numCols, new StructField("_change_type", StringType.STRING, false), changeTypeVector) + .withNewColumn( + numCols + 1, + new StructField("_commit_version", LongType.LONG, false), + commitVersionVector) + .withNewColumn( + numCols + 2, + new StructField("_commit_timestamp", TimestampType.TIMESTAMP, false), + commitTimestampVector); + } + + @SuppressWarnings("nullness") + private static io.delta.kernel.data.Row generateScanFileRow( + String path, Map partitionValues) { + StructType addFileSchema = + (StructType) InternalScanFileUtils.SCAN_FILE_SCHEMA.get("add").getDataType(); + MapValue partMapValue = VectorUtils.stringStringMapValue(partitionValues); + + Map addFileMap = new HashMap<>(); + addFileMap.put(addFileSchema.indexOf("path"), path); + addFileMap.put(addFileSchema.indexOf("partitionValues"), partMapValue); + addFileMap.put(addFileSchema.indexOf("size"), 0L); + addFileMap.put(addFileSchema.indexOf("modificationTime"), 0L); + addFileMap.put(addFileSchema.indexOf("dataChange"), true); + addFileMap.put(addFileSchema.indexOf("deletionVector"), null); + + io.delta.kernel.data.Row addFile = new GenericRow(addFileSchema, addFileMap); + + StructType scanFileSchema = InternalScanFileUtils.SCAN_FILE_SCHEMA; + Map scanFileMap = new HashMap<>(); + scanFileMap.put(scanFileSchema.indexOf("add"), addFile); + scanFileMap.put(scanFileSchema.indexOf("tableRoot"), "/"); + + return new GenericRow(scanFileSchema, scanFileMap); + } +} diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java index c511a7380dc1..8250d99fd425 100644 --- a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java +++ b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java @@ -18,9 +18,11 @@ package org.apache.beam.sdk.io.delta; import com.google.auto.value.AutoValue; +import io.delta.kernel.Snapshot; import io.delta.kernel.Table; import io.delta.kernel.defaults.engine.DefaultEngine; import io.delta.kernel.engine.Engine; +import io.delta.kernel.internal.TableImpl; import io.delta.kernel.types.ArrayType; import io.delta.kernel.types.BinaryType; import io.delta.kernel.types.BooleanType; @@ -60,6 +62,10 @@ public static ReadRows readRows() { return new AutoValue_DeltaIO_ReadRows.Builder().build(); } + public static ReadChanges readChanges() { + return new AutoValue_DeltaIO_ReadChanges.Builder().build(); + } + @AutoValue public abstract static class ReadRows extends PTransform> { @@ -174,4 +180,129 @@ static Schema.FieldType convertToBeamFieldType(DataType deltaType) { } } } + + @AutoValue + public abstract static class ReadChanges extends PTransform> { + 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 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 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 config) { + return toBuilder().setHadoopConfig(config).build(); + } + + @Override + public PCollection 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."); + } + + Configuration conf = new Configuration(); + Map hadoopConfig = getHadoopConfig(); + if (hadoopConfig != null) { + for (Map.Entry entry : hadoopConfig.entrySet()) { + conf.set(entry.getKey(), entry.getValue()); + } + } + Engine engine = DefaultEngine.create(conf); + Table table = Table.forPath(engine, path); + + TableImpl tableImpl = (TableImpl) table; + + long resolvedEndVersion; + Long endVersionVal = getEndVersion(); + String endTimestampVal = getEndTimestamp(); + if (endVersionVal != null) { + resolvedEndVersion = endVersionVal; + } else if (endTimestampVal != null) { + long endMillis = java.time.Instant.parse(endTimestampVal).toEpochMilli(); + resolvedEndVersion = tableImpl.getVersionBeforeOrAtTimestamp(engine, endMillis); + } else { + resolvedEndVersion = table.getLatestSnapshot(engine).getVersion(); + } + + Snapshot endSnapshot = table.getSnapshotAsOfVersion(engine, resolvedEndVersion); + StructType deltaSchema = endSnapshot.getSchema(); + if (deltaSchema == null) { + throw new IllegalStateException("Table schema is null."); + } + Schema beamSchema = + ReadRows.convertToBeamSchema( + deltaSchema + .add("_change_type", StringType.STRING, false) + .add("_commit_version", LongType.LONG, false) + .add("_commit_timestamp", TimestampType.TIMESTAMP, false)); + + return input + .apply("Create Path", Create.of(path)) + .apply( + "Plan CDF Files", + ParDo.of( + new CreateCDCReadTasksDoFn( + hadoopConfig, + getStartVersion(), + getStartTimestamp(), + getEndVersion(), + getEndTimestamp()))) + .apply("Read CDF Data", ParDo.of(new DeltaCDCSourceDoFn(hadoopConfig))) + .setRowSchema(beamSchema); + } + } } diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaSourceDoFn.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaSourceDoFn.java index bd53c3c9d045..fca19e30cb8f 100644 --- a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaSourceDoFn.java +++ b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaSourceDoFn.java @@ -213,7 +213,7 @@ public ColumnarBatch next() { } // Convert Delta `Row` to Beam `Row`. - private static Row toBeamRow(io.delta.kernel.data.Row deltaRow, Schema beamSchema) { + static Row toBeamRow(io.delta.kernel.data.Row deltaRow, Schema beamSchema) { Row.Builder builder = Row.withSchema(beamSchema); StructType deltaSchema = deltaRow.getSchema(); List fields = deltaSchema.fields(); diff --git a/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java b/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java index bd8bf8b3c8cc..ff8cc6096e11 100644 --- a/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java +++ b/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java @@ -48,12 +48,15 @@ import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Count; import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.joda.time.Instant; import org.junit.Assert; import org.junit.Rule; import org.junit.Test; @@ -810,4 +813,153 @@ 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)); + + File commitFile0 = new File(logDir, "00000000000000000000.json"); + String commitContent0 = + "{\"protocol\":{\"minReaderVersion\":1,\"minWriterVersion\":2}}\n" + + "{\"metaData\":{\"id\":\"test-id\",\"format\":{\"provider\":\"parquet\",\"options\":{}},\"schemaString\":\"{\\\"type\\\":\\\"struct\\\",\\\"fields\\\":[{\\\"name\\\":\\\"name\\\",\\\"type\\\":\\\"string\\\",\\\"nullable\\\":true,\\\"metadata\\\":{}}]}\",\"partitionColumns\":[],\"configuration\":{},\"createdAt\":123456789}}\n" + + "{\"add\":{\"path\":\"part-00000.parquet\",\"partitionValues\":{},\"size\":" + + partBytes.length + + ",\"modificationTime\":100000000000,\"dataChange\":true}}"; + Files.write(commitFile0.toPath(), commitContent0.getBytes(StandardCharsets.UTF_8)); + commitFile0.setLastModified(100000000000L); + + // 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)); + + File commitFile1 = new File(logDir, "00000000000000000001.json"); + String commitContent1 = + "{\"cdc\":{\"path\":\"change-00000.parquet\",\"partitionValues\":{},\"size\":" + + changeBytes.length + + ",\"dataChange\":true}}"; + Files.write(commitFile1.toPath(), commitContent1.getBytes(StandardCharsets.UTF_8)); + commitFile1.setLastModified(200000000000L); + + // 3. Read CDF data from table using ReadChanges + PCollection output = + readPipeline.apply( + DeltaIO.readChanges().from(tableDir.getAbsolutePath()).withStartVersion(0L)); + + // Construct expected rows + Schema expectedSchema = + 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 expRow1 = + Row.withSchema(expectedSchema) + .addValues("row-1", "insert", 0L, new Instant(100000000000L)) + .build(); + Row expRow2 = + Row.withSchema(expectedSchema) + .addValues("row-2", "insert", 0L, new Instant(100000000000L)) + .build(); + Row expRow3 = + Row.withSchema(expectedSchema) + .addValues("row-1", "update_preimage", 1L, new Instant(123456789000L)) + .build(); + Row expRow4 = + Row.withSchema(expectedSchema) + .addValues("row-1-updated", "update_postimage", 1L, new Instant(123456789000L)) + .build(); + Row expRow5 = + Row.withSchema(expectedSchema) + .addValues("row-2", "delete", 1L, new Instant(123456789000L)) + .build(); + + PCollection formattedOutput = + output.apply("Format ValueKind and Row", ParDo.of(new FormatValueKindAndRow())); + + PAssert.that(formattedOutput) + .containsInAnyOrder( + "INSERT:row-1:insert:0", + "INSERT:row-2:insert:0", + "UPDATE_BEFORE:row-1:update_preimage:1", + "UPDATE_AFTER:row-1-updated:update_postimage:1", + "DELETE:row-2:delete:1"); + + readPipeline.run().waitUntilFinish(); + } + + private static final class FormatValueKindAndRow extends DoFn { + @ProcessElement + public void process( + @Element Row row, ValueKind valueKind, OutputReceiver outputReceiver) { + outputReceiver.output( + valueKind.name() + + ":" + + row.getString("name") + + ":" + + row.getString("_change_type") + + ":" + + row.getInt64("_commit_version")); + } + } + + private byte[] writeParquetFile(File file, Schema schema, java.util.List rows) + throws Exception { + org.apache.avro.Schema avroSchema = + org.apache.beam.sdk.extensions.avro.schemas.utils.AvroUtils.toAvroSchema(schema); + org.apache.avro.generic.GenericRecord[] records = + new org.apache.avro.generic.GenericRecord[rows.size()]; + for (int i = 0; i < rows.size(); i++) { + records[i] = + org.apache.beam.sdk.extensions.avro.schemas.utils.AvroUtils.toGenericRecord( + rows.get(i), avroSchema); + } + org.apache.hadoop.fs.Path path = new org.apache.hadoop.fs.Path(file.getAbsolutePath()); + try (org.apache.parquet.hadoop.ParquetWriter writer = + org.apache.parquet.avro.AvroParquetWriter.builder( + path) + .withSchema(avroSchema) + .withConf(new org.apache.hadoop.conf.Configuration()) + .build()) { + for (org.apache.avro.generic.GenericRecord record : records) { + writer.write(record); + } + } + return java.nio.file.Files.readAllBytes(file.toPath()); + } } From ed28244170ddd9481345a8b859bb7aee5125cf67 Mon Sep 17 00:00:00 2001 From: Chamikara Jayalath Date: Wed, 22 Jul 2026 16:26:21 -0700 Subject: [PATCH 2/2] Update and add a test --- .../sdk/io/delta/CreateCDCReadTasksDoFn.java | 65 +++--- .../beam/sdk/io/delta/DeltaCDCSourceDoFn.java | 14 +- .../org/apache/beam/sdk/io/delta/DeltaIO.java | 7 +- .../apache/beam/sdk/io/delta/DeltaIOTest.java | 216 +++++++++++++----- 4 files changed, 214 insertions(+), 88 deletions(-) diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java index 16843b91906f..e35a075adb9a 100644 --- a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java +++ b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java @@ -45,9 +45,7 @@ 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. - */ +/** A DoFn that reads the Delta log and plans read tasks for Change Data Feed. */ class CreateCDCReadTasksDoFn extends DoFn { private static final long MAX_TASK_SIZE_BYTES = 1024L * 1024L * 1024L; // 1 GB private final @Nullable Map hadoopConfig; @@ -122,8 +120,9 @@ public void processElement(@Element String tablePath, OutputReceiver batchIter = range.getActions(engine, startSnapshot, actionSet)) { + try (CloseableIterator batchIter = + range.getActions(engine, startSnapshot, actionSet)) { Map commitActionsMap = new HashMap<>(); while (batchIter.hasNext()) { @@ -149,21 +149,26 @@ public void processElement(@Element String tablePath, OutputReceiver new CommitActionsInfo(version, timestamp)); + 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); + 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); + 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()) { @@ -199,8 +204,9 @@ public void processElement(@Element String tablePath, OutputReceiver rowGroupSizes = getRowGroupSizes(fullPath, conf); - DeltaCDCReadTask task = new DeltaCDCReadTask( - fullPath, - size, - partitionValues, - info.version, - info.timestamp, - isCDC, - rowGroupSizes, - serializableScanState); + DeltaCDCReadTask task = + new DeltaCDCReadTask( + fullPath, + size, + partitionValues, + info.version, + info.timestamp, + isCDC, + rowGroupSizes, + serializableScanState); if (size >= MAX_TASK_SIZE_BYTES) { if (!currentGroup.isEmpty()) { @@ -258,8 +265,8 @@ private List getRowGroupSizes(String pathStr, Configuration conf) { List 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( + org.apache.parquet.hadoop.metadata.ParquetMetadata metadata = + org.apache.parquet.hadoop.ParquetFileReader.readFooter( conf, hadoopPath, org.apache.parquet.format.converter.ParquetMetadataConverter.NO_FILTER); diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java index 41cfd5d2e716..81a44f28b290 100644 --- a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java +++ b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java @@ -117,6 +117,7 @@ public ProcessContinuation processElement( SerializableRow originalScanStateRow = task.getScanStateRow(); StructType logicalTableSchema = ScanStateRow.getLogicalSchema(originalScanStateRow); + Schema publicBeamSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema); StructType physicalTableSchema = ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow); StructType scanStateSchema = originalScanStateRow.getSchema(); @@ -224,7 +225,8 @@ public ColumnarBatch next() { throw new IllegalStateException("Field _change_type must not be null."); } ValueKind kind = getValueKind(changeType); - out.builder(beamRow).setValueKind(kind).output(); + Row publicRow = projectRow(beamRow, publicBeamSchema); + out.builder(publicRow).setValueKind(kind).output(); } } } @@ -234,7 +236,17 @@ public ColumnarBatch next() { 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; diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java index 8250d99fd425..f2e03bb701b3 100644 --- a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java +++ b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java @@ -283,12 +283,7 @@ public PCollection expand(PBegin input) { if (deltaSchema == null) { throw new IllegalStateException("Table schema is null."); } - Schema beamSchema = - ReadRows.convertToBeamSchema( - deltaSchema - .add("_change_type", StringType.STRING, false) - .add("_commit_version", LongType.LONG, false) - .add("_commit_timestamp", TimestampType.TIMESTAMP, false)); + Schema beamSchema = ReadRows.convertToBeamSchema(deltaSchema); return input .apply("Create Path", Create.of(path)) diff --git a/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java b/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java index ff8cc6096e11..8496fe645e0f 100644 --- a/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java +++ b/sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java @@ -31,6 +31,7 @@ import io.delta.kernel.types.StructType; import io.delta.kernel.types.TimestampType; import java.io.File; +import java.io.IOException; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.util.HashMap; @@ -56,6 +57,7 @@ import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.ValueKind; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.junit.Assert; import org.junit.Rule; @@ -829,15 +831,8 @@ public void testReadChanges() throws Exception { byte[] partBytes = writeParquetFile(partFile, tableSchema, java.util.Arrays.asList(tableRow1, tableRow2)); - File commitFile0 = new File(logDir, "00000000000000000000.json"); - String commitContent0 = - "{\"protocol\":{\"minReaderVersion\":1,\"minWriterVersion\":2}}\n" - + "{\"metaData\":{\"id\":\"test-id\",\"format\":{\"provider\":\"parquet\",\"options\":{}},\"schemaString\":\"{\\\"type\\\":\\\"struct\\\",\\\"fields\\\":[{\\\"name\\\":\\\"name\\\",\\\"type\\\":\\\"string\\\",\\\"nullable\\\":true,\\\"metadata\\\":{}}]}\",\"partitionColumns\":[],\"configuration\":{},\"createdAt\":123456789}}\n" - + "{\"add\":{\"path\":\"part-00000.parquet\",\"partitionValues\":{},\"size\":" - + partBytes.length - + ",\"modificationTime\":100000000000,\"dataChange\":true}}"; - Files.write(commitFile0.toPath(), commitContent0.getBytes(StandardCharsets.UTF_8)); - commitFile0.setLastModified(100000000000L); + 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 = @@ -866,75 +861,192 @@ public void testReadChanges() throws Exception { writeParquetFile( changeFile, cdcWriteSchema, java.util.Arrays.asList(cdcRow1, cdcRow2, cdcRow3)); - File commitFile1 = new File(logDir, "00000000000000000001.json"); - String commitContent1 = - "{\"cdc\":{\"path\":\"change-00000.parquet\",\"partitionValues\":{},\"size\":" - + changeBytes.length - + ",\"dataChange\":true}}"; - Files.write(commitFile1.toPath(), commitContent1.getBytes(StandardCharsets.UTF_8)); - commitFile1.setLastModified(200000000000L); + writeCommit( + logDir, + 1L, + 200000000000L, + null, + 0L, + null, + "change-00000.parquet", + changeBytes.length, + false); // 3. Read CDF data from table using ReadChanges PCollection output = readPipeline.apply( DeltaIO.readChanges().from(tableDir.getAbsolutePath()).withStartVersion(0L)); - // Construct expected rows - Schema expectedSchema = + PCollection 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 expRow1 = - Row.withSchema(expectedSchema) - .addValues("row-1", "insert", 0L, new Instant(100000000000L)) - .build(); - Row expRow2 = - Row.withSchema(expectedSchema) - .addValues("row-2", "insert", 0L, new Instant(100000000000L)) - .build(); - Row expRow3 = - Row.withSchema(expectedSchema) - .addValues("row-1", "update_preimage", 1L, new Instant(123456789000L)) + Row cdcRow1 = + Row.withSchema(cdcWriteSchema) + .addValues("row-1", "update_preimage", 1L, new Instant(200000000000L)) .build(); - Row expRow4 = - Row.withSchema(expectedSchema) - .addValues("row-1-updated", "update_postimage", 1L, new Instant(123456789000L)) + Row cdcRow2 = + Row.withSchema(cdcWriteSchema) + .addValues("row-1-updated", "update_postimage", 1L, new Instant(200000000000L)) .build(); - Row expRow5 = - Row.withSchema(expectedSchema) - .addValues("row-2", "delete", 1L, new Instant(123456789000L)) + 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 outputVersions = + readPipeline.apply( + "Read Changes Version Range", + DeltaIO.readChanges() + .from(tableDir.getAbsolutePath()) + .withStartVersion(0L) + .withEndVersion(2L)); - PCollection formattedOutput = - output.apply("Format ValueKind and Row", ParDo.of(new FormatValueKindAndRow())); + PCollection formattedVersions = + outputVersions.apply("Format Version Output", ParDo.of(new FormatValueKindAndRow())); - PAssert.that(formattedOutput) + 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 outputTimestamps = + filteringPipeline.apply( + "Read Changes Timestamp Range", + DeltaIO.readChanges() + .from(tableDir.getAbsolutePath()) + .withStartTimestamp(startTimestamp) + .withEndTimestamp(endTimestamp)); + + PCollection formattedTimestamps = + outputTimestamps.apply("Format Timestamp Output", ParDo.of(new FormatValueKindAndRow())); + + PAssert.that(formattedTimestamps) .containsInAnyOrder( - "INSERT:row-1:insert:0", - "INSERT:row-2:insert:0", - "UPDATE_BEFORE:row-1:update_preimage:1", - "UPDATE_AFTER:row-1-updated:update_postimage:1", - "DELETE:row-2:delete:1"); + "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 { @ProcessElement public void process( @Element Row row, ValueKind valueKind, OutputReceiver outputReceiver) { - outputReceiver.output( - valueKind.name() - + ":" - + row.getString("name") - + ":" - + row.getString("_change_type") - + ":" - + row.getInt64("_commit_version")); + outputReceiver.output(valueKind.name() + ":" + row.getString("name")); } }