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 95ea6cf1bd4b..feab82263245 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java @@ -19,20 +19,25 @@ import static org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergSchemaToBeamSchema; import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets.newHashSet; +import static org.apache.iceberg.types.Type.TypeID.LONG; +import static org.apache.iceberg.types.Type.TypeID.TIMESTAMP; import com.google.auto.value.AutoValue; import java.io.Serializable; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.Comparator; import java.util.HashSet; import java.util.LinkedHashSet; import java.util.List; +import java.util.Locale; import java.util.Set; +import java.util.concurrent.TimeUnit; import org.apache.beam.sdk.io.iceberg.IcebergIO.ReadRows.StartingStrategy; import org.apache.beam.sdk.io.iceberg.cdc.IcebergCdcMetadataColumns; import org.apache.beam.sdk.schemas.Schema; @@ -40,6 +45,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.MetadataColumns; import org.apache.iceberg.StructLike; import org.apache.iceberg.Table; import org.apache.iceberg.TableUtil; @@ -48,6 +54,7 @@ import org.apache.iceberg.expressions.Expression; import org.apache.iceberg.types.Comparators; import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types.NestedField; import org.apache.iceberg.util.SnapshotUtil; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; @@ -102,9 +109,9 @@ static org.apache.iceberg.Schema resolveSchema( @Nullable List keep, @Nullable List drop, @Nullable Set fieldsInFilter) { - ImmutableList.Builder selectedFieldsBuilder = ImmutableList.builder(); + Set selectedFields = new LinkedHashSet<>(); if (keep != null && !keep.isEmpty()) { - selectedFieldsBuilder.addAll(keep); + selectedFields.addAll(keep); } else if (drop != null && !drop.isEmpty()) { List paths = new ArrayList<>(TypeUtil.indexNameById(schema.asStruct()).values()); Collections.sort(paths); @@ -113,7 +120,7 @@ static org.apache.iceberg.Schema resolveSchema( boolean isParent = i + 1 < paths.size() && paths.get(i + 1).startsWith(path + "."); boolean isDrop = drop.stream().anyMatch(d -> path.equals(d) || path.startsWith(d + ".")); if (!isParent && !isDrop) { - selectedFieldsBuilder.add(path); + selectedFields.add(path); } } } else { @@ -124,9 +131,8 @@ static org.apache.iceberg.Schema resolveSchema( if (fieldsInFilter != null && !fieldsInFilter.isEmpty()) { fieldsInFilter.stream() .map(f -> schema.caseInsensitiveFindField(f).name()) - .forEach(selectedFieldsBuilder::add); + .forEach(selectedFields::add); } - ImmutableList selectedFields = selectedFieldsBuilder.build(); return selectedFields.isEmpty() ? schema : schema.select(selectedFields); } @@ -255,6 +261,12 @@ public Expression getFilter() { @Pure public abstract @Nullable List getDropFields(); + @Pure + public abstract @Nullable String getWatermarkColumn(); + + @Pure + public abstract @Nullable String getWatermarkColumnTimeUnit(); + @Pure public abstract @Nullable Duration getMaxSnapshotDiscoveryDelay(); @@ -284,6 +296,7 @@ public static Builder builder() { .setStartingStrategy(null) .setTag(null) .setBranch(null) + .setWatermarkColumn(null) .setMetadataColumns(ImmutableList.of()); } @@ -347,6 +360,10 @@ public Builder setTableIdentifier(String... names) { public abstract Builder setDropFields(@Nullable List fields); + public abstract Builder setWatermarkColumn(@Nullable String watermarkColumn); + + public abstract Builder setWatermarkColumnTimeUnit(@Nullable String timeUnit); + public abstract Builder setMaxSnapshotDiscoveryDelay(@Nullable Duration delay); public abstract Builder setMetadataColumns(List metadataColumns); @@ -357,6 +374,7 @@ public Builder setTableIdentifier(String... names) { @VisibleForTesting abstract Builder toBuilder(); + @SuppressWarnings("ReturnValueIgnored") void validate(Table table) { @Nullable List keep = getKeepFields(); @Nullable List drop = getDropFields(); @@ -368,16 +386,19 @@ void validate(Table table) { String param; if (keep != null) { param = "keep"; - fieldsSpecified = newHashSet(checkNotNull(keep)); + fieldsSpecified = newHashSet(checkArgumentNotNull(keep)); } else { // drop != null param = "drop"; - fieldsSpecified = newHashSet(checkNotNull(drop)); + fieldsSpecified = newHashSet(checkArgumentNotNull(drop)); } fieldsSpecified.removeIf(name -> table.schema().findField(name) != null); checkArgument( - fieldsSpecified.isEmpty(), - error(String.format("'%s' specifies unknown field(s): %s", param, fieldsSpecified))); + fieldsSpecified.isEmpty() + || fieldsSpecified.stream().allMatch(MetadataColumns::isMetadataColumn), + error("'%s' specifies unknown field(s): %s"), + param, + fieldsSpecified); } // TODO(#34168, ahmedabu98): fill these gaps for the existing batch source @@ -441,7 +462,6 @@ void validate(Table table) { checkArgument( getToTimestamp() == null || getToSnapshot() == null, error("only one of 'to_timestamp' or 'to_snapshot' can be set")); - @Nullable Long fromSnapshotId = ReadUtils.getFromSnapshotInclusive(table, this); @Nullable Long toSnapshotId = ReadUtils.getToSnapshot(table, this); if (fromSnapshotId != null) { @@ -464,11 +484,76 @@ void validate(Table table) { toSnapshotId); } + if (fromSnapshotId != null) { + checkArgumentNotNull( + table.snapshot(fromSnapshotId), + error("configured starting snapshot does not exist: '%s'"), + fromSnapshotId); + } + if (toSnapshotId != null) { + checkArgumentNotNull( + table.snapshot(toSnapshotId), + error("configured end snapshot does not exist: '%s'"), + toSnapshotId); + } + if (fromSnapshotId != null && toSnapshotId != null) { + checkArgument( + SnapshotUtil.isAncestorOf(table, toSnapshotId, fromSnapshotId), + error("fromSnapshot '%s' is not an ancestor of toSnapshot '%s'"), + fromSnapshotId, + toSnapshotId); + } + if (getPollInterval() != null) { checkArgument( Boolean.TRUE.equals(getStreaming()), error("'poll_interval_seconds' can only be set when streaming is true")); } + + @Nullable String watermarkColumn = getWatermarkColumn(); + if (watermarkColumn != null) { + checkArgument(getUseCdc(), error("'watermark_column' is only supported in CDC mode")); + NestedField field = table.schema().findField(watermarkColumn); + checkArgument( + field != null, error("'watermark_column' refers to unknown column: %s"), watermarkColumn); + checkArgument( + field.isRequired(), + error("'watermark_column' needs to be a non-nullable column: %s"), + watermarkColumn); + checkArgument( + field.type().typeId() == TIMESTAMP || field.type().typeId() == LONG, + error("'watermark_column' must be a timestamp-typed column, but '%s' has type %s"), + watermarkColumn, + field.type().typeId()); + checkArgumentNotNull( + getProjectedSchema().findField(watermarkColumn), + "'watermark_column' column should not be dropped."); + } + + @Nullable String watermarkColumnTimeUnit = getWatermarkColumnTimeUnit(); + if (watermarkColumnTimeUnit != null) { + checkArgument( + table + .schema() + .findField( + checkStateNotNull( + watermarkColumn, + "watermark_column_time_unit is configured without a specified watermark_column")) + .type() + .typeId() + == LONG, + error("watermark_column_time_unit is only applicable for LONG columns.")); + try { + TimeUnit.valueOf(watermarkColumnTimeUnit.toUpperCase(Locale.ENGLISH)); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException( + error( + String.format( + "watermark_column_time_unit '%s' is invalid. Please choose one of: %s", + watermarkColumnTimeUnit, Arrays.toString(TimeUnit.values()))), + e); + } + } } private void validateMetadataColumns(Table table) { diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java index 309205707a95..229d30602c64 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java @@ -46,6 +46,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.StructLike; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.catalog.TableIdentifierParser; import org.apache.iceberg.data.GenericRecord; @@ -109,8 +110,8 @@ private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) { return Schema.FieldType.STRING; case UUID: case BINARY: - return Schema.FieldType.BYTES; case FIXED: + return Schema.FieldType.BYTES; case DECIMAL: return Schema.FieldType.DECIMAL; case STRUCT: @@ -360,7 +361,8 @@ private static void copyFieldIntoRecord(Record rec, Types.NestedField field, Row .ifPresent(v -> rec.setField(name, UUID.nameUUIDFromBytes(v))); break; case FIXED: - throw new UnsupportedOperationException("Fixed-precision fields are not yet supported."); + Optional.ofNullable(value.getBytes(name)).ifPresent(v -> rec.setField(name, v)); + break; case BINARY: Optional.ofNullable(value.getBytes(name)) .ifPresent(v -> rec.setField(name, ByteBuffer.wrap(v))); @@ -471,120 +473,150 @@ private static Object getIcebergTimestampValue(Object beamValue, boolean shouldA } } + /** Converts a {@link StructLike} to a Beam {@link Row}. */ + public static Row structToRow(Schema schema, StructLike struct) { + checkState( + schema.getFieldCount() == struct.size(), + "Struct of size %s does not match expected schema size %s", + struct.size(), + schema.getFieldCount()); + Row.Builder rowBuilder = Row.withSchema(schema); + for (int i = 0; i < schema.getFieldCount(); i++) { + Schema.Field field = schema.getField(i); + @Nullable Object icebergValue = struct.get(i, Object.class); + addIcebergValue(rowBuilder, field, icebergValue); + } + return rowBuilder.build(); + } + /** Converts an Iceberg {@link Record} to a Beam {@link Row}. */ public static Row icebergRecordToBeamRow(Schema schema, Record record) { Row.Builder rowBuilder = Row.withSchema(schema); for (Schema.Field field : schema.getFields()) { - boolean isNullable = field.getType().getNullable(); @Nullable Object icebergValue = record.getField(field.getName()); - if (icebergValue == null) { - if (isNullable) { - rowBuilder.addValue(null); - continue; - } - throw new RuntimeException( - String.format("Received null value for required field '%s'.", field.getName())); + addIcebergValue(rowBuilder, field, icebergValue); + } + return rowBuilder.build(); + } + + private static void addIcebergValue( + Row.Builder rowBuilder, Schema.Field field, @Nullable Object icebergValue) { + boolean isNullable = field.getType().getNullable(); + if (icebergValue == null) { + if (isNullable) { + rowBuilder.addValue(null); + return; } - switch (field.getType().getTypeName()) { - case BYTE: - case INT16: - case INT32: - case INT64: - case DECIMAL: // Iceberg and Beam both use BigDecimal - case FLOAT: // Iceberg and Beam both use float - case DOUBLE: // Iceberg and Beam both use double - case STRING: // Iceberg and Beam both use String - case BOOLEAN: // Iceberg and Beam both use boolean - rowBuilder.addValue(icebergValue); - break; - case ARRAY: - checkState( - icebergValue instanceof List, - "Expected List type for field '%s' but received %s", - field.getName(), - icebergValue.getClass()); - List<@NonNull ?> beamList = (List<@NonNull ?>) icebergValue; - Schema.FieldType collectionType = - checkStateNotNull(field.getType().getCollectionElementType()); - // recurse on struct types - if (collectionType.getTypeName().isCompositeType()) { - Schema innerSchema = checkStateNotNull(collectionType.getRowSchema()); - beamList = - beamList.stream() - .map(v -> icebergRecordToBeamRow(innerSchema, (Record) v)) - .collect(Collectors.toList()); - } - rowBuilder.addValue(beamList); - break; - case ITERABLE: - checkState( - icebergValue instanceof Iterable, - "Expected Iterable type for field '%s' but received %s", - field.getName(), - icebergValue.getClass()); - Iterable<@NonNull ?> beamIterable = (Iterable<@NonNull ?>) icebergValue; - Schema.FieldType iterableCollectionType = - checkStateNotNull(field.getType().getCollectionElementType()); - // recurse on struct types - if (iterableCollectionType.getTypeName().isCompositeType()) { - Schema innerSchema = checkStateNotNull(iterableCollectionType.getRowSchema()); - ImmutableList.Builder builder = ImmutableList.builder(); - for (Record v : (Iterable<@NonNull Record>) icebergValue) { - builder.add(icebergRecordToBeamRow(innerSchema, v)); - } - beamIterable = builder.build(); + throw new RuntimeException( + String.format("Received null value for required field '%s'.", field.getName())); + } + switch (field.getType().getTypeName()) { + case BYTE: + case INT16: + case INT32: + case INT64: + case DECIMAL: // Iceberg and Beam both use BigDecimal + case FLOAT: // Iceberg and Beam both use float + case DOUBLE: // Iceberg and Beam both use double + case STRING: // Iceberg and Beam both use String + case BOOLEAN: // Iceberg and Beam both use boolean + rowBuilder.addValue(icebergValue); + break; + case ARRAY: + checkState( + icebergValue instanceof List, + "Expected List type for field '%s' but received %s", + field.getName(), + icebergValue.getClass()); + List<@NonNull ?> beamList = (List<@NonNull ?>) icebergValue; + Schema.FieldType collectionType = + checkStateNotNull(field.getType().getCollectionElementType()); + // recurse on struct types + if (collectionType.getTypeName().isCompositeType()) { + Schema innerSchema = checkStateNotNull(collectionType.getRowSchema()); + beamList = + beamList.stream() + .map(v -> icebergRecordToBeamRow(innerSchema, (Record) v)) + .collect(Collectors.toList()); + } + rowBuilder.addValue(beamList); + break; + case ITERABLE: + checkState( + icebergValue instanceof Iterable, + "Expected Iterable type for field '%s' but received %s", + field.getName(), + icebergValue.getClass()); + Iterable<@NonNull ?> beamIterable = (Iterable<@NonNull ?>) icebergValue; + Schema.FieldType iterableCollectionType = + checkStateNotNull(field.getType().getCollectionElementType()); + // recurse on struct types + if (iterableCollectionType.getTypeName().isCompositeType()) { + Schema innerSchema = checkStateNotNull(iterableCollectionType.getRowSchema()); + ImmutableList.Builder builder = ImmutableList.builder(); + for (Record v : (Iterable<@NonNull Record>) icebergValue) { + builder.add(icebergRecordToBeamRow(innerSchema, v)); } - rowBuilder.addValue(beamIterable); - break; - case MAP: - checkState( - icebergValue instanceof Map, - "Expected Map type for field '%s' but received %s", - field.getName(), - icebergValue.getClass()); - Map beamMap = (Map) icebergValue; - Schema.FieldType valueType = checkStateNotNull(field.getType().getMapValueType()); - // recurse on struct types - if (valueType.getTypeName().isCompositeType()) { - Schema innerSchema = checkStateNotNull(valueType.getRowSchema()); - ImmutableMap.Builder newMap = ImmutableMap.builder(); - for (Map.Entry entry : ((Map) icebergValue).entrySet()) { - Record rec = ((Record) entry.getValue()); - newMap.put( - checkStateNotNull(entry.getKey()), - icebergRecordToBeamRow(innerSchema, checkStateNotNull(rec))); - } - beamMap = newMap.build(); + beamIterable = builder.build(); + } + rowBuilder.addValue(beamIterable); + break; + case MAP: + checkState( + icebergValue instanceof Map, + "Expected Map type for field '%s' but received %s", + field.getName(), + icebergValue.getClass()); + Map beamMap = (Map) icebergValue; + Schema.FieldType valueType = checkStateNotNull(field.getType().getMapValueType()); + // recurse on struct types + if (valueType.getTypeName().isCompositeType()) { + Schema innerSchema = checkStateNotNull(valueType.getRowSchema()); + ImmutableMap.Builder newMap = ImmutableMap.builder(); + for (Map.Entry entry : ((Map) icebergValue).entrySet()) { + Record rec = ((Record) entry.getValue()); + newMap.put( + checkStateNotNull(entry.getKey()), + icebergRecordToBeamRow(innerSchema, checkStateNotNull(rec))); } - rowBuilder.addValue(beamMap); - break; - case DATETIME: - // Iceberg uses a long for micros. - // Beam DATETIME uses joda's DateTime, which only supports millis, - // so we do lose some precision here - rowBuilder.addValue(getBeamDateTimeValue(icebergValue)); - break; - case BYTES: - // Iceberg uses ByteBuffer; Beam uses byte[] - rowBuilder.addValue(((ByteBuffer) icebergValue).array()); - break; - case ROW: - Record nestedRecord = (Record) icebergValue; - Schema nestedSchema = - checkArgumentNotNull( - field.getType().getRowSchema(), - "Corrupted schema: Row type did not have associated nested schema."); - rowBuilder.addValue(icebergRecordToBeamRow(nestedSchema, nestedRecord)); - break; - case LOGICAL_TYPE: - rowBuilder.addValue(getLogicalTypeValue(icebergValue, field.getType())); - break; - default: + beamMap = newMap.build(); + } + rowBuilder.addValue(beamMap); + break; + case DATETIME: + // Iceberg uses a long for micros. + // Beam DATETIME uses joda's DateTime, which only supports millis, + // so we do lose some precision here + rowBuilder.addValue(getBeamDateTimeValue(icebergValue)); + break; + case BYTES: + // Beam uses byte[]. Iceberg represents `binary` as a ByteBuffer but `fixed` as a byte[]. + rowBuilder.addValue( + icebergValue instanceof byte[] + ? (byte[]) icebergValue + : ((ByteBuffer) icebergValue).array()); + break; + case ROW: + Schema nestedSchema = + checkArgumentNotNull( + field.getType().getRowSchema(), + "Corrupted schema: Row type did not have associated nested schema."); + if (icebergValue instanceof Record) { + rowBuilder.addValue(icebergRecordToBeamRow(nestedSchema, (Record) icebergValue)); + } else if (icebergValue instanceof StructLike) { + rowBuilder.addValue(structToRow(nestedSchema, (StructLike) icebergValue)); + } else { throw new UnsupportedOperationException( - "Unsupported Beam type: " + field.getType().getTypeName()); - } + "Unsupported row type: " + icebergValue.getClass()); + } + break; + case LOGICAL_TYPE: + rowBuilder.addValue(getLogicalTypeValue(icebergValue, field.getType())); + break; + default: + throw new UnsupportedOperationException( + "Unsupported Beam type: " + field.getType().getTypeName()); } - return rowBuilder.build(); } private static DateTime getBeamDateTimeValue(Object icebergValue) { diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java index 147e2adda1a7..8a3a543854d9 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcOutputUtils.java @@ -106,8 +106,7 @@ static Schema readBeamSchemaWithRowMetadata(List metadataColumns, Schema static Row outputRow( List metadataColumns, Schema outputSchema, - long commitSnapshotId, - long snapshotSequentNumber, + ChangelogDescriptor descriptor, ValueKind valueKind, Row dataAndRowMetadata) { if (metadataColumns.isEmpty() @@ -115,6 +114,9 @@ static Row outputRow( return dataAndRowMetadata; } + long commitSnapshotId = descriptor.getCommitSnapshotId(); + long snapshotSequentNumber = descriptor.getSnapshotSequenceNumber(); + List<@Nullable Object> values = new ArrayList<>(outputSchema.getFieldCount()); for (Schema.Field field : dataAndRowMetadata.getSchema().getFields()) { if (!metadataColumns.contains(field.getName())) { diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java index daa0a2c73fb8..383dad5a9919 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java @@ -70,6 +70,7 @@ */ public final class CdcReadUtils { private static final Logger LOG = LoggerFactory.getLogger(CdcReadUtils.class); + static final int COMPRESSED_TO_DECODED_BYTES_ESTIMATE = 4; /** * Maximum size of an equality delete set to push down as a Parquet residual {@code IN} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolver.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolver.java new file mode 100644 index 000000000000..0323c621e169 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolver.java @@ -0,0 +1,150 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.BiConsumer; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.iceberg.data.Record; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Helper class to reconcile CDC rows. Used by {@link ResolveChanges} (with Beam {@link Row}s) and + * {@link LocalResolveDoFn} (with Iceberg {@link Record}s). + * + *

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

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

General implementation: + * + *

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

Both input lists are inspected in their given order. + */ + final void resolve(List deletes, List inserts, BiConsumer emit) { + boolean hasDeletes = !deletes.isEmpty(); + boolean hasInserts = !inserts.isEmpty(); + + if (hasInserts && hasDeletes) { + // First, check if any (delete, insert) pairs are duplicates that should not be + // included in the output + boolean[] dupDeletes = new boolean[deletes.size()]; + boolean[] dupInserts = new boolean[inserts.size()]; + + // Map hash to insert-indices + Map> insertHashToIdx = new HashMap<>(); + for (int insertIdx = 0; insertIdx < inserts.size(); insertIdx++) { + int insertHash = nonPkHash(inserts.get(insertIdx)); + insertHashToIdx.computeIfAbsent(insertHash, k -> new ArrayList<>()).add(insertIdx); + } + for (int deleteIdx = 0; deleteIdx < deletes.size(); deleteIdx++) { + int deleteHash = nonPkHash(deletes.get(deleteIdx)); + @Nullable List candidates = insertHashToIdx.get(deleteHash); + if (candidates != null) { + // check if candidates are just duplicates (e.g. from CoW) + for (int idx = 0; idx < candidates.size(); idx++) { + int insertIdx = candidates.get(idx); + if (!dupInserts[insertIdx] + && nonPkEquals(deletes.get(deleteIdx), inserts.get(insertIdx))) { + // this (delete, insert) pair is a duplicate --> should be skipped + dupDeletes[deleteIdx] = true; + dupInserts[insertIdx] = true; + candidates.remove(idx); + break; + } + } + } + } + + // Emit matched pairs as UPDATE_BEFORE / UPDATE_AFTER. + int d = 0; + int i = 0; + while (d < deletes.size() && i < inserts.size()) { + // skip duplicates + while (d < deletes.size() && dupDeletes[d]) { + d++; + } + while (i < inserts.size() && dupInserts[i]) { + i++; + } + + if (d < deletes.size() && i < inserts.size()) { + emit.accept(ValueKind.UPDATE_BEFORE, deletes.get(d)); + emit.accept(ValueKind.UPDATE_AFTER, inserts.get(i)); + d++; + i++; + } + } + + // emit unmatched extras as DELETE / INSERT. + while (d < deletes.size()) { + if (!dupDeletes[d]) { + emit.accept(ValueKind.DELETE, deletes.get(d)); + } + d++; + } + while (i < inserts.size()) { + if (!dupInserts[i]) { + emit.accept(ValueKind.INSERT, inserts.get(i)); + } + i++; + } + } else if (hasInserts) { + for (T r : inserts) { + emit.accept(ValueKind.INSERT, r); + } + } else if (hasDeletes) { + for (T r : deletes) { + emit.accept(ValueKind.DELETE, r); + } + } + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcRowDescriptor.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcRowDescriptor.java new file mode 100644 index 000000000000..3bf3e8a619fd --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcRowDescriptor.java @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import com.google.auto.value.AutoValue; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.SchemaCoder; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptor; + +/** + * Shuffle key for bidirectional CDC rows. + * + *

The primary key isolates rows for update resolution. The snapshot sequence number and commit + * snapshot id carry commit-sourced metadata through {@link ResolveChanges}, where they can be + * appended to final output rows if requested. + */ +@DefaultSchema(AutoValueSchema.class) +@AutoValue +public abstract class CdcRowDescriptor { + @SuppressWarnings("nullness") + public static SchemaCoder coder(Schema identifierSchema) { + Schema descriptorSchema = + Schema.builder() + .addInt64Field("snapshotSequenceNumber") + .addInt64Field("commitSnapshotId") + .addRowField("primaryKey", identifierSchema) + .build(); + + return SchemaCoder.of( + descriptorSchema, + TypeDescriptor.of(CdcRowDescriptor.class), + descriptor -> + Row.withSchema(descriptorSchema) + .addValues( + descriptor.getSnapshotSequenceNumber(), + descriptor.getCommitSnapshotId(), + descriptor.getPrimaryKey()) + .build(), + row -> + CdcRowDescriptor.builder() + .setSnapshotSequenceNumber(row.getInt64("snapshotSequenceNumber")) + .setCommitSnapshotId(row.getInt64("commitSnapshotId")) + .setPrimaryKey(row.getRow("primaryKey")) + .build()); + } + + public static Builder builder() { + return new AutoValue_CdcRowDescriptor.Builder(); + } + + @SchemaFieldNumber("0") + public abstract long getSnapshotSequenceNumber(); + + @SchemaFieldNumber("1") + public abstract long getCommitSnapshotId(); + + @SchemaFieldNumber("2") + public abstract Row getPrimaryKey(); + + @AutoValue.Builder + public abstract static class Builder { + abstract Builder setSnapshotSequenceNumber(long sequenceNumber); + + abstract Builder setCommitSnapshotId(long snapshotId); + + abstract Builder setPrimaryKey(Row primaryKey); + + abstract CdcRowDescriptor build(); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ChangelogScanner.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ChangelogScanner.java index 17ab4c5d30cc..979d2f973081 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ChangelogScanner.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/ChangelogScanner.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.io.iceberg.cdc; import static java.lang.String.format; +import static org.apache.beam.sdk.io.iceberg.cdc.CdcReadUtils.COMPRESSED_TO_DECODED_BYTES_ESTIMATE; import static org.apache.beam.sdk.io.iceberg.cdc.SerializableChangelogTask.Type.ADDED_ROWS; import static org.apache.beam.sdk.io.iceberg.cdc.SerializableChangelogTask.getDataFile; import static org.apache.beam.sdk.io.iceberg.cdc.SerializableChangelogTask.getLength; @@ -662,10 +663,12 @@ static AnalysisResult analyzeFiles( * path otherwise. * *

For LOCAL routing, all bi-directional tasks for this snapshot/partition group are emitted as - * a batch so that the downstream {@link LocalResolveDoFn} can resolve them together in-memory. // - * * The total byte size may exceed {@code splitSize}, but the in-memory // * footprint is bounded - * by the overlap byte estimate (the local resolver still does per-record PK // * routing to avoid - * buffering records outside the overlap range). + * a batch so that the downstream {@link LocalResolveDoFn} can resolve them together in-memory. We + * only take this path when the group's estimated decoded size fits within {@code splitSize}. This + * bounds a single thread to roughly that footprint in the worst case (when metrics are missing or + * there is a very large overlap). The footprint is typically much smaller though: when PK bounds + * are available the resolver further prunes records outside the overlap range via per-record PK + * routing. * *

Returns the number of tasks routed to LOCAL so the caller can update counters. */ @@ -698,7 +701,7 @@ private void routeBidirectional( result.bidirectional.stream().map(t -> makeTask(t, table)).collect(Collectors.toList()); // If the batch is small enough, we can route to LOCAL (in-memory) resolver - if (totalBytes <= splitSize(table)) { + if (totalBytes * COMPRESSED_TO_DECODED_BYTES_ESTIMATE <= splitSize(table)) { Instant ts = Instant.ofEpochMilli(snapshot.timestampMillis()); multiOutputReceiver .get(SMALL_BIDIRECTIONAL_TASKS) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java new file mode 100644 index 000000000000..a3188b3a0245 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFn.java @@ -0,0 +1,245 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.apache.beam.sdk.io.iceberg.IcebergUtils.icebergSchemaToBeamSchema; +import static org.apache.beam.sdk.io.iceberg.cdc.SerializableChangelogTask.Type.ADDED_ROWS; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Objects; +import java.util.Set; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.TableCache; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.join.CoGroupByKey; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeMap; +import org.apache.iceberg.util.StructLikeUtil; +import org.apache.iceberg.util.StructProjection; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Resolves a small bi-directional changelog group entirely in memory. This is the equivalent of + * {@link ReadFromChangelogs} + {@link CoGroupByKey} + {@link ResolveChanges}. + * + *

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

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

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

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

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

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

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

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

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

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

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

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

If this DoFn is configured with {@link #bidirectional}, we compare against the Primary Key + * overlap range. If within the overlap, we key by (snapshotId, PK) and out to either {@link + * #BIDIRECTIONAL_INSERTS} or {@link #BIDIRECTIONAL_DELETES}. Otherwise (not in overlap), we + * output the record directly to {@link #UNIDIRECTIONAL_ROWS}. + */ + private void processTaskRecords( + ChangelogDescriptor descriptor, + SerializableChangelogTask task, + @Nullable Row overlapLowerRow, + @Nullable Row overlapUpperRow, + Table table, + MultiOutputReceiver outputReceiver) + throws IOException { + OverlapRange ovl = checkStateNotNull(overlap); + @Nullable StructLike overlapLower = ovl.toStructLike(overlapLowerRow); + @Nullable StructLike overlapUpper = ovl.toStructLike(overlapUpperRow); + + boolean isInsert = task.getType() == ADDED_ROWS; + TupleTag> taggedOutput = + isInsert ? BIDIRECTIONAL_INSERTS : BIDIRECTIONAL_DELETES; + ValueKind kind = isInsert ? ValueKind.INSERT : ValueKind.DELETE; + long commitSnapshotId = descriptor.getCommitSnapshotId(); + long commitSnapshotSequenceNumber = descriptor.getSnapshotSequenceNumber(); + + Schema readSchema = keyedOutput ? fullBeamRowSchema : projectedBeamRowSchema; + try (CloseableIterable records = + CdcReadUtils.changelogRecordsForTask(task, table, scanConfig, !keyedOutput)) { + for (Record rec : records) { + // uni-directional -- just output records (they are already projected by read pushdown) + if (!keyedOutput) { + Row row = icebergRecordToBeamRow(projectedBeamRowSchema, rec); + outputReceiver + .get(UNIDIRECTIONAL_ROWS) + .builder( + CdcOutputUtils.outputRow( + scanConfig.getMetadataColumns(), + outputBeamRowSchema, + descriptor, + kind, + row)) + .setValueKind(kind) + .output(); + continue; + } + + // bi-directional -- compare overlap + if (ovl.contains(rec, overlapLower, overlapUpper)) { + // inside overlap -- read full row and output KV + Row row = icebergRecordToBeamRow(readSchema, rec); + Row pk = structToRow(scanConfig.rowIdBeamSchema(), pkProjector().wrap(rec)); + outputReceiver + .get(taggedOutput) + .builder( + KV.of( + CdcRowDescriptor.builder() + .setCommitSnapshotId(commitSnapshotId) + .setSnapshotSequenceNumber(commitSnapshotSequenceNumber) + .setPrimaryKey(pk) + .build(), + row)) + .setValueKind(kind) + .output(); + + } else { + // outside overlap -- get projected record and output + StructLike projected = outputProjector().wrap(rec); + Row row = structToRow(projectedBeamRowSchema, projected); + outputReceiver + .get(UNIDIRECTIONAL_ROWS) + .builder( + CdcOutputUtils.outputRow( + scanConfig.getMetadataColumns(), + outputBeamRowSchema, + descriptor, + kind, + row)) + .setValueKind(kind) + .output(); + } + } + } + + trackMetrics(task.getType()); + } + + private StructProjection outputProjector() { + if (outputProjector == null) { + outputProjector = + StructProjection.create( + CdcOutputUtils.readSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), + TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()) + .schema()), + CdcOutputUtils.readSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), scanConfig.getProjectedSchema())); + } + return outputProjector; + } + + private StructProjection pkProjector() { + if (pkProjector == null) { + pkProjector = + StructProjection.create( + CdcOutputUtils.readSchemaWithRowMetadata( + scanConfig.getMetadataColumns(), + TableCache.get(scanConfig.getCatalogConfig(), scanConfig.getTableIdentifier()) + .schema()), + scanConfig.recordIdSchema()); + } + return pkProjector; + } + + private void trackMetrics(SerializableChangelogTask.Type type) { + switch (type) { + case ADDED_ROWS: + numAddedRowsScanTasksCompleted.inc(); + break; + case DELETED_ROWS: + numDeletedRowsScanTasksCompleted.inc(); + break; + case DELETED_FILE: + numDeletedDataFileScanTasksCompleted.inc(); + break; + } + } + + private String getKind(SerializableChangelogTask.Type taskType) { + switch (taskType) { + case ADDED_ROWS: + return "INSERT"; + case DELETED_ROWS: + return "DELETE"; + case DELETED_FILE: + default: + return "DELETE-DF"; + } + } + + @GetSize + public double getSize( + @Element KV> element, + @Restriction OffsetRange restriction) { + // TODO(ahmedabu98): can we make this estimate more accurate? + long size = 0; + + for (long l = restriction.getFrom(); l < restriction.getTo(); l++) { + SerializableChangelogTask task = element.getValue().get((int) l); + size += task.getLength() * COMPRESSED_TO_DECODED_BYTES_ESTIMATE; + size += + task.getAddedDeletes().stream() + .mapToLong(SerializableDeleteFile::getFileSizeInBytes) + .sum() + * COMPRESSED_TO_DECODED_BYTES_ESTIMATE; + size += + task.getExistingDeletes().stream() + .mapToLong(SerializableDeleteFile::getFileSizeInBytes) + .sum() + * COMPRESSED_TO_DECODED_BYTES_ESTIMATE; + } + + return size; + } + + @GetInitialRestriction + public OffsetRange getInitialRange( + @Element KV> element) { + return new OffsetRange(0, element.getValue().size()); + } + + @SplitRestriction + public void splitRestriction( + @Restriction OffsetRange restriction, OutputReceiver out) { + // Split into individual tasks for maximum initial parallelism + for (long i = restriction.getFrom(); i < restriction.getTo(); i++) { + out.output(new OffsetRange(i, i + 1)); + } + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java index 3da31ecc2061..de82bbca4252 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java @@ -262,7 +262,10 @@ public void testTimestampWithZone() { } @Test - public void testFixed() {} + public void testFixed() { + byte[] bytes = new byte[] {1, 2, 3, 4}; + checkRowValueToRecordValue(Schema.FieldType.BYTES, bytes, Types.FixedType.ofLength(4), bytes); + } @Test public void testBinary() { @@ -459,7 +462,10 @@ public void testTimestampWithZone() { } @Test - public void testFixed() {} + public void testFixed() { + byte[] bytes = new byte[] {1, 2, 3, 4}; + checkRecordValueToRowValue(Types.FixedType.ofLength(4), bytes, Schema.FieldType.BYTES, bytes); + } @Test public void testBinary() { diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolverTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolverTest.java new file mode 100644 index 000000000000..4f1340efb604 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/CdcResolverTest.java @@ -0,0 +1,118 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; +import static org.hamcrest.Matchers.empty; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link CdcResolver}. */ +@RunWith(JUnit4.class) +public class CdcResolverTest { + private static final TestResolver RESOLVER = new TestResolver(); + + @Test + public void duplicateDeleteInsertIsDropped() { + List emitted = + resolve( + Collections.singletonList(item("same", 7)), Collections.singletonList(item("same", 7))); + + assertThat(emitted, empty()); + } + + @Test + public void changedDeleteInsertBecomesUpdatePair() { + List emitted = + resolve( + Collections.singletonList(item("before", 1)), + Collections.singletonList(item("after", 2))); + + assertThat(emitted, contains("UPDATE_BEFORE:before", "UPDATE_AFTER:after")); + } + + @Test + public void duplicateUpdateAndSingletonsResolveByMultiplicity() { + List emitted = + resolve( + Arrays.asList(item("copy", 1), item("old", 2), item("deleted-only", 3)), + Arrays.asList(item("copy", 1), item("new", 4))); + + assertThat(emitted, contains("UPDATE_BEFORE:old", "UPDATE_AFTER:new", "DELETE:deleted-only")); + } + + @Test + public void hashCollisionOnlyConsumesEqualInsertOnce() { + List emitted = + resolve( + Arrays.asList(item("copy", 9), item("deleted-only", 9)), + Collections.singletonList(item("copy", 9))); + + assertThat(emitted, contains("DELETE:deleted-only")); + } + + @Test + public void hashMatchAloneDoesNotDeduplicate() { + List emitted = + resolve( + Collections.singletonList(item("before", 42)), + Collections.singletonList(item("after", 42))); + + assertThat(emitted, contains("UPDATE_BEFORE:before", "UPDATE_AFTER:after")); + } + + private static Item item(String nonPkValue, int hash) { + return new Item(nonPkValue, hash); + } + + private static List resolve(List deletes, List inserts) { + List emitted = new ArrayList<>(); + RESOLVER.resolve( + deletes, inserts, (kind, item) -> emitted.add(kind.name() + ":" + item.nonPkValue)); + return emitted; + } + + private static class TestResolver extends CdcResolver { + @Override + protected int nonPkHash(Item element) { + return element.hash; + } + + @Override + protected boolean nonPkEquals(Item delete, Item insert) { + return delete.nonPkValue.equals(insert.nonPkValue); + } + } + + private static class Item { + private final String nonPkValue; + private final int hash; + + private Item(String nonPkValue, int hash) { + this.nonPkValue = nonPkValue; + this.hash = hash; + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java new file mode 100644 index 000000000000..3a56dd9c83fc --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java @@ -0,0 +1,305 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; +import static org.hamcrest.Matchers.empty; +import static org.junit.Assert.assertEquals; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.iceberg.ChangelogOperation; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.expressions.ExpressionParser; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.types.Types; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.rules.TestName; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Integration tests for {@link LocalResolveDoFn}. */ +@RunWith(JUnit4.class) +public class LocalResolveDoFnTest { + private static final org.apache.iceberg.Schema CDC_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "visible", Types.StringType.get()), + Types.NestedField.optional(3, "hidden", Types.StringType.get())), + ImmutableSet.of(1)); + + private static final org.apache.iceberg.Schema FIXED_CDC_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "visible", Types.StringType.get()), + Types.NestedField.optional(3, "data", Types.FixedType.ofLength(4))), + ImmutableSet.of(1)); + + @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); + @Rule public TestName testName = new TestName(); + + @Test + public void copyOnWriteRewriteOfIdenticalRowsIsDropped() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId); + DataFile oldFile = + warehouse.writeRecords( + testName.getMethodName() + "-old.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "same-hidden"))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "same-hidden"))); + + List> output = + process( + scanConfig, + descriptor(tableId, 1L, 1L), + ImmutableList.of( + task(SerializableChangelogTask.Type.DELETED_FILE, oldFile, table, 300L), + task(SerializableChangelogTask.Type.ADDED_ROWS, newFile, table, 300L)), + new Instant(0L)); + + assertThat(output, empty()); + } + + @Test + public void hiddenOnlyUpdateIsResolvedBeforeProjection() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId); + DataFile oldFile = + warehouse.writeRecords( + testName.getMethodName() + "-old.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "old-hidden"))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "new-hidden"))); + Instant timestamp = new Instant(1234L); + + List> output = + process( + scanConfig, + descriptor(tableId, 1L, 1L), + ImmutableList.of( + task(SerializableChangelogTask.Type.DELETED_FILE, oldFile, table, 301L), + task(SerializableChangelogTask.Type.ADDED_ROWS, newFile, table, 301L)), + timestamp); + + assertThat( + output.stream().map(LocalResolveDoFnTest::kindAndProjectedRow).collect(Collectors.toList()), + contains("UPDATE_BEFORE:1:shown:2", "UPDATE_AFTER:1:shown:2")); + assertEquals( + ImmutableList.of(timestamp, timestamp), + output.stream().map(ValueInSingleWindow::getTimestamp).collect(Collectors.toList())); + } + + @Test + public void copyOnWriteRewriteWithFixedColumnIsDropped() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, FIXED_CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId); + // Distinct byte[] instances with identical content. The CoW no-op is only dropped when the + // `fixed` column is hashed/compared by content; an identity hashCode would leak a spurious + // UPDATE pair. + DataFile oldFile = + warehouse.writeRecords( + testName.getMethodName() + "-old.parquet", + table.schema(), + ImmutableList.of(fixedRecord(1L, "shown", new byte[] {1, 2, 3, 4}))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(fixedRecord(1L, "shown", new byte[] {1, 2, 3, 4}))); + + List> output = + process( + scanConfig, + descriptor(tableId, 1L, 1L), + ImmutableList.of( + task(SerializableChangelogTask.Type.DELETED_FILE, oldFile, table, 300L), + task(SerializableChangelogTask.Type.ADDED_ROWS, newFile, table, 300L)), + new Instant(0L)); + + assertThat(output, empty()); + } + + @Test + public void differingFixedColumnBecomesUpdatePair() throws Exception { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, FIXED_CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId); + // Same PK and projected fields, but the `fixed` column differs, so this must NOT be treated as + // a CoW duplicate -- guards against the fixed column being ignored during resolution. + DataFile oldFile = + warehouse.writeRecords( + testName.getMethodName() + "-old.parquet", + table.schema(), + ImmutableList.of(fixedRecord(1L, "shown", new byte[] {1, 2, 3, 4}))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(fixedRecord(1L, "shown", new byte[] {5, 6, 7, 8}))); + + List> output = + process( + scanConfig, + descriptor(tableId, 1L, 1L), + ImmutableList.of( + task(SerializableChangelogTask.Type.DELETED_FILE, oldFile, table, 302L), + task(SerializableChangelogTask.Type.ADDED_ROWS, newFile, table, 302L)), + new Instant(0L)); + + assertThat( + output.stream().map(LocalResolveDoFnTest::kindAndProjectedRow).collect(Collectors.toList()), + contains("UPDATE_BEFORE:1:shown:2", "UPDATE_AFTER:1:shown:2")); + } + + private List> process( + IcebergScanConfig scanConfig, + ChangelogDescriptor descriptor, + List tasks, + Instant timestamp) + throws Exception { + try (DoFnTester>, Row> tester = + DoFnTester.of(new LocalResolveDoFn(scanConfig))) { + tester.processTimestampedElement(TimestampedValue.of(KV.of(descriptor, tasks), timestamp)); + return tester.getMutableOutput(tester.getMainOutputTag()); + } + } + + private TableIdentifier tableId() { + return TableIdentifier.of("default", testName.getMethodName()); + } + + private IcebergScanConfig scanConfig(Table table, TableIdentifier tableId) { + return IcebergScanConfig.builder() + .setCatalogConfig( + IcebergCatalogConfig.builder() + .setCatalogName("name") + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build()) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(table.schema())) + .setKeepFields(ImmutableList.of("id", "visible")) + .setUseCdc(true) + .build(); + } + + private static ChangelogDescriptor descriptor( + TableIdentifier tableId, long lowerInclusive, long upperInclusive) { + org.apache.beam.sdk.schemas.Schema pkSchema = + org.apache.beam.sdk.schemas.Schema.builder().addInt64Field("id").build(); + return ChangelogDescriptor.builder() + .setTableIdentifierString(tableId.toString()) + .setSnapshotSequenceNumber(1) + .setCommitSnapshotId(1) + .setOverlapLower(Row.withSchema(pkSchema).addValue(lowerInclusive).build()) + .setOverlapUpper(Row.withSchema(pkSchema).addValue(upperInclusive).build()) + .build(); + } + + private static Record record(long id, String visible, String hidden) { + GenericRecord record = GenericRecord.create(CDC_SCHEMA); + record.setField("id", id); + record.setField("visible", visible); + record.setField("hidden", hidden); + return record; + } + + private static Record fixedRecord(long id, String visible, byte[] data) { + GenericRecord record = GenericRecord.create(FIXED_CDC_SCHEMA); + record.setField("id", id); + record.setField("visible", visible); + record.setField("data", data); + return record; + } + + private static SerializableChangelogTask task( + SerializableChangelogTask.Type type, DataFile dataFile, Table table, long snapshotId) { + return SerializableChangelogTask.builder() + .setType(type) + .setDataFile(dataFile, table.spec().partitionToPath(dataFile.partition()), true) + .setAddedDeletes(ImmutableList.of()) + .setExistingDeletes(ImmutableList.of()) + .setSpecId(table.spec().specId()) + .setOperation( + type == SerializableChangelogTask.Type.ADDED_ROWS + ? ChangelogOperation.INSERT + : ChangelogOperation.DELETE) + .setOrdinal(0) + .setCommitSnapshotId(snapshotId) + .setStart(0L) + .setLength(dataFile.fileSizeInBytes()) + .setJsonExpression(ExpressionParser.toJson(Expressions.alwaysTrue())) + .build(); + } + + private static Map tableProperties() { + return ImmutableMap.of(TableProperties.FORMAT_VERSION, "2"); + } + + private static String kindAndProjectedRow(ValueInSingleWindow value) { + ValueKind kind = value.getValueKind(); + Row row = value.getValue(); + return kind.name() + + ":" + + row.getInt64("id") + + ":" + + row.getString("visible") + + ":" + + row.getSchema().getFieldCount(); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRangeTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRangeTest.java new file mode 100644 index 000000000000..69e47182e173 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/OverlapRangeTest.java @@ -0,0 +1,161 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import java.io.IOException; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.rules.TestName; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link OverlapRange}. */ +@RunWith(JUnit4.class) +public class OverlapRangeTest { + private static final org.apache.iceberg.Schema SINGLE_PK_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "data", Types.StringType.get())), + ImmutableSet.of(1)); + + private static final org.apache.iceberg.Schema COMPOSITE_PK_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.optional(3, "data", Types.StringType.get()), + Types.NestedField.required(1, "account", Types.StringType.get()), + Types.NestedField.optional(4, "extra", Types.IntegerType.get()), + Types.NestedField.required(2, "sequence", Types.IntegerType.get())), + ImmutableSet.of(1, 2)); + + @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); + @Rule public final TestName testName = new TestName(); + + @Test + public void containsUsesInclusiveSingleColumnBounds() throws Exception { + OverlapRange range = overlapRange(SINGLE_PK_SCHEMA); + StructLike lower = range.toStructLike(pkRow(range.recordIdSchema(), 10)); + StructLike upper = range.toStructLike(pkRow(range.recordIdSchema(), 20)); + + assertFalse(range.contains(singlePkRecord(9), lower, upper)); + assertTrue(range.contains(singlePkRecord(10), lower, upper)); + assertTrue(range.contains(singlePkRecord(15), lower, upper)); + assertTrue(range.contains(singlePkRecord(20), lower, upper)); + assertFalse(range.contains(singlePkRecord(21), lower, upper)); + } + + @Test + public void containsUsesLexicographicCompositeBounds() throws Exception { + OverlapRange range = overlapRange(COMPOSITE_PK_SCHEMA); + StructLike lower = range.toStructLike(pkRow(range.recordIdSchema(), "a", 2)); + StructLike upper = range.toStructLike(pkRow(range.recordIdSchema(), "b", 1)); + + assertFalse(range.contains(compositePkRecord("a", 1), lower, upper)); + assertTrue(range.contains(compositePkRecord("a", 2), lower, upper)); + assertTrue(range.contains(compositePkRecord("a", 9), lower, upper)); + assertTrue(range.contains(compositePkRecord("b", 0), lower, upper)); + assertTrue(range.contains(compositePkRecord("b", 1), lower, upper)); + assertFalse(range.contains(compositePkRecord("b", 2), lower, upper)); + } + + @Test + public void nullBoundsAreConservative() throws Exception { + OverlapRange range = overlapRange(SINGLE_PK_SCHEMA); + StructLike lower = range.toStructLike(pkRow(range.recordIdSchema(), 10)); + StructLike upper = range.toStructLike(pkRow(range.recordIdSchema(), 20)); + + assertNull(range.toStructLike(null)); + assertTrue(range.contains(singlePkRecord(1), null, upper)); + assertTrue(range.contains(singlePkRecord(100), lower, null)); + assertTrue(range.contains(singlePkRecord(100), null, null)); + } + + @Test + public void recordIdProjectionUsesIdentifierFieldsFromFullRecord() throws Exception { + OverlapRange range = overlapRange(COMPOSITE_PK_SCHEMA); + StructLike lower = range.toStructLike(pkRow(range.recordIdSchema(), "acct", 7)); + StructLike upper = range.toStructLike(pkRow(range.recordIdSchema(), "acct", 7)); + + assertTrue(range.contains(compositePkRecord("acct", 7), lower, upper)); + + assertEquals("acct", range.recordIdProjection().get(0, String.class)); + assertEquals(7, (int) range.recordIdProjection().get(1, Integer.class)); + } + + private OverlapRange overlapRange(org.apache.iceberg.Schema schema) throws IOException { + TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName()); + IcebergCatalogConfig catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build(); + catalogConfig.catalog().createTable(tableId, schema); + IcebergScanConfig scanConfig = + IcebergScanConfig.builder() + .setCatalogConfig(catalogConfig) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(schema)) + .setUseCdc(true) + .build(); + return OverlapRange.forScanConfig(scanConfig); + } + + private static Row pkRow(Schema recordIdSchema, Object... values) { + return Row.withSchema(IcebergUtils.icebergSchemaToBeamSchema(recordIdSchema)) + .addValues(values) + .build(); + } + + private static Record singlePkRecord(int id) { + GenericRecord record = GenericRecord.create(SINGLE_PK_SCHEMA); + record.setField("id", id); + record.setField("data", "v" + id); + return record; + } + + private static Record compositePkRecord(String account, int sequence) { + GenericRecord record = GenericRecord.create(COMPOSITE_PK_SCHEMA); + record.setField("data", "payload"); + record.setField("account", account); + record.setField("extra", 100); + record.setField("sequence", sequence); + return record; + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogsTest.java new file mode 100644 index 000000000000..69591e6eaa7c --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ReadFromChangelogsTest.java @@ -0,0 +1,366 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.iceberg.cdc; + +import static org.junit.Assert.assertEquals; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; +import org.apache.beam.sdk.io.iceberg.IcebergUtils; +import org.apache.beam.sdk.io.iceberg.SerializableDeleteFile; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.iceberg.ChangelogOperation; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.deletes.PositionDelete; +import org.apache.iceberg.deletes.PositionDeleteWriter; +import org.apache.iceberg.encryption.EncryptedFiles; +import org.apache.iceberg.expressions.ExpressionParser; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.rules.TestName; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ReadFromChangelogs}. */ +@RunWith(JUnit4.class) +public class ReadFromChangelogsTest { + private static final org.apache.iceberg.Schema CDC_SCHEMA = + new org.apache.iceberg.Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "visible", Types.StringType.get()), + Types.NestedField.optional(3, "hidden", Types.StringType.get())), + ImmutableSet.of(1)); + + @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); + @Rule public TestName testName = new TestName(); + @Rule public TestPipeline pipeline = TestPipeline.create(); + + @Test + public void unidirectionalTasksEmitProjectedRowsOnly() throws IOException { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId, ImmutableList.of("id", "visible")); + + DataFile addedFile = + warehouse.writeRecords( + testName.getMethodName() + "-added.parquet", + table.schema(), + ImmutableList.of(record(10L, "added", "added-hidden"))); + DataFile deletedRowsFile = + warehouse.writeRecords( + testName.getMethodName() + "-deleted-rows.parquet", + table.schema(), + ImmutableList.of( + record(20L, "deleted-row", "deleted-row-hidden"), + record(21L, "not-deleted", "not-deleted-hidden"))); + DeleteFile addedPositionDelete = + writePositionDelete(table, deletedRowsFile, "deleted-rows-pos-delete.parquet", 0L); + DataFile deletedFile = + warehouse.writeRecords( + testName.getMethodName() + "-deleted-file.parquet", + table.schema(), + ImmutableList.of(record(30L, "deleted-file", "deleted-file-hidden"))); + + List tasks = + ImmutableList.of( + task( + SerializableChangelogTask.Type.ADDED_ROWS, + addedFile, + ImmutableList.of(), + ImmutableList.of(), + table, + 100L), + task( + SerializableChangelogTask.Type.DELETED_ROWS, + deletedRowsFile, + ImmutableList.of(addedPositionDelete), + ImmutableList.of(), + table, + 100L), + task( + SerializableChangelogTask.Type.DELETED_FILE, + deletedFile, + ImmutableList.of(), + ImmutableList.of(), + table, + 100L)); + + ReadFromChangelogs.Output output = + input(ImmutableList.of(KV.of(descriptor(), tasks)), ImmutableList.of()) + .apply(new ReadFromChangelogs(scanConfig)); + + assertEquals( + IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()), + output.uniDirectionalRows().getSchema()); + PAssert.that( + output.uniDirectionalRows().apply("Format Unidirectional", ParDo.of(new FormatRow()))) + .containsInAnyOrder( + "INSERT:10:added:2", "DELETE:20:deleted-row:2", "DELETE:30:deleted-file:2"); + PAssert.that(output.biDirectionalInserts()).empty(); + PAssert.that(output.biDirectionalDeletes()).empty(); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void bidirectionalTasksKeepFullRowsForDownstreamResolution() throws IOException { + TableIdentifier tableId = tableId(); + Table table = warehouse.createTable(tableId, CDC_SCHEMA, null, tableProperties()); + IcebergScanConfig scanConfig = scanConfig(table, tableId, ImmutableList.of("id", "visible")); + DataFile oldFile = + warehouse.writeRecords( + testName.getMethodName() + "-old.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "old-hidden"))); + DataFile newFile = + warehouse.writeRecords( + testName.getMethodName() + "-new.parquet", + table.schema(), + ImmutableList.of(record(1L, "shown", "new-hidden"))); + List tasks = + ImmutableList.of( + task( + SerializableChangelogTask.Type.DELETED_FILE, + oldFile, + ImmutableList.of(), + ImmutableList.of(), + table, + 200L), + task( + SerializableChangelogTask.Type.ADDED_ROWS, + newFile, + ImmutableList.of(), + ImmutableList.of(), + table, + 200L)); + + ReadFromChangelogs.Output output = + input(ImmutableList.of(), ImmutableList.of(KV.of(descriptor(200L, 200L, 1L, 1L), tasks))) + .apply(new ReadFromChangelogs(scanConfig)); + + PAssert.that(output.uniDirectionalRows()).empty(); + PAssert.that( + output.biDirectionalDeletes().apply("Format Deletes", ParDo.of(new FormatKeyedRow()))) + .containsInAnyOrder("DELETE:200:200:1:shown:old-hidden:3"); + PAssert.that( + output.biDirectionalInserts().apply("Format Inserts", ParDo.of(new FormatKeyedRow()))) + .containsInAnyOrder("INSERT:200:200:1:shown:new-hidden:3"); + + pipeline.run().waitUntilFinish(); + } + + private PCollectionTuple input( + List>> unidirectional, + List>> largeBidirectional) { + Schema rowIdBeamSchema = + IcebergUtils.icebergSchemaToBeamSchema( + TypeUtil.select(CDC_SCHEMA, CDC_SCHEMA.identifierFieldIds())); + KvCoder> coder = + ChangelogScanner.coder(rowIdBeamSchema); + PCollection>> uni = + unidirectional.isEmpty() + ? pipeline.apply("Empty Unidirectional", Create.empty(coder)) + : pipeline.apply("Create Unidirectional", Create.of(unidirectional).withCoder(coder)); + PCollection>> large = + largeBidirectional.isEmpty() + ? pipeline.apply("Empty Large Bidirectional", Create.empty(coder)) + : pipeline.apply( + "Create Large Bidirectional", Create.of(largeBidirectional).withCoder(coder)); + return PCollectionTuple.of(ChangelogScanner.UNIDIRECTIONAL_TASKS, uni) + .and(ChangelogScanner.LARGE_BIDIRECTIONAL_TASKS, large); + } + + private TableIdentifier tableId() { + return TableIdentifier.of("default", testName.getMethodName()); + } + + private IcebergScanConfig scanConfig( + Table table, TableIdentifier tableId, List keepFields) { + return IcebergScanConfig.builder() + .setCatalogConfig( + IcebergCatalogConfig.builder() + .setCatalogName("name") + .setCatalogProperties( + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .build()) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(table.schema())) + .setKeepFields(keepFields) + .setUseCdc(true) + .build(); + } + + private ChangelogDescriptor descriptor() { + return ChangelogDescriptor.builder() + .setTableIdentifierString(tableId().toString()) + .setSnapshotSequenceNumber(100L) + .setCommitSnapshotId(100L) + .build(); + } + + private ChangelogDescriptor descriptor( + long sequenceNumber, long snapshotId, long lowerInclusive, long upperInclusive) { + Schema pkSchema = Schema.builder().addInt64Field("id").build(); + return ChangelogDescriptor.builder() + .setTableIdentifierString(tableId().toString()) + .setSnapshotSequenceNumber(sequenceNumber) + .setCommitSnapshotId(snapshotId) + .setOverlapLower(Row.withSchema(pkSchema).addValue(lowerInclusive).build()) + .setOverlapUpper(Row.withSchema(pkSchema).addValue(upperInclusive).build()) + .build(); + } + + private static Record record(long id, String visible, String hidden) { + GenericRecord record = GenericRecord.create(CDC_SCHEMA); + record.setField("id", id); + record.setField("visible", visible); + record.setField("hidden", hidden); + return record; + } + + private static DeleteFile writePositionDelete( + Table table, DataFile dataFile, String filename, long... positions) throws IOException { + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec()); + PositionDeleteWriter writer = + appenderFactory.newPosDeleteWriter( + EncryptedFiles.plainAsEncryptedOutput( + table.io().newOutputFile(dataFile.location() + "." + filename)), + FileFormat.PARQUET, + null); + try { + for (long position : positions) { + writer.write(PositionDelete.create().set(dataFile.location(), position)); + } + } finally { + writer.close(); + } + return writer.toDeleteFile(); + } + + private static SerializableChangelogTask task( + SerializableChangelogTask.Type type, + DataFile dataFile, + List addedDeletes, + List existingDeletes, + Table table, + long snapshotId) { + return SerializableChangelogTask.builder() + .setType(type) + .setDataFile(dataFile, table.spec().partitionToPath(dataFile.partition()), true) + .setAddedDeletes(serializableDeletes(addedDeletes, table)) + .setExistingDeletes(serializableDeletes(existingDeletes, table)) + .setSpecId(table.spec().specId()) + .setOperation( + type == SerializableChangelogTask.Type.ADDED_ROWS + ? ChangelogOperation.INSERT + : ChangelogOperation.DELETE) + .setOrdinal(0) + .setCommitSnapshotId(snapshotId) + .setStart(0L) + .setLength(dataFile.fileSizeInBytes()) + .setJsonExpression(ExpressionParser.toJson(Expressions.alwaysTrue())) + .build(); + } + + private static List serializableDeletes( + List deletes, Table table) { + return deletes.stream() + .map( + delete -> + SerializableDeleteFile.from( + delete, table.spec().partitionToPath(delete.partition()), true)) + .collect(Collectors.toList()); + } + + private static Map tableProperties() { + return ImmutableMap.of(TableProperties.FORMAT_VERSION, "2"); + } + + private static class FormatRow extends DoFn { + @ProcessElement + public void process(@Element Row row, ValueKind kind, OutputReceiver out) { + out.output( + kind.name() + + ":" + + row.getInt64("id") + + ":" + + row.getString("visible") + + ":" + + row.getSchema().getFieldCount()); + } + } + + private static class FormatKeyedRow extends DoFn, String> { + @ProcessElement + public void process( + @Element KV element, ValueKind kind, OutputReceiver out) { + Row row = element.getValue(); + CdcRowDescriptor descriptor = element.getKey(); + out.output( + kind.name() + + ":" + + descriptor.getCommitSnapshotId() + + ":" + + descriptor.getSnapshotSequenceNumber() + + ":" + + descriptor.getPrimaryKey().getInt64("id") + + ":" + + row.getString("visible") + + ":" + + row.getString("hidden") + + ":" + + row.getSchema().getFieldCount()); + } + } +}