Add support for deletion vector in Iceberg

Author: ebyhr

docs/src/main/sphinx/connector/iceberg.md

@@ -945,7 +945,7 @@ connector using a {doc}`WITH </sql/create-table-as>` clause.
   - Optionally specifies the file system location URI for the table.
 * - `format_version`
   - Optionally specifies the format version of the Iceberg specification to use
-    for new tables; either `1` or `2`. Defaults to `2`. Version `2` is required
+    for new tables; either `1`, `2` or `3`. Defaults to `2`. Version `2` is required
     for row level deletes.
 * - `max_commit_retry`
   - Number of times to retry a commit before failing. Defaults to the value of 

plugin/trino-iceberg/pom.xml

@@ -223,6 +223,12 @@
             <artifactId>iceberg-core</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>org.apache.iceberg</groupId>
+            <artifactId>iceberg-data</artifactId>
+            <version>${dep.iceberg.version}</version>
+        </dependency>
+
         <dependency>
             <groupId>org.apache.iceberg</groupId>
             <artifactId>iceberg-nessie</artifactId>

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

@@ -13,22 +13,28 @@
  */
 package io.trino.plugin.iceberg;
 
+import com.google.common.collect.ImmutableList;
 import org.apache.iceberg.FileContent;
+import org.apache.iceberg.FileFormat;
 
 import java.util.List;
 import java.util.Optional;
+import java.util.OptionalLong;
 
 import static java.util.Objects.requireNonNull;
 
 public record CommitTaskData(
         String path,
-        IcebergFileFormat fileFormat,
+        FileFormat fileFormat,
         long fileSizeInBytes,
         MetricsWrapper metrics,
         String partitionSpecJson,
         Optional<String> partitionDataJson,
         FileContent content,
         Optional<String> referencedDataFile,
+        List<String> deletionVectorFiles,
+        OptionalLong deletionVectorContentOffset,
+        OptionalLong deletionVectorContentSize,
         Optional<List<Long>> fileSplitOffsets)
 {
     public CommitTaskData
@@ -40,6 +46,9 @@ public record CommitTaskData(
         requireNonNull(partitionDataJson, "partitionDataJson is null");
         requireNonNull(content, "content is null");
         requireNonNull(referencedDataFile, "referencedDataFile is null");
+        deletionVectorFiles = ImmutableList.copyOf(deletionVectorFiles);
         requireNonNull(fileSplitOffsets, "fileSplitOffsets is null");
+        requireNonNull(deletionVectorContentOffset, "deletionVectorContentOffset is null");
+        requireNonNull(deletionVectorContentSize, "deletionVectorContentSize is null");
     }
 }

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

@@ -18,6 +18,7 @@
 import io.trino.spi.Page;
 import io.trino.spi.TrinoException;
 import io.trino.spi.type.Type;
+import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.avro.Avro;
 import org.apache.iceberg.data.Record;
@@ -36,6 +37,7 @@
 import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_WRITER_OPEN_ERROR;
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.util.Objects.requireNonNull;
+import static org.apache.iceberg.FileFormat.AVRO;
 import static org.apache.iceberg.TableProperties.AVRO_COMPRESSION;
 
 public final class IcebergAvroFileWriter
@@ -46,6 +48,7 @@ public final class IcebergAvroFileWriter
     // Use static table name instead of the actual name because it becomes outdated once the table is renamed
     public static final String AVRO_TABLE_NAME = "table";
 
+    private final String location;
     private final Schema icebergSchema;
     private final List<Type> types;
     private final FileAppender<Record> avroWriter;
@@ -58,6 +61,7 @@ public IcebergAvroFileWriter(
             List<Type> types,
             HiveCompressionCodec hiveCompressionCodec)
     {
+        this.location = requireNonNull(file.location(), "location is null");
         this.rollbackAction = requireNonNull(rollbackAction, "rollbackAction null");
         this.icebergSchema = requireNonNull(icebergSchema, "icebergSchema is null");
         this.types = ImmutableList.copyOf(requireNonNull(types, "types is null"));
@@ -71,10 +75,22 @@ public IcebergAvroFileWriter(
                     .build();
         }
         catch (IOException e) {
-            throw new TrinoException(ICEBERG_WRITER_OPEN_ERROR, "Error creating Avro file: " + file.location(), e);
+            throw new TrinoException(ICEBERG_WRITER_OPEN_ERROR, "Error creating Avro file: " + location, e);
         }
     }
 
+    @Override
+    public FileFormat fileFormat()
+    {
+        return AVRO;
+    }
+
+    @Override
+    public String location()
+    {
+        return location;
+    }
+
     @Override
     public long getWrittenBytes()
     {

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

@@ -52,7 +52,8 @@
 public class IcebergConfig
 {
     public static final int FORMAT_VERSION_SUPPORT_MIN = 1;
-    public static final int FORMAT_VERSION_SUPPORT_MAX = 2;
+    private static final int FORMAT_VERSION_DEFAULT = 2;
+    public static final int FORMAT_VERSION_SUPPORT_MAX = 3;
     public static final String EXTENDED_STATISTICS_CONFIG = "iceberg.extended-statistics.enabled";
     public static final String EXTENDED_STATISTICS_DESCRIPTION = "Enable collection (ANALYZE) and use of extended statistics.";
     public static final String COLLECT_EXTENDED_STATISTICS_ON_WRITE_DESCRIPTION = "Collect extended statistics during writes";
@@ -74,7 +75,7 @@ public class IcebergConfig
     private boolean registerTableProcedureEnabled;
     private boolean addFilesProcedureEnabled;
     private Optional<String> hiveCatalogName = Optional.empty();
-    private int formatVersion = FORMAT_VERSION_SUPPORT_MAX;
+    private int formatVersion = FORMAT_VERSION_DEFAULT;
     private Duration expireSnapshotsMinRetention = new Duration(7, DAYS);
     private Duration removeOrphanFilesMinRetention = new Duration(7, DAYS);
     private DataSize targetMaxFileSize = DataSize.of(1, GIGABYTE);

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

@@ -13,7 +13,9 @@
  */
 package io.trino.plugin.iceberg;
 
+import com.google.common.collect.ImmutableList;
 import io.trino.plugin.hive.FileWriter;
+import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.Metrics;
 
 import java.util.List;
@@ -22,6 +24,15 @@
 public interface IcebergFileWriter
         extends FileWriter
 {
+    FileFormat fileFormat();
+
+    String location();
+
+    default List<String> rewrittenDeleteFiles()
+    {
+        return ImmutableList.of();
+    }
+
     FileMetrics getFileMetrics();
 
     record FileMetrics(Metrics metrics, Optional<List<Long>> splitOffsets) {}

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

@@ -33,25 +33,32 @@
 import io.trino.plugin.hive.HiveCompressionCodec;
 import io.trino.plugin.hive.NodeVersion;
 import io.trino.plugin.hive.orc.OrcWriterConfig;
+import io.trino.plugin.iceberg.delete.DeletionVectorWriter;
 import io.trino.plugin.iceberg.fileio.ForwardingOutputFile;
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
+import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.MetricsConfig;
+import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.Schema;
+import org.apache.iceberg.deletes.PositionDeleteIndex;
 import org.apache.iceberg.types.Types;
+import org.apache.iceberg.util.DeleteFileSet;
 import org.weakref.jmx.Managed;
 
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.IntStream;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.trino.plugin.hive.HiveCompressionCodecs.toCompressionCodec;
 import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME;
@@ -82,6 +89,7 @@
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
+import static org.apache.iceberg.MetadataColumns.DELETE_FILE_POS;
 import static org.apache.iceberg.TableProperties.DEFAULT_WRITE_METRICS_MODE;
 import static org.apache.iceberg.io.DeleteSchemaUtil.pathPosSchema;
 import static org.apache.iceberg.parquet.ParquetSchemaUtil.convert;
@@ -139,16 +147,37 @@ public IcebergFileWriter createPositionDeleteWriter(
             TrinoFileSystem fileSystem,
             Location outputPath,
             ConnectorSession session,
-            IcebergFileFormat fileFormat,
-            Map<String, String> storageProperties)
+            String dataFilePath,
+            FileFormat fileFormat,
+            PartitionSpec partitionSpec,
+            Optional<PartitionData> partition,
+            Map<String, String> storageProperties,
+            Map<String, DeleteFileSet> previousDeleteFiles)
     {
         return switch (fileFormat) {
+            case PUFFIN -> createDeletionVectorWriter(nodeVersion, fileSystem, outputPath, dataFilePath, partitionSpec, partition, previousDeleteFiles);
             case PARQUET -> createParquetWriter(FULL_METRICS_CONFIG, fileSystem, outputPath, POSITION_DELETE_SCHEMA, session, storageProperties);
             case ORC -> createOrcWriter(FULL_METRICS_CONFIG, fileSystem, outputPath, POSITION_DELETE_SCHEMA, session, storageProperties, DataSize.ofBytes(Integer.MAX_VALUE));
             case AVRO -> createAvroWriter(fileSystem, outputPath, POSITION_DELETE_SCHEMA, session);
+            case METADATA -> throw new IllegalArgumentException("Unexpected METADATA file format");
         };
     }
 
+    private static DeletionVectorWriter createDeletionVectorWriter(
+            NodeVersion nodeVersion,
+            TrinoFileSystem fileSystem,
+            Location outputPath,
+            String dataFilePath,
+            PartitionSpec partitionSpec,
+            Optional<PartitionData> partition,
+            Map<String, DeleteFileSet> previousDeleteFiles)
+    {
+        Function<CharSequence, PositionDeleteIndex> previousDeleteLoader = DeletionVectorWriter.create(fileSystem, previousDeleteFiles);
+        int positionChannel = POSITION_DELETE_SCHEMA.columns().indexOf(DELETE_FILE_POS);
+        checkState(positionChannel != -1, "positionChannel not found");
+        return new DeletionVectorWriter(nodeVersion, fileSystem, outputPath, dataFilePath, partitionSpec, partition, previousDeleteLoader::apply, positionChannel);
+    }
+
     private IcebergFileWriter createParquetWriter(
             MetricsConfig metricsConfig,
             TrinoFileSystem fileSystem,
@@ -234,6 +263,7 @@ private IcebergFileWriter createOrcWriter(
             }
 
             return new IcebergOrcFileWriter(
+                    outputPath,
                     metricsConfig,
                     icebergSchema,
                     orcDataSink,

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

@@ -30,6 +30,7 @@
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.io.LocationProvider;
 import org.apache.iceberg.types.Type;
+import org.apache.iceberg.util.DeleteFileSet;
 import org.roaringbitmap.longlong.ImmutableLongBitmapDataProvider;
 import org.roaringbitmap.longlong.LongBitmapDataProvider;
 import org.roaringbitmap.longlong.Roaring64Bitmap;
@@ -55,8 +56,10 @@ public class IcebergMergeSink
     private final LocationProvider locationProvider;
     private final IcebergFileWriterFactory fileWriterFactory;
     private final TrinoFileSystem fileSystem;
+    private final Map<String, DeleteFileSet> previousDeleteFiles;
     private final JsonCodec<CommitTaskData> jsonCodec;
     private final ConnectorSession session;
+    private final int formatVersion;
     private final IcebergFileFormat fileFormat;
     private final Map<String, String> storageProperties;
     private final Schema schema;
@@ -69,8 +72,10 @@ public IcebergMergeSink(
             LocationProvider locationProvider,
             IcebergFileWriterFactory fileWriterFactory,
             TrinoFileSystem fileSystem,
+            Map<String, DeleteFileSet> previousDeleteFiles,
             JsonCodec<CommitTaskData> jsonCodec,
             ConnectorSession session,
+            int formatVersion,
             IcebergFileFormat fileFormat,
             Map<String, String> storageProperties,
             Schema schema,
@@ -81,8 +86,10 @@ public IcebergMergeSink(
         this.locationProvider = requireNonNull(locationProvider, "locationProvider is null");
         this.fileWriterFactory = requireNonNull(fileWriterFactory, "fileWriterFactory is null");
         this.fileSystem = requireNonNull(fileSystem, "fileSystem is null");
+        this.previousDeleteFiles = ImmutableMap.copyOf(previousDeleteFiles);
         this.jsonCodec = requireNonNull(jsonCodec, "jsonCodec is null");
         this.session = requireNonNull(session, "session is null");
+        this.formatVersion = formatVersion;
         this.fileFormat = requireNonNull(fileFormat, "fileFormat is null");
         this.storageProperties = ImmutableMap.copyOf(requireNonNull(storageProperties, "storageProperties is null"));
         this.schema = requireNonNull(schema, "schema is null");
@@ -162,8 +169,10 @@ private PositionDeleteWriter createPositionDeleteWriter(String dataFilePath, Par
                 fileSystem,
                 jsonCodec,
                 session,
+                formatVersion,
                 fileFormat,
-                storageProperties);
+                storageProperties,
+                previousDeleteFiles);
     }
 
     private static Collection<Slice> writePositionDeletes(PositionDeleteWriter writer, ImmutableLongBitmapDataProvider rowsToDelete)