[core][flink][spark] Support PK-vector indexes for postpone buckets

docs/docs/primary-key-table/vector-index.md

@@ -39,7 +39,8 @@ is built separately from writes, see
 
 A table with a primary-key vector index must satisfy all of the following:
 
-- It is a primary-key table in fixed-bucket mode (`bucket > 0`).
+- It is a primary-key table in fixed-bucket mode (`bucket > 0`) or postpone-bucket mode
+  (`bucket = -2`).
 - `deletion-vectors.enabled` is `true`, except for `first-row`, where it must be `false`.
 - Its merge engine is `deduplicate`, `partial-update`, `aggregation`, or `first-row`.
 - The indexed column is a `VECTOR` whose element type is `FLOAT`.
@@ -115,6 +116,12 @@ The index segment records the source data files and maps ANN ordinals back to th
 positions. Compact-output data-file and index-file changes are committed atomically, so a reader
 never observes an index from a different compact-output snapshot.
 
+For a postpone-bucket table, foreground writes remain write-only. Fixed-bucket batch writes produce
+Level-0 files in real buckets, while postponed writes produce files in bucket `-2`. These rows do not
+become visible to normal reads or vector search until a batch compact runs. The background compact
+processes both kinds of pending files, builds their ANN indexes, and publishes the data and index
+changes in one atomic commit.
+
 ANN compaction is configured independently from data compaction. It does not inherit
 `vector.target-file-size`, `num-sorted-run.compaction-trigger`, or
 `compaction.delete-ratio-threshold`.

paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java

@@ -27,6 +27,8 @@
 import org.apache.paimon.manifest.ManifestEntry;
 import org.apache.paimon.utils.SnapshotManager;
 
+import javax.annotation.Nullable;
+
 import java.util.ArrayList;
 import java.util.List;
 
@@ -38,15 +40,35 @@ public class FileSystemWriteRestore implements WriteRestore {
     private final SnapshotManager snapshotManager;
     private final FileStoreScan scan;
     private final IndexFileHandler indexFileHandler;
+    private final @Nullable Long snapshotId;
 
     public FileSystemWriteRestore(
             CoreOptions options,
             SnapshotManager snapshotManager,
             FileStoreScan scan,
             IndexFileHandler indexFileHandler) {
+        this(options, snapshotManager, scan, indexFileHandler, null);
+    }
+
+    public FileSystemWriteRestore(
+            CoreOptions options,
+            SnapshotManager snapshotManager,
+            FileStoreScan scan,
+            IndexFileHandler indexFileHandler,
+            long snapshotId) {
+        this(options, snapshotManager, scan, indexFileHandler, Long.valueOf(snapshotId));
+    }
+
+    private FileSystemWriteRestore(
+            CoreOptions options,
+            SnapshotManager snapshotManager,
+            FileStoreScan scan,
+            IndexFileHandler indexFileHandler,
+            @Nullable Long snapshotId) {
         this.snapshotManager = snapshotManager;
         this.scan = scan;
         this.indexFileHandler = indexFileHandler;
+        this.snapshotId = snapshotId;
         if (options.manifestDeleteFileDropStats()) {
             if (this.scan != null) {
                 this.scan.dropStats();
@@ -71,7 +93,10 @@ public RestoreFiles restoreFiles(
             boolean scanVectorIndexPayloads) {
         // NOTE: don't use snapshotManager.latestSnapshot() here,
         // because we don't want to flood the catalog with high concurrency
-        Snapshot snapshot = snapshotManager.latestSnapshotFromFileSystem();
+        Snapshot snapshot =
+                snapshotId == null
+                        ? snapshotManager.latestSnapshotFromFileSystem()
+                        : snapshotManager.snapshot(snapshotId);
         if (snapshot == null) {
             return RestoreFiles.empty();
         }

paimon-core/src/main/java/org/apache/paimon/postpone/BucketFiles.java

@@ -89,7 +89,9 @@ public void update(CommitMessageImpl message) {
         changelogFiles.addAll(message.newFilesIncrement().changelogFiles());
         changelogFiles.addAll(message.compactIncrement().changelogFiles());
 
+        newIndexFiles.addAll(message.newFilesIncrement().newIndexFiles());
         newIndexFiles.addAll(message.compactIncrement().newIndexFiles());
+        deletedIndexFiles.addAll(message.newFilesIncrement().deletedIndexFiles());
         deletedIndexFiles.addAll(message.compactIncrement().deletedIndexFiles());
 
         toDelete.forEach((fileName, path) -> fileIO.deleteQuietly(path));

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

@@ -929,8 +929,9 @@ private static void validatePrimaryKeyVectorIndex(TableSchema schema, CoreOption
                 "Primary-key vector index with merge-engine = %s requires deletion-vectors.merge-on-read = false.",
                 options.mergeEngine());
         checkArgument(
-                options.bucket() > 0,
-                "Primary-key vector index requires fixed bucket mode (bucket > 0), but bucket is %s.",
+                options.bucket() > 0 || options.bucket() == BucketMode.POSTPONE_BUCKET,
+                "Primary-key vector index requires fixed or postpone bucket mode "
+                        + "(bucket > 0 or bucket = -2), but bucket is %s.",
                 options.bucket());
         checkArgument(
                 !options.pkClusteringOverride(),

paimon-core/src/main/java/org/apache/paimon/table/PostponeUtils.java

@@ -24,14 +24,20 @@
 
 import javax.annotation.Nullable;
 
+import java.io.Serializable;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 
 import static org.apache.paimon.CoreOptions.BUCKET;
+import static org.apache.paimon.CoreOptions.COMMIT_STRICT_MODE_LAST_SAFE_SNAPSHOT;
 import static org.apache.paimon.CoreOptions.WRITE_ONLY;
 
 /** Utils for postpone table. */
@@ -88,9 +94,23 @@ public static int determineBucketNum(
     }
 
     public static Map<BinaryRow, Integer> getKnownNumBuckets(FileStoreTable table) {
+        return getKnownNumBuckets(
+                table.store().newScan().onlyReadRealBuckets().readSimpleEntries());
+    }
+
+    public static Map<BinaryRow, Integer> getKnownNumBuckets(
+            FileStoreTable table, long snapshotId) {
+        return getKnownNumBuckets(
+                table.store()
+                        .newScan()
+                        .withSnapshot(snapshotId)
+                        .onlyReadRealBuckets()
+                        .readSimpleEntries());
+    }
+
+    private static Map<BinaryRow, Integer> getKnownNumBuckets(
+            List<SimpleFileEntry> simpleFileEntries) {
         Map<BinaryRow, Integer> knownNumBuckets = new HashMap<>();
-        List<SimpleFileEntry> simpleFileEntries =
-                table.store().newScan().onlyReadRealBuckets().readSimpleEntries();
         for (SimpleFileEntry entry : simpleFileEntries) {
             if (entry.totalBuckets() >= 0) {
                 Integer oldTotalBuckets =
@@ -109,11 +129,43 @@ public static Map<BinaryRow, Integer> getKnownNumBuckets(FileStoreTable table) {
         return knownNumBuckets;
     }
 
+    /** Returns real buckets containing active Level-0 files in the specified snapshot. */
+    public static List<CompactBucket> getLevel0Buckets(FileStoreTable table, long snapshotId) {
+        List<SimpleFileEntry> entries =
+                table.store()
+                        .newScan()
+                        .withSnapshot(snapshotId)
+                        .onlyReadRealBuckets()
+                        .readSimpleEntries();
+        Set<CompactBucket> buckets = new LinkedHashSet<>();
+        for (SimpleFileEntry entry : entries) {
+            if (entry.bucket() >= 0 && entry.totalBuckets() > 0 && entry.level() == 0) {
+                buckets.add(
+                        new CompactBucket(entry.partition(), entry.bucket(), entry.totalBuckets()));
+            }
+        }
+        return new ArrayList<>(buckets);
+    }
+
     /** Returns row counts of current active files in the postpone bucket. */
     public static Map<BinaryRow, Long> getPostponeRowCounts(FileStoreTable table) {
+        return getPostponeRowCounts(
+                table.newSnapshotReader()
+                        .withBucket(BucketMode.POSTPONE_BUCKET)
+                        .readFileIterator());
+    }
+
+    /** Returns row counts of active postpone files in the specified snapshot. */
+    public static Map<BinaryRow, Long> getPostponeRowCounts(FileStoreTable table, long snapshotId) {
+        return getPostponeRowCounts(
+                table.newSnapshotReader()
+                        .withSnapshot(snapshotId)
+                        .withBucket(BucketMode.POSTPONE_BUCKET)
+                        .readFileIterator());
+    }
+
+    private static Map<BinaryRow, Long> getPostponeRowCounts(Iterator<ManifestEntry> iterator) {
         Map<BinaryRow, Long> rowCounts = new HashMap<>();
-        Iterator<ManifestEntry> iterator =
-                table.newSnapshotReader().withBucket(BucketMode.POSTPONE_BUCKET).readFileIterator();
         while (iterator.hasNext()) {
             ManifestEntry entry = iterator.next();
             rowCounts.merge(entry.partition(), entry.file().rowCount(), Long::sum);
@@ -130,8 +182,64 @@ public static FileStoreTable tableForFixBucketWrite(FileStoreTable table) {
         return table.copy(batchWriteOptions);
     }
 
+    public static FileStoreTable tableForPostponeCompact(
+            FileStoreTable table, int numBuckets, long snapshotId) {
+        Map<String, String> compactOptions = new HashMap<>();
+        compactOptions.put(BUCKET.key(), String.valueOf(numBuckets));
+        compactOptions.put(WRITE_ONLY.key(), "false");
+        compactOptions.put(COMMIT_STRICT_MODE_LAST_SAFE_SNAPSHOT.key(), String.valueOf(snapshotId));
+        return table.copy(compactOptions);
+    }
+
     public static FileStoreTable tableForCommit(FileStoreTable table) {
         return table.copy(
                 Collections.singletonMap(BUCKET.key(), String.valueOf(BucketMode.POSTPONE_BUCKET)));
     }
+
+    /** A real bucket which requires background compaction. */
+    public static final class CompactBucket implements Serializable {
+
+        private static final long serialVersionUID = 1L;
+
+        private final BinaryRow partition;
+        private final int bucket;
+        private final int totalBuckets;
+
+        public CompactBucket(BinaryRow partition, int bucket, int totalBuckets) {
+            this.partition = partition.copy();
+            this.bucket = bucket;
+            this.totalBuckets = totalBuckets;
+        }
+
+        public BinaryRow partition() {
+            return partition;
+        }
+
+        public int bucket() {
+            return bucket;
+        }
+
+        public int totalBuckets() {
+            return totalBuckets;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (!(o instanceof CompactBucket)) {
+                return false;
+            }
+            CompactBucket that = (CompactBucket) o;
+            return bucket == that.bucket
+                    && totalBuckets == that.totalBuckets
+                    && Objects.equals(partition, that.partition);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(partition, bucket, totalBuckets);
+        }
+    }
 }

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

@@ -26,6 +26,7 @@
 import org.apache.paimon.manifest.FileKind;
 import org.apache.paimon.manifest.IndexManifestEntry;
 import org.apache.paimon.partition.PartitionPredicate;
+import org.apache.paimon.table.BucketMode;
 import org.apache.paimon.table.FileStoreTable;
 import org.apache.paimon.table.source.snapshot.SnapshotReader;
 import org.apache.paimon.table.source.snapshot.TimeTravelUtil;
@@ -74,6 +75,9 @@ public Plan scan() {
 
         SnapshotReader snapshotReader =
                 table.newSnapshotReader().withSnapshot(snapshot).withMode(ScanMode.ALL).keepStats();
+        if (table.coreOptions().bucket() == BucketMode.POSTPONE_BUCKET) {
+            snapshotReader.onlyReadRealBuckets();
+        }
         if (partitionFilter != null) {
             snapshotReader.withPartitionFilter(partitionFilter);
         }

paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java

@@ -32,11 +32,50 @@
 import static org.apache.paimon.data.BinaryRow.EMPTY_ROW;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 /** Tests for {@link FileSystemWriteRestore}. */
 class FileSystemWriteRestoreTest {
 
+    @Test
+    void testRestoreFromPinnedSnapshot() {
+        Snapshot pinned = mock(Snapshot.class);
+        Snapshot latest = mock(Snapshot.class);
+        SnapshotManager snapshotManager = mock(SnapshotManager.class);
+        when(snapshotManager.snapshot(5L)).thenReturn(pinned);
+        when(snapshotManager.latestSnapshotFromFileSystem()).thenReturn(latest);
+
+        FileStoreScan scan = mock(FileStoreScan.class);
+        FileStoreScan.Plan plan = mock(FileStoreScan.Plan.class);
+        when(scan.withSnapshot(pinned)).thenReturn(scan);
+        when(scan.withPartitionBucket(EMPTY_ROW, 0)).thenReturn(scan);
+        when(scan.plan()).thenReturn(plan);
+        when(plan.files()).thenReturn(Collections.emptyList());
+
+        IndexFileMeta ann = new IndexFileMeta("test-vector-ann", "ann", 1, 1, null, null, null);
+        IndexFileHandler indexFileHandler = mock(IndexFileHandler.class);
+        when(indexFileHandler.scanSourceIndexes(pinned, EMPTY_ROW, 0))
+                .thenReturn(Collections.singletonList(ann));
+
+        FileSystemWriteRestore restore =
+                new FileSystemWriteRestore(
+                        new CoreOptions(new HashMap<>()),
+                        snapshotManager,
+                        scan,
+                        indexFileHandler,
+                        5L);
+
+        RestoreFiles restored = restore.restoreFiles(EMPTY_ROW, 0, false, false, true);
+
+        assertThat(restored.snapshot()).isSameAs(pinned);
+        assertThat(restored.vectorIndexPayloads()).containsExactly(ann);
+        verify(scan).withSnapshot(pinned);
+        verify(indexFileHandler).scanSourceIndexes(pinned, EMPTY_ROW, 0);
+        verify(snapshotManager, never()).latestSnapshotFromFileSystem();
+    }
+
     @Test
     void testRestoreVectorIndexPayloadsWithoutDirectory() {
         Snapshot snapshot = mock(Snapshot.class);