diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests.json b/.github/trigger_files/IO_Iceberg_Integration_Tests.json index 7ab7bcd9a9c6..b73af5e61a43 100644 --- a/.github/trigger_files/IO_Iceberg_Integration_Tests.json +++ b/.github/trigger_files/IO_Iceberg_Integration_Tests.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 2 + "modification": 1 } diff --git a/.github/trigger_files/beam_CloudML_Benchmarks_Dataflow.json b/.github/trigger_files/beam_CloudML_Benchmarks_Dataflow.json index 37dd25bf9029..5d04b2c0a8c7 100644 --- a/.github/trigger_files/beam_CloudML_Benchmarks_Dataflow.json +++ b/.github/trigger_files/beam_CloudML_Benchmarks_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 3 + "modification": 5 } diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index c03ecf71f04d..98bf4bf95003 100644 --- a/.github/trigger_files/beam_PostCommit_Python.json +++ b/.github/trigger_files/beam_PostCommit_Python.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", "pr": "37345", - "modification": 50 + "modification": 51 } diff --git a/CHANGES.md b/CHANGES.md index e5975afc56e5..e20194f9ab09 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -61,7 +61,6 @@ * Python SDK now supports memory profiling with Memray ([#38853](https://github.com/apache/beam/issues/38853)). * (Python) Added [Qdrant](https://qdrant.tech/) VectorDatabaseWriteConfig implementation ([#38141](https://github.com/apache/beam/issues/38141)). -* (CodeQL) Enabled Code scanning alerts in GitHub repo. ([#38893](https://github.com/apache/beam/issues/38893)). ## I/Os 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 a942c9804c93..cac4a015346c 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,21 +19,26 @@ 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 com.google.auto.value.AutoValue; import java.io.Serializable; import java.util.ArrayList; import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashSet; import java.util.List; import java.util.Set; 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; 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.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.Table; +import org.apache.iceberg.TableUtil; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.expressions.Evaluator; import org.apache.iceberg.expressions.Expression; @@ -48,7 +53,6 @@ public abstract class IcebergScanConfig implements Serializable { private transient @MonotonicNonNull Table cachedTable; private transient org.apache.iceberg.@MonotonicNonNull Schema cachedProjectedSchema; private transient org.apache.iceberg.@MonotonicNonNull Schema cachedRequiredSchema; - private transient @MonotonicNonNull Evaluator cachedEvaluator; private transient @MonotonicNonNull Expression cachedFilter; public enum ScanType { @@ -142,15 +146,12 @@ public org.apache.iceberg.Schema getRequiredSchema() { @Pure @Nullable - public Evaluator getEvaluator() { + public Evaluator getEvaluator(org.apache.iceberg.Schema requiredSchema) { @Nullable Expression filter = getFilter(); if (filter == null) { return null; } - if (cachedEvaluator == null) { - cachedEvaluator = new Evaluator(getRequiredSchema().asStruct(), filter); - } - return cachedEvaluator; + return new Evaluator(requiredSchema.asStruct(), filter); } @Pure @@ -225,6 +226,9 @@ public Expression getFilter() { @Pure public abstract @Nullable List getDropFields(); + @Pure + public abstract List getMetadataColumns(); + @Pure public static Builder builder() { return new AutoValue_IcebergScanConfig.Builder() @@ -247,7 +251,8 @@ public static Builder builder() { .setPollInterval(null) .setStartingStrategy(null) .setTag(null) - .setBranch(null); + .setBranch(null) + .setMetadataColumns(ImmutableList.of()); } @AutoValue.Builder @@ -310,6 +315,8 @@ public Builder setTableIdentifier(String... names) { public abstract Builder setDropFields(@Nullable List fields); + public abstract Builder setMetadataColumns(List metadataColumns); + public abstract IcebergScanConfig build(); } @@ -363,6 +370,9 @@ void validate(Table table) { if (getStartingStrategy() != null) { invalidOptions.add("starting_strategy"); } + if (!getMetadataColumns().isEmpty()) { + invalidOptions.add("metadata_columns"); + } if (!invalidOptions.isEmpty()) { throw new IllegalArgumentException( error( @@ -370,6 +380,19 @@ void validate(Table table) { + "reading with Managed.ICEBERG_CDC: " + invalidOptions)); } + } else { + Set primaryKeyIds = new HashSet<>(table.schema().identifierFieldIds()); + checkState( + !primaryKeyIds.isEmpty(), + "Cannot read CDC records as the table schema does not specified any primary key fields."); + Set projectedFieldIds = TypeUtil.getProjectedIds(getProjectedSchema()); + primaryKeyIds.removeAll(projectedFieldIds); + checkArgument( + primaryKeyIds.isEmpty(), + "When reading CDC records, the projected schema must not drop primary key fields. " + + "The specified configuration drops the following PK fields: %s", + primaryKeyIds); + validateMetadataColumns(table); } if (getStartingStrategy() != null) { @@ -392,6 +415,47 @@ void validate(Table table) { } } + private void validateMetadataColumns(Table table) { + List metadataColumns = getMetadataColumns(); + if (metadataColumns.isEmpty()) { + return; + } + + Set uniqueMetadataColumns = new LinkedHashSet<>(metadataColumns); + checkArgument( + uniqueMetadataColumns.size() == metadataColumns.size(), + error("metadata_columns contains duplicate entries: %s"), + metadataColumns); + + List unsupportedMetadataColumns = new ArrayList<>(); + for (String metadataColumn : metadataColumns) { + if (!IcebergCdcMetadataColumns.isSupportedColumn(metadataColumn)) { + unsupportedMetadataColumns.add(metadataColumn); + } + } + checkArgument( + unsupportedMetadataColumns.isEmpty(), + error("unsupported metadata_columns: %s. Supported values are: %s"), + unsupportedMetadataColumns, + IcebergCdcMetadataColumns.SUPPORTED_COLUMNS); + + for (String metadataColumn : metadataColumns) { + checkArgument( + getProjectedSchema().findField(metadataColumn) == null, + error("metadata column '%s' conflicts with a projected data column"), + metadataColumn); + } + + boolean includesRowLineage = + metadataColumns.stream().anyMatch(IcebergCdcMetadataColumns::isRowMetadataColumn); + if (includesRowLineage) { + checkArgument( + TableUtil.formatVersion(table) >= 3, + error("row lineage metadata columns %s are only available for Iceberg format v3+ tables"), + metadataColumns); + } + } + private String error(String message) { return "Invalid source configuration: " + message; } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/PartitionUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/PartitionUtils.java index 805cc0672940..32a25439d850 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/PartitionUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/PartitionUtils.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.io.iceberg; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.iceberg.data.IdentityPartitionConverters.convertConstant; import java.util.List; import java.util.Map; @@ -25,11 +26,20 @@ import java.util.function.Function; import java.util.regex.Matcher; import java.util.regex.Pattern; +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.Maps; +import org.apache.iceberg.ChangelogScanTask; +import org.apache.iceberg.ContentFile; +import org.apache.iceberg.ContentScanTask; +import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.PartitionField; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; import org.apache.iceberg.expressions.Expressions; import org.apache.iceberg.expressions.Term; +import org.apache.iceberg.types.Types; import org.checkerframework.checker.nullness.qual.Nullable; class PartitionUtils { @@ -48,23 +58,28 @@ class PartitionUtils { Pattern, BiFunction> TRANSFORMATIONS = ImmutableMap.of( - HOUR, (builder, matcher) -> builder.hour(checkStateNotNull(matcher.group(1))), - DAY, (builder, matcher) -> builder.day(checkStateNotNull(matcher.group(1))), - MONTH, (builder, matcher) -> builder.month(checkStateNotNull(matcher.group(1))), - YEAR, (builder, matcher) -> builder.year(checkStateNotNull(matcher.group(1))), + HOUR, + (builder, matcher) -> builder.hour(checkStateNotNull(matcher.group(1))), + DAY, + (builder, matcher) -> builder.day(checkStateNotNull(matcher.group(1))), + MONTH, + (builder, matcher) -> builder.month(checkStateNotNull(matcher.group(1))), + YEAR, + (builder, matcher) -> builder.year(checkStateNotNull(matcher.group(1))), TRUNCATE, - (builder, matcher) -> - builder.truncate( - checkStateNotNull(matcher.group(1)), - Integer.parseInt(checkStateNotNull(matcher.group(2)))), + (builder, matcher) -> + builder.truncate( + checkStateNotNull(matcher.group(1)), + Integer.parseInt(checkStateNotNull(matcher.group(2)))), BUCKET, - (builder, matcher) -> - builder.bucket( - checkStateNotNull(matcher.group(1)), - Integer.parseInt(checkStateNotNull(matcher.group(2)))), - VOID, (builder, matcher) -> builder.alwaysNull(checkStateNotNull(matcher.group(1))), + (builder, matcher) -> + builder.bucket( + checkStateNotNull(matcher.group(1)), + Integer.parseInt(checkStateNotNull(matcher.group(2)))), + VOID, + (builder, matcher) -> builder.alwaysNull(checkStateNotNull(matcher.group(1))), IDENTITY, - (builder, matcher) -> builder.identity(checkStateNotNull(matcher.group(1)))); + (builder, matcher) -> builder.identity(checkStateNotNull(matcher.group(1)))); static PartitionSpec toPartitionSpec( @Nullable List fields, org.apache.beam.sdk.schemas.Schema beamSchema) { @@ -130,4 +145,61 @@ static Term toIcebergTerm(String field) { throw new IllegalArgumentException("Could not find a partition term for '" + field + "'."); } + + /** + * Copied over from Apache Iceberg's PartitionUtil. + * + *

Needed to accommodate CDC reads, where scans produce {@link ChangelogScanTask}s instead of + * {@link ContentScanTask}s. + */ + public static Map constantsMap( + PartitionSpec spec, ContentFile file, @Nullable Long fileSequenceNumber) { + Preconditions.checkState( + spec.specId() == file.specId(), + "File spec ID (%s) does not match PartitionSpec ID (%s)", + file.specId(), + spec.specId()); + StructLike partitionData = file.partition(); + + // use java.util.HashMap because partition data may contain null values + Map idToConstant = Maps.newHashMap(); + + // add first_row_id as _row_id + if (file.firstRowId() != null) { + idToConstant.put( + MetadataColumns.ROW_ID.fieldId(), + convertConstant(Types.LongType.get(), file.firstRowId())); + } + + // When reconstructing a DataFile, we lose the ability to attach its fileSequenceNumber, + // so we pipe it along the util methods to include it here. + fileSequenceNumber = + fileSequenceNumber != null ? fileSequenceNumber : file.fileSequenceNumber(); + idToConstant.put( + MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.fieldId(), + convertConstant(Types.LongType.get(), fileSequenceNumber)); + + // add _file + idToConstant.put( + MetadataColumns.FILE_PATH.fieldId(), + convertConstant(Types.StringType.get(), file.location())); + + // add _spec_id + idToConstant.put( + MetadataColumns.SPEC_ID.fieldId(), convertConstant(Types.IntegerType.get(), file.specId())); + + List partitionFields = spec.partitionType().fields(); + List fields = spec.fields(); + for (int pos = 0; pos < fields.size(); pos += 1) { + PartitionField field = fields.get(pos); + if (field.transform().isIdentity()) { + Object converted = + convertConstant(partitionFields.get(pos).type(), partitionData.get(pos, Object.class)); + idToConstant.put(field.sourceId(), converted); + } + } + + return idToConstant; + } } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java index 5eeeacda48e3..71114437731c 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java @@ -70,9 +70,7 @@ public void process( } FileScanTask task = fileScanTasks.get((int) l); Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()); - try (CloseableIterable fullIterable = - ReadUtils.createReader(task, table, scanConfig.getRequiredSchema())) { - CloseableIterable reader = ReadUtils.maybeApplyFilter(fullIterable, scanConfig); + try (CloseableIterable reader = ReadUtils.createReader(task, table, scanConfig)) { for (Record record : reader) { Row row = IcebergUtils.icebergRecordToBeamRow(beamSchema, record); diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadUtils.java index e7f50882f433..ea2dc7589c4e 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadUtils.java @@ -20,25 +20,22 @@ import static org.apache.iceberg.util.SnapshotUtil.ancestorsOf; import java.util.Collection; -import java.util.Collections; import java.util.Comparator; import java.util.List; import java.util.Map; -import java.util.Set; -import java.util.function.BiFunction; import java.util.stream.Collectors; import org.apache.beam.sdk.io.iceberg.IcebergIO.ReadRows.StartingStrategy; 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.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; import org.apache.hadoop.conf.Configuration; -import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.ContentFile; +import org.apache.iceberg.ContentScanTask; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.Snapshot; import org.apache.iceberg.Table; import org.apache.iceberg.TableProperties; -import org.apache.iceberg.data.IdentityPartitionConverters; import org.apache.iceberg.data.InternalRecordWrapper; import org.apache.iceberg.data.Record; import org.apache.iceberg.data.parquet.GenericParquetReaders; @@ -46,17 +43,12 @@ import org.apache.iceberg.encryption.EncryptedInputFile; import org.apache.iceberg.expressions.Evaluator; import org.apache.iceberg.expressions.Expression; -import org.apache.iceberg.expressions.Expressions; import org.apache.iceberg.hadoop.HadoopInputFile; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.InputFile; -import org.apache.iceberg.mapping.MappingUtil; import org.apache.iceberg.mapping.NameMapping; import org.apache.iceberg.mapping.NameMappingParser; import org.apache.iceberg.parquet.ParquetReader; -import org.apache.iceberg.types.Type; -import org.apache.iceberg.types.TypeUtil; -import org.apache.iceberg.util.PartitionUtil; import org.apache.iceberg.util.SnapshotUtil; import org.apache.parquet.HadoopReadOptions; import org.apache.parquet.ParquetReadOptions; @@ -73,13 +65,34 @@ public class ReadUtils { "parquet.read.support.class", "parquet.crypto.factory.class"); - static ParquetReader createReader(FileScanTask task, Table table, Schema schema) { - String filePath = task.file().path().toString(); + public static CloseableIterable createReader( + ContentScanTask task, Table table, IcebergScanConfig scanConfig) { + return createReader( + table, + scanConfig, + scanConfig.getRequiredSchema(), + task.spec(), + task.file(), + null, + task.start(), + task.length(), + task.residual()); + } + + public static CloseableIterable createReader( + Table table, + IcebergScanConfig scanConfig, + Schema requiredSchema, + PartitionSpec spec, + ContentFile file, + @Nullable Long fileSequenceNumber, + long start, + long length, + Expression residual) { EncryptedInputFile encryptedInput = - EncryptedFiles.encryptedInput(table.io().newInputFile(filePath), task.file().keyMetadata()); + EncryptedFiles.encryptedInput(table.io().newInputFile(file.location()), file.keyMetadata()); InputFile inputFile = table.encryption().decrypt(encryptedInput); - Map idToConstants = - ReadUtils.constantsMap(task, IdentityPartitionConverters::convertConstant, table.schema()); + Map idToConstants = PartitionUtils.constantsMap(spec, file, fileSequenceNumber); ParquetReadOptions.Builder optionsBuilder; if (inputFile instanceof HadoopInputFile) { @@ -94,71 +107,30 @@ static ParquetReader createReader(FileScanTask task, Table table, Schema } optionsBuilder = optionsBuilder - .withRange(task.start(), task.start() + task.length()) + .withRange(start, start + length) .withMaxAllocationInBytes(MAX_FILE_BUFFER_SIZE); @Nullable String nameMapping = table.properties().get(TableProperties.DEFAULT_NAME_MAPPING); NameMapping mapping = nameMapping != null ? NameMappingParser.fromJson(nameMapping) : NameMapping.empty(); - return new ParquetReader<>( - inputFile, - schema, - optionsBuilder.build(), - // TODO(ahmedabu98): Implement a Parquet-to-Beam Row reader, bypassing conversion to Iceberg - // Record - fileSchema -> GenericParquetReaders.buildReader(schema, fileSchema, idToConstants), - mapping, - task.residual(), - false, - true); + ParquetReader records = + new ParquetReader<>( + inputFile, + requiredSchema, + optionsBuilder.build(), + // TODO(ahmedabu98): Implement a Parquet-to-Beam Row reader, bypassing conversion to + // Iceberg Record + fileSchema -> + GenericParquetReaders.buildReader(requiredSchema, fileSchema, idToConstants), + mapping, + residual, + false, + true); + return maybeApplyFilter(records, scanConfig, requiredSchema); } - static ParquetReader createReader(InputFile inputFile, Schema schema) { - ParquetReadOptions.Builder optionsBuilder; - if (inputFile instanceof HadoopInputFile) { - // remove read properties already set that may conflict with this read - Configuration conf = new Configuration(((HadoopInputFile) inputFile).getConf()); - for (String property : READ_PROPERTIES_TO_REMOVE) { - conf.unset(property); - } - optionsBuilder = HadoopReadOptions.builder(conf); - } else { - optionsBuilder = ParquetReadOptions.builder(); - } - optionsBuilder = - optionsBuilder - .withRange(0, inputFile.getLength()) - .withMaxAllocationInBytes(MAX_FILE_BUFFER_SIZE); - - return new ParquetReader<>( - inputFile, - schema, - optionsBuilder.build(), - fileSchema -> GenericParquetReaders.buildReader(schema, fileSchema), - MappingUtil.create(schema), - Expressions.alwaysTrue(), - false, - true); - } - - static Map constantsMap( - FileScanTask task, - BiFunction converter, - org.apache.iceberg.Schema schema) { - PartitionSpec spec = task.spec(); - Set idColumns = spec.identitySourceIds(); - org.apache.iceberg.Schema partitionSchema = TypeUtil.select(schema, idColumns); - boolean projectsIdentityPartitionColumns = !partitionSchema.columns().isEmpty(); - - if (projectsIdentityPartitionColumns) { - return PartitionUtil.constantsMap(task, converter); - } else { - return Collections.emptyMap(); - } - } - - static @Nullable Long getFromSnapshotExclusive(Table table, IcebergScanConfig scanConfig) { + public static @Nullable Long getFromSnapshotInclusive(Table table, IcebergScanConfig scanConfig) { @Nullable StartingStrategy startingStrategy = scanConfig.getStartingStrategy(); boolean isStreaming = MoreObjects.firstNonNull(scanConfig.getStreaming(), false); if (startingStrategy == null) { @@ -179,6 +151,13 @@ static ParquetReader createReader(InputFile inputFile, Schema schema) { fromSnapshot = currentSnapshot.snapshotId(); } } + + return fromSnapshot; + } + + public static @Nullable Long getFromSnapshotExclusive(Table table, IcebergScanConfig scanConfig) { + @Nullable Long fromSnapshot = getFromSnapshotInclusive(table, scanConfig); + // incremental append scan can only be configured with an *exclusive* starting snapshot, // so we need to provide this snapshot's parent id. if (fromSnapshot != null) { @@ -189,7 +168,7 @@ static ParquetReader createReader(InputFile inputFile, Schema schema) { return fromSnapshot; } - static @Nullable Long getToSnapshot(Table table, IcebergScanConfig scanConfig) { + public static @Nullable Long getToSnapshot(Table table, IcebergScanConfig scanConfig) { // 1. fetch from to_snapshot @Nullable Long toSnapshot = scanConfig.getToSnapshot(); // 2. fetch from to_timestamp @@ -205,7 +184,7 @@ static ParquetReader createReader(InputFile inputFile, Schema schema) { * Returns a list of snapshots in the range (fromSnapshotId, toSnapshotId], ordered * chronologically. */ - static List snapshotsBetween( + public static List snapshotsBetween( Table table, String tableIdentifier, @Nullable Long fromSnapshotId, long toSnapshotId) { long from = MoreObjects.firstNonNull(fromSnapshotId, -1L); @SuppressWarnings("return") @@ -225,10 +204,14 @@ static List snapshotsBetween( public static CloseableIterable maybeApplyFilter( CloseableIterable iterable, IcebergScanConfig scanConfig) { - InternalRecordWrapper wrapper = - new InternalRecordWrapper(scanConfig.getRequiredSchema().asStruct()); + return maybeApplyFilter(iterable, scanConfig, scanConfig.getRequiredSchema()); + } + + public static CloseableIterable maybeApplyFilter( + CloseableIterable iterable, IcebergScanConfig scanConfig, Schema requiredSchema) { + InternalRecordWrapper wrapper = new InternalRecordWrapper(requiredSchema.asStruct()); Expression filter = scanConfig.getFilter(); - Evaluator evaluator = scanConfig.getEvaluator(); + Evaluator evaluator = scanConfig.getEvaluator(requiredSchema); if (filter != null && evaluator != null && filter.op() != Expression.Operation.TRUE) { return CloseableIterable.filter(iterable, record -> evaluator.eval(wrapper.wrap(record))); } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ScanTaskReader.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ScanTaskReader.java index b3485a7bcc4f..c9ad372a0751 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ScanTaskReader.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ScanTaskReader.java @@ -36,7 +36,6 @@ import org.apache.iceberg.TableProperties; import org.apache.iceberg.avro.Avro; import org.apache.iceberg.data.GenericDeleteFilter; -import org.apache.iceberg.data.IdentityPartitionConverters; import org.apache.iceberg.data.Record; import org.apache.iceberg.data.avro.DataReader; import org.apache.iceberg.data.orc.GenericOrcReader; @@ -121,8 +120,7 @@ public boolean advance() throws IOException { DataFile file = fileTask.file(); InputFile input = decryptor.getInputFile(fileTask); Map idToConstants = - ReadUtils.constantsMap( - fileTask, IdentityPartitionConverters::convertConstant, requiredSchema); + PartitionUtils.constantsMap(fileTask.spec(), fileTask.file(), null); CloseableIterable iterable; switch (file.format()) { diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableDataFile.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableDataFile.java index 9e75be0a1987..e1291601d149 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableDataFile.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableDataFile.java @@ -54,13 +54,13 @@ */ @DefaultSchema(AutoValueSchema.class) @AutoValue -abstract class SerializableDataFile { +public abstract class SerializableDataFile { public static Builder builder() { return new AutoValue_SerializableDataFile.Builder(); } @SchemaFieldNumber("0") - abstract String getPath(); + public abstract String getPath(); @SchemaFieldNumber("1") abstract String getFileFormat(); @@ -69,10 +69,10 @@ public static Builder builder() { abstract long getRecordCount(); @SchemaFieldNumber("3") - abstract long getFileSizeInBytes(); + public abstract long getFileSizeInBytes(); @SchemaFieldNumber("4") - abstract String getPartitionPath(); + public abstract String getPartitionPath(); @SchemaFieldNumber("5") abstract int getPartitionSpecId(); @@ -96,13 +96,22 @@ public static Builder builder() { abstract @Nullable Map getNanValueCounts(); @SchemaFieldNumber("12") - abstract @Nullable Map getLowerBounds(); + public abstract @Nullable Map getLowerBounds(); @SchemaFieldNumber("13") - abstract @Nullable Map getUpperBounds(); + public abstract @Nullable Map getUpperBounds(); + + @SchemaFieldNumber("14") + public abstract @Nullable Long getDataSequenceNumber(); + + @SchemaFieldNumber("15") + public abstract @Nullable Long getFileSequenceNumber(); + + @SchemaFieldNumber("16") + public abstract @Nullable Long getFirstRowId(); @AutoValue.Builder - abstract static class Builder { + public abstract static class Builder { abstract Builder setPath(String path); abstract Builder setFileFormat(String fileFormat); @@ -131,31 +140,49 @@ abstract static class Builder { abstract Builder setUpperBounds(@Nullable Map upperBounds); + abstract Builder setDataSequenceNumber(@Nullable Long number); + + abstract Builder setFileSequenceNumber(@Nullable Long number); + + abstract Builder setFirstRowId(@Nullable Long id); + abstract SerializableDataFile build(); } + public static SerializableDataFile from(DataFile f, String partitionPath) { + return from(f, partitionPath, true); + } + /** * Create a {@link SerializableDataFile} from a {@link DataFile} and its associated {@link * PartitionKey}. */ - static SerializableDataFile from(DataFile f, String partitionPath) { - - return SerializableDataFile.builder() - .setPath(f.location().toString()) - .setFileFormat(f.format().toString()) - .setRecordCount(f.recordCount()) - .setFileSizeInBytes(f.fileSizeInBytes()) - .setPartitionPath(partitionPath) - .setPartitionSpecId(f.specId()) - .setKeyMetadata(f.keyMetadata()) - .setSplitOffsets(f.splitOffsets()) - .setColumnSizes(f.columnSizes()) - .setValueCounts(f.valueCounts()) - .setNullValueCounts(f.nullValueCounts()) - .setNanValueCounts(f.nanValueCounts()) - .setLowerBounds(toByteArrayMap(f.lowerBounds())) - .setUpperBounds(toByteArrayMap(f.upperBounds())) - .build(); + public static SerializableDataFile from( + DataFile f, String partitionPath, boolean includeMetrics) { + SerializableDataFile.Builder builder = + SerializableDataFile.builder() + .setPath(f.location()) + .setFileFormat(f.format().toString()) + .setRecordCount(f.recordCount()) + .setFileSizeInBytes(f.fileSizeInBytes()) + .setPartitionPath(partitionPath) + .setPartitionSpecId(f.specId()) + .setKeyMetadata(f.keyMetadata()) + .setSplitOffsets(f.splitOffsets()) + .setColumnSizes(f.columnSizes()) + .setValueCounts(f.valueCounts()) + .setNullValueCounts(f.nullValueCounts()) + .setNanValueCounts(f.nanValueCounts()) + .setDataSequenceNumber(f.dataSequenceNumber()) + .setFileSequenceNumber(f.fileSequenceNumber()) + .setFirstRowId(f.firstRowId()); + if (includeMetrics) { + builder = + builder + .setLowerBounds(toByteArrayMap(f.lowerBounds())) + .setUpperBounds(toByteArrayMap(f.upperBounds())); + } + return builder.build(); } /** @@ -165,7 +192,7 @@ static SerializableDataFile from(DataFile f, String partitionPath) { * it from Beam-compatible types. */ @SuppressWarnings("nullness") - DataFile createDataFile(Map partitionSpecs) { + public DataFile createDataFile(Map partitionSpecs) { PartitionSpec partitionSpec = checkStateNotNull( partitionSpecs.get(getPartitionSpecId()), @@ -192,14 +219,14 @@ DataFile createDataFile(Map partitionSpecs) { .withFileSizeInBytes(getFileSizeInBytes()) .withMetrics(dataFileMetrics) .withSplitOffsets(getSplitOffsets()) + .withFirstRowId(getFirstRowId()) .build(); } // ByteBuddyUtils has trouble converting Map value type ByteBuffer // to byte[] and back to ByteBuffer, so we perform these conversions manually // TODO(https://github.com/apache/beam/issues/32701) - private static @Nullable Map toByteArrayMap( - @Nullable Map input) { + static @Nullable Map toByteArrayMap(@Nullable Map input) { if (input == null) { return null; } @@ -222,8 +249,7 @@ private static byte[] toByteArray(ByteBuffer buf) { return bytes; } - private static @Nullable Map toByteBufferMap( - @Nullable Map input) { + static @Nullable Map toByteBufferMap(@Nullable Map input) { if (input == null) { return null; } @@ -256,10 +282,13 @@ && getPartitionSpecId() == that.getPartitionSpecId() && Objects.equals(getNullValueCounts(), that.getNullValueCounts()) && Objects.equals(getNanValueCounts(), that.getNanValueCounts()) && mapEquals(getLowerBounds(), that.getLowerBounds()) - && mapEquals(getUpperBounds(), that.getUpperBounds()); + && mapEquals(getUpperBounds(), that.getUpperBounds()) + && Objects.equals(getDataSequenceNumber(), that.getDataSequenceNumber()) + && Objects.equals(getFileSequenceNumber(), that.getFileSequenceNumber()) + && Objects.equals(getFirstRowId(), that.getFirstRowId()); } - private static boolean mapEquals( + static boolean mapEquals( @Nullable Map map1, @Nullable Map map2) { if (map1 == null && map2 == null) { return true; @@ -297,13 +326,16 @@ public final int hashCode() { getColumnSizes(), getValueCounts(), getNullValueCounts(), - getNanValueCounts()); + getNanValueCounts(), + getDataSequenceNumber(), + getFileSequenceNumber(), + getFirstRowId()); hashCode = 31 * hashCode + computeMapByteHashCode(getLowerBounds()); hashCode = 31 * hashCode + computeMapByteHashCode(getUpperBounds()); return hashCode; } - private static int computeMapByteHashCode(@Nullable Map map) { + static int computeMapByteHashCode(@Nullable Map map) { if (map == null) { return 0; } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableDeleteFile.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableDeleteFile.java new file mode 100644 index 000000000000..ceb96d50f8aa --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/SerializableDeleteFile.java @@ -0,0 +1,335 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg; + +import static org.apache.beam.sdk.io.iceberg.SerializableDataFile.computeMapByteHashCode; +import static org.apache.beam.sdk.io.iceberg.SerializableDataFile.mapEquals; +import static org.apache.beam.sdk.io.iceberg.SerializableDataFile.toByteArrayMap; +import static org.apache.beam.sdk.io.iceberg.SerializableDataFile.toByteBufferMap; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import com.google.auto.value.AutoValue; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileContent; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.SortOrder; +import org.checkerframework.checker.nullness.qual.Nullable; + +@DefaultSchema(AutoValueSchema.class) +@AutoValue +public abstract class SerializableDeleteFile { + public static SerializableDeleteFile.Builder builder() { + return new AutoValue_SerializableDeleteFile.Builder(); + } + + @SchemaFieldNumber("0") + public abstract FileContent getContentType(); + + @SchemaFieldNumber("1") + public abstract String getLocation(); + + @SchemaFieldNumber("2") + public abstract String getFileFormat(); + + @SchemaFieldNumber("3") + public abstract long getRecordCount(); + + @SchemaFieldNumber("4") + public abstract long getFileSizeInBytes(); + + @SchemaFieldNumber("5") + public abstract String getPartitionPath(); + + @SchemaFieldNumber("6") + public abstract int getPartitionSpecId(); + + @SchemaFieldNumber("7") + public abstract @Nullable Integer getSortOrderId(); + + @SchemaFieldNumber("8") + public abstract @Nullable List getEqualityFieldIds(); + + @SchemaFieldNumber("9") + public abstract @Nullable ByteBuffer getKeyMetadata(); + + @SchemaFieldNumber("10") + public abstract @Nullable List getSplitOffsets(); + + @SchemaFieldNumber("11") + public abstract @Nullable Map getColumnSizes(); + + @SchemaFieldNumber("12") + public abstract @Nullable Map getValueCounts(); + + @SchemaFieldNumber("13") + public abstract @Nullable Map getNullValueCounts(); + + @SchemaFieldNumber("14") + public abstract @Nullable Map getNanValueCounts(); + + @SchemaFieldNumber("15") + public abstract @Nullable Map getLowerBounds(); + + @SchemaFieldNumber("16") + public abstract @Nullable Map getUpperBounds(); + + @SchemaFieldNumber("17") + public abstract @Nullable Long getContentOffset(); + + @SchemaFieldNumber("18") + public abstract @Nullable Long getContentSizeInBytes(); + + @SchemaFieldNumber("19") + public abstract @Nullable String getReferencedDataFile(); + + @SchemaFieldNumber("20") + public abstract @Nullable Long getDataSequenceNumber(); + + @SchemaFieldNumber("21") + public abstract @Nullable Long getFileSequenceNumber(); + + @AutoValue.Builder + abstract static class Builder { + abstract Builder setContentType(FileContent content); + + abstract Builder setLocation(String path); + + abstract Builder setFileFormat(String fileFormat); + + abstract Builder setRecordCount(long recordCount); + + abstract Builder setFileSizeInBytes(long fileSizeInBytes); + + abstract Builder setPartitionPath(String partitionPath); + + abstract Builder setPartitionSpecId(int partitionSpec); + + abstract Builder setSortOrderId(@Nullable Integer sortOrderId); + + abstract Builder setEqualityFieldIds(List equalityFieldIds); + + abstract Builder setKeyMetadata(ByteBuffer keyMetadata); + + abstract Builder setSplitOffsets(List splitOffsets); + + abstract Builder setColumnSizes(Map columnSizes); + + abstract Builder setValueCounts(Map valueCounts); + + abstract Builder setNullValueCounts(Map nullValueCounts); + + abstract Builder setNanValueCounts(Map nanValueCounts); + + abstract Builder setLowerBounds(@Nullable Map lowerBounds); + + abstract Builder setUpperBounds(@Nullable Map upperBounds); + + abstract Builder setContentOffset(@Nullable Long offset); + + abstract Builder setContentSizeInBytes(@Nullable Long sizeInBytes); + + abstract Builder setReferencedDataFile(@Nullable String dataFile); + + abstract Builder setDataSequenceNumber(@Nullable Long number); + + abstract Builder setFileSequenceNumber(@Nullable Long number); + + abstract SerializableDeleteFile build(); + } + + public static SerializableDeleteFile from( + DeleteFile deleteFile, String partitionPath, boolean includeMetrics) { + + SerializableDeleteFile.Builder builder = + SerializableDeleteFile.builder() + .setLocation(deleteFile.location()) + .setFileFormat(deleteFile.format().name()) + .setFileSizeInBytes(deleteFile.fileSizeInBytes()) + .setPartitionPath(partitionPath) + .setPartitionSpecId(deleteFile.specId()) + .setRecordCount(deleteFile.recordCount()) + .setColumnSizes(deleteFile.columnSizes()) + .setValueCounts(deleteFile.valueCounts()) + .setNullValueCounts(deleteFile.nullValueCounts()) + .setNanValueCounts(deleteFile.nanValueCounts()) + .setSplitOffsets(deleteFile.splitOffsets()) + .setKeyMetadata(deleteFile.keyMetadata()) + .setEqualityFieldIds(deleteFile.equalityFieldIds()) + .setSortOrderId(deleteFile.sortOrderId()) + .setContentOffset(deleteFile.contentOffset()) + .setContentSizeInBytes(deleteFile.contentSizeInBytes()) + .setReferencedDataFile(deleteFile.referencedDataFile()) + .setContentType(deleteFile.content()) + .setDataSequenceNumber(deleteFile.dataSequenceNumber()) + .setFileSequenceNumber(deleteFile.fileSequenceNumber()); + + if (includeMetrics) { + builder = + builder + .setLowerBounds(toByteArrayMap(deleteFile.lowerBounds())) + .setUpperBounds(toByteArrayMap(deleteFile.upperBounds())); + } + + return builder.build(); + } + + @SuppressWarnings("nullness") + public DeleteFile createDeleteFile( + Map partitionSpecs, @Nullable Map sortOrders) { + PartitionSpec partitionSpec = + checkStateNotNull( + partitionSpecs.get(getPartitionSpecId()), + "This DeleteFile was originally created with spec id '%s', " + + "but table only has spec ids: %s.", + getPartitionSpecId(), + partitionSpecs.keySet()); + + Metrics metrics = + new Metrics( + getRecordCount(), + getColumnSizes(), + getValueCounts(), + getNullValueCounts(), + getNanValueCounts(), + toByteBufferMap(getLowerBounds()), + toByteBufferMap(getUpperBounds())); + + FileMetadata.Builder deleteFileBuilder = + FileMetadata.deleteFileBuilder(partitionSpec) + .withPath(getLocation()) + .withFormat(getFileFormat()) + .withFileSizeInBytes(getFileSizeInBytes()) + .withRecordCount(getRecordCount()) + .withMetrics(metrics) + .withSplitOffsets(getSplitOffsets()) + .withEncryptionKeyMetadata(getKeyMetadata()) + .withPartitionPath(getPartitionPath()); + + switch (getContentType()) { + case POSITION_DELETES: + deleteFileBuilder = deleteFileBuilder.ofPositionDeletes(); + break; + case EQUALITY_DELETES: + List fieldIds = getEqualityFieldIds(); + int[] equalityFieldIds = new int[fieldIds != null ? fieldIds.size() : 0]; + if (fieldIds != null) { + for (int i = 0; i < fieldIds.size(); i++) { + equalityFieldIds[i] = fieldIds.get(i); + } + } + SortOrder sortOrder = SortOrder.unsorted(); + if (sortOrders != null) { + sortOrder = + checkStateNotNull( + sortOrders.get(getSortOrderId()), + "This DeleteFile was originally created with sort order id '%s', " + + "but table only has sort order ids: %s.", + getSortOrderId(), + sortOrders.keySet()); + } + deleteFileBuilder = + deleteFileBuilder.ofEqualityDeletes(equalityFieldIds).withSortOrder(sortOrder); + break; + default: + throw new IllegalStateException( + "Unexpected content type for DeleteFile: " + getContentType()); + } + + // needed for puffin files + if (getFileFormat().equalsIgnoreCase(FileFormat.PUFFIN.name())) { + deleteFileBuilder = + deleteFileBuilder + .withContentOffset(checkStateNotNull(getContentOffset())) + .withContentSizeInBytes(checkStateNotNull(getContentSizeInBytes())) + .withReferencedDataFile(checkStateNotNull(getReferencedDataFile())); + } + return deleteFileBuilder.build(); + } + + @Override + public final boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof SerializableDeleteFile)) { + return false; + } + SerializableDeleteFile that = (SerializableDeleteFile) o; + return getContentType().equals(that.getContentType()) + && getLocation().equals(that.getLocation()) + && getFileFormat().equals(that.getFileFormat()) + && getRecordCount() == that.getRecordCount() + && getFileSizeInBytes() == that.getFileSizeInBytes() + && getPartitionPath().equals(that.getPartitionPath()) + && getPartitionSpecId() == that.getPartitionSpecId() + && Objects.equals(getSortOrderId(), that.getSortOrderId()) + && Objects.equals(getEqualityFieldIds(), that.getEqualityFieldIds()) + && Objects.equals(getKeyMetadata(), that.getKeyMetadata()) + && Objects.equals(getSplitOffsets(), that.getSplitOffsets()) + && Objects.equals(getColumnSizes(), that.getColumnSizes()) + && Objects.equals(getValueCounts(), that.getValueCounts()) + && Objects.equals(getNullValueCounts(), that.getNullValueCounts()) + && Objects.equals(getNanValueCounts(), that.getNanValueCounts()) + && mapEquals(getLowerBounds(), that.getLowerBounds()) + && mapEquals(getUpperBounds(), that.getUpperBounds()) + && Objects.equals(getContentOffset(), that.getContentOffset()) + && Objects.equals(getContentSizeInBytes(), that.getContentSizeInBytes()) + && Objects.equals(getReferencedDataFile(), that.getReferencedDataFile()) + && Objects.equals(getDataSequenceNumber(), that.getDataSequenceNumber()) + && Objects.equals(getFileSequenceNumber(), that.getFileSequenceNumber()); + } + + @Override + public final int hashCode() { + int hashCode = + Objects.hash( + getContentType(), + getLocation(), + getFileFormat(), + getRecordCount(), + getFileSizeInBytes(), + getPartitionPath(), + getPartitionSpecId(), + getSortOrderId(), + getEqualityFieldIds(), + getKeyMetadata(), + getSplitOffsets(), + getColumnSizes(), + getValueCounts(), + getNullValueCounts(), + getNanValueCounts(), + getContentOffset(), + getContentSizeInBytes(), + getReferencedDataFile(), + getDataSequenceNumber(), + getFileSequenceNumber()); + hashCode = 31 * hashCode + computeMapByteHashCode(getLowerBounds()); + hashCode = 31 * hashCode + computeMapByteHashCode(getUpperBounds()); + return hashCode; + } +} 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 new file mode 100644 index 000000000000..147e2adda1a7 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java @@ -0,0 +1,178 @@ +/* + * 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.List; +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.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.iceberg.ChangelogOperation; +import org.apache.iceberg.types.Types; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Helpers for CDC schemas and output row construction. + * + *

CDC metadata is handled in two phases. Row metadata, such as {@code _row_id} and {@code + * _last_updated_sequence_number}, is added to intermediate read schemas so Iceberg readers can + * populate those values. Commit metadata, such as {@code _commit_snapshot_id} and {@code + * _commit_snapshot_sequence_number}, is carried separately in CDC descriptors. The {@code + * _change_type} metadata column comes from the resolved Beam output {@link ValueKind}. + * + *

The public output shape is assembled only when final Beam {@link Row}s are emitted. This keeps + * the read path table-shaped while still exposing all requested metadata as top-level output + * fields. + */ +final class CdcOutputUtils { + /** + * Returns the public CDC output schema: projected data fields followed by requested metadata + * columns in user-configured order. + */ + static Schema outputSchema(IcebergScanConfig scanConfig, Schema dataSchema) { + if (scanConfig.getMetadataColumns().isEmpty()) { + return dataSchema; + } + + Schema.Builder builder = Schema.builder().addFields(dataSchema.getFields()); + for (String metadataColumn : scanConfig.getMetadataColumns()) { + builder.addField(IcebergCdcMetadataColumns.beamField(metadataColumn)); + } + return builder.build(); + } + + /** + * Returns an Iceberg read schema that includes row metadata columns. + * + *

Commit metadata columns are not added here because Iceberg readers cannot populate them; + * those values are taken from {@link ChangelogDescriptor} or {@link CdcRowDescriptor} when output + * rows are built. + */ + static org.apache.iceberg.Schema readSchemaWithRowMetadata( + List metadataColumns, org.apache.iceberg.Schema dataSchema) { + List fields = new ArrayList<>(dataSchema.columns()); + for (String metadataColumn : metadataColumns) { + Types.NestedField rowMetadataField = + IcebergCdcMetadataColumns.icebergRowMetadataField(metadataColumn); + if (rowMetadataField != null && dataSchema.findField(rowMetadataField.fieldId()) == null) { + fields.add(rowMetadataField); + } + } + return new org.apache.iceberg.Schema(fields, dataSchema.identifierFieldIds()); + } + + /** + * Beam-schema equivalent of {@link #readSchemaWithRowMetadata(List, org.apache.iceberg.Schema)}. + */ + static Schema readBeamSchemaWithRowMetadata(List metadataColumns, Schema dataSchema) { + if (metadataColumns.stream().noneMatch(IcebergCdcMetadataColumns::isRowMetadataColumn)) { + return dataSchema; + } + + Schema.Builder builder = Schema.builder().addFields(dataSchema.getFields()); + for (String metadataColumn : metadataColumns) { + if (IcebergCdcMetadataColumns.isRowMetadataColumn(metadataColumn) + && !dataSchema.hasField(metadataColumn)) { + builder.addField(IcebergCdcMetadataColumns.beamField(metadataColumn)); + } + } + return builder.build(); + } + + /** + * Builds the final public Beam row. + * + *

{@code dataAndRowMetadata} may already include row metadata read from Iceberg. This method + * copies only data fields first, then appends every requested metadata column at the top level. + * That preserves configured column order and avoids exposing row metadata twice. + */ + static Row outputRow( + List metadataColumns, + Schema outputSchema, + long commitSnapshotId, + long snapshotSequentNumber, + ValueKind valueKind, + Row dataAndRowMetadata) { + if (metadataColumns.isEmpty() + || metadataColumns.stream().allMatch(IcebergCdcMetadataColumns::isRowMetadataColumn)) { + return dataAndRowMetadata; + } + + List<@Nullable Object> values = new ArrayList<>(outputSchema.getFieldCount()); + for (Schema.Field field : dataAndRowMetadata.getSchema().getFields()) { + if (!metadataColumns.contains(field.getName())) { + values.add(dataAndRowMetadata.getValue(field.getName())); + } + } + + for (String metadataColumn : metadataColumns) { + values.add( + metadataValue( + metadataColumn, + commitSnapshotId, + snapshotSequentNumber, + valueKind, + dataAndRowMetadata)); + } + return Row.withSchema(outputSchema).addValues(values).build(); + } + + static Schema readBeamSchemaWithRowMetadata( + List metadataColumns, org.apache.iceberg.Schema dataSchema) { + return IcebergUtils.icebergSchemaToBeamSchema( + readSchemaWithRowMetadata(metadataColumns, dataSchema)); + } + + private static @Nullable Object metadataValue( + String metadataColumn, + long commitSnapshotId, + long commitSnapshotSequenceNumber, + ValueKind valueKind, + Row dataAndRowMetadata) { + if (IcebergCdcMetadataColumns.CHANGE_TYPE.equals(metadataColumn)) { + return changelogOperation(valueKind).name(); + } + if (IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_ID.equals(metadataColumn)) { + return commitSnapshotId; + } + if (IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER.equals(metadataColumn)) { + return commitSnapshotSequenceNumber; + } + if (dataAndRowMetadata.getSchema().hasField(metadataColumn)) { + return dataAndRowMetadata.getValue(metadataColumn); + } + return null; + } + + private static ChangelogOperation changelogOperation(ValueKind valueKind) { + switch (valueKind) { + case INSERT: + return ChangelogOperation.INSERT; + case DELETE: + return ChangelogOperation.DELETE; + case UPDATE_BEFORE: + return ChangelogOperation.UPDATE_BEFORE; + case UPDATE_AFTER: + return ChangelogOperation.UPDATE_AFTER; + default: + throw new IllegalArgumentException("Unsupported CDC ValueKind: " + valueKind); + } + } +} 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 new file mode 100644 index 000000000000..daa0a2c73fb8 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java @@ -0,0 +1,698 @@ +/* + * 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.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.ReadUtils; +import org.apache.beam.sdk.io.iceberg.SerializableDeleteFile; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileContent; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; +import org.apache.iceberg.data.BaseDeleteLoader; +import org.apache.iceberg.data.DeleteFilter; +import org.apache.iceberg.data.DeleteLoader; +import org.apache.iceberg.data.InternalRecordWrapper; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.deletes.PositionDeleteIndex; +import org.apache.iceberg.expressions.Expression; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.io.SeekableInputStream; +import org.apache.iceberg.parquet.ParquetMetricsRowGroupFilter; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeSet; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.io.DelegatingSeekableInputStream; +import org.apache.parquet.schema.MessageType; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * 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. + */ +public final class CdcReadUtils { + private static final Logger LOG = LoggerFactory.getLogger(CdcReadUtils.class); + + /** + * Maximum size of an equality delete set to push down as a Parquet residual {@code IN} + * expression. Matches {@link ParquetMetricsRowGroupFilter#IN_PREDICATE_LIMIT}. + */ + private static final int IN_PREDICATE_LIMIT = 200; + + public static CloseableIterable createReader( + SerializableChangelogTask task, + Table table, + IcebergScanConfig scanConfig, + Schema outputSchema) { + return createReader(task, table, scanConfig, outputSchema, Expressions.alwaysTrue()); + } + + /** + * Same as {@link #createReader(SerializableChangelogTask, Table, IcebergScanConfig, Schema)} but + * ANDs {@code extraResidual} into the task's residual expression. The combined expression is + * passed to Iceberg's Parquet reader, which uses it as a row-group-level filter (skips row groups + * whose column statistics cannot match). The caller is still responsible for applying the + * residual at the row level. + * + *

This is used to push extra predicates (e.g. an equality-delete {@code IN} expression) down + * to the reader for cheap row-group skipping. + */ + public static CloseableIterable createReader( + SerializableChangelogTask task, + Table table, + IcebergScanConfig scanConfig, + Schema outputSchema, + Expression extraResidual) { + return createReader( + task, table, scanConfig, outputSchema, extraResidual, task.getStart(), task.getLength()); + } + + /** + * Same as {@link #createReader(SerializableChangelogTask, Table, IcebergScanConfig, Schema, + * Expression)} but reads the byte range {@code [start, start + length)} of the DataFile. + * Iceberg's Parquet reader selects the row groups whose starting offset falls within this range, + * allowing us to prune row-groups by byte-range. + * + *

Callers are responsible for ensuring the requested range stays within the task's assigned + * range, to avoid reading a section that is meant for another worker. + */ + public static CloseableIterable createReader( + SerializableChangelogTask task, + Table table, + IcebergScanConfig scanConfig, + Schema outputSchema, + Expression extraResidual, + long start, + long length) { + Expression baseResidual = task.getExpression(table.schema()); + Expression combined = + extraResidual.op() == Expression.Operation.TRUE + ? baseResidual + : Expressions.and(baseResidual, extraResidual); + return ReadUtils.createReader( + table, + scanConfig, + outputSchema, + checkStateNotNull(table.specs().get(task.getSpecId())), + task.getDataFile().createDataFile(table.specs()), + task.getDataFile().getFileSequenceNumber(), + start, + length, + combined); + } + + /** Returns a filter that skips records marked for deletion. */ + public static DeleteFilter genericDeleteFilter( + Table table, Schema outputSchema, String dataFilePath, List deletes) { + return new GenericDeleteFilter( + table.io(), + dataFilePath, + table.schema(), + outputSchema, + deletes.stream() + .map(sdf -> sdf.createDeleteFile(table.specs(), table.sortOrders())) + .collect(Collectors.toList())); + } + + /** Returns a delete reader that reuses delete structures already loaded by CDC planning. */ + public static DeleteReader genericDeleteReader( + Table table, + Schema outputSchema, + String dataFilePath, + List deletes, + DeleteReader.PreloadedDeletes preloadedDeletes) { + return new GenericDeleteReader( + table.io(), + dataFilePath, + table.schema(), + outputSchema, + deletes.stream() + .map(sdf -> sdf.createDeleteFile(table.specs(), table.sortOrders())) + .collect(Collectors.toList()), + preloadedDeletes); + } + + /** + * Opens the records that a CDC reader should process for a single {@link + * SerializableChangelogTask}, applying the appropriate delete-filter / delete-reader chain for + * the task's type: + * + *

    + *
  • {@code ADDED_ROWS}: Collect and return the records that became live in this commit: + *
      + *
    • 1. Iterate over records in the added DataFile + *
    • 2. Filter out records matched by any added deletes + *
    + *
  • {@code DELETED_ROWS}: Return records in the DataFile that are marked for deletion by new + * DeleteFiles, making sure to first ignore records that have already been marked by + * previous DeleteFiles: + *
      + *
    • 1. Iterate over records in the referenced DataFile + *
    • 2. Filter out records matched from existing deletes. + *
    • 3. Filter out records NOT matched from added deletes + *
    + *
  • {@code DELETED_FILE} — every record in the DataFile that wasn't already deleted by {@code + * existingDeletes}. + *
      + *
    • 1. Iterate over records in the referenced DataFile + *
    • 2. Filter out records matched from existing deletes. + *
    + *
+ * + *

Projection pushdown should not be used when reading bi-directional tasks because we need to + * compare all record columns to accurately identify updates. Otherwise, user-configured + * projection may drop a column that contains real updates. If this happens, the downstream + * resolver will mistakenly determine the (delete, insert) pair to be a duplicate. + * + *

If CDC metadata columns are requested, this method only adds row-sourced metadata columns + * ({@code _row_id}, {@code _last_updated_sequence_number}) to the Iceberg read schema. Changelog + * context columns are added later by {@link CdcOutputUtils#outputRow}. + */ + public static CloseableIterable changelogRecordsForTask( + SerializableChangelogTask task, + Table table, + IcebergScanConfig scanConfig, + boolean useProjectedSchema) { + String dataFilePath = task.getDataFile().getPath(); + Schema outputSchema = + CdcOutputUtils.readSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), + useProjectedSchema ? scanConfig.getRequiredSchema() : table.schema()); + switch (task.getType()) { + case ADDED_ROWS: + DeleteFilter addedDeletesFilter = + genericDeleteFilter(table, outputSchema, dataFilePath, task.getAddedDeletes()); + return addedDeletesFilter.filter( + createReader(task, table, scanConfig, addedDeletesFilter.requiredSchema())); + case DELETED_FILE: + DeleteFilter existingDeletesFilter = + genericDeleteFilter(table, outputSchema, dataFilePath, task.getExistingDeletes()); + return existingDeletesFilter.filter( + createReader(task, table, scanConfig, existingDeletesFilter.requiredSchema())); + case DELETED_ROWS: + return deletedRowsForTask(task, table, scanConfig, outputSchema); + default: + throw new IllegalStateException("Unknown ChangelogScanTask type: " + task.getType()); + } + } + + /** + * Builds the reader chain for a {@code DELETED_ROWS} task with row-group pushdown when possible. + * This helps the reader skip entire row groups. For unskipped row groups, the reader should still + * apply per-record position + equality checks at the row level. + * + *

We use two pushdown strategies, depending on the type of {@link DeleteFile} in the task + * (Position Delete vs. Equality Delete). The two strategies can be combined if both {@link + * DeleteFile} types are present. + * + *

    + *
  1. Byte-range pushdown for Position Deletes: pre-load the {@link + * PositionDeleteIndex}, read the Parquet footer, and compute a single contiguous byte range + * covering the row groups that contain at least one deleted position. + *
  2. IN-expression pushdown for Equality Deletes: build an Iceberg {@code IN} + * expression and pass it as a Parquet residual so the metrics row-group filter can skip + * non-matching row groups. + *
+ * + *

If Position and Equality deletes are both present, both strategies are used to get one + * contiguous range. We read only that range, skipping leading and trailing row groups that + * contain no deletions. + * + *

Note: Equality pushdown is only used when all delete files share a single equality field. + * Multi-column equality requires an exploded OR expression that Parquet's metrics filter handles + * poorly. + */ + private static CloseableIterable deletedRowsForTask( + SerializableChangelogTask task, + Table table, + IcebergScanConfig scanConfig, + Schema outputSchema) { + String dataFilePath = task.getDataFile().getPath(); + List addedDeletes = task.getAddedDeletes(); + + // Split into position vs equality. + List posFiles = new ArrayList<>(); + List eqFiles = new ArrayList<>(); + for (SerializableDeleteFile sd : addedDeletes) { + DeleteFile df = sd.createDeleteFile(table.specs(), table.sortOrders()); + if (df.content() == FileContent.POSITION_DELETES) { + posFiles.add(df); + } else if (df.content() == FileContent.EQUALITY_DELETES) { + eqFiles.add(df); + } + } + + // Strategy 1: byte-range pushdown around row groups with position deletes (+ eq + // matches). + DeleteReader.PreloadedDeletes preloadedDeletes = DeleteReader.PreloadedDeletes.empty(); + if (!posFiles.isEmpty()) { + @Nullable + PositionPushdownResult pushdown = + tryPositionByteRangePushdown( + task, table, scanConfig, outputSchema, posFiles, eqFiles, addedDeletes); + if (pushdown != null) { + if (pushdown.deletedRecords != null) { + return pushdown.deletedRecords; + } + preloadedDeletes = pushdown.preloadedDeletes; + } + // fall through to the default chain on failure + } + + // Strategy 2: equality IN-expression pushdown applied as a reader residual. + // Only safe when no position deletes are present. when both exist, the + // byte-range path above already incorporates the eq filter + Expression eqResidual = Expressions.alwaysTrue(); + if (posFiles.isEmpty() && !eqFiles.isEmpty()) { + EqualityPushdownResult eqPushdown = buildEqualityDeletePushdown(table, eqFiles); + eqResidual = eqPushdown.applicable ? eqPushdown.residual : Expressions.alwaysTrue(); + preloadedDeletes = eqPushdown.preloadedDeletes(null); + } + + DeleteFilter existingDeletesFilter = + genericDeleteFilter(table, outputSchema, dataFilePath, task.getExistingDeletes()); + DeleteReader addedDeletesReader = + genericDeleteReader(table, outputSchema, dataFilePath, addedDeletes, preloadedDeletes); + Schema requiredSchema = + TypeUtil.join(existingDeletesFilter.requiredSchema(), addedDeletesReader.requiredSchema()); + + CloseableIterable records = + createReader(task, table, scanConfig, requiredSchema, eqResidual); + CloseableIterable liveRecords = existingDeletesFilter.filter(records); + return addedDeletesReader.read(liveRecords); + } + + /** + * Path-A byte-range position-delete pushdown. Returns {@code null} if pushdown isn't applicable + * or any step fails, signaling to the caller to fall back. Returns an empty iterable if every row + * group is pruned. + */ + private static @Nullable PositionPushdownResult tryPositionByteRangePushdown( + SerializableChangelogTask task, + Table table, + IcebergScanConfig scanConfig, + Schema outputSchema, + List posFiles, + List eqFiles, + List addedDeletes) { + String dataFilePath = task.getDataFile().getPath(); + + // 1. pre-load the position index for this data file. + PositionDeleteIndex posIndex; + try { + DeleteLoader loader = new BaseDeleteLoader(df -> table.io().newInputFile(df.location())); + posIndex = loader.loadPositionDeletes(posFiles, dataFilePath); + } catch (RuntimeException e) { + LOG.info( + "Failed to pre-load position deletes for {}; falling back to default reader chain.", + dataFilePath, + e); + return null; + } + if (posIndex.isEmpty()) { + // the pos-delete files don't actually target this data file (rare but possible + // after metadata operations). Fall back so the eq pushdown does not run here either. + return PositionPushdownResult.fallback( + DeleteReader.PreloadedDeletes.of(posIndex, Collections.emptyMap())); + } + + // 2. optional equality filter (used to extend the byte range to include row groups + // whose stats match the equality IN values). + @Nullable ParquetMetricsRowGroupFilter eqFilter = null; + EqualityPushdownResult eqPushdown = EqualityPushdownResult.notApplicable(); + if (!eqFiles.isEmpty()) { + eqPushdown = buildEqualityDeletePushdown(table, eqFiles); + if (!eqPushdown.applicable) { + // eq deletes are present but we can't safely identify which row groups they target. + // A narrowed position-only range could drop eq-deleted rows, so fall back to the + // default full-range reader. DeleteReader will still apply residual per record. + return PositionPushdownResult.fallback(eqPushdown.preloadedDeletes(posIndex)); + } + eqFilter = new ParquetMetricsRowGroupFilter(table.schema(), eqPushdown.residual); + } + DeleteReader.PreloadedDeletes preloadedDeletes = eqPushdown.preloadedDeletes(posIndex); + + // 3. read the footer and compute the task byte range covering every row group that + // contains a position delete or matches the eq filter. + long taskStart = task.getStart(); + long taskEnd = taskStart + task.getLength(); + long minStart = Long.MAX_VALUE; + long maxEnd = Long.MIN_VALUE; + + try { + long[] sortedDeletePositions = sortedDeletePositions(posIndex); + InputFile inputFile = table.io().newInputFile(dataFilePath); + try (ParquetFileReader reader = ParquetFileReader.open(asParquetInputFile(inputFile))) { + ParquetMetadata footer = reader.getFooter(); + MessageType parquetSchema = footer.getFileMetaData().getSchema(); + + // track cumulative row count ourselves. not all Parquet writers will include + // it in BlockMetaData.getRowIndexOffset + long cumulativeRows = 0; + for (BlockMetaData rowGroup : footer.getBlocks()) { + long rgStartPos = cumulativeRows; + long rgEndPos = cumulativeRows + rowGroup.getRowCount(); + cumulativeRows = rgEndPos; + + long rgByteStart = rowGroup.getStartingPos(); + long rgByteEnd = rgByteStart + rowGroup.getCompressedSize(); + + // skip row groups outside this task's range. + if (rgByteEnd <= taskStart || rgByteStart >= taskEnd) { + continue; + } + + // if row group has a position and/or an equality delete, include it in the global range + boolean rowGroupHasPosDelete = anyInRange(sortedDeletePositions, rgStartPos, rgEndPos); + boolean rowGroupMatchesEq = + eqFilter != null && eqFilter.shouldRead(parquetSchema, rowGroup); + + if (rowGroupHasPosDelete || rowGroupMatchesEq) { + minStart = Math.min(minStart, rgByteStart); + maxEnd = Math.max(maxEnd, rgByteEnd); + } + } + } + } catch (IOException | RuntimeException e) { + LOG.info( + "Failed to read Parquet footer for {}; falling back to default reader chain.", + dataFilePath, + e); + return PositionPushdownResult.fallback(preloadedDeletes); + } + + long readStart = Math.max(minStart, taskStart); + long readEnd = Math.min(maxEnd, taskEnd); + if (readStart >= readEnd) { + // deletes don't target the portion of the DataFile covered by this read task. + return PositionPushdownResult.of(CloseableIterable.empty(), preloadedDeletes); + } + + // 4. Open the reader with the narrowed byte range. This range represents the union + // of "has position delete" + "matches eq stats" + DeleteFilter existingDeletesFilter = + genericDeleteFilter(table, outputSchema, dataFilePath, task.getExistingDeletes()); + DeleteReader addedDeletesReader = + genericDeleteReader(table, outputSchema, dataFilePath, addedDeletes, preloadedDeletes); + Schema requiredSchema = + TypeUtil.join(existingDeletesFilter.requiredSchema(), addedDeletesReader.requiredSchema()); + CloseableIterable records = + createReader( + task, + table, + scanConfig, + requiredSchema, + Expressions.alwaysTrue(), + readStart, + readEnd - readStart); + CloseableIterable liveRecords = existingDeletesFilter.filter(records); + return PositionPushdownResult.of(addedDeletesReader.read(liveRecords), preloadedDeletes); + } + + /** Materializes a sorted long[] of the positions in {@code posIndex} for binary-search lookup. */ + private static long[] sortedDeletePositions(PositionDeleteIndex posIndex) { + long cardinality = posIndex.cardinality(); + if (cardinality > Integer.MAX_VALUE) { + throw new IllegalStateException( + "Position delete index cardinality exceeds Integer.MAX_VALUE: " + cardinality); + } + long[] arr = new long[(int) cardinality]; + int[] idx = {0}; + posIndex.forEach(p -> arr[idx[0]++] = p); + // forEach is ordered for the bitmap-backed implementation, but the interface doesn't + // promise it, so sort defensively. Cheap relative to the I/O it gates. + Arrays.sort(arr); + return arr; + } + + /** Returns true iff {@code sortedDeletes} contains any value in {@code [start, end)}. */ + private static boolean anyInRange(long[] sortedDeletes, long startInclusive, long endExclusive) { + if (sortedDeletes.length == 0) { + return false; + } + int i = Arrays.binarySearch(sortedDeletes, startInclusive); + if (i < 0) { + i = -i - 1; // insertion point + } + return i < sortedDeletes.length && sortedDeletes[i] < endExclusive; + } + + /** + * Returns an {@code IN} expression suitable as a Parquet residual for the given equality-delete + * files, or {@link Expressions#alwaysTrue()} if pushdown is not applicable. See {@link + * #deletedRowsForTask} for the applicability rules. + */ + private static EqualityPushdownResult buildEqualityDeletePushdown( + Table table, List eqFiles) { + // All eq delete files in this task must share a single equality field id. + Set sharedIds = null; + for (DeleteFile df : eqFiles) { + Set ids = new HashSet<>(df.equalityFieldIds()); + if (sharedIds == null) { + sharedIds = ids; + } else if (!sharedIds.equals(ids)) { + return EqualityPushdownResult.notApplicable(); + } + } + if (sharedIds == null || sharedIds.size() != 1) { + return EqualityPushdownResult.notApplicable(); + } + + int fieldId = Iterables.getOnlyElement(sharedIds); + Types.NestedField field = table.schema().findField(fieldId); + if (field == null) { + return EqualityPushdownResult.notApplicable(); + } + Schema deleteSchema = TypeUtil.select(table.schema(), sharedIds); + + DeleteLoader loader = new BaseDeleteLoader(df -> table.io().newInputFile(df.location())); + StructLikeSet set; + try { + set = loader.loadEqualityDeletes(eqFiles, deleteSchema); + } catch (RuntimeException e) { + LOG.info( + "Failed to pre-load equality deletes for pushdown; falling back to per-record check.", e); + return EqualityPushdownResult.notApplicable(); + } + + Map, StructLikeSet> preloadedSets = new HashMap<>(); + preloadedSets.put(sharedIds, set); + + if (set.size() > IN_PREDICATE_LIMIT) { + return EqualityPushdownResult.notApplicable(preloadedSets); + } + Class javaClass = field.type().typeId().javaClass(); + List values = new ArrayList<>(set.size()); + for (StructLike s : set) { + @Nullable Object v = s.get(0, javaClass); + if (v == null) { + // Nulls don't match an IN-expression. pushing down would drop those deletions. + return EqualityPushdownResult.notApplicable(preloadedSets); + } + values.add(v); + } + if (values.isEmpty()) { + return EqualityPushdownResult.notApplicable(preloadedSets); + } + return EqualityPushdownResult.applicable(Expressions.in(field.name(), values), preloadedSets); + } + + private static final class PositionPushdownResult { + private final @Nullable CloseableIterable deletedRecords; + private final DeleteReader.PreloadedDeletes preloadedDeletes; + + private static PositionPushdownResult of( + CloseableIterable deletedRecords, DeleteReader.PreloadedDeletes preloadedDeletes) { + return new PositionPushdownResult(deletedRecords, preloadedDeletes); + } + + private static PositionPushdownResult fallback(DeleteReader.PreloadedDeletes preloadedDeletes) { + return new PositionPushdownResult(null, preloadedDeletes); + } + + private PositionPushdownResult( + @Nullable CloseableIterable records, + DeleteReader.PreloadedDeletes preloadedDeletes) { + this.deletedRecords = records; + this.preloadedDeletes = preloadedDeletes; + } + } + + private static final class EqualityPushdownResult { + private static final EqualityPushdownResult NOT_APPLICABLE = + new EqualityPushdownResult(Expressions.alwaysTrue(), Collections.emptyMap(), false); + + private final Expression residual; + private final Map, StructLikeSet> preloadedSets; + private final boolean applicable; + + private static EqualityPushdownResult applicable( + Expression residual, Map, StructLikeSet> preloadedSets) { + return new EqualityPushdownResult(residual, preloadedSets, true); + } + + private static EqualityPushdownResult notApplicable() { + return NOT_APPLICABLE; + } + + private static EqualityPushdownResult notApplicable( + Map, StructLikeSet> preloadedSets) { + if (preloadedSets.isEmpty()) { + return NOT_APPLICABLE; + } + return new EqualityPushdownResult(Expressions.alwaysTrue(), preloadedSets, false); + } + + private EqualityPushdownResult( + Expression residual, Map, StructLikeSet> preloadedSets, boolean applicable) { + this.residual = residual; + this.preloadedSets = preloadedSets; + this.applicable = applicable; + } + + private DeleteReader.PreloadedDeletes preloadedDeletes( + @Nullable PositionDeleteIndex positionDeleteIndex) { + return DeleteReader.PreloadedDeletes.of(positionDeleteIndex, preloadedSets); + } + } + + public static class GenericDeleteFilter extends DeleteFilter { + private final FileIO io; + private final InternalRecordWrapper asStructLike; + + @SuppressWarnings("method.invocation") + public GenericDeleteFilter( + FileIO io, + String dataFilePath, + Schema tableSchema, + Schema requiredSchema, + List deleteFiles) { + super(dataFilePath, deleteFiles, tableSchema, requiredSchema); + this.io = io; + this.asStructLike = new InternalRecordWrapper(requiredSchema().asStruct()); + } + + @Override + protected StructLike asStructLike(Record record) { + return asStructLike.wrap(record); + } + + @Override + protected InputFile getInputFile(String location) { + return io.newInputFile(location); + } + } + + public static class GenericDeleteReader extends DeleteReader { + private final FileIO io; + private final InternalRecordWrapper asStructLike; + + @SuppressWarnings("method.invocation") + public GenericDeleteReader( + FileIO io, + String dataFilePath, + Schema tableSchema, + Schema requiredSchema, + List deleteFiles, + DeleteReader.PreloadedDeletes preloadedDeletes) { + super(dataFilePath, deleteFiles, tableSchema, requiredSchema, true, preloadedDeletes); + this.io = io; + this.asStructLike = new InternalRecordWrapper(requiredSchema().asStruct()); + } + + @Override + protected StructLike asStructLike(Record record) { + return asStructLike.wrap(record); + } + + @Override + protected InputFile getInputFile(String location) { + return io.newInputFile(location); + } + } + + /** + * Adapter from Iceberg's {@link InputFile} to Parquet's {@link org.apache.parquet.io.InputFile}, + * for callers that need to open a Parquet file directly (e.g. to read the footer for row-group + * pruning decisions). Iceberg has an equivalent internal {@code ParquetIO} but it's + * package-private. + */ + public static org.apache.parquet.io.InputFile asParquetInputFile(InputFile icebergFile) { + return new IcebergParquetInputFile(icebergFile); + } + + private static final class IcebergParquetInputFile implements org.apache.parquet.io.InputFile { + private final InputFile delegate; + + IcebergParquetInputFile(InputFile delegate) { + this.delegate = delegate; + } + + @Override + public long getLength() { + return delegate.getLength(); + } + + @Override + public org.apache.parquet.io.SeekableInputStream newStream() { + return new IcebergParquetSeekableStream(delegate.newStream()); + } + } + + private static final class IcebergParquetSeekableStream extends DelegatingSeekableInputStream { + private final SeekableInputStream delegate; + + IcebergParquetSeekableStream(SeekableInputStream delegate) { + super(delegate); + this.delegate = delegate; + } + + @Override + public long getPos() throws java.io.IOException { + return delegate.getPos(); + } + + @Override + public void seek(long newPos) throws java.io.IOException { + delegate.seek(newPos); + } + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/DeleteReader.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/DeleteReader.java new file mode 100644 index 000000000000..e1b9a9c98583 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/DeleteReader.java @@ -0,0 +1,309 @@ +/* + * 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.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Predicate; +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.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Multimap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Multimaps; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; +import org.apache.iceberg.Accessor; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.data.BaseDeleteLoader; +import org.apache.iceberg.data.DeleteLoader; +import org.apache.iceberg.deletes.PositionDeleteIndex; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeSet; +import org.apache.iceberg.util.StructProjection; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Reads a {@link org.apache.iceberg.DataFile} and returns records marked deleted by the given + * {@link DeleteFile}s. + * + *

This is mostly a copy of {@link org.apache.iceberg.data.DeleteFilter}, but flipping the logic + * to output deleted records instead of filtering them out. + */ +public abstract class DeleteReader { + private static final Logger LOG = LoggerFactory.getLogger(DeleteReader.class); + + private final String filePath; + private final List posDeletes; + private final List eqDeletes; + private final PreloadedDeletes preloadedDeletes; + private final Schema requiredSchema; + private final Accessor posAccessor; + private volatile @Nullable DeleteLoader deleteLoader = null; + private @Nullable PositionDeleteIndex deleteRowPositions = null; + private @Nullable List> isInDeleteSets = null; + + protected DeleteReader( + String filePath, + List deletes, + Schema tableSchema, + Schema expectedSchema, + boolean needRowPosCol, + PreloadedDeletes preloadedDeletes) { + this.filePath = filePath; + this.preloadedDeletes = preloadedDeletes; + + ImmutableList.Builder posDeleteBuilder = ImmutableList.builder(); + ImmutableList.Builder eqDeleteBuilder = ImmutableList.builder(); + for (DeleteFile delete : deletes) { + switch (delete.content()) { + case POSITION_DELETES: + LOG.debug("Adding position delete file {} to reader", delete.location()); + posDeleteBuilder.add(delete); + break; + case EQUALITY_DELETES: + LOG.debug("Adding equality delete file {} to reader", delete.location()); + eqDeleteBuilder.add(delete); + break; + default: + throw new UnsupportedOperationException( + "Unknown delete file content: " + delete.content()); + } + } + + this.posDeletes = posDeleteBuilder.build(); + this.eqDeletes = eqDeleteBuilder.build(); + this.requiredSchema = + fileProjection(tableSchema, expectedSchema, posDeletes, eqDeletes, needRowPosCol); + this.posAccessor = requiredSchema.accessorForField(MetadataColumns.ROW_POSITION.fieldId()); + } + + public Schema requiredSchema() { + return requiredSchema; + } + + protected abstract StructLike asStructLike(T record); + + protected abstract InputFile getInputFile(String location); + + protected InputFile loadInputFile(DeleteFile deleteFile) { + return getInputFile(deleteFile.location()); + } + + protected long pos(T record) { + return (Long) posAccessor.get(asStructLike(record)); + } + + protected DeleteLoader newDeleteLoader() { + return new BaseDeleteLoader(this::loadInputFile); + } + + private DeleteLoader deleteLoader() { + if (deleteLoader == null) { + synchronized (this) { + if (deleteLoader == null) { + this.deleteLoader = newDeleteLoader(); + } + } + } + + return deleteLoader; + } + + /** + * Returns records that are deleted by either the position deletes or the equality + * deletes attached to this reader — i.e. the union of the two delete predicates. + * + *

Each delete-type predicate is built independently and defaults to "false" (no contribution + * to the union) when its side has no delete files. Both predicates are then OR-combined and + * applied in a single pass over {@code records}. This guarantees that: + * + *

    + *
  • A task with only position deletes emits all records whose position is in the index. + *
  • A task with only equality deletes emits all records matching any equality delete value. + *
  • A task with both emits the union of the two (without duplication). + *
+ */ + public CloseableIterable read(CloseableIterable records) { + Predicate isPosDeleted = + posDeletes.isEmpty() ? t -> false : positionDeletePredicate(deletedRowPositions()); + Predicate isEqDeleted = applyEqDeletes().stream().reduce(Predicate::or).orElse(t -> false); + return CloseableIterable.filter(records, isPosDeleted.or(isEqDeleted)); + } + + private Predicate positionDeletePredicate(PositionDeleteIndex positionIndex) { + return record -> positionIndex.isDeleted(pos(record)); + } + + private List> applyEqDeletes() { + if (isInDeleteSets != null) { + return isInDeleteSets; + } + + isInDeleteSets = Lists.newArrayList(); + if (eqDeletes.isEmpty()) { + return isInDeleteSets; + } + + Multimap, DeleteFile> filesByDeleteIds = + Multimaps.newMultimap(Maps.newHashMap(), Lists::newArrayList); + for (DeleteFile delete : eqDeletes) { + filesByDeleteIds.put(Sets.newHashSet(delete.equalityFieldIds()), delete); + } + + for (Map.Entry, Collection> entry : + filesByDeleteIds.asMap().entrySet()) { + Set ids = entry.getKey(); + Iterable deletes = entry.getValue(); + + Schema deleteSchema = TypeUtil.select(requiredSchema, ids); + + // a projection to select and reorder fields of the file schema to match the delete rows + StructProjection projectRow = StructProjection.create(requiredSchema, deleteSchema); + + StructLikeSet deleteSet = preloadedDeletes.equalityDeleteSet(ids); + if (deleteSet == null) { + deleteSet = deleteLoader().loadEqualityDeletes(deletes, deleteSchema); + } + StructLikeSet deleteSetForPredicate = deleteSet; + Predicate isInDeleteSet = + record -> deleteSetForPredicate.contains(projectRow.wrap(asStructLike(record))); + checkStateNotNull(isInDeleteSets).add(isInDeleteSet); + } + + return checkStateNotNull(isInDeleteSets); + } + + public PositionDeleteIndex deletedRowPositions() { + if (deleteRowPositions == null) { + deleteRowPositions = preloadedDeletes.positionDeleteIndex(); + if (deleteRowPositions == null && !posDeletes.isEmpty()) { + deleteRowPositions = deleteLoader().loadPositionDeletes(posDeletes, filePath); + } + } + + return checkStateNotNull(deleteRowPositions); + } + + /** Delete data already loaded by a planning/pushdown path for one task read. */ + public static final class PreloadedDeletes { + private static final PreloadedDeletes EMPTY = + new PreloadedDeletes(null, Collections.emptyMap()); + + private final @Nullable PositionDeleteIndex positionDeleteIndex; + private final Map, StructLikeSet> equalityDeleteSets; + + public static PreloadedDeletes empty() { + return EMPTY; + } + + public static PreloadedDeletes of( + @Nullable PositionDeleteIndex positionDeleteIndex, + Map, StructLikeSet> equalityDeleteSets) { + if (positionDeleteIndex == null && equalityDeleteSets.isEmpty()) { + return EMPTY; + } + return new PreloadedDeletes(positionDeleteIndex, equalityDeleteSets); + } + + private PreloadedDeletes( + @Nullable PositionDeleteIndex positionDeleteIndex, + Map, StructLikeSet> equalityDeleteSets) { + this.positionDeleteIndex = positionDeleteIndex; + Map, StructLikeSet> copied = new HashMap<>(); + for (Map.Entry, StructLikeSet> entry : equalityDeleteSets.entrySet()) { + copied.put(Collections.unmodifiableSet(Sets.newHashSet(entry.getKey())), entry.getValue()); + } + this.equalityDeleteSets = Collections.unmodifiableMap(copied); + } + + public @Nullable PositionDeleteIndex positionDeleteIndex() { + return positionDeleteIndex; + } + + public @Nullable StructLikeSet equalityDeleteSet(Set equalityFieldIds) { + return equalityDeleteSets.get(equalityFieldIds); + } + } + + private static Schema fileProjection( + Schema tableSchema, + Schema requestedSchema, + List posDeletes, + List eqDeletes, + boolean needRowPosCol) { + if (posDeletes.isEmpty() && eqDeletes.isEmpty()) { + return requestedSchema; + } + + Set requiredIds = Sets.newLinkedHashSet(); + if (needRowPosCol && !posDeletes.isEmpty()) { + requiredIds.add(MetadataColumns.ROW_POSITION.fieldId()); + } + + for (DeleteFile eqDelete : eqDeletes) { + requiredIds.addAll(eqDelete.equalityFieldIds()); + } + + Set missingIds = + Sets.newLinkedHashSet( + Sets.difference(requiredIds, TypeUtil.getProjectedIds(requestedSchema))); + + if (missingIds.isEmpty()) { + return requestedSchema; + } + + // TODO: support adding nested columns. this will currently fail when finding nested columns to + // add + List columns = Lists.newArrayList(requestedSchema.columns()); + for (int fieldId : missingIds) { + if (fieldId == MetadataColumns.ROW_POSITION.fieldId() + || fieldId == MetadataColumns.IS_DELETED.fieldId()) { + continue; // add _pos and _deleted at the end + } + + Types.NestedField field = tableSchema.asStruct().field(fieldId); + Preconditions.checkArgument(field != null, "Cannot find required field for ID %s", fieldId); + + columns.add(field); + } + + if (missingIds.contains(MetadataColumns.ROW_POSITION.fieldId())) { + columns.add(MetadataColumns.ROW_POSITION); + } + + if (missingIds.contains(MetadataColumns.IS_DELETED.fieldId())) { + columns.add(MetadataColumns.IS_DELETED); + } + + return new Schema(columns); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IcebergCdcMetadataColumns.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IcebergCdcMetadataColumns.java new file mode 100644 index 000000000000..407934fb188d --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IcebergCdcMetadataColumns.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 org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.schemas.Schema; +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.ImmutableSet; +import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.types.Types; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Supported top-level metadata columns for Beam Iceberg CDC reads. + * + *

The supported columns come from two sources: + * + *

    + *
  • Iceberg row metadata: {@code _row_id} and {@code _last_updated_sequence_number}. These are + * requested from the physical Iceberg reader and are only available for row-lineage tables + * (v3+). + *
  • Changelog context metadata: {@code _change_type}, {@code _commit_snapshot_id}, and {@code + * _commit_snapshot_sequence_number}. These are known from the changelog snapshot/task context + * and are appended when Beam output rows are built. + *
+ */ +@Internal +public final class IcebergCdcMetadataColumns { + public static final String CHANGE_TYPE = MetadataColumns.CHANGE_TYPE.name(); + public static final String COMMIT_SNAPSHOT_SEQUENCE_NUMBER = "_commit_snapshot_sequence_number"; + public static final String COMMIT_SNAPSHOT_ID = MetadataColumns.COMMIT_SNAPSHOT_ID.name(); + public static final String ROW_ID = MetadataColumns.ROW_ID.name(); + public static final String LAST_UPDATED_SEQUENCE_NUMBER = + MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name(); + + public static final ImmutableList SUPPORTED_COLUMNS = + ImmutableList.of( + CHANGE_TYPE, + COMMIT_SNAPSHOT_ID, + COMMIT_SNAPSHOT_SEQUENCE_NUMBER, + ROW_ID, + LAST_UPDATED_SEQUENCE_NUMBER); + + private static final ImmutableSet ROW_METADATA_COLUMNS = + ImmutableSet.of(ROW_ID, LAST_UPDATED_SEQUENCE_NUMBER); + + public static boolean isSupportedColumn(String name) { + return SUPPORTED_COLUMNS.contains(name); + } + + public static boolean isRowMetadataColumn(String name) { + return ROW_METADATA_COLUMNS.contains(name); + } + + public static Schema.Field beamField(String name) { + if (CHANGE_TYPE.equals(name)) { + return Schema.Field.of(name, Schema.FieldType.STRING); + } + if (COMMIT_SNAPSHOT_ID.equals(name) || COMMIT_SNAPSHOT_SEQUENCE_NUMBER.equals(name)) { + return Schema.Field.of(name, Schema.FieldType.INT64); + } + if (ROW_ID.equals(name) || LAST_UPDATED_SEQUENCE_NUMBER.equals(name)) { + return Schema.Field.nullable(name, Schema.FieldType.INT64); + } + throw new IllegalArgumentException("Unsupported CDC metadata column: " + name); + } + + /** Returns the Iceberg reader field for row-sourced metadata, or null for commit metadata. */ + public static Types.@Nullable NestedField icebergRowMetadataField(String name) { + if (ROW_ID.equals(name)) { + return MetadataColumns.ROW_ID; + } + if (LAST_UPDATED_SEQUENCE_NUMBER.equals(name)) { + return MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER; + } + return null; + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/SerializableChangelogTask.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/SerializableChangelogTask.java new file mode 100644 index 000000000000..9b6955d9e4a5 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/SerializableChangelogTask.java @@ -0,0 +1,282 @@ +/* + * 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 static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + +import com.google.auto.value.AutoValue; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.SerializableDataFile; +import org.apache.beam.sdk.io.iceberg.SerializableDeleteFile; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.NoSuchSchemaException; +import org.apache.beam.sdk.schemas.SchemaCoder; +import org.apache.beam.sdk.schemas.SchemaRegistry; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; +import org.apache.beam.sdk.schemas.annotations.SchemaIgnore; +import org.apache.iceberg.AddedRowsScanTask; +import org.apache.iceberg.ChangelogOperation; +import org.apache.iceberg.ChangelogScanTask; +import org.apache.iceberg.ContentScanTask; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.DeletedDataFileScanTask; +import org.apache.iceberg.DeletedRowsScanTask; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.expressions.Expression; +import org.apache.iceberg.expressions.ExpressionParser; + +@DefaultSchema(AutoValueSchema.class) +@AutoValue +public abstract class SerializableChangelogTask { + public enum Type { + ADDED_ROWS, + DELETED_ROWS, + DELETED_FILE + } + + public static SchemaCoder coder() { + try { + return SchemaRegistry.createDefault().getSchemaCoder(SerializableChangelogTask.class); + } catch (NoSuchSchemaException e) { + throw new RuntimeException(e); + } + } + + public static SerializableChangelogTask.Builder builder() { + return new AutoValue_SerializableChangelogTask.Builder() + .setExistingDeletes(Collections.emptyList()) + .setAddedDeletes(Collections.emptyList()); + } + + @SchemaFieldNumber("0") + public abstract Type getType(); + + @SchemaFieldNumber("1") + public abstract SerializableDataFile getDataFile(); + + @SchemaFieldNumber("2") + public abstract List getExistingDeletes(); + + @SchemaFieldNumber("3") + public abstract List getAddedDeletes(); + + @SchemaFieldNumber("4") + public abstract int getSpecId(); + + @SchemaFieldNumber("5") + public abstract ChangelogOperation getOperation(); + + @SchemaFieldNumber("6") + public abstract int getOrdinal(); + + @SchemaFieldNumber("7") + public abstract long getCommitSnapshotId(); + + @SchemaFieldNumber("8") + public abstract long getStart(); + + @SchemaFieldNumber("9") + public abstract long getLength(); + + @SchemaFieldNumber("10") + public abstract String getJsonExpression(); + + @SchemaIgnore + public Expression getExpression(Schema schema) { + return ExpressionParser.fromJson(getJsonExpression(), schema); + } + + public abstract Builder toBuilder(); + + @AutoValue.Builder + public abstract static class Builder { + abstract Builder setType(Type type); + + abstract Builder setDataFile(SerializableDataFile dataFile); + + @SchemaIgnore + public Builder setDataFile(DataFile df, String partitionPath, boolean includeMetrics) { + return setDataFile(SerializableDataFile.from(df, partitionPath, includeMetrics)); + } + + abstract Builder setExistingDeletes(List existingDeletes); + + abstract Builder setAddedDeletes(List addedDeletes); + + abstract Builder setSpecId(int specId); + + abstract Builder setOperation(ChangelogOperation operation); + + abstract Builder setOrdinal(int ordinal); + + abstract Builder setCommitSnapshotId(long commitSnapshotId); + + abstract Builder setStart(long start); + + abstract Builder setLength(long length); + + abstract Builder setJsonExpression(String expression); + + abstract SerializableChangelogTask build(); + } + + public static SerializableChangelogTask from( + ChangelogScanTask task, Map specs) { + return from(task, specs, false); + } + + public static SerializableChangelogTask from( + ChangelogScanTask task, Map specs, boolean includeMetrics) { + checkState( + task instanceof ContentScanTask, "Expected ChangelogScanTask to also be a ContentScanTask"); + ContentScanTask contentScanTask = (ContentScanTask) task; + PartitionSpec spec = contentScanTask.spec(); + SerializableChangelogTask.Builder builder = + SerializableChangelogTask.builder() + .setOperation(task.operation()) + .setOrdinal(task.changeOrdinal()) + .setCommitSnapshotId(task.commitSnapshotId()) + .setDataFile( + contentScanTask.file(), + spec.partitionToPath(contentScanTask.partition()), + includeMetrics) + .setSpecId(spec.specId()) + .setStart(contentScanTask.start()) + .setLength(contentScanTask.length()) + .setJsonExpression(ExpressionParser.toJson(contentScanTask.residual())); + + if (task instanceof AddedRowsScanTask) { + AddedRowsScanTask addedRowsTask = (AddedRowsScanTask) task; + builder = + builder + .setType(Type.ADDED_ROWS) + .setAddedDeletes( + toSerializableDeletes(addedRowsTask.deletes(), specs, includeMetrics)); + } else if (task instanceof DeletedRowsScanTask) { + DeletedRowsScanTask deletedRowsTask = (DeletedRowsScanTask) task; + builder = + builder + .setType(Type.DELETED_ROWS) + .setAddedDeletes( + toSerializableDeletes(deletedRowsTask.addedDeletes(), specs, includeMetrics)) + .setExistingDeletes( + toSerializableDeletes(deletedRowsTask.existingDeletes(), specs, includeMetrics)); + } else if (task instanceof DeletedDataFileScanTask) { + DeletedDataFileScanTask deletedFileTask = (DeletedDataFileScanTask) task; + builder = + builder + .setType(Type.DELETED_FILE) + .setExistingDeletes( + toSerializableDeletes(deletedFileTask.existingDeletes(), specs, includeMetrics)); + } else { + throw new IllegalStateException("Unknown ChangelogScanTask type: " + task.getClass()); + } + return builder.build(); + } + + static Type getType(ChangelogScanTask task) { + if (task instanceof AddedRowsScanTask) { + return Type.ADDED_ROWS; + } else if (task instanceof DeletedRowsScanTask) { + return Type.DELETED_ROWS; + } else if (task instanceof DeletedDataFileScanTask) { + return Type.DELETED_FILE; + } else { + throw new IllegalStateException("Unknown ChangelogScanTask type: " + task.getClass()); + } + } + + static long getTotalLength(List tasks) { + return tasks.stream().mapToLong(SerializableChangelogTask::getLength).sum(); + } + + static long getLength(ChangelogScanTask task) { + if (task instanceof AddedRowsScanTask) { + return ((AddedRowsScanTask) task).length(); + } else if (task instanceof DeletedRowsScanTask) { + return ((DeletedRowsScanTask) task).length(); + } else if (task instanceof DeletedDataFileScanTask) { + return ((DeletedDataFileScanTask) task).length(); + } + throw new IllegalStateException("Unknown ChangelogScanTask type: " + task.getClass()); + } + + static StructLike getPartition(ChangelogScanTask task) { + if (task instanceof AddedRowsScanTask) { + return ((AddedRowsScanTask) task).partition(); + } else if (task instanceof DeletedRowsScanTask) { + return ((DeletedRowsScanTask) task).partition(); + } else if (task instanceof DeletedDataFileScanTask) { + return ((DeletedDataFileScanTask) task).partition(); + } + throw new IllegalStateException("Unknown ChangelogScanTask type: " + task.getClass()); + } + + static PartitionSpec getSpec(ChangelogScanTask task) { + if (task instanceof AddedRowsScanTask) { + return ((AddedRowsScanTask) task).spec(); + } else if (task instanceof DeletedRowsScanTask) { + return ((DeletedRowsScanTask) task).spec(); + } else if (task instanceof DeletedDataFileScanTask) { + return ((DeletedDataFileScanTask) task).spec(); + } + throw new IllegalStateException("Unknown ChangelogScanTask type: " + task.getClass()); + } + + static DataFile getDataFile(ChangelogScanTask task) { + if (task instanceof AddedRowsScanTask) { + return ((AddedRowsScanTask) task).file(); + } else if (task instanceof DeletedRowsScanTask) { + return ((DeletedRowsScanTask) task).file(); + } else if (task instanceof DeletedDataFileScanTask) { + return ((DeletedDataFileScanTask) task).file(); + } + throw new IllegalStateException("Unknown ChangelogScanTask type: " + task.getClass()); + } + + static List getAddedDeleteFiles(ChangelogScanTask task) { + if (task instanceof AddedRowsScanTask) { + return ((AddedRowsScanTask) task).deletes(); + } else if (task instanceof DeletedRowsScanTask) { + return ((DeletedRowsScanTask) task).addedDeletes(); + } else if (task instanceof DeletedDataFileScanTask) { + return Collections.emptyList(); + } + throw new IllegalStateException("Unknown ChangelogScanTask type: " + task.getClass()); + } + + private static List toSerializableDeletes( + List dfs, Map specs, boolean includeMetrics) { + return dfs.stream() + .map( + df -> + SerializableDeleteFile.from( + df, + checkStateNotNull(specs.get(df.specId())).partitionToPath(df.partition()), + includeMetrics)) + .collect(Collectors.toList()); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/package-info.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/package-info.java new file mode 100644 index 000000000000..8285d91689be --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Iceberg CDC connectors. */ +package org.apache.beam.sdk.io.iceberg.cdc; diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/iceberg/BaseIncrementalChangelogScan.java b/sdks/java/io/iceberg/src/main/java/org/apache/iceberg/BaseIncrementalChangelogScan.java new file mode 100644 index 000000000000..6b12e16690ba --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/iceberg/BaseIncrementalChangelogScan.java @@ -0,0 +1,1014 @@ +/* + * 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.iceberg; + +import java.util.ArrayDeque; +import java.util.Collection; +import java.util.Comparator; +import java.util.Deque; +import java.util.List; +import java.util.Map; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.function.Supplier; +import java.util.stream.Collectors; +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.FluentIterable; +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.Lists; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; +import org.apache.iceberg.ManifestGroup.CreateTasksFunction; +import org.apache.iceberg.ManifestGroup.TaskContext; +import org.apache.iceberg.expressions.Evaluator; +import org.apache.iceberg.expressions.Expression; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.expressions.ManifestEvaluator; +import org.apache.iceberg.expressions.Projections; +import org.apache.iceberg.expressions.ResidualEvaluator; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.util.ContentFileUtil; +import org.apache.iceberg.util.Pair; +import org.apache.iceberg.util.PartitionMap; +import org.apache.iceberg.util.PartitionSet; +import org.apache.iceberg.util.SnapshotUtil; +import org.apache.iceberg.util.SortedMerge; +import org.apache.iceberg.util.TableScanUtil; +import org.apache.iceberg.util.Tasks; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Copied over from Iceberg PR #14264. + */ +@SuppressWarnings("nullness") +public class BaseIncrementalChangelogScan + extends BaseIncrementalScan< + IncrementalChangelogScan, ChangelogScanTask, ScanTaskGroup> + implements IncrementalChangelogScan { + private static final DeleteFileIndex EMPTY = createEmptyInstance(); + + private static DeleteFileIndex createEmptyInstance() { + try { + var constructor = + DeleteFileIndex.class.getDeclaredConstructor( + DeleteFileIndex.EqualityDeletes.class, + PartitionMap.class, + PartitionMap.class, + Map.class, + Map.class); + constructor.setAccessible(true); + return constructor.newInstance(null, null, null, null, null); + } catch (Exception e) { + throw new RuntimeException("Failed to initialize EMPTY DeleteFileIndex", e); + } + } + + private static final Logger LOG = LoggerFactory.getLogger(BaseIncrementalChangelogScan.class); + + public BaseIncrementalChangelogScan(Table table) { + this(table, table.schema(), TableScanContext.empty()); + } + + private BaseIncrementalChangelogScan(Table table, Schema schema, TableScanContext context) { + super(table, schema, context); + } + + @Override + protected IncrementalChangelogScan newRefinedScan( + Table newTable, Schema newSchema, TableScanContext newContext) { + return new BaseIncrementalChangelogScan(newTable, newSchema, newContext); + } + + // Private fields to track build call count and cache (accessed via package-private methods for + // testing) + private int existingDeleteIndexBuildCallCount = 0; + // Cache for the built index (null if not built yet) + private DeleteFileIndex cachedExistingDeleteIndex = null; + + @Override + protected CloseableIterable doPlanFiles( + Long fromSnapshotIdExclusive, long toSnapshotIdInclusive) { + + Deque changelogSnapshots = + orderedChangelogSnapshots(fromSnapshotIdExclusive, toSnapshotIdInclusive); + + if (changelogSnapshots.isEmpty()) { + return CloseableIterable.empty(); + } + + Set changelogSnapshotIds = toSnapshotIds(changelogSnapshots); + + Set newDataManifests = + FluentIterable.from(changelogSnapshots) + .transformAndConcat(snapshot -> snapshot.dataManifests(table().io())) + .filter(manifest -> changelogSnapshotIds.contains(manifest.snapshotId())) + .toSet(); + + // Build per-snapshot delete file indexes for added deletes + Map addedDeletesBySnapshot = buildAddedDeleteIndexes(changelogSnapshots); + + // Check if existing delete index is needed for equality deletes + boolean hasEqualityDeletes = + addedDeletesBySnapshot.values().stream() + .anyMatch(index -> !index.isEmpty() && index.hasEqualityDeletes()); + + // Build existing index early if needed for equality deletes, otherwise use lazy initialization + DeleteFileIndex existingDeleteIndex = + hasEqualityDeletes ? buildExistingDeleteIndexTracked(fromSnapshotIdExclusive) : EMPTY; + + ManifestGroup manifestGroup = + new ManifestGroup(table().io(), newDataManifests, ImmutableList.of()) + .specsById(table().specs()) + .caseSensitive(isCaseSensitive()) + .select(scanColumns()) + .filterData(filter()) + .filterManifestEntries(entry -> changelogSnapshotIds.contains(entry.snapshotId())) + .ignoreExisting() + .columnsToKeepStats(columnsToKeepStats()); + + if (shouldIgnoreResiduals()) { + manifestGroup = manifestGroup.ignoreResiduals(); + } + + if (newDataManifests.size() > 1 && shouldPlanWithExecutor()) { + manifestGroup = manifestGroup.planWith(planExecutor()); + } + + // Create a supplier that reuses already-built index or builds lazily when first DELETED entry + // is encountered + Supplier existingDeleteIndexSupplier = + () -> { + if (cachedExistingDeleteIndex != null) { + return cachedExistingDeleteIndex; + } + return buildExistingDeleteIndexTracked(fromSnapshotIdExclusive); + }; + + // Plan data file tasks (ADDED and DELETED) + Map> cumulativeDeletesMap = + buildCumulativeDeletesBySnapshot(changelogSnapshots, addedDeletesBySnapshot); + + CloseableIterable dataFileTasks = + manifestGroup.plan( + new CreateDataFileChangeTasks( + changelogSnapshots, + existingDeleteIndexSupplier, + addedDeletesBySnapshot, + cumulativeDeletesMap, + table().specs(), + isCaseSensitive())); + + // Find EXISTING data files affected by newly added delete files and create tasks for them + CloseableIterable deletedRowsTasks = + planDeletedRowsTasks( + changelogSnapshots, existingDeleteIndex, addedDeletesBySnapshot, changelogSnapshotIds); + + // Merge tasks from both iterables in order by changeOrdinal + Comparator byOrdinal = + Comparator.comparing(ChangelogScanTask::changeOrdinal) + .thenComparing(ChangelogScanTask::commitSnapshotId); + + return new SortedMerge<>(byOrdinal, ImmutableList.of(dataFileTasks, deletedRowsTasks)); + } + + @Override + public CloseableIterable> planTasks() { + return TableScanUtil.planTaskGroups( + planFiles(), targetSplitSize(), splitLookback(), splitOpenFileCost()); + } + + // builds a collection of changelog snapshots (oldest to newest) + // the order of the snapshots is important as it is used to determine change ordinals + private Deque orderedChangelogSnapshots(Long fromIdExcl, long toIdIncl) { + Deque changelogSnapshots = new ArrayDeque<>(); + + for (Snapshot snapshot : SnapshotUtil.ancestorsBetween(table(), toIdIncl, fromIdExcl)) { + if (!snapshot.operation().equals(DataOperations.REPLACE)) { + changelogSnapshots.addFirst(snapshot); + } + } + + return changelogSnapshots; + } + + private Set toSnapshotIds(Collection snapshots) { + return snapshots.stream().map(Snapshot::snapshotId).collect(Collectors.toSet()); + } + + private static Map computeSnapshotOrdinals(Deque snapshots) { + Map snapshotOrdinals = Maps.newHashMap(); + + int ordinal = 0; + + for (Snapshot snapshot : snapshots) { + snapshotOrdinals.put(snapshot.snapshotId(), ordinal++); + } + + return snapshotOrdinals; + } + + /** + * Builds a delete file index for existing deletes that were present before the start snapshot. + * These deletes should be applied to data files but should not generate DELETE changelog rows. + * Uses manifest pruning and caching to optimize performance. + */ + private DeleteFileIndex buildExistingDeleteIndex(Long fromSnapshotIdExclusive) { + if (fromSnapshotIdExclusive == null) { + return EMPTY; + } + Snapshot fromSnapshot = table().snapshot(fromSnapshotIdExclusive); + Preconditions.checkState( + fromSnapshot != null, "Cannot find starting snapshot: %s", fromSnapshotIdExclusive); + + List existingDeleteManifests = fromSnapshot.deleteManifests(table().io()); + if (existingDeleteManifests.isEmpty()) { + return EMPTY; + } + + // Prune manifests based on partition filter to avoid processing irrelevant manifests + List prunedManifests = pruneManifestsByPartition(existingDeleteManifests); + if (prunedManifests.isEmpty()) { + return EMPTY; + } + + // Load delete files from manifests + Iterable deleteFiles = loadDeleteFiles(prunedManifests, null); + + return DeleteFileIndex.builderFor(deleteFiles) + .specsById(table().specs()) + .caseSensitive(isCaseSensitive()) + .build(); + } + + /** + * Wrapper method that tracks build calls and caches the result for reuse. This ensures we only + * build the index once even if called from multiple places. + */ + private DeleteFileIndex buildExistingDeleteIndexTracked(Long fromSnapshotIdExclusive) { + if (cachedExistingDeleteIndex != null) { + return cachedExistingDeleteIndex; + } + existingDeleteIndexBuildCallCount++; + cachedExistingDeleteIndex = buildExistingDeleteIndex(fromSnapshotIdExclusive); + return cachedExistingDeleteIndex; + } + + // Visible for testing + int getExistingDeleteIndexBuildCallCount() { + return existingDeleteIndexBuildCallCount; + } + + // Visible for testing + boolean wasExistingDeleteIndexBuilt() { + return existingDeleteIndexBuildCallCount > 0; + } + + /** + * Builds per-snapshot delete file indexes for newly added delete files in each changelog + * snapshot. These deletes should generate DELETE changelog rows. Uses caching to avoid re-parsing + * manifests. + */ + private Map buildAddedDeleteIndexes(Deque changelogSnapshots) { + Map addedDeletesBySnapshot = Maps.newConcurrentMap(); + Tasks.foreach(changelogSnapshots) + .retry(3) + .stopOnFailure() + .throwFailureWhenFinished() + .executeWith(planExecutor()) + .onFailure( + (snapshot, exc) -> + LOG.warn( + "Failed to build delete index for snapshot {}", snapshot.snapshotId(), exc)) + .run( + snapshot -> { + List snapshotDeleteManifests = snapshot.deleteManifests(table().io()); + if (snapshotDeleteManifests.isEmpty()) { + addedDeletesBySnapshot.put(snapshot.snapshotId(), EMPTY); + return; + } + + // Filter to only include delete files added in this snapshot + List addedDeleteManifests = + snapshotDeleteManifests.stream() + .filter(manifest -> manifest.snapshotId().equals(snapshot.snapshotId())) + .collect(Collectors.toUnmodifiableList()); + + if (addedDeleteManifests.isEmpty()) { + addedDeletesBySnapshot.put(snapshot.snapshotId(), EMPTY); + } else { + // Load delete files from manifests + Iterable deleteFiles = + loadDeleteFiles(addedDeleteManifests, snapshot.snapshotId()); + + DeleteFileIndex index = + DeleteFileIndex.builderFor(deleteFiles) + .specsById(table().specs()) + .caseSensitive(isCaseSensitive()) + .build(); + addedDeletesBySnapshot.put(snapshot.snapshotId(), index); + } + }); + return addedDeletesBySnapshot; + } + + /** + * Plans tasks for EXISTING data files that are affected by newly added delete files. These files + * were not added or deleted in the changelog snapshot range, but have new delete files applied to + * them. + */ + private CloseableIterable planDeletedRowsTasks( + Deque changelogSnapshots, + DeleteFileIndex existingDeleteIndex, + Map addedDeletesBySnapshot, + Set changelogSnapshotIds) { + + Map snapshotOrdinals = computeSnapshotOrdinals(changelogSnapshots); + List tasks = Lists.newArrayList(); + + // Build a map of file statuses and collect affected partitions for each snapshot + Pair>, PartitionSet> fileStatusAndPartitions = + buildFileStatusBySnapshot(changelogSnapshots, changelogSnapshotIds); + Map> fileStatusBySnapshot = + fileStatusAndPartitions.first(); + PartitionSet affectedPartitions = fileStatusAndPartitions.second(); + + // Accumulate actual DeleteFile entries chronologically + List accumulatedDeletes = Lists.newArrayList(); + + // Start with deletes from before the changelog range + if (!existingDeleteIndex.isEmpty()) { + for (DeleteFile df : existingDeleteIndex.referencedDeleteFiles()) { + accumulatedDeletes.add(df); + } + } + + for (Snapshot snapshot : changelogSnapshots) { + DeleteFileIndex addedDeleteIndex = addedDeletesBySnapshot.get(snapshot.snapshotId()); + if (addedDeleteIndex.isEmpty()) { + continue; + } + + // Collect partitions of newly added delete files for pruning (important for the current + // snapshot) + for (DeleteFile df : addedDeleteIndex.referencedDeleteFiles()) { + affectedPartitions.add(df.specId(), df.partition()); + } + + DeleteFileIndex cumulativeDeleteIndex = + buildDeleteIndex(accumulatedDeletes, affectedPartitions); + + // Process data files for this snapshot + // Use a local set per snapshot to track processed files + Set alreadyProcessedPaths = Sets.newHashSet(); + processSnapshotForDeletedRowsTasks( + snapshot, + addedDeleteIndex, + cumulativeDeleteIndex, + fileStatusBySnapshot.get(snapshot.snapshotId()), + alreadyProcessedPaths, + snapshotOrdinals, + affectedPartitions, + tasks); + + // Accumulate this snapshot's added deletes for subsequent snapshots + for (DeleteFile df : addedDeleteIndex.referencedDeleteFiles()) { + accumulatedDeletes.add(df); + } + } + + return CloseableIterable.withNoopClose(tasks); + } + + /** + * Builds a map of file statuses for each snapshot, tracking which files were added or deleted in + * each snapshot. + */ + private Pair>, PartitionSet> + buildFileStatusBySnapshot( + Deque changelogSnapshots, Set changelogSnapshotIds) { + + Map> fileStatusBySnapshot = Maps.newConcurrentMap(); + java.util.Queue localPartitionsQueue = + new java.util.concurrent.ConcurrentLinkedQueue<>(); + + Tasks.foreach(changelogSnapshots) + .stopOnFailure() + .throwFailureWhenFinished() + .executeWith(planExecutor()) + .run( + snapshot -> { + Map fileStatuses = Maps.newHashMap(); + PartitionSet localAffected = PartitionSet.create(table().specs()); + + List changedDataManifests = + FluentIterable.from(snapshot.dataManifests(table().io())) + .filter(manifest -> manifest.snapshotId().equals(snapshot.snapshotId())) + .toList(); + + if (!changedDataManifests.isEmpty()) { + ManifestGroup changedGroup = + new ManifestGroup(table().io(), changedDataManifests, ImmutableList.of()) + .specsById(table().specs()) + .caseSensitive(isCaseSensitive()) + .select(scanColumns()) + .filterData(filter()) + .ignoreExisting() + .columnsToKeepStats(columnsToKeepStats()); + + try (CloseableIterable> entries = changedGroup.entries()) { + for (ManifestEntry entry : entries) { + if (changelogSnapshotIds.contains(entry.snapshotId())) { + fileStatuses.put(entry.file().location(), entry.status()); + localAffected.add(entry.file().specId(), entry.file().partition()); + } + } + } catch (Exception e) { + throw new RuntimeException( + "Failed to collect file statuses for snapshot " + snapshot.snapshotId(), e); + } + } + + fileStatusBySnapshot.put(snapshot.snapshotId(), fileStatuses); + localPartitionsQueue.add(localAffected); + }); + + PartitionSet globalAffected = PartitionSet.create(table().specs()); + for (PartitionSet local : localPartitionsQueue) { + globalAffected.addAll(local); + } + + return Pair.of(fileStatusBySnapshot, globalAffected); + } + + private List pruneManifestsByAffectedPartitions( + List manifests, PartitionSet affectedPartitions) { + if (affectedPartitions.isEmpty()) { + return manifests; + } + + Expression affectedExpr = buildAffectedPartitionExpression(affectedPartitions); + if (affectedExpr == Expressions.alwaysFalse()) { + return manifests; + } + + List pruned = Lists.newArrayList(); + for (ManifestFile manifest : manifests) { + PartitionSpec spec = table().specs().get(manifest.partitionSpecId()); + if (spec == null || spec.isUnpartitioned()) { + pruned.add(manifest); + } else if (manifestOverlapsFilter(manifest, spec, affectedExpr)) { + pruned.add(manifest); + } + } + return pruned; + } + + private Expression buildAffectedPartitionExpression(PartitionSet affectedPartitions) { + Expression combined = null; + + for (Pair pair : affectedPartitions) { + int specId = pair.first(); + StructLike partition = pair.second(); + PartitionSpec spec = table().specs().get(specId); + if (spec == null) { + continue; + } else if (spec.isUnpartitioned()) { + return Expressions.alwaysTrue(); // FALLBACK: Global delete exists, include ALL manifests! + } + + Expression specExpr = null; + for (int i = 0; i < spec.fields().size(); i++) { + org.apache.iceberg.PartitionField field = spec.fields().get(i); + Object value = partition.get(i, Object.class); + if (value != null) { + String columnName = table().schema().findColumnName(field.sourceId()); + if (columnName != null) { + Expression equalExpr = Expressions.equal(columnName, value); + specExpr = (specExpr == null) ? equalExpr : Expressions.and(specExpr, equalExpr); + } + } + } + + if (specExpr != null) { + combined = (combined == null) ? specExpr : Expressions.or(combined, specExpr); + } + } + + return combined != null ? combined : Expressions.alwaysFalse(); + } + + /** + * Builds a map of snapshot ID -> all delete files that were added in the scan range up to that + * snapshot, PRUNING files that were removed in the middle. + */ + private Map> buildCumulativeDeletesBySnapshot( + Deque snapshots, Map addedDeletesBySnapshot) { + Map> result = Maps.newHashMap(); + List accumulatedDeletes = Lists.newArrayList(); + + for (Snapshot snapshot : snapshots) { + // Save state first, so that this snapshot's tasks can use any deletes active up to this point + result.put(snapshot.snapshotId(), Lists.newArrayList(accumulatedDeletes)); + + // Check for removed deletes and prune from accumulatedDeletes for FUTURE snapshots + List changedDeletes = + FluentIterable.from(snapshot.deleteManifests(table().io())) + .filter(manifest -> manifest.snapshotId().equals(snapshot.snapshotId())) + .toList(); + + if (!changedDeletes.isEmpty()) { + Iterable removedDeletes = + loadRemovedDeleteFiles(changedDeletes, snapshot.snapshotId()); + Set removedPaths = Sets.newHashSet(); + for (DeleteFile rdf : removedDeletes) { + removedPaths.add(rdf.location()); + } + accumulatedDeletes.removeIf(df -> removedPaths.contains(df.location())); + } + + // Add new deletes for FUTURE snapshots + DeleteFileIndex addedDeleteIndex = addedDeletesBySnapshot.get(snapshot.snapshotId()); + if (addedDeleteIndex != null && !addedDeleteIndex.isEmpty()) { + for (DeleteFile df : addedDeleteIndex.referencedDeleteFiles()) { + accumulatedDeletes.add(df); + } + } + } + + return result; + } + + /** + * Builds a delete index from the accumulated list of delete files, pruning by affected + * partitions. + */ + private DeleteFileIndex buildDeleteIndex( + List accumulatedDeletes, PartitionSet affectedPartitions) { + if (accumulatedDeletes.isEmpty()) { + return EMPTY; + } + + List filteredDeletes = accumulatedDeletes; + if (!affectedPartitions.isEmpty()) { + filteredDeletes = Lists.newArrayList(); + for (DeleteFile df : accumulatedDeletes) { + PartitionSpec spec = table().specs().get(df.specId()); + if (spec == null || spec.isUnpartitioned()) { + filteredDeletes.add(df); // Always include unpartitioned deletes + } else if (affectedPartitions.contains(df.specId(), df.partition())) { + filteredDeletes.add(df); + } + } + } + + return DeleteFileIndex.builderFor(filteredDeletes) + .specsById(table().specs()) + .caseSensitive(isCaseSensitive()) + .build(); + } + + /** + * Processes data files for a snapshot to create DeletedRowsScanTask for existing files affected + * by new delete files. + */ + private void processSnapshotForDeletedRowsTasks( + Snapshot snapshot, + DeleteFileIndex addedDeleteIndex, + DeleteFileIndex cumulativeDeleteIndex, + Map currentSnapshotFiles, + Set alreadyProcessedPaths, + Map snapshotOrdinals, + PartitionSet affectedPartitions, + List tasks) { + + // Get all data files that exist in this snapshot, pruned by affected partitions + List allDataManifests = snapshot.dataManifests(table().io()); + List prunedManifests = + pruneManifestsByAffectedPartitions(allDataManifests, affectedPartitions); + + ManifestGroup allDataGroup = + new ManifestGroup(table().io(), prunedManifests, ImmutableList.of()) + .specsById(table().specs()) + .caseSensitive(isCaseSensitive()) + .select(scanColumns()) + .filterData(filter()) + .ignoreDeleted() + .columnsToKeepStats(columnsToKeepStats()); + + if (shouldIgnoreResiduals()) { + allDataGroup = allDataGroup.ignoreResiduals(); + } + + String schemaString = SchemaParser.toJson(schema()); + + // Cache per specId - same for all files with same specId + Map specStringCache = Maps.newHashMap(); + Map residualCache = Maps.newHashMap(); + Expression residualFilter = shouldIgnoreResiduals() ? Expressions.alwaysTrue() : filter(); + + try (CloseableIterable> entries = allDataGroup.entries()) { + for (ManifestEntry entry : entries) { + DataFile dataFile = entry.file(); + String filePath = dataFile.location(); + + // Skip if this file was ADDED or DELETED in this snapshot + // (those are handled by CreateDataFileChangeTasks) + if (currentSnapshotFiles.containsKey(filePath)) { + continue; + } + + // Skip if we already created a task for this file in this snapshot + // Note: alreadyProcessedPaths is local to this snapshot's processing + if (alreadyProcessedPaths.contains(filePath)) { + continue; + } + + // Check if this data file is affected by newly added delete files + DeleteFile[] addedDeletes = addedDeleteIndex.forEntry(entry); + if (addedDeletes.length == 0) { + continue; + } + + // This data file was EXISTING but has new delete files applied + // Get existing deletes from before this snapshot (cumulative) + DeleteFile[] existingDeletes = + cumulativeDeleteIndex.isEmpty() + ? new DeleteFile[0] + : cumulativeDeleteIndex.forEntry(entry); + + // Create a DeletedRowsScanTask + int changeOrdinal = snapshotOrdinals.get(snapshot.snapshotId()); + + // Use cached values (calculate once per specId) + int specId = dataFile.specId(); + String specString = + specStringCache.computeIfAbsent( + specId, id -> PartitionSpecParser.toJson(table().specs().get(id))); + ResidualEvaluator residuals = + residualCache.computeIfAbsent( + specId, + id -> { + PartitionSpec spec = table().specs().get(id); + return ResidualEvaluator.of(spec, residualFilter, isCaseSensitive()); + }); + + tasks.add( + new BaseDeletedRowsScanTask( + changeOrdinal, + snapshot.snapshotId(), + dataFile.copy(shouldKeepStats()), + addedDeletes, + existingDeletes, + schemaString, + specString, + residuals)); + + // Mark this file as processed for this snapshot + alreadyProcessedPaths.add(filePath); + } + } catch (Exception e) { + throw new RuntimeException("Failed to plan deleted rows tasks", e); + } + } + + private boolean shouldKeepStats() { + Set columns = columnsToKeepStats(); + return columns != null && !columns.isEmpty(); + } + + /** + * Loads delete files from manifests by parsing each manifest. + * + * @param manifests the delete manifests to load + * @return list of delete files + */ + private Iterable loadDeleteFiles( + List manifests, Long targetSnapshotId) { + Queue allDeleteFiles = new ConcurrentLinkedQueue<>(); + + Tasks.foreach(manifests) + .stopOnFailure() + .throwFailureWhenFinished() + .executeWith(planExecutor()) + .run( + manifest -> { + List deleteFiles = + loadDeleteFilesFromManifest(manifest, targetSnapshotId); + allDeleteFiles.addAll(deleteFiles); + }); + + return allDeleteFiles; + } + + private Iterable loadRemovedDeleteFiles( + List manifests, Long targetSnapshotId) { + Queue allDeleteFiles = new ConcurrentLinkedQueue<>(); + + Tasks.foreach(manifests) + .stopOnFailure() + .throwFailureWhenFinished() + .executeWith(planExecutor()) + .run( + manifest -> { + List deleteFiles = + loadRemovedDeleteFilesFromManifest(manifest, targetSnapshotId); + allDeleteFiles.addAll(deleteFiles); + }); + + return allDeleteFiles; + } + + private List loadRemovedDeleteFilesFromManifest( + ManifestFile manifest, Long targetSnapshotId) { + List deleteFiles = Lists.newArrayList(); + + try (ManifestReader reader = + ManifestFiles.readDeleteManifest(manifest, table().io(), table().specs())) { + for (ManifestEntry entry : reader.entries()) { + if (entry.status() == ManifestEntry.Status.DELETED + && entry.snapshotId().equals(targetSnapshotId)) { + DeleteFile file = entry.file(); + + if (!partitionMatchesFilter(file)) { + continue; + } + + Set columns = + file.content() == FileContent.POSITION_DELETES + ? Set.of(MetadataColumns.DELETE_FILE_PATH.fieldId()) + : Set.copyOf(file.equalityFieldIds()); + deleteFiles.add(ContentFileUtil.copy(file, true, columns)); + } + } + } catch (Exception e) { + throw new RuntimeException("Failed to read delete manifest: " + manifest.path(), e); + } + + return deleteFiles; + } + + /** + * Prunes delete manifests based on partition filter to avoid processing irrelevant manifests. + * This significantly improves performance when only a subset of partitions are relevant to the + * scan. + * + * @param manifests all delete manifests to consider + * @return list of manifests that might contain relevant delete files + */ + private List pruneManifestsByPartition(List manifests) { + Expression currentFilter = filter(); + + // If there's no filter, return all manifests + if (currentFilter == null || currentFilter.equals(Expressions.alwaysTrue())) { + return manifests; + } + + List prunedManifests = Lists.newArrayList(); + + for (ManifestFile manifest : manifests) { + PartitionSpec spec = table().specs().get(manifest.partitionSpecId()); + if (spec == null || spec.isUnpartitioned()) { + // Include unpartitioned manifests + prunedManifests.add(manifest); + } else if (manifestOverlapsFilter(manifest, spec, currentFilter)) { + // Check if manifest partition range overlaps with filter + prunedManifests.add(manifest); + } + } + + return prunedManifests; + } + + /** + * Checks if a manifest's partition range overlaps with the given filter. + * + * @param manifest the manifest to check + * @param spec the partition spec for the manifest + * @param filter the scan filter + * @return true if the manifest might contain matching partitions, false otherwise + */ + private boolean manifestOverlapsFilter( + ManifestFile manifest, PartitionSpec spec, Expression filter) { + try { + // Use inclusive projection to transform row filter to partition filter + Expression partitionFilter = Projections.inclusive(spec, isCaseSensitive()).project(filter); + + // Create evaluator for the partition filter + ManifestEvaluator evaluator = + ManifestEvaluator.forPartitionFilter(partitionFilter, spec, isCaseSensitive()); + + // Check if manifest could contain matching partitions + return evaluator.eval(manifest); + } catch (Exception e) { + // If evaluation fails, be conservative and include the manifest + return true; + } + } + + /** + * Checks if a delete file's partition overlaps with the current scan filter. This enables + * partition pruning to reduce memory footprint and planning overhead by skipping delete files + * that cannot possibly match any rows in the scan. + * + * @param file the delete file to check + * @return true if the delete file's partition might contain matching rows, false otherwise + */ + private boolean partitionMatchesFilter(DeleteFile file) { + // If there's no filter, all partitions match + Expression currentFilter = filter(); + if (currentFilter == null || currentFilter.equals(Expressions.alwaysTrue())) { + return true; + } + + // Get the partition spec for this delete file + PartitionSpec spec = table().specs().get(file.specId()); + if (spec == null || spec.isUnpartitioned()) { + // If spec not found or table is unpartitioned, be conservative and include the file + return true; + } + + try { + // Project the row filter to partition space using inclusive projection + // This transforms expressions on source columns to expressions on partition columns + Expression partitionFilter = + Projections.inclusive(spec, isCaseSensitive()).project(currentFilter); + + // Evaluate the projected filter against the delete file's partition + Evaluator evaluator = new Evaluator(spec.partitionType(), partitionFilter, isCaseSensitive()); + return evaluator.eval(file.partition()); + } catch (Exception e) { + // If evaluation fails, be conservative and include the file + return true; + } + } + + /** + * Loads delete files from a single manifest, parsing the manifest entries. + * + * @param manifest the delete manifest to load + * @return list of delete files from this manifest + */ + private List loadDeleteFilesFromManifest( + ManifestFile manifest, Long targetSnapshotId) { + List deleteFiles = Lists.newArrayList(); + + try (ManifestReader reader = + ManifestFiles.readDeleteManifest(manifest, table().io(), table().specs())) { + for (ManifestEntry entry : reader.entries()) { + if (entry.status() != ManifestEntry.Status.DELETED + && (targetSnapshotId == null || entry.snapshotId().equals(targetSnapshotId))) { + // Only include live delete files, copy with minimal stats to save memory + DeleteFile file = entry.file(); + + // Apply partition pruning - skip delete files that cannot match the scan filter + if (!partitionMatchesFilter(file)) { + continue; + } + + Set columns = + file.content() == FileContent.POSITION_DELETES + ? Set.of(MetadataColumns.DELETE_FILE_PATH.fieldId()) + : Set.copyOf(file.equalityFieldIds()); + deleteFiles.add(ContentFileUtil.copy(file, true, columns)); + } + } + } catch (Exception e) { + throw new RuntimeException("Failed to read delete manifest: " + manifest.path(), e); + } + + return deleteFiles; + } + + private static class CreateDataFileChangeTasks implements CreateTasksFunction { + private static final DeleteFile[] NO_DELETES = new DeleteFile[0]; + + private final Map snapshotOrdinals; + private final Supplier existingDeleteIndexSupplier; + private final Map addedDeletesBySnapshot; + private final Map> cumulativeDeletesMap; + private final Map specsById; + private final boolean caseSensitive; + + CreateDataFileChangeTasks( + Deque snapshots, + Supplier existingDeleteIndexSupplier, + Map addedDeletesBySnapshot, + Map> cumulativeDeletesMap, + Map specsById, + boolean caseSensitive) { + this.snapshotOrdinals = computeSnapshotOrdinals(snapshots); + this.existingDeleteIndexSupplier = existingDeleteIndexSupplier; + this.addedDeletesBySnapshot = addedDeletesBySnapshot; + this.cumulativeDeletesMap = cumulativeDeletesMap; + this.specsById = specsById; + this.caseSensitive = caseSensitive; + } + + @Override + public CloseableIterable apply( + CloseableIterable> entries, TaskContext context) { + + return CloseableIterable.transform( + entries, + entry -> { + long commitSnapshotId = entry.snapshotId(); + int changeOrdinal = snapshotOrdinals.get(commitSnapshotId); + DataFile dataFile = entry.file().copy(context.shouldKeepStats()); + + switch (entry.status()) { + case ADDED: + // For ADDED data files, attach delete files added in this snapshot + DeleteFile[] addedFileDeletes = getDeletesForAddedFile(entry, commitSnapshotId); + return new BaseAddedRowsScanTask( + changeOrdinal, + commitSnapshotId, + dataFile, + addedFileDeletes, + context.schemaAsString(), + context.specAsString(), + context.residuals()); + + case DELETED: + // For DELETED data files, attach ALL deletes that were present up to deletion + // This includes existing deletes AND deletes added in the scan range + DeleteFile[] deletedFileDeletes = getDeletesForDeletedFile(entry, commitSnapshotId); + return new BaseDeletedDataFileScanTask( + changeOrdinal, + commitSnapshotId, + dataFile, + deletedFileDeletes, + context.schemaAsString(), + context.specAsString(), + context.residuals()); + + default: + throw new IllegalArgumentException("Unexpected entry status: " + entry.status()); + } + }); + } + + /** + * Gets delete files that apply to an ADDED data file. Only includes deletes added in the same + * snapshot as the file. + */ + private DeleteFile[] getDeletesForAddedFile( + ManifestEntry entry, long commitSnapshotId) { + DeleteFileIndex addedDeleteIndex = addedDeletesBySnapshot.get(commitSnapshotId); + return addedDeleteIndex == null || addedDeleteIndex.isEmpty() + ? NO_DELETES + : addedDeleteIndex.forEntry(entry); + } + + /** + * Gets all delete files that were applied to a DELETED data file up to the point it was + * deleted. This includes existing deletes and all deletes added in the scan range up to (but + * not including) the deletion snapshot. + */ + private DeleteFile[] getDeletesForDeletedFile( + ManifestEntry entry, long deletionSnapshotId) { + + List allDeletes = Lists.newArrayList(); + + // Build existing delete index lazily when first DELETED entry is encountered + DeleteFileIndex existingDeleteIndex = existingDeleteIndexSupplier.get(); + DeleteFile[] existingDeletes = + existingDeleteIndex.isEmpty() ? NO_DELETES : existingDeleteIndex.forEntry(entry); + for (DeleteFile df : existingDeletes) { + allDeletes.add(df); + } + + // Add all deletes from snapshots in the scan range BEFORE the deletion + List cumulativeDeletes = cumulativeDeletesMap.get(deletionSnapshotId); + if (cumulativeDeletes != null && !cumulativeDeletes.isEmpty()) { + DeleteFileIndex tempIndex = + DeleteFileIndex.builderFor(cumulativeDeletes) + .specsById(specsById) + .caseSensitive(caseSensitive) + .build(); + DeleteFile[] applicable = tempIndex.forEntry(entry); + for (DeleteFile deleteFile : applicable) { + allDeletes.add(deleteFile); + } + } + + return allDeletes.isEmpty() ? NO_DELETES : allDeletes.toArray(new DeleteFile[0]); + } + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/iceberg/package-info.java b/sdks/java/io/iceberg/src/main/java/org/apache/iceberg/package-info.java new file mode 100644 index 000000000000..a32b307b45d3 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/iceberg/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Classes copied from unreleased Iceberg core. */ +package org.apache.iceberg; 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 9b08e1ff86e1..5849cbd00774 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 @@ -35,6 +35,7 @@ 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.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; @@ -52,6 +53,9 @@ public class IcebergCdcReadSchemaTransformProviderTest { @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + private static final org.apache.iceberg.Schema CDC_SCHEMA = + new org.apache.iceberg.Schema(TestFixtures.SCHEMA.columns(), ImmutableSet.of(1)); + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); @Rule public TestPipeline testPipeline = TestPipeline.create(); @@ -83,8 +87,8 @@ public void testSimpleScan() throws Exception { String identifier = "default.table_" + Long.toString(UUID.randomUUID().hashCode(), 16); TableIdentifier tableId = TableIdentifier.parse(identifier); - Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); - final Schema schema = IcebergUtils.icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + Table simpleTable = warehouse.createTable(tableId, CDC_SCHEMA); + final Schema schema = IcebergUtils.icebergSchemaToBeamSchema(simpleTable.schema()); List> expectedRecords = warehouse.commitData(simpleTable); @@ -122,8 +126,8 @@ public void testStreamingReadUsingManagedTransform() throws Exception { String identifier = "default.table_" + Long.toString(UUID.randomUUID().hashCode(), 16); TableIdentifier tableId = TableIdentifier.parse(identifier); - Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); - final Schema schema = IcebergUtils.icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + Table simpleTable = warehouse.createTable(tableId, CDC_SCHEMA); + final Schema schema = IcebergUtils.icebergSchemaToBeamSchema(simpleTable.schema()); List> expectedRecords = warehouse.commitData(simpleTable).subList(3, 9); List snapshots = Lists.newArrayList(simpleTable.snapshots()); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java index f79991cee571..d7c97efa19f3 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java @@ -54,6 +54,7 @@ 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.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; @@ -119,6 +120,14 @@ public static Iterable data() { @Parameter(0) public boolean useIncrementalScan; + private org.apache.iceberg.Schema schemaForMode( + org.apache.iceberg.Schema schema, Integer... identifierFieldIds) { + if (!useIncrementalScan) { + return schema; + } + return new org.apache.iceberg.Schema(schema.columns(), ImmutableSet.copyOf(identifierFieldIds)); + } + static class PrintRow extends PTransform, PCollection> { @Override @@ -143,7 +152,7 @@ public void process(@Element Row row, OutputReceiver output) { public void testFailWhenBothStartingSnapshotAndTimestampAreSet() { assumeTrue(useIncrementalScan); TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); - warehouse.createTable(tableId, TestFixtures.SCHEMA); + warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); IcebergIO.ReadRows read = IcebergIO.readRows(catalogConfig()) .from(tableId) @@ -161,7 +170,7 @@ public void testFailWhenBothStartingSnapshotAndTimestampAreSet() { public void testFailWhenBothEndingSnapshotAndTimestampAreSet() { assumeTrue(useIncrementalScan); TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); - warehouse.createTable(tableId, TestFixtures.SCHEMA); + warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); IcebergIO.ReadRows read = IcebergIO.readRows(catalogConfig()) .withCdc() @@ -179,7 +188,7 @@ public void testFailWhenBothEndingSnapshotAndTimestampAreSet() { public void testFailWhenStartingPointAndStartingStrategyAreSet() { assumeTrue(useIncrementalScan); TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); - warehouse.createTable(tableId, TestFixtures.SCHEMA); + warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); IcebergIO.ReadRows read = IcebergIO.readRows(catalogConfig()) .withCdc() @@ -197,7 +206,7 @@ public void testFailWhenStartingPointAndStartingStrategyAreSet() { public void testFailWhenPollIntervalIsSetOnBatchRead() { assumeTrue(useIncrementalScan); TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); - warehouse.createTable(tableId, TestFixtures.SCHEMA); + warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); IcebergIO.ReadRows read = IcebergIO.readRows(catalogConfig()) .withCdc() @@ -210,6 +219,32 @@ public void testFailWhenPollIntervalIsSetOnBatchRead() { read.expand(PBegin.in(testPipeline)); } + @Test + public void testCdcFailsWhenTableHasNoIdentifierFields() { + assumeTrue(useIncrementalScan); + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + warehouse.createTable(tableId, TestFixtures.SCHEMA); + IcebergIO.ReadRows read = IcebergIO.readRows(catalogConfig()).from(tableId).withCdc(); + + thrown.expect(IllegalStateException.class); + thrown.expectMessage("Cannot read CDC records"); + thrown.expectMessage("primary key fields"); + read.expand(PBegin.in(testPipeline)); + } + + @Test + public void testCdcFailsWhenProjectionDropsIdentifierFields() { + assumeTrue(useIncrementalScan); + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); + IcebergIO.ReadRows read = + IcebergIO.readRows(catalogConfig()).from(tableId).withCdc().dropping(singletonList("id")); + + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage("projected schema must not drop primary key fields"); + read.expand(PBegin.in(testPipeline)); + } + @Test public void testFailWhenDropAndKeepAreSet() { TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); @@ -307,8 +342,8 @@ public void testProjectedSchemaWithNestedFields() { public void testSimpleScan() throws Exception { TableIdentifier tableId = TableIdentifier.of("default", "table" + Long.toString(UUID.randomUUID().hashCode(), 16)); - Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); - final Schema schema = icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + Table simpleTable = warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); + final Schema schema = icebergSchemaToBeamSchema(simpleTable.schema()); List> expectedRecords = warehouse.commitData(simpleTable); @@ -339,8 +374,8 @@ public void testSimpleScan() throws Exception { public void testScanSelectedFields() throws Exception { TableIdentifier tableId = TableIdentifier.of("default", "table" + Long.toString(UUID.randomUUID().hashCode(), 16)); - Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); - final Schema schema = icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + Table simpleTable = warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); + final Schema schema = icebergSchemaToBeamSchema(simpleTable.schema()); List> expectedRecords = warehouse.commitData(simpleTable); @@ -368,6 +403,11 @@ public void testScanSelectedFields() throws Exception { return null; }); + if (useIncrementalScan) { + testPipeline.run(); + return; + } + // test drop fields read = read.keeping(null).dropping(singletonList("id")); PCollection outputDrop = @@ -387,7 +427,7 @@ public void testScanSelectedFields() throws Exception { public void testScanWithFilter() throws Exception { TableIdentifier tableId = TableIdentifier.of("default", "table" + Long.toString(UUID.randomUUID().hashCode(), 16)); - Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); + Table simpleTable = warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); List> expectedRecords = warehouse.commitData(simpleTable); @@ -440,6 +480,7 @@ public void testReadSchemaWithRandomlyOrderedIds() throws IOException { required(1, "a", Types.IntegerType.get()), required(2, "b", StructType.of(nestedSchema.columns())), required(5, "c", StringType.get())); + schema = schemaForMode(schema, 1); // hadoop catalog will re-order by breadth-first ordering Table simpleTable = warehouse.createTable(tableId, schema); @@ -484,7 +525,8 @@ public void testReadSchemaWithRandomlyOrderedIds() throws IOException { public void testIdentityColumnScan() throws Exception { TableIdentifier tableId = TableIdentifier.of("default", "table" + Long.toString(UUID.randomUUID().hashCode(), 16)); - Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); + org.apache.iceberg.Schema baseSchema = schemaForMode(TestFixtures.SCHEMA, 1); + Table simpleTable = warehouse.createTable(tableId, baseSchema); String identityColumnName = "identity"; String identityColumnValue = "some-value"; @@ -499,11 +541,7 @@ public void testIdentityColumnScan() throws Exception { .newFastAppend() .appendFile( warehouse.writeRecords( - "file1s1.parquet", - TestFixtures.SCHEMA, - spec, - partitionKey, - TestFixtures.FILE1SNAPSHOT1)) + "file1s1.parquet", baseSchema, spec, partitionKey, TestFixtures.FILE1SNAPSHOT1)) .commit(); final Schema schema = icebergSchemaToBeamSchema(simpleTable.schema()); @@ -609,9 +647,10 @@ public void testNameMappingScan() throws Exception { TableIdentifier tableId = TableIdentifier.of("default", "table" + Long.toString(UUID.randomUUID().hashCode(), 16)); + org.apache.iceberg.Schema tableSchema = schemaForMode(TestFixtures.NESTED_SCHEMA, 1); Table simpleTable = warehouse - .buildTable(tableId, TestFixtures.NESTED_SCHEMA) + .buildTable(tableId, tableSchema) .withProperties(tableProperties) .withPartitionSpec(PartitionSpec.unpartitioned()) .create(); @@ -625,7 +664,7 @@ public void testNameMappingScan() throws Exception { .withMetrics(metrics) .build(); - final Schema beamSchema = icebergSchemaToBeamSchema(TestFixtures.NESTED_SCHEMA); + final Schema beamSchema = icebergSchemaToBeamSchema(simpleTable.schema()); simpleTable.newFastAppend().appendFile(dataFile).commit(); @@ -637,7 +676,7 @@ public void testNameMappingScan() throws Exception { final Row[] expectedRows = recordData.stream() - .map(data -> icebergGenericRecord(TestFixtures.NESTED_SCHEMA.asStruct(), data)) + .map(data -> icebergGenericRecord(simpleTable.schema().asStruct(), data)) .map(record -> IcebergUtils.icebergRecordToBeamRow(beamSchema, record)) .toArray(Row[]::new); @@ -695,8 +734,8 @@ public void runWithStartingStrategy(@Nullable StartingStrategy strategy, boolean throws IOException { assumeTrue(useIncrementalScan); TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); - Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); - Schema schema = icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + Table simpleTable = warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); + Schema schema = icebergSchemaToBeamSchema(simpleTable.schema()); List> expectedRecords = warehouse.commitData(simpleTable); if ((strategy == StartingStrategy.LATEST) || (streaming && strategy == null)) { @@ -731,8 +770,8 @@ public void runReadWithBoundary(boolean useSnapshotBoundary, boolean streaming) throws IOException { assumeTrue(useIncrementalScan); TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); - Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); - Schema schema = icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + Table simpleTable = warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1)); + Schema schema = icebergSchemaToBeamSchema(simpleTable.schema()); // only read data committed in the second and third snapshots List> expectedRecords = warehouse.commitData(simpleTable).subList(3, 9); 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 a3217503564c..1319efa7229a 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 @@ -46,9 +46,14 @@ import org.apache.beam.sdk.values.PCollectionRowTuple; import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException; +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.data.Record; +import org.apache.iceberg.types.Types; import org.junit.ClassRule; import org.junit.Rule; import org.junit.Test; @@ -77,6 +82,13 @@ public class IcebergSchemaTransformTranslationTest { .build(); private static final Map CONFIG_PROPERTIES = ImmutableMap.builder().put("key", "value").put("key2", "value2").build(); + private static final org.apache.iceberg.Schema CDC_TRANSLATION_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.required(3, "event_micros", Types.LongType.get())), + ImmutableSet.of(1)); private static final Row WRITE_CONFIG_ROW = Row.withSchema(WRITE_PROVIDER.configurationSchema()) .withFieldValue("table", "test_table_identifier") @@ -104,6 +116,8 @@ public class IcebergSchemaTransformTranslationTest { .withFieldValue("to_timestamp", 456L) .withFieldValue("poll_interval_seconds", 123) .withFieldValue("streaming", true) + .withFieldValue("keep", ImmutableList.of("id", "event_micros")) + .withFieldValue("filter", "\"data\" = 'keep'") .build(); @Test @@ -269,7 +283,15 @@ public void testCdcReadTransformProtoTranslation() // First build a pipeline Pipeline p = Pipeline.create(); String identifier = "default.table_" + Long.toString(UUID.randomUUID().hashCode(), 16); - warehouse.createTable(TableIdentifier.parse(identifier), TestFixtures.SCHEMA); + Table table = warehouse.createTable(TableIdentifier.parse(identifier), CDC_TRANSLATION_SCHEMA); + table + .newFastAppend() + .appendFile( + warehouse.writeRecords( + "cdc-translation.parquet", + table.schema(), + Collections.singletonList(record(1L, "keep", 123L)))) + .commit(); Map properties = new HashMap<>(CATALOG_PROPERTIES); properties.put("warehouse", warehouse.location); @@ -278,6 +300,9 @@ public void testCdcReadTransformProtoTranslation() Row.fromRow(READ_CDC_CONFIG_ROW) .withFieldValue("table", identifier) .withFieldValue("catalog_properties", properties) + .withFieldValue("from_snapshot", table.currentSnapshot().snapshotId()) + .withFieldValue("to_snapshot", table.currentSnapshot().snapshotId()) + .withFieldValue("to_timestamp", null) .build(); IcebergCdcReadSchemaTransform readTransform = @@ -320,4 +345,10 @@ public void testCdcReadTransformProtoTranslation() assertEquals(transformConfigRow, readTransformFromSpec.getConfigurationRow()); } + + private static Record record(long id, String data, long eventMicros) { + return TestFixtures.createRecord( + CDC_TRANSLATION_SCHEMA, + ImmutableMap.of("id", id, "data", data, "event_micros", eventMicros)); + } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/PartitionUtilsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/PartitionUtilsTest.java index d80ec4f95310..740ede55811b 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/PartitionUtilsTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/PartitionUtilsTest.java @@ -23,17 +23,25 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; import java.util.Arrays; import java.util.List; +import java.util.Map; import java.util.Objects; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.MetadataColumns; import org.apache.iceberg.PartitionField; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.transforms.Days; import org.apache.iceberg.transforms.Hours; import org.apache.iceberg.transforms.Months; import org.apache.iceberg.transforms.Transform; +import org.apache.iceberg.types.Types; import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.Test; @@ -166,6 +174,66 @@ public void testAll() { assertEquals(expectedSpec, spec); } + @Test + public void testConstantsMapIncludesCdcMetadataAndIdentityConstants() throws Exception { + org.apache.iceberg.Schema icebergSchema = + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "category", Types.StringType.get())); + PartitionSpec spec = PartitionSpec.builderFor(icebergSchema).identity("category").build(); + DataFile file = + DataFiles.builder(spec) + .withFormat(FileFormat.PARQUET) + .withPath("file:///tmp/table/category=A/data.parquet") + .withPartitionPath("category=A") + .withFileSizeInBytes(100L) + .withRecordCount(2L) + .withFirstRowId(99L) + .build(); + setFileSequenceNumber(file, 42L); + + Map constants = PartitionUtils.constantsMap(spec, file, null); + + assertEquals(99L, constants.get(MetadataColumns.ROW_ID.fieldId())); + assertEquals(42L, constants.get(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.fieldId())); + assertEquals(file.location(), constants.get(MetadataColumns.FILE_PATH.fieldId())); + assertEquals(file.specId(), constants.get(MetadataColumns.SPEC_ID.fieldId())); + assertEquals("A", constants.get(2)); + } + + @Test + public void testConstantsMapUsesExplicitSequenceNumberWhenFileSequenceIsUnavailable() { + org.apache.iceberg.Schema icebergSchema = + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "category", Types.StringType.get())); + PartitionSpec spec = PartitionSpec.builderFor(icebergSchema).identity("category").build(); + DataFile file = + DataFiles.builder(spec) + .withFormat(FileFormat.PARQUET) + .withPath("file:///tmp/table/category=B/data.parquet") + .withPartitionPath("category=B") + .withFileSizeInBytes(100L) + .withRecordCount(2L) + .build(); + + Map constants = PartitionUtils.constantsMap(spec, file, 123L); + + assertEquals(123L, constants.get(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.fieldId())); + assertEquals("B", constants.get(2)); + } + + private static void setFileSequenceNumber(DataFile dataFile, long fileSequenceNumber) + throws Exception { + Method method = dataFile.getClass().getMethod("setFileSequenceNumber", Long.class); + method.setAccessible(true); + try { + method.invoke(dataFile, fileSequenceNumber); + } catch (InvocationTargetException e) { + throw (Exception) e.getCause(); + } + } + static class TestCase { private final String field; private @Nullable String name; diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadUtilsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadUtilsTest.java index 73a0fd19e893..df9c44b7b4f3 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadUtilsTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadUtilsTest.java @@ -30,6 +30,7 @@ import java.util.Objects; import java.util.stream.Collectors; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter; +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.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; @@ -40,7 +41,6 @@ import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.data.Record; import org.apache.iceberg.io.CloseableIterable; -import org.apache.iceberg.parquet.ParquetReader; import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.ClassRule; import org.junit.Rule; @@ -75,14 +75,25 @@ public void testCreateReader() throws IOException { .commit(); } + IcebergScanConfig scanConfig = + IcebergScanConfig.builder() + .setCatalogConfig( + IcebergCatalogConfig.builder() + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build()) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(simpleTable.schema())) + .build(); + int numFiles = 0; try (CloseableIterable iterable = simpleTable.newScan().planTasks()) { for (CombinedScanTask combinedScanTask : iterable) { for (FileScanTask fileScanTask : combinedScanTask.tasks()) { String fileName = Iterables.getLast(Splitter.on("/").split(fileScanTask.file().path())); List recordsRead = new ArrayList<>(); - try (ParquetReader reader = - ReadUtils.createReader(fileScanTask, simpleTable, simpleTable.schema())) { + try (CloseableIterable reader = + ReadUtils.createReader(fileScanTask, simpleTable, scanConfig)) { reader.forEach(recordsRead::add); } @@ -94,6 +105,36 @@ public void testCreateReader() throws IOException { assertEquals(data.size(), numFiles); } + @Test + public void testMaybeApplyFilterUsesRequiredSchemaWithFilterOnlyField() throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + Table simpleTable = warehouse.createTable(tableId, TestFixtures.SCHEMA); + + IcebergScanConfig scanConfig = + IcebergScanConfig.builder() + .setCatalogConfig( + IcebergCatalogConfig.builder() + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build()) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(simpleTable.schema())) + .setKeepFields(ImmutableList.of("data")) + .setFilterString("id = 2") + .build(); + + assertEquals(ImmutableList.of("data"), fieldNames(scanConfig.getProjectedSchema())); + assertEquals(ImmutableList.of("id", "data"), fieldNames(scanConfig.getRequiredSchema())); + + CloseableIterable filtered = + ReadUtils.maybeApplyFilter( + CloseableIterable.withNoopClose(TestFixtures.FILE1SNAPSHOT1), + scanConfig, + scanConfig.getRequiredSchema()); + + assertEquals(ImmutableList.of("falafel"), dataOf(filtered)); + } + @Test public void testSnapshotsBetween() throws IOException { TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); @@ -239,4 +280,16 @@ static TestCase of( return new TestCase(scanConfig, expectedSnapshotId, description); } } + + private static List fieldNames(org.apache.iceberg.Schema schema) { + return schema.columns().stream() + .map(org.apache.iceberg.types.Types.NestedField::name) + .collect(Collectors.toList()); + } + + private static List dataOf(CloseableIterable records) { + return ImmutableList.copyOf(records).stream() + .map(record -> (String) record.getField("data")) + .collect(Collectors.toList()); + } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableDataFileTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableDataFileTest.java index d4e7793718d8..5126822c06f6 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableDataFileTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableDataFileTest.java @@ -61,6 +61,9 @@ public class SerializableDataFileTest { .add("nanValueCounts") .add("lowerBounds") .add("upperBounds") + .add("dataSequenceNumber") + .add("fileSequenceNumber") + .add("firstRowId") .build(); @Test diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableDeleteFileTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableDeleteFileTest.java new file mode 100644 index 000000000000..29ef30c97efb --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/SerializableDeleteFileTest.java @@ -0,0 +1,222 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg; + +import static java.util.Collections.emptyMap; +import static java.util.Collections.singletonMap; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileContent; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.SortOrder; +import org.apache.iceberg.types.Types; +import org.junit.Test; + +/** Tests for {@link SerializableDeleteFile}. */ +public class SerializableDeleteFileTest { + private static final org.apache.iceberg.Schema SCHEMA = + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "category", Types.StringType.get())); + private static final PartitionSpec SPEC = + PartitionSpec.builderFor(SCHEMA).identity("category").build(); + + @Test + public void testPositionDeleteRoundTripPreservesMetadataUsedByCdcReads() throws Exception { + Map columnSizes = new HashMap<>(); + columnSizes.put(1, 11L); + Map valueCounts = new HashMap<>(); + valueCounts.put(1, 3L); + Map nullValueCounts = new HashMap<>(); + nullValueCounts.put(1, 0L); + Map nanValueCounts = new HashMap<>(); + nanValueCounts.put(1, 0L); + Map lowerBounds = new HashMap<>(); + lowerBounds.put(1, ByteBuffer.wrap(new byte[] {0x01})); + Map upperBounds = new HashMap<>(); + upperBounds.put(1, ByteBuffer.wrap(new byte[] {0x05})); + Metrics metrics = + new Metrics( + 3L, + columnSizes, + valueCounts, + nullValueCounts, + nanValueCounts, + lowerBounds, + upperBounds); + DeleteFile deleteFile = + FileMetadata.deleteFileBuilder(SPEC) + .ofPositionDeletes() + .withPath("gs://bucket/deletes/category=A/pos.parquet") + .withFormat(FileFormat.PARQUET) + .withPartitionPath("category=A") + .withFileSizeInBytes(256L) + .withMetrics(metrics) + .withSplitOffsets(Arrays.asList(4L, 128L)) + .withEncryptionKeyMetadata(ByteBuffer.wrap(new byte[] {0x0A, 0x0B})) + .build(); + setSequenceNumbers(deleteFile, 44L, 45L); + + SerializableDeleteFile serialized = SerializableDeleteFile.from(deleteFile, "category=A", true); + DeleteFile reconstructed = + serialized.createDeleteFile( + singletonMap(SPEC.specId(), SPEC), singletonMap(0, SortOrder.unsorted())); + + assertEquals(deleteFile.content(), reconstructed.content()); + assertEquals(deleteFile.location(), reconstructed.location()); + assertEquals(deleteFile.format(), reconstructed.format()); + assertEquals(deleteFile.recordCount(), reconstructed.recordCount()); + assertEquals(deleteFile.fileSizeInBytes(), reconstructed.fileSizeInBytes()); + assertEquals(deleteFile.partition(), reconstructed.partition()); + assertEquals(deleteFile.specId(), reconstructed.specId()); + assertEquals(deleteFile.keyMetadata(), reconstructed.keyMetadata()); + assertEquals(deleteFile.splitOffsets(), reconstructed.splitOffsets()); + assertEquals(deleteFile.columnSizes(), reconstructed.columnSizes()); + assertEquals(deleteFile.valueCounts(), reconstructed.valueCounts()); + assertEquals(deleteFile.nullValueCounts(), reconstructed.nullValueCounts()); + assertEquals(deleteFile.nanValueCounts(), reconstructed.nanValueCounts()); + assertEquals(deleteFile.lowerBounds(), reconstructed.lowerBounds()); + assertEquals(deleteFile.upperBounds(), reconstructed.upperBounds()); + assertEquals(Long.valueOf(44L), serialized.getDataSequenceNumber()); + assertEquals(Long.valueOf(45L), serialized.getFileSequenceNumber()); + assertNull(reconstructed.dataSequenceNumber()); + assertNull(reconstructed.fileSequenceNumber()); + } + + @Test + public void testEqualityDeleteRoundTripPreservesFieldIdsAndSortOrder() { + SortOrder sortOrder = SortOrder.builderFor(SCHEMA).asc("id").withOrderId(7).build(); + DeleteFile deleteFile = + FileMetadata.deleteFileBuilder(SPEC) + .ofEqualityDeletes(1, 2) + .withSortOrder(sortOrder) + .withPath("gs://bucket/deletes/category=A/eq.parquet") + .withFormat(FileFormat.PARQUET) + .withPartitionPath("category=A") + .withFileSizeInBytes(256L) + .withRecordCount(2L) + .build(); + + SerializableDeleteFile serialized = SerializableDeleteFile.from(deleteFile, "category=A", true); + DeleteFile reconstructed = + serialized.createDeleteFile(singletonMap(SPEC.specId(), SPEC), singletonMap(7, sortOrder)); + + assertEquals(FileContent.EQUALITY_DELETES, reconstructed.content()); + assertEquals(Arrays.asList(1, 2), reconstructed.equalityFieldIds()); + assertEquals(Integer.valueOf(7), reconstructed.sortOrderId()); + } + + @Test + public void testPuffinDeleteRoundTripPreservesDeletionVectorMetadata() { + DeleteFile deleteFile = + FileMetadata.deleteFileBuilder(SPEC) + .ofPositionDeletes() + .withPath("gs://bucket/deletes/category=A/dv.puffin") + .withFormat(FileFormat.PUFFIN) + .withPartitionPath("category=A") + .withFileSizeInBytes(512L) + .withRecordCount(1L) + .withContentOffset(64L) + .withContentSizeInBytes(128L) + .withReferencedDataFile("gs://bucket/data/category=A/data.parquet") + .build(); + + SerializableDeleteFile serialized = SerializableDeleteFile.from(deleteFile, "category=A", true); + DeleteFile reconstructed = + serialized.createDeleteFile( + singletonMap(SPEC.specId(), SPEC), singletonMap(0, SortOrder.unsorted())); + + assertEquals(FileFormat.PUFFIN, reconstructed.format()); + assertEquals(Long.valueOf(64L), reconstructed.contentOffset()); + assertEquals(Long.valueOf(128L), reconstructed.contentSizeInBytes()); + assertEquals("gs://bucket/data/category=A/data.parquet", reconstructed.referencedDataFile()); + } + + @Test + public void testCreateDeleteFileFailsClearlyForMissingPartitionSpec() { + DeleteFile deleteFile = + FileMetadata.deleteFileBuilder(SPEC) + .ofPositionDeletes() + .withPath("gs://bucket/deletes/category=A/pos.parquet") + .withFormat(FileFormat.PARQUET) + .withPartitionPath("category=A") + .withFileSizeInBytes(256L) + .withRecordCount(2L) + .build(); + SerializableDeleteFile serialized = SerializableDeleteFile.from(deleteFile, "category=A", true); + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, () -> serialized.createDeleteFile(emptyMap(), null)); + + assertTrue(thrown.getMessage().contains("created with spec id '" + SPEC.specId() + "'")); + } + + @Test + public void testCreateEqualityDeleteFileFailsClearlyForMissingSortOrder() { + SortOrder sortOrder = SortOrder.builderFor(SCHEMA).asc("id").withOrderId(7).build(); + DeleteFile deleteFile = + FileMetadata.deleteFileBuilder(SPEC) + .ofEqualityDeletes(1) + .withSortOrder(sortOrder) + .withPath("gs://bucket/deletes/category=A/eq.parquet") + .withFormat(FileFormat.PARQUET) + .withPartitionPath("category=A") + .withFileSizeInBytes(256L) + .withRecordCount(2L) + .build(); + SerializableDeleteFile serialized = SerializableDeleteFile.from(deleteFile, "category=A", true); + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, + () -> serialized.createDeleteFile(singletonMap(SPEC.specId(), SPEC), emptyMap())); + + assertTrue(thrown.getMessage().contains("sort order id '7'")); + } + + private static void setSequenceNumbers( + DeleteFile deleteFile, long dataSequenceNumber, long fileSequenceNumber) throws Exception { + invoke(deleteFile, "setDataSequenceNumber", dataSequenceNumber); + invoke(deleteFile, "setFileSequenceNumber", fileSequenceNumber); + } + + private static void invoke(DeleteFile deleteFile, String methodName, Long value) + throws Exception { + Method method = deleteFile.getClass().getMethod(methodName, Long.class); + method.setAccessible(true); + try { + method.invoke(deleteFile, value); + } catch (InvocationTargetException e) { + throw (Exception) e.getCause(); + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestDataWarehouse.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestDataWarehouse.java index dcb2d804d2e6..2e711219349c 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestDataWarehouse.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestDataWarehouse.java @@ -64,7 +64,7 @@ public class TestDataWarehouse extends ExternalResource { protected final Configuration hadoopConf; - protected String location; + public String location; protected Catalog catalog; protected boolean someTableHasBeenCreated = false; 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 74408d67ed86..f0c7ae925df7 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 @@ -318,7 +318,8 @@ public Row apply(Long num) { }; protected static final org.apache.iceberg.Schema ICEBERG_SCHEMA = - beamSchemaToIcebergSchema(BEAM_SCHEMA); + new org.apache.iceberg.Schema( + beamSchemaToIcebergSchema(BEAM_SCHEMA).columns(), Collections.singleton(1)); protected static final SimpleFunction RECORD_FUNC = new SimpleFunction() { @Override @@ -450,7 +451,7 @@ public void testReadWithColumnPruning_keep() throws Exception { List expectedRows = populateTable(table); - List fieldsToKeep = Arrays.asList("row", "str", "modulo_5", "nullable_long"); + List fieldsToKeep = Arrays.asList("row", "modulo_5", "nullable_long"); RowFilter rowFilter = new RowFilter(BEAM_SCHEMA).keep(fieldsToKeep); Map config = new HashMap<>(managedIcebergConfig(tableId())); @@ -543,7 +544,7 @@ public void testStreamingReadWithColumnPruning_drop() throws Exception { List expectedRows = populateTable(table); - List fieldsToDrop = Arrays.asList("row", "str", "modulo_5", "nullable_long"); + List fieldsToDrop = Arrays.asList("row", "modulo_5", "nullable_long"); RowFilter rowFilter = new RowFilter(BEAM_SCHEMA).drop(fieldsToDrop); Map config = new HashMap<>(managedIcebergConfig(tableId())); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtilsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtilsTest.java new file mode 100644 index 000000000000..546386073ff6 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtilsTest.java @@ -0,0 +1,381 @@ +/* + * 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.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.SerializableDeleteFile; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +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.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.EqualityDeleteWriter; +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.io.CloseableIterable; +import org.apache.iceberg.types.Types; +import org.checkerframework.checker.nullness.qual.Nullable; +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 CdcReadUtils}. */ +@RunWith(JUnit4.class) +public class CdcReadUtilsTest { + 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())), + 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 addedRowsFiltersPositionAndEqualityDeletesWithUnprojectedEqualityColumn() + throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + DataFile dataFile = + warehouse.writeRecords( + "cdc-added.parquet", + table.schema(), + ImmutableList.of( + record(0L, "keep-0"), + record(1L, "drop-by-pos"), + record(2L, "drop-by-data"), + record(3L, "keep-3"))); + + DeleteFile positionDelete = + writePositionDelete(table, dataFile, "cdc-added-pos-delete.parquet", 1L); + DeleteFile equalityDelete = + writeEqualityDelete(table, dataFile, "cdc-added-eq-delete.parquet", "drop-by-data"); + + SerializableChangelogTask task = + task( + SerializableChangelogTask.Type.ADDED_ROWS, + dataFile, + ImmutableList.of(positionDelete, equalityDelete), + ImmutableList.of(), + table); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig(table, tableId), true); + + assertEquals(ImmutableList.of(0L, 3L), idsOf(records)); + } + + @Test + public void addedRowsUsesFullTableSchemaWhenProjectionDisabled() throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + DataFile dataFile = + warehouse.writeRecords( + "cdc-added-full-schema.parquet", + table.schema(), + ImmutableList.of(record(0L, "non-projected-value"))); + SerializableChangelogTask task = + task( + SerializableChangelogTask.Type.ADDED_ROWS, + dataFile, + ImmutableList.of(), + ImmutableList.of(), + table); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig(table, tableId), false); + + // with projection disabled, all columns should be returned + assertEquals(ImmutableList.of("non-projected-value"), dataValuesOf(records)); + } + + @Test + public void deletedRowsExcludesRowsAlreadyHiddenByExistingDeletes() throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + DataFile dataFile = + warehouse.writeRecords( + "cdc-deleted-rows.parquet", + table.schema(), + ImmutableList.of( + record(0L, "already-hidden-by-position"), + record(1L, "new-position-delete"), + record(2L, "already-hidden-by-equality"), + record(3L, "new-equality-delete"), + record(4L, "still-live"))); + + DeleteFile existingPositionDelete = + writePositionDelete(table, dataFile, "cdc-existing-pos-delete.parquet", 0L); + DeleteFile existingEqualityDelete = + writeEqualityDelete( + table, dataFile, "cdc-existing-eq-delete.parquet", "already-hidden-by-equality"); + DeleteFile addedPositionDelete = + writePositionDelete(table, dataFile, "cdc-added-pos-delete.parquet", 1L); + DeleteFile addedEqualityDelete = + writeEqualityDelete(table, dataFile, "cdc-added-eq-delete.parquet", "new-equality-delete"); + + SerializableChangelogTask task = + task( + SerializableChangelogTask.Type.DELETED_ROWS, + dataFile, + ImmutableList.of(addedPositionDelete, addedEqualityDelete), + ImmutableList.of(existingPositionDelete, existingEqualityDelete), + table); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig(table, tableId), true); + + assertEquals(ImmutableList.of(1L, 3L), idsOf(records)); + } + + @Test + public void deletedRowsUsesFullTableSchemaWhenProjectionDisabled() throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + DataFile dataFile = + warehouse.writeRecords( + "cdc-deleted-rows-full-schema.parquet", + table.schema(), + ImmutableList.of(record(0L, "position-deleted-value"), record(1L, "still-live-value"))); + DeleteFile positionDelete = + writePositionDelete(table, dataFile, "cdc-deleted-rows-full-schema-pos-delete.parquet", 0L); + SerializableChangelogTask task = + task( + SerializableChangelogTask.Type.DELETED_ROWS, + dataFile, + ImmutableList.of(positionDelete), + ImmutableList.of(), + table); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig(table, tableId), false); + + // with projection disabled, all columns should be returned + assertEquals(ImmutableList.of("position-deleted-value"), dataValuesOf(records)); + } + + @Test + public void deletedFileEmitsOnlyRowsNotAlreadyHiddenByExistingDeletes() throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + DataFile dataFile = + warehouse.writeRecords( + "cdc-deleted-file.parquet", + table.schema(), + ImmutableList.of( + record(0L, "already-hidden-by-position"), + record(1L, "still-live-one"), + record(2L, "already-hidden-by-equality"), + record(3L, "still-live-two"))); + + DeleteFile existingPositionDelete = + writePositionDelete(table, dataFile, "cdc-deleted-file-pos-delete.parquet", 0L); + DeleteFile existingEqualityDelete = + writeEqualityDelete( + table, dataFile, "cdc-deleted-file-eq-delete.parquet", "already-hidden-by-equality"); + SerializableChangelogTask task = + task( + SerializableChangelogTask.Type.DELETED_FILE, + dataFile, + ImmutableList.of(), + ImmutableList.of(existingPositionDelete, existingEqualityDelete), + table); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig(table, tableId), true); + + assertEquals(ImmutableList.of(1L, 3L), idsOf(records)); + } + + @Test + public void deletedFileUsesFullTableSchemaWhenProjectionDisabled() throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + DataFile dataFile = + warehouse.writeRecords( + "cdc-deleted-file-full-schema.parquet", + table.schema(), + ImmutableList.of(record(0L, "deleted-file-value"))); + SerializableChangelogTask task = + task( + SerializableChangelogTask.Type.DELETED_FILE, + dataFile, + ImmutableList.of(), + ImmutableList.of(), + table); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig(table, tableId), false); + + // with projection disabled, all columns should be returned + assertEquals(ImmutableList.of("deleted-file-value"), dataValuesOf(records)); + } + + @Test + public void deletedRowsHandlesNullEqualityDeletesWithoutPushdown() throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + Table table = warehouse.createTable(tableId, CDC_SCHEMA); + DataFile dataFile = + warehouse.writeRecords( + "cdc-null-equality-delete.parquet", + table.schema(), + ImmutableList.of(record(0L, "live"), record(1L, null), record(2L, "also-live"))); + DeleteFile nullEqualityDelete = + writeEqualityDelete(table, dataFile, "cdc-null-eq-delete.parquet", null); + SerializableChangelogTask task = + task( + SerializableChangelogTask.Type.DELETED_ROWS, + dataFile, + ImmutableList.of(nullEqualityDelete), + ImmutableList.of(), + table); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig(table, tableId), true); + + assertEquals(ImmutableList.of(1L), idsOf(records)); + } + + private static Record record(long id, String data) { + GenericRecord record = GenericRecord.create(CDC_SCHEMA); + record.setField("id", id); + record.setField("data", data); + 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 (writer) { + for (long position : positions) { + writer.write(PositionDelete.create().set(dataFile.location(), position)); + } + } + return writer.toDeleteFile(); + } + + private static DeleteFile writeEqualityDelete( + Table table, DataFile dataFile, String filename, @Nullable String data) throws IOException { + org.apache.iceberg.Schema deleteSchema = table.schema().select("data"); + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec(), new int[] {2}, deleteSchema, null); + EqualityDeleteWriter writer = + appenderFactory.newEqDeleteWriter( + EncryptedFiles.plainAsEncryptedOutput( + table.io().newOutputFile(dataFile.location() + "." + filename)), + FileFormat.PARQUET, + null); + try (writer) { + GenericRecord deleteRecord = GenericRecord.create(deleteSchema); + deleteRecord.setField("data", data); + writer.write(deleteRecord); + } + return writer.toDeleteFile(); + } + + private IcebergScanConfig scanConfig(Table table, TableIdentifier tableId) { + return IcebergScanConfig.builder() + .setCatalogConfig( + IcebergCatalogConfig.builder() + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build()) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(table.schema())) + .setKeepFields(ImmutableList.of("id")) + .build(); + } + + private static SerializableChangelogTask task( + SerializableChangelogTask.Type type, + DataFile dataFile, + List addedDeletes, + List existingDeletes, + Table table) { + 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(1L) + .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 List idsOf(CloseableIterable records) { + return ImmutableList.copyOf(records).stream() + .map(record -> (Long) record.getField("id")) + .collect(Collectors.toList()); + } + + private static List dataValuesOf(CloseableIterable records) { + return ImmutableList.copyOf(records).stream() + .map(record -> (String) record.getField("data")) + .collect(Collectors.toList()); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/DeleteReaderTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/DeleteReaderTest.java new file mode 100644 index 000000000000..bf3aaf414579 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/DeleteReaderTest.java @@ -0,0 +1,419 @@ +/* + * 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.assertNotNull; +import static org.junit.Assert.assertNull; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +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.Sets; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.data.DeleteLoader; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.deletes.PositionDeleteIndex; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeSet; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Verifies that {@link DeleteReader#read} returns the union of records matched by position + * and equality deletes. + * + *

The tests stub the {@link DeleteLoader} so we exercise the predicate-composition logic + * directly without writing real delete files. End-to-end is covered by other tests. + */ +@RunWith(JUnit4.class) +public class DeleteReaderTest { + private static final Schema TABLE_SCHEMA = + new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "name", Types.StringType.get())); + + private static final DeleteFile POS_FILE = + FileMetadata.deleteFileBuilder(PartitionSpec.unpartitioned()) + .ofPositionDeletes() + .withPath("/test/pos.parquet") + .withFileSizeInBytes(100) + .withRecordCount(3) + .build(); + + private static final DeleteFile EQ_FILE_ID = + FileMetadata.deleteFileBuilder(PartitionSpec.unpartitioned()) + .ofEqualityDeletes(1) + .withPath("/test/eq.parquet") + .withFileSizeInBytes(100) + .withRecordCount(2) + .build(); + + private static final DeleteFile EQ_FILE_NAME = + FileMetadata.deleteFileBuilder(PartitionSpec.unpartitioned()) + .ofEqualityDeletes(2) + .withPath("/test/eq-name.parquet") + .withFileSizeInBytes(100) + .withRecordCount(2) + .build(); + + /** {@link DeleteReader} that returns a stubbed {@link DeleteLoader} for tests. */ + private static class StubDeleteReader extends DeleteReader { + private final DeleteLoader stub; + + StubDeleteReader(List deletes, DeleteLoader stub) { + super( + "/test/data.parquet", + deletes, + TABLE_SCHEMA, + TABLE_SCHEMA, + true, + PreloadedDeletes.empty()); + this.stub = stub; + } + + StubDeleteReader( + List deletes, + DeleteLoader stub, + DeleteReader.PreloadedDeletes preloadedDeletes) { + super("/test/data.parquet", deletes, TABLE_SCHEMA, TABLE_SCHEMA, true, preloadedDeletes); + this.stub = stub; + } + + StubDeleteReader( + List deletes, + DeleteLoader stub, + Schema requestedSchema, + boolean needRowPosCol) { + super( + "/test/data.parquet", + deletes, + TABLE_SCHEMA, + requestedSchema, + needRowPosCol, + PreloadedDeletes.empty()); + this.stub = stub; + } + + @Override + protected StructLike asStructLike(Record record) { + return record; + } + + @Override + protected InputFile getInputFile(String location) { + throw new UnsupportedOperationException("not used with a stubbed DeleteLoader"); + } + + @Override + protected DeleteLoader newDeleteLoader() { + return stub; + } + } + + /** {@link DeleteLoader} that returns pre-built indexes. */ + private static class StubLoader implements DeleteLoader { + private final PositionDeleteIndex posIndex; + private final Map, StructLikeSet> eqSets; + private int posLoadCount = 0; + private int eqLoadCount = 0; + + StubLoader(PositionDeleteIndex posIndex, StructLikeSet eqSet) { + this(posIndex, Collections.singletonMap(Collections.singleton(1), eqSet)); + } + + StubLoader(PositionDeleteIndex posIndex, Map, StructLikeSet> eqSets) { + this.posIndex = posIndex; + this.eqSets = eqSets; + } + + @Override + public PositionDeleteIndex loadPositionDeletes(Iterable files, CharSequence path) { + posLoadCount++; + return posIndex; + } + + @Override + public StructLikeSet loadEqualityDeletes(Iterable files, Schema schema) { + eqLoadCount++; + return eqSets.getOrDefault( + Sets.newHashSet(TypeUtil.getProjectedIds(new Schema(schema.asStruct().fields()))), + StructLikeSet.create(schema.asStruct())); + } + } + + /** A minimal HashSet-backed {@link PositionDeleteIndex} for tests. */ + private static PositionDeleteIndex posIndexOf(long... positions) { + Set backing = new HashSet<>(); + for (long p : positions) { + backing.add(p); + } + return new PositionDeleteIndex() { + @Override + public void delete(long pos) { + backing.add(pos); + } + + @Override + public void delete(long from, long to) { + for (long p = from; p < to; p++) { + backing.add(p); + } + } + + @Override + public boolean isDeleted(long pos) { + return backing.contains(pos); + } + + @Override + public boolean isEmpty() { + return backing.isEmpty(); + } + + @Override + public long cardinality() { + return backing.size(); + } + }; + } + + private static StructLikeSet eqSetOfIds(int... ids) { + Schema idSchema = TABLE_SCHEMA.select("id"); + StructLikeSet set = StructLikeSet.create(idSchema.asStruct()); + for (int id : ids) { + GenericRecord r = GenericRecord.create(idSchema); + r.setField("id", id); + set.add(r); + } + return set; + } + + private static StructLikeSet eqSetOfNames(String... names) { + Schema nameSchema = TABLE_SCHEMA.select("name"); + StructLikeSet set = StructLikeSet.create(nameSchema.asStruct()); + for (String name : names) { + GenericRecord r = GenericRecord.create(nameSchema); + r.setField("name", name); + set.add(r); + } + return set; + } + + /** Builds N records (id=0..N-1, name="v0".."vN-1") matching {@code readSchema}. */ + private static List records(Schema readSchema, int n) { + boolean hasPos = readSchema.findField("_pos") != null; + List recs = new ArrayList<>(n); + for (long i = 0; i < n; i++) { + GenericRecord r = GenericRecord.create(readSchema); + r.setField("id", (int) i); + r.setField("name", "v" + i); + if (hasPos) { + r.setField("_pos", i); + } + recs.add(r); + } + return recs; + } + + /** Sorted list of "id" values from the output, for stable assertions. */ + private static List idsOf(CloseableIterable records) { + return ImmutableList.copyOf(records).stream() + .map(r -> (Integer) r.getField("id")) + .sorted() + .collect(Collectors.toList()); + } + + /** With no delete files at all, {@code read()} emits nothing. */ + @Test + public void noDeletesEmitsNothing() { + DeleteLoader loader = new StubLoader(posIndexOf(), eqSetOfIds()); + DeleteReader reader = new StubDeleteReader(Collections.emptyList(), loader); + List input = records(reader.requiredSchema(), 5); + + CloseableIterable output = reader.read(CloseableIterable.withNoopClose(input)); + + assertEquals(Collections.emptyList(), idsOf(output)); + } + + /** Pos-only emits only the pos-deleted records. */ + @Test + public void posOnlyEmitsPosDeletedRecords() { + DeleteLoader loader = new StubLoader(posIndexOf(1L, 3L), eqSetOfIds()); + DeleteReader reader = new StubDeleteReader(ImmutableList.of(POS_FILE), loader); + List input = records(reader.requiredSchema(), 5); + + CloseableIterable output = reader.read(CloseableIterable.withNoopClose(input)); + + assertEquals(ImmutableList.of(1, 3), idsOf(output)); + } + + /** Only equality deletes, emits records matching the eq set. */ + @Test + public void eqOnlyEmitsEqDeletedRecords() { + DeleteLoader loader = new StubLoader(posIndexOf(), eqSetOfIds(2, 4)); + DeleteReader reader = new StubDeleteReader(ImmutableList.of(EQ_FILE_ID), loader); + List input = records(reader.requiredSchema(), 5); + + CloseableIterable output = reader.read(CloseableIterable.withNoopClose(input)); + + assertEquals(ImmutableList.of(2, 4), idsOf(output)); + } + + /** Pos-deletes plus equality deletes, emit the union without duplication. */ + @Test + public void posAndEqEmitUnion() { + DeleteLoader loader = new StubLoader(posIndexOf(0L, 4L), eqSetOfIds(2, 4)); + DeleteReader reader = + new StubDeleteReader(ImmutableList.of(POS_FILE, EQ_FILE_ID), loader); + List input = records(reader.requiredSchema(), 6); + + CloseableIterable output = reader.read(CloseableIterable.withNoopClose(input)); + + // id 4 is in both sides; it must appear exactly once. + assertEquals(ImmutableList.of(0, 2, 4), idsOf(output)); + } + + /** Preloaded position deletes are reused instead of loading the same delete files again. */ + @Test + public void preloadedPositionDeletesAvoidSecondLoad() { + StubLoader loader = new StubLoader(posIndexOf(), eqSetOfIds()); + PositionDeleteIndex preloadedPosIndex = posIndexOf(1L, 3L); + DeleteReader reader = + new StubDeleteReader( + ImmutableList.of(POS_FILE), + loader, + DeleteReader.PreloadedDeletes.of(preloadedPosIndex, Collections.emptyMap())); + List input = records(reader.requiredSchema(), 5); + + CloseableIterable output = reader.read(CloseableIterable.withNoopClose(input)); + + assertEquals(ImmutableList.of(1, 3), idsOf(output)); + assertEquals(0, loader.posLoadCount); + } + + /** Preloaded equality delete sets are reused instead of loading the same delete files again. */ + @Test + public void preloadedEqualityDeletesAvoidSecondLoad() { + StubLoader loader = new StubLoader(posIndexOf(), eqSetOfIds()); + Map, StructLikeSet> preloadedEqSets = new HashMap<>(); + preloadedEqSets.put(Collections.singleton(1), eqSetOfIds(2, 4)); + DeleteReader reader = + new StubDeleteReader( + ImmutableList.of(EQ_FILE_ID), + loader, + DeleteReader.PreloadedDeletes.of(null, preloadedEqSets)); + List input = records(reader.requiredSchema(), 5); + + CloseableIterable output = reader.read(CloseableIterable.withNoopClose(input)); + + assertEquals(ImmutableList.of(2, 4), idsOf(output)); + assertEquals(0, loader.eqLoadCount); + } + + @Test + public void requiredSchemaAddsUnprojectedEqualityDeleteField() { + Schema requestedSchema = TABLE_SCHEMA.select("id"); + DeleteLoader loader = + new StubLoader( + posIndexOf(), Collections.singletonMap(Collections.singleton(2), eqSetOfNames("v2"))); + DeleteReader reader = + new StubDeleteReader(ImmutableList.of(EQ_FILE_NAME), loader, requestedSchema, true); + + assertEquals( + ImmutableList.of("id", "name"), + reader.requiredSchema().columns().stream() + .map(Types.NestedField::name) + .collect(Collectors.toList())); + + List input = records(reader.requiredSchema(), 4); + CloseableIterable output = reader.read(CloseableIterable.withNoopClose(input)); + + assertEquals(ImmutableList.of(2), idsOf(output)); + } + + @Test + public void rowPositionColumnIsOnlyAddedWhenRequiredForPositionDeletes() { + DeleteLoader loader = new StubLoader(posIndexOf(1L), eqSetOfIds()); + + DeleteReader posReaderNeedsPos = + new StubDeleteReader(ImmutableList.of(POS_FILE), loader, TABLE_SCHEMA, true); + DeleteReader posReaderDoesNotNeedPos = + new StubDeleteReader(ImmutableList.of(POS_FILE), loader, TABLE_SCHEMA, false); + DeleteReader eqReader = + new StubDeleteReader(ImmutableList.of(EQ_FILE_ID), loader, TABLE_SCHEMA, true); + + assertNotNull(posReaderNeedsPos.requiredSchema().findField("_pos")); + assertNull(posReaderDoesNotNeedPos.requiredSchema().findField("_pos")); + assertNull(eqReader.requiredSchema().findField("_pos")); + } + + @Test + public void multipleEqualityDeleteGroupsAreOrCombined() { + Map, StructLikeSet> eqSets = new HashMap<>(); + eqSets.put(Collections.singleton(1), eqSetOfIds(1)); + eqSets.put(Collections.singleton(2), eqSetOfNames("v3")); + StubLoader loader = new StubLoader(posIndexOf(), eqSets); + DeleteReader reader = + new StubDeleteReader(ImmutableList.of(EQ_FILE_ID, EQ_FILE_NAME), loader); + + CloseableIterable output = + reader.read(CloseableIterable.withNoopClose(records(reader.requiredSchema(), 5))); + + assertEquals(ImmutableList.of(1, 3), idsOf(output)); + assertEquals(2, loader.eqLoadCount); + } + + @Test + public void preloadedEqualityDeleteKeysAreDefensivelyCopied() { + StructLikeSet idDeletes = eqSetOfIds(2); + Set mutableKey = new HashSet<>(Collections.singleton(1)); + Map, StructLikeSet> preloadedEqSets = new HashMap<>(); + preloadedEqSets.put(mutableKey, idDeletes); + + DeleteReader.PreloadedDeletes preloadedDeletes = + DeleteReader.PreloadedDeletes.of(null, preloadedEqSets); + mutableKey.add(2); + + assertEquals(idDeletes, preloadedDeletes.equalityDeleteSet(Collections.singleton(1))); + + StubLoader loader = new StubLoader(posIndexOf(), eqSetOfIds()); + DeleteReader reader = + new StubDeleteReader(ImmutableList.of(EQ_FILE_ID), loader, preloadedDeletes); + CloseableIterable output = + reader.read(CloseableIterable.withNoopClose(records(reader.requiredSchema(), 4))); + + assertEquals(ImmutableList.of(2), idsOf(output)); + assertEquals(0, loader.eqLoadCount); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/SerializableChangelogTaskTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/SerializableChangelogTaskTest.java new file mode 100644 index 000000000000..aa77d37af7af --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/SerializableChangelogTaskTest.java @@ -0,0 +1,256 @@ +/* + * 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.containsString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; + +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.io.iceberg.SerializableDataFile; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.iceberg.AddedRowsScanTask; +import org.apache.iceberg.ChangelogOperation; +import org.apache.iceberg.ChangelogScanTask; +import org.apache.iceberg.ContentScanTask; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.DeletedDataFileScanTask; +import org.apache.iceberg.DeletedRowsScanTask; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.expressions.Expression; +import org.apache.iceberg.expressions.ExpressionParser; +import org.apache.iceberg.expressions.Expressions; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link SerializableChangelogTask}. */ +@RunWith(JUnit4.class) +public class SerializableChangelogTaskTest { + private static final PartitionSpec SPEC = PartitionSpec.unpartitioned(); + private static final DataFile DATA_FILE = + DataFiles.builder(SPEC) + .withFormat(FileFormat.PARQUET) + .withPath("gs://bucket/data/file.parquet") + .withFileSizeInBytes(512L) + .withMetrics(new Metrics(3L, null, null, null, null, null, null)) + .build(); + private static final DeleteFile ADDED_DELETE = + FileMetadata.deleteFileBuilder(SPEC) + .ofPositionDeletes() + .withPath("gs://bucket/delete/added.parquet") + .withFileSizeInBytes(32L) + .withRecordCount(1L) + .build(); + private static final DeleteFile EXISTING_DELETE = + FileMetadata.deleteFileBuilder(SPEC) + .ofPositionDeletes() + .withPath("gs://bucket/delete/existing.parquet") + .withFileSizeInBytes(64L) + .withRecordCount(2L) + .build(); + + @Test + public void coderRoundTripPreservesTaskBasics() throws Exception { + SerializableChangelogTask task = + SerializableChangelogTask.builder() + .setType(SerializableChangelogTask.Type.ADDED_ROWS) + .setDataFile(SerializableDataFile.from(DATA_FILE, "", false)) + .setSpecId(SPEC.specId()) + .setOperation(ChangelogOperation.INSERT) + .setOrdinal(7) + .setCommitSnapshotId(123L) + .setStart(5L) + .setLength(99L) + .setJsonExpression(ExpressionParser.toJson(Expressions.alwaysTrue())) + .build(); + + SerializableChangelogTask decoded = CoderUtils.clone(SerializableChangelogTask.coder(), task); + + assertEquals(SerializableChangelogTask.Type.ADDED_ROWS, decoded.getType()); + assertEquals(ChangelogOperation.INSERT, decoded.getOperation()); + assertEquals(7, decoded.getOrdinal()); + assertEquals(123L, decoded.getCommitSnapshotId()); + assertEquals(5L, decoded.getStart()); + assertEquals(99L, decoded.getLength()); + assertEquals(DATA_FILE.location(), decoded.getDataFile().getPath()); + assertEquals(DATA_FILE.fileSizeInBytes(), decoded.getDataFile().getFileSizeInBytes()); + assertEquals(Collections.emptyList(), decoded.getExistingDeletes()); + assertEquals(Collections.emptyList(), decoded.getAddedDeletes()); + assertEquals(Expressions.alwaysTrue().toString(), decoded.getExpression(null).toString()); + } + + @Test + public void helperMethodsReadSupportedTaskTypes() { + FakeAddedRowsTask added = new FakeAddedRowsTask(ImmutableList.of(ADDED_DELETE)); + FakeDeletedRowsTask deletedRows = + new FakeDeletedRowsTask(ImmutableList.of(ADDED_DELETE), ImmutableList.of(EXISTING_DELETE)); + FakeDeletedDataFileTask deletedFile = + new FakeDeletedDataFileTask(ImmutableList.of(EXISTING_DELETE)); + + assertEquals( + SerializableChangelogTask.Type.ADDED_ROWS, SerializableChangelogTask.getType(added)); + assertEquals( + SerializableChangelogTask.Type.DELETED_ROWS, + SerializableChangelogTask.getType(deletedRows)); + assertEquals( + SerializableChangelogTask.Type.DELETED_FILE, + SerializableChangelogTask.getType(deletedFile)); + + assertEquals(22L, SerializableChangelogTask.getLength(added)); + assertEquals( + 44L, SerializableChangelogTask.getTotalLength(ImmutableList.of(added, deletedRows))); + assertSame(DATA_FILE, SerializableChangelogTask.getDataFile(deletedRows)); + assertSame(SPEC, SerializableChangelogTask.getSpec(deletedFile)); + assertSame(DATA_FILE.partition(), SerializableChangelogTask.getPartition(added)); + assertThat(SerializableChangelogTask.getAddedDeleteFiles(added), contains(ADDED_DELETE)); + assertThat(SerializableChangelogTask.getAddedDeleteFiles(deletedRows), contains(ADDED_DELETE)); + assertEquals( + Collections.emptyList(), SerializableChangelogTask.getAddedDeleteFiles(deletedFile)); + } + + @Test + public void unsupportedTaskTypeFailsClearly() { + ChangelogScanTask unsupported = + new ChangelogScanTask() { + @Override + public ChangelogOperation operation() { + return ChangelogOperation.INSERT; + } + + @Override + public int changeOrdinal() { + return 0; + } + + @Override + public long commitSnapshotId() { + return 0L; + } + }; + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, () -> SerializableChangelogTask.getLength(unsupported)); + + assertThat( + thrown.getMessage(), + containsString("Unknown ChangelogScanTask type: " + unsupported.getClass())); + } + + private abstract static class FakeContentTask + implements ChangelogScanTask, ContentScanTask { + @Override + public DataFile file() { + return DATA_FILE; + } + + @Override + public PartitionSpec spec() { + return SPEC; + } + + @Override + public StructLike partition() { + return DATA_FILE.partition(); + } + + @Override + public long start() { + return 11L; + } + + @Override + public long length() { + return 22L; + } + + @Override + public Expression residual() { + return Expressions.alwaysTrue(); + } + + @Override + public int changeOrdinal() { + return 2; + } + + @Override + public long commitSnapshotId() { + return 101L; + } + } + + private static class FakeAddedRowsTask extends FakeContentTask implements AddedRowsScanTask { + private final List deletes; + + FakeAddedRowsTask(List deletes) { + this.deletes = deletes; + } + + @Override + public List deletes() { + return deletes; + } + } + + private static class FakeDeletedRowsTask extends FakeContentTask implements DeletedRowsScanTask { + private final List addedDeletes; + private final List existingDeletes; + + FakeDeletedRowsTask(List addedDeletes, List existingDeletes) { + this.addedDeletes = addedDeletes; + this.existingDeletes = existingDeletes; + } + + @Override + public List addedDeletes() { + return addedDeletes; + } + + @Override + public List existingDeletes() { + return existingDeletes; + } + } + + private static class FakeDeletedDataFileTask extends FakeContentTask + implements DeletedDataFileScanTask { + private final List existingDeletes; + + FakeDeletedDataFileTask(List existingDeletes) { + this.existingDeletes = existingDeletes; + } + + @Override + public List existingDeletes() { + return existingDeletes; + } + } +} diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py index 4ef6c392254b..4e45d0324ee2 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py @@ -491,8 +491,6 @@ class TriggerCopyJobs(beam.DoFn): """ TRIGGER_DELETE_TEMP_TABLES = 'TriggerDeleteTempTables' - # https://docs.cloud.google.com/bigquery/quotas#copy_jobs - MAX_SOURCES_PER_COPY_JOB = 1200 def __init__( self, @@ -530,90 +528,96 @@ def process( self, element_list, job_name_prefix=None, unused_schema_mod_jobs=None): if isinstance(element_list, tuple): # Allow this for streaming update compatibility while fixing BEAM-24535. - element_list = [element_list] + self.process_one(element_list, job_name_prefix) + else: + for element in element_list: + self.process_one(element, job_name_prefix) - if not element_list: - return + def process_one(self, element, job_name_prefix): + destination, job_reference = element - first_destination = element_list[0][0] - copy_to_reference = bigquery_tools.parse_table_reference(first_destination) + copy_to_reference = bigquery_tools.parse_table_reference(destination) if copy_to_reference.projectId is None: copy_to_reference.projectId = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project - copy_from_references = [] - for destination, job_reference in element_list: - copy_from_reference = bigquery_tools.parse_table_reference(destination) - copy_from_reference.tableId = job_reference.jobId - if copy_from_reference.projectId is None: - copy_from_reference.projectId = vp.RuntimeValueProvider.get_value( - 'project', str, '') or self.project - copy_from_references.append(copy_from_reference) + copy_from_reference = bigquery_tools.parse_table_reference(destination) + copy_from_reference.tableId = job_reference.jobId + if copy_from_reference.projectId is None: + copy_from_reference.projectId = vp.RuntimeValueProvider.get_value( + 'project', str, '') or self.project - full_table_ref = bigquery_tools.get_hashable_destination(copy_to_reference) + _LOGGER.info( + "Triggering copy job from %s to %s", + copy_from_reference, + copy_to_reference) - is_first_time = full_table_ref not in self._observed_tables - if is_first_time: - self._observed_tables.add(full_table_ref) - if self.bq_io_metadata: - Lineage.sinks().add( - 'bigquery', - copy_to_reference.projectId, - copy_to_reference.datasetId, - copy_to_reference.tableId) - - # Split into chunks of MAX_SOURCES_PER_COPY_JOB - chunks = [ - copy_from_references[i:i + self.MAX_SOURCES_PER_COPY_JOB] - for i in range( - 0, len(copy_from_references), self.MAX_SOURCES_PER_COPY_JOB) - ] - - copy_job_name_base = '%s_%s' % ( - job_name_prefix, - _bq_uuid(bigquery_tools.get_hashable_destination(copy_to_reference))) + wait_for_job, write_disposition = ( + self._determine_write_disposition(copy_to_reference)) + + if not self.bq_io_metadata: + self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) project_id = ( copy_to_reference.projectId if self.load_job_project_id is None else self.load_job_project_id) + copy_job_name = '%s_%s' % ( + job_name_prefix, + _bq_uuid( + '%s:%s.%s' % ( + copy_from_reference.projectId, + copy_from_reference.datasetId, + copy_from_reference.tableId))) + job_reference = self.bq_wrapper._insert_copy_job( + project_id, + copy_job_name, + copy_from_reference, + copy_to_reference, + create_disposition=self.create_disposition, + write_disposition=write_disposition, + job_labels=self.bq_io_metadata.add_additional_bq_job_labels()) + + if wait_for_job: + self.bq_wrapper.wait_for_bq_job(job_reference, sleep_duration_sec=10) + self.pending_jobs.append( + GlobalWindows.windowed_value((destination, job_reference))) - for i, chunk in enumerate(chunks): - if i == 0 and is_first_time: - write_disposition = self.write_disposition - # Wait inline only if we have multiple chunks and write disposition is WRITE_TRUNCATE or WRITE_EMPTY. - # This ensures the first chunk initializes the table, and subsequent chunks (WRITE_APPEND) append to it. - wait_for_job = ( - self.write_disposition in ('WRITE_TRUNCATE', 'WRITE_EMPTY') and - len(chunks) > 1) - else: - write_disposition = 'WRITE_APPEND' - wait_for_job = False - - chunk_job_name = copy_job_name_base - if len(chunks) > 1: - chunk_job_name = f"{copy_job_name_base}_{i}" - - _LOGGER.info( - "Triggering copy job %s from %s to %s (write_disposition: %s)", - chunk_job_name, [str(r) for r in chunk], - copy_to_reference, - write_disposition) - - job_reference = self.bq_wrapper._insert_copy_job( - project_id, - chunk_job_name, - chunk, - copy_to_reference, - create_disposition=self.create_disposition, - write_disposition=write_disposition, - job_labels=self.bq_io_metadata.add_additional_bq_job_labels() - if self.bq_io_metadata else None) - - if wait_for_job: - self.bq_wrapper.wait_for_bq_job(job_reference, sleep_duration_sec=10) - - self.pending_jobs.append( - GlobalWindows.windowed_value((first_destination, job_reference))) + def _determine_write_disposition(self, copy_to_reference) -> tuple[bool, str]: + """ + Determines the write disposition for a BigQuery copy job, + based on destination. + + When the write_disposition for a job is WRITE_TRUNCATE, multiple copy jobs + to the same destination can interfere with each other, truncate data, and + write to the BigQuery table repeatedly. To prevent this, the first copy job + runs with the user's specified write_disposition, but subsequent jobs must + always use WRITE_APPEND. This ensures that subsequent copy jobs do not + clear out data appended by previous jobs. + + Args: + copy_to_reference: The reference to the destination table. + + Returns: + A tuple containing a boolean indicating whether to wait for the job to + complete and the write disposition to use for the job. + """ + full_table_ref = '%s:%s.%s' % ( + copy_to_reference.projectId, + copy_to_reference.datasetId, + copy_to_reference.tableId) + if full_table_ref not in self._observed_tables: + write_disposition = self.write_disposition + wait_for_job = True + self._observed_tables.add(full_table_ref) + Lineage.sinks().add( + 'bigquery', + copy_to_reference.projectId, + copy_to_reference.datasetId, + copy_to_reference.tableId) + else: + wait_for_job = False + write_disposition = 'WRITE_APPEND' + return wait_for_job, write_disposition def finish_bundle(self): for windowed_value in self.pending_jobs: @@ -740,7 +744,7 @@ def process( else: try: schema = bigquery_tools.table_schema_to_dict( - self.bq_wrapper.get_table( + bigquery_tools.BigQueryWrapper().get_table( project_id=table_reference.projectId, dataset_id=table_reference.datasetId, table_id=table_reference.tableId).schema) @@ -851,8 +855,7 @@ def process(self, element): if latest_partition.can_accept(file_size): latest_partition.add(file_path, file_size) else: - if latest_partition.files: - partitions.append(latest_partition.files) + partitions.append(latest_partition.files) latest_partition = PartitionFiles.Partition( self.max_partition_size, self.max_files_per_partition) latest_partition.add(file_path, file_size) @@ -1178,13 +1181,12 @@ def _load_data( # the truncation happens only once. See # https://github.com/apache/beam/issues/24535. finished_temp_tables_load_job_ids_list_pc = ( - finished_temp_tables_load_job_ids_pc - | beam.MapTuple( + finished_temp_tables_load_job_ids_pc | beam.MapTuple( lambda destination, job_reference: ( - bigquery_tools.get_hashable_destination(destination), + bigquery_tools.parse_table_reference(destination).tableId, (destination, job_reference))) | beam.GroupByKey() - | beam.MapTuple(lambda dest, batch: list(batch))) + | beam.MapTuple(lambda tableId, batch: list(batch))) else: # Loads can happen in parallel. finished_temp_tables_load_job_ids_list_pc = ( diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py index 47c1ce5ea1bb..191719e6a208 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py @@ -924,180 +924,69 @@ def dynamic_destination_resolver(element, *side_inputs): write_disposition=BigQueryDisposition.WRITE_TRUNCATE)) from apache_beam.io.gcp.internal.clients.bigquery import TableReference - mock_insert_copy_job.assert_has_calls([ - call( - 'project1', - mock.ANY, - [ + mock_insert_copy_job.assert_has_calls( + [ + call( + 'project1', + mock.ANY, TableReference( datasetId='dataset1', projectId='project1', tableId='job_name1'), + TableReference( + datasetId='dataset1', + projectId='project1', + tableId='table1'), + create_disposition=None, + write_disposition='WRITE_TRUNCATE', + job_labels={'step_name': 'bigquerybatchfileloads'}), + call( + 'project1', + mock.ANY, TableReference( datasetId='dataset1', projectId='project1', tableId='job_name1'), - ], - TableReference( - datasetId='dataset1', projectId='project1', tableId='table1'), - create_disposition=None, - write_disposition='WRITE_TRUNCATE', - job_labels={'step_name': 'bigquerybatchfileloads'}), - call( - 'project1', - mock.ANY, - [ + TableReference( + datasetId='dataset1', + projectId='project1', + tableId='table1'), + create_disposition=None, + write_disposition='WRITE_APPEND', + job_labels={'step_name': 'bigquerybatchfileloads'}), + call( + 'project1', + mock.ANY, TableReference( datasetId='dataset2', projectId='project1', tableId='job_name1'), - ], - TableReference( - datasetId='dataset2', projectId='project1', tableId='table1'), - create_disposition=None, - write_disposition='WRITE_TRUNCATE', - job_labels={'step_name': 'bigquerybatchfileloads'}), - call( - 'project1', - mock.ANY, - [ + TableReference( + datasetId='dataset2', + projectId='project1', + tableId='table1'), + create_disposition=None, + # Previously this was `WRITE_APPEND`. + write_disposition='WRITE_TRUNCATE', + job_labels={'step_name': 'bigquerybatchfileloads'}), + call( + 'project1', + mock.ANY, TableReference( datasetId='dataset3', projectId='project1', tableId='job_name1'), - ], - TableReference( - datasetId='dataset3', projectId='project1', tableId='table1'), - create_disposition=None, - write_disposition='WRITE_TRUNCATE', - job_labels={'step_name': 'bigquerybatchfileloads'}), - ], - any_order=True) - self.assertEqual(3, mock_insert_copy_job.call_count) - - @mock.patch( - 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper.wait_for_bq_job') - @mock.patch( - 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._insert_copy_job') - def test_copy_jobs_splitting( - self, mock_insert_copy_job, mock_wait_for_bq_job): - destination = 'project1:dataset1.table1' - - from apache_beam.io.gcp.bigquery_file_loads import TriggerCopyJobs - original_max_sources = TriggerCopyJobs.MAX_SOURCES_PER_COPY_JOB - TriggerCopyJobs.MAX_SOURCES_PER_COPY_JOB = 2 - - try: - job_reference = bigquery_api.JobReference() - job_reference.projectId = 'project1' - job_reference.jobId = 'job_name1' - result_job = mock.Mock() - result_job.jobReference = job_reference - - mock_job = mock.Mock() - mock_job.status.state = 'DONE' - mock_job.status.errorResult = None - mock_job.jobReference = job_reference - - bq_client = mock.Mock() - bq_client.jobs.Get.return_value = mock_job - bq_client.jobs.Insert.return_value = result_job - bq_client.tables.Delete.return_value = None - mock_insert_copy_job.return_value = job_reference - temp_dir = self._new_tempdir() - - with TestPipeline('FnApiRunner') as p: - _ = ( - p - | beam.Create([ - { - 'name': 'a' - }, - { - 'name': 'b' - }, - { - 'name': 'c' - }, - { - 'name': 'd' - }, - { - 'name': 'e' - }, - ], - reshuffle=False) - | bqfl.BigQueryBatchFileLoads( - destination, - custom_gcs_temp_location=temp_dir, - test_client=bq_client, - validate=False, - temp_file_format=bigquery_tools.FileFormat.JSON, - max_file_size=10, - max_partition_size=10, - max_files_per_partition=1, - write_disposition=BigQueryDisposition.WRITE_TRUNCATE)) - - self.assertEqual(3, mock_insert_copy_job.call_count) - - from apache_beam.io.gcp.internal.clients.bigquery import TableReference - expected_calls = [ - call( - 'project1', - mock.ANY, - [ - TableReference( - datasetId='dataset1', - projectId='project1', - tableId='job_name1'), - TableReference( - datasetId='dataset1', - projectId='project1', - tableId='job_name1'), - ], - TableReference( - datasetId='dataset1', projectId='project1', tableId='table1'), - create_disposition=None, - write_disposition='WRITE_TRUNCATE', - job_labels=mock.ANY), - call( - 'project1', - mock.ANY, - [ - TableReference( - datasetId='dataset1', - projectId='project1', - tableId='job_name1'), - TableReference( - datasetId='dataset1', - projectId='project1', - tableId='job_name1'), - ], - TableReference( - datasetId='dataset1', projectId='project1', tableId='table1'), - create_disposition=None, - write_disposition='WRITE_APPEND', - job_labels=mock.ANY), - call( - 'project1', - mock.ANY, - [ - TableReference( - datasetId='dataset1', - projectId='project1', - tableId='job_name1'), - ], - TableReference( - datasetId='dataset1', projectId='project1', tableId='table1'), - create_disposition=None, - write_disposition='WRITE_APPEND', - job_labels=mock.ANY), - ] - mock_insert_copy_job.assert_has_calls(expected_calls, any_order=True) - self.assertEqual(9, mock_wait_for_bq_job.call_count) - - finally: - TriggerCopyJobs.MAX_SOURCES_PER_COPY_JOB = original_max_sources + TableReference( + datasetId='dataset3', + projectId='project1', + tableId='table1'), + create_disposition=None, + # Previously this was `WRITE_APPEND`. + write_disposition='WRITE_TRUNCATE', + job_labels={'step_name': 'bigquerybatchfileloads'}), + ], + any_order=True) + self.assertEqual(4, mock_insert_copy_job.call_count) @parameterized.expand([ param(is_streaming=False, with_auto_sharding=False, compat_version=None), diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 491b7a39b0b7..8dd58cd55a01 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -506,22 +506,16 @@ def _insert_copy_job( reference = bigquery.JobReference() reference.jobId = job_id reference.projectId = project_id - - copy_config = bigquery.JobConfigurationTableCopy( - destinationTable=to_table_reference, - createDisposition=create_disposition, - writeDisposition=write_disposition, - ) - if isinstance(from_table_reference, list): - copy_config.sourceTables = from_table_reference - else: - copy_config.sourceTable = from_table_reference - request = bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=bigquery.Job( configuration=bigquery.JobConfiguration( - copy=copy_config, + copy=bigquery.JobConfigurationTableCopy( + destinationTable=to_table_reference, + sourceTable=from_table_reference, + createDisposition=create_disposition, + writeDisposition=write_disposition, + ), labels=_build_job_labels(job_labels), ), jobReference=reference, diff --git a/sdks/python/apache_beam/io/iobase.py b/sdks/python/apache_beam/io/iobase.py index afc977406af0..aa03280050fa 100644 --- a/sdks/python/apache_beam/io/iobase.py +++ b/sdks/python/apache_beam/io/iobase.py @@ -918,7 +918,10 @@ def __init__(self, source: SourceBase) -> None: """Initializes a Read transform. Args: - source: Data source to read from. + source: the data source to read from. May be a ``BoundedSource``, an + ``UnboundedSource``, or a ``PTransform`` (which is applied directly). + For any other source ``Read`` is treated as a primitive and relayed to + the runner implementation. """ super().__init__() self.source = source @@ -944,6 +947,11 @@ def expand(self, pbegin): | 'EmitSource' >> core.Map(lambda _: self.source).with_output_types(BoundedSource) | SDFBoundedSourceReader(display_data)) + # Local import to avoid a circular dependency. + from apache_beam.io.unbounded_source import ReadFromUnboundedSource + from apache_beam.io.unbounded_source import UnboundedSource + if isinstance(self.source, UnboundedSource): + return pbegin | ReadFromUnboundedSource(self.source) elif isinstance(self.source, ptransform.PTransform): # The Read transform can also admit a full PTransform as an input # rather than an anctual source. If the input is a PTransform, then @@ -993,6 +1001,10 @@ def to_runner_api_parameter( is_bounded=beam_runner_api_pb2.IsBounded.BOUNDED if self.source.is_bounded() else beam_runner_api_pb2.IsBounded.UNBOUNDED)) + # Local import to avoid a circular dependency. + from apache_beam.io.unbounded_source import UnboundedSource + if isinstance(self.source, UnboundedSource): + return super().to_runner_api_parameter(context) elif isinstance(self.source, ptransform.PTransform): return self.source.to_runner_api_parameter(context) raise NotImplementedError( diff --git a/sdks/python/apache_beam/io/iobase_test.py b/sdks/python/apache_beam/io/iobase_test.py index eb9617cfae34..dbedf4681f42 100644 --- a/sdks/python/apache_beam/io/iobase_test.py +++ b/sdks/python/apache_beam/io/iobase_test.py @@ -21,15 +21,16 @@ import unittest -import mock - import apache_beam as beam -from apache_beam.io.concat_source import ConcatSource -from apache_beam.io.concat_source_test import RangeSource +import mock from apache_beam.io import iobase from apache_beam.io import range_trackers +from apache_beam.io.concat_source import ConcatSource +from apache_beam.io.concat_source_test import RangeSource from apache_beam.io.iobase import SourceBundle from apache_beam.options.pipeline_options import DebugOptions +from apache_beam.portability import common_urns +from apache_beam.portability import python_urns from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to @@ -220,5 +221,45 @@ def test_sdf_wrap_range_source(self): self._run_sdf_wrapper_pipeline(RangeSource(0, 4), [0, 1, 2, 3]) +class UseSdfUnboundedSourcesTests(unittest.TestCase): + """Covers the UnboundedSource branch in + ``iobase.Read.expand()``. Uses ``UnboundedCountingSource`` from + ``unbounded_source_test`` as a finite fake source (no network). + """ + def test_read_end_to_end_unbounded(self): + from apache_beam.io.unbounded_source_test import UnboundedCountingSource + with beam.Pipeline() as p: + out = p | beam.io.Read(UnboundedCountingSource(5)) + assert_that(out, equal_to([0, 1, 2, 3, 4])) + + def test_read_unbounded_pcollection_is_unbounded(self): + from apache_beam.io.unbounded_source_test import UnboundedCountingSource + p = beam.Pipeline() + out = p | beam.io.Read(UnboundedCountingSource(3)) + self.assertFalse(out.is_bounded) + + def test_read_unbounded_serializes_as_expanded_composite(self): + from apache_beam.io.unbounded_source_test import UnboundedCountingSource + p = beam.Pipeline() + p | 'ReadIt' >> beam.io.Read(UnboundedCountingSource(3)) + + proto = p.to_runner_api(use_fake_coders=True) + transforms = proto.components.transforms.values() + deprecated_reads = [ + transform.unique_name for transform in transforms + if transform.spec.urn == common_urns.deprecated_primitives.READ.urn + ] + read_transforms = [ + transform for transform in proto.components.transforms.values() + if transform.unique_name == 'ReadIt' + ] + + self.assertEqual([], deprecated_reads) + self.assertEqual(1, len(read_transforms)) + self.assertEqual( + python_urns.GENERIC_COMPOSITE_TRANSFORM, read_transforms[0].spec.urn) + self.assertTrue(read_transforms[0].subtransforms) + + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/io/unbounded_source.py b/sdks/python/apache_beam/io/unbounded_source.py new file mode 100644 index 000000000000..3246602252e3 --- /dev/null +++ b/sdks/python/apache_beam/io/unbounded_source.py @@ -0,0 +1,996 @@ +# +# 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. +# + +"""Experimental ``UnboundedSource`` support for the Python SDK. + +``UnboundedSource`` support is currently experimental in the Python SDK; the +API may change in backwards-incompatible ways. + +An unbounded source reads an effectively infinite stream of records (message +queues, change-data-capture feeds, and similar) with checkpoint-based +pause/resume, watermark reporting, and bundle finalization. + +To define a source, implement :class:`UnboundedSource`, an +:class:`UnboundedReader`, and (when the reader has a resumable position) a +:class:`CheckpointMark`:: + + import apache_beam as beam + from apache_beam.io.unbounded_source import ( + CheckpointMark, UnboundedReader, UnboundedSource) + from apache_beam.utils.timestamp import MAX_TIMESTAMP + + class MyCheckpointMark(CheckpointMark): + def __init__(self, position): + self.position = position + + def finalize_checkpoint(self): + # Commit/acknowledge records up to ``position`` upstream, e.g. ack the + # consumed messages on a queue. + ... + + class MyReader(UnboundedReader): + def start(self): + # Position at the first record; return whether one is available. + ... + + def advance(self): + # Move to the next record; ``False`` means no data is available now. + ... + + def get_current(self): + ... + + def get_current_timestamp(self): + ... # event time of the current record + + def get_watermark(self): + # Lower bound on the timestamps of future records. Return + # ``MAX_TIMESTAMP`` to signal the reader has permanently finished. + ... + + def get_checkpoint_mark(self): + return MyCheckpointMark(...) + + class MySource(UnboundedSource): + def split(self, desired_num_splits, options=None): + # Return independent sub-sources, or ``[self]`` when not splittable. + return [self] + + def create_reader(self, options, checkpoint_mark): + # Build a reader; resume after ``checkpoint_mark`` when it is not None. + return MyReader(...) + + def get_checkpoint_mark_coder(self): + return ... # a Coder for MyCheckpointMark + +Read the source in a pipeline with :class:`apache_beam.io.Read`:: + + with beam.Pipeline() as p: + p | beam.io.Read(MySource()) | beam.Map(print) +""" + +import collections +import dataclasses +import logging +import threading +import time +from typing import Any +from typing import Callable +from typing import Iterable +from typing import Optional + +from apache_beam import coders +from apache_beam.coders.coders import BooleanCoder +from apache_beam.coders.coders import Coder +from apache_beam.coders.coders import NullableCoder +from apache_beam.coders.coders import TimestampCoder +from apache_beam.coders.coders import TupleCoder +from apache_beam.coders.coders import _MemoizingPickleCoder +from apache_beam.io import iobase +from apache_beam.io.watermark_estimators import ManualWatermarkEstimator +from apache_beam.runners import sdf_utils +from apache_beam.transforms import PTransform +from apache_beam.transforms import core +from apache_beam.transforms.window import TimestampedValue +from apache_beam.utils.timestamp import MAX_TIMESTAMP +from apache_beam.utils.timestamp import MIN_TIMESTAMP +from apache_beam.utils.timestamp import Duration +from apache_beam.utils.timestamp import Timestamp + +__all__ = [ + 'CheckpointMark', + 'UnboundedReader', + 'UnboundedSource', + 'ReadFromUnboundedSource', +] + +_LOGGER = logging.getLogger(__name__) + +# Sentinel used when a reader has no data available right now. +# This is distinct from end-of-stream. +_NO_DATA = object() + +_DEFAULT_POLL_INTERVAL_SECONDS = 1.0 +_DEFAULT_DESIRED_NUM_SPLITS = 20 +_DEFAULT_MAX_RECORDS_PER_BUNDLE = 10000 +_DEFAULT_MAX_READ_TIME_SECONDS = 10.0 +# A reader parked by a residual that never resumes is closed once idle this +# long; the cache also caps its entry count as a memory backstop. +_DEFAULT_READER_CACHE_IDLE_SECONDS = 60.0 +_DEFAULT_READER_CACHE_MAX_SIZE = 100 + +# Encodes a source to a structural cache key. Internally consistent across park +# and acquire; need not match the restriction wire coder. +_SOURCE_KEY_CODER = _MemoizingPickleCoder() + +# ------------------------------------------------------------------------------ +# Public abstract base classes. +# ------------------------------------------------------------------------------ + + +class CheckpointMark(object): + """A durable, serializable position in an :class:`UnboundedSource`. + + Produced by :meth:`UnboundedReader.get_checkpoint_mark`, encoded with + :meth:`UnboundedSource.get_checkpoint_mark_coder`, and used to resume a reader + (see :meth:`UnboundedSource.create_reader`). + """ + def finalize_checkpoint(self) -> None: + """Called once the runner has durably committed work up to this mark. + + Override to acknowledge/commit upstream (for example, ack the consumed + messages on a queue). The default is a no-op. + + The runner calls this at most once for a committed checkpoint mark. + Finalization is best effort; a mark may never be finalized. An exception + raised here is logged. On bundle retry an uncommitted mark may be re-cut + over an overlapping span, so this method must be idempotent (acknowledge by + absolute position). + """ + pass + + +class UnboundedReader(object): + """Reads records from an :class:`UnboundedSource`. + + Lifecycle: exactly one :meth:`start`, then any number of :meth:`advance` + calls; whenever one returns ``True`` the current record is available via + :meth:`get_current` / :meth:`get_current_timestamp`. A ``False`` return means + "no data available right now", which is distinct from end-of-stream: a reader + signals a permanent end by reporting a watermark of ``MAX_TIMESTAMP``. + """ + def start(self) -> bool: + """Positions at the first record; returns whether one is available.""" + raise NotImplementedError + + def advance(self) -> bool: + """Advances to the next record. ``False`` means no data is available now. + + Should not block. The wrapper enforces the per-bundle record and time caps + only between records, so a blocking ``start``/``advance`` can overrun the + time cap and stall the bundle. Return ``False`` when no data is currently + available instead of waiting. + """ + raise NotImplementedError + + def get_current(self) -> Any: + """Returns the record claimed by the last successful start/advance.""" + raise NotImplementedError + + def get_current_timestamp(self) -> Timestamp: + """Returns the event-time timestamp of the current record.""" + raise NotImplementedError + + def get_watermark(self) -> Timestamp: + """An approximate lower bound on timestamps of future records. + + Treated as monotonic by the wrapper. Return ``MAX_TIMESTAMP`` to signal that + this reader has permanently finished. + """ + raise NotImplementedError + + def get_checkpoint_mark(self) -> CheckpointMark: + """Returns a durable mark to resume from. Call only at a bundle boundary.""" + raise NotImplementedError + + def close(self) -> None: + """Releases reader resources. Default no-op.""" + pass + + +class UnboundedSource(iobase.SourceBase): + """A source producing an unbounded stream of records with checkpointing. + + Read it in a pipeline with :class:`apache_beam.io.Read`:: + + p | beam.io.Read(MyUnboundedSource()) + """ + def split(self, + desired_num_splits: int, + options: Optional[Any] = None) -> Iterable['UnboundedSource']: + """Splits into at most ``desired_num_splits`` independent sub-sources. + + Each returned sub-source must be independent and must not share mutable + state with siblings (the runner may execute them concurrently across + workers). Return ``[self]`` if the source cannot be split. Splitting is + performed once, before any checkpoint exists; once a reader has + checkpointed, the restriction is kept intact. + """ + raise NotImplementedError + + def create_reader( + self, options: Optional[Any], + checkpoint_mark: Optional[CheckpointMark]) -> UnboundedReader: + """Creates a reader, optionally resuming from ``checkpoint_mark``. + + Contract: + * When ``checkpoint_mark`` is ``None``, the returned reader's ``start()`` + produces the very first record of the source (or returns ``False`` if + none yet). + * When ``checkpoint_mark`` is not ``None``, the returned reader's + ``start()`` produces the first record strictly after the position + encoded by ``checkpoint_mark``. The reader must not re-deliver records + already covered by the prior bundle. + """ + raise NotImplementedError + + def get_checkpoint_mark_coder(self) -> Coder: + """Returns the coder for this source's :class:`CheckpointMark` instances. + + The SDK may call this while encoding or decoding source restrictions. + Implementations should be deterministic, side-effect free, and should not + perform I/O. + """ + raise NotImplementedError( + '%s must override get_checkpoint_mark_coder() to return a Coder for ' + 'its CheckpointMark subclass.' % type(self).__name__) + + def is_bounded(self) -> bool: + # SourceBase.is_bounded raises; an unbounded source is, by definition, not. + return False + + def default_output_coder(self) -> Coder: + # Permissive default; override for a tighter coder. + return coders.registry.get_coder(object) + + +# ------------------------------------------------------------------------------ +# SDF wrapper internals: a private implementation detail of +# ReadFromUnboundedSource. +# ------------------------------------------------------------------------------ + + +@dataclasses.dataclass(frozen=True) +class _UnboundedSourceRestriction(object): + """Durable SDF restriction describing where a reader should (re)start. + + Holds only serializable state -- never a live reader. ``is_done`` marks the + terminal (MAX-watermark) transition. ``finalization_checkpoint_mark`` is kept + separate from ``checkpoint_mark`` so a done primary can carry a commit hook + without polluting the resume state. + + Field roles: + * ``checkpoint_mark`` -- RESUME state. A reader rebuilt from this mark + must produce the first record strictly after it. + * ``finalization_checkpoint_mark`` -- COMMIT hook. Only set on a done + primary that was just cut this bundle. Registered with the runner's + bundle finalizer to acknowledge upstream. Independent of + ``checkpoint_mark`` so a residual's resume state can be ``None`` while + still recording what should be committed. + """ + source: UnboundedSource + checkpoint_mark: Optional[CheckpointMark] = None + watermark: Timestamp = MIN_TIMESTAMP + is_done: bool = False + finalization_checkpoint_mark: Optional[CheckpointMark] = None + + +class _UnboundedSourceRestrictionCoder(Coder): + """Encodes :class:`_UnboundedSourceRestriction` as a fixed 5-tuple. + + Stateless: at encode time the source's own + :meth:`UnboundedSource.get_checkpoint_mark_coder` is looked up from the + restriction; at decode time the source is decoded first and its coder + drives the checkpoint-mark decoding. This avoids passing source-specific + coder state into the coder's constructor, which in turn lets + :class:`_UnboundedSourceRestrictionProvider` and + :class:`_ReadFromUnboundedSourceDoFn` be module-level classes. + + Wire shape: source_bytes / checkpoint_bytes / watermark / done / + finalization_checkpoint_bytes -- the checkpoint and finalization bytes + are independently encoded with the (source-declared) checkpoint coder + wrapped in :class:`NullableCoder`. + """ + def __init__(self): + self._source_coder = _MemoizingPickleCoder() + self._bytes_coder = coders.BytesCoder() + self._tuple_coder = TupleCoder(( + self._bytes_coder, # source (pickled bytes) + self._bytes_coder, # checkpoint_mark (nullable-encoded bytes) + TimestampCoder(), # watermark + BooleanCoder(), # is_done + self._bytes_coder)) # finalization_checkpoint_mark (nullable-encoded) + + def _checkpoint_coder(self, source: UnboundedSource) -> Coder: + return NullableCoder(source.get_checkpoint_mark_coder()) + + def encode(self, restriction: '_UnboundedSourceRestriction') -> bytes: + source_bytes = self._source_coder.encode(restriction.source) + cp_coder = self._checkpoint_coder(restriction.source) + return self._tuple_coder.encode(( + source_bytes, + cp_coder.encode(restriction.checkpoint_mark), + restriction.watermark, + restriction.is_done, + cp_coder.encode(restriction.finalization_checkpoint_mark))) + + def decode(self, encoded: bytes) -> '_UnboundedSourceRestriction': + (source_bytes, checkpoint_bytes, watermark, is_done, + finalization_bytes) = self._tuple_coder.decode(encoded) + source = self._source_coder.decode(source_bytes) + cp_coder = self._checkpoint_coder(source) + return _UnboundedSourceRestriction( + source=source, + checkpoint_mark=cp_coder.decode(checkpoint_bytes), + watermark=watermark, + is_done=is_done, + finalization_checkpoint_mark=cp_coder.decode(finalization_bytes)) + + def is_deterministic(self) -> bool: + # Pickled source and checkpoint are not guaranteed deterministic. + return False + + +class _ReaderCache(object): + """Holds live readers between an SDF self-checkpoint and its resume. + + A fresh tracker is built for every bundle, so a reader parked at a + self-checkpoint would otherwise be closed and rebuilt from its checkpoint + mark on the next bundle. Parking it here lets the resuming bundle reclaim the + same started reader, keyed by the residual's structural ``(source, checkpoint + mark)``. ``acquire`` removes the entry, so two trackers never drive one + reader. + + A residual may be reassigned to another worker and never resume here; such a + reader is released once it falls idle past ``idle_seconds`` or when the entry + count exceeds ``max_size``, and the owning DoFn's teardown releases the rest. + One DoFn instance drives several trackers across threads, so access is locked. + """ + def __init__( + self, + idle_seconds: float = _DEFAULT_READER_CACHE_IDLE_SECONDS, + max_size: int = _DEFAULT_READER_CACHE_MAX_SIZE, + now: Optional[Callable[[], float]] = None): + self._idle_seconds = idle_seconds + self._max_size = max_size + self._now = now or time.monotonic + self._lock = threading.Lock() + # key -> (reader, started, parked_at); ordered oldest-first for eviction. + self._entries = collections.OrderedDict( + ) # type: collections.OrderedDict[Any, tuple[UnboundedReader, bool, float]] + + def acquire(self, key: Any) -> Optional[tuple['UnboundedReader', bool]]: + """Removes and returns ``(reader, started)`` for ``key``, or None.""" + with self._lock: + entry = self._entries.pop(key, None) + stale = self._evict_idle() + self._close_readers(stale) + if entry is None: + return None + return entry[0], entry[1] + + def park(self, key: Any, reader: 'UnboundedReader', started: bool) -> None: + """Stores ``reader`` under ``key`` for a later bundle to reclaim. A reader + already parked under ``key`` is closed.""" + with self._lock: + replaced = self._entries.pop(key, None) + self._entries[key] = (reader, started, self._now()) + stale = self._evict_idle() + while len(self._entries) > self._max_size: + _, oldest = self._entries.popitem(last=False) + stale.append(oldest[0]) + if replaced is not None and replaced[0] is not reader: + stale.append(replaced[0]) + self._close_readers(stale) + + def close_all(self) -> None: + """Closes every parked reader. Called from the owning DoFn's teardown.""" + with self._lock: + entries = list(self._entries.values()) + self._entries.clear() + self._close_readers(entry[0] for entry in entries) + + def _evict_idle(self) -> list: + # Caller holds the lock. Pops entries idle past the window (oldest first) + # and returns their readers for the caller to close after unlocking. + deadline = self._now() - self._idle_seconds + stale = [] + while self._entries: + key, entry = next(iter(self._entries.items())) + if entry[2] > deadline: + break + del self._entries[key] + stale.append(entry[0]) + return stale + + def _close_readers(self, readers) -> None: + for reader in readers: + try: + reader.close() + except Exception: # pylint: disable=broad-except + _LOGGER.warning('Error closing UnboundedReader', exc_info=True) + + +class _UnboundedSourceRestrictionTracker(iobase.RestrictionTracker): + """Drives an :class:`UnboundedReader` for one SDF restriction. + + Owns the live reader (lazily created, never serialized): both runner-initiated + ``defer_remainder`` self-checkpoints with ``try_split(0)``, which must + checkpoint the live reader. + + A self-checkpoint parks the reader in the DoFn's :class:`_ReaderCache` for the + next bundle to reclaim, keeping one started reader alive across bundles. The + DoFn injects ``_reader_cache`` at the start of ``process()``; with no cache + the tracker builds a fresh reader each bundle. + + ``process()`` only sees a ``RestrictionTrackerView``, which hides custom + methods and whose ``try_claim`` returns just a bool, so the freshly-read + record is handed back through a one-element holder list passed as the + ``try_claim`` *position* argument. + """ + def __init__( + self, + restriction: _UnboundedSourceRestriction, + options: Optional[Any] = None): + self._restriction = restriction + self._options = options + self._reader = None # type: Optional[UnboundedReader] + self._started = False + # True once a checkpoint has been cut this bundle (EOF or self-checkpoint). + self._checkpoint_taken = False + # Cross-bundle reader cache, injected by the DoFn; None disables caching. + self._reader_cache = None # type: Optional[_ReaderCache] + + def _ensure_reader(self) -> None: + if self._reader is not None: + return + cached = self._acquire_cached_reader() + if cached is not None: + # A parked reader is already started and positioned past its checkpoint. + self._reader, self._started = cached + return + self._reader = self._restriction.source.create_reader( + self._options, self._restriction.checkpoint_mark) + + def _cache_key(self, + restriction: _UnboundedSourceRestriction) -> Optional[Any]: + """Structural ``(source, checkpoint)`` key, or None when uncacheable. + + Built from the source pickle and the source's own checkpoint coder so a + parked reader and its resuming restriction map to the same entry. A None + key disables caching for that restriction; the resume then rebuilds from + the checkpoint mark under the source's ``create_reader`` contract. + """ + try: + source_bytes = _SOURCE_KEY_CODER.encode(restriction.source) + cp_coder = NullableCoder(restriction.source.get_checkpoint_mark_coder()) + return source_bytes, cp_coder.encode(restriction.checkpoint_mark) + except Exception: # pylint: disable=broad-except + return None + + def _acquire_cached_reader(self) -> Optional[tuple['UnboundedReader', bool]]: + if self._reader_cache is None: + return None + key = self._cache_key(self._restriction) + if key is None: + return None + return self._reader_cache.acquire(key) + + def _park_or_close_reader( + self, residual: _UnboundedSourceRestriction) -> None: + """Hands the live reader to the cache for ``residual`` to reclaim, or + closes it when no cache is available or the restriction is uncacheable.""" + if self._reader is None: + return + key = ( + self._cache_key(residual) if self._reader_cache is not None else None) + if key is None: + self._close_reader_if_open() + return + reader, self._reader = self._reader, None + self._reader_cache.park(key, reader, self._started) + + def _clone_checkpoint( + self, checkpoint: Optional[CheckpointMark]) -> Optional[CheckpointMark]: + """Returns an independent copy of a mark via the source's checkpoint coder. + + Used to keep the primary's finalize hook and the residual's resume state + from sharing one object, since a user ``finalize_checkpoint()`` may mutate + the mark. + """ + if checkpoint is None: + return None + coder = self._restriction.source.get_checkpoint_mark_coder() + return coder.decode(coder.encode(checkpoint)) + + def _close_reader_if_open(self) -> None: + """Idempotent reader release. Called by the EOF and split paths, and by + the DoFn's ``finally`` so an exception inside ``process()`` does not leak + sockets / file descriptors held by the live :class:`UnboundedReader`. + """ + if self._reader is None: + return + try: + self._reader.close() + except Exception: # pylint: disable=broad-except + _LOGGER.warning('Error closing UnboundedReader', exc_info=True) + finally: + self._reader = None + + def current_restriction(self) -> _UnboundedSourceRestriction: + return self._restriction + + def try_claim(self, out: list[Any]) -> bool: + """Advances the reader by one record. + + ``out[0]`` receives ``(value, record_timestamp, source_watermark)`` on the + has-data path, or the :data:`_NO_DATA` sentinel otherwise. The watermark is + the source's reported watermark, not the record's event time: the DoFn + advances the output watermark with the former and timestamps the record + with the latter. The argument doubles as the output holder (the + ``RestrictionTracker`` ABC treats it as a claim position), which the + threadsafe-tracker chain forwards opaquely. + """ + try: + return self._try_claim_inner(out) + except Exception: + # Reader state is now indeterminate; release it before re-raising. + self._close_reader_if_open() + raise + + def _try_claim_inner(self, out: list[Any]) -> bool: + if self._restriction.is_done: + out[0] = _NO_DATA + return False + self._ensure_reader() + if not self._started: + has_data = self._reader.start() + else: + has_data = self._reader.advance() + self._started = True + if has_data: + # Emit an available record before checking the watermark: a reader may + # report its last record and a MAX_TIMESTAMP watermark on the same call, + # and EOF is realized on the next data-less claim. + out[0] = ( + self._reader.get_current(), + self._reader.get_current_timestamp(), + self._reader.get_watermark()) + return True + watermark = self._reader.get_watermark() + if watermark >= MAX_TIMESTAMP: + # No data and watermark at MAX: cut a final checkpoint, close, mark done. + checkpoint = self._reader.get_checkpoint_mark() + self._close_reader_if_open() + self._restriction = dataclasses.replace( + self._restriction, + checkpoint_mark=None, # nothing left to resume from + watermark=MAX_TIMESTAMP, + is_done=True, + finalization_checkpoint_mark=checkpoint) + self._checkpoint_taken = True + out[0] = _NO_DATA + return False + # No data is available now. Refresh the watermark before deferring. + self._restriction = dataclasses.replace( + self._restriction, watermark=watermark) + out[0] = _NO_DATA + return True + + def try_split( + self, fraction_of_remainder + ) -> Optional[tuple[_UnboundedSourceRestriction, + _UnboundedSourceRestriction]]: + """Cuts a checkpoint, returning (primary, residual) or None. + + The cut checkpoint goes into ``primary.finalization_checkpoint_mark`` so + the DoFn can register a bundle-finalize callback for it. The same + checkpoint object also goes into ``residual.checkpoint_mark`` so the + resumed reader rebuilds at the correct position. The two fields are + independent on purpose (see :class:`_UnboundedSourceRestriction` + docstring): a runner that re-processes the primary alone must not see + a stale resume state, and a residual must not register finalize again + until ITS checkpoint is cut in a future bundle. + """ + try: + return self._try_split_inner(fraction_of_remainder) + except Exception: + self._close_reader_if_open() + raise + + def _try_split_inner(self, fraction_of_remainder): + # Only self-checkpoint (fraction 0) is supported; decline runner-initiated + # dynamic splits. + if fraction_of_remainder != 0: + return None + if self._reader is None or not self._started or self._restriction.is_done: + return None + checkpoint = self._reader.get_checkpoint_mark() + # The residual watermark is advisory; the SDF watermark estimator state is + # the authoritative cross-bundle watermark. + watermark = self._reader.get_watermark() + # Keep the two channels independent: the primary carries only the finalize + # hook, the residual only the resume state. The residual gets its own clone + # so a finalize_checkpoint() that mutates the primary's mark cannot corrupt + # the residual's resume position before the runner encodes it. + primary = dataclasses.replace( + self._restriction, + checkpoint_mark=None, + is_done=True, + finalization_checkpoint_mark=checkpoint) + residual = _UnboundedSourceRestriction( + source=self._restriction.source, + checkpoint_mark=self._clone_checkpoint(checkpoint), + watermark=watermark, + is_done=False, + finalization_checkpoint_mark=None) + self._restriction = primary + self._checkpoint_taken = True + # Park the reader so the resuming bundle reclaims it; on a cache miss the + # residual rebuilds one from its checkpoint mark. + self._park_or_close_reader(residual) + return primary, residual + + def check_done(self) -> bool: + # Called after every process(); must raise if work is left unaccounted for. + if self._restriction.is_done or self._checkpoint_taken: + return True + raise ValueError( + 'UnboundedSource restriction was neither finished nor checkpointed; ' + 'process() must self-checkpoint via defer_remainder() or run to EOF: ' + '%r' % (self._restriction, )) + + def current_progress(self) -> 'iobase.RestrictionProgress': + # Backlog-based progress is not implemented; report a coarse done/not-done + # signal via ``completed`` / ``remaining``. + if self._restriction.is_done: + return iobase.RestrictionProgress(completed=1.0, remaining=0.0) + return iobase.RestrictionProgress(completed=0.0, remaining=1.0) + + def is_bounded(self) -> bool: + return False + + +class _UnboundedSourceRestrictionProvider(core.RestrictionProvider): + """Wraps an :class:`UnboundedSource` element as an SDF restriction. + + Stateless module-level singleton (see :data:`_PROVIDER`): all + source-specific state (e.g. the source's checkpoint coder) is derived + per-call from the restriction's ``source`` field, which lets + :class:`_ReadFromUnboundedSourceDoFn` live at module level too. The provider + currently passes ``None`` for the ``options`` forwarded to + :meth:`UnboundedSource.split`. + """ + def __init__(self): + self._restriction_coder = _UnboundedSourceRestrictionCoder() + + def initial_restriction( + self, element: UnboundedSource) -> _UnboundedSourceRestriction: + if not isinstance(element, UnboundedSource): + raise TypeError( + 'ReadFromUnboundedSource expected an UnboundedSource element, got %r' + % (element, )) + return _UnboundedSourceRestriction(source=element) + + def create_tracker( + self, restriction: _UnboundedSourceRestriction + ) -> _UnboundedSourceRestrictionTracker: + return _UnboundedSourceRestrictionTracker(restriction) + + def split(self, element, + restriction) -> Iterable[_UnboundedSourceRestriction]: + if restriction.is_done or restriction.checkpoint_mark is not None: + yield restriction + return + + # ``source.split`` is user code and may refuse to split; fall back to a + # single restriction on error. + try: + split_sources = list( + restriction.source.split(_DEFAULT_DESIRED_NUM_SPLITS, None)) + except Exception: # pylint: disable=broad-except + _LOGGER.warning( + 'Exception while splitting UnboundedSource. Source not split.', + exc_info=True) + yield restriction + return + + if not split_sources: + yield restriction + return + + # A non-UnboundedSource split result is a contract violation, not a + # refusal, so fail loudly (outside the try/except above). + for split_source in split_sources: + if not isinstance(split_source, UnboundedSource): + raise TypeError( + 'UnboundedSource.split() produced %r, expected UnboundedSource' % + (split_source, )) + + for split_source in split_sources: + yield dataclasses.replace( + restriction, + source=split_source, + checkpoint_mark=None, + is_done=False, + finalization_checkpoint_mark=None) + + def restriction_size(self, element, restriction) -> int: + # TODO(https://github.com/apache/beam/issues/19137): implement backlog + # estimation. + return 1 + + def restriction_coder(self) -> Coder: + return self._restriction_coder + + def truncate(self, element, restriction): + # On drain, stop emitting new records. + return None + + +# Module-level stateless singleton, captured via ``RestrictionParam`` at the +# DoFn's class-definition time. +_PROVIDER = _UnboundedSourceRestrictionProvider() + + +class _FinalizeCheckpointOnce(object): + def __init__(self, checkpoint_mark: CheckpointMark): + self._checkpoint_mark = checkpoint_mark + # The lock keeps finalization idempotent if a runner ever invokes the + # callback more than once. + self._lock = threading.Lock() + self._finalized = False + + def __call__(self) -> None: + with self._lock: + if self._finalized: + return + self._finalized = True + # Finalization is best effort: log and swallow so a failing user override + # does not fail the bundle (matches CheckpointMark.finalize_checkpoint). + try: + self._checkpoint_mark.finalize_checkpoint() + except Exception: # pylint: disable=broad-except + _LOGGER.warning( + 'Error finalizing UnboundedSource checkpoint mark.', exc_info=True) + + +class _ReadFromUnboundedSourceDoFn(core.DoFn): + """SDF wrapper driving an :class:`UnboundedReader` for one restriction. + + Module-level so stdlib pickle and cloudpickle can serialise the DoFn. The + restriction provider is the module-level :data:`_PROVIDER` singleton. + """ + def __init__( + self, + poll_interval: float = _DEFAULT_POLL_INTERVAL_SECONDS, + max_records_per_bundle: int = _DEFAULT_MAX_RECORDS_PER_BUNDLE, + max_read_time_seconds: float = _DEFAULT_MAX_READ_TIME_SECONDS, + _now: Optional[Callable[[], float]] = None): + self._poll_interval = poll_interval + self._max_records_per_bundle = max_records_per_bundle + self._max_read_time_seconds = max_read_time_seconds + # Monotonic clock seam; tests inject a deterministic clock. + self._now = _now + # Per-worker reader cache; created in setup(), never serialized. + self._reader_cache = None # type: Optional[_ReaderCache] + + def setup(self): + self._reader_cache = _ReaderCache() + + def teardown(self): + if self._reader_cache is not None: + self._reader_cache.close_all() + self._reader_cache = None + + @core.DoFn.unbounded_per_element() + def process( + self, + unused_element, + bundle_finalizer=core.DoFn.BundleFinalizerParam, + tracker=core.DoFn.RestrictionParam(_PROVIDER), + watermark_estimator=core.DoFn.WatermarkEstimatorParam( + ManualWatermarkEstimator.default_provider())): + # Positional params (element, bundle finalizer) must precede the + # kwarg-injected ones (tracker, watermark estimator). + assert isinstance(tracker, sdf_utils.RestrictionTrackerView) + inner_tracker = _unwrap_tracker(tracker) + if inner_tracker is not None and self._reader_cache is not None: + # Let this bundle reclaim a reader parked by the prior bundle and re-park + # it on self-checkpoint. No cache means setup() was skipped. + inner_tracker._reader_cache = self._reader_cache + initial = tracker.current_restriction() + now = self._now or time.monotonic + records_emitted = 0 + # Armed on the first emitted record so reader startup is excluded. + read_deadline = None # type: Optional[float] + try: + while True: + holder = [None] + if not tracker.try_claim(holder): + # EOF: advance the estimator to the tracker's MAX watermark so + # downstream event-time windows can close. + _set_watermark_if_greater( + watermark_estimator, tracker.current_restriction().watermark) + break + record = holder[0] + if record is _NO_DATA: + # No data now: advance the watermark and self-checkpoint with the + # poll delay so an idle source backs off before resuming. + _set_watermark_if_greater( + watermark_estimator, tracker.current_restriction().watermark) + tracker.defer_remainder(Duration(seconds=self._poll_interval)) + break + # The third tuple field is the source watermark. The record timestamp + # remains the output event time. Emit the element before advancing the + # estimator so a reader that reports MAX on the same claim as its final + # record cannot push the output watermark past that record first. + value, record_timestamp, source_watermark = record + yield TimestampedValue(value, record_timestamp) + _set_watermark_if_greater(watermark_estimator, source_watermark) + records_emitted += 1 + if read_deadline is None: + read_deadline = now() + self._max_read_time_seconds + # A busy reader never hits the EOF or no-data branch. Bound the bundle + # by record count and elapsed time so the runner commits the checkpoint + # and runs finalization, then resume with no delay. The deadline is + # checked between records; a reader that blocks inside advance() can + # overrun it, so the record cap is the hard backstop. + reached_record_cap = records_emitted >= self._max_records_per_bundle + if reached_record_cap or now() >= read_deadline: + tracker.defer_remainder() + break + finally: + current = tracker.current_restriction() + try: + # Register finalization only when a checkpoint was cut this bundle. + # The SDK bundle finalizer applies no deadline, so finalization is + # unbounded best effort. + finalize_mark = current.finalization_checkpoint_mark + if current is not initial and finalize_mark is not None: + bundle_finalizer.register(_FinalizeCheckpointOnce(finalize_mark)) + finally: + # The EOF and self-checkpoint paths already closed or parked the + # reader, so this is a no-op there. It closes a reader still held when + # process() exits early, e.g. a downstream yield raised before any + # checkpoint. + if inner_tracker is not None: + inner_tracker._close_reader_if_open() + else: + _LOGGER.warning( + 'UnboundedSource DoFn could not close a reader because the SDF ' + 'tracker wrapper did not expose ' + '_UnboundedSourceRestrictionTracker (got %s). Reader resources ' + 'may remain open until garbage collection.', + type(tracker).__name__) + + +def _unwrap_tracker( + tracker: Any) -> Optional['_UnboundedSourceRestrictionTracker']: + """Returns the :class:`_UnboundedSourceRestrictionTracker` behind the SDF + view and threadsafe wrappers, or None when the chain is unexpected.""" + inner = tracker + if hasattr(inner, '_threadsafe_restriction_tracker'): + inner = inner._threadsafe_restriction_tracker + if hasattr(inner, '_restriction_tracker'): + inner = inner._restriction_tracker + if isinstance(inner, _UnboundedSourceRestrictionTracker): + return inner + return None + + +def _set_watermark_if_greater( + watermark_estimator, new_watermark: Timestamp) -> None: + # ManualWatermarkEstimator.set_watermark raises on regression, so only ever + # advance it (a regressing reader watermark is absorbed here). + current = watermark_estimator.current_watermark() + if current is None or new_watermark > current: + watermark_estimator.set_watermark(new_watermark) + + +class ReadFromUnboundedSource(PTransform): + """Reads an :class:`UnboundedSource`. + + Most users should prefer :class:`apache_beam.io.Read`, which dispatches an + ``UnboundedSource`` here automatically:: + + p | beam.io.Read(MyUnboundedSource()) + + Args: + source: the :class:`UnboundedSource` to read. + poll_interval: resume delay in seconds applied when the reader has no data, + which bounds how often an idle source is polled. Must be >= 0. + max_records_per_bundle: a busy reader self-checkpoints after emitting this + many records in one bundle. Must be >= 1. Defaults to 10000. + max_read_time_seconds: a busy reader self-checkpoints after this many + seconds in one bundle. Must be > 0. Defaults to 10.0. The deadline is + checked between records, so a reader that blocks inside ``advance()`` may + overrun it; ``max_records_per_bundle`` is the hard backstop. + + The bundle self-checkpoints as soon as either cap is reached. + """ + def __init__( + self, + source: UnboundedSource, + poll_interval: float = _DEFAULT_POLL_INTERVAL_SECONDS, + max_records_per_bundle: int = _DEFAULT_MAX_RECORDS_PER_BUNDLE, + max_read_time_seconds: float = _DEFAULT_MAX_READ_TIME_SECONDS): + if not isinstance(source, UnboundedSource): + raise TypeError('source must be an UnboundedSource, got %r' % (source, )) + if max_records_per_bundle < 1: + raise ValueError( + 'max_records_per_bundle must be >= 1, got %r' % + (max_records_per_bundle, )) + if max_read_time_seconds <= 0: + raise ValueError( + 'max_read_time_seconds must be > 0, got %r' % + (max_read_time_seconds, )) + if poll_interval < 0: + raise ValueError( + 'poll_interval must be >= 0, got %r' % (poll_interval, )) + super().__init__() + self._source = source + self._poll_interval = poll_interval + self._max_records_per_bundle = max_records_per_bundle + self._max_read_time_seconds = max_read_time_seconds + + def expand(self, pbegin): + source = self._source + output_coder = source.default_output_coder() + # The source is the SDF element used to derive the initial restriction. + # process() reads from the restriction, so it does not use the element + # directly. + output = ( + pbegin + | 'Create' >> core.Create([source]) + | 'ReadUnbounded' >> core.ParDo( + _ReadFromUnboundedSourceDoFn( + self._poll_interval, + self._max_records_per_bundle, + self._max_read_time_seconds))) + # Surface an element type only when the global registry already maps it to + # an equivalent coder. Avoid mutating ``coders.registry`` for a + # parameterized coder whose instance state would be lost. + try: + type_hint = output_coder.to_type_hint() + except NotImplementedError: + type_hint = None + if type_hint is not None: + try: + registered_coder = coders.registry.get_coder(type_hint) + except Exception: # pylint: disable=broad-except + _LOGGER.warning( + 'Could not look up the registered coder for element type %s.', + type_hint, + exc_info=True) + else: + if registered_coder == output_coder: + output.element_type = type_hint + return output + + def _infer_output_coder(self, input_type=None, input_coder=None): + return self._source.default_output_coder() diff --git a/sdks/python/apache_beam/io/unbounded_source_test.py b/sdks/python/apache_beam/io/unbounded_source_test.py new file mode 100644 index 000000000000..3b63d0d2de98 --- /dev/null +++ b/sdks/python/apache_beam/io/unbounded_source_test.py @@ -0,0 +1,1380 @@ +# +# 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. +# + +"""Tests for apache_beam.io.unbounded_source. + +Semantics are covered by deterministic unit tests; the end-to-end DirectRunner +tests assert ordering and termination only (no flaky defer-timing assertions). +""" + +import logging +import unittest + +from typing_extensions import override + +import apache_beam as beam +from apache_beam import coders +from apache_beam.io import unbounded_source as _unbounded_source_module +from apache_beam.io.unbounded_source import _NO_DATA +from apache_beam.io.unbounded_source import CheckpointMark +from apache_beam.io.unbounded_source import ReadFromUnboundedSource +from apache_beam.io.unbounded_source import UnboundedReader +from apache_beam.io.unbounded_source import UnboundedSource +from apache_beam.io.unbounded_source import _FinalizeCheckpointOnce +from apache_beam.io.unbounded_source import _ReaderCache +from apache_beam.io.unbounded_source import _ReadFromUnboundedSourceDoFn +from apache_beam.io.unbounded_source import _set_watermark_if_greater +from apache_beam.io.unbounded_source import _UnboundedSourceRestriction +from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionCoder +from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionProvider +from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionTracker +from apache_beam.io.watermark_estimators import ManualWatermarkEstimator +from apache_beam.runners import sdf_utils +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to +from apache_beam.transforms import core +from apache_beam.transforms.window import FixedWindows +from apache_beam.utils.timestamp import MAX_TIMESTAMP +from apache_beam.utils.timestamp import MIN_TIMESTAMP +from apache_beam.utils.timestamp import Timestamp + +# pylint: disable=expression-not-assigned + +# Realistic event-time base away from the Unix epoch. +_EVENT_TIME_BASE = Timestamp(1729987200) # 2024-10-27T00:00:00Z + +# ------------------------------------------------------------------------------ +# In-memory demo source emitting integers 0..count-1 with event time +# ``_EVENT_TIME_BASE + index``. It self-terminates at EOF, resumes from +# ``last_index + 1``, and splits into even/odd sub-sources when configured. +# ------------------------------------------------------------------------------ + + +class _CountingCheckpointMark(CheckpointMark): + def __init__(self, last_index, finalize_log=None): + self.last_index = last_index + self._finalize_log = finalize_log + + @override + def finalize_checkpoint(self): + if self._finalize_log is not None: + self._finalize_log.append(self.last_index) + + def __eq__(self, other): + return ( + isinstance(other, _CountingCheckpointMark) and + other.last_index == self.last_index) + + def __hash__(self): + return hash(self.last_index) + + def __repr__(self): + return '_CountingCheckpointMark(last_index=%r)' % (self.last_index, ) + + +class _CountingReader(UnboundedReader): + def __init__( + self, count, start_index, finalize_log=None, modulus=1, residue=0): + self._count = count + self._next = start_index + self._modulus = modulus + self._residue = residue + self._current = None + self._exhausted = False + self._finalize_log = finalize_log + self.closed = False + + def _read_next(self): + while self._next < self._count: + index = self._next + self._next += 1 + if index % self._modulus == self._residue: + self._current = index + return True + self._exhausted = True + return False + + @override + def start(self): + return self._read_next() + + @override + def advance(self): + return self._read_next() + + @override + def get_current(self): + return self._current + + @override + def get_current_timestamp(self): + return _EVENT_TIME_BASE + self._current + + @override + def get_watermark(self): + if self._exhausted: + return MAX_TIMESTAMP + if self._current is None: + return MIN_TIMESTAMP + return _EVENT_TIME_BASE + self._current + + @override + def get_checkpoint_mark(self): + last = self._current if self._current is not None else self._next - 1 + return _CountingCheckpointMark(last, finalize_log=self._finalize_log) + + @override + def close(self): + self.closed = True + + +class UnboundedCountingSource(UnboundedSource): + def __init__( + self, + count, + finalize_log=None, + is_splittable=False, + modulus=1, + residue=0): + self._count = count + self._finalize_log = finalize_log + self._is_splittable = is_splittable + self._modulus = modulus + self._residue = residue + self.last_reader = None + + @override + def split(self, desired_num_splits, options=None): + if not self._is_splittable or desired_num_splits < 2: + return [self] + # Split into independent even/odd sub-sources (each non-splittable). + return [ + UnboundedCountingSource( + self._count, + finalize_log=self._finalize_log, + modulus=2, + residue=residue) for residue in (0, 1) + ] + + @override + def create_reader(self, options, checkpoint_mark): + start_index = ( + 0 if checkpoint_mark is None else checkpoint_mark.last_index + 1) + self.last_reader = _CountingReader( + self._count, + start_index, + finalize_log=self._finalize_log, + modulus=self._modulus, + residue=self._residue) + return self.last_reader + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + +class _StringCountingReader(_CountingReader): + @override + def get_current(self): + return 'v%s' % self._current + + +class _StringCountingSource(UnboundedCountingSource): + @override + def create_reader(self, options, checkpoint_mark): + start_index = ( + 0 if checkpoint_mark is None else checkpoint_mark.last_index + 1) + self.last_reader = _StringCountingReader( + self._count, start_index, finalize_log=self._finalize_log) + return self.last_reader + + @override + def default_output_coder(self): + return coders.StrUtf8Coder() + + +class _PrefixStrCoder(coders.Coder): + def __init__(self, prefix): + self._prefix = prefix + + @override + def encode(self, value): + if not value.startswith(self._prefix): + raise ValueError('expected %r prefix' % self._prefix) + return value[len(self._prefix):].encode('utf-8') + + @override + def decode(self, value): + return self._prefix + value.decode('utf-8') + + @override + def is_deterministic(self): + return True + + @override + def to_type_hint(self): + return str + + +class _PrefixStringReader(_StringCountingReader): + @override + def get_current(self): + return 'prefix:%s' % super().get_current() + + +class _PrefixStringSource(_StringCountingSource): + @override + def create_reader(self, options, checkpoint_mark): + start_index = ( + 0 if checkpoint_mark is None else checkpoint_mark.last_index + 1) + self.last_reader = _PrefixStringReader( + self._count, start_index, finalize_log=self._finalize_log) + return self.last_reader + + @override + def default_output_coder(self): + return _PrefixStrCoder('prefix:') + + +class _NoDataReader(UnboundedReader): + """Always reports temporary absence of data with watermark below MAX.""" + @override + def start(self): + return False + + @override + def advance(self): + return False + + @override + def get_current(self): + raise AssertionError('no data available') + + @override + def get_current_timestamp(self): + raise AssertionError('no data available') + + @override + def get_watermark(self): + return _EVENT_TIME_BASE + + @override + def get_checkpoint_mark(self): + return _CountingCheckpointMark(-1) + + +class _NoDataSource(UnboundedSource): + @override + def split(self, desired_num_splits, options=None): + return [self] + + @override + def create_reader(self, options, checkpoint_mark): + return _NoDataReader() + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + +class _MutatingCheckpointMark(CheckpointMark): + """A mark that mutates itself on finalize, to test primary/residual mark + isolation across a checkpoint cut.""" + def __init__(self, last_index): + self.last_index = last_index + + @override + def finalize_checkpoint(self): + self.last_index = -999 + + +class _MutatingReader(UnboundedReader): + def __init__(self): + self._index = -1 + + @override + def start(self): + self._index = 0 + return True + + @override + def advance(self): + self._index += 1 + return True + + @override + def get_current(self): + return self._index + + @override + def get_current_timestamp(self): + return _EVENT_TIME_BASE + self._index + + @override + def get_watermark(self): + return _EVENT_TIME_BASE + self._index + + @override + def get_checkpoint_mark(self): + return _MutatingCheckpointMark(self._index) + + +class _MutatingSource(UnboundedSource): + @override + def split(self, desired_num_splits, options=None): + return [self] + + @override + def create_reader(self, options, checkpoint_mark): + return _MutatingReader() + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + +class _MaxOnLastReader(UnboundedReader): + """Returns its only record with a MAX_TIMESTAMP watermark on the same claim, + then EOF on the next claim.""" + @override + def start(self): + return True + + @override + def advance(self): + return False + + @override + def get_current(self): + return 7 + + @override + def get_current_timestamp(self): + return _EVENT_TIME_BASE + + @override + def get_watermark(self): + return MAX_TIMESTAMP + + @override + def get_checkpoint_mark(self): + return _CountingCheckpointMark(0) + + +class _MaxOnLastSource(UnboundedSource): + @override + def split(self, desired_num_splits, options=None): + return [self] + + @override + def create_reader(self, options, checkpoint_mark): + return _MaxOnLastReader() + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + +def _new_tracker(source, checkpoint=None): + restriction = _UnboundedSourceRestriction( + source=source, checkpoint_mark=checkpoint) + return _UnboundedSourceRestrictionTracker(restriction) + + +def _claim(tracker): + """Claims once; returns (claimed_bool, holder_value).""" + holder = [None] + claimed = tracker.try_claim(holder) + return claimed, holder[0] + + +# ------------------------------------------------------------------------------ +# Tests +# ------------------------------------------------------------------------------ + + +class AbcContractTest(unittest.TestCase): + def test_checkpointmark_default_finalize_is_noop(self): + self.assertIsNone(CheckpointMark().finalize_checkpoint()) + + def test_unboundedsource_is_bounded_false(self): + self.assertFalse(UnboundedCountingSource(3).is_bounded()) + + def test_reader_lifecycle_start_advance_eof(self): + reader = UnboundedCountingSource(3).create_reader(None, None) + self.assertTrue(reader.start()) + self.assertEqual(reader.get_current(), 0) + self.assertEqual(reader.get_current_timestamp(), _EVENT_TIME_BASE) + self.assertTrue(reader.advance()) + self.assertEqual(reader.get_current(), 1) + self.assertTrue(reader.advance()) + self.assertEqual(reader.get_current(), 2) + self.assertFalse(reader.advance()) + self.assertEqual(reader.get_watermark(), MAX_TIMESTAMP) + + +class RestrictionCoderTest(unittest.TestCase): + def test_roundtrip_no_checkpoint(self): + source = UnboundedCountingSource(3) + coder = _UnboundedSourceRestrictionCoder() + decoded = coder.decode( + coder.encode(_UnboundedSourceRestriction(source=source))) + self.assertIsNone(decoded.checkpoint_mark) + self.assertEqual(decoded.watermark, MIN_TIMESTAMP) + self.assertFalse(decoded.is_done) + reader = decoded.source.create_reader(None, None) + self.assertTrue(reader.start()) + self.assertEqual(reader.get_current(), 0) + + def test_roundtrip_with_checkpoint_resumes(self): + source = UnboundedCountingSource(5) + coder = _UnboundedSourceRestrictionCoder() + restriction = _UnboundedSourceRestriction( + source=source, + checkpoint_mark=_CountingCheckpointMark(1), + watermark=Timestamp(1), + is_done=False) + decoded = coder.decode(coder.encode(restriction)) + self.assertEqual(decoded.checkpoint_mark.last_index, 1) + self.assertEqual(decoded.watermark, Timestamp(1)) + self.assertFalse(decoded.is_done) + # A reader built from the decoded checkpoint resumes at the next index. + reader = decoded.source.create_reader(None, decoded.checkpoint_mark) + self.assertTrue(reader.start()) + self.assertEqual(reader.get_current(), 2) + + +class RestrictionProviderTest(unittest.TestCase): + def test_initial_split_calls_source_split(self): + split_log = [] + + class _NamedSource(UnboundedCountingSource): + def __init__(self, name): + super().__init__(0) + self.name = name + + @override + def split(self, desired_num_splits, options=None): + split_log.append((desired_num_splits, options)) + return [_NamedSource('a'), _NamedSource('b')] + + source = _NamedSource('root') + provider = _UnboundedSourceRestrictionProvider() + restriction = _UnboundedSourceRestriction( + source=source, watermark=Timestamp(7)) + + splits = list(provider.split(source, restriction)) + + # The provider is a stateless module-level singleton, so it always + # passes ``None`` as the ``options`` argument to ``UnboundedSource.split``. + self.assertEqual(split_log, [(20, None)]) + self.assertEqual([split.source.name for split in splits], ['a', 'b']) + self.assertEqual([split.watermark for split in splits], [Timestamp(7)] * 2) + self.assertTrue(all(split.checkpoint_mark is None for split in splits)) + self.assertTrue( + all(split.finalization_checkpoint_mark is None for split in splits)) + + def test_initial_split_does_not_split_checkpointed_restriction(self): + split_log = [] + + class _SplitSource(UnboundedCountingSource): + @override + def split(self, desired_num_splits, options=None): + split_log.append((desired_num_splits, options)) + return [self] + + source = _SplitSource(5) + provider = _UnboundedSourceRestrictionProvider() + restriction = _UnboundedSourceRestriction( + source=source, checkpoint_mark=_CountingCheckpointMark(2)) + + self.assertEqual(list(provider.split(source, restriction)), [restriction]) + self.assertEqual(split_log, []) + + def test_initial_split_falls_back_to_original_on_split_error(self): + class _BoomSource(UnboundedCountingSource): + @override + def split(self, desired_num_splits, options=None): + raise RuntimeError('split boom') + + source = _BoomSource(5) + provider = _UnboundedSourceRestrictionProvider() + restriction = _UnboundedSourceRestriction(source=source) + + self.assertEqual(list(provider.split(source, restriction)), [restriction]) + + def test_truncate_returns_none_for_drain(self): + # On drain the SDF stops emitting; truncate yields no residual. + provider = _UnboundedSourceRestrictionProvider() + source = UnboundedCountingSource(5) + restriction = _UnboundedSourceRestriction(source=source) + self.assertIsNone(provider.truncate(source, restriction)) + + def test_splittable_source_partitions_into_independent_subsources(self): + # A splittable source fans out into two sub-sources; reading each in + # isolation yields the even and the odd integers, and their union is the + # full sequence with no overlap. + source = UnboundedCountingSource(6, is_splittable=True) + provider = _UnboundedSourceRestrictionProvider() + restriction = _UnboundedSourceRestriction(source=source) + + splits = list(provider.split(source, restriction)) + self.assertEqual(len(splits), 2) + + shards = [] + for split in splits: + tracker = _UnboundedSourceRestrictionTracker(split) + shard = [] + while True: + claimed, record = _claim(tracker) + if not claimed: + break + if record is not _NO_DATA: + shard.append(record[0]) + shards.append(shard) + self.assertEqual(sorted(shards), [[0, 2, 4], [1, 3, 5]]) + + +class RestrictionTrackerTest(unittest.TestCase): + def test_claim_emits_in_order(self): + tracker = _new_tracker(UnboundedCountingSource(3)) + values = [] + while True: + claimed, record = _claim(tracker) + if not claimed: + break + self.assertIsNot(record, _NO_DATA) + values.append(record[0]) + self.assertEqual(values, [0, 1, 2]) + self.assertTrue(tracker.check_done()) + + def test_claim_emits_final_record_when_watermark_is_max(self): + # A reader may return its last record with a MAX_TIMESTAMP watermark on the + # same call; the record must still be emitted (EOF comes on the next claim). + class _FinalRecordReader(UnboundedReader): + @override + def start(self): + return True + + @override + def advance(self): + return False + + @override + def get_current(self): + return 'last' + + @override + def get_current_timestamp(self): + return _EVENT_TIME_BASE + + @override + def get_watermark(self): + return MAX_TIMESTAMP + + @override + def get_checkpoint_mark(self): + return _CountingCheckpointMark(0) + + class _FinalSource(UnboundedSource): + @override + def split(self, desired_num_splits, options=None): + return [self] + + @override + def create_reader(self, options, checkpoint_mark): + return _FinalRecordReader() + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + tracker = _new_tracker(_FinalSource()) + claimed, record = _claim(tracker) + self.assertTrue(claimed) + self.assertIsNot(record, _NO_DATA) + self.assertEqual(record[0], 'last') + # The next claim observes EOF and finishes (no second, phantom record). + claimed_again, _ = _claim(tracker) + self.assertFalse(claimed_again) + self.assertTrue(tracker.check_done()) + + def test_try_split_zero_produces_resumable_residual(self): + source = UnboundedCountingSource(5) + tracker = _new_tracker(source) + # Claim 0 and 1. + self.assertEqual(_claim(tracker)[1][0], 0) + self.assertEqual(_claim(tracker)[1][0], 1) + + split = tracker.try_split(0) + self.assertIsNotNone(split) + primary, residual = split + self.assertTrue(primary.is_done) + self.assertFalse(residual.is_done) + # Resume / finalize channel separation: primary carries only the + # finalize hook, residual carries only the resume state. + self.assertIsNone(primary.checkpoint_mark) + self.assertIsNotNone(primary.finalization_checkpoint_mark) + self.assertEqual(primary.finalization_checkpoint_mark.last_index, 1) + self.assertEqual(residual.checkpoint_mark.last_index, 1) + self.assertIsNone(residual.finalization_checkpoint_mark) + # check_done passes on the (now done) primary. + self.assertTrue(tracker.check_done()) + + # Resuming from the residual continues at index 2. + resumed = _new_tracker(source, checkpoint=residual.checkpoint_mark) + self.assertEqual(_claim(resumed)[1][0], 2) + + def test_try_split_isolates_residual_from_finalize_mutation(self): + # The primary's finalize hook and the residual's resume state must not + # share one object, so a mutating finalize_checkpoint() cannot corrupt the + # residual's resume position. + tracker = _new_tracker(_MutatingSource()) + _claim(tracker) # start -> index 0 + split = tracker.try_split(0) + self.assertIsNotNone(split) + primary, residual = split + self.assertIsNot( + primary.finalization_checkpoint_mark, residual.checkpoint_mark) + self.assertEqual(residual.checkpoint_mark.last_index, 0) + primary.finalization_checkpoint_mark.finalize_checkpoint() + self.assertEqual(residual.checkpoint_mark.last_index, 0) + + def test_try_split_nonzero_declined(self): + source = UnboundedCountingSource(5) + tracker = _new_tracker(source) + self.assertEqual(_claim(tracker)[1][0], 0) + + self.assertIsNone(tracker.try_split(0.5)) + self.assertFalse(tracker.current_restriction().is_done) + self.assertIsNotNone(tracker._reader) + self.assertEqual(_claim(tracker)[1][0], 1) + + def test_no_data_returns_sentinel_without_finishing(self): + tracker = _new_tracker(_NoDataSource()) + claimed, record = _claim(tracker) + self.assertTrue(claimed) + self.assertIs(record, _NO_DATA) + # A self-checkpoint is still possible (poll/resume path). + self.assertIsNotNone(tracker.try_split(0)) + + def test_check_done_raises_when_not_done(self): + tracker = _new_tracker(UnboundedCountingSource(3)) + with self.assertRaises(ValueError): + tracker.check_done() + + def test_is_bounded_false(self): + self.assertFalse(_new_tracker(UnboundedCountingSource(3)).is_bounded()) + + +class _RecordingBundleFinalizer: + def __init__(self): + self.registered = [] + + def register(self, callback): + self.registered.append(callback) + + +class _ManualClock: + """A deterministic monotonic clock for the time-cap tests.""" + def __init__(self, now=0.0): + self.now = now + + def __call__(self): + return self.now + + +class BundleCapTest(unittest.TestCase): + """A busy reader self-checkpoints once the per-bundle record or time cap is + reached, so the runner can commit progress and run finalization.""" + def _bundle(self, dofn, source, checkpoint=None, estimator=None): + """Builds the SDF tracker chain and returns the process() generator plus the + tracker, threadsafe tracker, finalizer, and watermark estimator.""" + tracker = _UnboundedSourceRestrictionTracker( + _UnboundedSourceRestriction(source=source, checkpoint_mark=checkpoint)) + threadsafe = sdf_utils.ThreadsafeRestrictionTracker(tracker) + view = sdf_utils.RestrictionTrackerView(threadsafe) + finalizer = _RecordingBundleFinalizer() + estimator = estimator or ManualWatermarkEstimator(None) + gen = dofn.process( + None, + bundle_finalizer=finalizer, + tracker=view, + watermark_estimator=estimator) + return gen, tracker, threadsafe, finalizer, estimator + + def test_record_cap_checkpoints_busy_source(self): + finalize_log = [] + dofn = _ReadFromUnboundedSourceDoFn( + poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9) + # 1000 records is effectively unbounded against a cap of 5. + gen, tracker, threadsafe, finalizer, estimator = self._bundle( + dofn, UnboundedCountingSource(1000, finalize_log=finalize_log)) + outputs = list(gen) + + self.assertEqual([tv.value for tv in outputs], [0, 1, 2, 3, 4]) + self.assertTrue(tracker.current_restriction().is_done) + self.assertTrue(tracker.check_done()) + # The estimator holds the last emitted record's source watermark. + self.assertEqual(estimator.current_watermark(), _EVENT_TIME_BASE + 4) + # Residual resumes after the cut and carries no finalize hook. + residual, _ = threadsafe.deferred_status() + self.assertEqual(residual.checkpoint_mark.last_index, 4) + self.assertIsNone(residual.finalization_checkpoint_mark) + # Exactly one finalizer is registered; firing it commits the cut index once. + self.assertEqual(len(finalizer.registered), 1) + finalizer.registered[0]() + finalizer.registered[0]() + self.assertEqual(finalize_log, [4]) + + def test_time_cap_checkpoints_busy_source(self): + clock = _ManualClock(1000.0) + dofn = _ReadFromUnboundedSourceDoFn( + poll_interval=0, + max_records_per_bundle=10**9, + max_read_time_seconds=5.0, + _now=clock) + gen, tracker, threadsafe, _, _ = self._bundle( + dofn, UnboundedCountingSource(1000)) + + # The deadline arms at 1000 + 5 after the first record and is checked + # between records, so records keep flowing until the clock passes it. + self.assertEqual(next(gen).value, 0) + self.assertEqual(next(gen).value, 1) + self.assertEqual(next(gen).value, 2) + clock.now = 1006.0 + with self.assertRaises(StopIteration): + next(gen) + + self.assertTrue(tracker.current_restriction().is_done) + residual, _ = threadsafe.deferred_status() + self.assertEqual(residual.checkpoint_mark.last_index, 2) + + def test_cap_residual_resumes_in_next_bundle(self): + dofn = _ReadFromUnboundedSourceDoFn( + poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9) + source = UnboundedCountingSource(1000) + # Bundle 1 emits 0-4 and cuts a residual at index 4. + gen1, _, threadsafe1, _, _ = self._bundle(dofn, source) + self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4]) + residual1, _ = threadsafe1.deferred_status() + + # Bundle 2 rebuilds the reader from the residual and emits 5-9. + gen2, _, threadsafe2, _, _ = self._bundle( + dofn, source, checkpoint=residual1.checkpoint_mark) + self.assertEqual([tv.value for tv in gen2], [5, 6, 7, 8, 9]) + residual2, _ = threadsafe2.deferred_status() + self.assertEqual(residual2.checkpoint_mark.last_index, 9) + + def test_busy_reader_is_reused_across_bundles(self): + # The self-checkpoint parks the reader; the resuming bundle reclaims the + # same started reader. + dofn = _ReadFromUnboundedSourceDoFn( + poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9) + dofn.setup() # creates the cross-bundle reader cache + source = UnboundedCountingSource(1000) + gen1, _, threadsafe1, _, _ = self._bundle(dofn, source) + self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4]) + reader1 = source.last_reader + self.assertFalse(reader1.closed) # parked, not closed + residual1, _ = threadsafe1.deferred_status() + + gen2, _, _, _, _ = self._bundle( + dofn, source, checkpoint=residual1.checkpoint_mark) + self.assertEqual([tv.value for tv in gen2], [5, 6, 7, 8, 9]) + # No new reader was created; source.last_reader is unchanged. + self.assertIs(source.last_reader, reader1) + + def test_teardown_closes_parked_readers(self): + dofn = _ReadFromUnboundedSourceDoFn( + poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9) + dofn.setup() + source = UnboundedCountingSource(1000) + gen, _, _, _, _ = self._bundle(dofn, source) + list(gen) # the bundle parks its reader + reader = source.last_reader + self.assertFalse(reader.closed) + dofn.teardown() + self.assertTrue(reader.closed) + + def test_eof_exactly_at_cap_resumes_then_finishes(self): + dofn = _ReadFromUnboundedSourceDoFn( + poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9) + source = UnboundedCountingSource(5) # exactly cap records + # Bundle 1 hits the cap on the last record before observing EOF. + gen1, t1, threadsafe1, _, _ = self._bundle(dofn, source) + self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4]) + self.assertTrue(t1.current_restriction().is_done) + residual1, _ = threadsafe1.deferred_status() + self.assertEqual(residual1.checkpoint_mark.last_index, 4) + + # Bundle 2 resumes at index 5, finds EOF, and finishes with no output. + gen2, t2, threadsafe2, _, _ = self._bundle( + dofn, source, checkpoint=residual1.checkpoint_mark) + self.assertEqual(list(gen2), []) + self.assertTrue(t2.current_restriction().is_done) + self.assertIsNone(threadsafe2.deferred_status()) + + def test_eof_before_cap_finishes_without_residual(self): + dofn = _ReadFromUnboundedSourceDoFn( + poll_interval=0, max_records_per_bundle=100, max_read_time_seconds=1e9) + gen, tracker, threadsafe, _, _ = self._bundle( + dofn, UnboundedCountingSource(3)) + + self.assertEqual([tv.value for tv in gen], [0, 1, 2]) + self.assertTrue(tracker.current_restriction().is_done) + self.assertIsNone(threadsafe.deferred_status()) + + def test_max_watermark_on_final_record_emitted_before_estimator_advances( + self): + # A reader that returns its final record with a MAX_TIMESTAMP watermark on + # the same claim must have that record emitted before the estimator reaches + # MAX, so the element is not stranded behind the output watermark. + dofn = _ReadFromUnboundedSourceDoFn(poll_interval=0) + gen, _, _, _, estimator = self._bundle(dofn, _MaxOnLastSource()) + + first = next(gen) + self.assertEqual(first.value, 7) + # The estimator has not yet been advanced to MAX when the record is yielded. + self.assertNotEqual(estimator.current_watermark(), MAX_TIMESTAMP) + # Draining the bundle then advances the estimator to MAX on EOF. + self.assertEqual(list(gen), []) + self.assertEqual(estimator.current_watermark(), MAX_TIMESTAMP) + + +class WatermarkTest(unittest.TestCase): + def test_set_watermark_is_monotonic(self): + estimator = ManualWatermarkEstimator(None) + _set_watermark_if_greater(estimator, Timestamp(5)) + self.assertEqual(estimator.current_watermark(), Timestamp(5)) + # A regression is ignored (would otherwise raise inside set_watermark). + _set_watermark_if_greater(estimator, Timestamp(3)) + self.assertEqual(estimator.current_watermark(), Timestamp(5)) + _set_watermark_if_greater(estimator, Timestamp(7)) + self.assertEqual(estimator.current_watermark(), Timestamp(7)) + + +class FinalizationTest(unittest.TestCase): + def test_finalize_checkpoint_callback_is_at_most_once(self): + finalize_log = [] + finalize_once = _FinalizeCheckpointOnce( + _CountingCheckpointMark(1, finalize_log=finalize_log)) + + finalize_once() + finalize_once() + + self.assertEqual(finalize_log, [1]) + + def test_finalize_checkpoint_invoked(self): + # Unit-level finalize test (the e2e finalize may run in a worker process); + # the hook lives on the primary, independent of the residual's resume state. + finalize_log = [] + source = UnboundedCountingSource(5, finalize_log=finalize_log) + tracker = _new_tracker(source) + _claim(tracker) # 0 + _claim(tracker) # 1 + primary, _ = tracker.try_split(0) + primary.finalization_checkpoint_mark.finalize_checkpoint() + self.assertEqual(finalize_log, [1]) + + +class EndToEndTest(unittest.TestCase): + def test_direct_runner_emits_all_in_order(self): + with TestPipeline() as p: + out = p | ReadFromUnboundedSource(UnboundedCountingSource(5)) + self.assertFalse(out.is_bounded) + assert_that(out, equal_to([0, 1, 2, 3, 4])) + + def test_eof_lets_event_time_window_fire(self): + # On EOF the DoFn advances the watermark estimator to MAX_TIMESTAMP so the + # downstream FixedWindow closes and the GroupByKey fires; otherwise the + # output would be empty. + with TestPipeline() as p: + out = ( + p + | ReadFromUnboundedSource(UnboundedCountingSource(5)) + | beam.WindowInto(FixedWindows(100)) + | beam.Map(lambda v: ('all', v)) + | beam.GroupByKey() + | beam.MapTuple(lambda _key, values: sorted(values))) + assert_that(out, equal_to([[0, 1, 2, 3, 4]])) + + def test_read_dispatches_through_iobase_read(self): + # ``beam.io.Read(source)`` must produce the same records as + # ``ReadFromUnboundedSource(source)``. + with TestPipeline() as p: + out = p | beam.io.Read(UnboundedCountingSource(5)) + self.assertFalse(out.is_bounded) + assert_that(out, equal_to([0, 1, 2, 3, 4])) + + def test_splittable_source_reads_all_records_across_splits(self): + # A splittable source fans out into even/odd sub-sources during initial + # SDF splitting; the union of all sub-source reads is the full sequence. + with TestPipeline() as p: + out = p | beam.io.Read(UnboundedCountingSource(6, is_splittable=True)) + assert_that(out, equal_to([0, 1, 2, 3, 4, 5])) + + def test_source_default_output_coder_sets_output_type(self): + with TestPipeline() as p: + out = p | ReadFromUnboundedSource(_StringCountingSource(2)) + self.assertEqual(out.element_type, str) + assert_that(out, equal_to(['v0', 'v1'])) + + def test_small_cap_self_checkpoints_and_resumes_to_eof(self): + # A small per-bundle cap forces several self-checkpoint/resume cycles + # through the runner before EOF, exercising the cross-bundle reader cache + # over real residual encode/decode. All records still arrive in order. + with TestPipeline() as p: + out = p | ReadFromUnboundedSource( + UnboundedCountingSource(20), max_records_per_bundle=5) + assert_that(out, equal_to(list(range(20)))) + + +class ReadFromUnboundedSourceCoderTest(unittest.TestCase): + def test_parameterized_output_coder_does_not_mutate_global_registry(self): + try: + p = beam.Pipeline() + out = p | ReadFromUnboundedSource(_PrefixStringSource(1)) + + self.assertNotEqual(out.element_type, str) + self.assertEqual(coders.registry.get_coder(str), coders.StrUtf8Coder()) + self.assertEqual( + ReadFromUnboundedSource(_PrefixStringSource(1))._infer_output_coder(), + _PrefixStrCoder('prefix:')) + finally: + coders.registry.register_coder(str, coders.StrUtf8Coder) + + +# ------------------------------------------------------------------------------ +# Reader lifecycle, watermark, and contract regression tests (reader close on +# every exit path, the NotImplementedError message, finalize idempotency). +# ------------------------------------------------------------------------------ + + +class ReaderCloseTest(unittest.TestCase): + """Reader lifecycle: close() must run on every tracker-driven exit path.""" + def test_tracker_closes_reader_on_eof(self): + source = UnboundedCountingSource(0) # immediately exhausted + tracker = _new_tracker(source) + holder = [None] + self.assertFalse(tracker.try_claim(holder)) + self.assertIsNone(tracker._reader) + self.assertTrue(source.last_reader.closed) + + def test_tracker_closes_reader_on_split_without_cache(self): + # With no cache injected, a self-checkpoint closes the reader. + source = UnboundedCountingSource(5) + tracker = _new_tracker(source) + _claim(tracker) # creates reader, claims 0 + reader = source.last_reader + self.assertFalse(reader.closed) + split = tracker.try_split(0) + self.assertIsNotNone(split) + self.assertIsNone(tracker._reader) + self.assertTrue(reader.closed) + + def test_tracker_parks_reader_on_split_with_cache(self): + # With a cache, a self-checkpoint parks the reader for the residual to + # reclaim. + source = UnboundedCountingSource(5) + cache = _ReaderCache() + tracker = _new_tracker(source) + tracker._reader_cache = cache + _claim(tracker) + reader = source.last_reader + _, residual = tracker.try_split(0) + self.assertIsNone(tracker._reader) + self.assertFalse(reader.closed) + # The residual's key reclaims the same started reader. + self.assertEqual( + cache.acquire(tracker._cache_key(residual)), (reader, True)) + + def test_close_helper_is_idempotent_and_safe_on_empty_tracker(self): + tracker = _new_tracker(UnboundedCountingSource(3)) + # No reader yet -- helper must be a no-op. + tracker._close_reader_if_open() + _claim(tracker) + reader = tracker._reader + tracker._close_reader_if_open() + self.assertTrue(reader.closed) + self.assertIsNone(tracker._reader) + # Second call is a no-op (no reader to close). + tracker._close_reader_if_open() + + def test_close_helper_swallows_reader_close_errors(self): + class _BoomReader(UnboundedReader): + @override + def start(self): + return True + + @override + def advance(self): + return False + + @override + def get_current(self): + return 'x' + + @override + def get_current_timestamp(self): + return _EVENT_TIME_BASE + + @override + def get_watermark(self): + return _EVENT_TIME_BASE + + @override + def get_checkpoint_mark(self): + return CheckpointMark() + + @override + def close(self): + raise RuntimeError('close blew up') + + class _BoomSource(UnboundedSource): + @override + def split(self, desired_num_splits, options=None): + return [self] + + @override + def create_reader(self, options, checkpoint_mark): + return _BoomReader() + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + tracker = _new_tracker(_BoomSource()) + _claim(tracker) + with self.assertLogs(_unbounded_source_module._LOGGER, 'WARNING') as logs: + tracker._close_reader_if_open() + self.assertTrue( + any('Error closing UnboundedReader' in line for line in logs.output)) + self.assertIsNone(tracker._reader) + + +class ReaderCacheTest(unittest.TestCase): + """The cache parks a reader under one key, hands it to the next acquirer, + bounds itself by idle time and entry count, and closes on teardown.""" + def _reader(self): + class _FakeReader(UnboundedReader): + def __init__(self): + self.closed = False + + @override + def close(self): + self.closed = True + + return _FakeReader() + + def test_park_then_acquire_returns_same_reader_and_started_flag(self): + cache = _ReaderCache() + reader = self._reader() + cache.park('k', reader, True) + self.assertEqual(cache.acquire('k'), (reader, True)) + # Acquire removes the entry so two trackers cannot share one reader. + self.assertIsNone(cache.acquire('k')) + self.assertFalse(reader.closed) + + def test_acquire_miss_returns_none(self): + self.assertIsNone(_ReaderCache().acquire('absent')) + + def test_park_replacing_same_key_closes_displaced_reader(self): + # Two parks under one key without an intervening acquire (e.g. a bundle + # retry) must not leak the first reader. + cache = _ReaderCache() + old, new = self._reader(), self._reader() + cache.park('k', old, True) + cache.park('k', new, True) + self.assertTrue(old.closed) + self.assertFalse(new.closed) + self.assertEqual(cache.acquire('k'), (new, True)) + + def test_idle_reader_is_closed_on_next_touch(self): + clock = _ManualClock(1000.0) + cache = _ReaderCache(idle_seconds=30.0, now=clock) + reader = self._reader() + cache.park('k', reader, True) + clock.now = 1031.0 # past the idle window + cache.park('other', self._reader(), False) # triggers idle eviction + self.assertTrue(reader.closed) + self.assertIsNone(cache.acquire('k')) + + def test_max_size_evicts_and_closes_oldest(self): + cache = _ReaderCache(max_size=2) + readers = [self._reader() for _ in range(3)] + cache.park('a', readers[0], True) + cache.park('b', readers[1], True) + cache.park('c', readers[2], True) # exceeds the cap, evicts 'a' + self.assertTrue(readers[0].closed) + self.assertIsNone(cache.acquire('a')) + self.assertIsNotNone(cache.acquire('b')) + self.assertIsNotNone(cache.acquire('c')) + + def test_close_all_closes_every_parked_reader(self): + cache = _ReaderCache() + readers = [self._reader() for _ in range(3)] + for i, reader in enumerate(readers): + cache.park(str(i), reader, True) + cache.close_all() + self.assertTrue(all(reader.closed for reader in readers)) + self.assertIsNone(cache.acquire('0')) + + def test_close_all_swallows_reader_close_errors(self): + class _BoomReader(UnboundedReader): + @override + def close(self): + raise RuntimeError('close blew up') + + cache = _ReaderCache() + cache.park('k', _BoomReader(), True) + with self.assertLogs(_unbounded_source_module._LOGGER, 'WARNING') as logs: + cache.close_all() + self.assertTrue( + any('Error closing UnboundedReader' in line for line in logs.output)) + + +class TrackerContractRegressionTest(unittest.TestCase): + """Tracker contract: source-watermark on the data path, finalize/resume + channel separation, and reader close on a reader-method failure.""" + def test_data_path_holder_carries_source_watermark(self): + class _LaggingReader(UnboundedReader): + @override + def start(self): + return True + + @override + def advance(self): + return False + + @override + def get_current(self): + return 'rec' + + @override + def get_current_timestamp(self): + return Timestamp(1000) # record event time + + @override + def get_watermark(self): + return Timestamp(990) # source watermark lags 10us behind + + @override + def get_checkpoint_mark(self): + return _CountingCheckpointMark(0) + + class _LaggingSource(UnboundedSource): + @override + def split(self, desired_num_splits, options=None): + return [self] + + @override + def create_reader(self, options, checkpoint_mark): + return _LaggingReader() + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + tracker = _new_tracker(_LaggingSource()) + claimed, record = _claim(tracker) + self.assertTrue(claimed) + self.assertIsNot(record, _NO_DATA) + value, record_timestamp, source_watermark = record + self.assertEqual(value, 'rec') + self.assertEqual(record_timestamp, Timestamp(1000)) + # Critical: watermark slot is the SOURCE watermark, NOT record timestamp. + self.assertEqual(source_watermark, Timestamp(990)) + self.assertNotEqual(source_watermark, record_timestamp) + + def test_split_separates_finalize_and_resume_channels(self): + source = UnboundedCountingSource(5) + tracker = _new_tracker(source) + _claim(tracker) # claim 0 so reader has progress + primary, residual = tracker.try_split(0) + # Primary carries ONLY the finalize hook -- no resume state. + self.assertIsNone(primary.checkpoint_mark) + self.assertIsNotNone(primary.finalization_checkpoint_mark) + self.assertTrue(primary.is_done) + # Residual carries ONLY the resume state -- no finalize hook (a future + # bundle that splits THIS residual will produce ITS own finalize mark). + self.assertIsNotNone(residual.checkpoint_mark) + self.assertIsNone(residual.finalization_checkpoint_mark) + self.assertFalse(residual.is_done) + # The two marks reference the same underlying checkpoint object. + self.assertEqual( + primary.finalization_checkpoint_mark.last_index, + residual.checkpoint_mark.last_index) + + def test_eof_populates_finalize_and_clears_resume(self): + # EOF transition: restriction.checkpoint_mark goes to None (no more + # records to resume from), finalization_checkpoint_mark carries the + # final commit hook. + source = UnboundedCountingSource(0) # immediately exhausted + tracker = _new_tracker(source) + holder = [None] + self.assertFalse(tracker.try_claim(holder)) + r = tracker.current_restriction() + self.assertTrue(r.is_done) + self.assertEqual(r.watermark, MAX_TIMESTAMP) + self.assertIsNone(r.checkpoint_mark) + self.assertIsNotNone(r.finalization_checkpoint_mark) + + def test_tracker_closes_reader_when_advance_raises(self): + # try_claim closes the reader before re-raising a reader-method failure, so + # the DoFn's finally need not traverse the SDF chain for these. + class _BoomReader(UnboundedReader): + def __init__(self): + self.closed = False + + @override + def start(self): + return True + + @override + def advance(self): + raise RuntimeError('advance boom') + + @override + def get_current(self): + return 'first' + + @override + def get_current_timestamp(self): + return _EVENT_TIME_BASE + + @override + def get_watermark(self): + return _EVENT_TIME_BASE + + @override + def get_checkpoint_mark(self): + return _CountingCheckpointMark(0) + + @override + def close(self): + self.closed = True + + class _BoomSource(UnboundedSource): + @override + def split(self, desired_num_splits, options=None): + return [self] + + @override + def create_reader(self, options, checkpoint_mark): + return _BoomReader() + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + src = _BoomSource() + tracker = _new_tracker(src) + # First claim succeeds (start returns True). + self.assertTrue(tracker.try_claim([None])) + reader_after_first = tracker._reader + self.assertIsNotNone(reader_after_first) + # The second claim's advance() raises; the tracker must close the reader + # before propagating. + with self.assertRaises(RuntimeError): + tracker.try_claim([None]) + self.assertTrue(reader_after_first.closed) + self.assertIsNone(tracker._reader) + + def test_tracker_closes_reader_when_get_watermark_raises(self): + # Reader method failures other than advance() also trigger close. + class _WatermarkBoomReader(UnboundedReader): + def __init__(self): + self.closed = False + + @override + def start(self): + return False # no data -> drops into get_watermark path + + @override + def advance(self): + return False + + @override + def get_current(self): + raise AssertionError + + @override + def get_current_timestamp(self): + raise AssertionError + + @override + def get_watermark(self): + raise RuntimeError('watermark boom') + + @override + def get_checkpoint_mark(self): + return _CountingCheckpointMark(0) + + @override + def close(self): + self.closed = True + + class _WatermarkBoomSource(UnboundedSource): + @override + def split(self, desired_num_splits, options=None): + return [self] + + @override + def create_reader(self, options, checkpoint_mark): + return _WatermarkBoomReader() + + @override + def get_checkpoint_mark_coder(self): + return coders.PickleCoder() + + src = _WatermarkBoomSource() + tracker = _new_tracker(src) + with self.assertRaises(RuntimeError): + tracker.try_claim([None]) + self.assertIsNone(tracker._reader) + + +class UnboundedSourceContractTest(unittest.TestCase): + def test_get_checkpoint_mark_coder_default_names_subclass(self): + class MySource(UnboundedSource): + pass + + with self.assertRaises(NotImplementedError) as cm: + MySource().get_checkpoint_mark_coder() + self.assertIn('MySource', str(cm.exception)) + + +class ReadFromUnboundedSourceValidationTest(unittest.TestCase): + def test_non_source_argument_raises(self): + with self.assertRaises(TypeError): + ReadFromUnboundedSource('not-a-source') # type: ignore[arg-type] + + def test_invalid_caps_raise(self): + source = UnboundedCountingSource(1) + with self.assertRaises(ValueError): + ReadFromUnboundedSource(source, max_records_per_bundle=0) + with self.assertRaises(ValueError): + ReadFromUnboundedSource(source, max_read_time_seconds=0) + with self.assertRaises(ValueError): + ReadFromUnboundedSource(source, poll_interval=-1) + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + unittest.main() diff --git a/sdks/python/apache_beam/testing/benchmarks/cloudml/constraints.txt b/sdks/python/apache_beam/testing/benchmarks/cloudml/constraints.txt new file mode 100644 index 000000000000..b2f76d200850 --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/cloudml/constraints.txt @@ -0,0 +1,42 @@ +# +# 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. +# + +# Constraints file to pin versions and avoid pip ResolutionTooDeep. +# This file is used with: pip install -c constraints.txt -r requirements.txt + +# Core dependencies +tfx_bsl==1.15.1 +tensorflow-transform==1.15.0 + +# TensorFlow ecosystem +tensorflow==2.15.1 +tensorflow-metadata==1.15.0 +tf-keras==2.15.1 + +# NumPy and data handling +numpy==1.26.4 +pyarrow==10.0.1 + +# Google Cloud (pin to avoid transitive resolution) +google-cloud-aiplatform==1.60.0 +google-api-core==2.19.1 + +# Note: google-auth is NOT constrained - let pip resolve it to satisfy +# apache-beam's google-genai requirement (>=2.48.1) + +# Note: tensorflow-serving-api is NOT constrained - let pip resolve it within +# the range specified in requirements.txt (>=2.15,<2.16) diff --git a/sdks/python/apache_beam/testing/benchmarks/cloudml/requirements.txt b/sdks/python/apache_beam/testing/benchmarks/cloudml/requirements.txt index ab94ec5e9acf..5f754e27148e 100644 --- a/sdks/python/apache_beam/testing/benchmarks/cloudml/requirements.txt +++ b/sdks/python/apache_beam/testing/benchmarks/cloudml/requirements.txt @@ -15,12 +15,15 @@ # limitations under the License. # -dill==0.4.1 -tfx_bsl==1.16.1 -tensorflow-transform==1.16.0 -tensorflow>=2.16,<2.17 +# Core TFT dependencies with version bounds. +# Note: To avoid pip ResolutionTooDeep errors, always install using the constraints file: +# pip install -c constraints.txt -r requirements.txt +dill>=0.3,<0.5 +tfx_bsl>=1.15,<1.17 +tensorflow-transform>=1.15,<1.17 +tensorflow>=2.15,<2.16 numpy>=1.22.0,<2.0 -tensorflow-metadata>=1.16.1,<1.17.0 +tensorflow-metadata>=1.15,<1.16 pyarrow>=10,<11 -tensorflow-serving-api>=2.16.1,<2.20 -tf-keras>=2.16.0,<2.17 +tensorflow-serving-api>=2.15,<2.16 +tf-keras>=2.15,<2.16 diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index 7c84700e29fa..480e2a62a2ef 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -573,13 +573,9 @@ task installTFTRequirements { exec { workingDir "$rootProject.projectDir/sdks/python/apache_beam/testing/benchmarks/cloudml/" executable 'sh' - // installGcpTest already installed apache-beam[gcp]. tensorflow-transform also - // lists that dependency, so a plain pip install -r can re-resolve the GCP extra - // and hit ResolutionTooDeep. Install TFT with --no-deps instead. - args '-c', ". ${envdir}/bin/activate && " + - "grep -v '^tensorflow-transform' requirements.txt > /tmp/cloudml_tft_base_requirements.txt && " + - "pip install -r /tmp/cloudml_tft_base_requirements.txt && " + - "pip install --no-deps tensorflow-transform==1.16.0" + // Use constraints file to pin versions while allowing pip to + // resolve compatible versions within the specified ranges in requirements.txt + args '-c', ". ${envdir}/bin/activate && pip install -c constraints.txt -r requirements.txt" } } }