From b0c7cbdbdca2bed4c781fa28ea13a135c3f9cafd Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Wed, 7 Oct 2026 11:38:33 -0700 Subject: [PATCH] allow dropping pk or watermark column --- .../sdk/io/iceberg/IcebergScanConfig.java | 78 +++++++++++++---- .../cdc/IncrementalChangelogSource.java | 34 +++++++- ...ergCdcReadSchemaTransformProviderTest.java | 7 +- .../sdk/io/iceberg/IcebergIOReadTest.java | 18 ---- .../sdk/io/iceberg/IcebergScanConfigTest.java | 77 ++++++++++++---- .../cdc/IncrementalChangelogSourceTest.java | 87 ++++++++++++++++++- 6 files changed, 239 insertions(+), 62 deletions(-) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java index bcd574afdaba..6aebf0aa9b26 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java @@ -65,6 +65,7 @@ public abstract class IcebergScanConfig implements Serializable { private transient @MonotonicNonNull Table cachedTable; private transient org.apache.iceberg.@MonotonicNonNull Schema cachedProjectedSchema; + private transient org.apache.iceberg.@MonotonicNonNull Schema cachedRequestedSchema; private transient org.apache.iceberg.@MonotonicNonNull Schema cachedRequiredSchema; private transient @MonotonicNonNull Expression cachedFilter; private transient org.apache.iceberg.@MonotonicNonNull Schema cachedRecordIdSchema; @@ -136,14 +137,25 @@ static org.apache.iceberg.Schema resolveSchema( return selectedFields.isEmpty() ? schema : schema.select(selectedFields); } - /** Returns the projected Schema after applying column pruning. */ + /** + * Returns the projected Schema after applying column pruning. For CDC reads, this also includes + * the primary key and watermark columns, even when they are not requested. + */ public org.apache.iceberg.Schema getProjectedSchema() { if (cachedProjectedSchema == null) { - cachedProjectedSchema = resolveSchema(getTable().schema(), getKeepFields(), getDropFields()); + cachedProjectedSchema = withCdcFields(getRequestedSchema()); } return cachedProjectedSchema; } + /** Returns the Schema of the columns selected by 'keep' or 'drop'. */ + public org.apache.iceberg.Schema getRequestedSchema() { + if (cachedRequestedSchema == null) { + cachedRequestedSchema = resolveSchema(getTable().schema(), getKeepFields(), getDropFields()); + } + return cachedRequestedSchema; + } + /** * Returns a Schema that includes all the fields required for a successful read. This includes * explicitly selected fields and fields referenced in the filter statement. @@ -151,15 +163,54 @@ public org.apache.iceberg.Schema getProjectedSchema() { public org.apache.iceberg.Schema getRequiredSchema() { if (cachedRequiredSchema == null) { cachedRequiredSchema = - resolveSchema( - getTable().schema(), - getKeepFields(), - getDropFields(), - FilterUtils.getReferencedFieldNames(getFilterString())); + withCdcFields( + resolveSchema( + getTable().schema(), + getKeepFields(), + getDropFields(), + FilterUtils.getReferencedFieldNames(getFilterString()))); } return cachedRequiredSchema; } + /** + * CDC reads also need the primary key, whose column stats the changelog scan uses to find update + * pairs, and the watermark column, which timestamps rows. Adds the top-level columns holding them + * if the projection leaves them out. + */ + private org.apache.iceberg.Schema withCdcFields(org.apache.iceberg.Schema projection) { + if (!getUseCdc()) { + return projection; + } + org.apache.iceberg.Schema tableSchema = getTable().schema(); + Set missing = new HashSet<>(tableSchema.identifierFieldIds()); + @Nullable String watermarkColumn = getWatermarkColumn(); + @Nullable NestedField watermarkField = + watermarkColumn != null ? tableSchema.findField(watermarkColumn) : null; + if (watermarkField != null) { + missing.add(watermarkField.fieldId()); + } + missing.removeAll(TypeUtil.getProjectedIds(projection)); + if (missing.isEmpty()) { + return projection; + } + + Set addedColumnIds = new HashSet<>(); + for (NestedField column : TypeUtil.select(tableSchema, missing).columns()) { + addedColumnIds.add(column.fieldId()); + } + List columns = new ArrayList<>(); + for (NestedField column : tableSchema.columns()) { + @Nullable NestedField projected = projection.asStruct().field(column.fieldId()); + if (addedColumnIds.contains(column.fieldId())) { + columns.add(column); + } else if (projected != null) { + columns.add(projected); + } + } + return new org.apache.iceberg.Schema(columns); + } + public org.apache.iceberg.Schema recordIdSchema() { if (cachedRecordIdSchema == null) { org.apache.iceberg.Schema fullSchema = getTable().schema(); @@ -438,17 +489,9 @@ void validate(Table table) { + invalidOptions)); } } else { - Set primaryKeyIds = new HashSet<>(table.schema().identifierFieldIds()); checkState( - !primaryKeyIds.isEmpty(), + !table.schema().identifierFieldIds().isEmpty(), "Cannot read CDC records as the table schema does not specified any primary key fields."); - Set projectedFieldIds = TypeUtil.getProjectedIds(getProjectedSchema()); - primaryKeyIds.removeAll(projectedFieldIds); - checkArgument( - primaryKeyIds.isEmpty(), - "When reading CDC records, the projected schema must not drop primary key fields. " - + "The specified configuration drops the following PK fields: %s", - primaryKeyIds); validateMetadataColumns(table); } @@ -527,9 +570,6 @@ void validate(Table table) { error("'watermark_column' must be a timestamp-typed column, but '%s' has type %s"), watermarkColumn, field.type().typeId()); - checkArgumentNotNull( - getProjectedSchema().findField(watermarkColumn), - "'watermark_column' column should not be dropped."); } @Nullable String watermarkColumnTimeUnit = getWatermarkColumnTimeUnit(); diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java index fe6240d5dc53..9a5d2413717c 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java @@ -31,7 +31,9 @@ import org.apache.beam.sdk.io.iceberg.IcebergUtils; import org.apache.beam.sdk.io.iceberg.ReadUtils; import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.transforms.Cast; import org.apache.beam.sdk.transforms.Create; +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; @@ -167,7 +169,21 @@ public PCollection expand(PBegin input) { watermarkColumn, scanConfig.getWatermarkColumnTimeUnit()))); } - return merged.setRowSchema(outputRowSchema); + Schema requestedRowSchema = + CdcOutputUtils.outputSchema( + scanConfig, + IcebergUtils.icebergSchemaToBeamSchema( + scanConfig.getRequestedSchema(), scanConfig.getUpdateCompatibilityVersion())); + if (requestedRowSchema.equals(outputRowSchema)) { + return merged.setRowSchema(outputRowSchema); + } + // the primary key or watermark column was read only for internal use + return merged + .setRowSchema(outputRowSchema) + .apply( + "Drop Internal Columns", + ParDo.of(new DropInternalColumns(outputRowSchema, requestedRowSchema))) + .setRowSchema(requestedRowSchema); } /** @@ -208,4 +224,20 @@ private PCollection boundedSnapshots(PBegin input) { .collect(Collectors.toList()); return input.apply("Create Snapshot Range", Create.timestamped(timestamped)); } + + /** Drops columns read only for internal use. Rows keep their ValueKind and timestamp. */ + private static class DropInternalColumns extends DoFn { + private final Schema inputSchema; + private final Schema outputSchema; + + DropInternalColumns(Schema inputSchema, Schema outputSchema) { + this.inputSchema = inputSchema; + this.outputSchema = outputSchema; + } + + @ProcessElement + public void process(@Element Row row, OutputReceiver out) { + out.output(Cast.castRow(row, inputSchema, outputSchema)); + } + } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProviderTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProviderTest.java index 9d12389184a7..54f67e0e30e6 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProviderTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProviderTest.java @@ -216,7 +216,8 @@ public void testManagedReadWithProjectionFilterWatermarkAndSnapshotRange() throw configMap.put("catalog_properties", properties); configMap.put("from_snapshot", table.currentSnapshot().snapshotId()); configMap.put("to_snapshot", table.currentSnapshot().snapshotId()); - configMap.put("keep", ImmutableList.of("id", "data", "event_micros")); + // leaves out the primary key and the watermark column + configMap.put("keep", ImmutableList.of("data")); configMap.put("filter", "\"category\" = 'include'"); configMap.put("watermark_column", "event_micros"); configMap.put( @@ -227,7 +228,7 @@ public void testManagedReadWithProjectionFilterWatermarkAndSnapshotRange() throw IcebergCdcMetadataColumns.ROW_ID, IcebergCdcMetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER)); - org.apache.iceberg.Schema projectedSchema = table.schema().select("id", "data", "event_micros"); + org.apache.iceberg.Schema projectedSchema = table.schema().select("data"); Schema recordSchema = IcebergUtils.icebergSchemaToBeamSchema(projectedSchema); Schema outputSchema = Schema.builder() @@ -249,9 +250,7 @@ public void testManagedReadWithProjectionFilterWatermarkAndSnapshotRange() throw Row record = IcebergUtils.icebergRecordToBeamRow(recordSchema, records.get(i)); return Row.withSchema(outputSchema) .addValues( - record.getInt64("id"), record.getString("data"), - record.getInt64("event_micros"), snapshotId, sequenceNumber, firstRowId + i, diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java index 7920fc37c84b..9a8076bbb7dd 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java @@ -238,19 +238,6 @@ public void testCdcFailsWhenTableHasNoIdentifierFields() { read.expand(PBegin.in(testPipeline)); } - @Test - public void testCdcFailsWhenProjectionDropsIdentifierFields() { - assumeTrue(useIncrementalScan); - TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); - warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); - IcebergIO.ReadRows read = - IcebergIO.readRows(catalogConfig()).from(tableId).withCdc().dropping(singletonList("id")); - - thrown.expect(IllegalArgumentException.class); - thrown.expectMessage("projected schema must not drop primary key fields"); - read.expand(PBegin.in(testPipeline)); - } - @Test public void testFailWhenDropAndKeepAreSet() { TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); @@ -410,11 +397,6 @@ public void testScanSelectedFields() throws Exception { return null; }); - if (useIncrementalScan) { - testPipeline.run(); - return; - } - // test drop fields read = read.keeping(null).dropping(singletonList("id")); PCollection outputDrop = diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java index e928d3cf8f01..989e08dd35d8 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java @@ -20,6 +20,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsString; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import java.util.List; @@ -76,31 +77,77 @@ public void cdcValidationRequiresIdentifierFields() { } @Test - public void cdcValidationRejectsProjectionDroppingIdentifierFields() { + public void cdcProjectionReadsKeyAndWatermarkLeftOutOfKeepOrDrop() { TableIdentifier tableId = uniqueTableId(); Table table = warehouse.createTable(tableId, CDC_SCHEMA); - IcebergScanConfig keepWithoutPk = + IcebergScanConfig keep = scanConfigBuilder(tableId, CDC_SCHEMA) .setUseCdc(true) .setKeepFields(ImmutableList.of("data")) + .setWatermarkColumn("event_time") + .setFilterString("\"category\" = 'x'") .build(); - IllegalArgumentException keepException = - assertThrows(IllegalArgumentException.class, () -> keepWithoutPk.validate(table)); - assertThat( - keepException.getMessage(), - containsString("projected schema must not drop primary key fields")); + keep.validate(table); + assertEquals(ImmutableList.of("data"), fieldNames(keep.getRequestedSchema())); + assertEquals( + ImmutableList.of("id", "data", "event_time"), fieldNames(keep.getProjectedSchema())); + assertEquals( + ImmutableList.of("id", "data", "category", "event_time"), + fieldNames(keep.getRequiredSchema())); + + IcebergScanConfig drop = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setDropFields(ImmutableList.of("id", "event_time")) + .setWatermarkColumn("event_time") + .build(); + drop.validate(table); + assertEquals( + ImmutableList.of( + "data", "category", "event_micros", "optional_time", "required_text", "nested"), + fieldNames(drop.getRequestedSchema())); + assertEquals(fieldNames(CDC_SCHEMA), fieldNames(drop.getProjectedSchema())); - IcebergScanConfig dropPk = + // a projection that already has both is unchanged + IcebergScanConfig unchanged = scanConfigBuilder(tableId, CDC_SCHEMA) .setUseCdc(true) - .setDropFields(ImmutableList.of("id")) + .setKeepFields(ImmutableList.of("id", "event_time")) + .setWatermarkColumn("event_time") .build(); - IllegalArgumentException dropException = - assertThrows(IllegalArgumentException.class, () -> dropPk.validate(table)); - assertThat( - dropException.getMessage(), - containsString("projected schema must not drop primary key fields")); + assertSame(unchanged.getRequestedSchema(), unchanged.getProjectedSchema()); + + // batch reads only project the requested columns + IcebergScanConfig batch = + scanConfigBuilder(tableId, CDC_SCHEMA).setKeepFields(ImmutableList.of("data")).build(); + assertEquals(ImmutableList.of("data"), fieldNames(batch.getProjectedSchema())); + } + + @Test + public void cdcProjectionReadsWholeColumnHoldingNestedKey() { + TableIdentifier tableId = uniqueTableId(); + Table table = + warehouse.createTable( + tableId, + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required( + 1, + "key", + Types.StructType.of( + Types.NestedField.required(2, "id", Types.LongType.get()), + Types.NestedField.optional(3, "region", Types.StringType.get()))), + Types.NestedField.optional(4, "data", Types.StringType.get())), + ImmutableSet.of(2))); + IcebergScanConfig scanConfig = + scanConfigBuilder(tableId, table.schema()) + .setUseCdc(true) + .setKeepFields(ImmutableList.of("data")) + .build(); + + // whole columns, as ResolveChanges emits them + assertEquals(table.schema().asStruct(), scanConfig.getProjectedSchema().asStruct()); } @Test @@ -219,8 +266,6 @@ public void watermarkColumnRejectsInvalidConfigurations() { true, ImmutableList.of("id", "required_text"), "must be a timestamp-typed column"); - assertInvalidWatermark( - tableId, table, "event_time", true, ImmutableList.of("id"), "should not be dropped"); } private void assertInvalidWatermark( diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java index 27a65a0be8fd..c055e609f765 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java @@ -460,6 +460,79 @@ public void watermarkColumnRestampsProjectedRows() throws Exception { pipeline.run().waitUntilFinish(); } + /** + * 'keep' and 'drop' may leave out the primary key and watermark column. The source still reads + * them to pair updates and timestamp rows. + */ + @Test + public void keyAndWatermarkCanBeLeftOutOfOutput() throws Exception { + LocalDateTime t1 = + LocalDateTime.ofInstant( + java.time.Instant.ofEpochMilli(System.currentTimeMillis() - 3_000L), ZoneOffset.UTC); + LocalDateTime t2 = t1.plusSeconds(1); + LocalDateTime t3 = t1.plusSeconds(2); + // SPLIT_SIZE=1 resolves updates through the CoGroupByKey, the default resolves them locally + Map reads = + ImmutableMap.of( + "Keep", + eventChanges("1", t1, t2, t3).setKeepFields(ImmutableList.of("data")).build(), + "Drop", + eventChanges(String.valueOf(TableProperties.SPLIT_SIZE_DEFAULT), t1, t2, t3) + .setDropFields(ImmutableList.of("id", "event_time")) + .build()); + + for (Map.Entry read : reads.entrySet()) { + PCollection rows = + pipeline.apply(read.getKey(), new IncrementalChangelogSource(read.getValue())); + assertEquals( + Schema.builder().addNullableField("data", Schema.FieldType.STRING).build(), + rows.getSchema()); + PAssert.that( + rows.apply(read.getKey() + " Format", ParDo.of(new FormatValueKindAndValues())) + .apply(read.getKey() + " Timestamps", Reify.timestamps())) + .containsInAnyOrder( + TimestampedValue.of("INSERT:[a]", millis(t1)), + TimestampedValue.of("INSERT:[b]", millis(t2)), + // only event_time changed, so pairing compared full rows + TimestampedValue.of("UPDATE_BEFORE:[a]", millis(t1)), + TimestampedValue.of("UPDATE_AFTER:[a]", millis(t3)), + TimestampedValue.of("DELETE:[b]", millis(t2)), + TimestampedValue.of("INSERT:[c]", millis(t3))); + } + + pipeline.run().waitUntilFinish(); + } + + /** + * Commits id 1 and 2, then an overwrite that changes only id 1's event_time, deletes 2, adds 3. + */ + private IcebergScanConfig.Builder eventChanges( + String splitSize, LocalDateTime t1, LocalDateTime t2, LocalDateTime t3) throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName() + splitSize); + Table table = + warehouse.createTable( + tableId, + EVENT_SCHEMA, + null, + ImmutableMap.of( + TableProperties.FORMAT_VERSION, "2", TableProperties.SPLIT_SIZE, splitSize)); + DataFile v1 = + commitAppend( + table, + splitSize + "-v1.parquet", + ImmutableList.of(eventRecord(1L, "a", t1), eventRecord(2L, "b", t2))); + commitOverwrite( + table, splitSize + "-v2.parquet", v1, eventRecord(1L, "a", t3), eventRecord(3L, "c", t3)); + return baseConfigBuilder(table, tableId) + .setWatermarkColumn("event_time") + .setStartingStrategy(StartingStrategy.EARLIEST) + .setToSnapshot(table.currentSnapshot().snapshotId()); + } + + private static Instant millis(LocalDateTime time) { + return new Instant(time.toInstant(ZoneOffset.UTC).toEpochMilli()); + } + /** * Bi-directional rows are windowed per snapshot with zero allowed lateness, so a record that * reaches the CoGroupByKey after the watermark has passed its snapshot's window is dropped @@ -516,12 +589,10 @@ private void assertNoRecordsDroppedAcrossSnapshots(boolean streaming) throws Exc } private DataFile commitOverwrite( - Table table, String fileName, DataFile replaced, Record replacement) throws IOException { + Table table, String fileName, DataFile replaced, Record... replacements) throws IOException { DataFile file = warehouse.writeRecords( - testName.getMethodName() + "-" + fileName, - table.schema(), - ImmutableList.of(replacement)); + testName.getMethodName() + "-" + fileName, table.schema(), Arrays.asList(replacements)); table.newOverwrite().deleteFile(replaced).addFile(file).commit(); table.refresh(); return file; @@ -583,6 +654,14 @@ public void process( } } + private static final class FormatValueKindAndValues extends DoFn { + @ProcessElement + public void process( + @Element Row row, ValueKind valueKind, OutputReceiver outputReceiver) { + outputReceiver.output(valueKind.name() + ":" + row.getValues()); + } + } + private static final class FormatDefaultedChange extends DoFn { @ProcessElement public void process(