Skip to content

Commit 82ab11e

Browse files
chamikaramjchamikaramj
committed
Update and add a test
1 parent 2feaa26 commit 82ab11e

4 files changed

Lines changed: 214 additions & 88 deletions

File tree

‎sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/CreateCDCReadTasksDoFn.java‎

Lines changed: 36 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,7 @@
4545
import org.apache.hadoop.conf.Configuration;
4646
import org.checkerframework.checker.nullness.qual.Nullable;
4747

48-
/**
49-
* A DoFn that reads the Delta log and plans read tasks for Change Data Feed.
50-
*/
48+
/** A DoFn that reads the Delta log and plans read tasks for Change Data Feed. */
5149
class CreateCDCReadTasksDoFn extends DoFn<String, DeltaCDCReadTask> {
5250
private static final long MAX_TASK_SIZE_BYTES = 1024L * 1024L * 1024L; // 1 GB
5351
private final @Nullable Map<String, String> hadoopConfig;
@@ -122,8 +120,9 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
122120
// 3. Load snapshot at resolvedStartVersion to initialize the CommitRange
123121
Snapshot startSnapshot = table.getSnapshotAsOfVersion(engine, resolvedStartVersion);
124122

125-
CommitRangeBuilder rangeBuilder = TableManager.loadCommitRange(
126-
tablePath, CommitRangeBuilder.CommitBoundary.atVersion(resolvedStartVersion));
123+
CommitRangeBuilder rangeBuilder =
124+
TableManager.loadCommitRange(
125+
tablePath, CommitRangeBuilder.CommitBoundary.atVersion(resolvedStartVersion));
127126
rangeBuilder.withEndBoundary(CommitRangeBuilder.CommitBoundary.atVersion(resolvedEndVersion));
128127
CommitRange range = rangeBuilder.build(engine);
129128

@@ -135,7 +134,8 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
135134
actionSet.add(DeltaAction.ADD);
136135

137136
// 4. Iterate over commits in the range and group actions by version
138-
try (CloseableIterator<ColumnarBatch> batchIter = range.getActions(engine, startSnapshot, actionSet)) {
137+
try (CloseableIterator<ColumnarBatch> batchIter =
138+
range.getActions(engine, startSnapshot, actionSet)) {
139139
Map<Long, CommitActionsInfo> commitActionsMap = new HashMap<>();
140140

141141
while (batchIter.hasNext()) {
@@ -149,21 +149,26 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
149149
long version = batch.getColumnVector(versionIdx).getLong(i);
150150
long timestamp = batch.getColumnVector(timestampIdx).getLong(i);
151151

152-
CommitActionsInfo info = commitActionsMap.computeIfAbsent(
153-
version, k -> new CommitActionsInfo(version, timestamp));
152+
CommitActionsInfo info =
153+
commitActionsMap.computeIfAbsent(
154+
version, k -> new CommitActionsInfo(version, timestamp));
154155

155156
if (cdcIdx >= 0 && !batch.getColumnVector(cdcIdx).isNullAt(i)) {
156-
Row cdcRow = (Row) VectorUtils.getValueAsObject(
157-
batch.getColumnVector(cdcIdx),
158-
batch.getSchema().at(cdcIdx).getDataType(),
159-
i);
157+
Row cdcRow =
158+
(Row)
159+
VectorUtils.getValueAsObject(
160+
batch.getColumnVector(cdcIdx),
161+
batch.getSchema().at(cdcIdx).getDataType(),
162+
i);
160163
info.cdcRows.add(cdcRow);
161164
}
162165
if (addIdx >= 0 && !batch.getColumnVector(addIdx).isNullAt(i)) {
163-
Row addRow = (Row) VectorUtils.getValueAsObject(
164-
batch.getColumnVector(addIdx),
165-
batch.getSchema().at(addIdx).getDataType(),
166-
i);
166+
Row addRow =
167+
(Row)
168+
VectorUtils.getValueAsObject(
169+
batch.getColumnVector(addIdx),
170+
batch.getSchema().at(addIdx).getDataType(),
171+
i);
167172
AddFile addFile = new AddFile(addRow);
168173
// Only consider add files that change data (ignore OPTIMIZE etc.)
169174
if (addFile.getDataChange()) {
@@ -199,8 +204,9 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
199204
if (isCDC) {
200205
relPath = fileRow.getString(AddCDCFile.FULL_SCHEMA.indexOf("path"));
201206
size = fileRow.getLong(AddCDCFile.FULL_SCHEMA.indexOf("size"));
202-
partitionValues = VectorUtils.toJavaMap(
203-
fileRow.getMap(AddCDCFile.FULL_SCHEMA.indexOf("partitionValues")));
207+
partitionValues =
208+
VectorUtils.toJavaMap(
209+
fileRow.getMap(AddCDCFile.FULL_SCHEMA.indexOf("partitionValues")));
204210
} else {
205211
AddFile addFile = new AddFile(fileRow);
206212
relPath = addFile.getPath();
@@ -211,15 +217,16 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
211217
String fullPath = new org.apache.hadoop.fs.Path(tablePath, relPath).toString();
212218
List<Long> rowGroupSizes = getRowGroupSizes(fullPath, conf);
213219

214-
DeltaCDCReadTask task = new DeltaCDCReadTask(
215-
fullPath,
216-
size,
217-
partitionValues,
218-
info.version,
219-
info.timestamp,
220-
isCDC,
221-
rowGroupSizes,
222-
serializableScanState);
220+
DeltaCDCReadTask task =
221+
new DeltaCDCReadTask(
222+
fullPath,
223+
size,
224+
partitionValues,
225+
info.version,
226+
info.timestamp,
227+
isCDC,
228+
rowGroupSizes,
229+
serializableScanState);
223230

224231
if (size >= MAX_TASK_SIZE_BYTES) {
225232
if (!currentGroup.isEmpty()) {
@@ -258,8 +265,8 @@ private List<Long> getRowGroupSizes(String pathStr, Configuration conf) {
258265
List<Long> sizes = new ArrayList<>();
259266
try {
260267
org.apache.hadoop.fs.Path hadoopPath = new org.apache.hadoop.fs.Path(pathStr);
261-
org.apache.parquet.hadoop.metadata.ParquetMetadata metadata = org.apache.parquet.hadoop.ParquetFileReader
262-
.readFooter(
268+
org.apache.parquet.hadoop.metadata.ParquetMetadata metadata =
269+
org.apache.parquet.hadoop.ParquetFileReader.readFooter(
263270
conf,
264271
hadoopPath,
265272
org.apache.parquet.format.converter.ParquetMetadataConverter.NO_FILTER);

‎sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaCDCSourceDoFn.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ public ProcessContinuation processElement(
117117

118118
SerializableRow originalScanStateRow = task.getScanStateRow();
119119
StructType logicalTableSchema = ScanStateRow.getLogicalSchema(originalScanStateRow);
120+
Schema publicBeamSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
120121
StructType physicalTableSchema = ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow);
121122

122123
StructType scanStateSchema = originalScanStateRow.getSchema();
@@ -224,7 +225,8 @@ public ColumnarBatch next() {
224225
throw new IllegalStateException("Field _change_type must not be null.");
225226
}
226227
ValueKind kind = getValueKind(changeType);
227-
out.builder(beamRow).setValueKind(kind).output();
228+
Row publicRow = projectRow(beamRow, publicBeamSchema);
229+
out.builder(publicRow).setValueKind(kind).output();
228230
}
229231
}
230232
}
@@ -234,7 +236,17 @@ public ColumnarBatch next() {
234236
return ProcessContinuation.stop();
235237
}
236238

239+
private static Row projectRow(Row row, Schema targetSchema) {
240+
Row.Builder builder = Row.withSchema(targetSchema);
241+
for (Schema.Field field : targetSchema.getFields()) {
242+
builder.addValue(row.getValue(field.getName()));
243+
}
244+
return builder.build();
245+
}
246+
237247
private static ValueKind getValueKind(String changeType) {
248+
// Maps Delta CDC change types to Beam's ValueKind enum.
249+
// https://docs.delta.io/delta-change-data-feed/#what-is-the-schema-for-the-change-data-feed
238250
switch (changeType) {
239251
case "insert":
240252
return ValueKind.INSERT;

‎sdks/java/io/delta/src/main/java/org/apache/beam/sdk/io/delta/DeltaIO.java‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -283,12 +283,7 @@ public PCollection<Row> expand(PBegin input) {
283283
if (deltaSchema == null) {
284284
throw new IllegalStateException("Table schema is null.");
285285
}
286-
Schema beamSchema =
287-
ReadRows.convertToBeamSchema(
288-
deltaSchema
289-
.add("_change_type", StringType.STRING, false)
290-
.add("_commit_version", LongType.LONG, false)
291-
.add("_commit_timestamp", TimestampType.TIMESTAMP, false));
286+
Schema beamSchema = ReadRows.convertToBeamSchema(deltaSchema);
292287

293288
return input
294289
.apply("Create Path", Create.of(path))

0 commit comments

Comments
 (0)