From 13c1da3ab165d67c41c28af66ebdd0d5dd8ee0a1 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Fri, 4 Sep 2026 22:30:40 -0700 Subject: [PATCH] assign cdc keys --- .../io/iceberg/cdc/sink/AssignCdcKeys.java | 362 +++++ .../cdc/sink/SingleTableDestinations.java | 79 ++ .../iceberg/cdc/sink/AssignCdcKeysTest.java | 1261 +++++++++++++++++ 3 files changed, 1702 insertions(+) create mode 100644 sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/sink/AssignCdcKeys.java create mode 100644 sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/sink/SingleTableDestinations.java create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/sink/AssignCdcKeysTest.java diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/sink/AssignCdcKeys.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/sink/AssignCdcKeys.java new file mode 100644 index 000000000000..2d4a75529760 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/sink/AssignCdcKeys.java @@ -0,0 +1,362 @@ +/* + * 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.sink; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.iceberg.DynamicDestinations; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +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.transforms.providers.ErrorHandling; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +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.ValueInSingleWindow; +import org.apache.beam.sdk.values.ValueKind; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * Assigns a sort key to input {@link Row}s and groups by destination and shard keys, outputting + * {@code KV, KV>}. + * + *

For each element this: + * + *

    + *
  1. resolves the destination string from the raw element; + *
  2. resolves the element's {@link ValueKind}; + *
  3. in upsert mode, drops {@code UPDATE_BEFORE} records; + *
  4. reads the sequence number from {@link CdcWriteConfig#getSequenceNumberColumn()}; + *
  5. takes the row to write from {@link DynamicDestinations#getData}, which excludes the control + * columns read above; + *
  6. resolves and validates the destination table through {@link TableSetup}; + *
  7. encodes the primary key to bytes, which feed both the shard hash and the sort key; + *
  8. computes the deterministic shard, according to {@code numShards} and {@code + * shardsPerPartition} + *
+ * + *

When {@link CdcWriteConfig#getErrorHandling()} is enabled, a record-level failure (unknown + * change type, missing/null sequence number, null equality value, an unresolvable destination) is + * diverted to the {@link #FAILED} output as an {@link ErrorHandling#errorSchema} row ({@code + * failed_row}, {@code error_message}). When error handling is disabled, the transform fails + * instead. + */ +final class AssignCdcKeys extends PTransform, PCollectionTuple> { + + static final TupleTag, KV>> KEYED = new TupleTag<>() {}; + static final TupleTag FAILED = new TupleTag() {}; + + private final IcebergCatalogConfig catalogConfig; + private final CdcWriteConfig config; + private final DynamicDestinations destinations; + private final String runId; + + AssignCdcKeys( + IcebergCatalogConfig catalogConfig, + CdcWriteConfig config, + DynamicDestinations destinations, + String runId) { + this.catalogConfig = catalogConfig; + this.config = config; + this.destinations = destinations; + this.runId = runId; + } + + @Override + public PCollectionTuple expand(PCollection input) { + Schema inputSchema = input.getSchema(); + Schema errorSchema = ErrorHandling.errorSchema(inputSchema); + Schema cdcDataSchema = destinations.getDataSchema(); + PCollectionTuple outputs = + input.apply( + "AssignKeys", + ParDo.of( + new AssignFn( + new TableSetup(catalogConfig, config, destinations, runId), + config, + destinations, + errorSchema)) + .withOutputTags(KEYED, TupleTagList.of(FAILED))); + outputs + .get(KEYED) + .setCoder( + KvCoder.of( + KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()), + KvCoder.of(ByteArrayCoder.of(), CdcRecordCoder.of(cdcDataSchema)))); + outputs.get(FAILED).setCoder(RowCoder.of(errorSchema)); + return outputs; + } + + /** Per-record entry point, running the eight steps listed in the main javadoc above. */ + private static final class AssignFn + extends DoFn, KV>> { + + private final TableSetup tableSetup; + private final CdcWriteConfig config; + private final DynamicDestinations destinations; + private final Schema errorSchema; + private final int numShards; + private final int shardsPerPartition; + private final Counter failedRecords = Metrics.counter(AssignCdcKeys.class, "failedRecords"); + private final Counter upsertUpdateBeforeDropped = + Metrics.counter(AssignCdcKeys.class, "upsertUpdateBeforeDropped"); + + /** The control columns' positions in the current source schema. */ + private @Nullable ControlColumns controls; + + AssignFn( + TableSetup tableSetup, + CdcWriteConfig config, + DynamicDestinations destinations, + Schema errorSchema) { + this.tableSetup = tableSetup; + this.config = config; + this.destinations = destinations; + this.errorSchema = errorSchema; + this.numShards = config.getNumShards(); + this.shardsPerPartition = config.getShardsPerPartition(); + } + + @ProcessElement + public void processElement( + @Element Row element, + ValueKind elementKind, + @Timestamp Instant timestamp, + BoundedWindow window, + PaneInfo pane, + MultiOutputReceiver out) { + try { + Schema schema = element.getSchema(); + String destString = + destinations.getTableStringIdentifier( + ValueInSingleWindow.of(element, timestamp, window, pane)); + + // Resolve the control columns' positions once per source schema. (The local lets the + // nullness checker prove non-nullness, which it cannot for the field.) + ControlColumns cols = controls; + if (cols == null || !cols.matches(schema)) { + cols = ControlColumns.of(schema, config); + controls = cols; + } + + ValueKind kind = resolveKind(element, cols, elementKind); + if (config.getUpsert() && kind == ValueKind.UPDATE_BEFORE) { + upsertUpdateBeforeDropped.inc(); + return; + } + long seq = readSeq(element, cols, kind); + + Row data = destinations.getData(element); + TableSetup.Dest dest = tableSetup.get(destString, data.getSchema()); + requireNonNullEqualityValues(dest, data); + byte[] pkBytes = encodePk(dest, data); + + out.get(KEYED) + .output( + KV.of( + KV.of(destString, shardFor(dest, data, pkBytes)), + KV.of(CdcSortKey.encode(pkBytes, seq, kind), CdcRecord.of(data, kind, seq)))); + } catch (TableSetup.TableConfigException e) { + throw e; + } catch (RuntimeException e) { + if (!config.getErrorHandling()) { + throw e; + } + failedRecords.inc(); + out.get(FAILED).output(ErrorHandling.errorRecord(errorSchema, element, e)); + } + } + + /** + * Resolves this element's {@link ValueKind}. When configured, uses the {@code + * change_type_column} value (mapped via {@code change_type_map} when configured). Otherwise, + * uses the element's native kind. + */ + private ValueKind resolveKind(Row element, ControlColumns cols, ValueKind elementKind) { + @Nullable String changeTypeColumn = config.getChangeTypeColumn(); + if (changeTypeColumn == null) { + return elementKind; + } + if (cols.changeTypeIndex < 0) { + throw new IllegalArgumentException( + "change_type_column '" + + changeTypeColumn + + "' not found in element schema " + + element.getSchema()); + } + @Nullable String raw = element.getString(cols.changeTypeIndex); + if (raw == null) { + throw new IllegalArgumentException( + "change_type_column '" + changeTypeColumn + "' is null for element " + element); + } + @Nullable Map changeTypeMap = config.getChangeTypeMap(); + String name = changeTypeMap != null ? changeTypeMap.getOrDefault(raw, raw) : raw; + try { + return ValueKind.valueOf(name); + } catch (IllegalArgumentException e) { + String mappedClause = name.equals(raw) ? "" : " (mapped to '" + name + "')"; + throw new IllegalArgumentException( + "change_type '" + + raw + + "'" + + mappedClause + + " is not a valid ValueKind name; must be one of " + + Arrays.toString(ValueKind.values()) + + ", or add a change_type_map entry for it.", + e); + } + } + + /** Reads the required non-null sequence number ({@code INT64}) from the full input row. */ + private long readSeq(Row element, ControlColumns cols, ValueKind kind) { + String seqColumn = config.getSequenceNumberColumn(); + Schema schema = element.getSchema(); + @Nullable Long value; + try { + value = cols.seqIndex < 0 ? null : element.getInt64(cols.seqIndex); + } catch (ClassCastException e) { + throw new IllegalArgumentException( + "sequence_number_column '" + + seqColumn + + "' must be INT64 (was: " + + schema.getField(seqColumn).getType() + + ")", + e); + } + if (value == null) { + throw new IllegalArgumentException( + "sequence_number_column '" + + seqColumn + + "' is missing or null for a " + + kind + + " record; every CDC record requires a non-null sequence number."); + } + return value; + } + + /** + * Computes the record's write shard. + * + *

If the table is unpartitioned or if {@code shards_per_partition == num_shards}, the plain + * primary-key shard is returned. + * + *

Otherwise computes the shard using {@link PartitionShardPlan}: each partition owns a block + * of {@code shards_per_partition} consecutive shards. A record's primary key maps to an offset + * within that block. + * + *

Must remain a pure function of the primary key: a key whose same-window records split + * across shards breaks same-commit dedup. + */ + private int shardFor(TableSetup.Dest dest, Row data, byte[] pkBytes) { + @Nullable PartitionShardPlan partitionShardPlan = dest.partitionShardPlan(); + if (partitionShardPlan == null) { + return TableSetup.shardFor(pkBytes, numShards); + } + int offset = Math.floorMod(TableSetup.pkHash(pkBytes), shardsPerPartition); + return partitionShardPlan.shardFor(data, offset, numShards); + } + + /** + * Rejects a projected row with a null equality value: it cannot define row identity, so it + * fails with a clear per-column error rather than an opaque coder failure (or, under + * partition-block sharding, a silently null partition value). + */ + private void requireNonNullEqualityValues(TableSetup.Dest dest, Row data) { + int[] positions = dest.pkFieldPositions(); + for (int i = 0; i < positions.length; i++) { + if (data.getValue(positions[i]) == null) { + throw new IllegalArgumentException( + "null value in equality column '" + + dest.pkSchema().getField(i).getName() + + "'; equality columns must be non-null to define row identity. Row: " + + data); + } + } + } + + /** Extracts the primary key from the projected data row and encodes it to bytes. */ + private byte[] encodePk(TableSetup.Dest dest, Row data) { + int[] pkPositions = dest.pkFieldPositions(); + List<@Nullable Object> pkValues = new ArrayList<>(pkPositions.length); + for (int position : pkPositions) { + pkValues.add(data.getValue(position)); + } + Row pk = Row.withSchema(dest.pkSchema()).attachValues(pkValues); + try { + return CoderUtils.encodeToByteArray(dest.pkCoder(), pk); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode primary key " + pk, e); + } + } + } + + /** The control columns' positions in a source row schema. */ + private static final class ControlColumns { + /** The source schema these positions were resolved against. */ + private final Schema schema; + + /** Position of the sequence-number column, or {@code -1} if the schema has none. */ + private final int seqIndex; + + /** Position of the change-type column, or {@code -1} if unconfigured or absent. */ + private final int changeTypeIndex; + + private ControlColumns(Schema schema, int seqIndex, int changeTypeIndex) { + this.schema = schema; + this.seqIndex = seqIndex; + this.changeTypeIndex = changeTypeIndex; + } + + static ControlColumns of(Schema schema, CdcWriteConfig config) { + @Nullable String changeTypeColumn = config.getChangeTypeColumn(); + return new ControlColumns( + schema, + indexOrAbsent(schema, config.getSequenceNumberColumn()), + changeTypeColumn == null ? -1 : indexOrAbsent(schema, changeTypeColumn)); + } + + private static int indexOrAbsent(Schema schema, String name) { + return schema.hasField(name) ? schema.indexOf(name) : -1; + } + + /** Whether these positions were resolved for {@code other}. */ + @SuppressWarnings("ReferenceEquality") + boolean matches(Schema other) { + return schema == other || schema.equals(other); + } + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/sink/SingleTableDestinations.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/sink/SingleTableDestinations.java new file mode 100644 index 000000000000..dc2c286db898 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/sink/SingleTableDestinations.java @@ -0,0 +1,79 @@ +/* + * 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.sink; + +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.io.iceberg.DynamicDestinations; +import org.apache.beam.sdk.io.iceberg.IcebergDestination; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.util.RowFilter; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.apache.iceberg.catalog.TableIdentifier; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** Routes every record to one table, dropping the control columns from the written row. */ +final class SingleTableDestinations implements DynamicDestinations { + + private final DynamicDestinations delegate; + private final RowFilter filter; + + private SingleTableDestinations(DynamicDestinations delegate, RowFilter filter) { + this.delegate = delegate; + this.filter = filter; + } + + static SingleTableDestinations of( + TableIdentifier tableId, Schema inputSchema, CdcWriteConfig config) { + List controlColumns = new ArrayList<>(); + @Nullable String changeTypeColumn = config.getChangeTypeColumn(); + if (changeTypeColumn != null && inputSchema.hasField(changeTypeColumn)) { + controlColumns.add(changeTypeColumn); + } + if (inputSchema.hasField(config.getSequenceNumberColumn())) { + controlColumns.add(config.getSequenceNumberColumn()); + } + RowFilter filter = new RowFilter(inputSchema); + if (!controlColumns.isEmpty()) { + filter = filter.drop(controlColumns); + } + return new SingleTableDestinations( + DynamicDestinations.singleTable(tableId, filter.outputSchema()), filter); + } + + @Override + public Schema getDataSchema() { + return filter.outputSchema(); + } + + @Override + public Row getData(Row element) { + return filter.filter(element); + } + + @Override + public IcebergDestination instantiateDestination(String destination) { + return delegate.instantiateDestination(destination); + } + + @Override + public String getTableStringIdentifier(ValueInSingleWindow element) { + return delegate.getTableStringIdentifier(element); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/sink/AssignCdcKeysTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/sink/AssignCdcKeysTest.java new file mode 100644 index 000000000000..dd6bc6448ec7 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/sink/AssignCdcKeysTest.java @@ -0,0 +1,1261 @@ +/* + * 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.sink; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.hamcrest.Matchers.hasSize; +import static org.hamcrest.Matchers.lessThan; +import static org.hamcrest.Matchers.nullValue; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.io.iceberg.DynamicDestinations; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.PortableIcebergDestinations; +import org.apache.beam.sdk.metrics.MetricNameFilter; +import org.apache.beam.sdk.metrics.MetricResult; +import org.apache.beam.sdk.metrics.MetricsFilter; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedBytes; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.types.Types; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.rules.TestName; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link AssignCdcKeys}. */ +@RunWith(JUnit4.class) +public class AssignCdcKeysTest { + + @Rule public transient TestPipeline p = TestPipeline.create(); + @Rule public transient TemporaryFolder tmp = new TemporaryFolder(); + @Rule public final TestName testName = new TestName(); + + private static final String SEQ_COL = CdcWriteConfig.DEFAULT_SEQUENCE_NUMBER_COLUMN; + private static final int NUM_SHARDS = 8; + + private static final org.apache.iceberg.Schema ICEBERG_SCHEMA = + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get()), + Types.NestedField.optional(3, "data", Types.StringType.get())); + + private static final Schema DATA_SCHEMA = + Schema.builder() + .addInt32Field("id") + .addNullableField("name", Schema.FieldType.STRING) + .addNullableField("data", Schema.FieldType.STRING) + .build(); + + /** Input schema = data columns + the default sequence-number column. */ + private static final Schema INPUT_SCHEMA = + Schema.builder().addFields(DATA_SCHEMA.getFields()).addInt64Field(SEQ_COL).build(); + + /** {@link #INPUT_SCHEMA} with an additional {@code op} change-type column. */ + private static final Schema INPUT_SCHEMA_WITH_OP = + Schema.builder().addFields(INPUT_SCHEMA.getFields()).addStringField("op").build(); + + /** {@link #INPUT_SCHEMA} but with a nullable sequence-number column. */ + private static final Schema NULLABLE_SEQ_SCHEMA = + Schema.builder() + .addFields(DATA_SCHEMA.getFields()) + .addNullableField(SEQ_COL, Schema.FieldType.INT64) + .build(); + + /** {@link #INPUT_SCHEMA} but with a nullable {@code id} (the equality column). */ + private static final Schema NULLABLE_ID_SCHEMA = + Schema.builder() + .addNullableField("id", Schema.FieldType.INT32) + .addNullableField("name", Schema.FieldType.STRING) + .addNullableField("data", Schema.FieldType.STRING) + .addInt64Field(SEQ_COL) + .build(); + + private Catalog catalog; + private IcebergCatalogConfig catalogConfig; + + @Before + public void setUp() { + catalog = CdcSinkTestUtils.hadoopCatalog(tmp.getRoot()); + catalogConfig = CdcSinkTestUtils.catalogConfig(tmp.getRoot()); + } + + private static TableIdentifier uniqueId(String prefix) { + return TableIdentifier.of("db", prefix + "_" + System.nanoTime()); + } + + private static CdcWriteConfig.Builder cdcWriteConfig() { + return CdcWriteConfig.builder() + .setSinkId("test-sink") + .setNumShards(NUM_SHARDS) + .setShardsPerPartition(NUM_SHARDS); + } + + private TableIdentifier createCanonicalTable() { + String prefix = testName.getMethodName(); + TableIdentifier id = uniqueId(prefix); + CdcSinkTestUtils.createTable( + catalog, id, ICEBERG_SCHEMA, ImmutableSet.of(1), 2, PartitionSpec.unpartitioned()); + return id; + } + + private static Row dataRow(int id, String name, String data, long seq) { + return Row.withSchema(INPUT_SCHEMA).addValues(id, name, data, seq).build(); + } + + private static Row dataRowWithOp(int id, String name, String data, long seq, String op) { + return Row.withSchema(INPUT_SCHEMA_WITH_OP).addValues(id, name, data, seq, op).build(); + } + + @SafeVarargs + private PCollection input(Schema schema, KV... rows) { + return CdcSinkTestUtils.withKinds(p.apply(Create.of(ImmutableList.copyOf(rows)))) + .setRowSchema(schema); + } + + private PCollectionTuple assignKeys( + PCollection in, CdcWriteConfig config, TableIdentifier id) { + return in.apply( + new AssignCdcKeys( + catalogConfig, + config, + SingleTableDestinations.of(id, in.getSchema(), config), + "test-runId")); + } + + /** Sums the committed values of the named {@link AssignCdcKeys} counter (0 if it never fired). */ + private static long counterTotal(PipelineResult result, String name) { + Iterable> counters = + result + .metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter(MetricNameFilter.named(AssignCdcKeys.class, name)) + .build()) + .getCounters(); + long total = 0; + for (MetricResult counter : counters) { + total += counter.getCommitted(); + } + return total; + } + + /** Asserts the pipeline fails and that some message in the failure's cause chain has tokens. */ + private void assertPipelineFailsMentioning(String... tokens) { + Pipeline.PipelineExecutionException e = + assertThrows(Pipeline.PipelineExecutionException.class, () -> p.run().waitUntilFinish()); + StringBuilder messages = new StringBuilder(); + for (Throwable t = e; t != null; t = t.getCause()) { + messages.append(t.getMessage()).append('\n'); + } + for (String token : tokens) { + assertThat(messages.toString(), containsString(token)); + } + } + + @Test + public void nativeKindsKeyRecordsWithShardSortKeyAndPayload() { + TableIdentifier id = createCanonicalTable(); + String dest = id.toString(); + + PCollection rows = + input( + INPUT_SCHEMA, + KV.of(ValueKind.INSERT, dataRow(1, "a", "x", 1L)), + KV.of(ValueKind.DELETE, dataRow(1, "a", "x", 2L))); + + PCollectionTuple outputs = assignKeys(rows, cdcWriteConfig().build(), id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + List, KV>> list = + ImmutableList.copyOf(iter); + assertThat(list, hasSize(2)); + Set shards = new HashSet<>(); + Set kinds = new HashSet<>(); + for (KV, KV> kv : list) { + assertThat(kv.getKey().getKey(), equalTo(dest)); + int shard = kv.getKey().getValue(); + assertThat(shard, greaterThanOrEqualTo(0)); + assertThat(shard, lessThan(NUM_SHARDS)); + shards.add(shard); + CdcRecord record = kv.getValue().getValue(); + kinds.add(record.getKind()); + // Sort key is exactly CdcSortKey.encode(pkBytes, seq, kind). + assertArrayEquals( + CdcSortKey.encode( + pkBytesForId(record.getData().getInt32("id")), + record.getSequenceNumber(), + record.getKind()), + kv.getValue().getKey()); + // Payload row is projected to the data columns, without the sequence column. + Row data = record.getData(); + assertThat(data.getSchema().getFieldNames(), contains("id", "name", "data")); + assertFalse(data.getSchema().hasField(SEQ_COL)); + assertThat(data.getInt32("id"), equalTo(1)); + assertThat(data.getString("name"), equalTo("a")); + assertThat(data.getString("data"), equalTo("x")); + assertThat( + record.getSequenceNumber(), + equalTo(record.getKind() == ValueKind.INSERT ? 1L : 2L)); + } + // Same primary key => same deterministic shard for both records. + assertThat(shards, hasSize(1)); + assertThat(kinds, containsInAnyOrder(ValueKind.INSERT, ValueKind.DELETE)); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + p.run().waitUntilFinish(); + } + + // ------------------------------------------------------------------------------------------- + // change_type_column paths + // ------------------------------------------------------------------------------------------- + + @Test + public void changeTypeColumnWithMapResolvesKindsAndIsStripped() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = + cdcWriteConfig() + .setChangeTypeColumn("op") + .setChangeTypeMap(ImmutableMap.of("c", "INSERT", "u", "UPDATE_AFTER", "d", "DELETE")) + .build(); + + // Native kind is INSERT for all three; the mapped change-type column must override it. + PCollectionTuple outputs = + assignKeys( + input( + INPUT_SCHEMA_WITH_OP, + KV.of(ValueKind.INSERT, dataRowWithOp(1, "a", "x", 1L, "c")), + KV.of(ValueKind.INSERT, dataRowWithOp(2, "b", "y", 2L, "u")), + KV.of(ValueKind.INSERT, dataRowWithOp(3, "c", "z", 3L, "d"))), + config, + id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + List, KV>> list = + ImmutableList.copyOf(iter); + assertThat(list, hasSize(3)); + List kinds = new ArrayList<>(); + for (KV, KV> kv : list) { + CdcRecord record = kv.getValue().getValue(); + kinds.add(record.getKind()); + // The change-type column is stripped by projection. + assertFalse(record.getData().getSchema().hasField("op")); + assertFalse(record.getData().getSchema().hasField(SEQ_COL)); + } + assertThat( + kinds, + containsInAnyOrder(ValueKind.INSERT, ValueKind.UPDATE_AFTER, ValueKind.DELETE)); + return null; + }); + p.run().waitUntilFinish(); + } + + @Test + public void changeTypeColumnUnmappedValueFallsThroughAsValueKindName() { + TableIdentifier id = createCanonicalTable(); + // "DELETE" is not a key of the map, so it falls through and parses as a ValueKind name. + CdcWriteConfig config = + cdcWriteConfig() + .setChangeTypeColumn("op") + .setChangeTypeMap(ImmutableMap.of("c", "INSERT")) + .build(); + + PCollectionTuple outputs = + assignKeys( + input( + INPUT_SCHEMA_WITH_OP, + KV.of(ValueKind.INSERT, dataRowWithOp(1, "a", "x", 1L, "DELETE"))), + config, + id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + List, KV>> list = + ImmutableList.copyOf(iter); + assertThat(list, hasSize(1)); + assertThat(list.get(0).getValue().getValue().getKind(), equalTo(ValueKind.DELETE)); + return null; + }); + p.run().waitUntilFinish(); + } + + @Test + public void unknownChangeTypeDivertedToFailedWithErrorHandling() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = + cdcWriteConfig().setChangeTypeColumn("op").setErrorHandling(true).build(); + Row poisoned = dataRowWithOp(1, "a", "x", 1L, "bogus"); + + PCollectionTuple outputs = + assignKeys(input(INPUT_SCHEMA_WITH_OP, KV.of(ValueKind.INSERT, poisoned)), config, id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)).empty(); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)) + .satisfies( + iter -> { + List failed = ImmutableList.copyOf(iter); + assertThat(failed, hasSize(1)); + String message = failed.get(0).getString("error_message"); + assertThat(message, containsString("bogus")); + // The message lists the valid ValueKind names ... + assertThat(message, containsString("INSERT")); + assertThat(message, containsString("UPDATE_BEFORE")); + assertThat(failed.get(0).getRow("failed_row"), equalTo(poisoned)); + return null; + }); + PipelineResult result = p.run(); + result.waitUntilFinish(); + assertThat(counterTotal(result, "failedRecords"), equalTo(1L)); + } + + /** + * A configured change-type column that is absent from the schema, or null in a row, diverts. + * {@link WriteCdcRows} rejects such schemas at construction; this pins the stage's own guard. + */ + @Test + public void changeTypeColumnAbsentOrNullDiverted() { + CdcWriteConfig config = + cdcWriteConfig().setChangeTypeColumn("op").setErrorHandling(true).build(); + + // facet: change_type_column configured but the input schema has no such column. + TableIdentifier absentId = createCanonicalTable(); + PCollectionTuple absent = + assignKeys( + input(INPUT_SCHEMA, KV.of(ValueKind.INSERT, dataRow(1, "a", "x", 1L))), + config, + absentId); + PAssert.that(absent.get(AssignCdcKeys.KEYED)).empty(); + PAssert.that(absent.get(AssignCdcKeys.FAILED)) + .satisfies( + iter -> { + List failed = ImmutableList.copyOf(iter); + assertThat(failed, hasSize(1)); + String message = failed.get(0).getString("error_message"); + assertThat(message, containsString("'op'")); + assertThat(message, containsString("not found")); + return null; + }); + + // facet: the column exists but the row's value is null. + TableIdentifier nullId = createCanonicalTable(); + Schema schema = + Schema.builder() + .addFields(INPUT_SCHEMA.getFields()) + .addNullableField("op", Schema.FieldType.STRING) + .build(); + Row row = Row.withSchema(schema).addValues(1, "a", "x", 1L, null).build(); + PCollectionTuple nullValue = + CdcSinkTestUtils.withKinds( + "KindsNullOp", + p.apply("CreateNullOp", Create.of(ImmutableList.of(KV.of(ValueKind.INSERT, row))))) + .setRowSchema(schema) + .apply( + "AssignNullOp", + new AssignCdcKeys( + catalogConfig, + config, + SingleTableDestinations.of(nullId, schema, config), + "test-runId")); + PAssert.that(nullValue.get(AssignCdcKeys.KEYED)).empty(); + PAssert.that(nullValue.get(AssignCdcKeys.FAILED)) + .satisfies( + iter -> { + List failed = ImmutableList.copyOf(iter); + assertThat(failed, hasSize(1)); + String message = failed.get(0).getString("error_message"); + assertThat(message, containsString("'op'")); + assertThat(message, containsString("null")); + return null; + }); + p.run().waitUntilFinish(); + } + + /** + * A MIXED batch: the good records must still be keyed and only the poison diverted; an all-poison + * batch cannot tell "divert the poison" from "divert everything". + */ + @Test + public void mixedBatchKeepsGoodRecordsAndDivertsOnlyPoison() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = + cdcWriteConfig().setChangeTypeColumn("op").setErrorHandling(true).build(); + + PCollectionTuple outputs = + assignKeys( + input( + INPUT_SCHEMA_WITH_OP, + KV.of(ValueKind.INSERT, dataRowWithOp(1, "a", "x", 1L, "INSERT")), + KV.of(ValueKind.INSERT, dataRowWithOp(2, "b", "y", 2L, "bogus")), + KV.of(ValueKind.INSERT, dataRowWithOp(3, "c", "z", 3L, "DELETE"))), + config, + id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + Map kindById = new HashMap<>(); + for (KV, KV> kv : iter) { + CdcRecord record = kv.getValue().getValue(); + kindById.put(record.getData().getInt32("id"), record.getKind()); + } + // The two healthy records survive, with their own resolved kinds. + assertThat( + kindById, equalTo(ImmutableMap.of(1, ValueKind.INSERT, 3, ValueKind.DELETE))); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)) + .satisfies( + iter -> { + List failed = ImmutableList.copyOf(iter); + assertThat(failed, hasSize(1)); + assertThat(failed.get(0).getRow("failed_row").getInt32("id"), equalTo(2)); + assertThat(failed.get(0).getString("error_message"), containsString("bogus")); + return null; + }); + PipelineResult result = p.run(); + result.waitUntilFinish(); + assertThat(counterTotal(result, "failedRecords"), equalTo(1L)); + } + + @Test + public void unknownChangeTypeFailsPipelineWithoutErrorHandling() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = cdcWriteConfig().setChangeTypeColumn("op").build(); + + assignKeys( + input( + INPUT_SCHEMA_WITH_OP, KV.of(ValueKind.INSERT, dataRowWithOp(1, "a", "x", 1L, "bogus"))), + config, + id); + + assertPipelineFailsMentioning("bogus"); + } + + // ------------------------------------------------------------------------------------------- + // Sequence number handling + // ------------------------------------------------------------------------------------------- + + /** + * A null sequence value diverts. {@link WriteCdcRows} rejects nullable declarations at + * construction; this pins the stage's own guard. + */ + @Test + public void nullSequenceValueDivertedToFailedWithErrorHandling() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = cdcWriteConfig().setErrorHandling(true).build(); + Row noSeq = Row.withSchema(NULLABLE_SEQ_SCHEMA).addValues(1, "a", "x", null).build(); + + PCollectionTuple outputs = + assignKeys(input(NULLABLE_SEQ_SCHEMA, KV.of(ValueKind.DELETE, noSeq)), config, id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)).empty(); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)) + .satisfies( + iter -> { + List failed = ImmutableList.copyOf(iter); + assertThat(failed, hasSize(1)); + assertThat(failed.get(0).getString("error_message"), containsString("sequence")); + return null; + }); + p.run().waitUntilFinish(); + } + + @Test + public void nullSequenceValueFailsPipelineWithoutErrorHandling() { + TableIdentifier id = createCanonicalTable(); + Row noSeq = Row.withSchema(NULLABLE_SEQ_SCHEMA).addValues(1, "a", "x", null).build(); + + // INSERT included: a null sequence value is poison on every kind (no defaulting). + assignKeys( + input(NULLABLE_SEQ_SCHEMA, KV.of(ValueKind.INSERT, noSeq)), cdcWriteConfig().build(), id); + + assertPipelineFailsMentioning("sequence"); + } + + // ------------------------------------------------------------------------------------------- + // Upsert mode + // ------------------------------------------------------------------------------------------- + + @Test + public void upsertDropsUpdateBeforeWithoutFailure() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = cdcWriteConfig().setUpsert(true).build(); + + PCollectionTuple outputs = + assignKeys( + input( + INPUT_SCHEMA, + KV.of(ValueKind.UPDATE_BEFORE, dataRow(1, "a", "old", 1L)), + KV.of(ValueKind.UPDATE_AFTER, dataRow(1, "a", "new", 2L))), + config, + id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + List, KV>> list = + ImmutableList.copyOf(iter); + assertThat(list, hasSize(1)); + CdcRecord record = list.get(0).getValue().getValue(); + assertThat(record.getKind(), equalTo(ValueKind.UPDATE_AFTER)); + assertThat(record.getData().getString("data"), equalTo("new")); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + PipelineResult result = p.run(); + result.waitUntilFinish(); + assertThat(counterTotal(result, "upsertUpdateBeforeDropped"), equalTo(1L)); + } + + @Test + public void upsertDropsUpdateBeforeWithNullSequenceSilently() { + TableIdentifier id = createCanonicalTable(); + // An upsert feed's before-image may carry no sequence number; it is dropped before the + // sequence is read, so it must neither be keyed nor diverted as a poison record. + CdcWriteConfig config = cdcWriteConfig().setUpsert(true).setErrorHandling(true).build(); + Row before = Row.withSchema(NULLABLE_SEQ_SCHEMA).addValues(1, "a", "old", null).build(); + + PCollectionTuple outputs = + assignKeys(input(NULLABLE_SEQ_SCHEMA, KV.of(ValueKind.UPDATE_BEFORE, before)), config, id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)).empty(); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + PipelineResult result = p.run(); + result.waitUntilFinish(); + assertThat(counterTotal(result, "upsertUpdateBeforeDropped"), equalTo(1L)); + assertThat(counterTotal(result, "failedRecords"), equalTo(0L)); + } + + // ------------------------------------------------------------------------------------------- + // Sort key: built from the resolved kind + // ------------------------------------------------------------------------------------------- + + /** + * The sort key's kind byte must come from the resolved kind, not the element's native {@link + * ValueKind}. + */ + @Test + public void sortKeyUsesResolvedKindNotElementNativeKind() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = + cdcWriteConfig() + .setChangeTypeColumn("op") + .setChangeTypeMap(ImmutableMap.of("b", "UPDATE_BEFORE", "u", "UPDATE_AFTER")) + .build(); + + // Both natively INSERT, both at sequence 5: ONLY the resolved kinds can order them. + PCollectionTuple outputs = + assignKeys( + input( + INPUT_SCHEMA_WITH_OP, + KV.of(ValueKind.INSERT, dataRowWithOp(1, "a", "after", 5L, "u")), + KV.of(ValueKind.INSERT, dataRowWithOp(1, "a", "before", 5L, "b"))), + config, + id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + Map keyByKind = new HashMap<>(); + for (KV, KV> kv : iter) { + keyByKind.put(kv.getValue().getValue().getKind(), kv.getValue().getKey()); + } + assertThat( + keyByKind.keySet(), + containsInAnyOrder(ValueKind.UPDATE_BEFORE, ValueKind.UPDATE_AFTER)); + byte[] before = keyByKind.get(ValueKind.UPDATE_BEFORE); + byte[] after = keyByKind.get(ValueKind.UPDATE_AFTER); + assertArrayEquals( + CdcSortKey.encode(pkBytesForId(1), 5L, ValueKind.UPDATE_BEFORE), before); + assertArrayEquals( + CdcSortKey.encode(pkBytesForId(1), 5L, ValueKind.UPDATE_AFTER), after); + // The shared pk and sequence make the kind byte the whole ordering: the before-image + // must sort strictly first under the byte comparator the shuffle sorter uses. + assertThat( + UnsignedBytes.lexicographicalComparator().compare(before, after), lessThan(0)); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + p.run().waitUntilFinish(); + } + + // ------------------------------------------------------------------------------------------- + // Sort key: prefixed by the encoded primary key + // ------------------------------------------------------------------------------------------- + + /** + * The sort key's leading bytes must be the length-prefixed {@code pkCoder} encoding of the + * record's primary key: the shuffle sorter makes one key's records contiguous by comparing that + * prefix, and the writer will read block boundaries off it without decoding rows. + */ + @Test + public void sortKeyPrefixIsTheEncodedPrimaryKey() { + TableIdentifier id = createCanonicalTable(); + + PCollectionTuple outputs = + assignKeys( + input( + INPUT_SCHEMA, + KV.of(ValueKind.INSERT, dataRow(1, "a", "x", 1L)), + KV.of(ValueKind.DELETE, dataRow(1, "a", "x", 2L)), + KV.of(ValueKind.INSERT, dataRow(2, "b", "y", 1L))), + cdcWriteConfig().build(), + id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + int count = 0; + for (KV, KV> kv : iter) { + count++; + byte[] key = kv.getValue().getKey(); + CdcRecord record = kv.getValue().getValue(); + byte[] pkBytes = pkBytesForId(record.getData().getInt32("id")); + // First 4 bytes carry the pk length big-endian, then the pkCoder bytes follow. + assertThat(ByteBuffer.wrap(key).getInt(), equalTo(pkBytes.length)); + assertArrayEquals(pkBytes, Arrays.copyOfRange(key, 4, 4 + pkBytes.length)); + assertArrayEquals( + CdcSortKey.encode(pkBytes, record.getSequenceNumber(), record.getKind()), key); + } + assertThat(count, equalTo(3)); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + p.run().waitUntilFinish(); + } + + // ------------------------------------------------------------------------------------------- + // Primary-key extraction + // ------------------------------------------------------------------------------------------- + + /** + * A nullable-declared equality column is rejected at resolution (table-level, bypassing error + * handling): the table's identifier columns are required. The per-record null guard still backs + * this up for rows whose schema drifts after resolution. + */ + @Test + public void nullableEqualityColumnSchemaRejectedAtResolution() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = cdcWriteConfig().setErrorHandling(true).build(); + Row nullId = Row.withSchema(NULLABLE_ID_SCHEMA).addValues(null, "a", "x", 1L).build(); + + assignKeys(input(NULLABLE_ID_SCHEMA, KV.of(ValueKind.INSERT, nullId)), config, id); + + assertPipelineFailsMentioning("'id'", "nullable in the input", "required in the table"); + } + + // ------------------------------------------------------------------------------------------- + // shards_per_partition (partition-block sharding) + // ------------------------------------------------------------------------------------------- + + /** + * Iceberg schema of the partitioned fixture: {@code (id INT, region STRING)} are both required + * and both equality columns, as partition-block sharding requires of a partition source column. + */ + private static final org.apache.iceberg.Schema PARTITIONED_ICEBERG_SCHEMA = + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "region", Types.StringType.get()), + Types.NestedField.optional(3, "name", Types.StringType.get())); + + private static final Schema PARTITIONED_DATA_SCHEMA = + IcebergUtils.icebergSchemaToBeamSchema(PARTITIONED_ICEBERG_SCHEMA); + + /** Input schema for {@link #PARTITIONED_ICEBERG_SCHEMA}: its columns plus the sequence column. */ + private static final Schema PARTITIONED_INPUT_SCHEMA = + Schema.builder() + .addFields(PARTITIONED_DATA_SCHEMA.getFields()) + .addInt64Field(SEQ_COL) + .build(); + + /** + * A table partitioned by {@code truncate(region, 2)}: a transform, deliberately, so a shard that + * followed the raw column value rather than the partition value would be visible. + */ + private TableIdentifier createTruncatePartitionedTable(String prefix) { + TableIdentifier id = uniqueId(prefix); + CdcSinkTestUtils.createTable( + catalog, + id, + PARTITIONED_ICEBERG_SCHEMA, + ImmutableSet.of(1, 2), + 2, + PartitionSpec.builderFor(PARTITIONED_ICEBERG_SCHEMA).truncate("region", 2).build()); + return id; + } + + private static Row partitionedRow(int id, String region, String name, long seq) { + return Row.withSchema(PARTITIONED_INPUT_SCHEMA).addValues(id, region, name, seq).build(); + } + + /** The {@code pkCoder} bytes of the canonical table's single-{@code id} primary key. */ + private static byte[] pkBytesForId(@Nullable Integer id) { + Schema pkSchema = Schema.builder().addInt32Field("id").build(); + try { + return CoderUtils.encodeToByteArray( + RowCoder.of(pkSchema), Row.withSchema(pkSchema).addValues(id).build()); + } catch (CoderException e) { + throw new RuntimeException(e); + } + } + + /** The primary-key shard the canonical (unpartitioned) table's single {@code id} column gives. */ + private static int pkShardForId(@Nullable Integer id) { + return TableSetup.shardFor(pkBytesForId(id), NUM_SHARDS); + } + + /** The rows used by every block-sharding test: 3 ids in each of 4 regions, over 2 partitions. */ + private PCollection regionRows() { + List> rows = new ArrayList<>(); + int id = 0; + for (String region : ImmutableList.of("us-east", "us-west", "eu-west", "eu-north")) { + for (int i = 0; i < 3; i++) { + rows.add(KV.of(ValueKind.INSERT, partitionedRow(++id, region, "n" + id, 1L))); + } + } + return CdcSinkTestUtils.withKinds(p.apply(Create.of(rows))) + .setRowSchema(PARTITIONED_INPUT_SCHEMA); + } + + /** + * The distinct shards each {@code truncate(region, 2)} partition's records were assigned, keyed + * by partition value. + */ + private static Map> shardsByPartition( + Iterable, KV>> keyed) { + Map> shards = new HashMap<>(); + for (KV, KV> kv : keyed) { + String region = checkStateNotNull(kv.getValue().getValue().getData().getString("region")); + shards + .computeIfAbsent(region.substring(0, 2), k -> new HashSet<>()) + .add(kv.getKey().getValue()); + } + return shards; + } + + /** + * At {@code shards_per_partition = 1} every record of a partition gets exactly ONE shard, + * collapsing a partition's files per commit window down to one. + */ + @Test + public void shardsPerPartitionOnePutsEachPartitionOnOneShard() { + TableIdentifier id = createTruncatePartitionedTable(testName.getMethodName()); + + PCollectionTuple outputs = + assignKeys(regionRows(), cdcWriteConfig().setShardsPerPartition(1).build(), id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + Map> shards = shardsByPartition(iter); + assertThat(shards.keySet(), containsInAnyOrder("us", "eu")); + // The transform decides, not the raw column: "us-east" and "us-west" share a shard. + assertThat(shards.get("us"), hasSize(1)); + assertThat(shards.get("eu"), hasSize(1)); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + p.run().waitUntilFinish(); + } + + /** + * The same input at the uncapped default: primary-key hashing scatters each partition across + * several shards, the per-partition file multiplication the cap removes. + */ + @Test + public void primaryKeyShardingScattersEachPartitionAcrossShards() { + TableIdentifier id = createTruncatePartitionedTable(testName.getMethodName()); + + PCollectionTuple outputs = assignKeys(regionRows(), cdcWriteConfig().build(), id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + Map> shards = shardsByPartition(iter); + assertThat(shards.get("us").size(), greaterThan(1)); + assertThat(shards.get("eu").size(), greaterThan(1)); + return null; + }); + p.run().waitUntilFinish(); + } + + /** + * At EVERY {@code shards_per_partition} setting (1, an interior 4-of-8 with the PK-derived offset + * live rather than pinned to zero, and the uncapped default) every record carrying one primary + * key lands on one shard, across all four change kinds and differing non-key columns. This is the + * invariant the delta writer's same-commit dedup rides on; partition columns are a subset of the + * equality columns, so it is the real invariant. + */ + @Test + public void everyRecordOfOneKeyLandsOnOneShardAtEveryShardsPerPartitionSetting() { + assertOneKeyOneShard(1); + assertOneKeyOneShard(4); + assertOneKeyOneShard(NUM_SHARDS); + p.run().waitUntilFinish(); + } + + /** Applies stage 1 at the given cap over one key's four kinds and asserts a single shard. */ + private void assertOneKeyOneShard(int shardsPerPartition) { + TableIdentifier id = + createTruncatePartitionedTable(testName.getMethodName() + shardsPerPartition); + PCollection in = + CdcSinkTestUtils.withKinds( + "KindsKeySpp" + shardsPerPartition, + p.apply( + "CreateKeySpp" + shardsPerPartition, + Create.of( + ImmutableList.of( + KV.of(ValueKind.INSERT, partitionedRow(7, "us-east", "a", 1L)), + KV.of(ValueKind.UPDATE_BEFORE, partitionedRow(7, "us-east", "a", 2L)), + KV.of(ValueKind.UPDATE_AFTER, partitionedRow(7, "us-east", "b", 2L)), + KV.of(ValueKind.DELETE, partitionedRow(7, "us-east", "b", 3L)))))) + .setRowSchema(PARTITIONED_INPUT_SCHEMA); + + CdcWriteConfig config = cdcWriteConfig().setShardsPerPartition(shardsPerPartition).build(); + PCollectionTuple outputs = + in.apply( + "AssignKeySpp" + shardsPerPartition, + new AssignCdcKeys( + catalogConfig, + config, + SingleTableDestinations.of(id, PARTITIONED_INPUT_SCHEMA, config), + "test-runId")); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + List, KV>> list = + ImmutableList.copyOf(iter); + assertThat(list, hasSize(4)); + Set shards = new HashSet<>(); + for (KV, KV> kv : list) { + shards.add(kv.getKey().getValue()); + } + assertThat(shards, hasSize(1)); + return null; + }); + } + + /** + * Determinism, pinned end to end: a change anywhere in the hash chain shows up here instead of + * silently resharding. A failure is not automatically a bug (an Iceberg {@code JavaHash} change + * moves these values harmlessly) but must be a conscious update: the same change could mean stage + * 1 and the writer drifted apart. + */ + @Test + public void partitionShardIsPinnedForKnownPartitionValues() { + TableIdentifier id = createTruncatePartitionedTable(testName.getMethodName()); + + // shards_per_partition = 1 must reproduce the pure partition-affine values (offset 0) exactly. + PCollectionTuple outputs = + assignKeys(regionRows(), cdcWriteConfig().setShardsPerPartition(1).build(), id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + Map> shards = shardsByPartition(iter); + assertThat(shards.get("us"), contains(3)); + assertThat(shards.get("eu"), contains(0)); + return null; + }); + p.run().waitUntilFinish(); + } + + /** + * An unpartitioned table ignores the cap (with a WARN) and keeps primary-key sharding: rejecting + * would fail a whole dynamic-destination fleet over one table, honoring it literally would funnel + * the table through a single shard. + */ + @Test + public void unpartitionedTableIgnoresShardsPerPartition() { + TableIdentifier id = createCanonicalTable(); + CdcWriteConfig config = cdcWriteConfig().setShardsPerPartition(1).build(); + + // The gate never builds a plan for an unpartitioned destination, whatever the cap says. + TableSetup setup = + new TableSetup( + catalogConfig, config, DynamicDestinations.singleTable(id, DATA_SCHEMA), "test-runId"); + assertThat(setup.get(id.toString(), DATA_SCHEMA).partitionShardPlan(), nullValue()); + + List> rows = new ArrayList<>(); + for (int i = 1; i <= 32; i++) { + rows.add(KV.of(ValueKind.INSERT, dataRow(i, "n" + i, "d", 1L))); + } + PCollection in = + CdcSinkTestUtils.withKinds(p.apply(Create.of(rows))).setRowSchema(INPUT_SCHEMA); + + PCollectionTuple outputs = assignKeys(in, config, id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + Set shards = new HashSet<>(); + for (KV, KV> kv : iter) { + // Identical to the primary-key shard: the cap was ignored, not applied. + assertThat( + kv.getKey().getValue(), + equalTo(pkShardForId(kv.getValue().getValue().getData().getInt32("id")))); + shards.add(kv.getKey().getValue()); + } + // The plain primary-key spread, not one funnelled shard. + assertThat(shards.size(), greaterThan(1)); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + p.run().waitUntilFinish(); + } + + /** + * The block guarantee: a partition occupies exactly {@code spp} CONSECUTIVE shards starting at + * its {@code spp = 1} shard. Enough distinct keys in one partition cover every offset residue, so + * the observed shard set must be exactly that block. + */ + @Test + public void partitionBlockOccupiesExactlySppConsecutiveShards() { + TableIdentifier id = createTruncatePartitionedTable(testName.getMethodName()); + int numShards = 16; + int spp = 4; + + // >= 64 distinct PKs, all in the single truncate(region, 2) partition "us". + List> rows = new ArrayList<>(); + for (int i = 1; i <= 64; i++) { + rows.add(KV.of(ValueKind.INSERT, partitionedRow(i, "us-east", "n" + i, 1L))); + } + PCollection in = + CdcSinkTestUtils.withKinds(p.apply(Create.of(rows))).setRowSchema(PARTITIONED_INPUT_SCHEMA); + + // base = the shard observed with spp=1, computed through the destination's real plan. + int base = + baseShardFor( + id, + Row.withSchema(PARTITIONED_DATA_SCHEMA).addValues(1, "us-east", "n1").build(), + numShards); + Set block = new HashSet<>(); + for (int k = 0; k < spp; k++) { + block.add(Math.floorMod(base + k, numShards)); + } + Set expectedShards = ImmutableSet.copyOf(block); + + PCollectionTuple outputs = + assignKeys( + in, cdcWriteConfig().setNumShards(numShards).setShardsPerPartition(spp).build(), id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + Set shards = new HashSet<>(); + for (KV, KV> kv : iter) { + shards.add(kv.getKey().getValue()); + } + // Exactly spp distinct shards, and they are {base..base+spp-1} (mod numShards). + assertThat(shards, equalTo(expectedShards)); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + p.run().waitUntilFinish(); + } + + /** + * {@code shards_per_partition == num_shards} bypasses the plan entirely: each record's shard is + * the PLAIN primary-key shard, bit-for-bit, not a block-derived shard that happens to spread. + */ + @Test + public void sppEqualToNumShardsBypassesThePlan() { + TableIdentifier id = createTruncatePartitionedTable(testName.getMethodName()); + int numShards = 16; + CdcWriteConfig config = + cdcWriteConfig().setNumShards(numShards).setShardsPerPartition(numShards).build(); + + // The gate never builds a plan when spp == num_shards, partitioned or not. + TableSetup setup = + new TableSetup( + catalogConfig, + config, + DynamicDestinations.singleTable(id, PARTITIONED_DATA_SCHEMA), + "test-runId"); + assertThat(setup.get(id.toString(), PARTITIONED_DATA_SCHEMA).partitionShardPlan(), nullValue()); + + List> rows = new ArrayList<>(); + Map byId = new HashMap<>(); + for (int i = 1; i <= 8; i++) { + rows.add(KV.of(ValueKind.INSERT, partitionedRow(i, "us-east", "n" + i, 1L))); + byId.put(i, partitionedPkShard(i, "us-east", numShards)); + } + Map expectedShardById = ImmutableMap.copyOf(byId); + PCollection in = + CdcSinkTestUtils.withKinds(p.apply(Create.of(rows))).setRowSchema(PARTITIONED_INPUT_SCHEMA); + + PCollectionTuple outputs = assignKeys(in, config, id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + for (KV, KV> kv : iter) { + int rowId = checkStateNotNull(kv.getValue().getValue().getData().getInt32("id")); + assertThat(kv.getKey().getValue(), equalTo(expectedShardById.get(rowId))); + } + return null; + }); + p.run().waitUntilFinish(); + } + + /** + * The composite shard function is part of the sink's cross-version contract: an in-place update + * mid-window with a changed function splits a primary key across shards and breaks same-commit + * dedup. This test confirms the breakage. + */ + @Test + public void blockShardsAreGoldenPinned() { + TableIdentifier id = createTruncatePartitionedTable(testName.getMethodName()); + Row fixed = partitionedRow(7, "us-east", "a", 1L); + + // pinned: changing any of these re-shards live pipelines on in-place update. + // spp=1 is the block base (avalanched partition hash mod 16); spp=4 offsets it by + // pkHash mod 4 (here +2, inside the block {3,4,5,6}); spp=16 bypasses the plan and is the + // plain primary-key shard. + assertBlockShardPinned(id, fixed, 1, 3); + assertBlockShardPinned(id, fixed, 4, 5); + assertBlockShardPinned(id, fixed, 16, 2); + p.run().waitUntilFinish(); + } + + /** Applies stage 1 at the given {@code shards_per_partition} and pins the single row's shard. */ + private void assertBlockShardPinned(TableIdentifier id, Row row, int spp, int expectedShard) { + PCollection in = + CdcSinkTestUtils.withKinds( + "KindsSpp" + spp, + p.apply( + "CreateSpp" + spp, Create.of(ImmutableList.of(KV.of(ValueKind.INSERT, row))))) + .setRowSchema(PARTITIONED_INPUT_SCHEMA); + CdcWriteConfig config = cdcWriteConfig().setNumShards(16).setShardsPerPartition(spp).build(); + PCollectionTuple outputs = + in.apply( + "AssignSpp" + spp, + new AssignCdcKeys( + catalogConfig, + config, + SingleTableDestinations.of(id, PARTITIONED_INPUT_SCHEMA, config), + "test-runId")); + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + List, KV>> list = + ImmutableList.copyOf(iter); + assertThat(list, hasSize(1)); + assertThat(list.get(0).getKey().getValue(), equalTo(expectedShard)); + return null; + }); + } + + /** + * The {@code shards_per_partition = 1} (base) shard for {@code inputRow}'s partition, computed + * through the destination's real {@link PartitionShardPlan}. + */ + private int baseShardFor(TableIdentifier id, Row dataRow, int numShards) { + TableSetup setup = + new TableSetup( + catalogConfig, + cdcWriteConfig().setNumShards(numShards).setShardsPerPartition(1).build(), + DynamicDestinations.singleTable(id, PARTITIONED_DATA_SCHEMA), + "test-runId"); + TableSetup.Dest dest = setup.get(id.toString(), PARTITIONED_DATA_SCHEMA); + PartitionShardPlan plan = checkStateNotNull(dest.partitionShardPlan()); + return plan.shardFor(dataRow, 0, numShards); + } + + /** The plain primary-key shard of the partitioned fixture's {@code (id, region)} key. */ + private static int partitionedPkShard(int id, String region, int numShards) { + Schema pkSchema = Schema.builder().addInt32Field("id").addStringField("region").build(); + try { + return TableSetup.shardFor( + CoderUtils.encodeToByteArray( + RowCoder.of(pkSchema), Row.withSchema(pkSchema).addValues(id, region).build()), + numShards); + } catch (CoderException e) { + throw new RuntimeException(e); + } + } + + // ------------------------------------------------------------------------------------------- + // Dynamic (templated) destinations + // ------------------------------------------------------------------------------------------- + + @Test + public void templateDestinationRoutesToMultipleTables() { + long suffix = System.nanoTime(); + String tableA = "tmpl_a" + suffix; + String tableB = "tmpl_b" + suffix; + // Tables whose columns are (id, dest): the routing column is also a data column. + CdcSinkTestUtils.createDestTables(catalog, tableA, tableB); + + Schema inputSchema = + Schema.builder() + .addInt32Field("id") + .addNullableField("dest", Schema.FieldType.STRING) + .addInt64Field(SEQ_COL) + .build(); + Row rowA = Row.withSchema(inputSchema).addValues(1, tableA, 1L).build(); + Row rowB = Row.withSchema(inputSchema).addValues(2, tableB, 1L).build(); + + PCollectionTuple outputs = + input(inputSchema, KV.of(ValueKind.INSERT, rowA), KV.of(ValueKind.INSERT, rowB)) + .apply( + new AssignCdcKeys( + catalogConfig, + cdcWriteConfig().build(), + new PortableIcebergDestinations( + "db.{dest}", + FileFormat.PARQUET.name(), + inputSchema, + null, + null, + null, + ImmutableList.of(SEQ_COL), + null, + null), + "test-runId")); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + List, KV>> list = + ImmutableList.copyOf(iter); + assertThat(list, hasSize(2)); + // Bind each record to its destination: row id=1 routed to tableA, id=2 to tableB. + Map idByDest = new HashMap<>(); + for (KV, KV> kv : list) { + idByDest.put( + kv.getKey().getKey(), kv.getValue().getValue().getData().getInt32("id")); + } + assertThat(idByDest, equalTo(ImmutableMap.of("db." + tableA, 1, "db." + tableB, 2))); + return null; + }); + PAssert.that(outputs.get(AssignCdcKeys.FAILED)).empty(); + p.run().waitUntilFinish(); + } + + // ------------------------------------------------------------------------------------------- + // Table-level failures bypass error handling + // ------------------------------------------------------------------------------------------- + + @Test + public void tableConfigExceptionPropagatesDespiteErrorHandling() { + TableIdentifier id = uniqueId("v1_table"); + CdcSinkTestUtils.createTable( + catalog, id, ICEBERG_SCHEMA, ImmutableSet.of(1), 1, PartitionSpec.unpartitioned()); + CdcWriteConfig config = cdcWriteConfig().setErrorHandling(true).build(); + + assignKeys(input(INPUT_SCHEMA, KV.of(ValueKind.INSERT, dataRow(1, "a", "x", 1L))), config, id); + + assertPipelineFailsMentioning("format-version", "append sink"); + } + + // ------------------------------------------------------------------------------------------- + // Shard distribution + // ------------------------------------------------------------------------------------------- + + @Test + public void samePkAlwaysSameShardAndDistinctPksSpread() { + TableIdentifier id = createCanonicalTable(); + + // 100 distinct primary keys, two records each (INSERT then DELETE). + List> rows = new ArrayList<>(); + for (int i = 0; i < 100; i++) { + rows.add(KV.of(ValueKind.INSERT, dataRow(i, "n" + i, "d", 1L))); + rows.add(KV.of(ValueKind.DELETE, dataRow(i, "n" + i, "d", 2L))); + } + PCollection in = + CdcSinkTestUtils.withKinds(p.apply(Create.of(rows))).setRowSchema(INPUT_SCHEMA); + + PCollectionTuple outputs = assignKeys(in, cdcWriteConfig().build(), id); + + PAssert.that(outputs.get(AssignCdcKeys.KEYED)) + .satisfies( + iter -> { + Map> shardsByPk = new HashMap<>(); + Set allShards = new HashSet<>(); + int count = 0; + for (KV, KV> kv : iter) { + count++; + int pk = kv.getValue().getValue().getData().getInt32("id"); + int shard = kv.getKey().getValue(); + shardsByPk.computeIfAbsent(pk, unused -> new HashSet<>()).add(shard); + allShards.add(shard); + } + assertThat(count, equalTo(200)); + // Every primary key maps to exactly one shard across its records. + for (Map.Entry> entry : shardsByPk.entrySet()) { + assertThat( + "pk " + entry.getKey() + " mapped to multiple shards", + entry.getValue(), + hasSize(1)); + } + // Distinct primary keys spread over more than one shard. + assertThat(allShards.size(), greaterThan(1)); + return null; + }); + p.run().waitUntilFinish(); + } +}