[Iceberg CDC source] Allow dropping PK or watermark columns

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,30 +137,80 @@ 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.
    */
   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<Integer> 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<Integer> addedColumnIds = new HashSet<>();
+    for (NestedField column : TypeUtil.select(tableSchema, missing).columns()) {
+      addedColumnIds.add(column.fieldId());
+    }
+    List<NestedField> 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<Integer> 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<Integer> 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();

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<Row> 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<Long> 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<Row, Row> {
+    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<Row> out) {
+      out.output(Cast.castRow(row, inputSchema, outputSchema));
+    }
+  }
 }

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,

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<Row> outputDrop =

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(