Skip to content

Commit b85ccc2

Browse files
chamikaramjchamikaramj
committed
Address reviewer comments
1 parent 4b75dc0 commit b85ccc2

6 files changed

Lines changed: 671 additions & 242 deletions

File tree

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

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
160160
batch.getColumnVector(cdcIdx),
161161
batch.getSchema().at(cdcIdx).getDataType(),
162162
i);
163-
info.cdcRows.add(cdcRow);
163+
info.cdcInfo.add(cdcRow);
164164
}
165165
if (addIdx >= 0 && !batch.getColumnVector(addIdx).isNullAt(i)) {
166166
Row addRow =
@@ -172,7 +172,7 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
172172
AddFile addFile = new AddFile(addRow);
173173
// Only consider add files that change data (ignore OPTIMIZE etc.)
174174
if (addFile.getDataChange()) {
175-
info.addRows.add(addRow);
175+
info.insertInfo.add(addRow);
176176
}
177177
}
178178
}
@@ -182,6 +182,9 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
182182
List<DeltaCDCReadTask> currentGroup = new ArrayList<>();
183183
long currentGroupSize = 0L;
184184

185+
// TODO: to prevent OOMs in true streaming executions, update DeltaReadTask to include a group
186+
// of files.
187+
185188
// Sort versions to process them in order
186189
List<Long> versions = new ArrayList<>(commitActionsMap.keySet());
187190
Collections.sort(versions);
@@ -191,17 +194,16 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
191194
if (info == null) {
192195
throw new IllegalStateException("CommitActionsInfo was not found for version " + version);
193196
}
194-
boolean hasCDC = !info.cdcRows.isEmpty();
197+
boolean hasCDC = !info.cdcInfo.isEmpty();
195198

196-
List<Row> rowsToProcess = hasCDC ? info.cdcRows : info.addRows;
197-
boolean isCDC = hasCDC;
199+
List<Row> rowsToProcess = hasCDC ? info.cdcInfo : info.insertInfo;
198200

199201
for (Row fileRow : rowsToProcess) {
200202
String relPath;
201203
long size;
202204
Map<String, String> partitionValues;
203205

204-
if (isCDC) {
206+
if (hasCDC) {
205207
relPath = fileRow.getString(AddCDCFile.FULL_SCHEMA.indexOf("path"));
206208
size = fileRow.getLong(AddCDCFile.FULL_SCHEMA.indexOf("size"));
207209
partitionValues =
@@ -224,7 +226,7 @@ public void processElement(@Element String tablePath, OutputReceiver<DeltaCDCRea
224226
partitionValues,
225227
info.version,
226228
info.timestamp,
227-
isCDC,
229+
hasCDC,
228230
rowGroupSizes,
229231
serializableScanState);
230232

@@ -281,9 +283,10 @@ private List<Long> getRowGroupSizes(String pathStr, Configuration conf) {
281283

282284
private static class CommitActionsInfo {
283285
final long version;
286+
284287
final long timestamp;
285-
final List<Row> cdcRows = new ArrayList<>();
286-
final List<Row> addRows = new ArrayList<>();
288+
final List<Row> cdcInfo = new ArrayList<>();
289+
final List<Row> insertInfo = new ArrayList<>();
287290

288291
CommitActionsInfo(long version, long timestamp) {
289292
this.version = version;

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

Lines changed: 4 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -117,24 +117,9 @@ public int hashCode() {
117117

118118
@Override
119119
public String toString() {
120-
return "DeltaCDCReadTask{"
121-
+ "path='"
122-
+ path
123-
+ '\''
124-
+ ", size="
125-
+ size
126-
+ ", partitionValues="
127-
+ partitionValues
128-
+ ", version="
129-
+ version
130-
+ ", timestamp="
131-
+ timestamp
132-
+ ", isCDC="
133-
+ isCDC
134-
+ ", rowGroupSizes="
135-
+ rowGroupSizes
136-
+ ", scanStateRow="
137-
+ scanStateRow
138-
+ '}';
120+
return String.format(
121+
"DeltaCDCReadTask{path='%s', size=%d, partitionValues=%s, version=%d, timestamp=%d, isCDC=%b, "
122+
+ "rowGroupSizes=%s, scanStateRow=%s}",
123+
path, size, partitionValues, version, timestamp, isCDC, rowGroupSizes, scanStateRow);
139124
}
140125
}

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

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ public void setUp() {
104104
}
105105

106106
@ProcessElement
107-
public ProcessContinuation processElement(
107+
public void processElement(
108108
@Element DeltaCDCReadTask task,
109109
RestrictionTracker<OffsetRange, Long> tracker,
110110
OutputReceiver<Row> out)
@@ -205,24 +205,26 @@ public ColumnarBatch next() {
205205

206206
while (logicalBatches.hasNext()) {
207207
FilteredColumnarBatch batch = logicalBatches.next();
208-
ColumnarBatch logicalBatch = batch.getData();
209208

210209
if (!task.isCDC()) {
211210
// For ADD files, we need to append the constant CDF columns:
212211
// _change_type = "insert", _commit_version = task.version, _commit_timestamp =
213212
// task.timestamp
214-
logicalBatch =
213+
ColumnarBatch logicalBatch =
215214
appendConstantCDFColumns(
216-
currentEngine, logicalBatch, task.getVersion(), task.getTimestamp());
215+
currentEngine, batch.getData(), task.getVersion(), task.getTimestamp());
216+
// Make sure we use selection vector to considered filtered out or deleted rows.
217+
batch = new FilteredColumnarBatch(logicalBatch, batch.getSelectionVector());
217218
}
218219

219-
try (CloseableIterator<io.delta.kernel.data.Row> logicalRows = logicalBatch.getRows()) {
220+
try (CloseableIterator<io.delta.kernel.data.Row> logicalRows = batch.getRows()) {
220221
while (logicalRows.hasNext()) {
221222
io.delta.kernel.data.Row deltaRow = logicalRows.next();
222223
Row beamRow = DeltaSourceDoFn.toBeamRow(deltaRow, beamSchema);
223-
String changeType = beamRow.getString("_change_type");
224+
String changeType = beamRow.getString(DeltaIO.CHANGE_TYPE_COLUMN);
224225
if (changeType == null) {
225-
throw new IllegalStateException("Field _change_type must not be null.");
226+
throw new IllegalStateException(
227+
"Field " + DeltaIO.CHANGE_TYPE_COLUMN + " must not be null.");
226228
}
227229
ValueKind kind = getValueKind(changeType);
228230
Row publicRow = projectRow(beamRow, publicBeamSchema);
@@ -232,11 +234,13 @@ public ColumnarBatch next() {
232234
}
233235
}
234236
}
235-
236-
return ProcessContinuation.stop();
237237
}
238238

239239
private static Row projectRow(Row row, Schema targetSchema) {
240+
if (row.getSchema().equals(targetSchema)) {
241+
// We can return the original Row since schemas are the same.
242+
return row;
243+
}
240244
Row.Builder builder = Row.withSchema(targetSchema);
241245
for (Schema.Field field : targetSchema.getFields()) {
242246
builder.addValue(row.getValue(field.getName()));
@@ -263,9 +267,9 @@ private static ValueKind getValueKind(String changeType) {
263267

264268
private static StructType appendCDFColumns(StructType schema) {
265269
return schema
266-
.add("_change_type", StringType.STRING, false)
267-
.add("_commit_version", LongType.LONG, false)
268-
.add("_commit_timestamp", TimestampType.TIMESTAMP, false);
270+
.add(DeltaIO.CHANGE_TYPE_COLUMN, StringType.STRING, false)
271+
.add(DeltaIO.COMMIT_VERSION_COLUMN, LongType.LONG, false)
272+
.add(DeltaIO.COMMIT_TIMESTAMP_COLUMN, TimestampType.TIMESTAMP, false);
269273
}
270274

271275
private ColumnarBatch appendConstantCDFColumns(
@@ -315,14 +319,16 @@ private ColumnarBatch appendConstantCDFColumns(
315319
int numCols = batch.getSchema().length();
316320
return batch
317321
.withNewColumn(
318-
numCols, new StructField("_change_type", StringType.STRING, false), changeTypeVector)
322+
numCols,
323+
new StructField(DeltaIO.CHANGE_TYPE_COLUMN, StringType.STRING, false),
324+
changeTypeVector)
319325
.withNewColumn(
320326
numCols + 1,
321-
new StructField("_commit_version", LongType.LONG, false),
327+
new StructField(DeltaIO.COMMIT_VERSION_COLUMN, LongType.LONG, false),
322328
commitVersionVector)
323329
.withNewColumn(
324330
numCols + 2,
325-
new StructField("_commit_timestamp", TimestampType.TIMESTAMP, false),
331+
new StructField(DeltaIO.COMMIT_TIMESTAMP_COLUMN, TimestampType.TIMESTAMP, false),
326332
commitTimestampVector);
327333
}
328334

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

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -52,23 +52,6 @@
5252
/**
5353
* A connector that reads from <a href="https://delta.io/">Delta Lake</a> tables.
5454
*
55-
* <p>{@link DeltaIO} is offered as a Managed transform. This class is subject to change and should
56-
* not be used directly. Instead, use it like so:
57-
*
58-
* <pre>{@code
59-
* Map<String, Object> config = Map.of(
60-
* "table", "gs://my-bucket/delta-table",
61-
* "hadoop_config", Map.of(
62-
* "fs.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem",
63-
* "fs.AbstractFileSystem.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFS",
64-
* "fs.gs.project.id", "my-project-id"));
65-
*
66-
* pipeline
67-
* .apply(Managed.read(Managed.DELTA_LAKE).withConfig(config))
68-
* .getSinglePCollection()
69-
* .apply(ParDo.of(...));
70-
* }</pre>
71-
*
7255
* <h2>Configuration Options</h2>
7356
*
7457
* Please check the <a href="https://beam.apache.org/documentation/io/managed-io/">Managed IO
@@ -80,6 +63,10 @@
8063
@Internal
8164
public class DeltaIO {
8265

66+
public static final String CHANGE_TYPE_COLUMN = "_change_type";
67+
public static final String COMMIT_VERSION_COLUMN = "_commit_version";
68+
public static final String COMMIT_TIMESTAMP_COLUMN = "_commit_timestamp";
69+
8370
/**
8471
* Reads rows from a Delta Lake table.
8572
*
@@ -90,6 +77,7 @@ public static ReadRows readRows() {
9077
return new AutoValue_DeltaIO_ReadRows.Builder().build();
9178
}
9279

80+
/** Reads change data feed (CDC) from a Delta Lake table. */
9381
public static ReadChanges readChanges() {
9482
return new AutoValue_DeltaIO_ReadChanges.Builder().build();
9583
}
@@ -282,6 +270,8 @@ public PCollection<Row> expand(PBegin input) {
282270
throw new IllegalArgumentException("Table path must be set.");
283271
}
284272
if (getStartVersion() == null && getStartTimestamp() == null) {
273+
// TODO: for unbounded reads, support using current HEAD or the latest snapshot
274+
// as the default starting point.
285275
throw new IllegalArgumentException("Either startVersion or startTimestamp must be set.");
286276
}
287277
if (getStartVersion() != null && getStartTimestamp() != null) {

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -130,11 +130,11 @@ static Builder builder() {
130130
abstract static class Builder {
131131
abstract Builder setTable(String table);
132132

133-
abstract Builder setVersion(@Nullable Long version);
133+
abstract Builder setVersion(Long version);
134134

135-
abstract Builder setTimestamp(@Nullable String timestamp);
135+
abstract Builder setTimestamp(String timestamp);
136136

137-
abstract Builder setHadoopConfig(@Nullable Map<String, String> hadoopConfig);
137+
abstract Builder setHadoopConfig(Map<String, String> hadoopConfig);
138138

139139
abstract Configuration build();
140140
}

0 commit comments

Comments
 (0)