diff --git a/core/src/main/java/org/apache/iceberg/io/BaseTaskWriter.java b/core/src/main/java/org/apache/iceberg/io/BaseTaskWriter.java index 4bde586a0f3b..e0e5638263cb 100644 --- a/core/src/main/java/org/apache/iceberg/io/BaseTaskWriter.java +++ b/core/src/main/java/org/apache/iceberg/io/BaseTaskWriter.java @@ -35,12 +35,10 @@ import org.apache.iceberg.deletes.PositionDelete; import org.apache.iceberg.deletes.SortingPositionOnlyDeleteWriter; import org.apache.iceberg.encryption.EncryptedOutputFile; -import org.apache.iceberg.relocated.com.google.common.base.MoreObjects; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.relocated.com.google.common.collect.Iterables; import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.iceberg.util.CharSequenceSet; -import org.apache.iceberg.util.StructLikeMap; import org.apache.iceberg.util.StructLikeUtil; import org.apache.iceberg.util.StructProjection; import org.apache.iceberg.util.Tasks; @@ -183,7 +181,9 @@ protected abstract class BaseEqualityDeltaWriter implements Closeable { private RollingFileWriter dataWriter; private RollingEqDeleteWriter eqDeleteWriter; private PartitioningWriter, DeleteWriteResult> posDeleteWriter; - private Map insertedRowMap; + private Map insertedRowMap; + private final PositionDeleteTracker insertedRowTracker; + private final boolean closeInsertedRowTracker; private boolean closePosDeleteWriter; protected BaseEqualityDeltaWriter(StructLike partition, Schema schema, Schema deleteSchema) { @@ -204,6 +204,16 @@ protected BaseEqualityDeltaWriter( Schema deleteSchema, DeleteGranularity deleteGranularity, PartitioningDVWriter posDeleteWriter) { + this(partition, schema, deleteSchema, deleteGranularity, posDeleteWriter, null); + } + + protected BaseEqualityDeltaWriter( + StructLike partition, + Schema schema, + Schema deleteSchema, + DeleteGranularity deleteGranularity, + PartitioningDVWriter posDeleteWriter, + PositionDeleteTracker insertedRowTracker) { Preconditions.checkNotNull(schema, "Iceberg table schema cannot be null."); Preconditions.checkNotNull(deleteSchema, "Equality-delete schema cannot be null."); this.structProjection = StructProjection.create(schema, deleteSchema); @@ -215,7 +225,12 @@ protected BaseEqualityDeltaWriter( posDeleteWriter != null ? posDeleteWriter : createPosDeleteWriter(partition, deleteGranularity); - this.insertedRowMap = StructLikeMap.create(deleteSchema.asStruct()); + this.insertedRowTracker = + insertedRowTracker != null + ? insertedRowTracker + : new PositionDeleteTracker(deleteSchema.asStruct()); + this.closeInsertedRowTracker = insertedRowTracker == null; + this.insertedRowMap = this.insertedRowTracker.insertedRows(); this.partitionKey = partition; } @@ -226,13 +241,14 @@ protected BaseEqualityDeltaWriter( protected abstract StructLike asStructLikeKey(T key); public void write(T row) throws IOException { - PathOffset pathOffset = PathOffset.of(dataWriter.currentPath(), dataWriter.currentRows()); + PositionDeleteTracker.PathOffset pathOffset = + PositionDeleteTracker.PathOffset.of(dataWriter.currentPath(), dataWriter.currentRows()); // Create a copied key from this row. StructLike copiedKey = StructLikeUtil.copy(structProjection.wrap(asStructLike(row))); // Adding a pos-delete to replace the old path-offset. - PathOffset previous = insertedRowMap.put(copiedKey, pathOffset); + PositionDeleteTracker.PathOffset previous = insertedRowMap.put(copiedKey, pathOffset); if (previous != null) { // TODO attach the previous row if has a positional-delete row schema in appender factory. writePosDelete(previous); @@ -260,8 +276,8 @@ private EncryptedOutputFile newOutputFile(StructLike partition) { } } - private void writePosDelete(PathOffset pathOffset) { - positionDelete.set(pathOffset.path, pathOffset.rowOffset); + private void writePosDelete(PositionDeleteTracker.PathOffset pathOffset) { + positionDelete.set(pathOffset.path(), pathOffset.rowOffset()); posDeleteWriter.write(positionDelete, spec, partitionKey); } @@ -271,7 +287,7 @@ private void writePosDelete(PathOffset pathOffset) { * @param key has the same columns with the equality fields. */ private boolean internalPosDelete(StructLike key) { - PathOffset previous = insertedRowMap.remove(key); + PositionDeleteTracker.PathOffset previous = insertedRowMap.remove(key); if (previous != null) { // TODO attach the previous row if has a positional-delete row schema in appender factory. @@ -327,8 +343,8 @@ public void close() throws IOException { } } - if (insertedRowMap != null) { - insertedRowMap.clear(); + if (closeInsertedRowTracker) { + insertedRowTracker.clear(); insertedRowMap = null; } @@ -351,28 +367,6 @@ public void close() throws IOException { } } - private static class PathOffset { - private final CharSequence path; - private final long rowOffset; - - private PathOffset(CharSequence path, long rowOffset) { - this.path = path; - this.rowOffset = rowOffset; - } - - private static PathOffset of(CharSequence path, long rowOffset) { - return new PathOffset(path, rowOffset); - } - - @Override - public String toString() { - return MoreObjects.toStringHelper(this) - .add("path", path) - .add("row_offset", rowOffset) - .toString(); - } - } - private abstract class BaseRollingWriter implements Closeable { private static final int ROWS_DIVISOR = 1000; private final StructLike partitionKey; diff --git a/core/src/main/java/org/apache/iceberg/io/PositionDeleteTracker.java b/core/src/main/java/org/apache/iceberg/io/PositionDeleteTracker.java new file mode 100644 index 000000000000..171847d135f6 --- /dev/null +++ b/core/src/main/java/org/apache/iceberg/io/PositionDeleteTracker.java @@ -0,0 +1,67 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.io; + +import java.util.Map; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeMap; + +/** Tracks data-file positions for equality keys written by related task writers. */ +public final class PositionDeleteTracker { + private final Map insertedRowMap; + + /** Creates a tracker for the equality-key struct type. */ + public PositionDeleteTracker(Types.StructType keyType) { + this.insertedRowMap = + StructLikeMap.create(Preconditions.checkNotNull(keyType, "Key type cannot be null")); + } + + Map insertedRows() { + return insertedRowMap; + } + + /** Clears all tracked rows. */ + public void clear() { + insertedRowMap.clear(); + } + + static final class PathOffset { + private final CharSequence path; + private final long rowOffset; + + private PathOffset(CharSequence path, long rowOffset) { + this.path = path; + this.rowOffset = rowOffset; + } + + CharSequence path() { + return path; + } + + long rowOffset() { + return rowOffset; + } + + static PathOffset of(CharSequence path, long rowOffset) { + return new PathOffset(path, rowOffset); + } + } +} diff --git a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/BaseDeltaTaskWriter.java b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/BaseDeltaTaskWriter.java index 33a09705e702..a1ac761419b4 100644 --- a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/BaseDeltaTaskWriter.java +++ b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/BaseDeltaTaskWriter.java @@ -36,6 +36,7 @@ import org.apache.iceberg.io.FileWriterFactory; import org.apache.iceberg.io.OutputFileFactory; import org.apache.iceberg.io.PartitioningDVWriter; +import org.apache.iceberg.io.PositionDeleteTracker; import org.apache.iceberg.relocated.com.google.common.collect.Sets; import org.apache.iceberg.types.TypeUtil; @@ -47,6 +48,7 @@ abstract class BaseDeltaTaskWriter extends BaseTaskWriter { private final RowDataWrapper keyWrapper; private final RowDataProjection keyProjection; private final boolean upsert; + private final PositionDeleteTracker positionDeleteTracker; BaseDeltaTaskWriter( PartitionSpec spec, @@ -59,7 +61,8 @@ abstract class BaseDeltaTaskWriter extends BaseTaskWriter { RowType flinkSchema, Set equalityFieldIds, boolean upsert, - boolean useDv) { + boolean useDv, + PositionDeleteTracker positionDeleteTracker) { super(spec, format, fileWriterFactory, fileFactory, io, targetFileSize, useDv); this.schema = schema; this.deleteSchema = TypeUtil.select(schema, Sets.newHashSet(equalityFieldIds)); @@ -69,6 +72,7 @@ abstract class BaseDeltaTaskWriter extends BaseTaskWriter { this.keyProjection = RowDataProjection.create(flinkSchema, schema.asStruct(), deleteSchema.asStruct()); this.upsert = upsert; + this.positionDeleteTracker = positionDeleteTracker; } abstract RowDataDeltaWriter route(RowData row); @@ -112,7 +116,13 @@ public void write(RowData row) throws IOException { protected class RowDataDeltaWriter extends BaseEqualityDeltaWriter { RowDataDeltaWriter(PartitionKey partition, PartitioningDVWriter dvFileWriter) { - super(partition, schema, deleteSchema, DeleteGranularity.FILE, dvFileWriter); + super( + partition, + schema, + deleteSchema, + DeleteGranularity.FILE, + dvFileWriter, + positionDeleteTracker); } @Override diff --git a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/PartitionedDeltaWriter.java b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/PartitionedDeltaWriter.java index 5e597d8e71a5..16133670525b 100644 --- a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/PartitionedDeltaWriter.java +++ b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/PartitionedDeltaWriter.java @@ -31,6 +31,7 @@ import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.FileWriterFactory; import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.io.PositionDeleteTracker; import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.apache.iceberg.util.Tasks; @@ -51,7 +52,8 @@ class PartitionedDeltaWriter extends BaseDeltaTaskWriter { RowType flinkSchema, Set equalityFieldIds, boolean upsert, - boolean useDv) { + boolean useDv, + PositionDeleteTracker positionDeleteTracker) { super( spec, format, @@ -63,7 +65,8 @@ class PartitionedDeltaWriter extends BaseDeltaTaskWriter { flinkSchema, equalityFieldIds, upsert, - useDv); + useDv, + positionDeleteTracker); this.partitionKey = new PartitionKey(spec, schema); } diff --git a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/RowDataTaskWriterFactory.java b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/RowDataTaskWriterFactory.java index bc3bc51cedc4..c00da957f23f 100644 --- a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/RowDataTaskWriterFactory.java +++ b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/RowDataTaskWriterFactory.java @@ -35,6 +35,7 @@ import org.apache.iceberg.io.FileWriterFactory; import org.apache.iceberg.io.OutputFileFactory; import org.apache.iceberg.io.PartitionedFanoutWriter; +import org.apache.iceberg.io.PositionDeleteTracker; import org.apache.iceberg.io.TaskWriter; import org.apache.iceberg.io.UnpartitionedWriter; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; @@ -184,6 +185,10 @@ public void initialize(int taskId, int attemptId) { @Override public TaskWriter create() { + return create(null); + } + + public TaskWriter create(PositionDeleteTracker positionDeleteTracker) { Preconditions.checkNotNull( outputFileFactory, "The outputFileFactory shouldn't be null if we have invoked the initialize()."); @@ -225,7 +230,8 @@ public TaskWriter create() { flinkSchema, equalityFieldIds, upsert, - useDv); + useDv, + positionDeleteTracker); } else { return new PartitionedDeltaWriter( spec, @@ -238,7 +244,8 @@ public TaskWriter create() { flinkSchema, equalityFieldIds, upsert, - useDv); + useDv, + positionDeleteTracker); } } } diff --git a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/UnpartitionedDeltaWriter.java b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/UnpartitionedDeltaWriter.java index 9d749d3062af..9f7dd90fee6a 100644 --- a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/UnpartitionedDeltaWriter.java +++ b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/UnpartitionedDeltaWriter.java @@ -28,6 +28,7 @@ import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.FileWriterFactory; import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.io.PositionDeleteTracker; class UnpartitionedDeltaWriter extends BaseDeltaTaskWriter { private final RowDataDeltaWriter writer; @@ -43,7 +44,8 @@ class UnpartitionedDeltaWriter extends BaseDeltaTaskWriter { RowType flinkSchema, Set equalityFieldIds, boolean upsert, - boolean useDv) { + boolean useDv, + PositionDeleteTracker positionDeleteTracker) { super( spec, format, @@ -55,7 +57,8 @@ class UnpartitionedDeltaWriter extends BaseDeltaTaskWriter { flinkSchema, equalityFieldIds, upsert, - useDv); + useDv, + positionDeleteTracker); this.writer = new RowDataDeltaWriter(null, dvFileWriter()); } diff --git a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicWriter.java b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicWriter.java index fcd0d082704a..245b99815012 100644 --- a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicWriter.java +++ b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicWriter.java @@ -23,6 +23,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Set; import java.util.concurrent.TimeUnit; import org.apache.flink.annotation.VisibleForTesting; @@ -38,12 +39,15 @@ import org.apache.iceberg.flink.FlinkWriteConf; import org.apache.iceberg.flink.sink.RowDataTaskWriterFactory; import org.apache.iceberg.flink.sink.SinkUtil; +import org.apache.iceberg.io.PositionDeleteTracker; import org.apache.iceberg.io.TaskWriter; import org.apache.iceberg.io.WriteResult; import org.apache.iceberg.relocated.com.google.common.base.MoreObjects; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -58,6 +62,7 @@ class DynamicWriter implements CommittingSinkWriter taskWriterFactories; private final Map> writers; + private final Map positionDeleteTrackers; private final Configuration flinkConfig; private final Map commonWriteProperties; private final DynamicWriterMetrics metrics; @@ -81,6 +86,7 @@ class DynamicWriter implements CommittingSinkWriter(cacheMaximumSize); this.writers = Maps.newHashMap(); + this.positionDeleteTrackers = Maps.newHashMap(); LOG.debug("DynamicIcebergSinkWriter created for subtask {} attemptId {}", subTaskId, attemptId); } @@ -142,8 +148,28 @@ public void write(DynamicRecordInternal element, Context context) element.spec()); }); + PositionDeleteTracker positionDeleteTracker = null; + if (element.upsertMode()) { + Set equalityFieldIds = writerKey.equalityFields(); + if (equalityFieldIds == null || equalityFieldIds.isEmpty()) { + Table table = catalog.loadTable(TableIdentifier.parse(writerKey.tableName())); + equalityFieldIds = getEqualityFields(table, equalityFieldIds); + } + Types.StructType equalityKeyType = + TypeUtil.select(element.schema(), equalityFieldIds).asStruct(); + PositionDeleteTrackerKey trackerKey = + new PositionDeleteTrackerKey( + writerKey.tableName(), + writerKey.branch(), + writerKey.specId(), + equalityKeyType); + positionDeleteTracker = + positionDeleteTrackers.computeIfAbsent( + trackerKey, key -> new PositionDeleteTracker(equalityKeyType)); + } + taskWriterFactory.initialize(subTaskId, attemptId); - return taskWriterFactory.create(); + return taskWriterFactory.create(positionDeleteTracker); }) .write(element.rowData()); metrics.mainMetricsGroup().getNumRecordsSendCounter().inc(); @@ -159,6 +185,7 @@ public void close() throws Exception { for (TaskWriter writer : writers.values()) { writer.close(); } + clearPositionDeleteTrackers(); } @Override @@ -196,10 +223,52 @@ public Collection prepareCommit() throws IOException { } writers.clear(); + clearPositionDeleteTrackers(); return result; } + private void clearPositionDeleteTrackers() { + positionDeleteTrackers.values().forEach(PositionDeleteTracker::clear); + positionDeleteTrackers.clear(); + } + + private static class PositionDeleteTrackerKey { + private final String tableName; + private final String branch; + private final int specId; + private final Types.StructType equalityKeyType; + + private PositionDeleteTrackerKey( + String tableName, String branch, int specId, Types.StructType equalityKeyType) { + this.tableName = tableName; + this.branch = branch; + this.specId = specId; + this.equalityKeyType = equalityKeyType; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + + if (!(other instanceof PositionDeleteTrackerKey that)) { + return false; + } + + return specId == that.specId + && Objects.equals(tableName, that.tableName) + && Objects.equals(branch, that.branch) + && Objects.equals(equalityKeyType, that.equalityKeyType); + } + + @Override + public int hashCode() { + return Objects.hash(tableName, branch, specId, equalityKeyType); + } + } + private static Set getEqualityFields(Table table, Set equalityFieldIds) { if (equalityFieldIds != null && !equalityFieldIds.isEmpty()) { return equalityFieldIds; diff --git a/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java b/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java index 7c4669ff9826..523b6999f2e7 100644 --- a/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java +++ b/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java @@ -56,11 +56,14 @@ import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; import org.apache.flink.streaming.api.graph.StreamNode; import org.apache.flink.table.catalog.ResolvedSchema; +import org.apache.flink.table.data.GenericRowData; import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; import org.apache.flink.table.data.util.DataFormatConverters; import org.apache.flink.table.types.DataType; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.types.Row; +import org.apache.flink.types.RowKind; import org.apache.flink.util.Collector; import org.apache.flink.util.ExceptionUtils; import org.apache.flink.util.function.SerializableSupplier; @@ -120,6 +123,7 @@ void before() { private static class DynamicIcebergDataImpl implements Serializable { Row rowProvided; + RowData rowDataProvided; Row rowExpected; Schema schemaProvided; Schema schemaExpected; @@ -218,7 +222,9 @@ public void generate(DynamicIcebergDataImpl row, Collector out) { tableIdentifier, branch, schema, - converter(schema).toInternal(row.rowProvided), + row.rowDataProvided != null + ? row.rowDataProvided + : converter(schema).toInternal(row.rowProvided), spec, spec.isPartitioned() ? DistributionMode.HASH : DistributionMode.NONE, row.writeParallelism); @@ -622,6 +628,67 @@ void testSchemaEvolutionAddField() throws Exception { runTest(rows, this.env, 1); } + @Test + void upsertAcrossSchemaEvolutionInCheckpoint() throws Exception { + String tableName = "schema_evolution_upsert"; + TableIdentifier tableId = TableIdentifier.of(DATABASE, tableName); + Schema schema = + new Schema( + Types.NestedField.required(1, "order_id", Types.StringType.get()), + Types.NestedField.required(2, "amount", Types.IntegerType.get())); + Schema evolvedSchema = + new Schema( + schema.columns().get(0), + schema.columns().get(1), + Types.NestedField.optional(3, "note", Types.StringType.get())); + PartitionSpec spec = PartitionSpec.unpartitioned(); + CATALOG_EXTENSION + .catalog() + .createTable(tableId, schema, spec, ImmutableMap.of(TableProperties.FORMAT_VERSION, "2")); + + DynamicIcebergDataImpl firstWrite = + new DynamicIcebergDataImpl( + schema, + tableName, + SnapshotRef.MAIN_BRANCH, + spec, + true, + Sets.newHashSet("order_id"), + false); + firstWrite.rowProvided = Row.of("K", 100); + DynamicIcebergDataImpl update = + new DynamicIcebergDataImpl( + evolvedSchema, + tableName, + SnapshotRef.MAIN_BRANCH, + spec, + true, + Sets.newHashSet("order_id"), + false); + GenericRowData updateRow = new GenericRowData(3); + updateRow.setField(0, StringData.fromString("K")); + updateRow.setField(1, 200); + updateRow.setField(2, null); + updateRow.setRowKind(RowKind.UPDATE_AFTER); + update.rowDataProvided = updateRow; + + env = + StreamExecutionEnvironment.getExecutionEnvironment( + MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG); + executeDynamicSink(Lists.newArrayList(firstWrite, update), env, true, 1, null); + + List actualRows; + try (CloseableIterable rows = + IcebergGenerics.read(CATALOG_EXTENSION.catalog().loadTable(tableId)).build()) { + actualRows = Lists.newArrayList(rows); + } + + assertThat(actualRows).hasSize(1); + assertThat(actualRows.get(0).get(0)).isEqualTo("K"); + assertThat(actualRows.get(0).get(1)).isEqualTo(200); + assertThat(actualRows.get(0).get(2)).isNull(); + } + @Test void testRowEvolutionNullMissingOptionalField() throws Exception { List rows =