ahmedabu98 commented on code in PR #39426:
URL: https://github.com/apache/beam/pull/39426#discussion_r3661451934
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java:
##########
@@ -174,4 +180,124 @@ static Schema.FieldType convertToBeamFieldType(DataType
deltaType) {
}
}
}
+
+ @AutoValue
+ public abstract static class ReadChanges extends PTransform<PBegin,
PCollection<Row>> {
+ public abstract @Nullable String getTablePath();
+
+ public abstract @Nullable Long getStartVersion();
+
+ public abstract @Nullable String getStartTimestamp();
+
+ public abstract @Nullable Long getEndVersion();
+
+ public abstract @Nullable String getEndTimestamp();
+
+ public abstract @Nullable Map<String, String> getHadoopConfig();
+
+ abstract Builder toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setTablePath(String tablePath);
+
+ abstract Builder setStartVersion(@Nullable Long startVersion);
+
+ abstract Builder setStartTimestamp(@Nullable String startTimestamp);
+
+ abstract Builder setEndVersion(@Nullable Long endVersion);
+
+ abstract Builder setEndTimestamp(@Nullable String endTimestamp);
+
+ abstract Builder setHadoopConfig(@Nullable Map<String, String>
hadoopConfig);
+
+ abstract ReadChanges build();
+ }
+
+ public ReadChanges from(String tablePath) {
+ return toBuilder().setTablePath(tablePath).build();
+ }
+
+ public ReadChanges withStartVersion(long startVersion) {
+ return toBuilder().setStartVersion(startVersion).build();
+ }
+
+ public ReadChanges withStartTimestamp(String startTimestamp) {
+ return toBuilder().setStartTimestamp(startTimestamp).build();
+ }
+
+ public ReadChanges withEndVersion(long endVersion) {
+ return toBuilder().setEndVersion(endVersion).build();
+ }
+
+ public ReadChanges withEndTimestamp(String endTimestamp) {
+ return toBuilder().setEndTimestamp(endTimestamp).build();
+ }
+
+ public ReadChanges withConfig(Map<String, String> config) {
+ return toBuilder().setHadoopConfig(config).build();
+ }
+
+ @Override
+ public PCollection<Row> expand(PBegin input) {
+ String path = getTablePath();
+ if (path == null) {
+ throw new IllegalArgumentException("Table path must be set.");
+ }
+ if (getStartVersion() == null && getStartTimestamp() == null) {
+ throw new IllegalArgumentException("Either startVersion or
startTimestamp must be set.");
+ }
+ if (getStartVersion() != null && getStartTimestamp() != null) {
+ throw new IllegalArgumentException("Cannot set both startVersion and
startTimestamp.");
+ }
+ if (getEndVersion() != null && getEndTimestamp() != null) {
+ throw new IllegalArgumentException("Cannot set both endVersion and
endTimestamp.");
+ }
Review Comment:
Same as above. If no end is specified, we can just stop at the very latest
change
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java:
##########
@@ -174,4 +180,124 @@ static Schema.FieldType convertToBeamFieldType(DataType
deltaType) {
}
}
}
+
+ @AutoValue
+ public abstract static class ReadChanges extends PTransform<PBegin,
PCollection<Row>> {
+ public abstract @Nullable String getTablePath();
+
+ public abstract @Nullable Long getStartVersion();
+
+ public abstract @Nullable String getStartTimestamp();
+
+ public abstract @Nullable Long getEndVersion();
+
+ public abstract @Nullable String getEndTimestamp();
+
+ public abstract @Nullable Map<String, String> getHadoopConfig();
+
+ abstract Builder toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setTablePath(String tablePath);
+
+ abstract Builder setStartVersion(@Nullable Long startVersion);
+
+ abstract Builder setStartTimestamp(@Nullable String startTimestamp);
+
+ abstract Builder setEndVersion(@Nullable Long endVersion);
+
+ abstract Builder setEndTimestamp(@Nullable String endTimestamp);
+
+ abstract Builder setHadoopConfig(@Nullable Map<String, String>
hadoopConfig);
+
+ abstract ReadChanges build();
+ }
+
+ public ReadChanges from(String tablePath) {
+ return toBuilder().setTablePath(tablePath).build();
+ }
+
+ public ReadChanges withStartVersion(long startVersion) {
+ return toBuilder().setStartVersion(startVersion).build();
+ }
+
+ public ReadChanges withStartTimestamp(String startTimestamp) {
+ return toBuilder().setStartTimestamp(startTimestamp).build();
+ }
+
+ public ReadChanges withEndVersion(long endVersion) {
+ return toBuilder().setEndVersion(endVersion).build();
+ }
+
+ public ReadChanges withEndTimestamp(String endTimestamp) {
+ return toBuilder().setEndTimestamp(endTimestamp).build();
+ }
+
+ public ReadChanges withConfig(Map<String, String> config) {
+ return toBuilder().setHadoopConfig(config).build();
+ }
+
+ @Override
+ public PCollection<Row> expand(PBegin input) {
+ String path = getTablePath();
+ if (path == null) {
+ throw new IllegalArgumentException("Table path must be set.");
+ }
+ if (getStartVersion() == null && getStartTimestamp() == null) {
+ throw new IllegalArgumentException("Either startVersion or
startTimestamp must be set.");
+ }
Review Comment:
If there's no start version specified, we can have a default starting
strategy. Either start at the very beginning, or at the very latest change.
Not blocking though, can add this in a future change
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java:
##########
@@ -0,0 +1,293 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import io.delta.kernel.CommitRange;
+import io.delta.kernel.CommitRangeBuilder;
+import io.delta.kernel.Scan;
+import io.delta.kernel.Snapshot;
+import io.delta.kernel.Table;
+import io.delta.kernel.TableManager;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.data.Row;
+import io.delta.kernel.defaults.engine.DefaultEngine;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.internal.DeltaLogActionUtils.DeltaAction;
+import io.delta.kernel.internal.TableImpl;
+import io.delta.kernel.internal.actions.AddCDCFile;
+import io.delta.kernel.internal.actions.AddFile;
+import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.utils.CloseableIterator;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.hadoop.conf.Configuration;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** A DoFn that reads the Delta log and plans read tasks for Change Data Feed.
*/
+class CreateCDCReadTasksDoFn extends DoFn<String, DeltaCDCReadTask> {
+ private static final long MAX_TASK_SIZE_BYTES = 1024L * 1024L * 1024L; // 1
GB
+ private final @Nullable Map<String, String> hadoopConfig;
+ private final @Nullable Long startVersion;
+ private final @Nullable String startTimestamp;
+ private final @Nullable Long endVersion;
+ private final @Nullable String endTimestamp;
+
+ public CreateCDCReadTasksDoFn(
+ @Nullable Map<String, String> hadoopConfig,
+ @Nullable Long startVersion,
+ @Nullable String startTimestamp,
+ @Nullable Long endVersion,
+ @Nullable String endTimestamp) {
+ this.hadoopConfig = hadoopConfig;
+ this.startVersion = startVersion;
+ this.startTimestamp = startTimestamp;
+ this.endVersion = endVersion;
+ this.endTimestamp = endTimestamp;
+ }
+
+ @ProcessElement
+ public void processElement(@Element String tablePath,
OutputReceiver<DeltaCDCReadTask> out)
+ throws Exception {
+ Configuration conf = new Configuration();
+ if (hadoopConfig != null) {
+ for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) {
+ conf.set(entry.getKey(), entry.getValue());
+ }
+ }
+ Engine engine = DefaultEngine.create(conf);
+ Table table = Table.forPath(engine, tablePath);
+ TableImpl tableImpl = (TableImpl) table;
+
+ // 1. Resolve starting and ending versions
+ long resolvedStartVersion;
+ if (startVersion != null) {
+ resolvedStartVersion = startVersion;
+ } else if (startTimestamp != null) {
+ long startMillis = Instant.parse(startTimestamp).toEpochMilli();
+ resolvedStartVersion = tableImpl.getVersionAtOrAfterTimestamp(engine,
startMillis);
+ } else {
+ throw new IllegalArgumentException("Starting version or timestamp must
be specified.");
+ }
+
+ long resolvedEndVersion;
+ if (endVersion != null) {
+ resolvedEndVersion = endVersion;
+ } else if (endTimestamp != null) {
+ long endMillis = Instant.parse(endTimestamp).toEpochMilli();
+ resolvedEndVersion = tableImpl.getVersionBeforeOrAtTimestamp(engine,
endMillis);
+ } else {
+ resolvedEndVersion = table.getLatestSnapshot(engine).getVersion();
+ }
+
+ if (resolvedStartVersion > resolvedEndVersion) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Resolved start version %d is greater than resolved end version
%d",
+ resolvedStartVersion, resolvedEndVersion));
+ }
+
+ // 2. Load snapshot at resolvedEndVersion to get the scanStateRow
+ // We use endVersion's schema because it represents the latest schema in
the
+ // read range
+ // which handles schema evolution (older files will just lack new columns).
+ Snapshot endSnapshot = table.getSnapshotAsOfVersion(engine,
resolvedEndVersion);
+ Scan scan = endSnapshot.getScanBuilder().build();
+ Row scanState = scan.getScanState(engine);
+ SerializableRow serializableScanState = new SerializableRow(scanState);
+
+ // 3. Load snapshot at resolvedStartVersion to initialize the CommitRange
+ Snapshot startSnapshot = table.getSnapshotAsOfVersion(engine,
resolvedStartVersion);
+
+ CommitRangeBuilder rangeBuilder =
+ TableManager.loadCommitRange(
+ tablePath,
CommitRangeBuilder.CommitBoundary.atVersion(resolvedStartVersion));
+
rangeBuilder.withEndBoundary(CommitRangeBuilder.CommitBoundary.atVersion(resolvedEndVersion));
+ CommitRange range = rangeBuilder.build(engine);
+
+ // We need both CDC and ADD actions.
+ // If a commit version has CDC files, we only read CDC files.
+ // If a commit version has no CDC files, we read ADD files (inserts).
+ Set<DeltaAction> actionSet = new HashSet<>();
+ actionSet.add(DeltaAction.CDC);
+ actionSet.add(DeltaAction.ADD);
+
+ // 4. Iterate over commits in the range and group actions by version
+ try (CloseableIterator<ColumnarBatch> batchIter =
+ range.getActions(engine, startSnapshot, actionSet)) {
+ Map<Long, CommitActionsInfo> commitActionsMap = new HashMap<>();
+
+ while (batchIter.hasNext()) {
+ ColumnarBatch batch = batchIter.next();
+ int versionIdx = batch.getSchema().indexOf("version");
+ int timestampIdx = batch.getSchema().indexOf("timestamp");
+ int cdcIdx = batch.getSchema().indexOf("cdc");
+ int addIdx = batch.getSchema().indexOf("add");
+
+ for (int i = 0; i < batch.getSize(); i++) {
+ long version = batch.getColumnVector(versionIdx).getLong(i);
+ long timestamp = batch.getColumnVector(timestampIdx).getLong(i);
+
+ CommitActionsInfo info =
+ commitActionsMap.computeIfAbsent(
+ version, k -> new CommitActionsInfo(version, timestamp));
+
+ if (cdcIdx >= 0 && !batch.getColumnVector(cdcIdx).isNullAt(i)) {
+ Row cdcRow =
+ (Row)
+ VectorUtils.getValueAsObject(
+ batch.getColumnVector(cdcIdx),
+ batch.getSchema().at(cdcIdx).getDataType(),
+ i);
+ info.cdcRows.add(cdcRow);
+ }
+ if (addIdx >= 0 && !batch.getColumnVector(addIdx).isNullAt(i)) {
+ Row addRow =
+ (Row)
+ VectorUtils.getValueAsObject(
+ batch.getColumnVector(addIdx),
+ batch.getSchema().at(addIdx).getDataType(),
+ i);
+ AddFile addFile = new AddFile(addRow);
+ // Only consider add files that change data (ignore OPTIMIZE etc.)
+ if (addFile.getDataChange()) {
+ info.addRows.add(addRow);
+ }
+ }
+ }
+ }
+
+ // 5. Emit tasks for each version
+ List<DeltaCDCReadTask> currentGroup = new ArrayList<>();
+ long currentGroupSize = 0L;
+
+ // Sort versions to process them in order
+ List<Long> versions = new ArrayList<>(commitActionsMap.keySet());
+ Collections.sort(versions);
+
+ for (long version : versions) {
+ CommitActionsInfo info = commitActionsMap.get(version);
+ if (info == null) {
+ throw new IllegalStateException("CommitActionsInfo was not found for
version " + version);
+ }
+ boolean hasCDC = !info.cdcRows.isEmpty();
+
+ List<Row> rowsToProcess = hasCDC ? info.cdcRows : info.addRows;
+ boolean isCDC = hasCDC;
Review Comment:
nit: unnecessary assignment to another variable?
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java:
##########
@@ -0,0 +1,353 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import static io.delta.kernel.internal.DeltaErrors.wrapEngineException;
+
+import io.delta.kernel.Scan;
+import io.delta.kernel.data.ColumnVector;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.MapValue;
+import io.delta.kernel.defaults.engine.DefaultEngine;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.engine.FileReadResult;
+import io.delta.kernel.expressions.ExpressionEvaluator;
+import io.delta.kernel.expressions.Literal;
+import io.delta.kernel.internal.InternalScanFileUtils;
+import io.delta.kernel.internal.data.GenericRow;
+import io.delta.kernel.internal.data.ScanStateRow;
+import io.delta.kernel.internal.util.Utils;
+import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.types.LongType;
+import io.delta.kernel.types.StringType;
+import io.delta.kernel.types.StructField;
+import io.delta.kernel.types.StructType;
+import io.delta.kernel.types.TimestampType;
+import io.delta.kernel.utils.CloseableIterator;
+import io.delta.kernel.utils.FileStatus;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.apache.hadoop.conf.Configuration;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and
reads Change Data Feed
+ * files, converting rows to Beam Rows.
+ */
[email protected]
+class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> {
Review Comment:
This DoFn parallelizes at the file level right? I think this should be fine
for batch, but note that it may not be scalable in DF streaming mode because
the runner opens too many threads per VM, each thread processing a file and
competing for memory leading to OOM.
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java:
##########
@@ -0,0 +1,353 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import static io.delta.kernel.internal.DeltaErrors.wrapEngineException;
+
+import io.delta.kernel.Scan;
+import io.delta.kernel.data.ColumnVector;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.MapValue;
+import io.delta.kernel.defaults.engine.DefaultEngine;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.engine.FileReadResult;
+import io.delta.kernel.expressions.ExpressionEvaluator;
+import io.delta.kernel.expressions.Literal;
+import io.delta.kernel.internal.InternalScanFileUtils;
+import io.delta.kernel.internal.data.GenericRow;
+import io.delta.kernel.internal.data.ScanStateRow;
+import io.delta.kernel.internal.util.Utils;
+import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.types.LongType;
+import io.delta.kernel.types.StringType;
+import io.delta.kernel.types.StructField;
+import io.delta.kernel.types.StructType;
+import io.delta.kernel.types.TimestampType;
+import io.delta.kernel.utils.CloseableIterator;
+import io.delta.kernel.utils.FileStatus;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.apache.hadoop.conf.Configuration;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and
reads Change Data Feed
+ * files, converting rows to Beam Rows.
+ */
[email protected]
+class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> {
+ @Nullable Map<String, String> hadoopConfig;
+ private transient @Nullable Engine engine;
+ private transient @Nullable Configuration conf;
+
+ public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) {
+ this.hadoopConfig = hadoopConfig;
+ }
+
+ private synchronized Configuration getConfiguration() {
+ Configuration localConf = conf;
+ if (localConf == null) {
+ localConf = new Configuration();
+ if (hadoopConfig != null) {
+ for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) {
+ localConf.set(entry.getKey(), entry.getValue());
+ }
+ }
+ conf = localConf;
+ }
+ return localConf;
+ }
+
+ private List<Long> getRowGroupSizes(DeltaCDCReadTask task) {
+ return task.getRowGroupSizes();
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) {
+ List<Long> rowGroupSizes = getRowGroupSizes(task);
+ return new OffsetRange(0L, rowGroupSizes.size());
+ }
+
+ @NewTracker
+ public DeltaReadTaskTracker newTracker(
+ @Restriction OffsetRange restriction, @Element DeltaCDCReadTask task) {
+ return new DeltaReadTaskTracker(restriction, getRowGroupSizes(task));
+ }
+
+ @Setup
+ public void setUp() {
+ engine = DefaultEngine.create(getConfiguration());
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element DeltaCDCReadTask task,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<Row> out)
+ throws Exception {
+
+ Engine currentEngine = engine;
+ if (currentEngine == null) {
+ throw new IllegalArgumentException("Expected the engine to not be null");
+ }
+
+ SerializableRow originalScanStateRow = task.getScanStateRow();
+ StructType logicalTableSchema =
ScanStateRow.getLogicalSchema(originalScanStateRow);
+ Schema publicBeamSchema =
DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
+ StructType physicalTableSchema =
ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow);
+
+ StructType scanStateSchema = originalScanStateRow.getSchema();
+
+ // 1. Build modified scanState and scanFile rows depending on whether we
read a CDC file or ADD
+ // file.
+ io.delta.kernel.data.Row scanStateRow;
+ StructType readPhysicalSchema;
+ StructType readLogicalSchema;
+ Schema beamSchema;
+
+ if (task.isCDC()) {
+ readLogicalSchema = appendCDFColumns(logicalTableSchema);
+ readPhysicalSchema = appendCDFColumns(physicalTableSchema);
+ beamSchema = DeltaIO.ReadRows.convertToBeamSchema(readLogicalSchema);
+
+ HashMap<Integer, Object> valueMap = new HashMap<>();
+
+ // Tracking row level lineage is not needed.
+ Map<String, String> config =
+ new HashMap<>(ScanStateRow.getConfiguration(originalScanStateRow));
+ config.put("delta.enableRowTracking", "false");
+
+ valueMap.put(
+ scanStateSchema.indexOf("configuration"),
VectorUtils.stringStringMapValue(config));
+ valueMap.put(scanStateSchema.indexOf("logicalSchemaJson"),
readLogicalSchema.toJson());
+ valueMap.put(scanStateSchema.indexOf("physicalSchemaJson"),
readPhysicalSchema.toJson());
+ valueMap.put(
+ scanStateSchema.indexOf("partitionColumns"),
+
originalScanStateRow.getArray(scanStateSchema.indexOf("partitionColumns")));
+ valueMap.put(
+ scanStateSchema.indexOf("minReaderVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minReaderVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("minWriterVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minWriterVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("tablePath"),
+
originalScanStateRow.getString(scanStateSchema.indexOf("tablePath")));
+
+ scanStateRow = new ScanStateRow(valueMap);
+ } else {
+ // For ADD files, we read the table schema and append the CDF columns
manually afterwards.
+ // readLogicalSchema = logicalTableSchema;
+ readPhysicalSchema = physicalTableSchema;
+ beamSchema =
DeltaIO.ReadRows.convertToBeamSchema(appendCDFColumns(logicalTableSchema));
+ scanStateRow = originalScanStateRow;
+ }
+
+ io.delta.kernel.data.Row scanFileRow =
+ generateScanFileRow(task.getPath(), task.getPartitionValues());
+ FileStatus fileStatus = FileStatus.of(task.getPath(), task.getSize(),
task.getTimestamp());
+
+ BeamParquetHandler parquetHandler =
+ new BeamParquetHandler(getConfiguration(),
currentEngine.getParquetHandler(), tracker);
+ BeamEngine beamEngine = new BeamEngine(currentEngine, parquetHandler);
+
+ long currentStartRgIndex = 0L;
+
+ try (CloseableIterator<FileReadResult> fileReadResults =
+ parquetHandler.readParquetFiles(
+ Utils.singletonCloseableIterator(fileStatus),
+ readPhysicalSchema,
+ Optional.empty(),
+ currentStartRgIndex)) {
+
+ CloseableIterator<ColumnarBatch> physicalData =
+ new CloseableIterator<ColumnarBatch>() {
+ @Override
+ public void close() throws java.io.IOException {}
+
+ @Override
+ public boolean hasNext() {
+ return fileReadResults.hasNext();
+ }
+
+ @Override
+ public ColumnarBatch next() {
+ return fileReadResults.next().getData();
+ }
+ };
+
+ try (CloseableIterator<FilteredColumnarBatch> logicalBatches =
+ Scan.transformPhysicalData(beamEngine, scanStateRow, scanFileRow,
physicalData)) {
+
+ while (logicalBatches.hasNext()) {
+ FilteredColumnarBatch batch = logicalBatches.next();
+ ColumnarBatch logicalBatch = batch.getData();
+
+ if (!task.isCDC()) {
+ // For ADD files, we need to append the constant CDF columns:
+ // _change_type = "insert", _commit_version = task.version,
_commit_timestamp =
+ // task.timestamp
+ logicalBatch =
+ appendConstantCDFColumns(
+ currentEngine, logicalBatch, task.getVersion(),
task.getTimestamp());
+ }
+
+ try (CloseableIterator<io.delta.kernel.data.Row> logicalRows =
logicalBatch.getRows()) {
+ while (logicalRows.hasNext()) {
+ io.delta.kernel.data.Row deltaRow = logicalRows.next();
+ Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema);
+ String changeType = beamRow.getString("_change_type");
+ if (changeType == null) {
+ throw new IllegalStateException("Field _change_type must not
be null.");
+ }
+ ValueKind kind = getValueKind(changeType);
+ Row publicRow = projectRow(beamRow, publicBeamSchema);
+ out.builder(publicRow).setValueKind(kind).output();
+ }
+ }
+ }
+ }
+ }
+
+ return ProcessContinuation.stop();
+ }
+
+ private static Row projectRow(Row row, Schema targetSchema) {
+ Row.Builder builder = Row.withSchema(targetSchema);
+ for (Schema.Field field : targetSchema.getFields()) {
+ builder.addValue(row.getValue(field.getName()));
+ }
+ return builder.build();
+ }
+
+ private static ValueKind getValueKind(String changeType) {
+ // Maps Delta CDC change types to Beam's ValueKind enum.
+ //
https://docs.delta.io/delta-change-data-feed/#what-is-the-schema-for-the-change-data-feed
+ switch (changeType) {
+ case "insert":
+ return ValueKind.INSERT;
+ case "delete":
+ return ValueKind.DELETE;
+ case "update_preimage":
+ return ValueKind.UPDATE_BEFORE;
+ case "update_postimage":
+ return ValueKind.UPDATE_AFTER;
+ default:
+ throw new IllegalArgumentException("Unsupported change type: " +
changeType);
+ }
+ }
+
+ private static StructType appendCDFColumns(StructType schema) {
+ return schema
+ .add("_change_type", StringType.STRING, false)
+ .add("_commit_version", LongType.LONG, false)
+ .add("_commit_timestamp", TimestampType.TIMESTAMP, false);
Review Comment:
Can we have these field names be final static variables at the top?
##########
sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java:
##########
@@ -810,4 +815,263 @@ public boolean tryClaim(Long i) {
}
}
}
+
+ @Test
+ public void testReadChanges() throws Exception {
+ File tableDir = tempFolder.newFolder("delta-table-changes");
+ File logDir = new File(tableDir, "_delta_log");
+ logDir.mkdirs();
+
+ // 1. Write parquet files for Version 0 (insert-only commit)
+ Schema tableSchema = Schema.builder().addField("name",
Schema.FieldType.STRING).build();
+ Row tableRow1 = Row.withSchema(tableSchema).addValues("row-1").build();
+ Row tableRow2 = Row.withSchema(tableSchema).addValues("row-2").build();
+
+ File partFile = new File(tableDir, "part-00000.parquet");
+ byte[] partBytes =
+ writeParquetFile(partFile, tableSchema,
java.util.Arrays.asList(tableRow1, tableRow2));
+
+ writeCommit(
+ logDir, 0L, 100000000000L, "part-00000.parquet", partBytes.length,
null, null, 0L, true);
+
+ // 2. Write cdc parquet file for Version 1 (commit with cdc actions)
+ Schema cdcWriteSchema =
+ Schema.builder()
+ .addField("name", Schema.FieldType.STRING)
+ .addField("_change_type", Schema.FieldType.STRING)
+ .addField("_commit_version", Schema.FieldType.INT64)
+ .addField("_commit_timestamp", Schema.FieldType.DATETIME)
+ .build();
+
+ Row cdcRow1 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-1", "update_preimage", 1L, new
Instant(123456789000L))
+ .build();
+ Row cdcRow2 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-1-updated", "update_postimage", 1L, new
Instant(123456789000L))
+ .build();
+ Row cdcRow3 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-2", "delete", 1L, new Instant(123456789000L))
+ .build();
+
+ File changeFile = new File(tableDir, "change-00000.parquet");
+ byte[] changeBytes =
+ writeParquetFile(
+ changeFile, cdcWriteSchema, java.util.Arrays.asList(cdcRow1,
cdcRow2, cdcRow3));
+
+ writeCommit(
+ logDir,
+ 1L,
+ 200000000000L,
+ null,
+ 0L,
+ null,
+ "change-00000.parquet",
+ changeBytes.length,
+ false);
+
+ // 3. Read CDF data from table using ReadChanges
+ PCollection<Row> output =
+ readPipeline.apply(
+
DeltaIO.readChanges().from(tableDir.getAbsolutePath()).withStartVersion(0L));
+
+ PCollection<String> formattedOutput =
+ output.apply("Format ValueKind and Row", ParDo.of(new
FormatValueKindAndRow()));
+
+ PAssert.that(formattedOutput)
+ .containsInAnyOrder(
+ "INSERT:row-1",
+ "INSERT:row-2",
+ "UPDATE_BEFORE:row-1",
+ "UPDATE_AFTER:row-1-updated",
+ "DELETE:row-2");
+
+ readPipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testReadChangesRanges() throws Exception {
+ File tableDir = tempFolder.newFolder("delta-table-changes-ranges");
+ File logDir = new File(tableDir, "_delta_log");
+ logDir.mkdirs();
+
+ Schema tableSchema = Schema.builder().addField("name",
Schema.FieldType.STRING).build();
+
+ // 1. Write parquet files for Version 0 (insert-only commit)
+ Row tableRow1 = Row.withSchema(tableSchema).addValues("row-1").build();
+ Row tableRow2 = Row.withSchema(tableSchema).addValues("row-2").build();
+ File partFile0 = new File(tableDir, "part-00000.parquet");
+ byte[] partBytes0 =
+ writeParquetFile(partFile0, tableSchema,
java.util.Arrays.asList(tableRow1, tableRow2));
+ writeCommit(
+ logDir, 0L, 100000000000L, "part-00000.parquet", partBytes0.length,
null, null, 0L, true);
+
+ // 2. Write parquet files for Version 1 (commit with updates and deletes)
+ Schema cdcWriteSchema =
+ Schema.builder()
+ .addField("name", Schema.FieldType.STRING)
+ .addField("_change_type", Schema.FieldType.STRING)
+ .addField("_commit_version", Schema.FieldType.INT64)
+ .addField("_commit_timestamp", Schema.FieldType.DATETIME)
+ .build();
+ Row cdcRow1 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-1", "update_preimage", 1L, new
Instant(200000000000L))
+ .build();
+ Row cdcRow2 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-1-updated", "update_postimage", 1L, new
Instant(200000000000L))
+ .build();
+ Row cdcRow3 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-2", "delete", 1L, new Instant(200000000000L))
+ .build();
+ File changeFile0 = new File(tableDir, "change-00000.parquet");
+ byte[] changeBytes0 =
+ writeParquetFile(
+ changeFile0, cdcWriteSchema, java.util.Arrays.asList(cdcRow1,
cdcRow2, cdcRow3));
+
+ Row tableRow1Updated =
Row.withSchema(tableSchema).addValues("row-1-updated").build();
+ File partFile1 = new File(tableDir, "part-00001.parquet");
+ byte[] partBytes1 =
+ writeParquetFile(partFile1, tableSchema,
java.util.Arrays.asList(tableRow1Updated));
+
+ writeCommit(
+ logDir,
+ 1L,
+ 200000000000L,
+ "part-00001.parquet",
+ partBytes1.length,
+ "part-00000.parquet",
+ "change-00000.parquet",
+ changeBytes0.length,
+ false);
+
+ // 3. Write parquet files for Version 2 (insert-only commit)
+ Row tableRow3 = Row.withSchema(tableSchema).addValues("row-3").build();
+ File partFile2 = new File(tableDir, "part-00002.parquet");
+ byte[] partBytes2 =
+ writeParquetFile(partFile2, tableSchema,
java.util.Arrays.asList(tableRow3));
+ writeCommit(
+ logDir, 2L, 300000000000L, "part-00002.parquet", partBytes2.length,
null, null, 0L, false);
+
+ // Test 1: Read changes between start version 0 and end version 2
+ PCollection<Row> outputVersions =
+ readPipeline.apply(
+ "Read Changes Version Range",
+ DeltaIO.readChanges()
+ .from(tableDir.getAbsolutePath())
+ .withStartVersion(0L)
+ .withEndVersion(2L));
Review Comment:
This test reads the whole range, so it can still pass if the range handling
logic happened to be incorrect. We should have start version > 0 and/or end
version < 2
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java:
##########
@@ -0,0 +1,353 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import static io.delta.kernel.internal.DeltaErrors.wrapEngineException;
+
+import io.delta.kernel.Scan;
+import io.delta.kernel.data.ColumnVector;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.MapValue;
+import io.delta.kernel.defaults.engine.DefaultEngine;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.engine.FileReadResult;
+import io.delta.kernel.expressions.ExpressionEvaluator;
+import io.delta.kernel.expressions.Literal;
+import io.delta.kernel.internal.InternalScanFileUtils;
+import io.delta.kernel.internal.data.GenericRow;
+import io.delta.kernel.internal.data.ScanStateRow;
+import io.delta.kernel.internal.util.Utils;
+import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.types.LongType;
+import io.delta.kernel.types.StringType;
+import io.delta.kernel.types.StructField;
+import io.delta.kernel.types.StructType;
+import io.delta.kernel.types.TimestampType;
+import io.delta.kernel.utils.CloseableIterator;
+import io.delta.kernel.utils.FileStatus;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.apache.hadoop.conf.Configuration;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and
reads Change Data Feed
+ * files, converting rows to Beam Rows.
+ */
[email protected]
+class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> {
+ @Nullable Map<String, String> hadoopConfig;
+ private transient @Nullable Engine engine;
+ private transient @Nullable Configuration conf;
+
+ public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) {
+ this.hadoopConfig = hadoopConfig;
+ }
+
+ private synchronized Configuration getConfiguration() {
+ Configuration localConf = conf;
+ if (localConf == null) {
+ localConf = new Configuration();
+ if (hadoopConfig != null) {
+ for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) {
+ localConf.set(entry.getKey(), entry.getValue());
+ }
+ }
+ conf = localConf;
+ }
+ return localConf;
+ }
+
+ private List<Long> getRowGroupSizes(DeltaCDCReadTask task) {
+ return task.getRowGroupSizes();
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) {
+ List<Long> rowGroupSizes = getRowGroupSizes(task);
+ return new OffsetRange(0L, rowGroupSizes.size());
+ }
+
+ @NewTracker
+ public DeltaReadTaskTracker newTracker(
+ @Restriction OffsetRange restriction, @Element DeltaCDCReadTask task) {
+ return new DeltaReadTaskTracker(restriction, getRowGroupSizes(task));
+ }
+
+ @Setup
+ public void setUp() {
+ engine = DefaultEngine.create(getConfiguration());
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element DeltaCDCReadTask task,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<Row> out)
+ throws Exception {
+
+ Engine currentEngine = engine;
+ if (currentEngine == null) {
+ throw new IllegalArgumentException("Expected the engine to not be null");
+ }
+
+ SerializableRow originalScanStateRow = task.getScanStateRow();
+ StructType logicalTableSchema =
ScanStateRow.getLogicalSchema(originalScanStateRow);
+ Schema publicBeamSchema =
DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
+ StructType physicalTableSchema =
ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow);
+
+ StructType scanStateSchema = originalScanStateRow.getSchema();
+
+ // 1. Build modified scanState and scanFile rows depending on whether we
read a CDC file or ADD
+ // file.
+ io.delta.kernel.data.Row scanStateRow;
+ StructType readPhysicalSchema;
+ StructType readLogicalSchema;
+ Schema beamSchema;
+
+ if (task.isCDC()) {
+ readLogicalSchema = appendCDFColumns(logicalTableSchema);
+ readPhysicalSchema = appendCDFColumns(physicalTableSchema);
+ beamSchema = DeltaIO.ReadRows.convertToBeamSchema(readLogicalSchema);
+
+ HashMap<Integer, Object> valueMap = new HashMap<>();
+
+ // Tracking row level lineage is not needed.
+ Map<String, String> config =
+ new HashMap<>(ScanStateRow.getConfiguration(originalScanStateRow));
+ config.put("delta.enableRowTracking", "false");
+
+ valueMap.put(
+ scanStateSchema.indexOf("configuration"),
VectorUtils.stringStringMapValue(config));
+ valueMap.put(scanStateSchema.indexOf("logicalSchemaJson"),
readLogicalSchema.toJson());
+ valueMap.put(scanStateSchema.indexOf("physicalSchemaJson"),
readPhysicalSchema.toJson());
+ valueMap.put(
+ scanStateSchema.indexOf("partitionColumns"),
+
originalScanStateRow.getArray(scanStateSchema.indexOf("partitionColumns")));
+ valueMap.put(
+ scanStateSchema.indexOf("minReaderVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minReaderVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("minWriterVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minWriterVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("tablePath"),
+
originalScanStateRow.getString(scanStateSchema.indexOf("tablePath")));
+
+ scanStateRow = new ScanStateRow(valueMap);
+ } else {
+ // For ADD files, we read the table schema and append the CDF columns
manually afterwards.
+ // readLogicalSchema = logicalTableSchema;
+ readPhysicalSchema = physicalTableSchema;
+ beamSchema =
DeltaIO.ReadRows.convertToBeamSchema(appendCDFColumns(logicalTableSchema));
+ scanStateRow = originalScanStateRow;
+ }
+
+ io.delta.kernel.data.Row scanFileRow =
+ generateScanFileRow(task.getPath(), task.getPartitionValues());
+ FileStatus fileStatus = FileStatus.of(task.getPath(), task.getSize(),
task.getTimestamp());
+
+ BeamParquetHandler parquetHandler =
+ new BeamParquetHandler(getConfiguration(),
currentEngine.getParquetHandler(), tracker);
+ BeamEngine beamEngine = new BeamEngine(currentEngine, parquetHandler);
+
+ long currentStartRgIndex = 0L;
+
+ try (CloseableIterator<FileReadResult> fileReadResults =
+ parquetHandler.readParquetFiles(
+ Utils.singletonCloseableIterator(fileStatus),
+ readPhysicalSchema,
+ Optional.empty(),
+ currentStartRgIndex)) {
+
+ CloseableIterator<ColumnarBatch> physicalData =
+ new CloseableIterator<ColumnarBatch>() {
+ @Override
+ public void close() throws java.io.IOException {}
+
+ @Override
+ public boolean hasNext() {
+ return fileReadResults.hasNext();
+ }
+
+ @Override
+ public ColumnarBatch next() {
+ return fileReadResults.next().getData();
+ }
+ };
+
+ try (CloseableIterator<FilteredColumnarBatch> logicalBatches =
+ Scan.transformPhysicalData(beamEngine, scanStateRow, scanFileRow,
physicalData)) {
+
+ while (logicalBatches.hasNext()) {
+ FilteredColumnarBatch batch = logicalBatches.next();
+ ColumnarBatch logicalBatch = batch.getData();
+
+ if (!task.isCDC()) {
+ // For ADD files, we need to append the constant CDF columns:
+ // _change_type = "insert", _commit_version = task.version,
_commit_timestamp =
+ // task.timestamp
+ logicalBatch =
+ appendConstantCDFColumns(
+ currentEngine, logicalBatch, task.getVersion(),
task.getTimestamp());
+ }
+
+ try (CloseableIterator<io.delta.kernel.data.Row> logicalRows =
logicalBatch.getRows()) {
+ while (logicalRows.hasNext()) {
+ io.delta.kernel.data.Row deltaRow = logicalRows.next();
+ Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema);
+ String changeType = beamRow.getString("_change_type");
+ if (changeType == null) {
+ throw new IllegalStateException("Field _change_type must not
be null.");
+ }
+ ValueKind kind = getValueKind(changeType);
+ Row publicRow = projectRow(beamRow, publicBeamSchema);
+ out.builder(publicRow).setValueKind(kind).output();
+ }
+ }
+ }
+ }
+ }
+
+ return ProcessContinuation.stop();
+ }
+
+ private static Row projectRow(Row row, Schema targetSchema) {
+ Row.Builder builder = Row.withSchema(targetSchema);
+ for (Schema.Field field : targetSchema.getFields()) {
+ builder.addValue(row.getValue(field.getName()));
+ }
+ return builder.build();
+ }
Review Comment:
If schemas are equal we can just return the original row
##########
sdks/java/io/delta/src/test/java/org/apache/beam/sdk/io/delta/DeltaIOTest.java:
##########
@@ -810,4 +815,263 @@ public boolean tryClaim(Long i) {
}
}
}
+
+ @Test
+ public void testReadChanges() throws Exception {
+ File tableDir = tempFolder.newFolder("delta-table-changes");
+ File logDir = new File(tableDir, "_delta_log");
+ logDir.mkdirs();
+
+ // 1. Write parquet files for Version 0 (insert-only commit)
+ Schema tableSchema = Schema.builder().addField("name",
Schema.FieldType.STRING).build();
+ Row tableRow1 = Row.withSchema(tableSchema).addValues("row-1").build();
+ Row tableRow2 = Row.withSchema(tableSchema).addValues("row-2").build();
+
+ File partFile = new File(tableDir, "part-00000.parquet");
+ byte[] partBytes =
+ writeParquetFile(partFile, tableSchema,
java.util.Arrays.asList(tableRow1, tableRow2));
+
+ writeCommit(
+ logDir, 0L, 100000000000L, "part-00000.parquet", partBytes.length,
null, null, 0L, true);
+
+ // 2. Write cdc parquet file for Version 1 (commit with cdc actions)
+ Schema cdcWriteSchema =
+ Schema.builder()
+ .addField("name", Schema.FieldType.STRING)
+ .addField("_change_type", Schema.FieldType.STRING)
+ .addField("_commit_version", Schema.FieldType.INT64)
+ .addField("_commit_timestamp", Schema.FieldType.DATETIME)
+ .build();
+
+ Row cdcRow1 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-1", "update_preimage", 1L, new
Instant(123456789000L))
+ .build();
+ Row cdcRow2 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-1-updated", "update_postimage", 1L, new
Instant(123456789000L))
+ .build();
+ Row cdcRow3 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-2", "delete", 1L, new Instant(123456789000L))
+ .build();
+
+ File changeFile = new File(tableDir, "change-00000.parquet");
+ byte[] changeBytes =
+ writeParquetFile(
+ changeFile, cdcWriteSchema, java.util.Arrays.asList(cdcRow1,
cdcRow2, cdcRow3));
+
+ writeCommit(
+ logDir,
+ 1L,
+ 200000000000L,
+ null,
+ 0L,
+ null,
+ "change-00000.parquet",
+ changeBytes.length,
+ false);
+
+ // 3. Read CDF data from table using ReadChanges
+ PCollection<Row> output =
+ readPipeline.apply(
+
DeltaIO.readChanges().from(tableDir.getAbsolutePath()).withStartVersion(0L));
+
+ PCollection<String> formattedOutput =
+ output.apply("Format ValueKind and Row", ParDo.of(new
FormatValueKindAndRow()));
+
+ PAssert.that(formattedOutput)
+ .containsInAnyOrder(
+ "INSERT:row-1",
+ "INSERT:row-2",
+ "UPDATE_BEFORE:row-1",
+ "UPDATE_AFTER:row-1-updated",
+ "DELETE:row-2");
+
+ readPipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testReadChangesRanges() throws Exception {
+ File tableDir = tempFolder.newFolder("delta-table-changes-ranges");
+ File logDir = new File(tableDir, "_delta_log");
+ logDir.mkdirs();
+
+ Schema tableSchema = Schema.builder().addField("name",
Schema.FieldType.STRING).build();
+
+ // 1. Write parquet files for Version 0 (insert-only commit)
+ Row tableRow1 = Row.withSchema(tableSchema).addValues("row-1").build();
+ Row tableRow2 = Row.withSchema(tableSchema).addValues("row-2").build();
+ File partFile0 = new File(tableDir, "part-00000.parquet");
+ byte[] partBytes0 =
+ writeParquetFile(partFile0, tableSchema,
java.util.Arrays.asList(tableRow1, tableRow2));
+ writeCommit(
+ logDir, 0L, 100000000000L, "part-00000.parquet", partBytes0.length,
null, null, 0L, true);
+
+ // 2. Write parquet files for Version 1 (commit with updates and deletes)
+ Schema cdcWriteSchema =
+ Schema.builder()
+ .addField("name", Schema.FieldType.STRING)
+ .addField("_change_type", Schema.FieldType.STRING)
+ .addField("_commit_version", Schema.FieldType.INT64)
+ .addField("_commit_timestamp", Schema.FieldType.DATETIME)
+ .build();
+ Row cdcRow1 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-1", "update_preimage", 1L, new
Instant(200000000000L))
+ .build();
+ Row cdcRow2 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-1-updated", "update_postimage", 1L, new
Instant(200000000000L))
+ .build();
+ Row cdcRow3 =
+ Row.withSchema(cdcWriteSchema)
+ .addValues("row-2", "delete", 1L, new Instant(200000000000L))
+ .build();
+ File changeFile0 = new File(tableDir, "change-00000.parquet");
+ byte[] changeBytes0 =
+ writeParquetFile(
+ changeFile0, cdcWriteSchema, java.util.Arrays.asList(cdcRow1,
cdcRow2, cdcRow3));
+
+ Row tableRow1Updated =
Row.withSchema(tableSchema).addValues("row-1-updated").build();
+ File partFile1 = new File(tableDir, "part-00001.parquet");
+ byte[] partBytes1 =
+ writeParquetFile(partFile1, tableSchema,
java.util.Arrays.asList(tableRow1Updated));
+
+ writeCommit(
+ logDir,
+ 1L,
+ 200000000000L,
+ "part-00001.parquet",
+ partBytes1.length,
+ "part-00000.parquet",
+ "change-00000.parquet",
+ changeBytes0.length,
+ false);
+
+ // 3. Write parquet files for Version 2 (insert-only commit)
+ Row tableRow3 = Row.withSchema(tableSchema).addValues("row-3").build();
+ File partFile2 = new File(tableDir, "part-00002.parquet");
+ byte[] partBytes2 =
+ writeParquetFile(partFile2, tableSchema,
java.util.Arrays.asList(tableRow3));
+ writeCommit(
+ logDir, 2L, 300000000000L, "part-00002.parquet", partBytes2.length,
null, null, 0L, false);
+
+ // Test 1: Read changes between start version 0 and end version 2
+ PCollection<Row> outputVersions =
+ readPipeline.apply(
+ "Read Changes Version Range",
+ DeltaIO.readChanges()
+ .from(tableDir.getAbsolutePath())
+ .withStartVersion(0L)
+ .withEndVersion(2L));
+
+ PCollection<String> formattedVersions =
+ outputVersions.apply("Format Version Output", ParDo.of(new
FormatValueKindAndRow()));
+
+ PAssert.that(formattedVersions)
+ .containsInAnyOrder(
+ "INSERT:row-1",
+ "INSERT:row-2",
+ "UPDATE_BEFORE:row-1",
+ "UPDATE_AFTER:row-1-updated",
+ "DELETE:row-2",
+ "INSERT:row-3");
+
+ // Test 2: Read changes between start timestamp (after version 0) and end
timestamp (after
+ // version 2)
+ String startTimestamp =
java.time.Instant.ofEpochMilli(150000000000L).toString();
+ String endTimestamp =
java.time.Instant.ofEpochMilli(350000000000L).toString();
+
+ PCollection<Row> outputTimestamps =
+ filteringPipeline.apply(
+ "Read Changes Timestamp Range",
+ DeltaIO.readChanges()
+ .from(tableDir.getAbsolutePath())
+ .withStartTimestamp(startTimestamp)
+ .withEndTimestamp(endTimestamp));
+
+ PCollection<String> formattedTimestamps =
+ outputTimestamps.apply("Format Timestamp Output", ParDo.of(new
FormatValueKindAndRow()));
+
+ PAssert.that(formattedTimestamps)
+ .containsInAnyOrder(
+ "UPDATE_BEFORE:row-1", "UPDATE_AFTER:row-1-updated",
"DELETE:row-2", "INSERT:row-3");
+
+ readPipeline.run().waitUntilFinish();
+ filteringPipeline.run().waitUntilFinish();
+ }
+
+ private void writeCommit(
+ File logDir,
+ long version,
+ long timestamp,
+ @Nullable String addPath,
+ long addSize,
+ @Nullable String removePath,
+ @Nullable String cdcPath,
+ long cdcSize,
+ boolean writeMetadata)
+ throws IOException {
Review Comment:
Can we commit using the DeltaLake API? to make sure it works against the
actual API and get notified if a library upgrade breaks things
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java:
##########
@@ -0,0 +1,353 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import static io.delta.kernel.internal.DeltaErrors.wrapEngineException;
+
+import io.delta.kernel.Scan;
+import io.delta.kernel.data.ColumnVector;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.MapValue;
+import io.delta.kernel.defaults.engine.DefaultEngine;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.engine.FileReadResult;
+import io.delta.kernel.expressions.ExpressionEvaluator;
+import io.delta.kernel.expressions.Literal;
+import io.delta.kernel.internal.InternalScanFileUtils;
+import io.delta.kernel.internal.data.GenericRow;
+import io.delta.kernel.internal.data.ScanStateRow;
+import io.delta.kernel.internal.util.Utils;
+import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.types.LongType;
+import io.delta.kernel.types.StringType;
+import io.delta.kernel.types.StructField;
+import io.delta.kernel.types.StructType;
+import io.delta.kernel.types.TimestampType;
+import io.delta.kernel.utils.CloseableIterator;
+import io.delta.kernel.utils.FileStatus;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.apache.hadoop.conf.Configuration;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and
reads Change Data Feed
+ * files, converting rows to Beam Rows.
+ */
[email protected]
+class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> {
+ @Nullable Map<String, String> hadoopConfig;
+ private transient @Nullable Engine engine;
+ private transient @Nullable Configuration conf;
+
+ public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) {
+ this.hadoopConfig = hadoopConfig;
+ }
+
+ private synchronized Configuration getConfiguration() {
+ Configuration localConf = conf;
+ if (localConf == null) {
+ localConf = new Configuration();
+ if (hadoopConfig != null) {
+ for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) {
+ localConf.set(entry.getKey(), entry.getValue());
+ }
+ }
+ conf = localConf;
+ }
+ return localConf;
+ }
+
+ private List<Long> getRowGroupSizes(DeltaCDCReadTask task) {
+ return task.getRowGroupSizes();
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) {
+ List<Long> rowGroupSizes = getRowGroupSizes(task);
+ return new OffsetRange(0L, rowGroupSizes.size());
+ }
+
+ @NewTracker
+ public DeltaReadTaskTracker newTracker(
+ @Restriction OffsetRange restriction, @Element DeltaCDCReadTask task) {
+ return new DeltaReadTaskTracker(restriction, getRowGroupSizes(task));
+ }
+
+ @Setup
+ public void setUp() {
+ engine = DefaultEngine.create(getConfiguration());
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element DeltaCDCReadTask task,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<Row> out)
+ throws Exception {
+
+ Engine currentEngine = engine;
+ if (currentEngine == null) {
+ throw new IllegalArgumentException("Expected the engine to not be null");
+ }
+
+ SerializableRow originalScanStateRow = task.getScanStateRow();
+ StructType logicalTableSchema =
ScanStateRow.getLogicalSchema(originalScanStateRow);
+ Schema publicBeamSchema =
DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
+ StructType physicalTableSchema =
ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow);
+
+ StructType scanStateSchema = originalScanStateRow.getSchema();
+
+ // 1. Build modified scanState and scanFile rows depending on whether we
read a CDC file or ADD
+ // file.
+ io.delta.kernel.data.Row scanStateRow;
+ StructType readPhysicalSchema;
+ StructType readLogicalSchema;
+ Schema beamSchema;
+
+ if (task.isCDC()) {
+ readLogicalSchema = appendCDFColumns(logicalTableSchema);
+ readPhysicalSchema = appendCDFColumns(physicalTableSchema);
+ beamSchema = DeltaIO.ReadRows.convertToBeamSchema(readLogicalSchema);
+
+ HashMap<Integer, Object> valueMap = new HashMap<>();
+
+ // Tracking row level lineage is not needed.
+ Map<String, String> config =
+ new HashMap<>(ScanStateRow.getConfiguration(originalScanStateRow));
+ config.put("delta.enableRowTracking", "false");
+
+ valueMap.put(
+ scanStateSchema.indexOf("configuration"),
VectorUtils.stringStringMapValue(config));
+ valueMap.put(scanStateSchema.indexOf("logicalSchemaJson"),
readLogicalSchema.toJson());
+ valueMap.put(scanStateSchema.indexOf("physicalSchemaJson"),
readPhysicalSchema.toJson());
+ valueMap.put(
+ scanStateSchema.indexOf("partitionColumns"),
+
originalScanStateRow.getArray(scanStateSchema.indexOf("partitionColumns")));
+ valueMap.put(
+ scanStateSchema.indexOf("minReaderVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minReaderVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("minWriterVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minWriterVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("tablePath"),
+
originalScanStateRow.getString(scanStateSchema.indexOf("tablePath")));
+
+ scanStateRow = new ScanStateRow(valueMap);
+ } else {
+ // For ADD files, we read the table schema and append the CDF columns
manually afterwards.
+ // readLogicalSchema = logicalTableSchema;
+ readPhysicalSchema = physicalTableSchema;
+ beamSchema =
DeltaIO.ReadRows.convertToBeamSchema(appendCDFColumns(logicalTableSchema));
+ scanStateRow = originalScanStateRow;
+ }
+
+ io.delta.kernel.data.Row scanFileRow =
+ generateScanFileRow(task.getPath(), task.getPartitionValues());
+ FileStatus fileStatus = FileStatus.of(task.getPath(), task.getSize(),
task.getTimestamp());
+
+ BeamParquetHandler parquetHandler =
+ new BeamParquetHandler(getConfiguration(),
currentEngine.getParquetHandler(), tracker);
+ BeamEngine beamEngine = new BeamEngine(currentEngine, parquetHandler);
+
+ long currentStartRgIndex = 0L;
+
+ try (CloseableIterator<FileReadResult> fileReadResults =
+ parquetHandler.readParquetFiles(
+ Utils.singletonCloseableIterator(fileStatus),
+ readPhysicalSchema,
+ Optional.empty(),
+ currentStartRgIndex)) {
+
+ CloseableIterator<ColumnarBatch> physicalData =
+ new CloseableIterator<ColumnarBatch>() {
+ @Override
+ public void close() throws java.io.IOException {}
+
+ @Override
+ public boolean hasNext() {
+ return fileReadResults.hasNext();
+ }
+
+ @Override
+ public ColumnarBatch next() {
+ return fileReadResults.next().getData();
+ }
+ };
+
+ try (CloseableIterator<FilteredColumnarBatch> logicalBatches =
+ Scan.transformPhysicalData(beamEngine, scanStateRow, scanFileRow,
physicalData)) {
+
+ while (logicalBatches.hasNext()) {
+ FilteredColumnarBatch batch = logicalBatches.next();
+ ColumnarBatch logicalBatch = batch.getData();
+
+ if (!task.isCDC()) {
+ // For ADD files, we need to append the constant CDF columns:
+ // _change_type = "insert", _commit_version = task.version,
_commit_timestamp =
+ // task.timestamp
+ logicalBatch =
+ appendConstantCDFColumns(
+ currentEngine, logicalBatch, task.getVersion(),
task.getTimestamp());
+ }
+
+ try (CloseableIterator<io.delta.kernel.data.Row> logicalRows =
logicalBatch.getRows()) {
+ while (logicalRows.hasNext()) {
+ io.delta.kernel.data.Row deltaRow = logicalRows.next();
+ Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema);
+ String changeType = beamRow.getString("_change_type");
+ if (changeType == null) {
+ throw new IllegalStateException("Field _change_type must not
be null.");
+ }
+ ValueKind kind = getValueKind(changeType);
+ Row publicRow = projectRow(beamRow, publicBeamSchema);
+ out.builder(publicRow).setValueKind(kind).output();
+ }
+ }
+ }
+ }
+ }
+
+ return ProcessContinuation.stop();
+ }
+
+ private static Row projectRow(Row row, Schema targetSchema) {
+ Row.Builder builder = Row.withSchema(targetSchema);
+ for (Schema.Field field : targetSchema.getFields()) {
+ builder.addValue(row.getValue(field.getName()));
+ }
+ return builder.build();
+ }
+
+ private static ValueKind getValueKind(String changeType) {
+ // Maps Delta CDC change types to Beam's ValueKind enum.
+ //
https://docs.delta.io/delta-change-data-feed/#what-is-the-schema-for-the-change-data-feed
+ switch (changeType) {
+ case "insert":
+ return ValueKind.INSERT;
+ case "delete":
+ return ValueKind.DELETE;
+ case "update_preimage":
+ return ValueKind.UPDATE_BEFORE;
+ case "update_postimage":
+ return ValueKind.UPDATE_AFTER;
+ default:
+ throw new IllegalArgumentException("Unsupported change type: " +
changeType);
+ }
+ }
+
+ private static StructType appendCDFColumns(StructType schema) {
+ return schema
+ .add("_change_type", StringType.STRING, false)
+ .add("_commit_version", LongType.LONG, false)
+ .add("_commit_timestamp", TimestampType.TIMESTAMP, false);
Review Comment:
What values can `_change_type` take?
Also is `_commit_version` increasingly monotonic?
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java:
##########
@@ -174,4 +180,124 @@ static Schema.FieldType convertToBeamFieldType(DataType
deltaType) {
}
}
}
+
+ @AutoValue
+ public abstract static class ReadChanges extends PTransform<PBegin,
PCollection<Row>> {
+ public abstract @Nullable String getTablePath();
+
+ public abstract @Nullable Long getStartVersion();
+
+ public abstract @Nullable String getStartTimestamp();
+
+ public abstract @Nullable Long getEndVersion();
+
+ public abstract @Nullable String getEndTimestamp();
+
+ public abstract @Nullable Map<String, String> getHadoopConfig();
+
+ abstract Builder toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setTablePath(String tablePath);
+
+ abstract Builder setStartVersion(@Nullable Long startVersion);
+
+ abstract Builder setStartTimestamp(@Nullable String startTimestamp);
+
+ abstract Builder setEndVersion(@Nullable Long endVersion);
+
+ abstract Builder setEndTimestamp(@Nullable String endTimestamp);
+
+ abstract Builder setHadoopConfig(@Nullable Map<String, String>
hadoopConfig);
+
+ abstract ReadChanges build();
+ }
+
+ public ReadChanges from(String tablePath) {
+ return toBuilder().setTablePath(tablePath).build();
+ }
+
+ public ReadChanges withStartVersion(long startVersion) {
+ return toBuilder().setStartVersion(startVersion).build();
+ }
+
+ public ReadChanges withStartTimestamp(String startTimestamp) {
+ return toBuilder().setStartTimestamp(startTimestamp).build();
+ }
+
+ public ReadChanges withEndVersion(long endVersion) {
+ return toBuilder().setEndVersion(endVersion).build();
+ }
+
+ public ReadChanges withEndTimestamp(String endTimestamp) {
+ return toBuilder().setEndTimestamp(endTimestamp).build();
+ }
+
+ public ReadChanges withConfig(Map<String, String> config) {
+ return toBuilder().setHadoopConfig(config).build();
+ }
+
+ @Override
+ public PCollection<Row> expand(PBegin input) {
+ String path = getTablePath();
+ if (path == null) {
+ throw new IllegalArgumentException("Table path must be set.");
+ }
+ if (getStartVersion() == null && getStartTimestamp() == null) {
+ throw new IllegalArgumentException("Either startVersion or
startTimestamp must be set.");
+ }
+ if (getStartVersion() != null && getStartTimestamp() != null) {
+ throw new IllegalArgumentException("Cannot set both startVersion and
startTimestamp.");
+ }
+ if (getEndVersion() != null && getEndTimestamp() != null) {
+ throw new IllegalArgumentException("Cannot set both endVersion and
endTimestamp.");
+ }
Review Comment:
I think your logic in the Create DoFn already handles this
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java:
##########
@@ -0,0 +1,293 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import io.delta.kernel.CommitRange;
+import io.delta.kernel.CommitRangeBuilder;
+import io.delta.kernel.Scan;
+import io.delta.kernel.Snapshot;
+import io.delta.kernel.Table;
+import io.delta.kernel.TableManager;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.data.Row;
+import io.delta.kernel.defaults.engine.DefaultEngine;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.internal.DeltaLogActionUtils.DeltaAction;
+import io.delta.kernel.internal.TableImpl;
+import io.delta.kernel.internal.actions.AddCDCFile;
+import io.delta.kernel.internal.actions.AddFile;
+import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.utils.CloseableIterator;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.hadoop.conf.Configuration;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** A DoFn that reads the Delta log and plans read tasks for Change Data Feed.
*/
+class CreateCDCReadTasksDoFn extends DoFn<String, DeltaCDCReadTask> {
+ private static final long MAX_TASK_SIZE_BYTES = 1024L * 1024L * 1024L; // 1
GB
+ private final @Nullable Map<String, String> hadoopConfig;
+ private final @Nullable Long startVersion;
+ private final @Nullable String startTimestamp;
+ private final @Nullable Long endVersion;
+ private final @Nullable String endTimestamp;
+
+ public CreateCDCReadTasksDoFn(
+ @Nullable Map<String, String> hadoopConfig,
+ @Nullable Long startVersion,
+ @Nullable String startTimestamp,
+ @Nullable Long endVersion,
+ @Nullable String endTimestamp) {
+ this.hadoopConfig = hadoopConfig;
+ this.startVersion = startVersion;
+ this.startTimestamp = startTimestamp;
+ this.endVersion = endVersion;
+ this.endTimestamp = endTimestamp;
+ }
+
+ @ProcessElement
+ public void processElement(@Element String tablePath,
OutputReceiver<DeltaCDCReadTask> out)
+ throws Exception {
+ Configuration conf = new Configuration();
+ if (hadoopConfig != null) {
+ for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) {
+ conf.set(entry.getKey(), entry.getValue());
+ }
+ }
+ Engine engine = DefaultEngine.create(conf);
+ Table table = Table.forPath(engine, tablePath);
+ TableImpl tableImpl = (TableImpl) table;
+
+ // 1. Resolve starting and ending versions
+ long resolvedStartVersion;
+ if (startVersion != null) {
+ resolvedStartVersion = startVersion;
+ } else if (startTimestamp != null) {
+ long startMillis = Instant.parse(startTimestamp).toEpochMilli();
+ resolvedStartVersion = tableImpl.getVersionAtOrAfterTimestamp(engine,
startMillis);
+ } else {
+ throw new IllegalArgumentException("Starting version or timestamp must
be specified.");
+ }
+
+ long resolvedEndVersion;
+ if (endVersion != null) {
+ resolvedEndVersion = endVersion;
+ } else if (endTimestamp != null) {
+ long endMillis = Instant.parse(endTimestamp).toEpochMilli();
+ resolvedEndVersion = tableImpl.getVersionBeforeOrAtTimestamp(engine,
endMillis);
+ } else {
+ resolvedEndVersion = table.getLatestSnapshot(engine).getVersion();
+ }
+
+ if (resolvedStartVersion > resolvedEndVersion) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Resolved start version %d is greater than resolved end version
%d",
+ resolvedStartVersion, resolvedEndVersion));
+ }
+
+ // 2. Load snapshot at resolvedEndVersion to get the scanStateRow
+ // We use endVersion's schema because it represents the latest schema in
the
+ // read range
+ // which handles schema evolution (older files will just lack new columns).
+ Snapshot endSnapshot = table.getSnapshotAsOfVersion(engine,
resolvedEndVersion);
+ Scan scan = endSnapshot.getScanBuilder().build();
+ Row scanState = scan.getScanState(engine);
+ SerializableRow serializableScanState = new SerializableRow(scanState);
+
+ // 3. Load snapshot at resolvedStartVersion to initialize the CommitRange
+ Snapshot startSnapshot = table.getSnapshotAsOfVersion(engine,
resolvedStartVersion);
+
+ CommitRangeBuilder rangeBuilder =
+ TableManager.loadCommitRange(
+ tablePath,
CommitRangeBuilder.CommitBoundary.atVersion(resolvedStartVersion));
+
rangeBuilder.withEndBoundary(CommitRangeBuilder.CommitBoundary.atVersion(resolvedEndVersion));
+ CommitRange range = rangeBuilder.build(engine);
+
+ // We need both CDC and ADD actions.
+ // If a commit version has CDC files, we only read CDC files.
+ // If a commit version has no CDC files, we read ADD files (inserts).
+ Set<DeltaAction> actionSet = new HashSet<>();
+ actionSet.add(DeltaAction.CDC);
+ actionSet.add(DeltaAction.ADD);
+
+ // 4. Iterate over commits in the range and group actions by version
+ try (CloseableIterator<ColumnarBatch> batchIter =
+ range.getActions(engine, startSnapshot, actionSet)) {
+ Map<Long, CommitActionsInfo> commitActionsMap = new HashMap<>();
+
+ while (batchIter.hasNext()) {
+ ColumnarBatch batch = batchIter.next();
+ int versionIdx = batch.getSchema().indexOf("version");
+ int timestampIdx = batch.getSchema().indexOf("timestamp");
+ int cdcIdx = batch.getSchema().indexOf("cdc");
+ int addIdx = batch.getSchema().indexOf("add");
+
+ for (int i = 0; i < batch.getSize(); i++) {
+ long version = batch.getColumnVector(versionIdx).getLong(i);
+ long timestamp = batch.getColumnVector(timestampIdx).getLong(i);
+
+ CommitActionsInfo info =
+ commitActionsMap.computeIfAbsent(
+ version, k -> new CommitActionsInfo(version, timestamp));
+
+ if (cdcIdx >= 0 && !batch.getColumnVector(cdcIdx).isNullAt(i)) {
+ Row cdcRow =
+ (Row)
+ VectorUtils.getValueAsObject(
+ batch.getColumnVector(cdcIdx),
+ batch.getSchema().at(cdcIdx).getDataType(),
+ i);
+ info.cdcRows.add(cdcRow);
+ }
+ if (addIdx >= 0 && !batch.getColumnVector(addIdx).isNullAt(i)) {
+ Row addRow =
+ (Row)
+ VectorUtils.getValueAsObject(
+ batch.getColumnVector(addIdx),
+ batch.getSchema().at(addIdx).getDataType(),
+ i);
+ AddFile addFile = new AddFile(addRow);
+ // Only consider add files that change data (ignore OPTIMIZE etc.)
+ if (addFile.getDataChange()) {
+ info.addRows.add(addRow);
+ }
+ }
+ }
+ }
+
+ // 5. Emit tasks for each version
+ List<DeltaCDCReadTask> currentGroup = new ArrayList<>();
+ long currentGroupSize = 0L;
+
+ // Sort versions to process them in order
+ List<Long> versions = new ArrayList<>(commitActionsMap.keySet());
+ Collections.sort(versions);
+
+ for (long version : versions) {
+ CommitActionsInfo info = commitActionsMap.get(version);
+ if (info == null) {
+ throw new IllegalStateException("CommitActionsInfo was not found for
version " + version);
+ }
+ boolean hasCDC = !info.cdcRows.isEmpty();
+
+ List<Row> rowsToProcess = hasCDC ? info.cdcRows : info.addRows;
+ boolean isCDC = hasCDC;
+
+ for (Row fileRow : rowsToProcess) {
+ String relPath;
+ long size;
+ Map<String, String> partitionValues;
+
+ if (isCDC) {
+ relPath =
fileRow.getString(AddCDCFile.FULL_SCHEMA.indexOf("path"));
+ size = fileRow.getLong(AddCDCFile.FULL_SCHEMA.indexOf("size"));
+ partitionValues =
+ VectorUtils.toJavaMap(
+
fileRow.getMap(AddCDCFile.FULL_SCHEMA.indexOf("partitionValues")));
+ } else {
+ AddFile addFile = new AddFile(fileRow);
+ relPath = addFile.getPath();
+ size = addFile.getSize();
+ partitionValues =
VectorUtils.toJavaMap(addFile.getPartitionValues());
+ }
+
+ String fullPath = new org.apache.hadoop.fs.Path(tablePath,
relPath).toString();
+ List<Long> rowGroupSizes = getRowGroupSizes(fullPath, conf);
+
+ DeltaCDCReadTask task =
+ new DeltaCDCReadTask(
+ fullPath,
+ size,
+ partitionValues,
+ info.version,
+ info.timestamp,
+ isCDC,
+ rowGroupSizes,
+ serializableScanState);
+
+ if (size >= MAX_TASK_SIZE_BYTES) {
+ if (!currentGroup.isEmpty()) {
+ emitGroup(currentGroup, out);
+ currentGroup = new ArrayList<>();
+ currentGroupSize = 0L;
+ }
+ out.output(task);
+ } else {
+ if (currentGroupSize + size > MAX_TASK_SIZE_BYTES) {
+ emitGroup(currentGroup, out);
+ currentGroup = new ArrayList<>();
+ currentGroup.add(task);
+ currentGroupSize = size;
+ } else {
+ currentGroup.add(task);
+ currentGroupSize += size;
+ }
+ }
+ }
+ }
+
+ if (!currentGroup.isEmpty()) {
+ emitGroup(currentGroup, out);
+ }
+ }
+ }
+
+ private void emitGroup(List<DeltaCDCReadTask> group,
OutputReceiver<DeltaCDCReadTask> out) {
+ for (DeltaCDCReadTask task : group) {
+ out.output(task);
+ }
+ }
+
+ private List<Long> getRowGroupSizes(String pathStr, Configuration conf) {
+ List<Long> sizes = new ArrayList<>();
+ try {
+ org.apache.hadoop.fs.Path hadoopPath = new
org.apache.hadoop.fs.Path(pathStr);
+ org.apache.parquet.hadoop.metadata.ParquetMetadata metadata =
+ org.apache.parquet.hadoop.ParquetFileReader.readFooter(
+ conf,
+ hadoopPath,
+
org.apache.parquet.format.converter.ParquetMetadataConverter.NO_FILTER);
+ for (org.apache.parquet.hadoop.metadata.BlockMetaData block :
metadata.getBlocks()) {
+ sizes.add(block.getTotalByteSize());
+ }
+ } catch (java.io.IOException e) {
+ throw new RuntimeException("Failed to read Parquet footer for " +
pathStr, e);
+ }
+ return sizes;
+ }
+
+ private static class CommitActionsInfo {
+ final long version;
+ final long timestamp;
+ final List<Row> cdcRows = new ArrayList<>();
+ final List<Row> addRows = new ArrayList<>();
Review Comment:
Naming is throwing me off a little. So these aren't data rows right? Each
row is metadata about a particular CDC or INSERT task?
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java:
##########
@@ -0,0 +1,353 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import static io.delta.kernel.internal.DeltaErrors.wrapEngineException;
+
+import io.delta.kernel.Scan;
+import io.delta.kernel.data.ColumnVector;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.MapValue;
+import io.delta.kernel.defaults.engine.DefaultEngine;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.engine.FileReadResult;
+import io.delta.kernel.expressions.ExpressionEvaluator;
+import io.delta.kernel.expressions.Literal;
+import io.delta.kernel.internal.InternalScanFileUtils;
+import io.delta.kernel.internal.data.GenericRow;
+import io.delta.kernel.internal.data.ScanStateRow;
+import io.delta.kernel.internal.util.Utils;
+import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.types.LongType;
+import io.delta.kernel.types.StringType;
+import io.delta.kernel.types.StructField;
+import io.delta.kernel.types.StructType;
+import io.delta.kernel.types.TimestampType;
+import io.delta.kernel.utils.CloseableIterator;
+import io.delta.kernel.utils.FileStatus;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.apache.hadoop.conf.Configuration;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and
reads Change Data Feed
+ * files, converting rows to Beam Rows.
+ */
[email protected]
+class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> {
+ @Nullable Map<String, String> hadoopConfig;
+ private transient @Nullable Engine engine;
+ private transient @Nullable Configuration conf;
+
+ public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) {
+ this.hadoopConfig = hadoopConfig;
+ }
+
+ private synchronized Configuration getConfiguration() {
+ Configuration localConf = conf;
+ if (localConf == null) {
+ localConf = new Configuration();
+ if (hadoopConfig != null) {
+ for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) {
+ localConf.set(entry.getKey(), entry.getValue());
+ }
+ }
+ conf = localConf;
+ }
+ return localConf;
+ }
+
+ private List<Long> getRowGroupSizes(DeltaCDCReadTask task) {
+ return task.getRowGroupSizes();
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) {
+ List<Long> rowGroupSizes = getRowGroupSizes(task);
+ return new OffsetRange(0L, rowGroupSizes.size());
+ }
+
+ @NewTracker
+ public DeltaReadTaskTracker newTracker(
+ @Restriction OffsetRange restriction, @Element DeltaCDCReadTask task) {
+ return new DeltaReadTaskTracker(restriction, getRowGroupSizes(task));
+ }
+
+ @Setup
+ public void setUp() {
+ engine = DefaultEngine.create(getConfiguration());
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element DeltaCDCReadTask task,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<Row> out)
+ throws Exception {
+
+ Engine currentEngine = engine;
+ if (currentEngine == null) {
+ throw new IllegalArgumentException("Expected the engine to not be null");
+ }
+
+ SerializableRow originalScanStateRow = task.getScanStateRow();
+ StructType logicalTableSchema =
ScanStateRow.getLogicalSchema(originalScanStateRow);
+ Schema publicBeamSchema =
DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
+ StructType physicalTableSchema =
ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow);
+
+ StructType scanStateSchema = originalScanStateRow.getSchema();
+
+ // 1. Build modified scanState and scanFile rows depending on whether we
read a CDC file or ADD
+ // file.
+ io.delta.kernel.data.Row scanStateRow;
+ StructType readPhysicalSchema;
+ StructType readLogicalSchema;
+ Schema beamSchema;
+
+ if (task.isCDC()) {
+ readLogicalSchema = appendCDFColumns(logicalTableSchema);
+ readPhysicalSchema = appendCDFColumns(physicalTableSchema);
+ beamSchema = DeltaIO.ReadRows.convertToBeamSchema(readLogicalSchema);
+
+ HashMap<Integer, Object> valueMap = new HashMap<>();
+
+ // Tracking row level lineage is not needed.
+ Map<String, String> config =
+ new HashMap<>(ScanStateRow.getConfiguration(originalScanStateRow));
+ config.put("delta.enableRowTracking", "false");
+
+ valueMap.put(
+ scanStateSchema.indexOf("configuration"),
VectorUtils.stringStringMapValue(config));
+ valueMap.put(scanStateSchema.indexOf("logicalSchemaJson"),
readLogicalSchema.toJson());
+ valueMap.put(scanStateSchema.indexOf("physicalSchemaJson"),
readPhysicalSchema.toJson());
+ valueMap.put(
+ scanStateSchema.indexOf("partitionColumns"),
+
originalScanStateRow.getArray(scanStateSchema.indexOf("partitionColumns")));
+ valueMap.put(
+ scanStateSchema.indexOf("minReaderVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minReaderVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("minWriterVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minWriterVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("tablePath"),
+
originalScanStateRow.getString(scanStateSchema.indexOf("tablePath")));
+
+ scanStateRow = new ScanStateRow(valueMap);
+ } else {
+ // For ADD files, we read the table schema and append the CDF columns
manually afterwards.
+ // readLogicalSchema = logicalTableSchema;
+ readPhysicalSchema = physicalTableSchema;
+ beamSchema =
DeltaIO.ReadRows.convertToBeamSchema(appendCDFColumns(logicalTableSchema));
+ scanStateRow = originalScanStateRow;
+ }
+
+ io.delta.kernel.data.Row scanFileRow =
+ generateScanFileRow(task.getPath(), task.getPartitionValues());
+ FileStatus fileStatus = FileStatus.of(task.getPath(), task.getSize(),
task.getTimestamp());
+
+ BeamParquetHandler parquetHandler =
+ new BeamParquetHandler(getConfiguration(),
currentEngine.getParquetHandler(), tracker);
+ BeamEngine beamEngine = new BeamEngine(currentEngine, parquetHandler);
+
+ long currentStartRgIndex = 0L;
+
+ try (CloseableIterator<FileReadResult> fileReadResults =
+ parquetHandler.readParquetFiles(
+ Utils.singletonCloseableIterator(fileStatus),
+ readPhysicalSchema,
+ Optional.empty(),
+ currentStartRgIndex)) {
+
+ CloseableIterator<ColumnarBatch> physicalData =
+ new CloseableIterator<ColumnarBatch>() {
+ @Override
+ public void close() throws java.io.IOException {}
+
+ @Override
+ public boolean hasNext() {
+ return fileReadResults.hasNext();
+ }
+
+ @Override
+ public ColumnarBatch next() {
+ return fileReadResults.next().getData();
+ }
+ };
+
+ try (CloseableIterator<FilteredColumnarBatch> logicalBatches =
+ Scan.transformPhysicalData(beamEngine, scanStateRow, scanFileRow,
physicalData)) {
+
+ while (logicalBatches.hasNext()) {
+ FilteredColumnarBatch batch = logicalBatches.next();
+ ColumnarBatch logicalBatch = batch.getData();
+
+ if (!task.isCDC()) {
+ // For ADD files, we need to append the constant CDF columns:
+ // _change_type = "insert", _commit_version = task.version,
_commit_timestamp =
+ // task.timestamp
+ logicalBatch =
+ appendConstantCDFColumns(
+ currentEngine, logicalBatch, task.getVersion(),
task.getTimestamp());
+ }
+
+ try (CloseableIterator<io.delta.kernel.data.Row> logicalRows =
logicalBatch.getRows()) {
+ while (logicalRows.hasNext()) {
+ io.delta.kernel.data.Row deltaRow = logicalRows.next();
+ Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema);
+ String changeType = beamRow.getString("_change_type");
+ if (changeType == null) {
+ throw new IllegalStateException("Field _change_type must not
be null.");
+ }
+ ValueKind kind = getValueKind(changeType);
+ Row publicRow = projectRow(beamRow, publicBeamSchema);
+ out.builder(publicRow).setValueKind(kind).output();
+ }
+ }
+ }
+ }
+ }
+
+ return ProcessContinuation.stop();
+ }
+
+ private static Row projectRow(Row row, Schema targetSchema) {
+ Row.Builder builder = Row.withSchema(targetSchema);
+ for (Schema.Field field : targetSchema.getFields()) {
+ builder.addValue(row.getValue(field.getName()));
+ }
+ return builder.build();
+ }
+
+ private static ValueKind getValueKind(String changeType) {
+ // Maps Delta CDC change types to Beam's ValueKind enum.
+ //
https://docs.delta.io/delta-change-data-feed/#what-is-the-schema-for-the-change-data-feed
+ switch (changeType) {
+ case "insert":
+ return ValueKind.INSERT;
+ case "delete":
+ return ValueKind.DELETE;
+ case "update_preimage":
+ return ValueKind.UPDATE_BEFORE;
+ case "update_postimage":
+ return ValueKind.UPDATE_AFTER;
+ default:
+ throw new IllegalArgumentException("Unsupported change type: " +
changeType);
+ }
+ }
+
+ private static StructType appendCDFColumns(StructType schema) {
+ return schema
+ .add("_change_type", StringType.STRING, false)
+ .add("_commit_version", LongType.LONG, false)
+ .add("_commit_timestamp", TimestampType.TIMESTAMP, false);
Review Comment:
hmmm I see that delta uses `update_preimage` and `update_postimage` for
UPDATE_BEFORE/AFTER. This might not play too well with Iceberg CDC sink but I
think should be okay because we can rely on the ValueKind metadata instead
##########
sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java:
##########
@@ -0,0 +1,353 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.delta;
+
+import static io.delta.kernel.internal.DeltaErrors.wrapEngineException;
+
+import io.delta.kernel.Scan;
+import io.delta.kernel.data.ColumnVector;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.MapValue;
+import io.delta.kernel.defaults.engine.DefaultEngine;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.engine.FileReadResult;
+import io.delta.kernel.expressions.ExpressionEvaluator;
+import io.delta.kernel.expressions.Literal;
+import io.delta.kernel.internal.InternalScanFileUtils;
+import io.delta.kernel.internal.data.GenericRow;
+import io.delta.kernel.internal.data.ScanStateRow;
+import io.delta.kernel.internal.util.Utils;
+import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.types.LongType;
+import io.delta.kernel.types.StringType;
+import io.delta.kernel.types.StructField;
+import io.delta.kernel.types.StructType;
+import io.delta.kernel.types.TimestampType;
+import io.delta.kernel.utils.CloseableIterator;
+import io.delta.kernel.utils.FileStatus;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.apache.hadoop.conf.Configuration;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A Splittable DoFn that processes {@link DeltaCDCReadTask} elements and
reads Change Data Feed
+ * files, converting rows to Beam Rows.
+ */
[email protected]
+class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> {
+ @Nullable Map<String, String> hadoopConfig;
+ private transient @Nullable Engine engine;
+ private transient @Nullable Configuration conf;
+
+ public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) {
+ this.hadoopConfig = hadoopConfig;
+ }
+
+ private synchronized Configuration getConfiguration() {
+ Configuration localConf = conf;
+ if (localConf == null) {
+ localConf = new Configuration();
+ if (hadoopConfig != null) {
+ for (Map.Entry<String, String> entry : hadoopConfig.entrySet()) {
+ localConf.set(entry.getKey(), entry.getValue());
+ }
+ }
+ conf = localConf;
+ }
+ return localConf;
+ }
+
+ private List<Long> getRowGroupSizes(DeltaCDCReadTask task) {
+ return task.getRowGroupSizes();
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element DeltaCDCReadTask task) {
+ List<Long> rowGroupSizes = getRowGroupSizes(task);
+ return new OffsetRange(0L, rowGroupSizes.size());
+ }
+
+ @NewTracker
+ public DeltaReadTaskTracker newTracker(
+ @Restriction OffsetRange restriction, @Element DeltaCDCReadTask task) {
+ return new DeltaReadTaskTracker(restriction, getRowGroupSizes(task));
+ }
+
+ @Setup
+ public void setUp() {
+ engine = DefaultEngine.create(getConfiguration());
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ @Element DeltaCDCReadTask task,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<Row> out)
+ throws Exception {
+
+ Engine currentEngine = engine;
+ if (currentEngine == null) {
+ throw new IllegalArgumentException("Expected the engine to not be null");
+ }
+
+ SerializableRow originalScanStateRow = task.getScanStateRow();
+ StructType logicalTableSchema =
ScanStateRow.getLogicalSchema(originalScanStateRow);
+ Schema publicBeamSchema =
DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
+ StructType physicalTableSchema =
ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow);
+
+ StructType scanStateSchema = originalScanStateRow.getSchema();
+
+ // 1. Build modified scanState and scanFile rows depending on whether we
read a CDC file or ADD
+ // file.
+ io.delta.kernel.data.Row scanStateRow;
+ StructType readPhysicalSchema;
+ StructType readLogicalSchema;
+ Schema beamSchema;
+
+ if (task.isCDC()) {
+ readLogicalSchema = appendCDFColumns(logicalTableSchema);
+ readPhysicalSchema = appendCDFColumns(physicalTableSchema);
+ beamSchema = DeltaIO.ReadRows.convertToBeamSchema(readLogicalSchema);
+
+ HashMap<Integer, Object> valueMap = new HashMap<>();
+
+ // Tracking row level lineage is not needed.
+ Map<String, String> config =
+ new HashMap<>(ScanStateRow.getConfiguration(originalScanStateRow));
+ config.put("delta.enableRowTracking", "false");
+
+ valueMap.put(
+ scanStateSchema.indexOf("configuration"),
VectorUtils.stringStringMapValue(config));
+ valueMap.put(scanStateSchema.indexOf("logicalSchemaJson"),
readLogicalSchema.toJson());
+ valueMap.put(scanStateSchema.indexOf("physicalSchemaJson"),
readPhysicalSchema.toJson());
+ valueMap.put(
+ scanStateSchema.indexOf("partitionColumns"),
+
originalScanStateRow.getArray(scanStateSchema.indexOf("partitionColumns")));
+ valueMap.put(
+ scanStateSchema.indexOf("minReaderVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minReaderVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("minWriterVersion"),
+
originalScanStateRow.getInt(scanStateSchema.indexOf("minWriterVersion")));
+ valueMap.put(
+ scanStateSchema.indexOf("tablePath"),
+
originalScanStateRow.getString(scanStateSchema.indexOf("tablePath")));
+
+ scanStateRow = new ScanStateRow(valueMap);
+ } else {
+ // For ADD files, we read the table schema and append the CDF columns
manually afterwards.
+ // readLogicalSchema = logicalTableSchema;
+ readPhysicalSchema = physicalTableSchema;
+ beamSchema =
DeltaIO.ReadRows.convertToBeamSchema(appendCDFColumns(logicalTableSchema));
+ scanStateRow = originalScanStateRow;
+ }
+
+ io.delta.kernel.data.Row scanFileRow =
+ generateScanFileRow(task.getPath(), task.getPartitionValues());
+ FileStatus fileStatus = FileStatus.of(task.getPath(), task.getSize(),
task.getTimestamp());
+
+ BeamParquetHandler parquetHandler =
+ new BeamParquetHandler(getConfiguration(),
currentEngine.getParquetHandler(), tracker);
+ BeamEngine beamEngine = new BeamEngine(currentEngine, parquetHandler);
+
+ long currentStartRgIndex = 0L;
+
+ try (CloseableIterator<FileReadResult> fileReadResults =
+ parquetHandler.readParquetFiles(
+ Utils.singletonCloseableIterator(fileStatus),
+ readPhysicalSchema,
+ Optional.empty(),
+ currentStartRgIndex)) {
+
+ CloseableIterator<ColumnarBatch> physicalData =
+ new CloseableIterator<ColumnarBatch>() {
+ @Override
+ public void close() throws java.io.IOException {}
+
+ @Override
+ public boolean hasNext() {
+ return fileReadResults.hasNext();
+ }
+
+ @Override
+ public ColumnarBatch next() {
+ return fileReadResults.next().getData();
+ }
+ };
+
+ try (CloseableIterator<FilteredColumnarBatch> logicalBatches =
+ Scan.transformPhysicalData(beamEngine, scanStateRow, scanFileRow,
physicalData)) {
+
+ while (logicalBatches.hasNext()) {
+ FilteredColumnarBatch batch = logicalBatches.next();
+ ColumnarBatch logicalBatch = batch.getData();
+
+ if (!task.isCDC()) {
+ // For ADD files, we need to append the constant CDF columns:
+ // _change_type = "insert", _commit_version = task.version,
_commit_timestamp =
+ // task.timestamp
+ logicalBatch =
+ appendConstantCDFColumns(
+ currentEngine, logicalBatch, task.getVersion(),
task.getTimestamp());
+ }
+
+ try (CloseableIterator<io.delta.kernel.data.Row> logicalRows =
logicalBatch.getRows()) {
+ while (logicalRows.hasNext()) {
+ io.delta.kernel.data.Row deltaRow = logicalRows.next();
+ Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema);
+ String changeType = beamRow.getString("_change_type");
+ if (changeType == null) {
+ throw new IllegalStateException("Field _change_type must not
be null.");
+ }
+ ValueKind kind = getValueKind(changeType);
+ Row publicRow = projectRow(beamRow, publicBeamSchema);
+ out.builder(publicRow).setValueKind(kind).output();
+ }
+ }
+ }
+ }
+ }
+
+ return ProcessContinuation.stop();
Review Comment:
I wonder if we need `ProcessContinuation` here at all, since we never call
`resume()` ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
