Fix Iceberg merge, update, delete, for tables with equality deletes

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergColumnHandle.java

@@ -42,8 +42,7 @@ public class IcebergColumnHandle
     private static final int INSTANCE_SIZE = instanceSize(IcebergColumnHandle.class);
 
     // Iceberg reserved row ids begin at INTEGER.MAX_VALUE and count down. Starting with MIN_VALUE here to avoid conflicts.
-    public static final int TRINO_UPDATE_ROW_ID = Integer.MIN_VALUE;
-    public static final int TRINO_MERGE_ROW_ID = Integer.MIN_VALUE + 1;
+    public static final int TRINO_MERGE_ROW_ID = Integer.MIN_VALUE;
     public static final String TRINO_ROW_ID_NAME = "$row_id";
 
     public static final int TRINO_MERGE_PARTITION_SPEC_ID = Integer.MIN_VALUE + 2;
@@ -182,12 +181,6 @@ public boolean isRowPositionColumn()
         return id == ROW_POSITION.fieldId();
     }
 
-    @JsonIgnore
-    public boolean isUpdateRowIdColumn()
-    {
-        return id == TRINO_UPDATE_ROW_ID;
-    }
-
     @JsonIgnore
     public boolean isMergeRowIdColumn()
     {

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSource.java

@@ -30,13 +30,13 @@
 import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Throwables.throwIfInstanceOf;
 import static io.trino.plugin.base.util.Closables.closeAllSuppress;
 import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_BAD_DATA;
-import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
 public class IcebergPageSource
@@ -46,10 +46,9 @@ public class IcebergPageSource
     private final ConnectorPageSource delegate;
     private final Optional<ReaderProjectionsAdapter> projectionsAdapter;
     private final Supplier<Optional<RowPredicate>> deletePredicate;
-    // An array with one element per field in the $row_id column. The value in the array points to the
-    // channel where the data can be read from.
-    private int[] rowIdChildColumnIndexes = new int[0];
+    private final Function<Block, RowBlock> rowIdBlockFactory;
     // The $row_id's index in 'expectedColumns', or -1 if there isn't one
+    // this column with contain row position populated in the source, and must be wrapped with constant data for full row id
     private int rowIdColumnIndex = -1;
     // Maps the Iceberg field ids of unmodified columns to their indexes in updateRowIdChildColumnIndexes
 
@@ -58,7 +57,8 @@ public IcebergPageSource(
             List<IcebergColumnHandle> requiredColumns,
             ConnectorPageSource delegate,
             Optional<ReaderProjectionsAdapter> projectionsAdapter,
-            Supplier<Optional<RowPredicate>> deletePredicate)
+            Supplier<Optional<RowPredicate>> deletePredicate,
+            Function<Block, RowBlock> rowIdBlockFactory)
     {
         // expectedColumns should contain columns which should be in the final Page
         // requiredColumns should include all expectedColumns as well as any columns needed by the DeleteFilter
@@ -70,22 +70,15 @@ public IcebergPageSource(
             checkArgument(expectedColumn.equals(requiredColumns.get(i)), "Expected columns must be a prefix of required columns");
             expectedColumnIndexes[i] = i;
 
-            if (expectedColumn.isUpdateRowIdColumn() || expectedColumn.isMergeRowIdColumn()) {
+            if (expectedColumn.isMergeRowIdColumn()) {
                 this.rowIdColumnIndex = i;
-
-                Map<Integer, Integer> fieldIdToColumnIndex = mapFieldIdsToIndex(requiredColumns);
-                List<ColumnIdentity> rowIdFields = expectedColumn.getColumnIdentity().getChildren();
-                this.rowIdChildColumnIndexes = new int[rowIdFields.size()];
-                for (int columnIndex = 0; columnIndex < rowIdFields.size(); columnIndex++) {
-                    int fieldId = rowIdFields.get(columnIndex).getId();
-                    rowIdChildColumnIndexes[columnIndex] = requireNonNull(fieldIdToColumnIndex.get(fieldId), () -> format("Column %s not found in requiredColumns", fieldId));
-                }
             }
         }
 
         this.delegate = requireNonNull(delegate, "delegate is null");
         this.projectionsAdapter = requireNonNull(projectionsAdapter, "projectionsAdapter is null");
         this.deletePredicate = requireNonNull(deletePredicate, "deletePredicate is null");
+        this.rowIdBlockFactory = requireNonNull(rowIdBlockFactory, "rowIdBlockFactory is null");
     }
 
     @Override
@@ -161,21 +154,11 @@ private Page withRowIdBlock(Page page)
             return page;
         }
 
-        Block[] rowIdFields = new Block[rowIdChildColumnIndexes.length];
-        for (int childIndex = 0; childIndex < rowIdChildColumnIndexes.length; childIndex++) {
-            rowIdFields[childIndex] = page.getBlock(rowIdChildColumnIndexes[childIndex]);
-        }
-
+        RowBlock rowIdBlock = rowIdBlockFactory.apply(page.getBlock(rowIdColumnIndex));
         Block[] fullPage = new Block[page.getChannelCount()];
         for (int channel = 0; channel < page.getChannelCount(); channel++) {
-            if (channel == rowIdColumnIndex) {
-                fullPage[channel] = RowBlock.fromFieldBlocks(page.getPositionCount(), rowIdFields);
-                continue;
-            }
-
-            fullPage[channel] = page.getBlock(channel);
+            fullPage[channel] = channel == rowIdColumnIndex ? rowIdBlock : page.getBlock(channel);
         }
-
         return new Page(page.getPositionCount(), fullPage);
     }