diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests.json b/.github/trigger_files/IO_Iceberg_Integration_Tests.json index 37dd25bf9029..b73af5e61a43 100644 --- a/.github/trigger_files/IO_Iceberg_Integration_Tests.json +++ b/.github/trigger_files/IO_Iceberg_Integration_Tests.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 3 + "modification": 1 } diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests_Dataflow.json b/.github/trigger_files/IO_Iceberg_Integration_Tests_Dataflow.json index 3a009261f4f9..5abe02fc09c7 100644 --- a/.github/trigger_files/IO_Iceberg_Integration_Tests_Dataflow.json +++ b/.github/trigger_files/IO_Iceberg_Integration_Tests_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 2 + "modification": 1 } diff --git a/CHANGES.md b/CHANGES.md index fcb011d1489f..4ea769676172 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -68,6 +68,7 @@ * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). * Add ArrowFlight IO (Java) ([#20116](https://github.com/apache/beam/issues/20116)). * (Python) JmsIO (IBM MQ, ActiveMQ, and other providers) is now supported in Python via cross-language ([#30716](https://github.com/apache/beam/issues/30716)). +* Added a full Iceberg batch and streaming changelog source (CDC) ([#38831](https://github.com/apache/beam/issues/38831)) ## New Features / Improvements diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProvider.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProvider.java index e029a85a812f..30930e880607 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProvider.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProvider.java @@ -117,7 +117,10 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { .streaming(configuration.getStreaming()) .keeping(configuration.getKeep()) .dropping(configuration.getDrop()) - .withFilter(configuration.getFilter()); + .withFilter(configuration.getFilter()) + .withWatermarkColumn(configuration.getWatermarkColumn()) + .withWatermarkColumnTimeUnit(configuration.getWatermarkColumnTimeUnit()) + .withMetadataColumns(configuration.getIncludeMetadataColumns()); @Nullable Integer pollIntervalSeconds = configuration.getPollIntervalSeconds(); if (pollIntervalSeconds != null) { @@ -193,6 +196,26 @@ static Builder builder() { "A subset of column names to exclude from reading. If null or empty, all columns will be read.") abstract @Nullable List getDrop(); + @SchemaFieldDescription( + "Column used to derive the source's output watermark. " + + "Must be an existing, required, top-level column of type 'long' or 'timestamp'. " + + "If not set, the watermark advances according to snapshot commit timestamp.") + abstract @Nullable String getWatermarkColumn(); + + @SchemaFieldDescription( + "Time unit used to interpret watermark column of type LONG. One of NANOSECONDS, MICROSECONDS, " + + "MILLISECONDS, SECONDS, MINUTES, HOURS, DAYS. Defaults to MICROSECONDS.") + abstract @Nullable String getWatermarkColumnTimeUnit(); + + @SchemaFieldDescription( + "List of top-level metadata columns to include with CDC output rows. Supported columns: \n" + + "- `_change_type`\n" + + "- `_row_id`\n" + + "- `_last_updated_sequence_number`\n" + + "- `_commit_snapshot_id`\n" + + "- `_commit_snapshot_sequence_number`\n") + abstract @Nullable List getIncludeMetadataColumns(); + @AutoValue.Builder abstract static class Builder { abstract Builder setTable(String table); @@ -223,6 +246,12 @@ abstract static class Builder { abstract Builder setFilter(String filter); + abstract Builder setWatermarkColumn(String watermarkColumn); + + abstract Builder setWatermarkColumnTimeUnit(String timeUnit); + + abstract Builder setIncludeMetadataColumns(List metadataColumns); + abstract Configuration build(); } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java index ee5755898b7f..78a72ccdbb8a 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java @@ -25,6 +25,7 @@ import java.util.Map; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.io.iceberg.cdc.IncrementalChangelogSource; import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.transforms.PTransform; @@ -33,6 +34,7 @@ import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.iceberg.DistributionMode; import org.apache.iceberg.Table; import org.apache.iceberg.catalog.Catalog; @@ -576,6 +578,7 @@ public static ReadRows readRows(IcebergCatalogConfig catalogConfig) { return new AutoValue_IcebergIO_ReadRows.Builder() .setCatalogConfig(catalogConfig) .setUseCdc(false) + .setMetadataColumns(ImmutableList.of()) .build(); } @@ -612,6 +615,12 @@ public enum StartingStrategy { abstract @Nullable String getFilter(); + abstract @Nullable String getWatermarkColumn(); + + abstract @Nullable String getWatermarkColumnTimeUnit(); + + abstract List getMetadataColumns(); + abstract Builder toBuilder(); @AutoValue.Builder @@ -642,6 +651,12 @@ abstract static class Builder { abstract Builder setFilter(@Nullable String filter); + abstract Builder setWatermarkColumn(@Nullable String watermarkColumn); + + abstract Builder setWatermarkColumnTimeUnit(@Nullable String timeUnit); + + abstract Builder setMetadataColumns(List metadataColumns); + abstract ReadRows build(); } @@ -693,6 +708,31 @@ public ReadRows withFilter(@Nullable String filter) { return toBuilder().setFilter(filter).build(); } + public ReadRows withWatermarkColumn(@Nullable String watermarkColumn) { + return toBuilder().setWatermarkColumn(watermarkColumn).build(); + } + + public ReadRows withWatermarkColumnTimeUnit(@Nullable String timeUnit) { + return toBuilder().setWatermarkColumnTimeUnit(timeUnit).build(); + } + + /** + * Appends top-level metadata columns to CDC output rows. + * + *

Supported values are {@code _change_type}, {@code _commit_snapshot_id}, {@code + * _commit_snapshot_sequence_number}, {@code _row_id}, and {@code + * _last_updated_sequence_number}. The row metadata columns are read from Iceberg data files and + * require a row-lineage table. The changelog metadata columns come from the emitted change kind + * and snapshot context and are appended when final Beam rows are emitted. + * + *

This option is only valid {@link #withCdc()}. + */ + public ReadRows withMetadataColumns(@Nullable List metadataColumns) { + return toBuilder() + .setMetadataColumns(metadataColumns == null ? ImmutableList.of() : metadataColumns) + .build(); + } + @Override public PCollection expand(PBegin input) { TableIdentifier tableId = @@ -728,12 +768,15 @@ public PCollection expand(PBegin input) { .setKeepFields(getKeep()) .setDropFields(getDrop()) .setFilterString(getFilter()) + .setWatermarkColumn(getWatermarkColumn()) + .setWatermarkColumnTimeUnit(getWatermarkColumnTimeUnit()) + .setMetadataColumns(getMetadataColumns()) .build(); scanConfig.validate(table); PTransform> source = getUseCdc() - ? new IncrementalScanSource(scanConfig) + ? new IncrementalChangelogSource(scanConfig) : Read.from(new ScanSource(scanConfig)); return input.apply(source); 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 45ec21f0ca51..bcd574afdaba 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 @@ -271,9 +271,6 @@ public Expression getFilter() { @Pure public abstract @Nullable String getWatermarkColumnTimeUnit(); - @Pure - public abstract @Nullable Duration getMaxSnapshotDiscoveryDelay(); - @Pure public abstract List getMetadataColumns(); @@ -371,8 +368,6 @@ public abstract Builder setUpdateCompatibilityVersion( public abstract Builder setWatermarkColumnTimeUnit(@Nullable String timeUnit); - public abstract Builder setMaxSnapshotDiscoveryDelay(@Nullable Duration delay); - public abstract Builder setMetadataColumns(List metadataColumns); public abstract IcebergScanConfig build(); diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IncrementalScanSource.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IncrementalScanSource.java deleted file mode 100644 index 98870095e171..000000000000 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IncrementalScanSource.java +++ /dev/null @@ -1,102 +0,0 @@ -/* - * 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; - -import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; - -import java.util.List; -import org.apache.beam.sdk.coders.KvCoder; -import org.apache.beam.sdk.coders.ListCoder; -import org.apache.beam.sdk.coders.StringUtf8Coder; -import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.transforms.PTransform; -import org.apache.beam.sdk.transforms.ParDo; -import org.apache.beam.sdk.transforms.Redistribute; -import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.PBegin; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.Row; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; -import org.apache.iceberg.Table; -import org.checkerframework.checker.nullness.qual.Nullable; -import org.joda.time.Duration; - -/** - * An Iceberg source that reads a table incrementally using range(s) of table snapshots. The bounded - * source creates a single range, while the unbounded implementation continuously polls for new - * snapshots at the specified interval. - */ -class IncrementalScanSource extends PTransform> { - private static final Duration DEFAULT_POLL_INTERVAL = Duration.standardSeconds(60); - private final IcebergScanConfig scanConfig; - - IncrementalScanSource(IcebergScanConfig scanConfig) { - this.scanConfig = scanConfig; - } - - @Override - public PCollection expand(PBegin input) { - Table table = - TableCache.get( - scanConfig.getCatalogConfig(), - IcebergUtils.parseTableIdentifier(scanConfig.getTableIdentifier())); - - PCollection>> snapshots = - MoreObjects.firstNonNull(scanConfig.getStreaming(), false) - ? unboundedSnapshots(input) - : boundedSnapshots(input, table); - - return snapshots - .setCoder(KvCoder.of(StringUtf8Coder.of(), ListCoder.of(SnapshotInfo.getCoder()))) - .apply(Redistribute.byKey()) - .apply("Create Read Tasks", ParDo.of(new CreateReadTasksDoFn(scanConfig))) - .setCoder(KvCoder.of(ReadTaskDescriptor.getCoder(), ReadTask.getCoder())) - .apply(Redistribute.arbitrarily()) - .apply("Read Rows From Tasks", ParDo.of(new ReadFromTasks(scanConfig))) - .setRowSchema( - IcebergUtils.icebergSchemaToBeamSchema( - scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion())); - } - - /** Continuously watches for new snapshots. */ - private PCollection>> unboundedSnapshots(PBegin input) { - Duration pollInterval = - MoreObjects.firstNonNull(scanConfig.getPollInterval(), DEFAULT_POLL_INTERVAL); - return input.apply("Watch for Snapshots", new WatchForSnapshots(scanConfig, pollInterval)); - } - - /** Creates a fixed snapshot range. */ - private PCollection>> boundedSnapshots(PBegin input, Table table) { - checkStateNotNull( - table.currentSnapshot().snapshotId(), - "Table %s does not have any snapshots to read from.", - scanConfig.getTableIdentifier()); - - @Nullable Long from = ReadUtils.getFromSnapshotExclusive(table, scanConfig); - // if no end snapshot is provided, we read up to the current snapshot. - long to = - MoreObjects.firstNonNull( - ReadUtils.getToSnapshot(table, scanConfig), table.currentSnapshot().snapshotId()); - return input.apply( - "Create Snapshot Range", - Create.of( - KV.of( - scanConfig.getTableIdentifier(), - ReadUtils.snapshotsBetween(table, scanConfig.getTableIdentifier(), from, to)))); - } -} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java deleted file mode 100644 index 438e2de464d6..000000000000 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java +++ /dev/null @@ -1,98 +0,0 @@ -/* - * 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; - -import java.io.IOException; -import java.util.List; -import java.util.concurrent.ExecutionException; -import org.apache.beam.sdk.io.range.OffsetRange; -import org.apache.beam.sdk.metrics.Counter; -import org.apache.beam.sdk.metrics.Metrics; -import org.apache.beam.sdk.schemas.Schema; -import org.apache.beam.sdk.transforms.DoFn; -import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; -import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.Row; -import org.apache.iceberg.FileScanTask; -import org.apache.iceberg.Table; -import org.apache.iceberg.data.Record; -import org.apache.iceberg.io.CloseableIterable; - -/** - * Bounded read implementation. - * - *

For each {@link ReadTask}, reads Iceberg {@link Record}s, and converts to Beam {@link Row}s. - * - *

Implemented as an SDF to leverage communicating bundle size (i.e. {@link DoFn.GetSize}) to the - * runner, to help with scaling decisions. - */ -@DoFn.BoundedPerElement -class ReadFromTasks extends DoFn, Row> { - private final IcebergScanConfig scanConfig; - private final Counter scanTasksCompleted = - Metrics.counter(ReadFromTasks.class, "scanTasksCompleted"); - - ReadFromTasks(IcebergScanConfig scanConfig) { - this.scanConfig = scanConfig; - } - - @ProcessElement - public void process( - @Element KV element, - RestrictionTracker tracker, - OutputReceiver out) - throws IOException, ExecutionException, InterruptedException { - ReadTask readTask = element.getValue(); - Table table = TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()); - - List fileScanTasks = readTask.getFileScanTasks(); - - for (long l = tracker.currentRestriction().getFrom(); - l < tracker.currentRestriction().getTo(); - l++) { - if (!tracker.tryClaim(l)) { - return; - } - FileScanTask task = fileScanTasks.get((int) l); - Schema beamSchema = - IcebergUtils.icebergSchemaToBeamSchema( - scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion()); - try (CloseableIterable reader = ReadUtils.createReader(task, table, scanConfig)) { - - for (Record record : reader) { - Row row = IcebergUtils.icebergRecordToBeamRow(beamSchema, record); - out.output(row); - } - } - scanTasksCompleted.inc(); - } - } - - @GetSize - public double getSize( - @Element KV element, @Restriction OffsetRange restriction) { - // TODO(ahmedabu98): this is actually the file byte size, likely compressed. - // find a way to output the actual Beam Row byte size. - return element.getValue().getSize(restriction.getFrom(), restriction.getTo()); - } - - @GetInitialRestriction - public OffsetRange getInitialRange(@Element KV element) { - return new OffsetRange(0, element.getValue().getFileScanTaskJsons().size()); - } -} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WatchForSnapshots.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WatchForSnapshots.java deleted file mode 100644 index 8bd436c55700..000000000000 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WatchForSnapshots.java +++ /dev/null @@ -1,190 +0,0 @@ -/* - * 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; - -import static org.apache.beam.sdk.transforms.Watch.Growth.PollResult; - -import java.util.List; -import java.util.stream.Collectors; -import org.apache.beam.sdk.coders.ListCoder; -import org.apache.beam.sdk.metrics.Counter; -import org.apache.beam.sdk.metrics.Gauge; -import org.apache.beam.sdk.metrics.Metrics; -import org.apache.beam.sdk.state.StateSpec; -import org.apache.beam.sdk.state.StateSpecs; -import org.apache.beam.sdk.state.ValueState; -import org.apache.beam.sdk.transforms.Create; -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.Watch; -import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.PBegin; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.TimestampedValue; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Objects; -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.Iterables; -import org.apache.iceberg.Snapshot; -import org.apache.iceberg.Table; -import org.checkerframework.checker.nullness.qual.Nullable; -import org.joda.time.Duration; -import org.joda.time.Instant; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Keeps watch over an Iceberg table and continuously outputs a range of snapshots, at the specified - * interval. - * - *

A downstream transform will create a list of read tasks for each range. - */ -class WatchForSnapshots extends PTransform>>> { - private static final Logger LOG = LoggerFactory.getLogger(WatchForSnapshots.class); - private final Duration pollInterval; - private final IcebergScanConfig scanConfig; - - WatchForSnapshots(IcebergScanConfig scanConfig, Duration pollInterval) { - this.pollInterval = pollInterval; - this.scanConfig = scanConfig; - } - - @Override - public PCollection>> expand(PBegin input) { - return input - .apply(Create.of(scanConfig.getTableIdentifier())) - .apply( - "Scan Table Snapshots", - Watch.growthOf(new SnapshotPollFn(scanConfig)) - .withPollInterval(pollInterval) - .withOutputCoder(ListCoder.of(SnapshotInfo.getCoder()))) - .apply("Persist Snapshot Progress", ParDo.of(new PersistSnapshotProgress())); - } - - /** - * Periodically scans the table for new snapshots, emitting a list for each new snapshot range. - * - *

This tracks progress locally but is not resilient to retries -- upon worker failure, it will - * restart from the initial starting strategy. Resilience is handled downstream by {@link - * PersistSnapshotProgress}. - */ - private static class SnapshotPollFn extends Watch.Growth.PollFn> { - private final IcebergScanConfig scanConfig; - private @Nullable Long fromSnapshotId; - - SnapshotPollFn(IcebergScanConfig scanConfig) { - this.scanConfig = scanConfig; - } - - @Override - public PollResult> apply(String tableIdentifier, Context c) { - Table table = TableCache.getRefreshed(scanConfig.getCatalogConfig(), tableIdentifier); - - @Nullable Long userSpecifiedToSnapshot = ReadUtils.getToSnapshot(table, scanConfig); - boolean isComplete = userSpecifiedToSnapshot != null; - if (fromSnapshotId == null) { - // first scan, initialize starting point with user config - fromSnapshotId = ReadUtils.getFromSnapshotExclusive(table, scanConfig); - } - - Snapshot currentSnapshot = table.currentSnapshot(); - if (currentSnapshot == null || Objects.equal(currentSnapshot.snapshotId(), fromSnapshotId)) { - // no new snapshots since last poll. return empty result. - return getPollResult(null, isComplete); - } - - Long currentSnapshotId = currentSnapshot.snapshotId(); - // if no upper bound is specified, we poll up to the current snapshot - long toSnapshotId = MoreObjects.firstNonNull(userSpecifiedToSnapshot, currentSnapshotId); - - List snapshots = - ReadUtils.snapshotsBetween(table, tableIdentifier, fromSnapshotId, toSnapshotId); - - fromSnapshotId = currentSnapshotId; - return getPollResult(snapshots, isComplete); - } - - private PollResult> getPollResult( - @Nullable List snapshots, boolean isComplete) { - ImmutableList.Builder>> timestampedSnapshots = - ImmutableList.builder(); - if (snapshots != null) { - // watermark based on the oldest observed snapshot in this poll interval - Instant watermark = Instant.ofEpochMilli(snapshots.get(0).getTimestampMillis()); - timestampedSnapshots.add(TimestampedValue.of(snapshots, watermark)); - } - - return isComplete - ? PollResult.complete(timestampedSnapshots.build()) // stop at specified snapshot - : PollResult.incomplete(timestampedSnapshots.build()); // continue forever - } - } - - /** - * Stateful DoFn that persists the latest observed snapshot ID to state, making sure we pick up - * where we left off in case of a worker crash. - */ - // Ideally, Watch.Growth would support state out of the box, but that is a bigger change. - static class PersistSnapshotProgress - extends DoFn>, KV>> { - private final Gauge latestSnapshot = Metrics.gauge(SnapshotPollFn.class, "latestSnapshot"); - private final Counter snapshotsObserved = - Metrics.counter(SnapshotPollFn.class, "snapshotsObserved"); - - @StateId("latestObservedSnapshotId") - @SuppressWarnings("UnusedVariable") - private final StateSpec> latestObservedSnapshotId = StateSpecs.value(); - - @ProcessElement - public void process( - @Element KV> element, - final @AlwaysFetched @StateId("latestObservedSnapshotId") ValueState - latestObservedSnapshotId, - OutputReceiver>> out) { - List snapshots = element.getValue(); - - @Nullable Long latest = latestObservedSnapshotId.read(); - if (latest != null) { - int newSnapshotIndex = 0; - for (int i = 0; i < snapshots.size(); i++) { - if (snapshots.get(i).getSnapshotId() == latest) { - newSnapshotIndex = i + 1; - break; - } - } - if (newSnapshotIndex > 0) { - snapshots = snapshots.subList(newSnapshotIndex, snapshots.size()); - } - } - - SnapshotInfo checkpoint = Iterables.getLast(snapshots); - out.output(KV.of(element.getKey(), snapshots)); - LOG.info( - "New poll fetched {} snapshots: {}. Checkpointing at snapshot {} of timestamp {}.", - snapshots.size(), - snapshots.stream().map(SnapshotInfo::getSnapshotId).collect(Collectors.toList()), - checkpoint.getSnapshotId(), - checkpoint.getTimestampMillis()); - - latestObservedSnapshotId.write(checkpoint.getSnapshotId()); - latestSnapshot.set(checkpoint.getSnapshotId()); - snapshotsObserved.inc(snapshots.size()); - } - } -} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumn.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumn.java new file mode 100644 index 000000000000..0b312904fe26 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumn.java @@ -0,0 +1,99 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static java.util.concurrent.TimeUnit.MICROSECONDS; + +import java.time.LocalDateTime; +import java.util.concurrent.TimeUnit; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.util.Preconditions; +import org.apache.beam.sdk.values.Row; +import org.apache.iceberg.util.DateTimeUtil; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; + +/** + * Re-stamps each output row using the configured {@code watermarkColumn}'s value, so the source's + * output watermark advances per record rather than per snapshot. + * + *

If the configured column's value on a record is null or missing, this DoFn is a pass-through, + * preserving the snapshot commit timestamp. + * + *

The {@link #getAllowedTimestampSkew()} return is intentionally generous — the user's watermark + * column may produce values well before the snapshot commit time (event-time data can lag + * wall-clock by hours or days). Restricting the skew here would force the source to drop legitimate + * output. + */ +class ApplyWatermarkColumn extends DoFn { + private final String watermarkColumn; + private final TimeUnit timeUnit; + + ApplyWatermarkColumn(String watermarkColumn, @Nullable String timeUnit) { + this.watermarkColumn = watermarkColumn; + this.timeUnit = timeUnit != null ? TimeUnit.valueOf(timeUnit.toUpperCase()) : MICROSECONDS; + } + + @ProcessElement + public void process(@Element Row row, OutputReceiver out) { + @Nullable + Instant instant = + getInstant(row.getValue(watermarkColumn), row.getSchema().getField(watermarkColumn)); + if (instant != null) { + out.outputWithTimestamp(row, instant); + } else { + out.output(row); + } + } + + private @Nullable Instant getInstant(@Nullable Object value, Schema.Field field) { + if (value == null) { + return null; + } + switch (field.getType().getTypeName()) { + case INT64: + return Instant.ofEpochMilli(timeUnit.toMillis((Long) value)); + case DATETIME: + return (Instant) value; + case LOGICAL_TYPE: + String logicalType = + Preconditions.checkStateNotNull(field.getType().getLogicalType()).getIdentifier(); + if (logicalType.equals(SqlTypes.DATETIME.getIdentifier())) { + return Instant.ofEpochMilli( + MICROSECONDS.toMillis(DateTimeUtil.microsFromTimestamp((LocalDateTime) value))); + } else if (logicalType.equals(SqlTypes.TIMESTAMP.getIdentifier()) + || logicalType.equals(org.apache.beam.sdk.schemas.logicaltypes.Timestamp.IDENTIFIER)) { + return Instant.ofEpochMilli( + MICROSECONDS.toMillis(DateTimeUtil.microsFromInstant((java.time.Instant) value))); + } else { + throw new UnsupportedOperationException("Unexpected logical type: " + logicalType); + } + default: + throw new UnsupportedOperationException("Unexpected Beam type: " + field.getType()); + } + } + + @Override + public Duration getAllowedTimestampSkew() { + // Generous skew to cover backfill of historical data and late-arriving CDC patterns. + return Duration.standardDays(365); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java index 8a3a543854d9..a421b3af276c 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java @@ -96,6 +96,21 @@ static Schema readBeamSchemaWithRowMetadata(List metadataColumns, Schema return builder.build(); } + static Row outputRow( + List metadataColumns, + Schema outputSchema, + ChangelogDescriptor descriptor, + ValueKind valueKind, + Row dataAndRowMetadata) { + return outputRow( + metadataColumns, + outputSchema, + descriptor.getCommitSnapshotId(), + descriptor.getSnapshotSequenceNumber(), + valueKind, + dataAndRowMetadata); + } + /** * Builds the final public Beam row. * @@ -106,7 +121,8 @@ static Schema readBeamSchemaWithRowMetadata(List metadataColumns, Schema static Row outputRow( List metadataColumns, Schema outputSchema, - ChangelogDescriptor descriptor, + long commitSnapshotId, + long snapshotSequenceNumber, ValueKind valueKind, Row dataAndRowMetadata) { if (metadataColumns.isEmpty() @@ -114,9 +130,6 @@ static Row outputRow( return dataAndRowMetadata; } - long commitSnapshotId = descriptor.getCommitSnapshotId(); - long snapshotSequentNumber = descriptor.getSnapshotSequenceNumber(); - List<@Nullable Object> values = new ArrayList<>(outputSchema.getFieldCount()); for (Schema.Field field : dataAndRowMetadata.getSchema().getFields()) { if (!metadataColumns.contains(field.getName())) { @@ -129,7 +142,7 @@ static Row outputRow( metadataValue( metadataColumn, commitSnapshotId, - snapshotSequentNumber, + snapshotSequenceNumber, valueKind, dataAndRowMetadata)); } @@ -137,9 +150,11 @@ static Row outputRow( } static Schema readBeamSchemaWithRowMetadata( - List metadataColumns, org.apache.iceberg.Schema dataSchema) { + List metadataColumns, + org.apache.iceberg.Schema dataSchema, + @Nullable String updateCompatibilityVersion) { return IcebergUtils.icebergSchemaToBeamSchema( - readSchemaWithRowMetadata(metadataColumns, dataSchema)); + readSchemaWithRowMetadata(metadataColumns, dataSchema), updateCompatibilityVersion); } private static @Nullable Object metadataValue( diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolver.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolver.java index be2191688965..284c10abe44b 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolver.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolver.java @@ -80,12 +80,23 @@ abstract class CdcResolver { * Resolves a Primary Key group of deletes and inserts. Caller provides {@code emit} which decides * how to materialize each output. * - *

In the rare case of duplicate PKs within a snapshot, one side may hold more than one record. - * When this happens, we re-order the lists by {@link #nonPkHash} so the result is deterministic. + *

The dominant case (unique identifier values) is exactly one delete and one insert, decided + * directly by a single {@link #nonPkEquals} with no hashing. In the rare case of duplicate PKs + * within a snapshot, one side may hold more than one record. When this happens, we re-order the + * lists by {@link #nonPkHash} so the result is deterministic. */ final void resolve(List deletes, List inserts, BiConsumer emit) { - // Fast path: with unique identifier values each side holds at most one record, so there is - // only one possible pairing and nothing to order. + if (deletes.size() == 1 && inserts.size() == 1) { + // No-op if non-PK fields are equal, otherwise we emit an update pair + T delete = deletes.get(0); + T insert = inserts.get(0); + if (!nonPkEquals(delete, insert)) { + emit.accept(ValueKind.UPDATE_BEFORE, delete); + emit.accept(ValueKind.UPDATE_AFTER, insert); + } + return; + } + if (deletes.size() > 1 || inserts.size() > 1) { resolveOrdered(sortedByNonPkHash(deletes), sortedByNonPkHash(inserts), emit); } else { 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 new file mode 100644 index 000000000000..fe6240d5dc53 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java @@ -0,0 +1,211 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.apache.beam.sdk.io.iceberg.cdc.ChangelogScanner.LARGE_BIDIRECTIONAL_TASKS; +import static org.apache.beam.sdk.io.iceberg.cdc.ChangelogScanner.SMALL_BIDIRECTIONAL_TASKS; +import static org.apache.beam.sdk.io.iceberg.cdc.ChangelogScanner.UNIDIRECTIONAL_TASKS; +import static org.apache.beam.sdk.io.iceberg.cdc.ResolveChanges.DELETES; +import static org.apache.beam.sdk.io.iceberg.cdc.ResolveChanges.INSERTS; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.ReadUtils; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Redistribute; +import org.apache.beam.sdk.transforms.join.CoGroupByKey; +import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple; +import org.apache.beam.sdk.transforms.windowing.AfterWatermark; +import org.apache.beam.sdk.transforms.windowing.DefaultTrigger; +import org.apache.beam.sdk.transforms.windowing.GlobalWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.apache.iceberg.Table; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; + +/** + * An Iceberg source that incrementally reads a table's changelogs, processing one snapshot at a + * time. + * + *

Each snapshot is resolved independently. For a given primary key, this source emits the net + * change the snapshot produces, and not its intermediate states. + * + *

Implications: if a writer batches several transitions for the same PK into one snapshot (for + * example {@code A → B, then B → C}), only the endpoints survive. The intermediate {@code B} is + * dropped. A round-trip within a single snapshot (for example {@code A → B → A}) is also dropped. + * + *

The streaming path uses {@link WatchForSnapshotsSdf} for proper per-snapshot watermarks. The + * bounded path creates the snapshot range up front. + */ +public class IncrementalChangelogSource extends PTransform> { + private final IcebergScanConfig scanConfig; + + public IncrementalChangelogSource(IcebergScanConfig scanConfig) { + this.scanConfig = scanConfig; + } + + @Override + public PCollection expand(PBegin input) { + // emit one SnapshotInfo per element, with element timestamp -> snapshot commit time. + PCollection snapshots = + MoreObjects.firstNonNull(scanConfig.getStreaming(), false) + ? unboundedSnapshots(input) + : boundedSnapshots(input); + + // process one snapshot at a time and produce batches of changelog scan tasks. + // tasks are emitted to three outputs: + // 1. unidirectional tasks: we know these won't have any updates + // 2. small bidirectional tasks: these may contain an update, but the batch is small enough to + // resolve in-memory + // 2. large bidirectional tasks: may contain an update, but are too large for in-memory + // resolution. will + // need to run these output rows through a CoGBK + PCollectionTuple changelogTasks = + snapshots.apply( + "Create Changelog Tasks", + ParDo.of(new ChangelogScanner(scanConfig)) + .withOutputTags( + UNIDIRECTIONAL_TASKS, + TupleTagList.of(LARGE_BIDIRECTIONAL_TASKS).and(SMALL_BIDIRECTIONAL_TASKS))); + KvCoder> tasksCoder = + ChangelogScanner.coder(scanConfig.rowIdBeamSchema()); + changelogTasks.get(UNIDIRECTIONAL_TASKS).setCoder(tasksCoder); + changelogTasks.get(SMALL_BIDIRECTIONAL_TASKS).setCoder(tasksCoder); + changelogTasks.get(LARGE_BIDIRECTIONAL_TASKS).setCoder(tasksCoder); + + Schema projectedRowSchema = + IcebergUtils.icebergSchemaToBeamSchema( + scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion()); + Schema outputRowSchema = CdcOutputUtils.outputSchema(scanConfig, projectedRowSchema); + + // reads UNIDIRECTIONAL and BIDIRECTIONAL tags and produces rows. + ReadFromChangelogs.Output outputRows = changelogTasks.apply(new ReadFromChangelogs(scanConfig)); + + // Small overlapping groups get resolved entirely in memory with no shuffle. + PCollection smallBidirectionalCdcRows = + changelogTasks + .get(SMALL_BIDIRECTIONAL_TASKS) + .apply("Redistribute Small Bidirectional Changes", Redistribute.arbitrarily()) + .apply("Resolve Locally", ParDo.of(new LocalResolveDoFn(scanConfig))) + .setRowSchema(outputRowSchema); + + // BIDIRECTIONAL records go through a CoGBK and ResolveChanges + // We window locally using a custom WindowFn based on the snapsot's commit time. Each snapshot + // exists in its own window. + // We re-window the resolved output back to GlobalWindows before the final Flatten + // to align with the other branches. + Window> keyedWindowing = + Window.>into(new SnapshotWindowFn()) + .triggering(AfterWatermark.pastEndOfWindow()) + .withAllowedLateness(Duration.ZERO) + .discardingFiredPanes(); + PCollection> keyedInserts = + outputRows.biDirectionalInserts().apply("Window Inserts", keyedWindowing); + PCollection> keyedDeletes = + outputRows.biDirectionalDeletes().apply("Window Deletes", keyedWindowing); + PCollection biDirectionalCdcRows = + KeyedPCollectionTuple.of(INSERTS, keyedInserts) + .and(DELETES, keyedDeletes) + .apply("CoGroupBy Primary Key", CoGroupByKey.create()) + .apply("Resolve Delete-Insert Pairs", ParDo.of(new ResolveChanges(scanConfig))) + .setRowSchema(outputRowSchema) + .apply( + "Re-window to Global", + Window.into(new GlobalWindows()) + .triggering(DefaultTrigger.of()) + .discardingFiredPanes()); + + // Merge all three paths into a single output. All three are in GlobalWindows. + PCollection merged = + PCollectionList.of(outputRows.uniDirectionalRows()) + .and(smallBidirectionalCdcRows) + .and(biDirectionalCdcRows) + .apply(Flatten.pCollections()); + + // If the user configures a watermark column, restamp each record by + // that column's value. Output watermark then advances per-record rather than per-snapshot. + @Nullable String watermarkColumn = scanConfig.getWatermarkColumn(); + if (watermarkColumn != null) { + merged = + merged.apply( + "Apply Watermark Column", + ParDo.of( + new ApplyWatermarkColumn( + watermarkColumn, scanConfig.getWatermarkColumnTimeUnit()))); + } + + return merged.setRowSchema(outputRowSchema); + } + + /** + * Continuously watches the Iceberg table for new snapshots via {@link WatchForSnapshotsSdf} and + * emits per snapshot. + */ + private PCollection unboundedSnapshots(PBegin input) { + return input + .apply("Impulse", Create.of("")) + .apply("Watch for Snapshots", ParDo.of(new WatchForSnapshotsSdf(scanConfig))); + } + + /** + * Reads the full snapshot range up front and emits each snapshot individually, each carrying its + * own commit time as the element timestamp. + */ + private PCollection boundedSnapshots(PBegin input) { + Table table = + scanConfig + .getCatalogConfig() + .catalog() + .loadTable(IcebergUtils.parseTableIdentifier(scanConfig.getTableIdentifier())); + checkStateNotNull( + table.currentSnapshot(), + "Table %s does not have any snapshots to read from.", + scanConfig.getTableIdentifier()); + + @Nullable Long from = ReadUtils.getFromSnapshotExclusive(table, scanConfig); + long to = + MoreObjects.firstNonNull( + ReadUtils.getToSnapshot(table, scanConfig), table.currentSnapshot().snapshotId()); + List> timestamped = + ReadUtils.snapshotsBetween(table, scanConfig.getTableIdentifier(), from, to).stream() + .map( + s -> + TimestampedValue.of( + s.getSnapshotId(), Instant.ofEpochMilli(s.getTimestampMillis()))) + .collect(Collectors.toList()); + return input.apply("Create Snapshot Range", Create.timestamped(timestamped)); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java index a3188b3a0245..683a7ed86132 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java @@ -82,10 +82,14 @@ class LocalResolveDoFn extends DoFn, Row> { + static final TupleTag DELETES = new TupleTag<>() {}; + static final TupleTag INSERTS = new TupleTag<>() {}; + private final IcebergScanConfig scanConfig; + private final RowFilter rowFilter; + private final Schema outputSchema; + // Positions and types of the non-PK data fields in the input row schema, precomputed once so + // the per-record hash/equals loops need no name lookups. The input schema is fixed by the + // CoGroupByKey's coder, so positions are stable across elements. + private final int[] nonPkIndices; + private final Schema.FieldType[] nonPkTypes; + private transient @MonotonicNonNull RowResolver resolver; + + ResolveChanges(IcebergScanConfig scanConfig) { + this.scanConfig = scanConfig; + Schema inputSchema = + CdcOutputUtils.readBeamSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getSchema()); + this.rowFilter = + new RowFilter(inputSchema) + .keep( + CdcOutputUtils.readSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getProjectedSchema()) + .columns().stream() + .map(Types.NestedField::name) + .collect(Collectors.toList())); + this.outputSchema = + CdcOutputUtils.outputSchema( + scanConfig, + IcebergUtils.icebergSchemaToBeamSchema( + scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion())); + + Set pkFields = new HashSet<>(scanConfig.rowIdBeamSchema().getFieldNames()); + List metadataColumns = scanConfig.getMetadataColumns(); + List indices = new ArrayList<>(); + List types = new ArrayList<>(); + List fields = inputSchema.getFields(); + for (int i = 0; i < fields.size(); i++) { + Schema.Field field = fields.get(i); + String name = field.getName(); + if (pkFields.contains(name) + || (IcebergCdcMetadataColumns.isSupportedColumn(name) + && metadataColumns.contains(name))) { + continue; + } + indices.add(i); + types.add(field.getType()); + } + this.nonPkIndices = indices.stream().mapToInt(Integer::intValue).toArray(); + this.nonPkTypes = types.toArray(new Schema.FieldType[0]); + } + + @Setup + public void setup() { + this.resolver = new RowResolver(nonPkIndices, nonPkTypes); + } + + @ProcessElement + public void processElement( + @Element KV element, + @Timestamp Instant timestamp, + OutputReceiver out) { + CdcRowDescriptor descriptor = element.getKey(); + CoGbkResult result = element.getValue(); + + // should be okay to materialize these lists. a PK collision will likely be a handful of records + // at most + List deletes = Lists.newArrayList(result.getAll(DELETES)); + List inserts = Lists.newArrayList(result.getAll(INSERTS)); + + checkStateNotNull(resolver) + .resolve( + deletes, + inserts, + (kind, row) -> { + Row projectedRow = rowFilter.filter(row); + out.builder( + CdcOutputUtils.outputRow( + scanConfig.getMetadataColumns(), + outputSchema, + descriptor.getCommitSnapshotId(), + descriptor.getSnapshotSequenceNumber(), + kind, + projectedRow)) + .setValueKind(kind) + .setTimestamp(timestamp) + .output(); + }); + } + + /** Resolver specialization over Beam Rows, using precomputed non-PK field positions. */ + private static final class RowResolver extends CdcResolver { + private final int[] nonPkIndices; + private final Schema.FieldType[] nonPkTypes; + + RowResolver(int[] nonPkIndices, Schema.FieldType[] nonPkTypes) { + this.nonPkIndices = nonPkIndices; + this.nonPkTypes = nonPkTypes; + } + + @Override + protected int nonPkHash(Row element) { + int hash = 1; + for (int i = 0; i < nonPkIndices.length; i++) { + hash = + 31 * hash + Row.Equals.deepHashCode(element.getValue(nonPkIndices[i]), nonPkTypes[i]); + } + return hash; + } + + @Override + protected boolean nonPkEquals(Row delete, Row insert) { + // compare non-PK, we already know PK values are equal + for (int i = 0; i < nonPkIndices.length; i++) { + int idx = nonPkIndices[i]; + // return early if two values are not equal + if (!Row.Equals.deepEquals(insert.getValue(idx), delete.getValue(idx), nonPkTypes[i])) { + return false; + } + } + return true; + } + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/SnapshotWindowFn.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/SnapshotWindowFn.java new file mode 100644 index 000000000000..06dbc1740bdb --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/SnapshotWindowFn.java @@ -0,0 +1,87 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import java.util.Collection; +import java.util.Collections; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.transforms.windowing.NonMergingWindowFn; +import org.apache.beam.sdk.transforms.windowing.WindowFn; +import org.apache.beam.sdk.transforms.windowing.WindowMappingFn; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; + +/** + * A {@link WindowFn} that assigns each element to a 1-millisecond {@link IntervalWindow} anchored + * at the element's event timestamp. + * + *

We set the element's timestamp as its snapshot commit timestamp. All tasks/records from the + * same snapshot land in the same window. + * + *

With the per-snapshot watermark from {@link WatchForSnapshotsSdf}, the CoGroupByKey fires when + * a snapshot is fully drained. The watermark advances past the snapshot's commit time only after + * every downstream stage has finished processing that snapshot's records. + * + *

Two snapshots committed within the same millisecond may collapse into the same window. But + * that's okay because {@link ReadFromChangelogs} includes snapshot sequence number in the key + * before routing to the CoGBK, so it won't produce incorrect joins. + */ +public class SnapshotWindowFn extends NonMergingWindowFn { + private static final Duration WINDOW_LENGTH = Duration.millis(1); + + @Override + public Collection assignWindows(AssignContext c) { + Instant ts = c.timestamp(); + return Collections.singletonList(new IntervalWindow(ts, ts.plus(WINDOW_LENGTH))); + } + + @Override + public boolean isCompatible(WindowFn other) { + return other instanceof SnapshotWindowFn; + } + + @Override + public Coder windowCoder() { + return IntervalWindow.getCoder(); + } + + @Override + public WindowMappingFn getDefaultWindowMappingFn() { + // Just return a window covering the main-input window's end timestamp. + return new WindowMappingFn<>() { + @Override + public IntervalWindow getSideInputWindow(BoundedWindow mainWindow) { + Instant end = mainWindow.maxTimestamp(); + return new IntervalWindow(end, end.plus(WINDOW_LENGTH)); + } + }; + } + + @Override + public boolean equals(@Nullable Object obj) { + return obj instanceof SnapshotWindowFn; + } + + @Override + public int hashCode() { + return SnapshotWindowFn.class.hashCode(); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/WatchForSnapshotsSdf.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/WatchForSnapshotsSdf.java index 473fea7bfa81..658c019b30e7 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/WatchForSnapshotsSdf.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/WatchForSnapshotsSdf.java @@ -62,10 +62,15 @@ * {@code @ProcessElement} claims the sequence numbers of newly discovered snapshots in * chronological order. * - *

Uses a {@link Manual} watermark estimator. After emitting a snapshot, the watermark is set to - * that snapshot's commit time. On empty polls, the watermark is bumped to {@code now() - - * MAX_SNAPSHOT_DISCOVERY_DELAY} to prevent downstream windows from stalling indefinitely during - * quiet periods. + *

Uses a {@link Manual} watermark estimator. After emitting a snapshot, the watermark is set + * just past that snapshot's timestamp, so its {@link SnapshotWindowFn} window can fire as + * soon as its records drain. On empty polls, the watermark is bumped to {@code now() - + * pollInterval} to prevent downstream windows from stalling indefinitely during quiet periods. + * + *

If a snapshot is somehow discovered after the watermark has already moved past its commit time + * (e.g. idle bump ran ahead of a slow/downed catalog), it will be emitted with its timestamp + * clamped to the current watermark instead. This guarantees that no snapshot is ever late, and no + * records are silently dropped. */ @DoFn.UnboundedPerElement class WatchForSnapshotsSdf extends DoFn { @@ -75,10 +80,10 @@ class WatchForSnapshotsSdf extends DoFn { private static final Counter snapshotsEmitted = Metrics.counter(WatchForSnapshotsSdf.class, "snapshotsEmitted"); + private static final Counter lateDiscoveredSnapshots = + Metrics.counter(WatchForSnapshotsSdf.class, "lateDiscoveredSnapshots"); private static final Gauge latestEmittedSnapshotId = Metrics.gauge(WatchForSnapshotsSdf.class, "latestEmittedSnapshotId"); - // TODO(ahmedabu98): consider exposing this as a config option - private static final Duration MAX_SNAPSHOT_DISCOVERY_DELAY = Duration.standardMinutes(5); private static final Long POLL_FOREVER = Long.MAX_VALUE; private final IcebergScanConfig scanConfig; @@ -231,35 +236,45 @@ public ProcessContinuation process( if (!tracker.tryClaim(snap.getSequenceNumber())) { return ProcessContinuation.stop(); } - Instant ts = Instant.ofEpochMilli(snap.getTimestampMillis()); + Instant commitTs = Instant.ofEpochMilli(snap.getTimestampMillis()); + Instant ts = commitTs; + if (ts.isBefore(watermark.currentWatermark())) { + // The watermark already moved past this snapshot's commit time (e.g. the idle bump ran + // ahead of a slow discovery). Use the current watermark so the snapshot is not dropped + ts = watermark.currentWatermark(); + lateDiscoveredSnapshots.inc(); + LOG.warn( + "Snapshot {} (commit ts: {}) was discovered after the watermark already advanced " + + "to {}. Emitting it with the current watermark.", + snap.getSnapshotId(), + commitTs, + ts); + } out.outputWithTimestamp(snap.getSnapshotId(), ts); - if (watermark.currentWatermark().isBefore(ts)) { - watermark.setWatermark(ts); - } + // Advance just past `ts` so this snapshot's window can fire as + // soon as its records drain + watermark.setWatermark(ts.plus(Duration.millis(1))); snapshotsEmitted.inc(); latestEmittedSnapshotId.set(snap.getSnapshotId()); LOG.info( - "Emitted snapshot {} (sequence id: {}, commit ts: {})", + "Emitted snapshot {} (sequence id: {}, timestamp: {})", snap.getSnapshotId(), snap.getSequenceNumber(), ts); } - return pauseOrStop(watermark, bounded); + return continueOrStop(bounded); } /** - * On an empty poll, bump the watermark to {@code now() - MAX_SNAPSHOT_DISCOVERY_DELAY} so - * downstream windows can still fire. Returns {@code stop()} when end snapshot has been reached, - * otherwise {@code resume()} after the poll interval. + * On an empty poll, bump the watermark to {@code now() - pollInterval} so downstream windows and + * timers can make progress while the table is quiet. Returns {@code stop()} when end snapshot has + * been reached, otherwise {@code resume()} after the poll interval. */ private ProcessContinuation pauseOrStop( ManualWatermarkEstimator watermark, boolean bounded) { - Duration delay = - MoreObjects.firstNonNull( - scanConfig.getMaxSnapshotDiscoveryDelay(), MAX_SNAPSHOT_DISCOVERY_DELAY); - Instant idleWatermark = Instant.now().minus(delay); + Instant idleWatermark = Instant.now().minus(pollInterval); if (watermark.currentWatermark().isBefore(idleWatermark)) { LOG.info( "Sitting idle for {} seconds. Bumping watermark to {}", @@ -268,6 +283,10 @@ private ProcessContinuation pauseOrStop( idleWatermark); watermark.setWatermark(idleWatermark); } + return continueOrStop(bounded); + } + + private ProcessContinuation continueOrStop(boolean bounded) { return bounded ? ProcessContinuation.stop() : ProcessContinuation.resume().withResumeDelay(pollInterval); 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 5849cbd00774..9d12389184a7 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 @@ -28,6 +28,8 @@ import java.util.Map; import java.util.UUID; import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.apache.beam.sdk.io.iceberg.cdc.IcebergCdcMetadataColumns; import org.apache.beam.sdk.managed.Managed; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.testing.PAssert; @@ -35,13 +37,17 @@ import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionRowTuple; import org.apache.beam.sdk.values.Row; +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.collect.Lists; import org.apache.iceberg.CatalogUtil; import org.apache.iceberg.Snapshot; import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types; import org.junit.ClassRule; import org.junit.Rule; import org.junit.Test; @@ -56,6 +62,15 @@ public class IcebergCdcReadSchemaTransformProviderTest { private static final org.apache.iceberg.Schema CDC_SCHEMA = new org.apache.iceberg.Schema(TestFixtures.SCHEMA.columns(), ImmutableSet.of(1)); + private static final org.apache.iceberg.Schema CDC_CONFIG_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "data", Types.StringType.get()), + Types.NestedField.optional(3, "category", Types.StringType.get()), + Types.NestedField.required(4, "event_micros", Types.LongType.get())), + ImmutableSet.of(1)); + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); @Rule public TestPipeline testPipeline = TestPipeline.create(); @@ -77,6 +92,14 @@ public void testBuildTransformWithRow() { .withFieldValue("to_timestamp", 456L) .withFieldValue("starting_strategy", "earliest") .withFieldValue("poll_interval_seconds", 789) + .withFieldValue("keep", ImmutableList.of("id", "data", "event_micros")) + .withFieldValue("filter", "\"category\" = 'include'") + .withFieldValue("watermark_column", "event_micros") + .withFieldValue( + "include_metadata_columns", + ImmutableList.of( + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID, + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER)) .build(); new IcebergCdcReadSchemaTransformProvider().from(config); @@ -163,4 +186,95 @@ public void testStreamingReadUsingManagedTransform() throws Exception { testPipeline.run(); } + + @Test + public void testManagedReadWithProjectionFilterWatermarkAndSnapshotRange() throws Exception { + String identifier = "default.table_" + Long.toString(UUID.randomUUID().hashCode(), 16); + TableIdentifier tableId = TableIdentifier.parse(identifier); + + Table table = + warehouse.createTable( + tableId, CDC_CONFIG_SCHEMA, null, ImmutableMap.of(TableProperties.FORMAT_VERSION, "3")); + long eventMicros = (System.currentTimeMillis() - 1_000L) * 1_000L; + List records = + ImmutableList.of( + record(1L, "keep-a", "include", eventMicros), + record(2L, "drop", "exclude", eventMicros + 1_000L), + record(3L, "keep-b", "include", eventMicros + 2_000L)); + table + .newFastAppend() + .appendFile(warehouse.writeRecords("cdc-managed-config.parquet", table.schema(), records)) + .commit(); + + Map properties = new HashMap<>(); + properties.put("type", CatalogUtil.ICEBERG_CATALOG_TYPE_HADOOP); + properties.put("warehouse", warehouse.location); + + Map configMap = new HashMap<>(); + configMap.put("table", identifier); + configMap.put("catalog_name", "test-name"); + 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")); + configMap.put("filter", "\"category\" = 'include'"); + configMap.put("watermark_column", "event_micros"); + configMap.put( + "include_metadata_columns", + ImmutableList.of( + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID, + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER, + IcebergCdcMetadataColumns.ROW_ID, + IcebergCdcMetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER)); + + org.apache.iceberg.Schema projectedSchema = table.schema().select("id", "data", "event_micros"); + Schema recordSchema = IcebergUtils.icebergSchemaToBeamSchema(projectedSchema); + Schema outputSchema = + Schema.builder() + .addFields(recordSchema.getFields()) + .addInt64Field(IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID) + .addInt64Field(IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER) + .addNullableField(IcebergCdcMetadataColumns.ROW_ID, Schema.FieldType.INT64) + .addNullableField( + IcebergCdcMetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER, Schema.FieldType.INT64) + .build(); + long snapshotId = table.currentSnapshot().snapshotId(); + long sequenceNumber = table.currentSnapshot().sequenceNumber(); + long firstRowId = table.currentSnapshot().firstRowId(); + List expectedRows = + IntStream.range(0, records.size()) + .filter(i -> "include".equals(records.get(i).getField("category"))) + .mapToObj( + i -> { + 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, + sequenceNumber) + .build(); + }) + .collect(Collectors.toList()); + + PCollection output = + testPipeline + .apply(Managed.read(Managed.ICEBERG_CDC).withConfig(configMap)) + .getSinglePCollection(); + + assertThat(output.isBounded(), equalTo(BOUNDED)); + assertThat(output.getSchema(), equalTo(outputSchema)); + PAssert.that(output).containsInAnyOrder(expectedRows); + + testPipeline.run(); + } + + private static Record record(long id, String data, String category, long eventMicros) { + return TestFixtures.createRecord( + CDC_CONFIG_SCHEMA, + ImmutableMap.of("id", id, "data", data, "category", category, "event_micros", eventMicros)); + } } 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 new file mode 100644 index 000000000000..e928d3cf8f01 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java @@ -0,0 +1,270 @@ +/* + * 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; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.cdc.IcebergCdcMetadataColumns; +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.iceberg.CatalogUtil; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.types.Types; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +/** Tests for {@link IcebergScanConfig}. */ +public class IcebergScanConfigTest { + @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + + private static final org.apache.iceberg.Schema CDC_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "data", Types.StringType.get()), + Types.NestedField.optional(3, "category", Types.StringType.get()), + Types.NestedField.required(4, "event_time", Types.TimestampType.withoutZone()), + Types.NestedField.required(5, "event_micros", Types.LongType.get()), + Types.NestedField.optional(6, "optional_time", Types.TimestampType.withoutZone()), + Types.NestedField.required(7, "required_text", Types.StringType.get()), + Types.NestedField.required( + 8, + "nested", + Types.StructType.of( + Types.NestedField.required( + 9, "nested_time", Types.TimestampType.withoutZone())))), + ImmutableSet.of(1)); + + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); + + @Test + public void cdcValidationRequiresIdentifierFields() { + TableIdentifier tableId = uniqueTableId(); + Table table = warehouse.createTable(tableId, TestFixtures.SCHEMA); + IcebergScanConfig scanConfig = + scanConfigBuilder(tableId, TestFixtures.SCHEMA).setUseCdc(true).build(); + + IllegalStateException thrown = + assertThrows(IllegalStateException.class, () -> scanConfig.validate(table)); + assertThat(thrown.getMessage(), containsString("Cannot read CDC records")); + assertThat(thrown.getMessage(), containsString("primary key fields")); + } + + @Test + public void cdcValidationRejectsProjectionDroppingIdentifierFields() { + TableIdentifier tableId = uniqueTableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + + IcebergScanConfig keepWithoutPk = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setKeepFields(ImmutableList.of("data")) + .build(); + IllegalArgumentException keepException = + assertThrows(IllegalArgumentException.class, () -> keepWithoutPk.validate(table)); + assertThat( + keepException.getMessage(), + containsString("projected schema must not drop primary key fields")); + + IcebergScanConfig dropPk = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setDropFields(ImmutableList.of("id")) + .build(); + IllegalArgumentException dropException = + assertThrows(IllegalArgumentException.class, () -> dropPk.validate(table)); + assertThat( + dropException.getMessage(), + containsString("projected schema must not drop primary key fields")); + } + + @Test + public void requiredSchemaIncludesFilterOnlyFieldsWithoutChangingProjection() { + TableIdentifier tableId = uniqueTableId(); + warehouse.createTable(tableId, CDC_SCHEMA); + IcebergScanConfig scanConfig = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setKeepFields(ImmutableList.of("id")) + .setFilterString("\"data\" = 'keep' AND \"category\" = 'include'") + .build(); + + assertEquals(ImmutableList.of("id"), fieldNames(scanConfig.getProjectedSchema())); + assertEquals( + ImmutableSet.of("id", "data", "category"), + ImmutableSet.copyOf(fieldNames(scanConfig.getRequiredSchema()))); + } + + @Test + public void metadataColumnsRequireCdcMode() { + TableIdentifier tableId = uniqueTableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + IcebergScanConfig scanConfig = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setMetadataColumns(ImmutableList.of(IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID)) + .build(); + + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> scanConfig.validate(table)); + assertThat(thrown.getMessage(), containsString("metadata_columns")); + } + + @Test + public void metadataColumnsRejectUnsupportedAndDuplicateNames() { + TableIdentifier tableId = uniqueTableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + + IcebergScanConfig unsupported = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setMetadataColumns(ImmutableList.of("_missing_metadata")) + .build(); + IllegalArgumentException unsupportedThrown = + assertThrows(IllegalArgumentException.class, () -> unsupported.validate(table)); + assertThat(unsupportedThrown.getMessage(), containsString("unsupported metadata_columns")); + + IcebergScanConfig duplicate = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setMetadataColumns( + ImmutableList.of( + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID, + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID)) + .build(); + IllegalArgumentException duplicateThrown = + assertThrows(IllegalArgumentException.class, () -> duplicate.validate(table)); + assertThat(duplicateThrown.getMessage(), containsString("duplicate")); + } + + @Test + public void rowLineageMetadataRequiresFormatV3Table() { + TableIdentifier tableId = uniqueTableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + IcebergScanConfig scanConfig = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setMetadataColumns(ImmutableList.of(IcebergCdcMetadataColumns.ROW_ID)) + .build(); + + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> scanConfig.validate(table)); + assertThat(thrown.getMessage(), containsString("format v3+")); + } + + @Test + public void watermarkColumnAcceptsRequiredTimestampAndLongColumns() { + TableIdentifier tableId = uniqueTableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setKeepFields(ImmutableList.of("id", "event_time")) + .setWatermarkColumn("event_time") + .build() + .validate(table); + + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(true) + .setKeepFields(ImmutableList.of("id", "event_micros")) + .setWatermarkColumn("event_micros") + .build() + .validate(table); + } + + @Test + public void watermarkColumnRejectsInvalidConfigurations() { + TableIdentifier tableId = uniqueTableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + + assertInvalidWatermark( + tableId, table, "event_time", false, ImmutableList.of("id", "event_time"), "CDC mode"); + assertInvalidWatermark( + tableId, table, "missing", true, ImmutableList.of("id"), "unknown column"); + assertInvalidWatermark( + tableId, + table, + "optional_time", + true, + ImmutableList.of("id", "optional_time"), + "non-nullable"); + assertInvalidWatermark( + tableId, + table, + "required_text", + 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( + TableIdentifier tableId, + Table table, + String watermarkColumn, + boolean useCdc, + List keepFields, + String expectedMessage) { + IcebergScanConfig scanConfig = + scanConfigBuilder(tableId, CDC_SCHEMA) + .setUseCdc(useCdc) + .setKeepFields(keepFields) + .setWatermarkColumn(watermarkColumn) + .build(); + + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> scanConfig.validate(table)); + assertThat(thrown.getMessage(), containsString(expectedMessage)); + } + + private IcebergScanConfig.Builder scanConfigBuilder( + TableIdentifier tableId, org.apache.iceberg.Schema schema) { + return IcebergScanConfig.builder() + .setCatalogConfig( + IcebergCatalogConfig.builder() + .setCatalogName("name") + .setCatalogProperties( + ImmutableMap.of( + "type", + CatalogUtil.ICEBERG_CATALOG_TYPE_HADOOP, + "warehouse", + warehouse.location)) + .build()) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(schema)); + } + + private static List fieldNames(org.apache.iceberg.Schema schema) { + return schema.columns().stream().map(Types.NestedField::name).collect(Collectors.toList()); + } + + private static TableIdentifier uniqueTableId() { + return TableIdentifier.of( + "default", "table_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergSchemaTransformTranslationTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergSchemaTransformTranslationTest.java index 1319efa7229a..675b4aafe76a 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergSchemaTransformTranslationTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergSchemaTransformTranslationTest.java @@ -118,6 +118,7 @@ public class IcebergSchemaTransformTranslationTest { .withFieldValue("streaming", true) .withFieldValue("keep", ImmutableList.of("id", "event_micros")) .withFieldValue("filter", "\"data\" = 'keep'") + .withFieldValue("watermark_column", "event_micros") .build(); @Test diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java index 5c28f0192a61..2435cdb231da 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/catalog/IcebergCatalogBaseIT.java @@ -56,6 +56,7 @@ import org.apache.beam.sdk.extensions.gcp.util.GcsUtil; import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.cdc.IcebergCdcMetadataColumns; import org.apache.beam.sdk.managed.Managed; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; @@ -63,7 +64,9 @@ 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.transforms.DoFn; import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.PeriodicImpulse; import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.transforms.SimpleFunction; @@ -75,10 +78,16 @@ import org.apache.beam.sdk.values.PCollection.IsBounded; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.beam.sdk.values.ValueKind; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; 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.iceberg.AppendFiles; +import org.apache.iceberg.ChangelogOperation; import org.apache.iceberg.CombinedScanTask; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.NullOrder; import org.apache.iceberg.PartitionSpec; @@ -86,15 +95,22 @@ import org.apache.iceberg.SortDirection; import org.apache.iceberg.SortOrder; import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; import org.apache.iceberg.TableScan; import org.apache.iceberg.catalog.Catalog; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.SupportsNamespaces; import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.IdentityPartitionConverters; import org.apache.iceberg.data.Record; import org.apache.iceberg.data.parquet.GenericParquetReaders; import org.apache.iceberg.data.parquet.GenericParquetWriter; +import org.apache.iceberg.deletes.EqualityDeleteWriter; +import org.apache.iceberg.deletes.PositionDelete; +import org.apache.iceberg.deletes.PositionDeleteWriter; +import org.apache.iceberg.encryption.EncryptedFiles; import org.apache.iceberg.encryption.InputFilesDecryptor; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.DataWriter; @@ -276,6 +292,8 @@ public void cleanUp() throws Exception { .addLogicalTypeField("date", SqlTypes.DATE) .addLogicalTypeField("time", SqlTypes.TIME) .build(); + private static final Schema CDC_BEAM_SCHEMA = + Schema.builder().addInt64Field("id").addStringField("data").build(); private static final SimpleFunction ROW_FUNC = new SimpleFunction() { @@ -321,6 +339,9 @@ public Row apply(Long num) { protected static final org.apache.iceberg.Schema ICEBERG_SCHEMA = new org.apache.iceberg.Schema( beamSchemaToIcebergSchema(BEAM_SCHEMA).columns(), Collections.singleton(1)); + private static final org.apache.iceberg.Schema CDC_ICEBERG_SCHEMA = + new org.apache.iceberg.Schema( + beamSchemaToIcebergSchema(CDC_BEAM_SCHEMA).columns(), Collections.singleton(1)); protected static final SimpleFunction RECORD_FUNC = new SimpleFunction() { @Override @@ -561,6 +582,131 @@ public void testStreamingReadWithColumnPruning_drop() throws Exception { pipeline.run().waitUntilFinish(); } + @Test + public void testStreamingCdcReadMixedDeleteAndOverwriteSnapshots() throws Exception { + Table table = createCdcTable(); + DataFile firstFile = + commitCdcAppend( + table, + "cdc-first-data.parquet", + Arrays.asList( + cdcRecord(table.schema(), 1L, "first-file-update-before"), + cdcRecord(table.schema(), 2L, "first-file-delete"))); + Snapshot firstSnapshot = checkStateNotNull(table.currentSnapshot()); + + DataFile secondFile = + commitCdcAppend( + table, + "cdc-second-data.parquet", + Arrays.asList( + cdcRecord(table.schema(), 3L, "second-file-unchanged"), + cdcRecord(table.schema(), 4L, "second-file-update-before"))); + Snapshot secondSnapshot = checkStateNotNull(table.currentSnapshot()); + + DeleteFile equalityDelete = writeCdcEqualityDelete(table, "cdc-equality-delete.parquet", 2L); + DeleteFile positionDelete = + writeCdcPositionDelete(table, "cdc-position-delete.parquet", firstFile, 0L); + DataFile thirdFile = + writeCdcDataFile( + table, + "cdc-third-data.parquet", + Collections.singletonList(cdcRecord(table.schema(), 1L, "third-file-update-after"))); + table + .newRowDelta() + .addDeletes(equalityDelete) + .addDeletes(positionDelete) + .addRows(thirdFile) + .commit(); + table.refresh(); + Snapshot thirdSnapshot = checkStateNotNull(table.currentSnapshot()); + + DataFile fourthFile = + writeCdcDataFile( + table, + "cdc-fourth-data.parquet", + Arrays.asList( + cdcRecord(table.schema(), 3L, "second-file-unchanged"), + cdcRecord(table.schema(), 4L, "fourth-file-update-after"))); + table.newOverwrite().deleteFile(secondFile).addFile(fourthFile).commit(); + table.refresh(); + Snapshot fourthSnapshot = checkStateNotNull(table.currentSnapshot()); + + Map config = new HashMap<>(managedIcebergConfig(tableId())); + config.put("from_snapshot", firstSnapshot.snapshotId()); + config.put("to_snapshot", fourthSnapshot.snapshotId()); + config.put("streaming", true); + + PCollection rows = + pipeline.apply(Managed.read(ICEBERG_CDC).withConfig(config)).getSinglePCollection(); + + PCollection changes = rows.apply("Format CDC Changes", ParDo.of(new FormatCdcChange())); + + assertThat(rows.isBounded(), equalTo(UNBOUNDED)); + assertEquals(CDC_BEAM_SCHEMA, rows.getSchema()); + PAssert.that(changes) + .containsInAnyOrder( + cdcChange(ValueKind.INSERT, firstSnapshot, 1L, "first-file-update-before"), + cdcChange(ValueKind.INSERT, firstSnapshot, 2L, "first-file-delete"), + cdcChange(ValueKind.INSERT, secondSnapshot, 3L, "second-file-unchanged"), + cdcChange(ValueKind.INSERT, secondSnapshot, 4L, "second-file-update-before"), + cdcChange(ValueKind.UPDATE_BEFORE, thirdSnapshot, 1L, "first-file-update-before"), + cdcChange(ValueKind.UPDATE_AFTER, thirdSnapshot, 1L, "third-file-update-after"), + cdcChange(ValueKind.DELETE, thirdSnapshot, 2L, "first-file-delete"), + cdcChange(ValueKind.UPDATE_BEFORE, fourthSnapshot, 4L, "second-file-update-before"), + cdcChange(ValueKind.UPDATE_AFTER, fourthSnapshot, 4L, "fourth-file-update-after")); + pipeline.run().waitUntilFinish(); + } + + @Test + public void testCdcReadWithMetadataColumns() throws Exception { + Table table = + catalog.createTable( + TableIdentifier.parse(tableId()), + CDC_ICEBERG_SCHEMA, + PartitionSpec.unpartitioned(), + ImmutableMap.of(TableProperties.FORMAT_VERSION, "3")); + commitCdcAppend( + table, + "metadata-columns.parquet", + Arrays.asList(cdcRecord(table.schema(), 1L, "one"), cdcRecord(table.schema(), 2L, "two"))); + Snapshot snapshot = checkStateNotNull(table.currentSnapshot()); + long firstRowId = checkStateNotNull(snapshot.firstRowId()); + + Map config = new HashMap<>(managedIcebergConfig(tableId())); + config.put("to_snapshot", snapshot.snapshotId()); + config.put( + "include_metadata_columns", + Arrays.asList( + IcebergCdcMetadataColumns.CHANGE_TYPE, + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID, + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER, + IcebergCdcMetadataColumns.ROW_ID, + IcebergCdcMetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER)); + + Schema outputSchema = + Schema.builder() + .addInt64Field("id") + .addStringField("data") + .addStringField(IcebergCdcMetadataColumns.CHANGE_TYPE) + .addInt64Field(IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID) + .addInt64Field(IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER) + .addNullableField(IcebergCdcMetadataColumns.ROW_ID, Schema.FieldType.INT64) + .addNullableField( + IcebergCdcMetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER, Schema.FieldType.INT64) + .build(); + + PCollection rows = + pipeline.apply(Managed.read(ICEBERG_CDC).withConfig(config)).getSinglePCollection(); + + assertEquals(BOUNDED, rows.isBounded()); + assertEquals(outputSchema, rows.getSchema()); + PAssert.that(rows) + .containsInAnyOrder( + cdcMetadataRow(1L, "one", snapshot, firstRowId, outputSchema), + cdcMetadataRow(2L, "two", snapshot, firstRowId + 1, outputSchema)); + pipeline.run().waitUntilFinish(); + } + @Test public void testBatchReadBetweenSnapshots() throws Exception { runReadBetween(true, false); @@ -622,7 +768,9 @@ public void testWriteReadWithFilter() throws IOException { @Test public void testReadWriteStreaming() throws IOException { - Table table = catalog.createTable(TableIdentifier.parse(tableId()), ICEBERG_SCHEMA); + org.apache.iceberg.Schema schemaWithPk = + new org.apache.iceberg.Schema(ICEBERG_SCHEMA.columns(), ImmutableSet.of(1)); + Table table = catalog.createTable(TableIdentifier.parse(tableId()), schemaWithPk); List expectedRows = populateTable(table); Map config = managedIcebergConfig(tableId()); @@ -1036,7 +1184,9 @@ && checkStateNotNull(rec.getBoolean("bool_field")) == bool) } public void runReadBetween(boolean useSnapshotBoundary, boolean streaming) throws Exception { - Table table = catalog.createTable(TableIdentifier.parse(tableId()), ICEBERG_SCHEMA); + org.apache.iceberg.Schema schemaWithPk = + new org.apache.iceberg.Schema(ICEBERG_SCHEMA.columns(), ImmutableSet.of(1)); + Table table = catalog.createTable(TableIdentifier.parse(tableId()), schemaWithPk); populateTable(table, "a"); // first snapshot Thread.sleep(AFTER_UPDATE_SLEEP_MS); @@ -1069,6 +1219,133 @@ public void runReadBetween(boolean useSnapshotBoundary, boolean streaming) throw pipeline.run().waitUntilFinish(); } + private Table createCdcTable() { + return catalog.createTable( + TableIdentifier.parse(tableId()), + CDC_ICEBERG_SCHEMA, + PartitionSpec.unpartitioned(), + ImmutableMap.of( + TableProperties.FORMAT_VERSION, + "2", + TableProperties.SPLIT_SIZE, + "1", + TableProperties.DEFAULT_WRITE_METRICS_MODE, + "full")); + } + + private static String cdcChange(ValueKind valueKind, Snapshot snapshot, long id, String data) { + return String.format("%s:%d:%d:%s", valueKind, snapshot.timestampMillis(), id, data); + } + + private static Row cdcMetadataRow( + long id, String data, Snapshot snapshot, long rowId, Schema outputSchema) { + return Row.withSchema(outputSchema) + .addValues( + id, + data, + ChangelogOperation.INSERT.name(), + snapshot.snapshotId(), + snapshot.sequenceNumber(), + rowId, + snapshot.sequenceNumber()) + .build(); + } + + private static Record cdcRecord(org.apache.iceberg.Schema schema, long id, String data) { + GenericRecord record = GenericRecord.create(schema); + record.setField("id", id); + if (schema.findField("data") != null) { + record.setField("data", data); + } + return record; + } + + private DataFile commitCdcAppend(Table table, String filename, List records) + throws IOException { + DataFile dataFile = writeCdcDataFile(table, filename, records); + table.newFastAppend().appendFile(dataFile).commit(); + table.refresh(); + return dataFile; + } + + private DataFile writeCdcDataFile(Table table, String filename, List records) + throws IOException { + OutputFile file = + table.io().newOutputFile(table.location() + "/" + UUID.randomUUID() + "-" + filename); + DataWriter writer = + Parquet.writeData(file) + .schema(table.schema()) + .createWriterFunc(GenericParquetWriter::create) + .overwrite() + .withSpec(table.spec()) + .build(); + + try (writer) { + for (Record record : records) { + writer.write(record); + } + } + + return writer.toDataFile(); + } + + private DeleteFile writeCdcEqualityDelete(Table table, String filename, long id) + throws IOException { + org.apache.iceberg.Schema deleteSchema = table.schema().select("id"); + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec(), new int[] {1}, deleteSchema, null); + EqualityDeleteWriter writer = + appenderFactory.newEqDeleteWriter( + EncryptedFiles.plainAsEncryptedOutput( + table + .io() + .newOutputFile(table.location() + "/" + UUID.randomUUID() + "-" + filename)), + FileFormat.PARQUET, + null); + + try (writer) { + writer.write(cdcRecord(deleteSchema, id, null)); + } + + return writer.toDeleteFile(); + } + + private DeleteFile writeCdcPositionDelete( + Table table, String filename, DataFile dataFile, long... positions) throws IOException { + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec()); + PositionDeleteWriter writer = + appenderFactory.newPosDeleteWriter( + EncryptedFiles.plainAsEncryptedOutput( + table + .io() + .newOutputFile(table.location() + "/" + UUID.randomUUID() + "-" + filename)), + FileFormat.PARQUET, + null); + + try (writer) { + for (long position : positions) { + writer.write(PositionDelete.create().set(dataFile.location(), position)); + } + } + + return writer.toDeleteFile(); + } + + private static final class FormatCdcChange extends DoFn { + @ProcessElement + public void process( + @Element Row row, + ValueKind valueKind, + @DoFn.Timestamp Instant timestamp, + OutputReceiver outputReceiver) { + outputReceiver.output( + String.format( + "%s:%d:%d:%s", + valueKind, timestamp.getMillis(), row.getInt64("id"), row.getString("data"))); + } + } + @Test public void testWriteWithTableProperties() throws IOException { Map config = new HashMap<>(managedIcebergConfig(tableId())); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumnTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumnTest.java new file mode 100644 index 000000000000..b28f5c2c91f8 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumnTest.java @@ -0,0 +1,158 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.junit.Assert.assertThrows; + +import java.time.LocalDateTime; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; +import org.apache.beam.sdk.schemas.logicaltypes.Timestamp; +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.transforms.DoFnTester; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Reify; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Instant; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link ApplyWatermarkColumn}. */ +@RunWith(JUnit4.class) +public class ApplyWatermarkColumnTest { + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + @Test + public void stampsRowsFromSupportedWatermarkTypes() { + assertWatermarkTimestamp( + "jodaDateTime", + Schema.builder().addStringField("id").addDateTimeField("wm").build(), + Row.withSchema(Schema.builder().addStringField("id").addDateTimeField("wm").build()) + .addValues("joda", new Instant(1_234L)) + .build(), + new Instant(1_234L)); + + assertWatermarkTimestamp( + "longMicros", + Schema.builder().addStringField("id").addInt64Field("wm").build(), + Row.withSchema(Schema.builder().addStringField("id").addInt64Field("wm").build()) + .addValues("long", 1_234_567L) + .build(), + new Instant(1_234L)); + + assertWatermarkTimestamp( + "localDateTime", + Schema.builder().addStringField("id").addLogicalTypeField("wm", SqlTypes.DATETIME).build(), + Row.withSchema( + Schema.builder() + .addStringField("id") + .addLogicalTypeField("wm", SqlTypes.DATETIME) + .build()) + .addValues("ldt", LocalDateTime.of(1969, 12, 31, 23, 59, 59, 123_000_000)) + .build(), + new Instant(-877L)); + + assertWatermarkTimestamp( + "javaInstant", + Schema.builder().addStringField("id").addLogicalTypeField("wm", Timestamp.MICROS).build(), + Row.withSchema( + Schema.builder() + .addStringField("id") + .addLogicalTypeField("wm", Timestamp.MICROS) + .build()) + .addValues("instant", java.time.Instant.parse("1969-12-31T23:59:59.123Z")) + .build(), + new Instant(-877L)); + + pipeline.run(); + } + + @Test + public void nullWatermarkValuePreservesInputTimestamp() { + Schema schema = + Schema.of( + Schema.Field.of("id", Schema.FieldType.STRING), + Schema.Field.nullable("wm", Schema.FieldType.DATETIME)); + Row row = Row.withSchema(schema).addValues("null", null).build(); + Instant inputTimestamp = new Instant(99L); + + PCollection output = + pipeline + .apply( + Create.timestamped(TimestampedValue.of(row, inputTimestamp)) + .withCoder(RowCoder.of(schema))) + .apply(ParDo.of(new ApplyWatermarkColumn("wm", "microseconds"))); + output.setCoder(RowCoder.of(schema)); + + PAssert.that(output.apply(Reify.timestamps())) + .containsInAnyOrder(TimestampedValue.of(row, inputTimestamp)); + + pipeline.run(); + } + + @Test + public void unsupportedWatermarkTypeThrows() { + Schema schema = Schema.builder().addStringField("wm").build(); + Row row = Row.withSchema(schema).addValue("2026-05-24T00:00:00Z").build(); + + assertThrows( + UnsupportedOperationException.class, + () -> { + try (DoFnTester tester = + DoFnTester.of(new ApplyWatermarkColumn("wm", "microseconds"))) { + tester.processElement(row); + } + }); + } + + @Test + public void missingWatermarkColumnThrows() { + Schema schema = Schema.builder().addStringField("other").build(); + Row row = Row.withSchema(schema).addValue("value").build(); + + assertThrows( + IllegalArgumentException.class, + () -> { + try (DoFnTester tester = + DoFnTester.of(new ApplyWatermarkColumn("wm", "microseconds"))) { + tester.processElement(row); + } + }); + } + + private void assertWatermarkTimestamp( + String name, Schema schema, Row row, Instant expectedTimestamp) { + PCollection output = + pipeline + .apply(name + "Create", Create.of(row).withCoder(RowCoder.of(schema))) + .apply( + name + "ApplyWatermark", ParDo.of(new ApplyWatermarkColumn("wm", "microseconds"))); + output.setCoder(RowCoder.of(schema)); + + PCollection> timestamps = + output.apply(name + "ReifyTimestamp", Reify.timestamps()); + PAssert.that(timestamps).containsInAnyOrder(TimestampedValue.of(row, expectedTimestamp)); + } +} 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 new file mode 100644 index 000000000000..14e7208dad0c --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java @@ -0,0 +1,514 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.apache.iceberg.types.Types.NestedField.optional; +import static org.apache.iceberg.types.Types.NestedField.required; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.junit.Assert.assertEquals; + +import java.io.IOException; +import java.time.LocalDateTime; +import java.time.ZoneOffset; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergIO.ReadRows.StartingStrategy; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +import org.apache.beam.sdk.io.iceberg.TestFixtures; +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.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Reify; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollection.IsBounded; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TimestampedValue; +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.collect.Lists; +import org.apache.iceberg.ChangelogOperation; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types; +import org.joda.time.Instant; +import org.junit.ClassRule; +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; + +/** Integration tests for {@link IncrementalChangelogSource}. */ +@RunWith(JUnit4.class) +public class IncrementalChangelogSourceTest { + private static final org.apache.iceberg.Schema CDC_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + required(1, "id", Types.LongType.get()), optional(2, "data", Types.StringType.get())), + ImmutableSet.of(1)); + + private static final org.apache.iceberg.Schema EVENT_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + required(1, "id", Types.LongType.get()), + optional(2, "data", Types.StringType.get()), + required(3, "event_time", Types.TimestampType.withoutZone())), + ImmutableSet.of(1)); + + @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); + @Rule public TestName testName = new TestName(); + @Rule public TestPipeline pipeline = TestPipeline.create(); + + @Test + public void boundedSnapshotRangeEmitsOnlyRequestedSnapshotsWithProjectedSchema() + throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + commitAppend(table, "s1.parquet", records(1L, "one")); + commitAppend(table, "s2.parquet", records(2L, "two")); + commitAppend(table, "s3.parquet", records(3L, "three")); + commitAppend(table, "s4.parquet", records(4L, "four")); + List snapshots = Lists.newArrayList(table.snapshots()); + + IcebergScanConfig scanConfig = + baseConfigBuilder(table, tableId) + .setKeepFields(ImmutableList.of("id")) + .setFromSnapshotInclusive(snapshots.get(1).snapshotId()) + .setToSnapshot(snapshots.get(2).snapshotId()) + .build(); + Schema projectedSchema = Schema.builder().addInt64Field("id").build(); + + PCollection rows = pipeline.apply(new IncrementalChangelogSource(scanConfig)); + + assertThat(rows.isBounded(), equalTo(IsBounded.BOUNDED)); + assertEquals(projectedSchema, rows.getSchema()); + PAssert.that(rows) + .containsInAnyOrder( + Row.withSchema(projectedSchema).addValue(2L).build(), + Row.withSchema(projectedSchema).addValue(3L).build()); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void metadataColumnsAreAppendedToProjectedRecord() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tablePropertiesV3()); + DataFile file1 = + commitAppend( + table, + testName.getMethodName() + "file1.parquet", + Arrays.asList( + record(1L, "one"), record(2L, "two"), record(3L, "three"), record(4L, "four"))); + + DataFile file2 = + warehouse.writeRecords( + testName.getMethodName() + "file2.parquet", + table.schema(), + PartitionSpec.unpartitioned(), + null, + ImmutableList.of(record(3L, "three_new"), record(4L, "four_new"))); + table.newOverwrite().deleteFile(file1).addFile(file2).commit(); + table.refresh(); + + List snapshots = Lists.newArrayList(table.snapshots()); + long snap1Id = snapshots.get(0).snapshotId(); + long snap1Seq = snapshots.get(0).sequenceNumber(); + long file1Seq = snapshots.get(0).sequenceNumber(); + long snap1FirstRowId = snapshots.get(0).firstRowId(); + + long snap2Id = snapshots.get(1).snapshotId(); + long snap2Seq = snapshots.get(1).sequenceNumber(); + long file2Seq = snapshots.get(1).sequenceNumber(); + long snap2FirstRowId = snapshots.get(1).firstRowId(); + + IcebergScanConfig scanConfig = + baseConfigBuilder(table, tableId) + .setKeepFields(ImmutableList.of("id")) + .setFromSnapshotInclusive(snapshots.get(0).snapshotId()) + .setToSnapshot(snapshots.get(1).snapshotId()) + .setMetadataColumns( + ImmutableList.of( + IcebergCdcMetadataColumns.CHANGE_TYPE, + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID, + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER, + IcebergCdcMetadataColumns.ROW_ID, + IcebergCdcMetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER)) + .build(); + Schema outputSchema = + Schema.builder() + .addInt64Field("id") + .addStringField(IcebergCdcMetadataColumns.CHANGE_TYPE) + .addInt64Field(IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID) + .addInt64Field(IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER) + .addNullableField(IcebergCdcMetadataColumns.ROW_ID, Schema.FieldType.INT64) + .addNullableField( + IcebergCdcMetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER, Schema.FieldType.INT64) + .build(); + + PCollection rows = pipeline.apply(new IncrementalChangelogSource(scanConfig)); + + assertEquals(outputSchema, rows.getSchema()); + PAssert.that(rows) + .containsInAnyOrder( + // snapshot 1: insert data file 1 + row( + 1L, + ChangelogOperation.INSERT, + snap1Id, + snap1Seq, + snap1FirstRowId, + file1Seq, + outputSchema), + row( + 2L, + ChangelogOperation.INSERT, + snap1Id, + snap1Seq, + snap1FirstRowId + 1, + file1Seq, + outputSchema), + row( + 3L, + ChangelogOperation.INSERT, + snap1Id, + snap1Seq, + snap1FirstRowId + 2, + file1Seq, + outputSchema), + row( + 4L, + ChangelogOperation.INSERT, + snap1Id, + snap1Seq, + snap1FirstRowId + 3, + file1Seq, + outputSchema), + // snapshot 2: delete data file 1 + row( + 1L, + ChangelogOperation.DELETE, + snap2Id, + snap2Seq, + snap1FirstRowId, + file1Seq, + outputSchema), + row( + 2L, + ChangelogOperation.DELETE, + snap2Id, + snap2Seq, + snap1FirstRowId + 1, + file1Seq, + outputSchema), + row( + 3L, + ChangelogOperation.UPDATE_BEFORE, + snap2Id, + snap2Seq, + snap1FirstRowId + 2, + file1Seq, + outputSchema), + row( + 4L, + ChangelogOperation.UPDATE_BEFORE, + snap2Id, + snap2Seq, + snap1FirstRowId + 3, + file1Seq, + outputSchema), + // snapshot 2: insert data file 2 + row( + 3L, + ChangelogOperation.UPDATE_AFTER, + snap2Id, + snap2Seq, + snap2FirstRowId, + file2Seq, + outputSchema), + row( + 4L, + ChangelogOperation.UPDATE_AFTER, + snap2Id, + snap2Seq, + snap2FirstRowId + 1, + file2Seq, + outputSchema)); + + pipeline.run().waitUntilFinish(); + } + + private Row row( + long id, + ChangelogOperation operation, + long snapshotId, + long snapshotSequence, + long rowId, + long lastUpdatedSequence, + Schema outputSchema) { + return Row.withSchema(outputSchema) + .addValues(id, operation.name(), snapshotId, snapshotSequence, rowId, lastUpdatedSequence) + .build(); + } + + @Test + public void streamingSnapshotRangeTerminatesWithoutDuplicates() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + commitAppend(table, "s1.parquet", records(1L, "one")); + commitAppend(table, "s2.parquet", records(2L, "two")); + commitAppend(table, "s3.parquet", records(3L, "three")); + + IcebergScanConfig scanConfig = + baseConfigBuilder(table, tableId) + .setStreaming(true) + .setStartingStrategy(StartingStrategy.EARLIEST) + .setToSnapshot(table.currentSnapshot().snapshotId()) + .build(); + Schema rowSchema = IcebergUtils.icebergSchemaToBeamSchema(CDC_SCHEMA); + + PCollection rows = pipeline.apply(new IncrementalChangelogSource(scanConfig)); + + assertThat(rows.isBounded(), equalTo(IsBounded.UNBOUNDED)); + PAssert.that(rows) + .containsInAnyOrder( + Row.withSchema(rowSchema).addValues(1L, "one").build(), + Row.withSchema(rowSchema).addValues(2L, "two").build(), + Row.withSchema(rowSchema).addValues(3L, "three").build()); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void overwriteUpdatePairsAreResolvedWithinSnapshotWindow() throws Exception { + TableIdentifier tableId = tableId(); + Table table = + warehouse.createTable( + tableId, + CDC_SCHEMA, + null, + ImmutableMap.of(TableProperties.FORMAT_VERSION, "2", TableProperties.SPLIT_SIZE, "1")); + DataFile oldFile = + commitAppend( + table, "old.parquet", ImmutableList.of(record(1L, "before"), record(2L, "same"))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(record(1L, "after"), record(2L, "same"))); + table.newOverwrite().deleteFile(oldFile).addFile(newFile).commit(); + table.refresh(); + + IcebergScanConfig scanConfig = + baseConfigBuilder(table, tableId) + .setFromSnapshotInclusive(table.currentSnapshot().snapshotId()) + .setToSnapshot(table.currentSnapshot().snapshotId()) + .build(); + + PCollection changes = + pipeline + .apply(new IncrementalChangelogSource(scanConfig)) + .apply("Format Changes", ParDo.of(new FormatValueKindAndRow())); + + PAssert.that(changes).containsInAnyOrder("UPDATE_BEFORE:1:before", "UPDATE_AFTER:1:after"); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void watermarkColumnRestampsProjectedRows() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, EVENT_SCHEMA, null, tableProperties()); + java.time.Instant firstEventInstant = + java.time.Instant.ofEpochMilli(System.currentTimeMillis() - 1_000L); + java.time.Instant secondEventInstant = firstEventInstant.plusMillis(5_000L); + LocalDateTime firstEvent = LocalDateTime.ofInstant(firstEventInstant, ZoneOffset.UTC); + LocalDateTime secondEvent = LocalDateTime.ofInstant(secondEventInstant, ZoneOffset.UTC); + commitAppend( + table, + "events.parquet", + ImmutableList.of(eventRecord(1L, "one", firstEvent), eventRecord(2L, "two", secondEvent))); + + IcebergScanConfig scanConfig = + baseConfigBuilder(table, tableId) + .setKeepFields(ImmutableList.of("id", "event_time")) + .setWatermarkColumn("event_time") + .setToSnapshot(table.currentSnapshot().snapshotId()) + .build(); + Schema projectedSchema = + Schema.builder() + .addInt64Field("id") + .addLogicalTypeField( + "event_time", org.apache.beam.sdk.schemas.logicaltypes.SqlTypes.DATETIME) + .build(); + + PCollection rows = pipeline.apply(new IncrementalChangelogSource(scanConfig)); + + assertEquals(projectedSchema, rows.getSchema()); + PAssert.that(rows.apply(Reify.timestamps())) + .containsInAnyOrder( + TimestampedValue.of( + Row.withSchema(projectedSchema).addValues(1L, firstEvent).build(), + new Instant(firstEventInstant.toEpochMilli())), + TimestampedValue.of( + Row.withSchema(projectedSchema).addValues(2L, secondEvent).build(), + new Instant(secondEventInstant.toEpochMilli()))); + + pipeline.run().waitUntilFinish(); + } + + /** + * 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 + * silently. Drive several consecutive snapshots through the shuffle path and assert that + * advancing past one snapshot never discards an earlier one's records. + */ + @Test + public void consecutiveSnapshotsThroughShuffleDropNoRecords() throws Exception { + assertNoRecordsDroppedAcrossSnapshots(false); + } + + /** Same, but driven by the per-snapshot watermark from {@link WatchForSnapshotsSdf}. */ + @Test + public void consecutiveSnapshotsThroughShuffleDropNoRecordsWhenStreaming() throws Exception { + assertNoRecordsDroppedAcrossSnapshots(true); + } + + private void assertNoRecordsDroppedAcrossSnapshots(boolean streaming) throws Exception { + TableIdentifier tableId = tableId(); + // SPLIT_SIZE=1 forces every bi-directional group onto the CoGroupByKey path rather than + // LocalResolveDoFn, so the windowing/lateness behaviour is what's under test. + Table table = + warehouse.createTable( + tableId, + CDC_SCHEMA, + null, + ImmutableMap.of(TableProperties.FORMAT_VERSION, "2", TableProperties.SPLIT_SIZE, "1")); + + DataFile v1 = commitAppend(table, "v1.parquet", ImmutableList.of(record(1L, "v1"))); + DataFile v2 = commitOverwrite(table, "v2.parquet", v1, record(1L, "v2")); + commitOverwrite(table, "v3.parquet", v2, record(1L, "v3")); + + IcebergScanConfig scanConfig = + baseConfigBuilder(table, tableId) + .setStreaming(streaming) + .setStartingStrategy(StartingStrategy.EARLIEST) + .setToSnapshot(table.currentSnapshot().snapshotId()) + .build(); + + PCollection changes = + pipeline + .apply(new IncrementalChangelogSource(scanConfig)) + .apply("Format Changes", ParDo.of(new FormatValueKindAndRow())); + + PAssert.that(changes) + .containsInAnyOrder( + "INSERT:1:v1", + "UPDATE_BEFORE:1:v1", + "UPDATE_AFTER:1:v2", + "UPDATE_BEFORE:1:v2", + "UPDATE_AFTER:1:v3"); + + pipeline.run().waitUntilFinish(); + } + + private DataFile commitOverwrite( + Table table, String fileName, DataFile replaced, Record replacement) throws IOException { + DataFile file = + warehouse.writeRecords( + testName.getMethodName() + "-" + fileName, + table.schema(), + ImmutableList.of(replacement)); + table.newOverwrite().deleteFile(replaced).addFile(file).commit(); + table.refresh(); + return file; + } + + private TableIdentifier tableId() { + return TableIdentifier.of("default", testName.getMethodName()); + } + + private IcebergScanConfig.Builder baseConfigBuilder(Table table, TableIdentifier tableId) { + return IcebergScanConfig.builder() + .setCatalogConfig( + IcebergCatalogConfig.builder() + .setCatalogName("name") + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build()) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(table.schema())) + .setUseCdc(true); + } + + private static Map tableProperties() { + return ImmutableMap.of(TableProperties.FORMAT_VERSION, "2"); + } + + private static Map tablePropertiesV3() { + return ImmutableMap.of(TableProperties.FORMAT_VERSION, "3"); + } + + private DataFile commitAppend(Table table, String fileName, List records) + throws IOException { + DataFile file = + warehouse.writeRecords(testName.getMethodName() + "-" + fileName, table.schema(), records); + table.newFastAppend().appendFile(file).commit(); + table.refresh(); + return file; + } + + private static List records(long id, String data) { + return ImmutableList.of(record(id, data)); + } + + private static Record record(long id, String data) { + return TestFixtures.createRecord(CDC_SCHEMA, ImmutableMap.of("id", id, "data", data)); + } + + private static Record eventRecord(long id, String data, LocalDateTime eventTime) { + return TestFixtures.createRecord( + EVENT_SCHEMA, ImmutableMap.of("id", id, "data", data, "event_time", eventTime)); + } + + private static final class FormatValueKindAndRow extends DoFn { + @ProcessElement + public void process( + @Element Row row, ValueKind valueKind, OutputReceiver outputReceiver) { + outputReceiver.output( + valueKind.name() + ":" + row.getInt64("id") + ":" + row.getString("data")); + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java new file mode 100644 index 000000000000..4d9e709b8d42 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java @@ -0,0 +1,222 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; +import static org.hamcrest.Matchers.empty; +import static org.junit.Assert.assertEquals; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.transforms.join.CoGbkResult; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.sdk.values.ValueInSingleWindow; +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.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.types.Types; +import org.joda.time.Instant; +import org.junit.ClassRule; +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; + +/** Unit tests for {@link ResolveChanges}. */ +@RunWith(JUnit4.class) +public class ResolveChangesTest { + private static final org.apache.iceberg.Schema SIMPLE_ICEBERG_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "data", Types.StringType.get())), + ImmutableSet.of(1)); + private static final Schema SIMPLE_BEAM_SCHEMA = + IcebergUtils.icebergSchemaToBeamSchema(SIMPLE_ICEBERG_SCHEMA); + private static final Schema PK_SCHEMA = Schema.builder().addInt32Field("id").build(); + + @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); + @Rule public final TestName testName = new TestName(); + + @Test + public void fullRowDuplicateDeleteInsertEmitsNothing() throws Exception { + Row duplicate = simpleRow(1, "duplicate"); + + List> output = + process( + SIMPLE_ICEBERG_SCHEMA, + pkRow(1), + Collections.singletonList(duplicate), + Collections.singletonList(duplicate), + new Instant(0L)); + + assertThat(output, empty()); + } + + @Test + public void updatePairAndExtraRowsPreserveKindsAndTimestamp() throws Exception { + Instant timestamp = new Instant(123L); + Row before = simpleRow(1, "before"); + Row after = simpleRow(1, "after"); + Row extraDelete = simpleRow(1, "deleted-only"); + + List> output = + process( + SIMPLE_ICEBERG_SCHEMA, + pkRow(1), + Arrays.asList(before, extraDelete), + Collections.singletonList(after), + timestamp); + + // Simulates an unusual case where two deletes share a PK. Which one pairs with the insert is + // arbitrary because Iceberg doesn't track lineage within a commit; the resolver orders both + // sides by nonPkHash so the choice is at least deterministic. These values are picked so that + // ordering matches the intuitive reading ("before" hashes below "deleted-only"). + assertThat( + output.stream().map(ResolveChangesTest::kindAndData).collect(Collectors.toList()), + contains("UPDATE_BEFORE:before", "UPDATE_AFTER:after", "DELETE:deleted-only")); + assertEquals( + Collections.nCopies(3, timestamp), + output.stream().map(ValueInSingleWindow::getTimestamp).collect(Collectors.toList())); + } + + @Test + public void duplicateDetectionUsesDeepEqualityForNestedValues() throws Exception { + org.apache.iceberg.Schema icebergSchema = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional( + 2, + "nested", + Types.StructType.of( + Types.NestedField.optional(3, "name", Types.StringType.get()))), + Types.NestedField.optional( + 4, "items", Types.ListType.ofOptional(5, Types.StringType.get())), + Types.NestedField.optional( + 6, + "attrs", + Types.MapType.ofOptional( + 7, 8, Types.StringType.get(), Types.IntegerType.get())), + Types.NestedField.optional(9, "payload", Types.BinaryType.get()), + Types.NestedField.optional(10, "nullable", Types.StringType.get())), + ImmutableSet.of(1)); + Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(icebergSchema); + Schema nestedSchema = beamSchema.getField("nested").getType().getRowSchema(); + Row delete = + Row.withSchema(beamSchema) + .addValues( + 1, + Row.withSchema(nestedSchema).addValue("same").build(), + ImmutableList.of("a", "b"), + ImmutableMap.of("x", 1), + new byte[] {1, 2, 3}, + null) + .build(); + Row insert = + Row.withSchema(beamSchema) + .addValues( + 1, + Row.withSchema(nestedSchema).addValue("same").build(), + ImmutableList.of("a", "b"), + ImmutableMap.of("x", 1), + new byte[] {1, 2, 3}, + null) + .build(); + + List> output = + process( + icebergSchema, + pkRow(1), + Collections.singletonList(delete), + Collections.singletonList(insert), + new Instant(0L)); + + assertThat(output, empty()); + } + + private List> process( + org.apache.iceberg.Schema icebergSchema, + Row pk, + List deletes, + List inserts, + Instant timestamp) + throws Exception { + CoGbkResult result = + CoGbkResult.of(ResolveChanges.DELETES, deletes).and(ResolveChanges.INSERTS, inserts); + try (DoFnTester, Row> tester = + DoFnTester.of(new ResolveChanges(scanConfig(icebergSchema)))) { + tester.processTimestampedElement( + TimestampedValue.of( + KV.of( + CdcRowDescriptor.builder() + .setCommitSnapshotId(123) + .setSnapshotSequenceNumber(456) + .setPrimaryKey(pk) + .build(), + result), + timestamp)); + return tester.getMutableOutput(tester.getMainOutputTag()); + } + } + + private IcebergScanConfig scanConfig(org.apache.iceberg.Schema icebergSchema) { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + IcebergCatalogConfig catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build(); + catalogConfig.catalog().createTable(tableId, icebergSchema); + return IcebergScanConfig.builder() + .setCatalogConfig(catalogConfig) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(icebergSchema)) + .setUseCdc(true) + .build(); + } + + private static Row simpleRow(int id, String data) { + return Row.withSchema(SIMPLE_BEAM_SCHEMA).addValues(id, data).build(); + } + + private static Row pkRow(int id) { + return Row.withSchema(PK_SCHEMA).addValue(id).build(); + } + + private static String kindAndData(ValueInSingleWindow value) { + ValueKind kind = value.getValueKind(); + return kind.name() + ":" + value.getValue().getString("data"); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/SnapshotWindowFnTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/SnapshotWindowFnTest.java new file mode 100644 index 000000000000..c3c081c1eed7 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/SnapshotWindowFnTest.java @@ -0,0 +1,93 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotEquals; + +import java.util.Collection; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link SnapshotWindowFn}. */ +@RunWith(JUnit4.class) +public class SnapshotWindowFnTest { + @Test + public void identicalTimestampsShareWindowAndAdjacentTimestampsDoNot() throws Exception { + SnapshotWindowFn fn = new SnapshotWindowFn(); + Instant timestamp = new Instant(1_000L); + + IntervalWindow first = onlyWindow(fn, timestamp); + IntervalWindow second = onlyWindow(fn, timestamp); + IntervalWindow adjacent = onlyWindow(fn, timestamp.plus(Duration.millis(1))); + + assertEquals(new IntervalWindow(timestamp, timestamp.plus(Duration.millis(1))), first); + assertEquals(first, second); + assertNotEquals(first, adjacent); + assertEquals( + new IntervalWindow(timestamp.plus(Duration.millis(1)), timestamp.plus(Duration.millis(2))), + adjacent); + } + + @Test + public void sideInputMappingStartsAtMainWindowMaxTimestamp() { + SnapshotWindowFn fn = new SnapshotWindowFn(); + IntervalWindow mainWindow = new IntervalWindow(new Instant(10L), new Instant(20L)); + + IntervalWindow sideInputWindow = fn.getDefaultWindowMappingFn().getSideInputWindow(mainWindow); + + assertEquals( + new IntervalWindow( + mainWindow.maxTimestamp(), mainWindow.maxTimestamp().plus(Duration.millis(1L))), + sideInputWindow); + } + + @SuppressWarnings("NonCanonicalType") + private static IntervalWindow onlyWindow(SnapshotWindowFn fn, Instant timestamp) + throws Exception { + Collection windows = + fn.assignWindows( + fn.new AssignContext() { + @Override + public Object element() { + return "element"; + } + + @Override + public Instant timestamp() { + return timestamp; + } + + @Override + public BoundedWindow window() { + return GlobalWindow.INSTANCE; + } + }); + assertThat( + windows, contains(new IntervalWindow(timestamp, timestamp.plus(Duration.millis(1L))))); + return windows.iterator().next(); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/WatchForSnapshotsSdfTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/WatchForSnapshotsSdfTest.java index c62a9d6fb4ed..02bd33c3bcec 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/WatchForSnapshotsSdfTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/WatchForSnapshotsSdfTest.java @@ -112,8 +112,11 @@ public void earliestStreamingRestrictionEmitsSnapshotsInSequenceOrder() throws E assertTrue(continuation.shouldResume()); assertEquals(Duration.millis(1L), continuation.resumeDelay()); assertThat(actualSnapshotIds, contains(expectedSnapshotIds)); + // The watermark lands just past the last snapshot's timestamp, so its 1ms window can fire as + // soon as its records drain. assertEquals( - Instant.ofEpochMilli(snapshots.get(snapshots.size() - 1).timestampMillis()), + Instant.ofEpochMilli(snapshots.get(snapshots.size() - 1).timestampMillis()) + .plus(Duration.millis(1L)), watermark.currentWatermark()); } @@ -189,12 +192,12 @@ public void streamingDefaultStartsAtLatestSnapshotAndEarliestStartsAtFirst() thr public void emptyTableReturnsResumeAndAdvancesIdleWatermark() { TableIdentifier tableId = tableId(); Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + Duration pollInterval = Duration.millis(25L); WatchForSnapshotsSdf sdf = new WatchForSnapshotsSdf( scanConfigBuilder(table, tableId) .setStreaming(true) - .setMaxSnapshotDiscoveryDelay(Duration.ZERO) - .setPollInterval(Duration.millis(25L)) + .setPollInterval(pollInterval) .build()); ManualWatermarkEstimator watermark = sdf.newWatermarkEstimator(sdf.initialWatermarkState()); @@ -205,12 +208,48 @@ public void emptyTableReturnsResumeAndAdvancesIdleWatermark() { sdf.process(sdf.newTracker(sdf.initialRestriction()), watermark, out); assertTrue(continuation.shouldResume()); - assertEquals(Duration.millis(25L), continuation.resumeDelay()); + assertEquals(pollInterval, continuation.resumeDelay()); assertThat(out.values, empty()); - assertThat(watermark.currentWatermark(), greaterThan(beforeProcess.minus(Duration.millis(1L)))); + // The idle bump advances the watermark to now() - pollInterval. + assertThat( + watermark.currentWatermark(), + greaterThan(beforeProcess.minus(pollInterval).minus(Duration.millis(1L)))); assertThat(watermark.currentWatermark(), lessThanOrEqualTo(Instant.now())); } + /** + * A snapshot discovered after the watermark has already advanced past its commit time (e.g. the + * idle bump ran ahead) must be emitted with a clamped timestamp rather than behind the watermark, + * where its records would be silently dropped as late data. + */ + @Test + public void lateDiscoveredSnapshotsAreClampedToTheWatermark() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + commitAppend(table, "s1.parquet", records("one", 1L)); + commitAppend(table, "s2.parquet", records("two", 2L)); + + WatchForSnapshotsSdf sdf = + new WatchForSnapshotsSdf( + scanConfigBuilder(table, tableId) + .setStreaming(true) + .setStartingStrategy(StartingStrategy.EARLIEST) + .setPollInterval(Duration.millis(1L)) + .build()); + // Seed a watermark ahead of both commit timestamps, as if the idle bump ran before discovery. + Instant seededWatermark = Instant.now().plus(Duration.standardHours(1)); + ManualWatermarkEstimator watermark = sdf.newWatermarkEstimator(seededWatermark); + CapturingOutputReceiver out = new CapturingOutputReceiver(); + + sdf.process(sdf.newTracker(sdf.initialRestriction()), watermark, out); + + // Neither snapshot is emitted behind the seeded watermark; each lands 1ms after the previous. + assertEquals(2, out.values.size()); + assertEquals(seededWatermark, out.values.get(0).getTimestamp()); + assertEquals(seededWatermark.plus(Duration.millis(1L)), out.values.get(1).getTimestamp()); + assertEquals(seededWatermark.plus(Duration.millis(2L)), watermark.currentWatermark()); + } + private TableIdentifier tableId() { return TableIdentifier.of("default", testName.getMethodName()); }