From 0c74b47c26b1cbf292aafd89c6d6d19745f13859 Mon Sep 17 00:00:00 2001 From: Chamikara Jayalath Date: Tue, 21 Jul 2026 22:52:59 -0700 Subject: [PATCH 1/3] 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 d9c928ef6590..e11b74f3b790 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; @@ -88,6 +90,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> { @@ -211,4 +217,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 4b75dc0d8bbabf5e30c07b808ee6c0c7db0ca6fa Mon Sep 17 00:00:00 2001 From: Chamikara Jayalath Date: Wed, 22 Jul 2026 16:26:21 -0700 Subject: [PATCH 2/3] 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 e11b74f3b790..14ff9f54e2db 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 @@ -320,12 +320,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")); } } From b85ccc290814c5d27805c1620d7b9439714f2d74 Mon Sep 17 00:00:00 2001 From: Chamikara Jayalath Date: Tue, 28 Jul 2026 22:20:06 -0700 Subject: [PATCH 3/3] Address reviewer comments --- .../sdk/io/delta/CreateCDCReadTasksDoFn.java | 21 +- .../beam/sdk/io/delta/DeltaCDCReadTask.java | 23 +- .../beam/sdk/io/delta/DeltaCDCSourceDoFn.java | 36 +- .../org/apache/beam/sdk/io/delta/DeltaIO.java | 24 +- .../DeltaReadSchemaTransformProvider.java | 6 +- .../apache/beam/sdk/io/delta/DeltaIOTest.java | 803 ++++++++++++++---- 6 files changed, 671 insertions(+), 242 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 e35a075adb9a..9a0049a8651e 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 @@ -160,7 +160,7 @@ public void processElement(@Element String tablePath, OutputReceiver= 0 && !batch.getColumnVector(addIdx).isNullAt(i)) { Row addRow = @@ -172,7 +172,7 @@ public void processElement(@Element String tablePath, OutputReceiver currentGroup = new ArrayList<>(); long currentGroupSize = 0L; + // TODO: to prevent OOMs in true streaming executions, update DeltaReadTask to include a group + // of files. + // Sort versions to process them in order List versions = new ArrayList<>(commitActionsMap.keySet()); Collections.sort(versions); @@ -191,17 +194,16 @@ public void processElement(@Element String tablePath, OutputReceiver rowsToProcess = hasCDC ? info.cdcRows : info.addRows; - boolean isCDC = hasCDC; + List rowsToProcess = hasCDC ? info.cdcInfo : info.insertInfo; for (Row fileRow : rowsToProcess) { String relPath; long size; Map partitionValues; - if (isCDC) { + if (hasCDC) { relPath = fileRow.getString(AddCDCFile.FULL_SCHEMA.indexOf("path")); size = fileRow.getLong(AddCDCFile.FULL_SCHEMA.indexOf("size")); partitionValues = @@ -224,7 +226,7 @@ public void processElement(@Element String tablePath, OutputReceiver getRowGroupSizes(String pathStr, Configuration conf) { private static class CommitActionsInfo { final long version; + final long timestamp; - final List cdcRows = new ArrayList<>(); - final List addRows = new ArrayList<>(); + final List cdcInfo = new ArrayList<>(); + final List insertInfo = new ArrayList<>(); CommitActionsInfo(long version, long timestamp) { this.version = version; 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 index f2adf354dc09..24c594343def 100644 --- 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 @@ -117,24 +117,9 @@ public int hashCode() { @Override public String toString() { - return "DeltaCDCReadTask{" - + "path='" - + path - + '\'' - + ", size=" - + size - + ", partitionValues=" - + partitionValues - + ", version=" - + version - + ", timestamp=" - + timestamp - + ", isCDC=" - + isCDC - + ", rowGroupSizes=" - + rowGroupSizes - + ", scanStateRow=" - + scanStateRow - + '}'; + return String.format( + "DeltaCDCReadTask{path='%s', size=%d, partitionValues=%s, version=%d, timestamp=%d, isCDC=%b, " + + "rowGroupSizes=%s, scanStateRow=%s}", + path, size, partitionValues, version, timestamp, isCDC, rowGroupSizes, 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 index 81a44f28b290..cf10adc9865c 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 @@ -104,7 +104,7 @@ public void setUp() { } @ProcessElement - public ProcessContinuation processElement( + public void processElement( @Element DeltaCDCReadTask task, RestrictionTracker tracker, OutputReceiver out) @@ -205,24 +205,26 @@ public ColumnarBatch next() { 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 = + ColumnarBatch logicalBatch = appendConstantCDFColumns( - currentEngine, logicalBatch, task.getVersion(), task.getTimestamp()); + currentEngine, batch.getData(), task.getVersion(), task.getTimestamp()); + // Make sure we use selection vector to considered filtered out or deleted rows. + batch = new FilteredColumnarBatch(logicalBatch, batch.getSelectionVector()); } - try (CloseableIterator logicalRows = logicalBatch.getRows()) { + try (CloseableIterator logicalRows = batch.getRows()) { while (logicalRows.hasNext()) { io.delta.kernel.data.Row deltaRow = logicalRows.next(); Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema); - String changeType = beamRow.getString("_change_type"); + String changeType = beamRow.getString(DeltaIO.CHANGE_TYPE_COLUMN); if (changeType == null) { - throw new IllegalStateException("Field _change_type must not be null."); + throw new IllegalStateException( + "Field " + DeltaIO.CHANGE_TYPE_COLUMN + " must not be null."); } ValueKind kind = getValueKind(changeType); Row publicRow = projectRow(beamRow, publicBeamSchema); @@ -232,11 +234,13 @@ public ColumnarBatch next() { } } } - - return ProcessContinuation.stop(); } private static Row projectRow(Row row, Schema targetSchema) { + if (row.getSchema().equals(targetSchema)) { + // We can return the original Row since schemas are the same. + return row; + } Row.Builder builder = Row.withSchema(targetSchema); for (Schema.Field field : targetSchema.getFields()) { builder.addValue(row.getValue(field.getName())); @@ -263,9 +267,9 @@ private static ValueKind getValueKind(String 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); + .add(DeltaIO.CHANGE_TYPE_COLUMN, StringType.STRING, false) + .add(DeltaIO.COMMIT_VERSION_COLUMN, LongType.LONG, false) + .add(DeltaIO.COMMIT_TIMESTAMP_COLUMN, TimestampType.TIMESTAMP, false); } private ColumnarBatch appendConstantCDFColumns( @@ -315,14 +319,16 @@ private ColumnarBatch appendConstantCDFColumns( int numCols = batch.getSchema().length(); return batch .withNewColumn( - numCols, new StructField("_change_type", StringType.STRING, false), changeTypeVector) + numCols, + new StructField(DeltaIO.CHANGE_TYPE_COLUMN, StringType.STRING, false), + changeTypeVector) .withNewColumn( numCols + 1, - new StructField("_commit_version", LongType.LONG, false), + new StructField(DeltaIO.COMMIT_VERSION_COLUMN, LongType.LONG, false), commitVersionVector) .withNewColumn( numCols + 2, - new StructField("_commit_timestamp", TimestampType.TIMESTAMP, false), + new StructField(DeltaIO.COMMIT_TIMESTAMP_COLUMN, TimestampType.TIMESTAMP, false), commitTimestampVector); } 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 14ff9f54e2db..3ac2c7a84a8f 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 @@ -52,23 +52,6 @@ /** * A connector that reads from Delta Lake tables. * - *

{@link DeltaIO} is offered as a Managed transform. This class is subject to change and should - * not be used directly. Instead, use it like so: - * - * 

{@code
- * Map config = Map.of(
- *         "table", "gs://my-bucket/delta-table",
- *         "hadoop_config", Map.of(
- *                 "fs.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem",
- *                 "fs.AbstractFileSystem.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFS",
- *                 "fs.gs.project.id", "my-project-id"));
- *
- * pipeline
- *     .apply(Managed.read(Managed.DELTA_LAKE).withConfig(config))
- *     .getSinglePCollection()
- *     .apply(ParDo.of(...));
- * }
- * *

Configuration Options

* * Please check the Managed IO @@ -80,6 +63,10 @@ @Internal public class DeltaIO { + public static final String CHANGE_TYPE_COLUMN = "_change_type"; + public static final String COMMIT_VERSION_COLUMN = "_commit_version"; + public static final String COMMIT_TIMESTAMP_COLUMN = "_commit_timestamp"; + /** * Reads rows from a Delta Lake table. * @@ -90,6 +77,7 @@ public static ReadRows readRows() { return new AutoValue_DeltaIO_ReadRows.Builder().build(); } + /** Reads change data feed (CDC) from a Delta Lake table. */ public static ReadChanges readChanges() { return new AutoValue_DeltaIO_ReadChanges.Builder().build(); } @@ -282,6 +270,8 @@ public PCollection expand(PBegin input) { throw new IllegalArgumentException("Table path must be set."); } if (getStartVersion() == null && getStartTimestamp() == null) { + // TODO: for unbounded reads, support using current HEAD or the latest snapshot + // as the default starting point. throw new IllegalArgumentException("Either startVersion or startTimestamp must be set."); } if (getStartVersion() != null && getStartTimestamp() != null) { diff --git a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaReadSchemaTransformProvider.java b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaReadSchemaTransformProvider.java index 42ca3f24def9..48dc3a2c7488 100644 --- a/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaReadSchemaTransformProvider.java +++ b/sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaReadSchemaTransformProvider.java @@ -130,11 +130,11 @@ static Builder builder() { abstract static class Builder { abstract Builder setTable(String table); - abstract Builder setVersion(@Nullable Long version); + abstract Builder setVersion(Long version); - abstract Builder setTimestamp(@Nullable String timestamp); + abstract Builder setTimestamp(String timestamp); - abstract Builder setHadoopConfig(@Nullable Map hadoopConfig); + abstract Builder setHadoopConfig(Map hadoopConfig); abstract Configuration build(); } 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 8496fe645e0f..f00b34be4609 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 @@ -17,9 +17,23 @@ */ package org.apache.beam.sdk.io.delta; +import io.delta.kernel.DataWriteContext; +import io.delta.kernel.Operation; +import io.delta.kernel.Table; +import io.delta.kernel.Transaction; +import io.delta.kernel.TransactionBuilder; +import io.delta.kernel.TransactionCommitResult; +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.defaults.internal.data.DefaultColumnarBatch; +import io.delta.kernel.engine.Engine; import io.delta.kernel.types.ArrayType; import io.delta.kernel.types.BinaryType; import io.delta.kernel.types.BooleanType; +import io.delta.kernel.types.DataType; import io.delta.kernel.types.DateType; import io.delta.kernel.types.DoubleType; import io.delta.kernel.types.FloatType; @@ -30,12 +44,19 @@ import io.delta.kernel.types.StructField; import io.delta.kernel.types.StructType; import io.delta.kernel.types.TimestampType; +import io.delta.kernel.utils.CloseableIterable; +import io.delta.kernel.utils.CloseableIterator; +import io.delta.kernel.utils.DataFileStatus; import java.io.File; -import java.io.IOException; +import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; +import java.util.Optional; import org.apache.avro.generic.GenericRecord; import org.apache.beam.sdk.extensions.avro.coders.AvroCoder; import org.apache.beam.sdk.extensions.avro.schemas.utils.AvroUtils; @@ -127,9 +148,6 @@ public void testPrintScanStateSchema() throws Exception { io.delta.kernel.Snapshot snapshot = table.getLatestSnapshot(engine); io.delta.kernel.Scan scan = snapshot.getScanBuilder().build(); - io.delta.kernel.data.Row scanState = scan.getScanState(engine); - System.err.println("SCAN STATE SCHEMA: " + scanState.getSchema().toString()); - try (io.delta.kernel.utils.CloseableIterator scanFiles = scan.getScanFiles(engine)) { while (scanFiles.hasNext()) { @@ -305,53 +323,8 @@ public void testFullPipelineRead() throws Exception { writePipeline.run().waitUntilFinish(); - System.out.println("FILES IN TABLE DIR:"); - for (File f : tableDir.listFiles()) { - System.out.println( - " - " + f.getName() + " (size=" + f.length() + ", isDir=" + f.isDirectory() + ")"); - if (f.isDirectory()) { - for (File sub : f.listFiles()) { - System.out.println(" - " + sub.getName() + " (size=" + sub.length() + ")"); - } - } - } - File parquetFile = new File(tableDir, "part-00000.parquet"); byte[] fileBytes = Files.readAllBytes(parquetFile.toPath()); - System.out.println("PARQUET FILE LENGTH: " + fileBytes.length); - if (fileBytes.length >= 8) { - System.out.println( - "PARQUET FIRST 4 BYTES: " - + fileBytes[0] - + ", " - + fileBytes[1] - + ", " - + fileBytes[2] - + ", " - + fileBytes[3] - + " ('" - + (char) fileBytes[0] - + (char) fileBytes[1] - + (char) fileBytes[2] - + (char) fileBytes[3] - + "')"); - int len = fileBytes.length; - System.out.println( - "PARQUET LAST 4 BYTES: " - + fileBytes[len - 4] - + ", " - + fileBytes[len - 3] - + ", " - + fileBytes[len - 2] - + ", " - + fileBytes[len - 1] - + " ('" - + (char) fileBytes[len - 4] - + (char) fileBytes[len - 3] - + (char) fileBytes[len - 2] - + (char) fileBytes[len - 1] - + "')"); - } // 2. Create the Delta log File logDir = new File(tableDir, "_delta_log"); @@ -397,28 +370,19 @@ private byte[] writeParquetFile(File file, Row row) throws Exception { @Test public void testManagedDeltaRead() throws Exception { File tableDir = tempFolder.newFolder("managed-delta-table"); + Engine engine = DefaultEngine.create(new org.apache.hadoop.conf.Configuration()); - // 1. Write a Parquet file to simulate a Delta table Schema schema = Schema.builder().addField("name", Schema.FieldType.STRING).build(); Row row = Row.withSchema(schema).addValues("test-name").build(); - writeParquetFile(new File(tableDir, "part-00000.parquet"), row); - - // 2. Create the Delta log - File logDir = new File(tableDir, "_delta_log"); - logDir.mkdirs(); - File commitFile = new File(logDir, "00000000000000000000.json"); - - File parquetFile = new File(tableDir, "part-00000.parquet"); - byte[] fileBytes = Files.readAllBytes(parquetFile.toPath()); + StructType deltaSchema = new StructType().add("name", StringType.STRING); - String commitContent = - "{\"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\":" - + fileBytes.length - + ",\"modificationTime\":123456789,\"dataChange\":true}}"; - - Files.write(commitFile.toPath(), commitContent.getBytes(StandardCharsets.UTF_8)); + writeAppendCommit( + engine, + tableDir.getAbsolutePath(), + 0L, + 123456789L, + deltaSchema, + Collections.singletonList(row)); // 3. Read it using Managed PCollection output = @@ -819,29 +783,36 @@ 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(); + Engine engine = DefaultEngine.create(new org.apache.hadoop.conf.Configuration()); // 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(); + StructType deltaSchema = new StructType().add("name", StringType.STRING); - 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); + writeAppendCommit( + engine, + tableDir.getAbsolutePath(), + 0L, + 100000000000L, + deltaSchema, + java.util.Arrays.asList(tableRow1, tableRow2)); // 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) + .addField(DeltaIO.CHANGE_TYPE_COLUMN, Schema.FieldType.STRING) + .addField(DeltaIO.COMMIT_VERSION_COLUMN, Schema.FieldType.INT64) + .addField(DeltaIO.COMMIT_TIMESTAMP_COLUMN, Schema.FieldType.DATETIME) .build(); + StructType cdcWriteDeltaSchema = + new StructType() + .add("name", StringType.STRING) + .add(DeltaIO.CHANGE_TYPE_COLUMN, StringType.STRING) + .add(DeltaIO.COMMIT_VERSION_COLUMN, LongType.LONG) + .add(DeltaIO.COMMIT_TIMESTAMP_COLUMN, TimestampType.TIMESTAMP); Row cdcRow1 = Row.withSchema(cdcWriteSchema) @@ -856,21 +827,16 @@ public void testReadChanges() throws Exception { .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, + writeCdcCommit( + engine, + tableDir.getAbsolutePath(), 1L, 200000000000L, + deltaSchema, null, - 0L, null, - "change-00000.parquet", - changeBytes.length, - false); + java.util.Arrays.asList(cdcRow1, cdcRow2, cdcRow3), + cdcWriteDeltaSchema); // 3. Read CDF data from table using ReadChanges PCollection output = @@ -894,28 +860,37 @@ public void testReadChanges() throws Exception { @Test public void testReadChangesRanges() throws Exception { File tableDir = tempFolder.newFolder("delta-table-changes-ranges"); - File logDir = new File(tableDir, "_delta_log"); - logDir.mkdirs(); + Engine engine = DefaultEngine.create(new org.apache.hadoop.conf.Configuration()); Schema tableSchema = Schema.builder().addField("name", Schema.FieldType.STRING).build(); + StructType deltaSchema = new StructType().add("name", StringType.STRING); // 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); + writeAppendCommit( + engine, + tableDir.getAbsolutePath(), + 0L, + 100000000000L, + deltaSchema, + java.util.Arrays.asList(tableRow1, tableRow2)); // 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) + .addField(DeltaIO.CHANGE_TYPE_COLUMN, Schema.FieldType.STRING) + .addField(DeltaIO.COMMIT_VERSION_COLUMN, Schema.FieldType.INT64) + .addField(DeltaIO.COMMIT_TIMESTAMP_COLUMN, Schema.FieldType.DATETIME) .build(); + StructType cdcWriteDeltaSchema = + new StructType() + .add("name", StringType.STRING) + .add(DeltaIO.CHANGE_TYPE_COLUMN, StringType.STRING) + .add(DeltaIO.COMMIT_VERSION_COLUMN, LongType.LONG) + .add(DeltaIO.COMMIT_TIMESTAMP_COLUMN, TimestampType.TIMESTAMP); + Row cdcRow1 = Row.withSchema(cdcWriteSchema) .addValues("row-1", "update_preimage", 1L, new Instant(200000000000L)) @@ -928,34 +903,27 @@ public void testReadChangesRanges() throws Exception { 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, + + writeCdcCommit( + engine, + tableDir.getAbsolutePath(), 1L, 200000000000L, - "part-00001.parquet", - partBytes1.length, - "part-00000.parquet", - "change-00000.parquet", - changeBytes0.length, - false); + deltaSchema, + null, + null, + java.util.Arrays.asList(cdcRow1, cdcRow2, cdcRow3), + cdcWriteDeltaSchema); // 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); + writeAppendCommit( + engine, + tableDir.getAbsolutePath(), + 2L, + 300000000000L, + deltaSchema, + java.util.Arrays.asList(tableRow3)); // Test 1: Read changes between start version 0 and end version 2 PCollection outputVersions = @@ -1002,44 +970,132 @@ public void testReadChangesRanges() throws Exception { 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); + @Test + public void testReadChangesPartialRange() throws Exception { + File tableDir = tempFolder.newFolder("delta-table-changes-partial-range"); + Engine engine = DefaultEngine.create(new org.apache.hadoop.conf.Configuration()); + + Schema tableSchema = Schema.builder().addField("name", Schema.FieldType.STRING).build(); + StructType deltaSchema = new StructType().add("name", StringType.STRING); + + // 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(); + writeAppendCommit( + engine, + tableDir.getAbsolutePath(), + 0L, + 100000000000L, + deltaSchema, + java.util.Arrays.asList(tableRow1, tableRow2)); + + // 2. Write parquet files for Version 1 (commit with updates and deletes) + Schema cdcWriteSchema = + Schema.builder() + .addField("name", Schema.FieldType.STRING) + .addField(DeltaIO.CHANGE_TYPE_COLUMN, Schema.FieldType.STRING) + .addField(DeltaIO.COMMIT_VERSION_COLUMN, Schema.FieldType.INT64) + .addField(DeltaIO.COMMIT_TIMESTAMP_COLUMN, Schema.FieldType.DATETIME) + .build(); + StructType cdcWriteDeltaSchema = + new StructType() + .add("name", StringType.STRING) + .add(DeltaIO.CHANGE_TYPE_COLUMN, StringType.STRING) + .add(DeltaIO.COMMIT_VERSION_COLUMN, LongType.LONG) + .add(DeltaIO.COMMIT_TIMESTAMP_COLUMN, TimestampType.TIMESTAMP); + + Row cdcRow1 = + Row.withSchema(cdcWriteSchema) + .addValues("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(); + + writeCdcCommit( + engine, + tableDir.getAbsolutePath(), + 1L, + 200000000000L, + deltaSchema, + null, + null, + java.util.Arrays.asList(cdcRow1, cdcRow2, cdcRow3), + cdcWriteDeltaSchema); + + // 3. Write parquet files for Version 2 (insert-only commit) + Row tableRow3 = Row.withSchema(tableSchema).addValues("row-3").build(); + writeAppendCommit( + engine, + tableDir.getAbsolutePath(), + 2L, + 300000000000L, + deltaSchema, + java.util.Arrays.asList(tableRow3)); + + // 4. Write parquet files for Version 3 (commit with updates and deletes) + Row cdcRow4 = + Row.withSchema(cdcWriteSchema) + .addValues("row-3", "update_preimage", 3L, new Instant(400000000000L)) + .build(); + Row cdcRow5 = + Row.withSchema(cdcWriteSchema) + .addValues("row-3-updated", "update_postimage", 3L, new Instant(400000000000L)) + .build(); + Row cdcRow6 = + Row.withSchema(cdcWriteSchema) + .addValues("row-1-updated", "delete", 3L, new Instant(400000000000L)) + .build(); + + writeCdcCommit( + engine, + tableDir.getAbsolutePath(), + 3L, + 400000000000L, + deltaSchema, + null, + null, + java.util.Arrays.asList(cdcRow4, cdcRow5, cdcRow6), + cdcWriteDeltaSchema); + + // 5. Write parquet files for Version 4 (insert-only commit) + Row tableRow4 = Row.withSchema(tableSchema).addValues("row-4").build(); + writeAppendCommit( + engine, + tableDir.getAbsolutePath(), + 4L, + 500000000000L, + deltaSchema, + java.util.Arrays.asList(tableRow4)); + + // Read changes between start version 1 and end version 3 + PCollection outputVersions = + readPipeline.apply( + "Read Changes Partial Version Range", + DeltaIO.readChanges() + .from(tableDir.getAbsolutePath()) + .withStartVersion(1L) + .withEndVersion(3L)); + + PCollection formattedVersions = + outputVersions.apply("Format Version Output", ParDo.of(new FormatValueKindAndRow())); + + PAssert.that(formattedVersions) + .containsInAnyOrder( + "UPDATE_BEFORE:row-1", + "UPDATE_AFTER:row-1-updated", + "DELETE:row-2", + "INSERT:row-3", + "UPDATE_BEFORE:row-3", + "UPDATE_AFTER:row-3-updated", + "DELETE:row-1-updated"); + + readPipeline.run().waitUntilFinish(); } private static final class FormatValueKindAndRow extends DoFn { @@ -1050,28 +1106,417 @@ public void process( } } - private byte[] writeParquetFile(File file, Schema schema, java.util.List rows) + private List writeAppendCommit( + Engine engine, + String tablePath, + long expectedVersion, + long timestamp, + StructType deltaSchema, + List beamRows) 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); + + Table table = Table.forPath(engine, tablePath); + TransactionBuilder txnBuilder = + table.createTransactionBuilder(engine, "DeltaIOTest", Operation.WRITE); + if (expectedVersion == 0) { + txnBuilder = + txnBuilder + .withSchema(engine, deltaSchema) + .withTableProperties( + engine, Collections.singletonMap("delta.enableChangeDataFeed", "true")); } - 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); + Transaction txn = txnBuilder.build(engine); + io.delta.kernel.data.Row txnState = txn.getTransactionState(engine); + + ColumnVector[] vectors = new ColumnVector[deltaSchema.fields().size()]; + for (int i = 0; i < deltaSchema.fields().size(); i++) { + StructField field = deltaSchema.fields().get(i); + vectors[i] = createColumnVector(beamRows, i, field.getDataType()); + } + + ColumnarBatch columnarBatch = new DefaultColumnarBatch(beamRows.size(), deltaSchema, vectors); + FilteredColumnarBatch filteredBatch = + new FilteredColumnarBatch(columnarBatch, Optional.empty()); + + CloseableIterator data = + io.delta.kernel.internal.util.Utils.toCloseableIterator( + Collections.singletonList(filteredBatch).iterator()); + + CloseableIterator physicalData = + Transaction.transformLogicalData(engine, txnState, data, Collections.emptyMap()); + + DataWriteContext writeContext = + Transaction.getWriteContext(engine, txnState, Collections.emptyMap()); + + CloseableIterator dataFiles = + engine + .getParquetHandler() + .writeParquetFiles( + writeContext.getTargetDirectory(), + physicalData, + writeContext.getStatisticsColumns()); + + List writtenFiles = new ArrayList<>(); + List filesList = new ArrayList<>(); + while (dataFiles.hasNext()) { + DataFileStatus file = dataFiles.next(); + filesList.add(file); + writtenFiles.add(new File(file.getPath()).getName()); + } + CloseableIterator dataFilesCopy = + io.delta.kernel.internal.util.Utils.toCloseableIterator(filesList.iterator()); + + CloseableIterator dataActions = + Transaction.generateAppendActions(engine, txnState, dataFilesCopy, writeContext); + + TransactionCommitResult result = + txn.commit(engine, CloseableIterable.inMemoryIterable(dataActions)); + org.junit.Assert.assertEquals(expectedVersion, result.getVersion()); + File commitFile = + new File(new File(tablePath, "_delta_log"), String.format("%020d.json", expectedVersion)); + commitFile.setLastModified(timestamp); + return writtenFiles; + } + + private void writeCdcCommit( + Engine engine, + String tablePath, + long expectedVersion, + long timestamp, + StructType deltaSchema, + @Nullable List addBeamRows, + @Nullable String removePath, + @Nullable List cdcBeamRows, + StructType cdcWriteSchema) + throws Exception { + + Table table = Table.forPath(engine, tablePath); + TransactionBuilder txnBuilder = + table.createTransactionBuilder(engine, "DeltaIOTest", Operation.WRITE); + Transaction txn = txnBuilder.build(engine); + io.delta.kernel.data.Row txnState = txn.getTransactionState(engine); + + StructType customSingleActionSchema = getCustomSingleActionSchema(); + List commitActions = new ArrayList<>(); + + if (addBeamRows != null && !addBeamRows.isEmpty()) { + ColumnVector[] vectors = new ColumnVector[deltaSchema.fields().size()]; + for (int i = 0; i < deltaSchema.fields().size(); i++) { + StructField field = deltaSchema.fields().get(i); + vectors[i] = createColumnVector(addBeamRows, i, field.getDataType()); + } + ColumnarBatch columnarBatch = + new DefaultColumnarBatch(addBeamRows.size(), deltaSchema, vectors); + FilteredColumnarBatch filteredBatch = + new FilteredColumnarBatch(columnarBatch, Optional.empty()); + CloseableIterator data = + io.delta.kernel.internal.util.Utils.toCloseableIterator( + Collections.singletonList(filteredBatch).iterator()); + CloseableIterator physicalData = + Transaction.transformLogicalData(engine, txnState, data, Collections.emptyMap()); + DataWriteContext writeContext = + Transaction.getWriteContext(engine, txnState, Collections.emptyMap()); + CloseableIterator dataFiles = + engine + .getParquetHandler() + .writeParquetFiles( + writeContext.getTargetDirectory(), + physicalData, + writeContext.getStatisticsColumns()); + CloseableIterator addActions = + Transaction.generateAppendActions(engine, txnState, dataFiles, writeContext); + while (addActions.hasNext()) { + commitActions.add(addActions.next()); } } - return java.nio.file.Files.readAllBytes(file.toPath()); + + if (removePath != null) { + StructType removeSchema = + (StructType) + io.delta.kernel.internal.actions.SingleAction.FULL_SCHEMA + .fields() + .get(io.delta.kernel.internal.actions.SingleAction.REMOVE_FILE_ORDINAL) + .getDataType(); + io.delta.kernel.data.Row removeAction = + createRemoveAction(removeSchema, removePath, timestamp); + commitActions.add(createSingleAction(customSingleActionSchema, "remove", removeAction)); + } + + if (cdcBeamRows != null && !cdcBeamRows.isEmpty()) { + ColumnVector[] vectors = new ColumnVector[cdcWriteSchema.fields().size()]; + for (int i = 0; i < cdcWriteSchema.fields().size(); i++) { + StructField field = cdcWriteSchema.fields().get(i); + vectors[i] = createColumnVector(cdcBeamRows, i, field.getDataType()); + } + ColumnarBatch columnarBatch = + new DefaultColumnarBatch(cdcBeamRows.size(), cdcWriteSchema, vectors); + FilteredColumnarBatch filteredBatch = + new FilteredColumnarBatch(columnarBatch, Optional.empty()); + CloseableIterator data = + io.delta.kernel.internal.util.Utils.toCloseableIterator( + Collections.singletonList(filteredBatch).iterator()); + + String cdcDir = new File(tablePath, "_change_data").getAbsolutePath(); + + CloseableIterator cdcFiles = + engine.getParquetHandler().writeParquetFiles(cdcDir, data, Collections.emptyList()); + + StructType cdcActionSchema = CDC_ACTION_SCHEMA; + while (cdcFiles.hasNext()) { + DataFileStatus cdcFile = cdcFiles.next(); + String relativeCdcPath = "_change_data/" + new File(cdcFile.getPath()).getName(); + io.delta.kernel.data.Row cdcAction = + createCdcAction(cdcActionSchema, relativeCdcPath, cdcFile.getSize()); + commitActions.add(createSingleAction(customSingleActionSchema, "cdc", cdcAction)); + } + } + + TransactionCommitResult result = + txn.commit( + engine, + CloseableIterable.inMemoryIterable( + io.delta.kernel.internal.util.Utils.toCloseableIterator(commitActions.iterator()))); + org.junit.Assert.assertEquals(expectedVersion, result.getVersion()); + File commitFile = + new File(new File(tablePath, "_delta_log"), String.format("%020d.json", expectedVersion)); + commitFile.setLastModified(timestamp); + } + + private static final StructType CDC_ACTION_SCHEMA = + new StructType() + .add("path", StringType.STRING, false) + .add("partitionValues", new MapType(StringType.STRING, StringType.STRING, false), false) + .add("size", LongType.LONG, false) + .add("dataChange", BooleanType.BOOLEAN, false); + + private static StructType getCustomSingleActionSchema() { + StructType originalSchema = io.delta.kernel.internal.actions.SingleAction.FULL_SCHEMA; + List fields = new ArrayList<>(); + for (StructField field : originalSchema.fields()) { + if (field.getName().equals("cdc")) { + fields.add(new StructField("cdc", CDC_ACTION_SCHEMA, true)); + } else { + fields.add(field); + } + } + return new StructType(fields); + } + + private static io.delta.kernel.data.Row createSingleAction( + StructType customSingleActionSchema, String actionName, io.delta.kernel.data.Row actionRow) { + Map values = new HashMap<>(); + values.put(actionName, actionRow); + return new TestRow(customSingleActionSchema, values); + } + + private static final MapValue EMPTY_MAP_VALUE = + new MapValue() { + @Override + public int getSize() { + return 0; + } + + @Override + public ColumnVector getKeys() { + return new ColumnVector() { + @Override + public DataType getDataType() { + return StringType.STRING; + } + + @Override + public int getSize() { + return 0; + } + + @Override + public void close() {} + + @Override + public boolean isNullAt(int rowId) { + return true; + } + }; + } + + @Override + public ColumnVector getValues() { + return new ColumnVector() { + @Override + public DataType getDataType() { + return StringType.STRING; + } + + @Override + public int getSize() { + return 0; + } + + @Override + public void close() {} + + @Override + public boolean isNullAt(int rowId) { + return true; + } + }; + } + }; + + private static io.delta.kernel.data.Row createRemoveAction( + StructType removeSchema, String path, long deletionTimestamp) { + Map values = new HashMap<>(); + values.put("path", path); + values.put("deletionTimestamp", deletionTimestamp); + values.put("dataChange", true); + values.put("size", 100L); + return new TestRow(removeSchema, values); + } + + private static io.delta.kernel.data.Row createCdcAction( + StructType cdcSchema, String path, long size) { + Map values = new HashMap<>(); + values.put("path", path); + values.put("partitionValues", EMPTY_MAP_VALUE); + values.put("size", size); + values.put("dataChange", true); + return new TestRow(cdcSchema, values); + } + + private static ColumnVector createColumnVector( + List rows, int fieldIndex, DataType dataType) { + return new ColumnVector() { + @Override + public DataType getDataType() { + return dataType; + } + + @Override + public int getSize() { + return rows.size(); + } + + @Override + public void close() {} + + @Override + public boolean isNullAt(int rowId) { + return rows.get(rowId).getValue(fieldIndex) == null; + } + + @Override + public boolean getBoolean(int rowId) { + return rows.get(rowId).getBoolean(fieldIndex); + } + + @Override + public int getInt(int rowId) { + return rows.get(rowId).getInt32(fieldIndex); + } + + @Override + public long getLong(int rowId) { + if (dataType instanceof TimestampType) { + org.joda.time.Instant instant = rows.get(rowId).getDateTime(fieldIndex).toInstant(); + return instant.getMillis() * 1000L; + } + return rows.get(rowId).getInt64(fieldIndex); + } + + @Override + public String getString(int rowId) { + return rows.get(rowId).getString(fieldIndex); + } + }; + } + + private static class TestRow implements io.delta.kernel.data.Row { + private final StructType schema; + private final Map values; + + public TestRow(StructType schema, Map values) { + this.schema = schema; + this.values = values; + } + + @Override + public StructType getSchema() { + return schema; + } + + private Object getVal(int ord) { + String name = schema.fields().get(ord).getName(); + return values.get(name); + } + + @Override + public boolean isNullAt(int ord) { + return getVal(ord) == null; + } + + @Override + public boolean getBoolean(int ord) { + return (Boolean) getVal(ord); + } + + @Override + public byte getByte(int ord) { + return (Byte) getVal(ord); + } + + @Override + public short getShort(int ord) { + return (Short) getVal(ord); + } + + @Override + public int getInt(int ord) { + return (Integer) getVal(ord); + } + + @Override + public long getLong(int ord) { + return (Long) getVal(ord); + } + + @Override + public float getFloat(int ord) { + return (Float) getVal(ord); + } + + @Override + public double getDouble(int ord) { + return (Double) getVal(ord); + } + + @Override + public String getString(int ord) { + return (String) getVal(ord); + } + + @Override + public byte[] getBinary(int ord) { + return (byte[]) getVal(ord); + } + + @Override + public BigDecimal getDecimal(int ord) { + return (BigDecimal) getVal(ord); + } + + @Override + public io.delta.kernel.data.Row getStruct(int ord) { + return (io.delta.kernel.data.Row) getVal(ord); + } + + @Override + public io.delta.kernel.data.ArrayValue getArray(int ord) { + return (io.delta.kernel.data.ArrayValue) getVal(ord); + } + + @Override + public io.delta.kernel.data.MapValue getMap(int ord) { + return (io.delta.kernel.data.MapValue) getVal(ord); + } } }