diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests.json b/.github/trigger_files/IO_Iceberg_Integration_Tests.json index b73af5e61a43..059704e8d120 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": 1 + "modification": 0 } diff --git a/sdks/java/io/iceberg/build.gradle b/sdks/java/io/iceberg/build.gradle index 8142c5f5b90b..f874df14ee45 100644 --- a/sdks/java/io/iceberg/build.gradle +++ b/sdks/java/io/iceberg/build.gradle @@ -50,6 +50,7 @@ dependencies { implementation library.java.slf4j_api implementation library.java.joda_time implementation "org.apache.parquet:parquet-column:$parquet_version" + implementation "org.apache.parquet:parquet-common:$parquet_version" implementation "org.apache.parquet:parquet-hadoop:$parquet_version" implementation "org.apache.parquet:parquet-common:$parquet_version" implementation project(":sdks:java:io:parquet") 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 31ff57a668bb..fcc40f972e8c 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 @@ -118,12 +118,19 @@ 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) { readRows = readRows.withPollInterval(Duration.standardSeconds(pollIntervalSeconds)); } + @Nullable Long maxDelay = configuration.getMaxSnapshotDiscoveryDelay(); + if (maxDelay != null) { + readRows = readRows.withMaxSnapshotDiscoveryDelay(Duration.standardSeconds(maxDelay)); + } PCollection output = input.getPipeline().apply(readRows); @@ -194,6 +201,32 @@ 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( + "Maximum expected snapshot discovery delay in seconds. While idle, the source may advance " + + "the watermark to now() minus this delay; snapshots discovered later with older commit " + + "timestamps may be treated as late by downstream windowing. Default: 600 seconds.") + abstract @Nullable Long getMaxSnapshotDiscoveryDelay(); + + @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); @@ -224,6 +257,14 @@ abstract static class Builder { abstract Builder setFilter(String filter); + abstract Builder setWatermarkColumn(String watermarkColumn); + + abstract Builder setWatermarkColumnTimeUnit(String timeUnit); + + abstract Builder setMaxSnapshotDiscoveryDelay(Long seconds); + + 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 04feea5037d1..0f26ece5aff5 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.schemas.Schema; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.values.PBegin; @@ -32,6 +33,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; @@ -544,6 +546,7 @@ public static ReadRows readRows(IcebergCatalogConfig catalogConfig) { return new AutoValue_IcebergIO_ReadRows.Builder() .setCatalogConfig(catalogConfig) .setUseCdc(false) + .setMetadataColumns(ImmutableList.of()) .build(); } @@ -580,6 +583,14 @@ public enum StartingStrategy { abstract @Nullable String getFilter(); + abstract @Nullable String getWatermarkColumn(); + + abstract @Nullable String getWatermarkColumnTimeUnit(); + + abstract @Nullable Duration getMaxSnapshotDiscoveryDelay(); + + abstract List getMetadataColumns(); + abstract Builder toBuilder(); @AutoValue.Builder @@ -610,6 +621,14 @@ abstract static class Builder { abstract Builder setFilter(@Nullable String filter); + abstract Builder setWatermarkColumn(@Nullable String watermarkColumn); + + abstract Builder setWatermarkColumnTimeUnit(@Nullable String timeUnit); + + abstract Builder setMaxSnapshotDiscoveryDelay(@Nullable Duration delay); + + abstract Builder setMetadataColumns(List metadataColumns); + abstract ReadRows build(); } @@ -661,6 +680,35 @@ 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(); + } + + public ReadRows withMaxSnapshotDiscoveryDelay(@Nullable Duration delay) { + return toBuilder().setMaxSnapshotDiscoveryDelay(delay).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 = @@ -685,12 +733,17 @@ public PCollection expand(PBegin input) { .setKeepFields(getKeep()) .setDropFields(getDrop()) .setFilterString(getFilter()) + .setWatermarkColumn(getWatermarkColumn()) + .setWatermarkColumnTimeUnit(getWatermarkColumnTimeUnit()) + .setMaxSnapshotDiscoveryDelay(getMaxSnapshotDiscoveryDelay()) + .setMetadataColumns(getMetadataColumns()) .build(); scanConfig.validate(table); PTransform> source = getUseCdc() - ? new IncrementalScanSource(scanConfig) + ? new IncrementalChangelogSource(scanConfig) + // ? new IncrementalScanSource(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 bdcf652a5dfc..6e1126ab4d74 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 @@ -19,20 +19,24 @@ import static org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergSchemaToBeamSchema; import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets.newHashSet; +import static org.apache.iceberg.types.Type.TypeID.LONG; +import static org.apache.iceberg.types.Type.TypeID.TIMESTAMP; import com.google.auto.value.AutoValue; import java.io.Serializable; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.Comparator; import java.util.HashSet; import java.util.LinkedHashSet; import java.util.List; import java.util.Set; +import java.util.concurrent.TimeUnit; import org.apache.beam.sdk.io.iceberg.IcebergIO.ReadRows.StartingStrategy; import org.apache.beam.sdk.io.iceberg.cdc.IcebergCdcMetadataColumns; import org.apache.beam.sdk.schemas.Schema; @@ -40,6 +44,7 @@ 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.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.MetadataColumns; import org.apache.iceberg.StructLike; import org.apache.iceberg.Table; import org.apache.iceberg.TableUtil; @@ -48,6 +53,7 @@ import org.apache.iceberg.expressions.Expression; import org.apache.iceberg.types.Comparators; import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types.NestedField; import org.apache.iceberg.util.SnapshotUtil; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; @@ -100,9 +106,9 @@ static org.apache.iceberg.Schema resolveSchema( @Nullable List keep, @Nullable List drop, @Nullable Set fieldsInFilter) { - ImmutableList.Builder selectedFieldsBuilder = ImmutableList.builder(); + Set selectedFields = new LinkedHashSet<>(); if (keep != null && !keep.isEmpty()) { - selectedFieldsBuilder.addAll(keep); + selectedFields.addAll(keep); } else if (drop != null && !drop.isEmpty()) { List paths = new ArrayList<>(TypeUtil.indexNameById(schema.asStruct()).values()); Collections.sort(paths); @@ -111,7 +117,7 @@ static org.apache.iceberg.Schema resolveSchema( boolean isParent = i + 1 < paths.size() && paths.get(i + 1).startsWith(path + "."); boolean isDrop = drop.stream().anyMatch(d -> path.equals(d) || path.startsWith(d + ".")); if (!isParent && !isDrop) { - selectedFieldsBuilder.add(path); + selectedFields.add(path); } } } else { @@ -122,9 +128,8 @@ static org.apache.iceberg.Schema resolveSchema( if (fieldsInFilter != null && !fieldsInFilter.isEmpty()) { fieldsInFilter.stream() .map(f -> schema.caseInsensitiveFindField(f).name()) - .forEach(selectedFieldsBuilder::add); + .forEach(selectedFields::add); } - ImmutableList selectedFields = selectedFieldsBuilder.build(); return selectedFields.isEmpty() ? schema : schema.select(selectedFields); } @@ -253,6 +258,12 @@ public Expression getFilter() { @Pure public abstract @Nullable List getDropFields(); + @Pure + public abstract @Nullable String getWatermarkColumn(); + + @Pure + public abstract @Nullable String getWatermarkColumnTimeUnit(); + @Pure public abstract @Nullable Duration getMaxSnapshotDiscoveryDelay(); @@ -282,6 +293,7 @@ public static Builder builder() { .setStartingStrategy(null) .setTag(null) .setBranch(null) + .setWatermarkColumn(null) .setMetadataColumns(ImmutableList.of()); } @@ -345,6 +357,10 @@ public Builder setTableIdentifier(String... names) { public abstract Builder setDropFields(@Nullable List fields); + public abstract Builder setWatermarkColumn(@Nullable String watermarkColumn); + + public abstract Builder setWatermarkColumnTimeUnit(@Nullable String timeUnit); + public abstract Builder setMaxSnapshotDiscoveryDelay(@Nullable Duration delay); public abstract Builder setMetadataColumns(List metadataColumns); @@ -355,6 +371,7 @@ public Builder setTableIdentifier(String... names) { @VisibleForTesting abstract Builder toBuilder(); + @SuppressWarnings("ReturnValueIgnored") void validate(Table table) { @Nullable List keep = getKeepFields(); @Nullable List drop = getDropFields(); @@ -366,16 +383,19 @@ void validate(Table table) { String param; if (keep != null) { param = "keep"; - fieldsSpecified = newHashSet(checkNotNull(keep)); + fieldsSpecified = newHashSet(checkArgumentNotNull(keep)); } else { // drop != null param = "drop"; - fieldsSpecified = newHashSet(checkNotNull(drop)); + fieldsSpecified = newHashSet(checkArgumentNotNull(drop)); } fieldsSpecified.removeIf(name -> table.schema().findField(name) != null); checkArgument( - fieldsSpecified.isEmpty(), - error(String.format("'%s' specifies unknown field(s): %s", param, fieldsSpecified))); + fieldsSpecified.isEmpty() + || fieldsSpecified.stream().allMatch(MetadataColumns::isMetadataColumn), + error("'%s' specifies unknown field(s): %s"), + param, + fieldsSpecified); } // TODO(#34168, ahmedabu98): fill these gaps for the existing batch source @@ -439,7 +459,6 @@ void validate(Table table) { checkArgument( getToTimestamp() == null || getToSnapshot() == null, error("only one of 'to_timestamp' or 'to_snapshot' can be set")); - @Nullable Long fromSnapshotId = ReadUtils.getFromSnapshotInclusive(table, this); @Nullable Long toSnapshotId = ReadUtils.getToSnapshot(table, this); if (fromSnapshotId != null) { @@ -462,11 +481,76 @@ void validate(Table table) { toSnapshotId); } + if (fromSnapshotId != null) { + checkArgumentNotNull( + table.snapshot(fromSnapshotId), + error("configured starting snapshot does not exist: '%s'"), + fromSnapshotId); + } + if (toSnapshotId != null) { + checkArgumentNotNull( + table.snapshot(toSnapshotId), + error("configured end snapshot does not exist: '%s'"), + toSnapshotId); + } + if (fromSnapshotId != null && toSnapshotId != null) { + checkArgument( + SnapshotUtil.isAncestorOf(table, toSnapshotId, fromSnapshotId), + error("fromSnapshot '%s' is not an ancestor of toSnapshot '%s'"), + fromSnapshotId, + toSnapshotId); + } + if (getPollInterval() != null) { checkArgument( Boolean.TRUE.equals(getStreaming()), error("'poll_interval_seconds' can only be set when streaming is true")); } + + @Nullable String watermarkColumn = getWatermarkColumn(); + if (watermarkColumn != null) { + checkArgument(getUseCdc(), error("'watermark_column' is only supported in CDC mode")); + NestedField field = table.schema().findField(watermarkColumn); + checkArgument( + field != null, error("'watermark_column' refers to unknown column: %s"), watermarkColumn); + checkArgument( + field.isRequired(), + error("'watermark_column' needs to be a non-nullable column: %s"), + watermarkColumn); + checkArgument( + field.type().typeId() == TIMESTAMP || field.type().typeId() == LONG, + error("'watermark_column' must be a timestamp-typed column, but '%s' has type %s"), + watermarkColumn, + field.type().typeId()); + checkArgumentNotNull( + getProjectedSchema().findField(watermarkColumn), + "'watermark_column' column should not be dropped."); + } + + @Nullable String watermarkColumnTimeUnit = getWatermarkColumnTimeUnit(); + if (watermarkColumnTimeUnit != null) { + checkArgument( + table + .schema() + .findField( + checkStateNotNull( + watermarkColumn, + "watermark_column_time_unit is configured without a specified watermark_column")) + .type() + .typeId() + == LONG, + error("watermark_column_time_unit is only applicable for LONG columns.")); + try { + TimeUnit.valueOf(watermarkColumnTimeUnit.toUpperCase()); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException( + error( + String.format( + "watermark_column_time_unit '%s' is invalid. Please choose one of: %s", + watermarkColumnTimeUnit, Arrays.toString(TimeUnit.values()))), + e); + } + } } private void validateMetadataColumns(Table table) { diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java index d0d24532ff39..2b7a8f956bae 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java @@ -46,6 +46,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; 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.iceberg.StructLike; import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.Record; import org.apache.iceberg.types.Type; @@ -469,120 +470,141 @@ private static Object getIcebergTimestampValue(Object beamValue, boolean shouldA } } + /** Converts a {@link StructLike} to a Beam {@link Row}. */ + public static Row structToRow(Schema schema, StructLike struct) { + checkState( + schema.getFieldCount() == struct.size(), + "Struct of size %s does not match expected schema size %s", + struct.size(), + schema.getFieldCount()); + Row.Builder rowBuilder = Row.withSchema(schema); + for (int i = 0; i < schema.getFieldCount(); i++) { + Schema.Field field = schema.getField(i); + @Nullable Object icebergValue = struct.get(i, Object.class); + addIcebergValue(rowBuilder, field, icebergValue); + } + return rowBuilder.build(); + } + /** Converts an Iceberg {@link Record} to a Beam {@link Row}. */ public static Row icebergRecordToBeamRow(Schema schema, Record record) { Row.Builder rowBuilder = Row.withSchema(schema); for (Schema.Field field : schema.getFields()) { - boolean isNullable = field.getType().getNullable(); @Nullable Object icebergValue = record.getField(field.getName()); - if (icebergValue == null) { - if (isNullable) { - rowBuilder.addValue(null); - continue; - } - throw new RuntimeException( - String.format("Received null value for required field '%s'.", field.getName())); + addIcebergValue(rowBuilder, field, icebergValue); + } + return rowBuilder.build(); + } + + private static void addIcebergValue( + Row.Builder rowBuilder, Schema.Field field, @Nullable Object icebergValue) { + boolean isNullable = field.getType().getNullable(); + if (icebergValue == null) { + if (isNullable) { + rowBuilder.addValue(null); + return; } - switch (field.getType().getTypeName()) { - case BYTE: - case INT16: - case INT32: - case INT64: - case DECIMAL: // Iceberg and Beam both use BigDecimal - case FLOAT: // Iceberg and Beam both use float - case DOUBLE: // Iceberg and Beam both use double - case STRING: // Iceberg and Beam both use String - case BOOLEAN: // Iceberg and Beam both use boolean - rowBuilder.addValue(icebergValue); - break; - case ARRAY: - checkState( - icebergValue instanceof List, - "Expected List type for field '%s' but received %s", - field.getName(), - icebergValue.getClass()); - List<@NonNull ?> beamList = (List<@NonNull ?>) icebergValue; - Schema.FieldType collectionType = - checkStateNotNull(field.getType().getCollectionElementType()); - // recurse on struct types - if (collectionType.getTypeName().isCompositeType()) { - Schema innerSchema = checkStateNotNull(collectionType.getRowSchema()); - beamList = - beamList.stream() - .map(v -> icebergRecordToBeamRow(innerSchema, (Record) v)) - .collect(Collectors.toList()); - } - rowBuilder.addValue(beamList); - break; - case ITERABLE: - checkState( - icebergValue instanceof Iterable, - "Expected Iterable type for field '%s' but received %s", - field.getName(), - icebergValue.getClass()); - Iterable<@NonNull ?> beamIterable = (Iterable<@NonNull ?>) icebergValue; - Schema.FieldType iterableCollectionType = - checkStateNotNull(field.getType().getCollectionElementType()); - // recurse on struct types - if (iterableCollectionType.getTypeName().isCompositeType()) { - Schema innerSchema = checkStateNotNull(iterableCollectionType.getRowSchema()); - ImmutableList.Builder builder = ImmutableList.builder(); - for (Record v : (Iterable<@NonNull Record>) icebergValue) { - builder.add(icebergRecordToBeamRow(innerSchema, v)); - } - beamIterable = builder.build(); + throw new RuntimeException( + String.format("Received null value for required field '%s'.", field.getName())); + } + switch (field.getType().getTypeName()) { + case BYTE: + case INT16: + case INT32: + case INT64: + case DECIMAL: // Iceberg and Beam both use BigDecimal + case FLOAT: // Iceberg and Beam both use float + case DOUBLE: // Iceberg and Beam both use double + case STRING: // Iceberg and Beam both use String + case BOOLEAN: // Iceberg and Beam both use boolean + rowBuilder.addValue(icebergValue); + break; + case ARRAY: + checkState( + icebergValue instanceof List, + "Expected List type for field '%s' but received %s", + field.getName(), + icebergValue.getClass()); + List<@NonNull ?> beamList = (List<@NonNull ?>) icebergValue; + Schema.FieldType collectionType = + checkStateNotNull(field.getType().getCollectionElementType()); + // recurse on struct types + if (collectionType.getTypeName().isCompositeType()) { + Schema innerSchema = checkStateNotNull(collectionType.getRowSchema()); + beamList = + beamList.stream() + .map(v -> icebergRecordToBeamRow(innerSchema, (Record) v)) + .collect(Collectors.toList()); + } + rowBuilder.addValue(beamList); + break; + case ITERABLE: + checkState( + icebergValue instanceof Iterable, + "Expected Iterable type for field '%s' but received %s", + field.getName(), + icebergValue.getClass()); + Iterable<@NonNull ?> beamIterable = (Iterable<@NonNull ?>) icebergValue; + Schema.FieldType iterableCollectionType = + checkStateNotNull(field.getType().getCollectionElementType()); + // recurse on struct types + if (iterableCollectionType.getTypeName().isCompositeType()) { + Schema innerSchema = checkStateNotNull(iterableCollectionType.getRowSchema()); + ImmutableList.Builder builder = ImmutableList.builder(); + for (Record v : (Iterable<@NonNull Record>) icebergValue) { + builder.add(icebergRecordToBeamRow(innerSchema, v)); } - rowBuilder.addValue(beamIterable); - break; - case MAP: - checkState( - icebergValue instanceof Map, - "Expected Map type for field '%s' but received %s", - field.getName(), - icebergValue.getClass()); - Map beamMap = (Map) icebergValue; - Schema.FieldType valueType = checkStateNotNull(field.getType().getMapValueType()); - // recurse on struct types - if (valueType.getTypeName().isCompositeType()) { - Schema innerSchema = checkStateNotNull(valueType.getRowSchema()); - ImmutableMap.Builder newMap = ImmutableMap.builder(); - for (Map.Entry entry : ((Map) icebergValue).entrySet()) { - Record rec = ((Record) entry.getValue()); - newMap.put( - checkStateNotNull(entry.getKey()), - icebergRecordToBeamRow(innerSchema, checkStateNotNull(rec))); - } - beamMap = newMap.build(); + beamIterable = builder.build(); + } + rowBuilder.addValue(beamIterable); + break; + case MAP: + checkState( + icebergValue instanceof Map, + "Expected Map type for field '%s' but received %s", + field.getName(), + icebergValue.getClass()); + Map beamMap = (Map) icebergValue; + Schema.FieldType valueType = checkStateNotNull(field.getType().getMapValueType()); + // recurse on struct types + if (valueType.getTypeName().isCompositeType()) { + Schema innerSchema = checkStateNotNull(valueType.getRowSchema()); + ImmutableMap.Builder newMap = ImmutableMap.builder(); + for (Map.Entry entry : ((Map) icebergValue).entrySet()) { + Record rec = ((Record) entry.getValue()); + newMap.put( + checkStateNotNull(entry.getKey()), + icebergRecordToBeamRow(innerSchema, checkStateNotNull(rec))); } - rowBuilder.addValue(beamMap); - break; - case DATETIME: - // Iceberg uses a long for micros. - // Beam DATETIME uses joda's DateTime, which only supports millis, - // so we do lose some precision here - rowBuilder.addValue(getBeamDateTimeValue(icebergValue)); - break; - case BYTES: - // Iceberg uses ByteBuffer; Beam uses byte[] - rowBuilder.addValue(((ByteBuffer) icebergValue).array()); - break; - case ROW: - Record nestedRecord = (Record) icebergValue; - Schema nestedSchema = - checkArgumentNotNull( - field.getType().getRowSchema(), - "Corrupted schema: Row type did not have associated nested schema."); - rowBuilder.addValue(icebergRecordToBeamRow(nestedSchema, nestedRecord)); - break; - case LOGICAL_TYPE: - rowBuilder.addValue(getLogicalTypeValue(icebergValue, field.getType())); - break; - default: - throw new UnsupportedOperationException( - "Unsupported Beam type: " + field.getType().getTypeName()); - } + beamMap = newMap.build(); + } + rowBuilder.addValue(beamMap); + break; + case DATETIME: + // Iceberg uses a long for micros. + // Beam DATETIME uses joda's DateTime, which only supports millis, + // so we do lose some precision here + rowBuilder.addValue(getBeamDateTimeValue(icebergValue)); + break; + case BYTES: + // Iceberg uses ByteBuffer; Beam uses byte[] + rowBuilder.addValue(((ByteBuffer) icebergValue).array()); + break; + case ROW: + Record nestedRecord = (Record) icebergValue; + Schema nestedSchema = + checkArgumentNotNull( + field.getType().getRowSchema(), + "Corrupted schema: Row type did not have associated nested schema."); + rowBuilder.addValue(icebergRecordToBeamRow(nestedSchema, nestedRecord)); + break; + case LOGICAL_TYPE: + rowBuilder.addValue(getLogicalTypeValue(icebergValue, field.getType())); + break; + default: + throw new UnsupportedOperationException( + "Unsupported Beam type: " + field.getType().getTypeName()); } - return rowBuilder.build(); } private static DateTime getBeamDateTimeValue(Object icebergValue) { 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 147e2adda1a7..fe3556a06e11 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. * diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java index daa0a2c73fb8..43a9472b9e22 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java @@ -67,6 +67,9 @@ * Read-side helpers specific to the CDC source. Keeps {@link ReadUtils} focused on the * general-purpose append-only read path; everything that takes a {@link SerializableChangelogTask}, * references {@link DeleteReader}, or implements the delete-pushdown row-group skipping lives here. + * + *

This class still delegates to {@link ReadUtils} for the low-level Parquet reader construction + * — the goal is decoupling, not duplication. */ public final class CdcReadUtils { private static final Logger LOG = LoggerFactory.getLogger(CdcReadUtils.class); 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 new file mode 100644 index 000000000000..0323c621e169 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolver.java @@ -0,0 +1,150 @@ +/* + * 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.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.BiConsumer; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.iceberg.data.Record; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Helper class to reconcile CDC rows. Used by {@link ResolveChanges} (with Beam {@link Row}s) and + * {@link LocalResolveDoFn} (with Iceberg {@link Record}s). + * + *

We determine the output ValueKind as follows: + * + *

    + *
  • (delete, insert) pairs become {@code UPDATE_BEFORE} + {@code UPDATE_AFTER} + *
  • singletons remain {@code DELETE} or {@code INSERT} + *
  • matching delete+insert with identical non-PK fields are considered Copy-on-Write side + * effects and are dropped + *
+ * + *

General implementation: + * + *

    + *
  1. Hash-index inserts by their non-PK field hash, for efficient Copy-on-Write detection. + *
  2. Skip matching (delete, insert) pairs with identical non-PK columns. A CoW operation deletes + * and rewrites the whole file (minus some records that are actually marked for deletion). + * Unchanged records are no-ops and should not be mistaken for updates. + *
  3. Walk the remaining deletes and inserts, emitting matched pairs as {@link + * ValueKind#UPDATE_BEFORE} / {@link ValueKind#UPDATE_AFTER}. + *
  4. Emit any unmatched extras as {@link ValueKind#DELETE} / {@link ValueKind#INSERT}. + *
+ */ +abstract class CdcResolver { + /** Hashes the non-PK fields of an element. Used as the index for O(n+m) CoW deduplication. */ + protected abstract int nonPkHash(T element); + + /** + * Returns true if two records (already known to share a PK) share identical non-PK fields. Called + * only when the two elements collide in the {@link #nonPkHash} index, so the implementation can + * stay simple (linear scan of non-PK fields). + */ + protected abstract boolean nonPkEquals(T delete, T insert); + + /** + * Resolves a Primary Key group of deletes and inserts. Caller provides {@code emit} which decides + * how to materialize each output. + * + *

Both input lists are inspected in their given order. + */ + final void resolve(List deletes, List inserts, BiConsumer emit) { + boolean hasDeletes = !deletes.isEmpty(); + boolean hasInserts = !inserts.isEmpty(); + + if (hasInserts && hasDeletes) { + // First, check if any (delete, insert) pairs are duplicates that should not be + // included in the output + boolean[] dupDeletes = new boolean[deletes.size()]; + boolean[] dupInserts = new boolean[inserts.size()]; + + // Map hash to insert-indices + Map> insertHashToIdx = new HashMap<>(); + for (int insertIdx = 0; insertIdx < inserts.size(); insertIdx++) { + int insertHash = nonPkHash(inserts.get(insertIdx)); + insertHashToIdx.computeIfAbsent(insertHash, k -> new ArrayList<>()).add(insertIdx); + } + for (int deleteIdx = 0; deleteIdx < deletes.size(); deleteIdx++) { + int deleteHash = nonPkHash(deletes.get(deleteIdx)); + @Nullable List candidates = insertHashToIdx.get(deleteHash); + if (candidates != null) { + // check if candidates are just duplicates (e.g. from CoW) + for (int idx = 0; idx < candidates.size(); idx++) { + int insertIdx = candidates.get(idx); + if (!dupInserts[insertIdx] + && nonPkEquals(deletes.get(deleteIdx), inserts.get(insertIdx))) { + // this (delete, insert) pair is a duplicate --> should be skipped + dupDeletes[deleteIdx] = true; + dupInserts[insertIdx] = true; + candidates.remove(idx); + break; + } + } + } + } + + // Emit matched pairs as UPDATE_BEFORE / UPDATE_AFTER. + int d = 0; + int i = 0; + while (d < deletes.size() && i < inserts.size()) { + // skip duplicates + while (d < deletes.size() && dupDeletes[d]) { + d++; + } + while (i < inserts.size() && dupInserts[i]) { + i++; + } + + if (d < deletes.size() && i < inserts.size()) { + emit.accept(ValueKind.UPDATE_BEFORE, deletes.get(d)); + emit.accept(ValueKind.UPDATE_AFTER, inserts.get(i)); + d++; + i++; + } + } + + // emit unmatched extras as DELETE / INSERT. + while (d < deletes.size()) { + if (!dupDeletes[d]) { + emit.accept(ValueKind.DELETE, deletes.get(d)); + } + d++; + } + while (i < inserts.size()) { + if (!dupInserts[i]) { + emit.accept(ValueKind.INSERT, inserts.get(i)); + } + i++; + } + } else if (hasInserts) { + for (T r : inserts) { + emit.accept(ValueKind.INSERT, r); + } + } else if (hasDeletes) { + for (T r : deletes) { + emit.accept(ValueKind.DELETE, r); + } + } + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcRowDescriptor.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcRowDescriptor.java new file mode 100644 index 000000000000..3bf3e8a619fd --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcRowDescriptor.java @@ -0,0 +1,89 @@ +/* + * 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 com.google.auto.value.AutoValue; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.SchemaCoder; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptor; + +/** + * Shuffle key for bidirectional CDC rows. + * + *

The primary key isolates rows for update resolution. The snapshot sequence number and commit + * snapshot id carry commit-sourced metadata through {@link ResolveChanges}, where they can be + * appended to final output rows if requested. + */ +@DefaultSchema(AutoValueSchema.class) +@AutoValue +public abstract class CdcRowDescriptor { + @SuppressWarnings("nullness") + public static SchemaCoder coder(Schema identifierSchema) { + Schema descriptorSchema = + Schema.builder() + .addInt64Field("snapshotSequenceNumber") + .addInt64Field("commitSnapshotId") + .addRowField("primaryKey", identifierSchema) + .build(); + + return SchemaCoder.of( + descriptorSchema, + TypeDescriptor.of(CdcRowDescriptor.class), + descriptor -> + Row.withSchema(descriptorSchema) + .addValues( + descriptor.getSnapshotSequenceNumber(), + descriptor.getCommitSnapshotId(), + descriptor.getPrimaryKey()) + .build(), + row -> + CdcRowDescriptor.builder() + .setSnapshotSequenceNumber(row.getInt64("snapshotSequenceNumber")) + .setCommitSnapshotId(row.getInt64("commitSnapshotId")) + .setPrimaryKey(row.getRow("primaryKey")) + .build()); + } + + public static Builder builder() { + return new AutoValue_CdcRowDescriptor.Builder(); + } + + @SchemaFieldNumber("0") + public abstract long getSnapshotSequenceNumber(); + + @SchemaFieldNumber("1") + public abstract long getCommitSnapshotId(); + + @SchemaFieldNumber("2") + public abstract Row getPrimaryKey(); + + @AutoValue.Builder + public abstract static class Builder { + abstract Builder setSnapshotSequenceNumber(long sequenceNumber); + + abstract Builder setCommitSnapshotId(long snapshotId); + + abstract Builder setPrimaryKey(Row primaryKey); + + abstract CdcRowDescriptor build(); + } +} 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..ae3214bc5501 --- /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.apache.iceberg.catalog.TableIdentifier; +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()); + 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(TableIdentifier.parse(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 new file mode 100644 index 000000000000..f353678f2a59 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java @@ -0,0 +1,241 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergSchemaToBeamSchema; +import static org.apache.beam.sdk.io.iceberg.cdc.SerializableChangelogTask.Type.ADDED_ROWS; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.TableCache; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.join.CoGroupByKey; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.types.Comparators; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeMap; +import org.apache.iceberg.util.StructLikeUtil; +import org.apache.iceberg.util.StructProjection; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Resolves a small bi-directional changelog group entirely in memory. This is the equivalent of + * {@link ReadFromChangelogs} + {@link CoGroupByKey} + {@link ResolveChanges}. + * + *

All tasks in a changelog group belong to the same Iceberg {@link Snapshot}. The upstream + * {@link ChangelogScanner} routes here only when the total size of the bi-directional group fits + * within {@link TableProperties#SPLIT_SIZE}. + * + *

The incoming batch's overlap region has already been computed in the scanning phase by {@link + * ChangelogScanner}. In this DoFn, we just process each task and route records: + * + *

    + *
  • Records whose PK falls outside the overlap range cannot have an opposing-side match, + * so they are emitted directly with {@code INSERT} or {@code DELETE} kind. + *
  • Records whose PK falls inside the overlap range are stashed in a {@link + * StructLikeMap} keyed by PK, then resolved by {@link CdcResolver}. + *
+ */ +class LocalResolveDoFn extends DoFn>, Row> { + private final IcebergScanConfig scanConfig; + private final org.apache.beam.sdk.schemas.Schema projectedBeamSchema; + private final org.apache.beam.sdk.schemas.Schema outputBeamSchema; + + private transient @MonotonicNonNull OverlapRange overlap; + private transient @MonotonicNonNull List nonPkFields; + private transient @MonotonicNonNull StructProjection projector; + + LocalResolveDoFn(IcebergScanConfig scanConfig) { + this.scanConfig = scanConfig; + this.projectedBeamSchema = + CdcOutputUtils.readBeamSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getProjectedSchema()); + this.outputBeamSchema = + CdcOutputUtils.outputSchema( + scanConfig, icebergSchemaToBeamSchema(scanConfig.getProjectedSchema())); + } + + @Setup + public void setup() { + Schema tableSchema = + TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()).schema(); + Schema fullReadSchema = + CdcOutputUtils.readSchemaWithRowMetadata(scanConfig.getMetadataColumns(), tableSchema); + this.overlap = OverlapRange.forScanConfig(scanConfig); + Set pkFieldNames = new HashSet<>(overlap.recordIdSchema().identifierFieldNames()); + // The dedup logic only inspects non-PK fields, so precompute them once. + List nonPk = new ArrayList<>(); + for (Types.NestedField f : tableSchema.columns()) { + if (!pkFieldNames.contains(f.name())) { + nonPk.add(f); + } + } + this.nonPkFields = nonPk; + this.projector = + StructProjection.create( + fullReadSchema, + CdcOutputUtils.readSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getProjectedSchema())); + } + + @ProcessElement + public void process( + @Element KV> element, + OutputReceiver out) + throws IOException { + ChangelogDescriptor descriptor = element.getKey(); + Table table = TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()); + OverlapRange ovl = checkStateNotNull(overlap); + + // {PK: (inserts | deletes)} for in-overlap records that need resolution. + // Records outside the overlap are emitted directly + StructLikeMap pkGroups = StructLikeMap.create(ovl.recordIdSchema().asStruct()); + + @Nullable StructLike overlapLower = ovl.toStructLike(descriptor.getOverlapLower()); + @Nullable StructLike overlapUpper = ovl.toStructLike(descriptor.getOverlapUpper()); + for (SerializableChangelogTask task : element.getValue()) { + readAndRoute(descriptor, task, table, overlapLower, overlapUpper, pkGroups, out); + } + + resolveAndEmit(descriptor, pkGroups, table.schema(), out); + } + + /** + * Processes a {@link SerializableChangelogTask} and routes each record. + * + *
    + *
  • Out of overlap: emit directly + *
  • Inside overlap: stash in {@code pkGroups} to resolve in {@link #resolveAndEmit} + *
+ */ + private void readAndRoute( + ChangelogDescriptor descriptor, + SerializableChangelogTask task, + Table table, + @Nullable StructLike overlapLower, + @Nullable StructLike overlapUpper, + StructLikeMap pkGroups, + OutputReceiver out) + throws IOException { + OverlapRange ovl = checkStateNotNull(overlap); + boolean isInsert = task.getType() == ADDED_ROWS; + try (CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig, false)) { + for (Record rec : records) { + if (ovl.contains(rec, overlapLower, overlapUpper)) { // needs resolution + StructLike pk = StructLikeUtil.copy(ovl.recordIdProjection()); + PkGroup group = pkGroups.computeIfAbsent(pk, k -> new PkGroup()); + if (isInsert) { + group.inserts.add(rec); + } else { + group.deletes.add(rec); + } + } else { // safe to emit directly + emit(descriptor, rec, isInsert ? ValueKind.INSERT : ValueKind.DELETE, out); + } + } + } + } + + /** Resolves each PK group using {@link CdcResolver}. */ + private void resolveAndEmit( + ChangelogDescriptor descriptor, + StructLikeMap pkGroups, + Schema fullSchema, + OutputReceiver out) { + CdcResolver resolver = new RecordResolver(checkStateNotNull(nonPkFields), fullSchema); + for (PkGroup group : pkGroups.values()) { + resolver.resolve( + group.deletes, + group.inserts, + (kind, rec) -> { + emit(descriptor, rec, kind, out); + }); + } + } + + /** Resolver specialization that hashes Iceberg Record non-PK fields. */ + private static final class RecordResolver extends CdcResolver { + private final List nonPkFields; + private final Comparator nonPkComparator; + private final StructProjection left; + private final StructProjection right; + + RecordResolver(List nonPkFields, Schema recSchema) { + this.nonPkFields = nonPkFields; + Set nonPkFieldIds = + nonPkFields.stream().map(Types.NestedField::fieldId).collect(Collectors.toSet()); + this.left = StructProjection.create(recSchema, nonPkFieldIds); + this.right = StructProjection.create(recSchema, nonPkFieldIds); + this.nonPkComparator = + Comparators.forType(TypeUtil.select(recSchema, nonPkFieldIds).asStruct()); + } + + @Override + protected int nonPkHash(Record rec) { + int hash = 1; + for (Types.NestedField field : nonPkFields) { + hash = 31 * hash + Objects.hashCode(rec.getField(field.name())); + } + return hash; + } + + @Override + protected boolean nonPkEquals(Record delete, Record insert) { + return nonPkComparator.compare(left.wrap(delete), right.wrap(insert)) == 0; + } + } + + /** Prune to get the final projected record then output as a Beam Row. */ + private void emit( + ChangelogDescriptor descriptor, Record rec, ValueKind kind, OutputReceiver out) { + StructLike projected = checkStateNotNull(projector).wrap(rec); + Row record = IcebergUtils.structToRow(projectedBeamSchema, projected); + out.builder( + CdcOutputUtils.outputRow( + scanConfig.getMetadataColumns(), outputBeamSchema, descriptor, kind, record)) + .setValueKind(kind) + .output(); + } + + /** Two parallel lists of inserts/deletes that share a primary key. */ + private static final class PkGroup { + final List inserts = new ArrayList<>(); + final List deletes = new ArrayList<>(); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRange.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRange.java new file mode 100644 index 000000000000..f291b4c8c0ad --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRange.java @@ -0,0 +1,93 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.util.Comparator; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.TableCache; +import org.apache.beam.sdk.values.Row; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.util.StructProjection; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Primary-key-projection and overlap-range comparison helper. + * + *

Used by {@link LocalResolveDoFn} and {@link ReadFromChangelogs} to decide whether a record's + * PK falls within an overlap of two opposing tasks. If so, the record needs to be compared with + * others to determine if it is part of an update pair. + */ +final class OverlapRange { + private final Schema recordIdSchema; + private final StructProjection recordIdProjection; + private final Comparator idComp; + + private OverlapRange( + Schema recordIdSchema, StructProjection recordIdProjection, Comparator idComp) { + this.recordIdSchema = recordIdSchema; + this.recordIdProjection = recordIdProjection; + this.idComp = idComp; + } + + static OverlapRange forScanConfig(IcebergScanConfig scanConfig) { + Schema fullSchema = + TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()).schema(); + StructProjection projection = StructProjection.create(fullSchema, scanConfig.recordIdSchema()); + return new OverlapRange( + scanConfig.recordIdSchema(), projection, scanConfig.recordIdComparator()); + } + + StructProjection recordIdProjection() { + return recordIdProjection; + } + + Schema recordIdSchema() { + return recordIdSchema; + } + + /** Converts a Beam Row (overlap bound) back to an Iceberg {@link StructLike}. */ + @Nullable + StructLike toStructLike(@Nullable Row beamBound) { + if (beamBound == null) { + return null; + } + return IcebergUtils.beamRowToIcebergRecord(recordIdSchema, beamBound); + } + + /** + * Wraps the record to project its Primary Key, then checks if the PK within the overlap {@code + * [lower, upper]} (inclusive). Can be paired with a subsequent {@link #recordIdProjection()} call + * to fetch the PK value. + * + *

If either bound is null, we conservatively assume it falls within the overlap. + */ + boolean contains(Record rec, @Nullable StructLike lower, @Nullable StructLike upper) { + checkStateNotNull(recordIdProjection).wrap(rec); + + if (lower == null || upper == null) { + return true; + } + return idComp.compare(recordIdProjection, lower) >= 0 + && idComp.compare(recordIdProjection, upper) <= 0; + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogs.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogs.java new file mode 100644 index 000000000000..2ff7ff10dbc5 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogs.java @@ -0,0 +1,493 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergRecordToBeamRow; +import static org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergSchemaToBeamSchema; +import static org.apache.beam.sdk.io.iceberg.IcebergUtils.structToRow; +import static org.apache.beam.sdk.io.iceberg.cdc.ChangelogScanner.LARGE_BIDIRECTIONAL_TASKS; +import static org.apache.beam.sdk.io.iceberg.cdc.ChangelogScanner.UNIDIRECTIONAL_TASKS; +import static org.apache.beam.sdk.io.iceberg.cdc.SerializableChangelogTask.Type.ADDED_ROWS; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.SerializableDeleteFile; +import org.apache.beam.sdk.io.iceberg.TableCache; +import org.apache.beam.sdk.io.range.OffsetRange; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.SchemaCoder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Redistribute; +import org.apache.beam.sdk.transforms.join.CoGroupByKey; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.PInput; +import org.apache.beam.sdk.values.POutput; +import org.apache.beam.sdk.values.PValue; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.ChangelogScanTask; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.util.StructProjection; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A {@link PTransform} that processes batches of {@link ChangelogScanTask}s and routes them + * accordingly: + * + *

    + *
  • Records from Uni-directional batches are directly emitted, as INSERT or DELETE kind + *
  • Records from Bi-directional batches are compared against the Primary Key overlap range: + *
      + *
    • if outside the overlap, emit directly as INSERT or DELETE kind + *
    • if inside the overlap, key by (snapshot seq#, pk) and route to downstream {@link + * CoGroupByKey} and final resolution by {@link ResolveChanges} + *
    + *
+ * + *

We first key bi-directional rows by (snapshot sequence number, primary key) before sending to + * {@link CoGroupByKey} to ensure they stay isolated from other PKs or snapshots. Inserts are routed + * to + * + *

A {@link ChangelogScanTask} comes in three types: + * + *

    + *
  1. AddedRowsScanTask: Indicates records have been inserted by a new DataFile. + *
  2. DeletedRowsScanTask: Indicates records have been deleted using a DeleteFile. + *
  3. DeletedDataFileScanTask: Indicates a whole DataFile has been deleted. + *
+ * + *

Each of these types need to be processed differently. More details in {@link + * CdcReadUtils#changelogRecordsForTask}. + * + *

CDC metadata has two entry points in this transform. Row metadata columns are requested from + * the Iceberg reader by {@link CdcReadUtils} and travel inside intermediate rows until final output + * assembly. Snapshot metadata columns come from the {@link ChangelogDescriptor} / {@link + * CdcRowDescriptor} carried with each task or shuffled row, and {@code _change_type} comes from the + * emitted change kind. Final user-visible rows are assembled by {@link CdcOutputUtils#outputRow}, + * which appends all requested metadata as top-level columns in the configured order. + */ +public class ReadFromChangelogs extends PTransform { + private static final Counter numAddedRowsScanTasksCompleted = + Metrics.counter(ReadFromChangelogs.class, "numAddedRowsScanTasksCompleted"); + private static final Counter numDeletedRowsScanTasksCompleted = + Metrics.counter(ReadFromChangelogs.class, "numDeletedRowsScanTasksCompleted"); + private static final Counter numDeletedDataFileScanTasksCompleted = + Metrics.counter(ReadFromChangelogs.class, "numDeletedDataFileScanTasksCompleted"); + + private static final TupleTag UNIDIRECTIONAL_ROWS = new TupleTag<>(); + private static final TupleTag> BIDIRECTIONAL_INSERTS = new TupleTag<>(); + private static final TupleTag> BIDIRECTIONAL_DELETES = new TupleTag<>(); + + private final IcebergScanConfig scanConfig; + + ReadFromChangelogs(IcebergScanConfig scanConfig) { + this.scanConfig = scanConfig; + } + + @Override + public org.apache.beam.sdk.io.iceberg.cdc.ReadFromChangelogs.Output expand( + PCollectionTuple input) { + Schema fullRowSchema = + CdcOutputUtils.readBeamSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getSchema()); + Schema projectedRowSchema = + IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()); + Schema outputRowSchema = CdcOutputUtils.outputSchema(scanConfig, projectedRowSchema); + + // === UNIDIRECTIONAL tasks === + // (i.e. only deletes, or only inserts) + // take the fast approach of just reading and emitting CDC records + PCollection uniDirectionalRows = + input + .get(UNIDIRECTIONAL_TASKS) + .apply("Redistribute Uni-Directional Changes", Redistribute.arbitrarily()) + .apply( + "Read Uni-Directional Changes", + ParDo.of(ReadDoFn.unidirectional(scanConfig)) + .withOutputTags(UNIDIRECTIONAL_ROWS, TupleTagList.empty())) + .get(UNIDIRECTIONAL_ROWS) + .setRowSchema(outputRowSchema); + + // === BIDIRECTIONAL tasks === + // (i.e. a task group containing a mix of deletes and inserts) + // read and route records according to their PK (see class java doc) + PCollectionTuple biDirectionalRows = + input + .get(LARGE_BIDIRECTIONAL_TASKS) + .apply("Redistribute Large Bi-Directional Changes", Redistribute.arbitrarily()) + .apply( + "Read Bi-Directional Changes", + ParDo.of(ReadDoFn.bidirectional(scanConfig)) + .withOutputTags( + BIDIRECTIONAL_INSERTS, + TupleTagList.of(BIDIRECTIONAL_DELETES).and(UNIDIRECTIONAL_ROWS))); + // Collect pruned (non-overlapping) rows from bi-directional reader + PCollection nonOverlappingRowsFromBiDirTasks = + biDirectionalRows.get(UNIDIRECTIONAL_ROWS).setRowSchema(outputRowSchema); + + // Flatten uni-directional rows from both sources + PCollection allUniDirectionalRows = + PCollectionList.of(uniDirectionalRows) + .and(nonOverlappingRowsFromBiDirTasks) + .apply("Flatten Uni-Directional Rows", Flatten.pCollections()); + + // Reify to preserve each record's timestamp (CoGBK overwrites timestamps with the window's + // end-of-window) + // Note: element timestamps are snapshot commit timestamp + KvCoder keyedOutputCoder = + KvCoder.of( + CdcRowDescriptor.coder(scanConfig.rowIdBeamSchema()), SchemaCoder.of(fullRowSchema)); + PCollection> keyedInsertsWithTimestamps = + biDirectionalRows.get(BIDIRECTIONAL_INSERTS).setCoder(keyedOutputCoder); + PCollection> keyedDeletesWithTimestamps = + biDirectionalRows.get(BIDIRECTIONAL_DELETES).setCoder(keyedOutputCoder); + + return new org.apache.beam.sdk.io.iceberg.cdc.ReadFromChangelogs.Output( + input.getPipeline(), + allUniDirectionalRows, + keyedInsertsWithTimestamps, + keyedDeletesWithTimestamps); + } + + public static class Output implements POutput { + private final Pipeline pipeline; + private final PCollection uniDirectionalRows; + private final PCollection> biDirectionalInserts; + private final PCollection> biDirectionalDeletes; + + Output( + Pipeline p, + PCollection uniDirectionalRows, + PCollection> biDirectionalInserts, + PCollection> biDirectionalDeletes) { + this.pipeline = p; + this.uniDirectionalRows = uniDirectionalRows; + this.biDirectionalInserts = biDirectionalInserts; + this.biDirectionalDeletes = biDirectionalDeletes; + } + + PCollection uniDirectionalRows() { + return uniDirectionalRows; + } + + PCollection> biDirectionalInserts() { + return biDirectionalInserts; + } + + PCollection> biDirectionalDeletes() { + return biDirectionalDeletes; + } + + @Override + public Pipeline getPipeline() { + return pipeline; + } + + @Override + public Map, PValue> expand() { + return ImmutableMap.of( + UNIDIRECTIONAL_ROWS, + uniDirectionalRows, + BIDIRECTIONAL_INSERTS, + biDirectionalInserts, + BIDIRECTIONAL_DELETES, + biDirectionalDeletes); + } + + @Override + public void finishSpecifyingOutput( + String transformName, PInput input, PTransform transform) {} + } + + @DoFn.BoundedPerElement + private static class ReadDoFn + extends DoFn>, OutT> { + private final IcebergScanConfig scanConfig; + private final boolean keyedOutput; + private final Schema projectedBeamRowSchema; + private final Schema outputBeamRowSchema; + private final Schema fullBeamRowSchema; + private transient @MonotonicNonNull OverlapRange overlap; + private transient @MonotonicNonNull StructProjection outputProjector; + private transient @MonotonicNonNull StructProjection pkProjector; + + /** Used for uni-directional changes. Records are output immediately as-is. */ + static ReadDoFn unidirectional(IcebergScanConfig scanConfig) { + return new ReadDoFn<>(scanConfig, false); + } + + /** + * Used for bi-directional changes. Records are keyed by (snapshot sequence number, primary key) + * and sent to a CoGBK. + */ + static ReadDoFn> bidirectional(IcebergScanConfig scanConfig) { + return new ReadDoFn<>(scanConfig, true); + } + + private ReadDoFn(IcebergScanConfig scanConfig, boolean keyedOutput) { + this.scanConfig = scanConfig; + this.keyedOutput = keyedOutput; + + this.projectedBeamRowSchema = + CdcOutputUtils.readBeamSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getProjectedSchema()); + this.outputBeamRowSchema = + CdcOutputUtils.outputSchema( + scanConfig, icebergSchemaToBeamSchema(scanConfig.getProjectedSchema())); + this.fullBeamRowSchema = + CdcOutputUtils.readBeamSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getSchema()); + } + + @Setup + public void setup() { + this.overlap = OverlapRange.forScanConfig(scanConfig); + } + + @ProcessElement + public void process( + @Element KV> element, + RestrictionTracker tracker, + MultiOutputReceiver out) + throws IOException { + Table table = TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()); + + List tasks = element.getValue(); + ChangelogDescriptor descriptor = element.getKey(); + @Nullable Row overlapLower = descriptor.getOverlapLower(); + @Nullable Row overlapUpper = descriptor.getOverlapUpper(); + + for (long l = tracker.currentRestriction().getFrom(); + l < tracker.currentRestriction().getTo(); + l++) { + if (!tracker.tryClaim(l)) { + return; + } + + SerializableChangelogTask task = tasks.get((int) l); + processTaskRecords(descriptor, task, overlapLower, overlapUpper, table, out); + } + } + + /** + * Processes a ChangelogScanTask and routes records accordingly: + * + *

If this DoFn is configured with {@link #unidirectional}, we simply read records and output + * directly to {@link #UNIDIRECTIONAL_ROWS}. + * + *

If this DoFn is configured with {@link #bidirectional}, we compare against the Primary Key + * overlap range. If within the overlap, we key by (snapshotId, PK) and out to either {@link + * #BIDIRECTIONAL_INSERTS} or {@link #BIDIRECTIONAL_DELETES}. Otherwise (not in overlap), we + * output the record directly to {@link #UNIDIRECTIONAL_ROWS}. + */ + private void processTaskRecords( + ChangelogDescriptor descriptor, + SerializableChangelogTask task, + @Nullable Row overlapLowerRow, + @Nullable Row overlapUpperRow, + Table table, + MultiOutputReceiver outputReceiver) + throws IOException { + OverlapRange ovl = checkStateNotNull(overlap); + @Nullable StructLike overlapLower = ovl.toStructLike(overlapLowerRow); + @Nullable StructLike overlapUpper = ovl.toStructLike(overlapUpperRow); + + boolean isInsert = task.getType() == ADDED_ROWS; + TupleTag> taggedOutput = + isInsert ? BIDIRECTIONAL_INSERTS : BIDIRECTIONAL_DELETES; + ValueKind kind = isInsert ? ValueKind.INSERT : ValueKind.DELETE; + long commitSnapshotId = descriptor.getCommitSnapshotId(); + long commitSnapshotSequenceNumber = descriptor.getSnapshotSequenceNumber(); + + Schema readSchema = keyedOutput ? fullBeamRowSchema : projectedBeamRowSchema; + try (CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig, !keyedOutput)) { + for (Record rec : records) { + // uni-directional -- just output records (they are already projected by read pushdown) + if (!keyedOutput) { + Row row = icebergRecordToBeamRow(projectedBeamRowSchema, rec); + outputReceiver + .get(UNIDIRECTIONAL_ROWS) + .builder( + CdcOutputUtils.outputRow( + scanConfig.getMetadataColumns(), + outputBeamRowSchema, + descriptor, + kind, + row)) + .setValueKind(kind) + .output(); + continue; + } + + // bi-directional -- compare overlap + if (ovl.contains(rec, overlapLower, overlapUpper)) { + // inside overlap -- read full row and output KV + Row row = icebergRecordToBeamRow(readSchema, rec); + Row pk = structToRow(scanConfig.rowIdBeamSchema(), pkProjector().wrap(rec)); + outputReceiver + .get(taggedOutput) + .builder( + KV.of( + CdcRowDescriptor.builder() + .setCommitSnapshotId(commitSnapshotId) + .setSnapshotSequenceNumber(commitSnapshotSequenceNumber) + .setPrimaryKey(pk) + .build(), + row)) + .setValueKind(kind) + .output(); + + } else { + // outside overlap -- get projected record and output + StructLike projected = outputProjector().wrap(rec); + Row row = structToRow(projectedBeamRowSchema, projected); + outputReceiver + .get(UNIDIRECTIONAL_ROWS) + .builder( + CdcOutputUtils.outputRow( + scanConfig.getMetadataColumns(), + outputBeamRowSchema, + descriptor, + kind, + row)) + .setValueKind(kind) + .output(); + } + } + } + + trackMetrics(task.getType()); + } + + private StructProjection outputProjector() { + if (outputProjector == null) { + outputProjector = + StructProjection.create( + CdcOutputUtils.readSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), + TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()) + .schema()), + CdcOutputUtils.readSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getProjectedSchema())); + } + return outputProjector; + } + + private StructProjection pkProjector() { + if (pkProjector == null) { + pkProjector = + StructProjection.create( + TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()) + .schema(), + scanConfig.recordIdSchema()); + } + return pkProjector; + } + + private void trackMetrics(SerializableChangelogTask.Type type) { + switch (type) { + case ADDED_ROWS: + numAddedRowsScanTasksCompleted.inc(); + break; + case DELETED_ROWS: + numDeletedRowsScanTasksCompleted.inc(); + break; + case DELETED_FILE: + numDeletedDataFileScanTasksCompleted.inc(); + break; + } + } + + private String getKind(SerializableChangelogTask.Type taskType) { + switch (taskType) { + case ADDED_ROWS: + return "INSERT"; + case DELETED_ROWS: + return "DELETE"; + case DELETED_FILE: + default: + return "DELETE-DF"; + } + } + + private static final int COMPRESSED_TO_DECODED_BYTES_ESTIMATE = 4; + + @GetSize + public double getSize( + @Element KV> element, + @Restriction OffsetRange restriction) { + // TODO(ahmedabu98): can we make this estimate more accurate? + long size = 0; + + for (long l = restriction.getFrom(); l < restriction.getTo(); l++) { + SerializableChangelogTask task = element.getValue().get((int) l); + size += task.getDataFile().getFileSizeInBytes() * COMPRESSED_TO_DECODED_BYTES_ESTIMATE; + size += + task.getAddedDeletes().stream() + .mapToLong(SerializableDeleteFile::getFileSizeInBytes) + .sum() + * COMPRESSED_TO_DECODED_BYTES_ESTIMATE; + size += + task.getExistingDeletes().stream() + .mapToLong(SerializableDeleteFile::getFileSizeInBytes) + .sum() + * COMPRESSED_TO_DECODED_BYTES_ESTIMATE; + } + + return size; + } + + @GetInitialRestriction + public OffsetRange getInitialRange( + @Element KV> element) { + return new OffsetRange(0, element.getValue().size()); + } + + @SplitRestriction + public void splitRestriction( + @Restriction OffsetRange restriction, OutputReceiver out) { + // Split into individual tasks for maximum initial parallelism + for (long i = restriction.getFrom(); i < restriction.getTo(); i++) { + out.output(new OffsetRange(i, i + 1)); + } + } + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChanges.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChanges.java new file mode 100644 index 000000000000..e1c31162729d --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChanges.java @@ -0,0 +1,145 @@ +/* + * 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.HashSet; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.join.CoGbkResult; +import org.apache.beam.sdk.util.RowFilter; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.iceberg.types.Types; +import org.joda.time.Instant; + +/** + * Receives a {@link CoGbkResult} containing inserts and deletes sharing the same snapshot sequence + * number and Primary Key, and uses {@link CdcResolver} to identify logical updates. + */ +class ResolveChanges 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; + + ResolveChanges(IcebergScanConfig scanConfig) { + this.scanConfig = scanConfig; + this.rowFilter = + new RowFilter( + CdcOutputUtils.readBeamSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getSchema())) + .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())); + } + + @ProcessElement + public void processElement( + @Element KV element, + @Timestamp Instant timestamp, + OutputReceiver out) { + CdcRowDescriptor descriptor = element.getKey(); + + Set pkFields = new HashSet<>(descriptor.getPrimaryKey().getSchema().getFieldNames()); + 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)); + + new RowResolver(pkFields, scanConfig.getMetadataColumns()) + .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(); + }); + } + + private static final class RowResolver extends CdcResolver { + private final Set pkFields; + private final List metadataColumns; + + RowResolver(Set pkFields, List metadataColumns) { + this.pkFields = pkFields; + this.metadataColumns = metadataColumns; + } + + @Override + protected int nonPkHash(Row element) { + int hash = 1; + for (String field : element.getSchema().getFieldNames()) { + if (pkFields.contains(field) + || (IcebergCdcMetadataColumns.isSupportedColumn(field) + && metadataColumns.contains(field))) { + continue; + } + hash = + 31 * hash + + Row.Equals.deepHashCode( + element.getValue(field), element.getSchema().getField(field).getType()); + } + return hash; + } + + @Override + protected boolean nonPkEquals(Row delete, Row insert) { + Schema schema = insert.getSchema(); + for (String field : schema.getFieldNames()) { + // we already know PK values are equal + if (pkFields.contains(field) + || (IcebergCdcMetadataColumns.isSupportedColumn(field) + && metadataColumns.contains(field))) { + continue; + } + // return early if two values are not equal + if (!Row.Equals.deepEquals( + insert.getValue(field), delete.getValue(field), schema.getField(field).getType())) { + 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/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..682f5f8e1447 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,15 @@ 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("max_snapshot_discovery_delay", 321L) + .withFieldValue( + "include_metadata_columns", + ImmutableList.of( + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID, + IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER)) .build(); new IcebergCdcReadSchemaTransformProvider().from(config); @@ -163,4 +187,96 @@ 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("max_snapshot_discovery_delay", 30L); + 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..64f16f83ab5d 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,8 @@ public class IcebergSchemaTransformTranslationTest { .withFieldValue("streaming", true) .withFieldValue("keep", ImmutableList.of("id", "event_micros")) .withFieldValue("filter", "\"data\" = 'keep'") + .withFieldValue("watermark_column", "event_micros") + .withFieldValue("max_snapshot_discovery_delay", 321L) .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 f0c7ae925df7..b03f9fcdc89c 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,13 +56,16 @@ 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; 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; @@ -74,10 +77,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; @@ -85,15 +94,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; @@ -277,6 +293,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() { @@ -320,6 +338,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 @@ -560,6 +581,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); @@ -621,7 +767,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()); @@ -1035,7 +1183,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); @@ -1068,6 +1218,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, + @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/CdcResolverTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolverTest.java new file mode 100644 index 000000000000..4f1340efb604 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolverTest.java @@ -0,0 +1,118 @@ +/* + * 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 java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link CdcResolver}. */ +@RunWith(JUnit4.class) +public class CdcResolverTest { + private static final TestResolver RESOLVER = new TestResolver(); + + @Test + public void duplicateDeleteInsertIsDropped() { + List emitted = + resolve( + Collections.singletonList(item("same", 7)), Collections.singletonList(item("same", 7))); + + assertThat(emitted, empty()); + } + + @Test + public void changedDeleteInsertBecomesUpdatePair() { + List emitted = + resolve( + Collections.singletonList(item("before", 1)), + Collections.singletonList(item("after", 2))); + + assertThat(emitted, contains("UPDATE_BEFORE:before", "UPDATE_AFTER:after")); + } + + @Test + public void duplicateUpdateAndSingletonsResolveByMultiplicity() { + List emitted = + resolve( + Arrays.asList(item("copy", 1), item("old", 2), item("deleted-only", 3)), + Arrays.asList(item("copy", 1), item("new", 4))); + + assertThat(emitted, contains("UPDATE_BEFORE:old", "UPDATE_AFTER:new", "DELETE:deleted-only")); + } + + @Test + public void hashCollisionOnlyConsumesEqualInsertOnce() { + List emitted = + resolve( + Arrays.asList(item("copy", 9), item("deleted-only", 9)), + Collections.singletonList(item("copy", 9))); + + assertThat(emitted, contains("DELETE:deleted-only")); + } + + @Test + public void hashMatchAloneDoesNotDeduplicate() { + List emitted = + resolve( + Collections.singletonList(item("before", 42)), + Collections.singletonList(item("after", 42))); + + assertThat(emitted, contains("UPDATE_BEFORE:before", "UPDATE_AFTER:after")); + } + + private static Item item(String nonPkValue, int hash) { + return new Item(nonPkValue, hash); + } + + private static List resolve(List deletes, List inserts) { + List emitted = new ArrayList<>(); + RESOLVER.resolve( + deletes, inserts, (kind, item) -> emitted.add(kind.name() + ":" + item.nonPkValue)); + return emitted; + } + + private static class TestResolver extends CdcResolver { + @Override + protected int nonPkHash(Item element) { + return element.hash; + } + + @Override + protected boolean nonPkEquals(Item delete, Item insert) { + return delete.nonPkValue.equals(insert.nonPkValue); + } + } + + private static class Item { + private final String nonPkValue; + private final int hash; + + private Item(String nonPkValue, int hash) { + this.nonPkValue = nonPkValue; + this.hash = hash; + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ChangelogScannerTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ChangelogScannerTest.java index 3c43b3ecc29a..3051a5c1ff5a 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ChangelogScannerTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ChangelogScannerTest.java @@ -49,9 +49,11 @@ import org.apache.iceberg.DeletedDataFileScanTask; import org.apache.iceberg.FileFormat; import org.apache.iceberg.Metrics; +import org.apache.iceberg.PartitionKey; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.StructLike; +import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.Record; import org.apache.iceberg.expressions.Expression; import org.apache.iceberg.expressions.ExpressionParser; @@ -83,6 +85,8 @@ public class ChangelogScannerTest { private static final Schema SINGLE_RECORD_ID_SCHEMA = recordIdSchema(SINGLE_PK_SCHEMA); private static final Schema COMPOSITE_RECORD_ID_SCHEMA = recordIdSchema(COMPOSITE_PK_SCHEMA); private static final PartitionSpec UNPARTITIONED_SPEC = PartitionSpec.unpartitioned(); + private static final PartitionSpec IDENTITY_ID_SPEC = + PartitionSpec.builderFor(SINGLE_PK_SCHEMA).identity("id").build(); @Test public void analyzeFilesPrunesNonOverlappingOpposingTasksToUnidirectional() { @@ -308,6 +312,15 @@ private static Map bounds(Schema schema, Map 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(); + } + + 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/LocalResolveDoFnTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java new file mode 100644 index 000000000000..140fec81785a --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java @@ -0,0 +1,226 @@ +/* + * 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.List; +import java.util.Map; +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.transforms.DoFnTester; +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.ChangelogOperation; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.expressions.ExpressionParser; +import org.apache.iceberg.expressions.Expressions; +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 LocalResolveDoFn}. */ +@RunWith(JUnit4.class) +public class LocalResolveDoFnTest { + 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, "visible", Types.StringType.get()), + Types.NestedField.optional(3, "hidden", Types.StringType.get())), + 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(); + + @Test + public void copyOnWriteRewriteOfIdenticalRowsIsDropped() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId); + DataFile oldFile = + warehouse.writeRecords( + testName.getMethodName() + "-old.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "same-hidden"))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "same-hidden"))); + + List> output = + process( + scanConfig, + descriptor(tableId, 1L, 1L), + ImmutableList.of( + task(SerializableChangelogTask.Type.DELETED_FILE, oldFile, table, 300L), + task(SerializableChangelogTask.Type.ADDED_ROWS, newFile, table, 300L)), + new Instant(0L)); + + assertThat(output, empty()); + } + + @Test + public void hiddenOnlyUpdateIsResolvedBeforeProjection() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId); + DataFile oldFile = + warehouse.writeRecords( + testName.getMethodName() + "-old.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "old-hidden"))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "new-hidden"))); + Instant timestamp = new Instant(1234L); + + List> output = + process( + scanConfig, + descriptor(tableId, 1L, 1L), + ImmutableList.of( + task(SerializableChangelogTask.Type.DELETED_FILE, oldFile, table, 301L), + task(SerializableChangelogTask.Type.ADDED_ROWS, newFile, table, 301L)), + timestamp); + + assertThat( + output.stream().map(LocalResolveDoFnTest::kindAndProjectedRow).collect(Collectors.toList()), + contains("UPDATE_BEFORE:1:shown:2", "UPDATE_AFTER:1:shown:2")); + assertEquals( + ImmutableList.of(timestamp, timestamp), + output.stream().map(ValueInSingleWindow::getTimestamp).collect(Collectors.toList())); + } + + private List> process( + IcebergScanConfig scanConfig, + ChangelogDescriptor descriptor, + List tasks, + Instant timestamp) + throws Exception { + try (DoFnTester>, Row> tester = + DoFnTester.of(new LocalResolveDoFn(scanConfig))) { + tester.processTimestampedElement(TimestampedValue.of(KV.of(descriptor, tasks), timestamp)); + return tester.getMutableOutput(tester.getMainOutputTag()); + } + } + + private TableIdentifier tableId() { + return TableIdentifier.of("default", testName.getMethodName()); + } + + private IcebergScanConfig scanConfig(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())) + .setKeepFields(ImmutableList.of("id", "visible")) + .setUseCdc(true) + .build(); + } + + private static ChangelogDescriptor descriptor( + TableIdentifier tableId, long lowerInclusive, long upperInclusive) { + org.apache.beam.sdk.schemas.Schema pkSchema = + org.apache.beam.sdk.schemas.Schema.builder().addInt64Field("id").build(); + return ChangelogDescriptor.builder() + .setTableIdentifierString(tableId.toString()) + .setSnapshotSequenceNumber(1) + .setCommitSnapshotId(1) + .setOverlapLower(Row.withSchema(pkSchema).addValue(lowerInclusive).build()) + .setOverlapUpper(Row.withSchema(pkSchema).addValue(upperInclusive).build()) + .build(); + } + + private static Record record(long id, String visible, String hidden) { + GenericRecord record = GenericRecord.create(CDC_SCHEMA); + record.setField("id", id); + record.setField("visible", visible); + record.setField("hidden", hidden); + return record; + } + + private static SerializableChangelogTask task( + SerializableChangelogTask.Type type, DataFile dataFile, Table table, long snapshotId) { + return SerializableChangelogTask.builder() + .setType(type) + .setDataFile(dataFile, table.spec().partitionToPath(dataFile.partition()), true) + .setAddedDeletes(ImmutableList.of()) + .setExistingDeletes(ImmutableList.of()) + .setSpecId(table.spec().specId()) + .setOperation( + type == SerializableChangelogTask.Type.ADDED_ROWS + ? ChangelogOperation.INSERT + : ChangelogOperation.DELETE) + .setOrdinal(0) + .setCommitSnapshotId(snapshotId) + .setStart(0L) + .setLength(dataFile.fileSizeInBytes()) + .setJsonExpression(ExpressionParser.toJson(Expressions.alwaysTrue())) + .build(); + } + + private static Map tableProperties() { + return ImmutableMap.of(TableProperties.FORMAT_VERSION, "2"); + } + + private static String kindAndProjectedRow(ValueInSingleWindow value) { + ValueKind kind = value.getValueKind(); + Row row = value.getValue(); + return kind.name() + + ":" + + row.getInt64("id") + + ":" + + row.getString("visible") + + ":" + + row.getSchema().getFieldCount(); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRangeTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRangeTest.java new file mode 100644 index 000000000000..69e47182e173 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRangeTest.java @@ -0,0 +1,161 @@ +/* + * 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.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import java.io.IOException; +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.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.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types; +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 OverlapRange}. */ +@RunWith(JUnit4.class) +public class OverlapRangeTest { + private static final org.apache.iceberg.Schema SINGLE_PK_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 org.apache.iceberg.Schema COMPOSITE_PK_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.optional(3, "data", Types.StringType.get()), + Types.NestedField.required(1, "account", Types.StringType.get()), + Types.NestedField.optional(4, "extra", Types.IntegerType.get()), + Types.NestedField.required(2, "sequence", Types.IntegerType.get())), + ImmutableSet.of(1, 2)); + + @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 containsUsesInclusiveSingleColumnBounds() throws Exception { + OverlapRange range = overlapRange(SINGLE_PK_SCHEMA); + StructLike lower = range.toStructLike(pkRow(range.recordIdSchema(), 10)); + StructLike upper = range.toStructLike(pkRow(range.recordIdSchema(), 20)); + + assertFalse(range.contains(singlePkRecord(9), lower, upper)); + assertTrue(range.contains(singlePkRecord(10), lower, upper)); + assertTrue(range.contains(singlePkRecord(15), lower, upper)); + assertTrue(range.contains(singlePkRecord(20), lower, upper)); + assertFalse(range.contains(singlePkRecord(21), lower, upper)); + } + + @Test + public void containsUsesLexicographicCompositeBounds() throws Exception { + OverlapRange range = overlapRange(COMPOSITE_PK_SCHEMA); + StructLike lower = range.toStructLike(pkRow(range.recordIdSchema(), "a", 2)); + StructLike upper = range.toStructLike(pkRow(range.recordIdSchema(), "b", 1)); + + assertFalse(range.contains(compositePkRecord("a", 1), lower, upper)); + assertTrue(range.contains(compositePkRecord("a", 2), lower, upper)); + assertTrue(range.contains(compositePkRecord("a", 9), lower, upper)); + assertTrue(range.contains(compositePkRecord("b", 0), lower, upper)); + assertTrue(range.contains(compositePkRecord("b", 1), lower, upper)); + assertFalse(range.contains(compositePkRecord("b", 2), lower, upper)); + } + + @Test + public void nullBoundsAreConservative() throws Exception { + OverlapRange range = overlapRange(SINGLE_PK_SCHEMA); + StructLike lower = range.toStructLike(pkRow(range.recordIdSchema(), 10)); + StructLike upper = range.toStructLike(pkRow(range.recordIdSchema(), 20)); + + assertNull(range.toStructLike(null)); + assertTrue(range.contains(singlePkRecord(1), null, upper)); + assertTrue(range.contains(singlePkRecord(100), lower, null)); + assertTrue(range.contains(singlePkRecord(100), null, null)); + } + + @Test + public void recordIdProjectionUsesIdentifierFieldsFromFullRecord() throws Exception { + OverlapRange range = overlapRange(COMPOSITE_PK_SCHEMA); + StructLike lower = range.toStructLike(pkRow(range.recordIdSchema(), "acct", 7)); + StructLike upper = range.toStructLike(pkRow(range.recordIdSchema(), "acct", 7)); + + assertTrue(range.contains(compositePkRecord("acct", 7), lower, upper)); + + assertEquals("acct", range.recordIdProjection().get(0, String.class)); + assertEquals(7, (int) range.recordIdProjection().get(1, Integer.class)); + } + + private OverlapRange overlapRange(org.apache.iceberg.Schema schema) throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + IcebergCatalogConfig catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build(); + catalogConfig.catalog().createTable(tableId, schema); + IcebergScanConfig scanConfig = + IcebergScanConfig.builder() + .setCatalogConfig(catalogConfig) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(schema)) + .setUseCdc(true) + .build(); + return OverlapRange.forScanConfig(scanConfig); + } + + private static Row pkRow(Schema recordIdSchema, Object... values) { + return Row.withSchema(IcebergUtils.icebergSchemaToBeamSchema(recordIdSchema)) + .addValues(values) + .build(); + } + + private static Record singlePkRecord(int id) { + GenericRecord record = GenericRecord.create(SINGLE_PK_SCHEMA); + record.setField("id", id); + record.setField("data", "v" + id); + return record; + } + + private static Record compositePkRecord(String account, int sequence) { + GenericRecord record = GenericRecord.create(COMPOSITE_PK_SCHEMA); + record.setField("data", "payload"); + record.setField("account", account); + record.setField("extra", 100); + record.setField("sequence", sequence); + return record; + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogsTest.java new file mode 100644 index 000000000000..69591e6eaa7c --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogsTest.java @@ -0,0 +1,366 @@ +/* + * 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.assertEquals; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.beam.sdk.coders.KvCoder; +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.SerializableDeleteFile; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.iceberg.ChangelogOperation; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.deletes.PositionDelete; +import org.apache.iceberg.deletes.PositionDeleteWriter; +import org.apache.iceberg.encryption.EncryptedFiles; +import org.apache.iceberg.expressions.ExpressionParser; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; +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; + +/** Tests for {@link ReadFromChangelogs}. */ +@RunWith(JUnit4.class) +public class ReadFromChangelogsTest { + 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, "visible", Types.StringType.get()), + Types.NestedField.optional(3, "hidden", Types.StringType.get())), + 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 unidirectionalTasksEmitProjectedRowsOnly() throws IOException { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId, ImmutableList.of("id", "visible")); + + DataFile addedFile = + warehouse.writeRecords( + testName.getMethodName() + "-added.parquet", + table.schema(), + ImmutableList.of(record(10L, "added", "added-hidden"))); + DataFile deletedRowsFile = + warehouse.writeRecords( + testName.getMethodName() + "-deleted-rows.parquet", + table.schema(), + ImmutableList.of( + record(20L, "deleted-row", "deleted-row-hidden"), + record(21L, "not-deleted", "not-deleted-hidden"))); + DeleteFile addedPositionDelete = + writePositionDelete(table, deletedRowsFile, "deleted-rows-pos-delete.parquet", 0L); + DataFile deletedFile = + warehouse.writeRecords( + testName.getMethodName() + "-deleted-file.parquet", + table.schema(), + ImmutableList.of(record(30L, "deleted-file", "deleted-file-hidden"))); + + List tasks = + ImmutableList.of( + task( + SerializableChangelogTask.Type.ADDED_ROWS, + addedFile, + ImmutableList.of(), + ImmutableList.of(), + table, + 100L), + task( + SerializableChangelogTask.Type.DELETED_ROWS, + deletedRowsFile, + ImmutableList.of(addedPositionDelete), + ImmutableList.of(), + table, + 100L), + task( + SerializableChangelogTask.Type.DELETED_FILE, + deletedFile, + ImmutableList.of(), + ImmutableList.of(), + table, + 100L)); + + ReadFromChangelogs.Output output = + input(ImmutableList.of(KV.of(descriptor(), tasks)), ImmutableList.of()) + .apply(new ReadFromChangelogs(scanConfig)); + + assertEquals( + IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()), + output.uniDirectionalRows().getSchema()); + PAssert.that( + output.uniDirectionalRows().apply("Format Unidirectional", ParDo.of(new FormatRow()))) + .containsInAnyOrder( + "INSERT:10:added:2", "DELETE:20:deleted-row:2", "DELETE:30:deleted-file:2"); + PAssert.that(output.biDirectionalInserts()).empty(); + PAssert.that(output.biDirectionalDeletes()).empty(); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void bidirectionalTasksKeepFullRowsForDownstreamResolution() throws IOException { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId, ImmutableList.of("id", "visible")); + DataFile oldFile = + warehouse.writeRecords( + testName.getMethodName() + "-old.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "old-hidden"))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "new-hidden"))); + List tasks = + ImmutableList.of( + task( + SerializableChangelogTask.Type.DELETED_FILE, + oldFile, + ImmutableList.of(), + ImmutableList.of(), + table, + 200L), + task( + SerializableChangelogTask.Type.ADDED_ROWS, + newFile, + ImmutableList.of(), + ImmutableList.of(), + table, + 200L)); + + ReadFromChangelogs.Output output = + input(ImmutableList.of(), ImmutableList.of(KV.of(descriptor(200L, 200L, 1L, 1L), tasks))) + .apply(new ReadFromChangelogs(scanConfig)); + + PAssert.that(output.uniDirectionalRows()).empty(); + PAssert.that( + output.biDirectionalDeletes().apply("Format Deletes", ParDo.of(new FormatKeyedRow()))) + .containsInAnyOrder("DELETE:200:200:1:shown:old-hidden:3"); + PAssert.that( + output.biDirectionalInserts().apply("Format Inserts", ParDo.of(new FormatKeyedRow()))) + .containsInAnyOrder("INSERT:200:200:1:shown:new-hidden:3"); + + pipeline.run().waitUntilFinish(); + } + + private PCollectionTuple input( + List>> unidirectional, + List>> largeBidirectional) { + Schema rowIdBeamSchema = + IcebergUtils.icebergSchemaToBeamSchema( + TypeUtil.select(CDC_SCHEMA, CDC_SCHEMA.identifierFieldIds())); + KvCoder> coder = + ChangelogScanner.coder(rowIdBeamSchema); + PCollection>> uni = + unidirectional.isEmpty() + ? pipeline.apply("Empty Unidirectional", Create.empty(coder)) + : pipeline.apply("Create Unidirectional", Create.of(unidirectional).withCoder(coder)); + PCollection>> large = + largeBidirectional.isEmpty() + ? pipeline.apply("Empty Large Bidirectional", Create.empty(coder)) + : pipeline.apply( + "Create Large Bidirectional", Create.of(largeBidirectional).withCoder(coder)); + return PCollectionTuple.of(ChangelogScanner.UNIDIRECTIONAL_TASKS, uni) + .and(ChangelogScanner.LARGE_BIDIRECTIONAL_TASKS, large); + } + + private TableIdentifier tableId() { + return TableIdentifier.of("default", testName.getMethodName()); + } + + private IcebergScanConfig scanConfig( + Table table, TableIdentifier tableId, List keepFields) { + return IcebergScanConfig.builder() + .setCatalogConfig( + IcebergCatalogConfig.builder() + .setCatalogName("name") + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build()) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(table.schema())) + .setKeepFields(keepFields) + .setUseCdc(true) + .build(); + } + + private ChangelogDescriptor descriptor() { + return ChangelogDescriptor.builder() + .setTableIdentifierString(tableId().toString()) + .setSnapshotSequenceNumber(100L) + .setCommitSnapshotId(100L) + .build(); + } + + private ChangelogDescriptor descriptor( + long sequenceNumber, long snapshotId, long lowerInclusive, long upperInclusive) { + Schema pkSchema = Schema.builder().addInt64Field("id").build(); + return ChangelogDescriptor.builder() + .setTableIdentifierString(tableId().toString()) + .setSnapshotSequenceNumber(sequenceNumber) + .setCommitSnapshotId(snapshotId) + .setOverlapLower(Row.withSchema(pkSchema).addValue(lowerInclusive).build()) + .setOverlapUpper(Row.withSchema(pkSchema).addValue(upperInclusive).build()) + .build(); + } + + private static Record record(long id, String visible, String hidden) { + GenericRecord record = GenericRecord.create(CDC_SCHEMA); + record.setField("id", id); + record.setField("visible", visible); + record.setField("hidden", hidden); + return record; + } + + private static DeleteFile writePositionDelete( + Table table, DataFile dataFile, String filename, long... positions) throws IOException { + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec()); + PositionDeleteWriter writer = + appenderFactory.newPosDeleteWriter( + EncryptedFiles.plainAsEncryptedOutput( + table.io().newOutputFile(dataFile.location() + "." + filename)), + FileFormat.PARQUET, + null); + try { + for (long position : positions) { + writer.write(PositionDelete.create().set(dataFile.location(), position)); + } + } finally { + writer.close(); + } + return writer.toDeleteFile(); + } + + private static SerializableChangelogTask task( + SerializableChangelogTask.Type type, + DataFile dataFile, + List addedDeletes, + List existingDeletes, + Table table, + long snapshotId) { + return SerializableChangelogTask.builder() + .setType(type) + .setDataFile(dataFile, table.spec().partitionToPath(dataFile.partition()), true) + .setAddedDeletes(serializableDeletes(addedDeletes, table)) + .setExistingDeletes(serializableDeletes(existingDeletes, table)) + .setSpecId(table.spec().specId()) + .setOperation( + type == SerializableChangelogTask.Type.ADDED_ROWS + ? ChangelogOperation.INSERT + : ChangelogOperation.DELETE) + .setOrdinal(0) + .setCommitSnapshotId(snapshotId) + .setStart(0L) + .setLength(dataFile.fileSizeInBytes()) + .setJsonExpression(ExpressionParser.toJson(Expressions.alwaysTrue())) + .build(); + } + + private static List serializableDeletes( + List deletes, Table table) { + return deletes.stream() + .map( + delete -> + SerializableDeleteFile.from( + delete, table.spec().partitionToPath(delete.partition()), true)) + .collect(Collectors.toList()); + } + + private static Map tableProperties() { + return ImmutableMap.of(TableProperties.FORMAT_VERSION, "2"); + } + + private static class FormatRow extends DoFn { + @ProcessElement + public void process(@Element Row row, ValueKind kind, OutputReceiver out) { + out.output( + kind.name() + + ":" + + row.getInt64("id") + + ":" + + row.getString("visible") + + ":" + + row.getSchema().getFieldCount()); + } + } + + private static class FormatKeyedRow extends DoFn, String> { + @ProcessElement + public void process( + @Element KV element, ValueKind kind, OutputReceiver out) { + Row row = element.getValue(); + CdcRowDescriptor descriptor = element.getKey(); + out.output( + kind.name() + + ":" + + descriptor.getCommitSnapshotId() + + ":" + + descriptor.getSnapshotSequenceNumber() + + ":" + + descriptor.getPrimaryKey().getInt64("id") + + ":" + + row.getString("visible") + + ":" + + row.getString("hidden") + + ":" + + row.getSchema().getFieldCount()); + } + } +} 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..8dbb3817543d --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java @@ -0,0 +1,218 @@ +/* + * 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, "old"); + Row after = simpleRow(1, "new"); + Row extraDelete = simpleRow(1, "deleted-only"); + + List> output = + process( + SIMPLE_ICEBERG_SCHEMA, + pkRow(1), + Arrays.asList(before, extraDelete), + Collections.singletonList(after), + timestamp); + + assertThat( + output.stream().map(ResolveChangesTest::kindAndData).collect(Collectors.toList()), + contains("UPDATE_BEFORE:old", "UPDATE_AFTER:new", "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(); + } +}