gemini-code-assist[bot] commented on code in PR #38837:
URL: https://github.com/apache/beam/pull/38837#discussion_r3565653078
##########
sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java:
##########
@@ -469,120 +470,141 @@ private static Object getIcebergTimestampValue(Object
beamValue, boolean shouldA
}
}
+ /** Converts a {@link StructLike} to a Beam {@link Row}. */
+ public static Row structToRow(Schema schema, StructLike struct) {
+ checkState(
+ schema.getFieldCount() == struct.size(),
+ "Struct of size %s does not match expected schema size %s",
+ struct.size(),
+ schema.getFieldCount());
+ Row.Builder rowBuilder = Row.withSchema(schema);
+ for (int i = 0; i < schema.getFieldCount(); i++) {
+ Schema.Field field = schema.getField(i);
+ @Nullable Object icebergValue = struct.get(i, Object.class);
+ addIcebergValue(rowBuilder, field, icebergValue);
+ }
+ return rowBuilder.build();
+ }
+
/** Converts an Iceberg {@link Record} to a Beam {@link Row}. */
public static Row icebergRecordToBeamRow(Schema schema, Record record) {
Row.Builder rowBuilder = Row.withSchema(schema);
for (Schema.Field field : schema.getFields()) {
- boolean isNullable = field.getType().getNullable();
@Nullable Object icebergValue = record.getField(field.getName());
- if (icebergValue == null) {
- if (isNullable) {
- rowBuilder.addValue(null);
- continue;
- }
- throw new RuntimeException(
- String.format("Received null value for required field '%s'.",
field.getName()));
+ addIcebergValue(rowBuilder, field, icebergValue);
+ }
+ return rowBuilder.build();
+ }
+
+ private static void addIcebergValue(
+ Row.Builder rowBuilder, Schema.Field field, @Nullable Object
icebergValue) {
+ boolean isNullable = field.getType().getNullable();
+ if (icebergValue == null) {
+ if (isNullable) {
+ rowBuilder.addValue(null);
+ return;
}
- switch (field.getType().getTypeName()) {
- case BYTE:
- case INT16:
- case INT32:
- case INT64:
- case DECIMAL: // Iceberg and Beam both use BigDecimal
- case FLOAT: // Iceberg and Beam both use float
- case DOUBLE: // Iceberg and Beam both use double
- case STRING: // Iceberg and Beam both use String
- case BOOLEAN: // Iceberg and Beam both use boolean
- rowBuilder.addValue(icebergValue);
- break;
- case ARRAY:
- checkState(
- icebergValue instanceof List,
- "Expected List type for field '%s' but received %s",
- field.getName(),
- icebergValue.getClass());
- List<@NonNull ?> beamList = (List<@NonNull ?>) icebergValue;
- Schema.FieldType collectionType =
- checkStateNotNull(field.getType().getCollectionElementType());
- // recurse on struct types
- if (collectionType.getTypeName().isCompositeType()) {
- Schema innerSchema =
checkStateNotNull(collectionType.getRowSchema());
- beamList =
- beamList.stream()
- .map(v -> icebergRecordToBeamRow(innerSchema, (Record) v))
- .collect(Collectors.toList());
- }
- rowBuilder.addValue(beamList);
- break;
- case ITERABLE:
- checkState(
- icebergValue instanceof Iterable,
- "Expected Iterable type for field '%s' but received %s",
- field.getName(),
- icebergValue.getClass());
- Iterable<@NonNull ?> beamIterable = (Iterable<@NonNull ?>)
icebergValue;
- Schema.FieldType iterableCollectionType =
- checkStateNotNull(field.getType().getCollectionElementType());
- // recurse on struct types
- if (iterableCollectionType.getTypeName().isCompositeType()) {
- Schema innerSchema =
checkStateNotNull(iterableCollectionType.getRowSchema());
- ImmutableList.Builder<Row> builder = ImmutableList.builder();
- for (Record v : (Iterable<@NonNull Record>) icebergValue) {
- builder.add(icebergRecordToBeamRow(innerSchema, v));
- }
- beamIterable = builder.build();
+ throw new RuntimeException(
+ String.format("Received null value for required field '%s'.",
field.getName()));
+ }
+ switch (field.getType().getTypeName()) {
+ case BYTE:
+ case INT16:
+ case INT32:
+ case INT64:
+ case DECIMAL: // Iceberg and Beam both use BigDecimal
+ case FLOAT: // Iceberg and Beam both use float
+ case DOUBLE: // Iceberg and Beam both use double
+ case STRING: // Iceberg and Beam both use String
+ case BOOLEAN: // Iceberg and Beam both use boolean
+ rowBuilder.addValue(icebergValue);
+ break;
+ case ARRAY:
+ checkState(
+ icebergValue instanceof List,
+ "Expected List type for field '%s' but received %s",
+ field.getName(),
+ icebergValue.getClass());
+ List<@NonNull ?> beamList = (List<@NonNull ?>) icebergValue;
+ Schema.FieldType collectionType =
+ checkStateNotNull(field.getType().getCollectionElementType());
+ // recurse on struct types
+ if (collectionType.getTypeName().isCompositeType()) {
+ Schema innerSchema =
checkStateNotNull(collectionType.getRowSchema());
+ beamList =
+ beamList.stream()
+ .map(v -> icebergRecordToBeamRow(innerSchema, (Record) v))
+ .collect(Collectors.toList());
+ }
+ rowBuilder.addValue(beamList);
+ break;
+ case ITERABLE:
+ checkState(
+ icebergValue instanceof Iterable,
+ "Expected Iterable type for field '%s' but received %s",
+ field.getName(),
+ icebergValue.getClass());
+ Iterable<@NonNull ?> beamIterable = (Iterable<@NonNull ?>)
icebergValue;
+ Schema.FieldType iterableCollectionType =
+ checkStateNotNull(field.getType().getCollectionElementType());
+ // recurse on struct types
+ if (iterableCollectionType.getTypeName().isCompositeType()) {
+ Schema innerSchema =
checkStateNotNull(iterableCollectionType.getRowSchema());
+ ImmutableList.Builder<Row> builder = ImmutableList.builder();
+ for (Record v : (Iterable<@NonNull Record>) icebergValue) {
+ builder.add(icebergRecordToBeamRow(innerSchema, v));
}
- rowBuilder.addValue(beamIterable);
- break;
- case MAP:
- checkState(
- icebergValue instanceof Map,
- "Expected Map type for field '%s' but received %s",
- field.getName(),
- icebergValue.getClass());
- Map<?, ?> beamMap = (Map<?, ?>) icebergValue;
- Schema.FieldType valueType =
checkStateNotNull(field.getType().getMapValueType());
- // recurse on struct types
- if (valueType.getTypeName().isCompositeType()) {
- Schema innerSchema = checkStateNotNull(valueType.getRowSchema());
- ImmutableMap.Builder<Object, Row> newMap = ImmutableMap.builder();
- for (Map.Entry<?, ?> entry : ((Map<?, ?>)
icebergValue).entrySet()) {
- Record rec = ((Record) entry.getValue());
- newMap.put(
- checkStateNotNull(entry.getKey()),
- icebergRecordToBeamRow(innerSchema, checkStateNotNull(rec)));
- }
- beamMap = newMap.build();
+ beamIterable = builder.build();
+ }
+ rowBuilder.addValue(beamIterable);
+ break;
+ case MAP:
+ checkState(
+ icebergValue instanceof Map,
+ "Expected Map type for field '%s' but received %s",
+ field.getName(),
+ icebergValue.getClass());
+ Map<?, ?> beamMap = (Map<?, ?>) icebergValue;
+ Schema.FieldType valueType =
checkStateNotNull(field.getType().getMapValueType());
+ // recurse on struct types
+ if (valueType.getTypeName().isCompositeType()) {
+ Schema innerSchema = checkStateNotNull(valueType.getRowSchema());
+ ImmutableMap.Builder<Object, Row> newMap = ImmutableMap.builder();
+ for (Map.Entry<?, ?> entry : ((Map<?, ?>) icebergValue).entrySet()) {
+ Record rec = ((Record) entry.getValue());
+ newMap.put(
+ checkStateNotNull(entry.getKey()),
+ icebergRecordToBeamRow(innerSchema, checkStateNotNull(rec)));
}
- rowBuilder.addValue(beamMap);
- break;
- case DATETIME:
- // Iceberg uses a long for micros.
- // Beam DATETIME uses joda's DateTime, which only supports millis,
- // so we do lose some precision here
- rowBuilder.addValue(getBeamDateTimeValue(icebergValue));
- break;
- case BYTES:
- // Iceberg uses ByteBuffer; Beam uses byte[]
- rowBuilder.addValue(((ByteBuffer) icebergValue).array());
- break;
- case ROW:
- Record nestedRecord = (Record) icebergValue;
- Schema nestedSchema =
- checkArgumentNotNull(
- field.getType().getRowSchema(),
- "Corrupted schema: Row type did not have associated nested
schema.");
- rowBuilder.addValue(icebergRecordToBeamRow(nestedSchema,
nestedRecord));
- break;
- case LOGICAL_TYPE:
- rowBuilder.addValue(getLogicalTypeValue(icebergValue,
field.getType()));
- break;
- default:
- throw new UnsupportedOperationException(
- "Unsupported Beam type: " + field.getType().getTypeName());
- }
+ beamMap = newMap.build();
+ }
+ rowBuilder.addValue(beamMap);
+ break;
+ case DATETIME:
+ // Iceberg uses a long for micros.
+ // Beam DATETIME uses joda's DateTime, which only supports millis,
+ // so we do lose some precision here
+ rowBuilder.addValue(getBeamDateTimeValue(icebergValue));
+ break;
+ case BYTES:
+ // Iceberg uses ByteBuffer; Beam uses byte[]
+ rowBuilder.addValue(((ByteBuffer) icebergValue).array());
+ break;
+ case ROW:
+ Record nestedRecord = (Record) icebergValue;
+ Schema nestedSchema =
+ checkArgumentNotNull(
+ field.getType().getRowSchema(),
+ "Corrupted schema: Row type did not have associated nested
schema.");
+ rowBuilder.addValue(icebergRecordToBeamRow(nestedSchema,
nestedRecord));
Review Comment:
Casting `icebergValue` directly to `Record` will throw a
`ClassCastException` if the nested object is a `StructLike` (such as a
`StructProjection`) but not a `Record`. Since `structToRow` can be called with
`StructProjection` containing nested structs, we should check the type of
`icebergValue` and handle both `Record` and `StructLike` appropriately.
```suggestion
Schema nestedSchema =
checkArgumentNotNull(
field.getType().getRowSchema(),
"Corrupted schema: Row type did not have associated nested
schema.");
if (icebergValue instanceof Record) {
rowBuilder.addValue(icebergRecordToBeamRow(nestedSchema, (Record)
icebergValue));
} else if (icebergValue instanceof StructLike) {
rowBuilder.addValue(structToRow(nestedSchema, (StructLike)
icebergValue));
} else {
throw new UnsupportedOperationException("Unsupported row type: " +
icebergValue.getClass());
}
```
##########
sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogs.java:
##########
@@ -0,0 +1,493 @@
+/*
+ * 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.iceberg.cdc;
+
+import static
org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergRecordToBeamRow;
+import static
org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergSchemaToBeamSchema;
+import static org.apache.beam.sdk.io.iceberg.IcebergUtils.structToRow;
+import static
org.apache.beam.sdk.io.iceberg.cdc.ChangelogScanner.LARGE_BIDIRECTIONAL_TASKS;
+import static
org.apache.beam.sdk.io.iceberg.cdc.ChangelogScanner.UNIDIRECTIONAL_TASKS;
+import static
org.apache.beam.sdk.io.iceberg.cdc.SerializableChangelogTask.Type.ADDED_ROWS;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.io.iceberg.IcebergScanConfig;
+import org.apache.beam.sdk.io.iceberg.IcebergUtils;
+import org.apache.beam.sdk.io.iceberg.SerializableDeleteFile;
+import org.apache.beam.sdk.io.iceberg.TableCache;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.Flatten;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Redistribute;
+import org.apache.beam.sdk.transforms.join.CoGroupByKey;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionList;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PInput;
+import org.apache.beam.sdk.values.POutput;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueKind;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.ChangelogScanTask;
+import org.apache.iceberg.StructLike;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.util.StructProjection;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A {@link PTransform} that processes batches of {@link ChangelogScanTask}s
and routes them
+ * accordingly:
+ *
+ * <ul>
+ * <li>Records from Uni-directional batches are directly emitted, as INSERT
or DELETE kind
+ * <li>Records from Bi-directional batches are compared against the Primary
Key overlap range:
+ * <ul>
+ * <li>if outside the overlap, emit directly as INSERT or DELETE kind
+ * <li>if inside the overlap, key by (snapshot seq#, pk) and route to
downstream {@link
+ * CoGroupByKey} and final resolution by {@link ResolveChanges}
+ * </ul>
+ * </ul>
+ *
+ * <p>We first key bi-directional rows by (snapshot sequence number, primary
key) before sending to
+ * {@link CoGroupByKey} to ensure they stay isolated from other PKs or
snapshots. Inserts are routed
+ * to
+ *
+ * <p>A {@link ChangelogScanTask} comes in three types:
+ *
+ * <ol>
+ * <li><b>AddedRowsScanTask</b>: Indicates records have been inserted by a
new DataFile.
+ * <li><b>DeletedRowsScanTask</b>: Indicates records have been deleted using
a DeleteFile.
+ * <li><b>DeletedDataFileScanTask</b>: Indicates a whole DataFile has been
deleted.
+ * </ol>
+ *
+ * <p>Each of these types need to be processed differently. More details in
{@link
+ * CdcReadUtils#changelogRecordsForTask}.
+ *
+ * <p>CDC metadata has two entry points in this transform. Row metadata
columns are requested from
+ * the Iceberg reader by {@link CdcReadUtils} and travel inside intermediate
rows until final output
+ * assembly. Snapshot metadata columns come from the {@link
ChangelogDescriptor} / {@link
+ * CdcRowDescriptor} carried with each task or shuffled row, and {@code
_change_type} comes from the
+ * emitted change kind. Final user-visible rows are assembled by {@link
CdcOutputUtils#outputRow},
+ * which appends all requested metadata as top-level columns in the configured
order.
+ */
+public class ReadFromChangelogs extends PTransform<PCollectionTuple,
ReadFromChangelogs.Output> {
+ private static final Counter numAddedRowsScanTasksCompleted =
+ Metrics.counter(ReadFromChangelogs.class,
"numAddedRowsScanTasksCompleted");
+ private static final Counter numDeletedRowsScanTasksCompleted =
+ Metrics.counter(ReadFromChangelogs.class,
"numDeletedRowsScanTasksCompleted");
+ private static final Counter numDeletedDataFileScanTasksCompleted =
+ Metrics.counter(ReadFromChangelogs.class,
"numDeletedDataFileScanTasksCompleted");
+
+ private static final TupleTag<Row> UNIDIRECTIONAL_ROWS = new TupleTag<>();
+ private static final TupleTag<KV<CdcRowDescriptor, Row>>
BIDIRECTIONAL_INSERTS = new TupleTag<>();
+ private static final TupleTag<KV<CdcRowDescriptor, Row>>
BIDIRECTIONAL_DELETES = new TupleTag<>();
+
+ private final IcebergScanConfig scanConfig;
+
+ ReadFromChangelogs(IcebergScanConfig scanConfig) {
+ this.scanConfig = scanConfig;
+ }
+
+ @Override
+ public org.apache.beam.sdk.io.iceberg.cdc.ReadFromChangelogs.Output expand(
+ PCollectionTuple input) {
+ Schema fullRowSchema =
+ CdcOutputUtils.readBeamSchemaWithRowMetadata(
+ scanConfig.getMetadataColumns(), scanConfig.getSchema());
+ Schema projectedRowSchema =
+
IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema());
+ Schema outputRowSchema = CdcOutputUtils.outputSchema(scanConfig,
projectedRowSchema);
+
+ // === UNIDIRECTIONAL tasks ===
+ // (i.e. only deletes, or only inserts)
+ // take the fast approach of just reading and emitting CDC records
+ PCollection<Row> uniDirectionalRows =
+ input
+ .get(UNIDIRECTIONAL_TASKS)
+ .apply("Redistribute Uni-Directional Changes",
Redistribute.arbitrarily())
+ .apply(
+ "Read Uni-Directional Changes",
+ ParDo.of(ReadDoFn.unidirectional(scanConfig))
+ .withOutputTags(UNIDIRECTIONAL_ROWS, TupleTagList.empty()))
+ .get(UNIDIRECTIONAL_ROWS)
+ .setRowSchema(outputRowSchema);
+
+ // === BIDIRECTIONAL tasks ===
+ // (i.e. a task group containing a mix of deletes and inserts)
+ // read and route records according to their PK (see class java doc)
+ PCollectionTuple biDirectionalRows =
+ input
+ .get(LARGE_BIDIRECTIONAL_TASKS)
+ .apply("Redistribute Large Bi-Directional Changes",
Redistribute.arbitrarily())
+ .apply(
+ "Read Bi-Directional Changes",
+ ParDo.of(ReadDoFn.bidirectional(scanConfig))
+ .withOutputTags(
+ BIDIRECTIONAL_INSERTS,
+
TupleTagList.of(BIDIRECTIONAL_DELETES).and(UNIDIRECTIONAL_ROWS)));
+ // Collect pruned (non-overlapping) rows from bi-directional reader
+ PCollection<Row> nonOverlappingRowsFromBiDirTasks =
+
biDirectionalRows.get(UNIDIRECTIONAL_ROWS).setRowSchema(outputRowSchema);
+
+ // Flatten uni-directional rows from both sources
+ PCollection<Row> allUniDirectionalRows =
+ PCollectionList.of(uniDirectionalRows)
+ .and(nonOverlappingRowsFromBiDirTasks)
+ .apply("Flatten Uni-Directional Rows", Flatten.pCollections());
+
+ // Reify to preserve each record's timestamp (CoGBK overwrites timestamps
with the window's
+ // end-of-window)
+ // Note: element timestamps are snapshot commit timestamp
+ KvCoder<CdcRowDescriptor, Row> keyedOutputCoder =
+ KvCoder.of(
+ CdcRowDescriptor.coder(scanConfig.rowIdBeamSchema()),
SchemaCoder.of(fullRowSchema));
+ PCollection<KV<CdcRowDescriptor, Row>> keyedInsertsWithTimestamps =
+
biDirectionalRows.get(BIDIRECTIONAL_INSERTS).setCoder(keyedOutputCoder);
+ PCollection<KV<CdcRowDescriptor, Row>> keyedDeletesWithTimestamps =
+
biDirectionalRows.get(BIDIRECTIONAL_DELETES).setCoder(keyedOutputCoder);
+
+ return new org.apache.beam.sdk.io.iceberg.cdc.ReadFromChangelogs.Output(
+ input.getPipeline(),
+ allUniDirectionalRows,
+ keyedInsertsWithTimestamps,
+ keyedDeletesWithTimestamps);
+ }
+
+ public static class Output implements POutput {
+ private final Pipeline pipeline;
+ private final PCollection<Row> uniDirectionalRows;
+ private final PCollection<KV<CdcRowDescriptor, Row>> biDirectionalInserts;
+ private final PCollection<KV<CdcRowDescriptor, Row>> biDirectionalDeletes;
+
+ Output(
+ Pipeline p,
+ PCollection<Row> uniDirectionalRows,
+ PCollection<KV<CdcRowDescriptor, Row>> biDirectionalInserts,
+ PCollection<KV<CdcRowDescriptor, Row>> biDirectionalDeletes) {
+ this.pipeline = p;
+ this.uniDirectionalRows = uniDirectionalRows;
+ this.biDirectionalInserts = biDirectionalInserts;
+ this.biDirectionalDeletes = biDirectionalDeletes;
+ }
+
+ PCollection<Row> uniDirectionalRows() {
+ return uniDirectionalRows;
+ }
+
+ PCollection<KV<CdcRowDescriptor, Row>> biDirectionalInserts() {
+ return biDirectionalInserts;
+ }
+
+ PCollection<KV<CdcRowDescriptor, Row>> biDirectionalDeletes() {
+ return biDirectionalDeletes;
+ }
+
+ @Override
+ public Pipeline getPipeline() {
+ return pipeline;
+ }
+
+ @Override
+ public Map<TupleTag<?>, PValue> expand() {
+ return ImmutableMap.of(
+ UNIDIRECTIONAL_ROWS,
+ uniDirectionalRows,
+ BIDIRECTIONAL_INSERTS,
+ biDirectionalInserts,
+ BIDIRECTIONAL_DELETES,
+ biDirectionalDeletes);
+ }
+
+ @Override
+ public void finishSpecifyingOutput(
+ String transformName, PInput input, PTransform<?, ?> transform) {}
+ }
+
+ @DoFn.BoundedPerElement
+ private static class ReadDoFn<OutT>
+ extends DoFn<KV<ChangelogDescriptor, List<SerializableChangelogTask>>,
OutT> {
+ private final IcebergScanConfig scanConfig;
+ private final boolean keyedOutput;
+ private final Schema projectedBeamRowSchema;
+ private final Schema outputBeamRowSchema;
+ private final Schema fullBeamRowSchema;
+ private transient @MonotonicNonNull OverlapRange overlap;
+ private transient @MonotonicNonNull StructProjection outputProjector;
+ private transient @MonotonicNonNull StructProjection pkProjector;
+
+ /** Used for uni-directional changes. Records are output immediately
as-is. */
+ static ReadDoFn<Row> unidirectional(IcebergScanConfig scanConfig) {
+ return new ReadDoFn<>(scanConfig, false);
+ }
+
+ /**
+ * Used for bi-directional changes. Records are keyed by (snapshot
sequence number, primary key)
+ * and sent to a CoGBK.
+ */
+ static ReadDoFn<KV<CdcRowDescriptor, Row>> bidirectional(IcebergScanConfig
scanConfig) {
+ return new ReadDoFn<>(scanConfig, true);
+ }
+
+ private ReadDoFn(IcebergScanConfig scanConfig, boolean keyedOutput) {
+ this.scanConfig = scanConfig;
+ this.keyedOutput = keyedOutput;
+
+ this.projectedBeamRowSchema =
+ CdcOutputUtils.readBeamSchemaWithRowMetadata(
+ scanConfig.getMetadataColumns(),
scanConfig.getProjectedSchema());
+ this.outputBeamRowSchema =
+ CdcOutputUtils.outputSchema(
+ scanConfig,
icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()));
+ this.fullBeamRowSchema =
+ CdcOutputUtils.readBeamSchemaWithRowMetadata(
+ scanConfig.getMetadataColumns(), scanConfig.getSchema());
+ }
+
+ @Setup
+ public void setup() {
+ this.overlap = OverlapRange.forScanConfig(scanConfig);
+ }
+
+ @ProcessElement
+ public void process(
+ @Element KV<ChangelogDescriptor, List<SerializableChangelogTask>>
element,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ MultiOutputReceiver out)
+ throws IOException {
+ Table table = TableCache.get(scanConfig.getCatalogConfig(),
scanConfig.getTableIdentifier());
+
+ List<SerializableChangelogTask> tasks = element.getValue();
+ ChangelogDescriptor descriptor = element.getKey();
+ @Nullable Row overlapLower = descriptor.getOverlapLower();
+ @Nullable Row overlapUpper = descriptor.getOverlapUpper();
+
+ for (long l = tracker.currentRestriction().getFrom();
+ l < tracker.currentRestriction().getTo();
+ l++) {
+ if (!tracker.tryClaim(l)) {
+ return;
+ }
+
+ SerializableChangelogTask task = tasks.get((int) l);
+ processTaskRecords(descriptor, task, overlapLower, overlapUpper,
table, out);
+ }
+ }
+
+ /**
+ * Processes a ChangelogScanTask and routes records accordingly:
+ *
+ * <p>If this DoFn is configured with {@link #unidirectional}, we simply
read records and output
+ * directly to {@link #UNIDIRECTIONAL_ROWS}.
+ *
+ * <p>If this DoFn is configured with {@link #bidirectional}, we compare
against the Primary Key
+ * overlap range. If within the overlap, we key by (snapshotId, PK) and
out to either {@link
+ * #BIDIRECTIONAL_INSERTS} or {@link #BIDIRECTIONAL_DELETES}. Otherwise
(not in overlap), we
+ * output the record directly to {@link #UNIDIRECTIONAL_ROWS}.
+ */
+ private void processTaskRecords(
+ ChangelogDescriptor descriptor,
+ SerializableChangelogTask task,
+ @Nullable Row overlapLowerRow,
+ @Nullable Row overlapUpperRow,
+ Table table,
+ MultiOutputReceiver outputReceiver)
+ throws IOException {
+ OverlapRange ovl = checkStateNotNull(overlap);
+ @Nullable StructLike overlapLower = ovl.toStructLike(overlapLowerRow);
+ @Nullable StructLike overlapUpper = ovl.toStructLike(overlapUpperRow);
+
+ boolean isInsert = task.getType() == ADDED_ROWS;
+ TupleTag<KV<CdcRowDescriptor, Row>> taggedOutput =
+ isInsert ? BIDIRECTIONAL_INSERTS : BIDIRECTIONAL_DELETES;
+ ValueKind kind = isInsert ? ValueKind.INSERT : ValueKind.DELETE;
+ long commitSnapshotId = descriptor.getCommitSnapshotId();
+ long commitSnapshotSequenceNumber =
descriptor.getSnapshotSequenceNumber();
+
+ Schema readSchema = keyedOutput ? fullBeamRowSchema :
projectedBeamRowSchema;
+ try (CloseableIterable<Record> records =
+ CdcReadUtils.changelogRecordsForTask(task, table, scanConfig,
!keyedOutput)) {
+ for (Record rec : records) {
+ // uni-directional -- just output records (they are already
projected by read pushdown)
+ if (!keyedOutput) {
+ Row row = icebergRecordToBeamRow(projectedBeamRowSchema, rec);
+ outputReceiver
+ .get(UNIDIRECTIONAL_ROWS)
+ .builder(
+ CdcOutputUtils.outputRow(
+ scanConfig.getMetadataColumns(),
+ outputBeamRowSchema,
+ descriptor,
+ kind,
+ row))
+ .setValueKind(kind)
+ .output();
+ continue;
+ }
+
+ // bi-directional -- compare overlap
+ if (ovl.contains(rec, overlapLower, overlapUpper)) {
+ // inside overlap -- read full row and output KV
+ Row row = icebergRecordToBeamRow(readSchema, rec);
+ Row pk = structToRow(scanConfig.rowIdBeamSchema(),
pkProjector().wrap(rec));
+ outputReceiver
+ .get(taggedOutput)
+ .builder(
+ KV.of(
+ CdcRowDescriptor.builder()
+ .setCommitSnapshotId(commitSnapshotId)
+
.setSnapshotSequenceNumber(commitSnapshotSequenceNumber)
+ .setPrimaryKey(pk)
+ .build(),
+ row))
+ .setValueKind(kind)
+ .output();
+
+ } else {
+ // outside overlap -- get projected record and output
+ StructLike projected = outputProjector().wrap(rec);
+ Row row = structToRow(projectedBeamRowSchema, projected);
+ outputReceiver
+ .get(UNIDIRECTIONAL_ROWS)
+ .builder(
+ CdcOutputUtils.outputRow(
+ scanConfig.getMetadataColumns(),
+ outputBeamRowSchema,
+ descriptor,
+ kind,
+ row))
+ .setValueKind(kind)
+ .output();
+ }
+ }
+ }
+
+ trackMetrics(task.getType());
+ }
+
+ private StructProjection outputProjector() {
+ if (outputProjector == null) {
+ outputProjector =
+ StructProjection.create(
+ CdcOutputUtils.readSchemaWithRowMetadata(
+ scanConfig.getMetadataColumns(),
+ TableCache.get(scanConfig.getCatalogConfig(),
scanConfig.getTableIdentifier())
+ .schema()),
+ CdcOutputUtils.readSchemaWithRowMetadata(
+ scanConfig.getMetadataColumns(),
scanConfig.getProjectedSchema()));
+ }
+ return outputProjector;
+ }
+
+ private StructProjection pkProjector() {
+ if (pkProjector == null) {
+ pkProjector =
+ StructProjection.create(
+ TableCache.get(scanConfig.getCatalogConfig(),
scanConfig.getTableIdentifier())
+ .schema(),
+ scanConfig.recordIdSchema());
Review Comment:
The `pkProjector` is created using `table.schema()`, but the records it
wraps have `fullReadSchema` (which includes metadata columns). This schema
mismatch will cause a position mismatch when `StructProjection` attempts to
access fields of the wrapped record. We should construct the source schema for
`pkProjector` using `CdcOutputUtils.readSchemaWithRowMetadata` to match the
schema of the records being processed, consistent with `outputProjector()`.
```java
StructProjection.create(
CdcOutputUtils.readSchemaWithRowMetadata(
scanConfig.getMetadataColumns(),
TableCache.get(scanConfig.getCatalogConfig(),
scanConfig.getTableIdentifier())
.schema()),
scanConfig.recordIdSchema());
```
##########
sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java:
##########
@@ -0,0 +1,241 @@
+/*
+ * 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.iceberg.cdc;
+
+import static
org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergSchemaToBeamSchema;
+import static
org.apache.beam.sdk.io.iceberg.cdc.SerializableChangelogTask.Type.ADDED_ROWS;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.io.iceberg.IcebergScanConfig;
+import org.apache.beam.sdk.io.iceberg.IcebergUtils;
+import org.apache.beam.sdk.io.iceberg.TableCache;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.join.CoGroupByKey;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.StructLike;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.types.Comparators;
+import org.apache.iceberg.types.TypeUtil;
+import org.apache.iceberg.types.Types;
+import org.apache.iceberg.util.StructLikeMap;
+import org.apache.iceberg.util.StructLikeUtil;
+import org.apache.iceberg.util.StructProjection;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Resolves a small bi-directional changelog group entirely in memory. This is
the equivalent of
+ * {@link ReadFromChangelogs} + {@link CoGroupByKey} + {@link ResolveChanges}.
+ *
+ * <p>All tasks in a changelog group belong to the same Iceberg {@link
Snapshot}. The upstream
+ * {@link ChangelogScanner} routes here only when the total size of the
bi-directional group fits
+ * within {@link TableProperties#SPLIT_SIZE}.
+ *
+ * <p>The incoming batch's overlap region has already been computed in the
scanning phase by {@link
+ * ChangelogScanner}. In this DoFn, we just process each task and route
records:
+ *
+ * <ul>
+ * <li>Records whose PK falls <b>outside</b> the overlap range cannot have
an opposing-side match,
+ * so they are emitted directly with {@code INSERT} or {@code DELETE}
kind.
+ * <li>Records whose PK falls <b>inside</b> the overlap range are stashed in
a {@link
+ * StructLikeMap} keyed by PK, then resolved by {@link CdcResolver}.
+ * </ul>
+ */
+class LocalResolveDoFn extends DoFn<KV<ChangelogDescriptor,
List<SerializableChangelogTask>>, Row> {
+ private final IcebergScanConfig scanConfig;
+ private final org.apache.beam.sdk.schemas.Schema projectedBeamSchema;
+ private final org.apache.beam.sdk.schemas.Schema outputBeamSchema;
+
+ private transient @MonotonicNonNull OverlapRange overlap;
+ private transient @MonotonicNonNull List<Types.NestedField> nonPkFields;
+ private transient @MonotonicNonNull StructProjection projector;
+
+ LocalResolveDoFn(IcebergScanConfig scanConfig) {
+ this.scanConfig = scanConfig;
+ this.projectedBeamSchema =
+ CdcOutputUtils.readBeamSchemaWithRowMetadata(
+ scanConfig.getMetadataColumns(), scanConfig.getProjectedSchema());
+ this.outputBeamSchema =
+ CdcOutputUtils.outputSchema(
+ scanConfig,
icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()));
+ }
+
+ @Setup
+ public void setup() {
+ Schema tableSchema =
+ TableCache.get(scanConfig.getCatalogConfig(),
scanConfig.getTableIdentifier()).schema();
+ Schema fullReadSchema =
+
CdcOutputUtils.readSchemaWithRowMetadata(scanConfig.getMetadataColumns(),
tableSchema);
+ this.overlap = OverlapRange.forScanConfig(scanConfig);
+ Set<String> pkFieldNames = new
HashSet<>(overlap.recordIdSchema().identifierFieldNames());
+ // The dedup logic only inspects non-PK fields, so precompute them once.
+ List<Types.NestedField> nonPk = new ArrayList<>();
+ for (Types.NestedField f : tableSchema.columns()) {
+ if (!pkFieldNames.contains(f.name())) {
+ nonPk.add(f);
+ }
+ }
+ this.nonPkFields = nonPk;
+ this.projector =
+ StructProjection.create(
+ fullReadSchema,
+ CdcOutputUtils.readSchemaWithRowMetadata(
+ scanConfig.getMetadataColumns(),
scanConfig.getProjectedSchema()));
+ }
+
+ @ProcessElement
+ public void process(
+ @Element KV<ChangelogDescriptor, List<SerializableChangelogTask>>
element,
+ OutputReceiver<Row> out)
+ throws IOException {
+ ChangelogDescriptor descriptor = element.getKey();
+ Table table = TableCache.get(scanConfig.getCatalogConfig(),
scanConfig.getTableIdentifier());
+ OverlapRange ovl = checkStateNotNull(overlap);
+
+ // {PK: (inserts | deletes)} for in-overlap records that need resolution.
+ // Records outside the overlap are emitted directly
+ StructLikeMap<PkGroup> pkGroups =
StructLikeMap.create(ovl.recordIdSchema().asStruct());
+
+ @Nullable StructLike overlapLower =
ovl.toStructLike(descriptor.getOverlapLower());
+ @Nullable StructLike overlapUpper =
ovl.toStructLike(descriptor.getOverlapUpper());
+ for (SerializableChangelogTask task : element.getValue()) {
+ readAndRoute(descriptor, task, table, overlapLower, overlapUpper,
pkGroups, out);
+ }
+
+ resolveAndEmit(descriptor, pkGroups, table.schema(), out);
Review Comment:
The records in `pkGroups` are read with `projected = false`, meaning they
have `fullReadSchema` (which includes metadata columns). However,
`resolveAndEmit` is called with `table.schema()` (which does not include
metadata columns). This schema mismatch will cause a position mismatch when
`RecordResolver`'s `StructProjection` attempts to compare non-PK fields. We
should pass the schema with metadata columns to `resolveAndEmit`.
```java
resolveAndEmit(
descriptor,
pkGroups,
CdcOutputUtils.readSchemaWithRowMetadata(scanConfig.getMetadataColumns(),
table.schema()),
out);
```
##########
sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRange.java:
##########
@@ -0,0 +1,93 @@
+/*
+ * 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.iceberg.cdc;
+
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
+import java.util.Comparator;
+import org.apache.beam.sdk.io.iceberg.IcebergScanConfig;
+import org.apache.beam.sdk.io.iceberg.IcebergUtils;
+import org.apache.beam.sdk.io.iceberg.TableCache;
+import org.apache.beam.sdk.values.Row;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.StructLike;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.util.StructProjection;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Primary-key-projection and overlap-range comparison helper.
+ *
+ * <p>Used by {@link LocalResolveDoFn} and {@link ReadFromChangelogs} to
decide whether a record's
+ * PK falls within an overlap of two opposing tasks. If so, the record needs
to be compared with
+ * others to determine if it is part of an update pair.
+ */
+final class OverlapRange {
+ private final Schema recordIdSchema;
+ private final StructProjection recordIdProjection;
+ private final Comparator<StructLike> idComp;
+
+ private OverlapRange(
+ Schema recordIdSchema, StructProjection recordIdProjection,
Comparator<StructLike> idComp) {
+ this.recordIdSchema = recordIdSchema;
+ this.recordIdProjection = recordIdProjection;
+ this.idComp = idComp;
+ }
+
+ static OverlapRange forScanConfig(IcebergScanConfig scanConfig) {
+ Schema fullSchema =
+ TableCache.get(scanConfig.getCatalogConfig(),
scanConfig.getTableIdentifier()).schema();
+ StructProjection projection = StructProjection.create(fullSchema,
scanConfig.recordIdSchema());
Review Comment:
The records evaluated by `contains` have `fullReadSchema` (which includes
metadata columns) because they are read with `projected = false`. However,
`recordIdProjection` is created using `table.schema()` (which does not include
metadata columns). This schema mismatch will cause a position mismatch when
`StructProjection` attempts to access fields of the wrapped record. We should
construct `fullSchema` using `CdcOutputUtils.readSchemaWithRowMetadata` to
include the metadata columns.
```suggestion
Schema tableSchema =
TableCache.get(scanConfig.getCatalogConfig(),
scanConfig.getTableIdentifier()).schema();
Schema fullSchema =
CdcOutputUtils.readSchemaWithRowMetadata(scanConfig.getMetadataColumns(),
tableSchema);
StructProjection projection = StructProjection.create(fullSchema,
scanConfig.recordIdSchema());
```
--
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]
