Skip to content

Commit 3691832

Browse files
chamikaramjchamikaramj
committed
Adds the Delta Lake CDC read transforms to the Managed I/O API
1 parent f073408 commit 3691832

8 files changed

Lines changed: 1116 additions & 7 deletions

File tree

‎model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/external_transforms.proto‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,8 @@ message ManagedTransforms {
107107
"beam:schematransform:org.apache.beam:sql_server_write:v1"];
108108
DELTA_LAKE_READ = 13 [(org.apache.beam.model.pipeline.v1.beam_urn) =
109109
"beam:schematransform:org.apache.beam:delta_lake_read:v1"];
110+
DELTA_LAKE_CDC_READ = 14 [(org.apache.beam.model.pipeline.v1.beam_urn) =
111+
"beam:schematransform:org.apache.beam:delta_lake_cdc_read:v1"];
110112
}
111113
}
112114

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,11 +61,14 @@
6161
@DoFn.BoundedPerElement
6262
class DeltaCDCSourceDoFn extends DoFn<DeltaCDCReadTask, Row> {
6363
@Nullable Map<String, String> hadoopConfig;
64+
private final @Nullable List<String> metadataColumns;
6465
private transient @Nullable Engine engine;
6566
private transient @Nullable Configuration conf;
6667

67-
public DeltaCDCSourceDoFn(@Nullable Map<String, String> hadoopConfig) {
68+
public DeltaCDCSourceDoFn(
69+
@Nullable Map<String, String> hadoopConfig, @Nullable List<String> metadataColumns) {
6870
this.hadoopConfig = hadoopConfig;
71+
this.metadataColumns = metadataColumns;
6972
}
7073

7174
private synchronized Configuration getConfiguration() {
@@ -117,7 +120,8 @@ public void processElement(
117120

118121
SerializableRow originalScanStateRow = task.getScanStateRow();
119122
StructType logicalTableSchema = ScanStateRow.getLogicalSchema(originalScanStateRow);
120-
Schema publicBeamSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
123+
Schema baseSchema = DeltaIO.ReadRows.convertToBeamSchema(logicalTableSchema);
124+
Schema publicBeamSchema = DeltaIO.buildPublicBeamSchema(baseSchema, metadataColumns);
121125
StructType physicalTableSchema = ScanStateRow.getPhysicalDataReadSchema(originalScanStateRow);
122126

123127
StructType scanStateSchema = originalScanStateRow.getSchema();

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

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.sdk.io.delta;
19+
20+
import static org.apache.beam.sdk.io.delta.DeltaCdcReadSchemaTransformProvider.Configuration;
21+
import static org.apache.beam.sdk.util.construction.BeamUrns.getUrn;
22+
23+
import com.google.auto.service.AutoService;
24+
import com.google.auto.value.AutoValue;
25+
import java.util.Collections;
26+
import java.util.List;
27+
import java.util.Map;
28+
import org.apache.beam.model.pipeline.v1.ExternalTransforms;
29+
import org.apache.beam.sdk.schemas.AutoValueSchema;
30+
import org.apache.beam.sdk.schemas.NoSuchSchemaException;
31+
import org.apache.beam.sdk.schemas.SchemaRegistry;
32+
import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
33+
import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
34+
import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
35+
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
36+
import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
37+
import org.apache.beam.sdk.values.PCollection;
38+
import org.apache.beam.sdk.values.PCollectionRowTuple;
39+
import org.apache.beam.sdk.values.Row;
40+
import org.checkerframework.checker.nullness.qual.Nullable;
41+
42+
/**
43+
* SchemaTransform implementation for {@link DeltaIO#readChanges}. Reads change records from Delta
44+
* Lake and outputs a {@link org.apache.beam.sdk.values.PCollection} of Beam {@link
45+
* org.apache.beam.sdk.values.Row}s.
46+
*/
47+
@AutoService(SchemaTransformProvider.class)
48+
public class DeltaCdcReadSchemaTransformProvider
49+
extends TypedSchemaTransformProvider<Configuration> {
50+
static final String OUTPUT_TAG = "output";
51+
52+
@Override
53+
protected SchemaTransform from(Configuration configuration) {
54+
return new DeltaCdcReadSchemaTransform(configuration);
55+
}
56+
57+
@Override
58+
public List<String> outputCollectionNames() {
59+
return Collections.singletonList(OUTPUT_TAG);
60+
}
61+
62+
@Override
63+
public String identifier() {
64+
return getUrn(ExternalTransforms.ManagedTransforms.Urns.DELTA_LAKE_CDC_READ);
65+
}
66+
67+
static class DeltaCdcReadSchemaTransform extends SchemaTransform {
68+
private final Configuration configuration;
69+
70+
DeltaCdcReadSchemaTransform(Configuration configuration) {
71+
this.configuration =
72+
java.util.Objects.requireNonNull(configuration, "configuration cannot be null");
73+
}
74+
75+
Row getConfigurationRow() {
76+
try {
77+
return SchemaRegistry.createDefault()
78+
.getToRowFunction(Configuration.class)
79+
.apply(configuration)
80+
.sorted()
81+
.toSnakeCase();
82+
} catch (NoSuchSchemaException e) {
83+
throw new RuntimeException(e);
84+
}
85+
}
86+
87+
@Override
88+
public PCollectionRowTuple expand(PCollectionRowTuple input) {
89+
DeltaIO.ReadChanges read = DeltaIO.readChanges().from(configuration.getTable());
90+
Long startVersion = configuration.getStartVersion();
91+
if (startVersion != null) {
92+
read = read.withStartVersion(startVersion);
93+
}
94+
String startTimestamp = configuration.getStartTimestamp();
95+
if (startTimestamp != null) {
96+
read = read.withStartTimestamp(startTimestamp);
97+
}
98+
Long endVersion = configuration.getEndVersion();
99+
if (endVersion != null) {
100+
read = read.withEndVersion(endVersion);
101+
}
102+
String endTimestamp = configuration.getEndTimestamp();
103+
if (endTimestamp != null) {
104+
read = read.withEndTimestamp(endTimestamp);
105+
}
106+
Map<String, String> hadoopConfig = configuration.getHadoopConfig();
107+
if (hadoopConfig != null) {
108+
read = read.withConfig(hadoopConfig);
109+
}
110+
List<String> includeMetadataColumns = configuration.getIncludeMetadataColumns();
111+
if (includeMetadataColumns != null && !includeMetadataColumns.isEmpty()) {
112+
read = read.withMetadataColumns(includeMetadataColumns.toArray(new String[0]));
113+
}
114+
115+
PCollection<Row> output = input.getPipeline().apply(read);
116+
117+
return PCollectionRowTuple.of(OUTPUT_TAG, output);
118+
}
119+
}
120+
121+
@DefaultSchema(AutoValueSchema.class)
122+
@AutoValue
123+
public abstract static class Configuration {
124+
static Builder builder() {
125+
return new AutoValue_DeltaCdcReadSchemaTransformProvider_Configuration.Builder();
126+
}
127+
128+
@SchemaFieldDescription("Identifier of the Delta Lake table.")
129+
abstract String getTable();
130+
131+
@SchemaFieldDescription("Start version of the Delta Lake table to read changes from.")
132+
@Nullable
133+
abstract Long getStartVersion();
134+
135+
@SchemaFieldDescription("Start timestamp of the Delta Lake table to read changes from.")
136+
@Nullable
137+
abstract String getStartTimestamp();
138+
139+
@SchemaFieldDescription("End version of the Delta Lake table to read changes up to.")
140+
@Nullable
141+
abstract Long getEndVersion();
142+
143+
@SchemaFieldDescription("End timestamp of the Delta Lake table to read changes up to.")
144+
@Nullable
145+
abstract String getEndTimestamp();
146+
147+
@SchemaFieldDescription("Properties passed to the Hadoop Configuration.")
148+
@Nullable
149+
abstract Map<String, String> getHadoopConfig();
150+
151+
@SchemaFieldDescription("Metadata columns to include in the output rows.")
152+
@Nullable
153+
abstract List<String> getIncludeMetadataColumns();
154+
155+
@AutoValue.Builder
156+
abstract static class Builder {
157+
abstract Builder setTable(String table);
158+
159+
abstract Builder setStartVersion(Long startVersion);
160+
161+
abstract Builder setStartTimestamp(String startTimestamp);
162+
163+
abstract Builder setEndVersion(Long endVersion);
164+
165+
abstract Builder setEndTimestamp(String endTimestamp);
166+
167+
abstract Builder setHadoopConfig(Map<String, String> hadoopConfig);
168+
169+
abstract Builder setIncludeMetadataColumns(List<String> includeMetadataColumns);
170+
171+
abstract Configuration build();
172+
}
173+
}
174+
}

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

Lines changed: 42 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@
3737
import io.delta.kernel.types.StructField;
3838
import io.delta.kernel.types.StructType;
3939
import io.delta.kernel.types.TimestampType;
40+
import java.util.Arrays;
41+
import java.util.List;
4042
import java.util.Map;
4143
import org.apache.beam.sdk.annotations.Internal;
4244
import org.apache.beam.sdk.schemas.Schema;
@@ -206,6 +208,26 @@ static Schema.FieldType convertToBeamFieldType(DataType deltaType) {
206208
}
207209
}
208210

211+
static Schema buildPublicBeamSchema(Schema baseSchema, @Nullable List<String> metadataColumns) {
212+
if (metadataColumns == null || metadataColumns.isEmpty()) {
213+
return baseSchema;
214+
}
215+
Schema.Builder builder = Schema.builder();
216+
for (Schema.Field field : baseSchema.getFields()) {
217+
builder.addField(field);
218+
}
219+
for (String col : metadataColumns) {
220+
if (col.equals(CHANGE_TYPE_COLUMN)) {
221+
builder.addField(CHANGE_TYPE_COLUMN, Schema.FieldType.STRING);
222+
} else if (col.equals(COMMIT_VERSION_COLUMN)) {
223+
builder.addField(COMMIT_VERSION_COLUMN, Schema.FieldType.INT64);
224+
} else if (col.equals(COMMIT_TIMESTAMP_COLUMN)) {
225+
builder.addField(COMMIT_TIMESTAMP_COLUMN, Schema.FieldType.DATETIME);
226+
}
227+
}
228+
return builder.build();
229+
}
230+
209231
@AutoValue
210232
public abstract static class ReadChanges extends PTransform<PBegin, PCollection<Row>> {
211233
public abstract @Nullable String getTablePath();
@@ -218,6 +240,8 @@ public abstract static class ReadChanges extends PTransform<PBegin, PCollection<
218240

219241
public abstract @Nullable String getEndTimestamp();
220242

243+
public abstract @Nullable List<String> getMetadataColumns();
244+
221245
public abstract @Nullable Map<String, String> getHadoopConfig();
222246

223247
abstract Builder toBuilder();
@@ -234,6 +258,8 @@ abstract static class Builder {
234258

235259
abstract Builder setEndTimestamp(@Nullable String endTimestamp);
236260

261+
abstract Builder setMetadataColumns(@Nullable List<String> metadataColumns);
262+
237263
abstract Builder setHadoopConfig(@Nullable Map<String, String> hadoopConfig);
238264

239265
abstract ReadChanges build();
@@ -259,6 +285,17 @@ public ReadChanges withEndTimestamp(String endTimestamp) {
259285
return toBuilder().setEndTimestamp(endTimestamp).build();
260286
}
261287

288+
public ReadChanges withMetadataColumns(String... metadataColumns) {
289+
for (String col : metadataColumns) {
290+
if (!col.equals(CHANGE_TYPE_COLUMN)
291+
&& !col.equals(COMMIT_VERSION_COLUMN)
292+
&& !col.equals(COMMIT_TIMESTAMP_COLUMN)) {
293+
throw new IllegalArgumentException("Unsupported metadata column: " + col);
294+
}
295+
}
296+
return toBuilder().setMetadataColumns(Arrays.asList(metadataColumns)).build();
297+
}
298+
262299
public ReadChanges withConfig(Map<String, String> config) {
263300
return toBuilder().setHadoopConfig(config).build();
264301
}
@@ -310,7 +347,8 @@ public PCollection<Row> expand(PBegin input) {
310347
if (deltaSchema == null) {
311348
throw new IllegalStateException("Table schema is null.");
312349
}
313-
Schema beamSchema = ReadRows.convertToBeamSchema(deltaSchema);
350+
Schema baseSchema = ReadRows.convertToBeamSchema(deltaSchema);
351+
Schema publicBeamSchema = buildPublicBeamSchema(baseSchema, getMetadataColumns());
314352

315353
return input
316354
.apply("Create Path", Create.of(path))
@@ -323,8 +361,9 @@ public PCollection<Row> expand(PBegin input) {
323361
getStartTimestamp(),
324362
getEndVersion(),
325363
getEndTimestamp())))
326-
.apply("Read CDF Data", ParDo.of(new DeltaCDCSourceDoFn(hadoopConfig)))
327-
.setRowSchema(beamSchema);
364+
.apply(
365+
"Read CDF Data", ParDo.of(new DeltaCDCSourceDoFn(hadoopConfig, getMetadataColumns())))
366+
.setRowSchema(publicBeamSchema);
328367
}
329368
}
330369
}

0 commit comments

Comments
 (0)