[spark] spark delete with deletion vector (apache#3390)

Author: YannByron

paimon-common/src/main/java/org/apache/paimon/utils/RoaringBitmap32.java

@@ -23,6 +23,7 @@
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
+import java.util.Objects;
 
 /** A compressed bitmap for 32-bit integer. */
 public class RoaringBitmap32 {
@@ -67,4 +68,16 @@ public void serialize(DataOutput out) throws IOException {
     public void deserialize(DataInput in) throws IOException {
         roaringBitmap.deserialize(in);
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        RoaringBitmap32 that = (RoaringBitmap32) o;
+        return Objects.equals(this.roaringBitmap, that.roaringBitmap);
+    }
 }

paimon-core/src/main/java/org/apache/paimon/deletionvectors/BitmapDeletionVector.java

@@ -24,6 +24,7 @@
 import java.io.DataInput;
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.util.Objects;
 
 /**
  * A {@link DeletionVector} based on {@link RoaringBitmap32}, it only supports files with row count
@@ -104,4 +105,16 @@ private void checkPosition(long position) {
                     "The file has too many rows, RoaringBitmap32 only supports files with row count not exceeding 2147483647.");
         }
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        BitmapDeletionVector that = (BitmapDeletionVector) o;
+        return Objects.equals(this.roaringBitmap, that.roaringBitmap);
+    }
 }

paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionVectorsMaintainer.java

@@ -22,6 +22,7 @@
 import org.apache.paimon.data.BinaryRow;
 import org.apache.paimon.index.IndexFileHandler;
 import org.apache.paimon.index.IndexFileMeta;
+import org.apache.paimon.manifest.IndexManifestEntry;
 
 import javax.annotation.Nullable;
 
@@ -30,6 +31,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 import static org.apache.paimon.deletionvectors.DeletionVectorsIndexFile.DELETION_VECTORS_INDEX;
 
@@ -138,6 +140,20 @@ public DeletionVectorsMaintainer createOrRestore(
             return createOrRestore(deletionVectors);
         }
 
+        @VisibleForTesting
+        public DeletionVectorsMaintainer createOrRestore(
+                @Nullable Long snapshotId, BinaryRow partition) {
+            List<IndexFileMeta> indexFiles =
+                    snapshotId == null
+                            ? Collections.emptyList()
+                            : handler.scan(snapshotId, DELETION_VECTORS_INDEX, partition).stream()
+                                    .map(IndexManifestEntry::indexFile)
+                                    .collect(Collectors.toList());
+            Map<String, DeletionVector> deletionVectors =
+                    new HashMap<>(handler.readAllDeletionVectors(indexFiles));
+            return createOrRestore(deletionVectors);
+        }
+
         public DeletionVectorsMaintainer create() {
             return createOrRestore(new HashMap<>());
         }

paimon-core/src/main/java/org/apache/paimon/index/IndexFileHandler.java

@@ -21,10 +21,12 @@
 import org.apache.paimon.Snapshot;
 import org.apache.paimon.data.BinaryRow;
 import org.apache.paimon.deletionvectors.DeletionVector;
+import org.apache.paimon.deletionvectors.DeletionVectorIndexFileMaintainer;
 import org.apache.paimon.deletionvectors.DeletionVectorsIndexFile;
 import org.apache.paimon.fs.Path;
 import org.apache.paimon.manifest.IndexManifestEntry;
 import org.apache.paimon.manifest.IndexManifestFile;
+import org.apache.paimon.table.source.DeletionFile;
 import org.apache.paimon.utils.IntIterator;
 import org.apache.paimon.utils.Pair;
 import org.apache.paimon.utils.PathFactory;
@@ -185,6 +187,11 @@ public void deleteManifest(String indexManifest) {
         indexManifestFile.delete(indexManifest);
     }
 
+    public DeletionVectorIndexFileMaintainer createDVIndexFileMaintainer(
+            Map<String, DeletionFile> dataFileToDeletionFiles) {
+        return new DeletionVectorIndexFileMaintainer(this, dataFileToDeletionFiles);
+    }
+
     public Map<String, DeletionVector> readAllDeletionVectors(List<IndexFileMeta> fileMetas) {
         Map<String, DeletionVector> deletionVectors = new HashMap<>();
         for (IndexFileMeta indexFile : fileMetas) {

paimon-core/src/main/java/org/apache/paimon/schema/SchemaValidation.java

@@ -184,7 +184,7 @@ public static void validateTableSchema(TableSchema schema) {
                                         field));
 
         if (options.deletionVectorsEnabled()) {
-            validateForDeletionVectors(schema, options);
+            validateForDeletionVectors(options);
         }
     }
 
@@ -461,11 +461,7 @@ private static void validateDefaultValues(TableSchema schema) {
         }
     }
 
-    private static void validateForDeletionVectors(TableSchema schema, CoreOptions options) {
-        checkArgument(
-                !schema.primaryKeys().isEmpty(),
-                "Deletion vectors mode is only supported for tables with primary keys.");
-
+    private static void validateForDeletionVectors(CoreOptions options) {
         checkArgument(
                 options.changelogProducer() == ChangelogProducer.NONE
                         || options.changelogProducer() == ChangelogProducer.LOOKUP,

paimon-core/src/main/java/org/apache/paimon/table/source/DeletionFile.java

@@ -26,6 +26,7 @@
 import javax.annotation.Nullable;
 
 import java.io.IOException;
+import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -44,7 +45,9 @@
  * </ul>
  */
 @Public
-public class DeletionFile {
+public class DeletionFile implements Serializable {
+
+    private static final long serialVersionUID = 1L;
 
     private final String path;
     private final long offset;

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonSplitScan.scala

@@ -18,19 +18,23 @@
 
 package org.apache.paimon.spark
 
+import org.apache.paimon.spark.schema.PaimonMetadataColumn
 import org.apache.paimon.table.Table
 import org.apache.paimon.table.source.{DataSplit, Split}
 
 import org.apache.spark.sql.connector.read.{Batch, Scan}
 import org.apache.spark.sql.types.StructType
 
 /** For internal use only. */
-case class PaimonSplitScan(table: Table, dataSplits: Array[DataSplit]) extends Scan {
-
+case class PaimonSplitScan(
+    table: Table,
+    dataSplits: Array[DataSplit],
+    metadataColumns: Seq[PaimonMetadataColumn] = Seq.empty)
+  extends Scan {
   override def readSchema(): StructType = SparkTypeUtils.fromPaimonRowType(table.rowType())
 
   override def toBatch: Batch = {
-    PaimonBatch(dataSplits.asInstanceOf[Array[Split]], table.newReadBuilder)
+    PaimonBatch(dataSplits.asInstanceOf[Array[Split]], table.newReadBuilder, metadataColumns)
   }
 
 }

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/DeleteFromPaimonTableCommand.scala

@@ -20,13 +20,12 @@ package org.apache.paimon.spark.commands
 
 import org.apache.paimon.CoreOptions
 import org.apache.paimon.CoreOptions.MergeEngine
-import org.apache.paimon.spark.{InsertInto, SparkTable}
 import org.apache.paimon.spark.PaimonSplitScan
 import org.apache.paimon.spark.catalyst.Compatibility
 import org.apache.paimon.spark.catalyst.analysis.expressions.ExpressionHelper
 import org.apache.paimon.spark.leafnode.PaimonLeafRunnableCommand
 import org.apache.paimon.spark.schema.SparkSystemColumns.ROW_KIND_COL
-import org.apache.paimon.table.FileStoreTable
+import org.apache.paimon.table.{BucketMode, FileStoreTable}
 import org.apache.paimon.table.sink.{BatchWriteBuilder, CommitMessage}
 import org.apache.paimon.types.RowKind
 import org.apache.paimon.utils.RowDataPartitionComputer
@@ -56,7 +55,7 @@ case class DeleteFromPaimonTableCommand(
 
   override def run(sparkSession: SparkSession): Seq[Row] = {
 
-    val commit = table.store.newCommit(UUID.randomUUID.toString)
+    val commit = fileStore.newCommit(UUID.randomUUID.toString)
     if (condition == null || condition == TrueLiteral) {
       commit.truncateTable(BatchWriteBuilder.COMMIT_IDENTIFIER)
     } else {
@@ -100,7 +99,7 @@ case class DeleteFromPaimonTableCommand(
         val commitMessages = if (usePrimaryKeyDelete()) {
           performPrimaryKeyDelete(sparkSession)
         } else {
-          performDeleteCopyOnWrite(sparkSession)
+          performNonPrimaryKeyDelete(sparkSession)
         }
         writer.commit(commitMessages)
       }
@@ -119,39 +118,70 @@ case class DeleteFromPaimonTableCommand(
     writer.write(df)
   }
 
-  def performDeleteCopyOnWrite(sparkSession: SparkSession): Seq[CommitMessage] = {
+  def performNonPrimaryKeyDelete(sparkSession: SparkSession): Seq[CommitMessage] = {
+    val pathFactory = fileStore.pathFactory()
     // Step1: the candidate data splits which are filtered by Paimon Predicate.
     val candidateDataSplits = findCandidateDataSplits(condition, relation.output)
-    val fileNameToMeta = candidateFileMap(candidateDataSplits)
+    val dataFilePathToMeta = candidateFileMap(candidateDataSplits)
 
-    // Step2: extract out the exactly files, which must have at least one record to be updated.
-    val touchedFilePaths = findTouchedFiles(candidateDataSplits, condition, relation, sparkSession)
+    if (deletionVectorsEnabled) {
+      // Step2: collect all the deletion vectors that marks the deleted rows.
+      val deletionVectors = collectDeletionVectors(
+        candidateDataSplits,
+        dataFilePathToMeta,
+        condition,
+        relation,
+        sparkSession)
+
+      deletionVectors.cache()
+      try {
+        // Step3: write these deletion vectors.
+        val newIndexCommitMsg = writer.persistDeletionVectors(deletionVectors)
+
+        // Step4: mark the touched index files as DELETE if needed.
+        val rewriteIndexCommitMsg = fileStore.bucketMode() match {
+          case BucketMode.BUCKET_UNAWARE =>
+            val indexEntries = getDeletedIndexFiles(dataFilePathToMeta, deletionVectors)
+            writer.buildCommitMessageFromIndexManifestEntry(indexEntries)
+          case _ =>
+            Seq.empty[CommitMessage]
+        }
 
-    // Step3: the smallest range of data files that need to be rewritten.
-    val touchedFiles = touchedFilePaths.map {
-      file => fileNameToMeta.getOrElse(file, throw new RuntimeException(s"Missing file: $file"))
-    }
+        newIndexCommitMsg ++ rewriteIndexCommitMsg
+      } finally {
+        deletionVectors.unpersist()
+      }
 
-    // Step4: build a dataframe that contains the unchanged data, and write out them.
-    val touchedDataSplits = SparkDataFileMeta.convertToDataSplits(
-      touchedFiles,
-      rawConvertible = true,
-      table.store().pathFactory())
-    val toRewriteScanRelation = Filter(
-      Not(condition),
-      Compatibility.createDataSourceV2ScanRelation(
-        relation,
-        PaimonSplitScan(table, touchedDataSplits),
-        relation.output))
-    val data = createDataset(sparkSession, toRewriteScanRelation)
+    } else {
+      // Step2: extract out the exactly files, which must have at least one record to be updated.
+      val touchedFilePaths =
+        findTouchedFiles(candidateDataSplits, condition, relation, sparkSession)
+
+      // Step3: the smallest range of data files that need to be rewritten.
+      val touchedFiles = touchedFilePaths.map {
+        file =>
+          dataFilePathToMeta.getOrElse(file, throw new RuntimeException(s"Missing file: $file"))
+      }
 
-    // only write new files, should have no compaction
-    val addCommitMessage = writer.writeOnly().write(data)
+      // Step4: build a dataframe that contains the unchanged data, and write out them.
+      val touchedDataSplits =
+        SparkDataFileMeta.convertToDataSplits(touchedFiles, rawConvertible = true, pathFactory)
+      val toRewriteScanRelation = Filter(
+        Not(condition),
+        Compatibility.createDataSourceV2ScanRelation(
+          relation,
+          PaimonSplitScan(table, touchedDataSplits),
+          relation.output))
+      val data = createDataset(sparkSession, toRewriteScanRelation)
 
-    // Step5: convert the deleted files that need to be wrote to commit message.
-    val deletedCommitMessage = buildDeletedCommitMessage(touchedFiles)
+      // only write new files, should have no compaction
+      val addCommitMessage = writer.writeOnly().write(data)
 
-    addCommitMessage ++ deletedCommitMessage
+      // Step5: convert the deleted files that need to be wrote to commit message.
+      val deletedCommitMessage = buildDeletedCommitMessage(touchedFiles)
+
+      addCommitMessage ++ deletedCommitMessage
+    }
   }
 
 }