This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 26f415a1a2 [core][flink][spark] Support PK-vector indexes for postpone
buckets (#8594)
26f415a1a2 is described below
commit 26f415a1a20c63ceac1cf2d9e829d5516273d81a
Author: Jingsong Lee <[email protected]>
AuthorDate: Mon Jul 13 17:48:38 2026 +0800
[core][flink][spark] Support PK-vector indexes for postpone buckets (#8594)
- Allow primary-key vector indexes in postpone-bucket mode (`bucket =
-2`) while continuing to reject dynamic buckets.
- Capture one snapshot for postpone planning and writer restore.
- Make Flink and Spark background compaction process both bucket `-2`
rows and real-bucket Level-0 files.
- Build and publish ANN index changes together with compacted data in
one strict atomic commit.
- Exclude postponed data from vector search until compaction and
document the visibility behavior.
---
docs/docs/primary-key-table/vector-index.md | 9 +-
.../paimon/operation/FileSystemWriteRestore.java | 27 ++-
.../org/apache/paimon/postpone/BucketFiles.java | 2 +
.../org/apache/paimon/schema/SchemaValidation.java | 5 +-
.../org/apache/paimon/table/PostponeUtils.java | 116 ++++++++++-
.../paimon/table/source/PrimaryKeyVectorScan.java | 4 +
.../operation/FileSystemWriteRestoreTest.java | 39 ++++
.../apache/paimon/postpone/BucketFilesTest.java | 72 +++++++
.../PrimaryKeyVectorIndexValidationTest.java | 12 +-
.../org/apache/paimon/table/PostponeUtilsTest.java | 106 ++++++++++
.../table/source/PrimaryKeyVectorScanTest.java | 34 +++
.../apache/paimon/flink/action/CompactAction.java | 116 ++++++++---
.../postpone/PostponeBucketCompactOperator.java | 134 ++++++++++++
.../postpone/PostponeBucketCompactSplitSource.java | 44 +++-
.../PostponeBucketCompactSplitSourceTest.java | 57 +++++
.../procedure/VectorSearchProcedureITCase.java | 86 ++++++++
.../procedure/SparkPostponeCompactProcedure.scala | 231 ++++++++++++++-------
.../spark/sql/PrimaryKeyVectorSearchTest.scala | 69 ++++++
18 files changed, 1044 insertions(+), 119 deletions(-)
diff --git a/docs/docs/primary-key-table/vector-index.md
b/docs/docs/primary-key-table/vector-index.md
index e6adaa55f2..899e6f0d32 100644
--- a/docs/docs/primary-key-table/vector-index.md
+++ b/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`.
diff --git
a/paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java
b/paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java
index 46e174cc21..3a2017ef56 100644
---
a/paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java
+++
b/paimon-core/src/main/java/org/apache/paimon/operation/FileSystemWriteRestore.java
@@ -27,6 +27,8 @@ import org.apache.paimon.io.DataFileMeta;
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 class FileSystemWriteRestore implements WriteRestore {
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();
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/postpone/BucketFiles.java
b/paimon-core/src/main/java/org/apache/paimon/postpone/BucketFiles.java
index 51212f708e..d019de4421 100644
--- a/paimon-core/src/main/java/org/apache/paimon/postpone/BucketFiles.java
+++ b/paimon-core/src/main/java/org/apache/paimon/postpone/BucketFiles.java
@@ -89,7 +89,9 @@ public class BucketFiles {
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));
diff --git
a/paimon-core/src/main/java/org/apache/paimon/schema/SchemaValidation.java
b/paimon-core/src/main/java/org/apache/paimon/schema/SchemaValidation.java
index 4f68c0b996..746338314d 100644
--- a/paimon-core/src/main/java/org/apache/paimon/schema/SchemaValidation.java
+++ b/paimon-core/src/main/java/org/apache/paimon/schema/SchemaValidation.java
@@ -929,8 +929,9 @@ public class SchemaValidation {
"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(),
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/PostponeUtils.java
b/paimon-core/src/main/java/org/apache/paimon/table/PostponeUtils.java
index 62fb41a9e9..8a600eee9e 100644
--- a/paimon-core/src/main/java/org/apache/paimon/table/PostponeUtils.java
+++ b/paimon-core/src/main/java/org/apache/paimon/table/PostponeUtils.java
@@ -24,14 +24,20 @@ import org.apache.paimon.manifest.SimpleFileEntry;
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 class PostponeUtils {
}
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 class PostponeUtils {
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 class PostponeUtils {
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);
+ }
+ }
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java
index bb88799617..221d4afaaf 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java
@@ -26,6 +26,7 @@ import org.apache.paimon.io.DataFileMeta;
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 class PrimaryKeyVectorScan implements VectorScan {
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);
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java
b/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java
index 401439be91..830fa5ded0 100644
---
a/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/operation/FileSystemWriteRestoreTest.java
@@ -32,11 +32,50 @@ import java.util.HashMap;
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);
diff --git
a/paimon-core/src/test/java/org/apache/paimon/postpone/BucketFilesTest.java
b/paimon-core/src/test/java/org/apache/paimon/postpone/BucketFilesTest.java
new file mode 100644
index 0000000000..6c60075850
--- /dev/null
+++ b/paimon-core/src/test/java/org/apache/paimon/postpone/BucketFilesTest.java
@@ -0,0 +1,72 @@
+/*
+ * 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.paimon.postpone;
+
+import org.apache.paimon.fs.FileIO;
+import org.apache.paimon.index.IndexFileMeta;
+import org.apache.paimon.io.CompactIncrement;
+import org.apache.paimon.io.DataFilePathFactory;
+import org.apache.paimon.io.DataIncrement;
+import org.apache.paimon.table.sink.CommitMessageImpl;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+
+import static org.apache.paimon.data.BinaryRow.EMPTY_ROW;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+
+/** Tests for {@link BucketFiles}. */
+class BucketFilesTest {
+
+ @Test
+ void testPreservesIndexChangesFromDataAndCompactIncrements() {
+ IndexFileMeta dataAdd = index("data-add");
+ IndexFileMeta dataDelete = index("data-delete");
+ IndexFileMeta compactAdd = index("compact-add");
+ IndexFileMeta compactDelete = index("compact-delete");
+ DataIncrement dataIncrement =
+ new DataIncrement(
+ Collections.emptyList(),
+ Collections.emptyList(),
+ Collections.emptyList(),
+ Collections.singletonList(dataAdd),
+ Collections.singletonList(dataDelete));
+ CompactIncrement compactIncrement =
+ new CompactIncrement(
+ Collections.emptyList(),
+ Collections.emptyList(),
+ Collections.emptyList(),
+ Collections.singletonList(compactAdd),
+ Collections.singletonList(compactDelete));
+
+ BucketFiles bucketFiles =
+ new BucketFiles(mock(DataFilePathFactory.class),
mock(FileIO.class));
+ bucketFiles.update(new CommitMessageImpl(EMPTY_ROW, 0, 1,
dataIncrement, compactIncrement));
+
+ CompactIncrement result = bucketFiles.makeMessage(EMPTY_ROW,
0).compactIncrement();
+ assertThat(result.newIndexFiles()).containsExactly(dataAdd,
compactAdd);
+ assertThat(result.deletedIndexFiles()).containsExactly(dataDelete,
compactDelete);
+ }
+
+ private static IndexFileMeta index(String name) {
+ return new IndexFileMeta(name, "ann", 1, 1, null, null, null);
+ }
+}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/schema/PrimaryKeyVectorIndexValidationTest.java
b/paimon-core/src/test/java/org/apache/paimon/schema/PrimaryKeyVectorIndexValidationTest.java
index 8341d6bc4e..af6155411b 100644
---
a/paimon-core/src/test/java/org/apache/paimon/schema/PrimaryKeyVectorIndexValidationTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/schema/PrimaryKeyVectorIndexValidationTest.java
@@ -171,12 +171,20 @@ class PrimaryKeyVectorIndexValidationTest {
}
@Test
- void testRequiresFixedBucket() {
+ void testSupportsPostponeBucket() {
+ Map<String, String> options = enabledOptions();
+ options.put(CoreOptions.BUCKET.key(), "-2");
+
+ assertThatCode(() ->
validateTableSchema(schema(options))).doesNotThrowAnyException();
+ }
+
+ @Test
+ void testRejectsDynamicBucket() {
Map<String, String> options = enabledOptions();
options.put(CoreOptions.BUCKET.key(), "-1");
assertThatThrownBy(() -> validateTableSchema(schema(options)))
- .hasMessageContaining("requires fixed bucket mode");
+ .hasMessageContaining("requires fixed or postpone bucket
mode");
}
@Test
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/PostponeUtilsTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/PostponeUtilsTest.java
index 50d9bc2077..251f06c5f6 100644
--- a/paimon-core/src/test/java/org/apache/paimon/table/PostponeUtilsTest.java
+++ b/paimon-core/src/test/java/org/apache/paimon/table/PostponeUtilsTest.java
@@ -18,20 +18,116 @@
package org.apache.paimon.table;
+import org.apache.paimon.FileStore;
import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.data.BinaryRowWriter;
+import org.apache.paimon.io.DataFileMeta;
+import org.apache.paimon.manifest.ManifestEntry;
+import org.apache.paimon.manifest.SimpleFileEntry;
+import org.apache.paimon.operation.FileStoreScan;
+import org.apache.paimon.table.source.snapshot.SnapshotReader;
import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import java.util.Arrays;
+import java.util.Collections;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Answers.RETURNS_SELF;
+import static org.mockito.ArgumentMatchers.anyMap;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
/** Tests for {@link PostponeUtils}. */
public class PostponeUtilsTest {
+ @Test
+ public void testGetKnownNumBucketsFromSnapshot() {
+ BinaryRow partition = partition(1);
+ SimpleFileEntry entry = mock(SimpleFileEntry.class);
+ when(entry.partition()).thenReturn(partition);
+ when(entry.totalBuckets()).thenReturn(4);
+
+ FileStoreScan scan = mock(FileStoreScan.class, RETURNS_SELF);
+
when(scan.readSimpleEntries()).thenReturn(Collections.singletonList(entry));
+ FileStore store = mock(FileStore.class);
+ when(store.newScan()).thenReturn(scan);
+ FileStoreTable table = mock(FileStoreTable.class);
+ when(table.store()).thenReturn(store);
+
+ assertThat(PostponeUtils.getKnownNumBuckets(table,
5L)).containsEntry(partition, 4);
+ verify(scan).withSnapshot(5L);
+ verify(scan).onlyReadRealBuckets();
+ }
+
+ @Test
+ public void testGetPostponeRowCountsFromSnapshot() {
+ BinaryRow partition = partition(1);
+ DataFileMeta file = mock(DataFileMeta.class);
+ when(file.rowCount()).thenReturn(10L);
+ ManifestEntry entry = mock(ManifestEntry.class);
+ when(entry.partition()).thenReturn(partition);
+ when(entry.file()).thenReturn(file);
+
+ SnapshotReader reader = mock(SnapshotReader.class, RETURNS_SELF);
+
when(reader.readFileIterator()).thenReturn(Collections.singletonList(entry).iterator());
+ FileStoreTable table = mock(FileStoreTable.class);
+ when(table.newSnapshotReader()).thenReturn(reader);
+
+ assertThat(PostponeUtils.getPostponeRowCounts(table,
5L)).containsEntry(partition, 10L);
+ verify(reader).withSnapshot(5L);
+ verify(reader).withBucket(BucketMode.POSTPONE_BUCKET);
+ }
+
+ @Test
+ public void testGetLevel0BucketsFromSnapshot() {
+ BinaryRow partition = partition(1);
+ SimpleFileEntry level0 = fileEntry(partition, 0, 2, 0);
+ SimpleFileEntry duplicate = fileEntry(partition, 0, 2, 0);
+ SimpleFileEntry compacted = fileEntry(partition, 1, 2, 1);
+ SimpleFileEntry postpone = fileEntry(partition, -2, -2, 0);
+
+ FileStoreScan scan = mock(FileStoreScan.class, RETURNS_SELF);
+ when(scan.readSimpleEntries())
+ .thenReturn(Arrays.asList(level0, duplicate, compacted,
postpone));
+ FileStore store = mock(FileStore.class);
+ when(store.newScan()).thenReturn(scan);
+ FileStoreTable table = mock(FileStoreTable.class);
+ when(table.store()).thenReturn(store);
+
+ List<PostponeUtils.CompactBucket> buckets =
PostponeUtils.getLevel0Buckets(table, 5L);
+
+ assertThat(buckets).hasSize(1);
+ assertThat(buckets.get(0).partition()).isEqualTo(partition);
+ assertThat(buckets.get(0).bucket()).isEqualTo(0);
+ assertThat(buckets.get(0).totalBuckets()).isEqualTo(2);
+ verify(scan).withSnapshot(5L);
+ verify(scan).onlyReadRealBuckets();
+ }
+
+ @Test
+ public void testTableForPostponeCompact() {
+ FileStoreTable table = mock(FileStoreTable.class);
+ FileStoreTable copied = mock(FileStoreTable.class);
+ when(table.copy(anyMap())).thenReturn(copied);
+
+ assertThat(PostponeUtils.tableForPostponeCompact(table, 4,
5L)).isSameAs(copied);
+
+ @SuppressWarnings("unchecked")
+ ArgumentCaptor<Map<String, String>> options =
ArgumentCaptor.forClass(Map.class);
+ verify(table).copy(options.capture());
+ assertThat(options.getValue())
+ .containsEntry("bucket", "4")
+ .containsEntry("write-only", "false")
+ .containsEntry("commit.strict-mode.last-safe-snapshot", "5");
+ }
+
@Test
public void testComputeBucketNumByRowCount() {
assertThat(PostponeUtils.computeBucketNumByRowCount(0,
100)).isEqualTo(1);
@@ -96,4 +192,14 @@ public class PostponeUtilsTest {
writer.complete();
return row;
}
+
+ private static SimpleFileEntry fileEntry(
+ BinaryRow partition, int bucket, int totalBuckets, int level) {
+ SimpleFileEntry entry = mock(SimpleFileEntry.class);
+ when(entry.partition()).thenReturn(partition);
+ when(entry.bucket()).thenReturn(bucket);
+ when(entry.totalBuckets()).thenReturn(totalBuckets);
+ when(entry.level()).thenReturn(level);
+ return entry;
+ }
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java
index d3db372da3..25b823139c 100644
---
a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java
@@ -33,6 +33,7 @@ import org.apache.paimon.manifest.FileSource;
import org.apache.paimon.manifest.IndexManifestEntry;
import org.apache.paimon.options.Options;
import org.apache.paimon.stats.SimpleStats;
+import org.apache.paimon.table.BucketMode;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.source.snapshot.SnapshotReader;
import org.apache.paimon.utils.Filter;
@@ -55,11 +56,44 @@ import static org.mockito.Answers.RETURNS_SELF;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
/** Tests snapshot-consistent planning for bucket-local primary-key vector
search. */
class PrimaryKeyVectorScanTest {
+ @Test
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ void testPostponeTableOnlyScansRealBuckets() {
+ Options options = new Options();
+ options.set(CoreOptions.BUCKET, BucketMode.POSTPONE_BUCKET);
+ options.set(CoreOptions.PK_VECTOR_INDEX_COLUMNS, "embedding");
+ options.setString("fields.embedding.pk-vector.index.type", "ivf-pq");
+
+ FileStoreTable table = mock(FileStoreTable.class);
+ Snapshot snapshot = mock(Snapshot.class);
+ when(snapshot.id()).thenReturn(11L);
+ when(table.coreOptions()).thenReturn(new CoreOptions(options));
+ when(table.latestSnapshot()).thenReturn(Optional.of(snapshot));
+
+ SnapshotReader reader = mock(SnapshotReader.class, RETURNS_SELF);
+ SnapshotReader.Plan snapshotPlan = mock(SnapshotReader.Plan.class,
CALLS_REAL_METHODS);
+ when(snapshotPlan.splits()).thenReturn(Collections.emptyList());
+ when(reader.read()).thenReturn(snapshotPlan);
+ when(table.newSnapshotReader()).thenReturn(reader);
+
+ IndexFileHandler indexFileHandler = mock(IndexFileHandler.class);
+ when(indexFileHandler.scan(eq(snapshot), any(Filter.class)))
+ .thenReturn(Collections.emptyList());
+ FileStore store = mock(FileStore.class);
+ when(store.newIndexFileHandler()).thenReturn(indexFileHandler);
+ when(table.store()).thenReturn(store);
+
+ new PrimaryKeyVectorScan(table, 7, "ivf-pq", null).scan();
+
+ verify(reader).onlyReadRealBuckets();
+ }
+
@Test
@SuppressWarnings({"unchecked", "rawtypes"})
void testScansOneSnapshotAndFiltersVectorIdentity() {
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactAction.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactAction.java
index cbdcf825f2..146c7b7a62 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactAction.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactAction.java
@@ -19,12 +19,14 @@
package org.apache.paimon.flink.action;
import org.apache.paimon.CoreOptions;
+import org.apache.paimon.Snapshot;
import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.flink.FlinkConnectorOptions;
import org.apache.paimon.flink.compact.AppendTableCompact;
import org.apache.paimon.flink.compact.DataEvolutionTableCompact;
import org.apache.paimon.flink.compact.IncrementalClusterCompact;
+import org.apache.paimon.flink.postpone.PostponeBucketCompactOperator;
import org.apache.paimon.flink.postpone.PostponeBucketCompactSplitSource;
import
org.apache.paimon.flink.postpone.RewritePostponeBucketCommittableOperator;
import org.apache.paimon.flink.predicate.SimpleSqlPredicateConvertor;
@@ -36,6 +38,7 @@ import org.apache.paimon.flink.sink.FlinkSinkBuilder;
import org.apache.paimon.flink.sink.FlinkStreamPartitioner;
import org.apache.paimon.flink.sink.RowDataChannelComputer;
import org.apache.paimon.flink.source.CompactorSourceBuilder;
+import org.apache.paimon.flink.utils.JavaTypeInfo;
import org.apache.paimon.options.Options;
import org.apache.paimon.partition.PartitionPredicate;
import org.apache.paimon.predicate.PartitionPredicateVisitor;
@@ -45,6 +48,8 @@ import
org.apache.paimon.predicate.PredicateProjectionConverter;
import org.apache.paimon.table.BucketMode;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.PostponeUtils;
+import org.apache.paimon.table.PostponeUtils.CompactBucket;
+import org.apache.paimon.table.sink.ChannelComputer;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.InternalRowPartitionComputer;
import org.apache.paimon.utils.Pair;
@@ -67,9 +72,11 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
+import java.util.Set;
import static
org.apache.paimon.partition.PartitionPredicate.createBinaryPartitions;
import static
org.apache.paimon.partition.PartitionPredicate.createPartitionPredicate;
@@ -293,42 +300,45 @@ public class CompactAction extends TableActionBase {
int defaultBucketNum =
options.get(CoreOptions.POSTPONE_DEFAULT_BUCKET_NUM);
Optional<Long> targetRowNumPerBucket =
options.getOptional(CoreOptions.POSTPONE_TARGET_ROW_NUM_PER_BUCKET);
- Map<BinaryRow, Integer> knownNumBuckets =
PostponeUtils.getKnownNumBuckets(table);
+ Optional<Snapshot> optionalSnapshot = table.latestSnapshot();
+ if (!optionalSnapshot.isPresent()) {
+ return buildNothingToCompact(env);
+ }
+ long snapshotId = optionalSnapshot.get().id();
+ Map<BinaryRow, Integer> knownNumBuckets =
+ PostponeUtils.getKnownNumBuckets(table, snapshotId);
Map<BinaryRow, Long> postponeRowCounts =
targetRowNumPerBucket.isPresent()
- ? PostponeUtils.getPostponeRowCounts(table)
+ ? PostponeUtils.getPostponeRowCounts(table, snapshotId)
: Collections.emptyMap();
- // change bucket to a positive value, so we can scan files from the
bucket = -2 directory
- Map<String, String> bucketOptions = new HashMap<>(table.options());
- bucketOptions.put(CoreOptions.BUCKET.key(),
String.valueOf(defaultBucketNum));
- FileStoreTable fileStoreTable =
table.copy(table.schema().copy(bucketOptions));
-
- List<BinaryRow> partitions =
- fileStoreTable
- .newSnapshotReader()
+ List<BinaryRow> postponePartitions =
+ table.newSnapshotReader()
+ .withSnapshot(snapshotId)
.withBucket(BucketMode.POSTPONE_BUCKET)
.partitions();
- if (partitions.isEmpty()) {
- if (this.forceStartFlinkJob) {
- env.fromSequence(0, 0)
- .name("Nothing to Compact Source")
- .sinkTo(new DiscardingSink<>());
- return true;
- } else {
- return false;
- }
+
+ Map<BinaryRow, List<CompactBucket>> compactBucketsByPartition = new
LinkedHashMap<>();
+ for (CompactBucket bucket : PostponeUtils.getLevel0Buckets(table,
snapshotId)) {
+ compactBucketsByPartition
+ .computeIfAbsent(bucket.partition(), ignored -> new
ArrayList<>())
+ .add(bucket);
+ }
+ Set<BinaryRow> affectedPartitions = new
LinkedHashSet<>(postponePartitions);
+ affectedPartitions.addAll(compactBucketsByPartition.keySet());
+ if (affectedPartitions.isEmpty()) {
+ return buildNothingToCompact(env);
}
InternalRowPartitionComputer partitionComputer =
new InternalRowPartitionComputer(
- fileStoreTable.coreOptions().partitionDefaultName(),
- fileStoreTable.store().partitionType(),
- fileStoreTable.partitionKeys().toArray(new String[0]),
- fileStoreTable.coreOptions().legacyPartitionName());
+ table.coreOptions().partitionDefaultName(),
+ table.store().partitionType(),
+ table.partitionKeys().toArray(new String[0]),
+ table.coreOptions().legacyPartitionName());
String commitUser = CoreOptions.createCommitUser(options);
List<DataStream<Committable>> dataStreams = new ArrayList<>();
- for (BinaryRow partition : partitions) {
+ for (BinaryRow partition : affectedPartitions) {
int bucketNum =
PostponeUtils.determineBucketNum(
partition,
@@ -336,10 +346,8 @@ public class CompactAction extends TableActionBase {
targetRowNumPerBucket,
postponeRowCounts,
defaultBucketNum);
-
- bucketOptions = new HashMap<>(table.options());
- bucketOptions.put(CoreOptions.BUCKET.key(),
String.valueOf(bucketNum));
- FileStoreTable realTable =
table.copy(table.schema().copy(bucketOptions));
+ FileStoreTable realTable =
+ PostponeUtils.tableForPostponeCompact(table, bucketNum,
snapshotId);
LinkedHashMap<String, String> partitionSpec =
partitionComputer.generatePartValues(partition);
@@ -348,6 +356,7 @@ public class CompactAction extends TableActionBase {
env,
realTable,
partitionSpec,
+ snapshotId,
options.get(FlinkConnectorOptions.SCAN_PARALLELISM));
DataStream<InternalRow> partitioned =
@@ -358,9 +367,29 @@ public class CompactAction extends TableActionBase {
table.catalogEnvironment().catalogContext()),
new RowDataChannelComputer(realTable.schema()),
null);
- FixedBucketSink sink = new FixedBucketSink(realTable, null);
+ List<CompactBucket> compactBuckets =
+ compactBucketsByPartition.getOrDefault(partition,
Collections.emptyList());
+ DataStream<CompactBucket> partitionedCompactBuckets =
+ FlinkStreamPartitioner.partition(
+ env.fromCollection(
+ compactBuckets, new
JavaTypeInfo<>(CompactBucket.class))
+ .name(
+ String.format(
+ "Level-0 compact buckets:
%s - %s",
+ table.fullName(),
partitionSpec))
+ .forceNonParallel(),
+ new CompactBucketChannelComputer(),
+ partitioned.getParallelism());
DataStream<Committable> written =
- sink.doWrite(partitioned, commitUser,
partitioned.getParallelism())
+ partitioned
+ .connect(partitionedCompactBuckets)
+ .transform(
+ String.format(
+ "Write and compact postpone
buckets: %s - %s",
+ table.fullName(), partitionSpec),
+ new CommittableTypeInfo(),
+ new PostponeBucketCompactOperator(
+ realTable, commitUser, snapshotId))
.forward()
.transform(
"Rewrite compact committable",
@@ -370,6 +399,8 @@ public class CompactAction extends TableActionBase {
dataStreams.add(sourcePair.getRight());
}
+ FileStoreTable fileStoreTable =
+ PostponeUtils.tableForPostponeCompact(table, defaultBucketNum,
snapshotId);
FixedBucketSink sink = new FixedBucketSink(fileStoreTable, null);
DataStream<Committable> dataStream = dataStreams.get(0);
for (int i = 1; i < dataStreams.size(); i++) {
@@ -379,6 +410,31 @@ public class CompactAction extends TableActionBase {
return true;
}
+ private boolean buildNothingToCompact(StreamExecutionEnvironment env) {
+ if (this.forceStartFlinkJob) {
+ env.fromSequence(0, 0).name("Nothing to Compact
Source").sinkTo(new DiscardingSink<>());
+ return true;
+ }
+ return false;
+ }
+
+ private static class CompactBucketChannelComputer implements
ChannelComputer<CompactBucket> {
+
+ private static final long serialVersionUID = 1L;
+
+ private transient int numChannels;
+
+ @Override
+ public void setup(int numChannels) {
+ this.numChannels = numChannels;
+ }
+
+ @Override
+ public int channel(CompactBucket bucket) {
+ return ChannelComputer.select(bucket.partition(), bucket.bucket(),
numChannels);
+ }
+ }
+
@Override
public void run() throws Exception {
if (buildImpl()) {
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactOperator.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactOperator.java
new file mode 100644
index 0000000000..8055908fe3
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactOperator.java
@@ -0,0 +1,134 @@
+/*
+ * 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.paimon.flink.postpone;
+
+import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.disk.IOManagerImpl;
+import org.apache.paimon.flink.sink.Committable;
+import org.apache.paimon.flink.utils.BoundedTwoInputOperator;
+import org.apache.paimon.flink.utils.RuntimeContextUtils;
+import org.apache.paimon.operation.FileSystemWriteRestore;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.PostponeUtils.CompactBucket;
+import org.apache.paimon.table.sink.CommitMessage;
+import org.apache.paimon.table.sink.FixedBucketRowKeyExtractor;
+import org.apache.paimon.table.sink.TableWriteImpl;
+import org.apache.paimon.utils.Pair;
+
+import org.apache.flink.streaming.api.operators.InputSelection;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+/** Writes postponed rows and compacts existing Level-0 buckets with the same
table writer. */
+public class PostponeBucketCompactOperator
+ extends BoundedTwoInputOperator<InternalRow, CompactBucket,
Committable> {
+
+ private static final long serialVersionUID = 1L;
+
+ private final FileStoreTable table;
+ private final String commitUser;
+ private final long snapshotId;
+
+ private transient IOManagerImpl ioManager;
+ private transient TableWriteImpl<?> write;
+ private transient FixedBucketRowKeyExtractor rowKeyExtractor;
+ private transient Set<Pair<BinaryRow, Integer>> buckets;
+ private transient boolean[] endedInputs;
+
+ public PostponeBucketCompactOperator(FileStoreTable table, String
commitUser, long snapshotId) {
+ this.table = table;
+ this.commitUser = commitUser;
+ this.snapshotId = snapshotId;
+ }
+
+ @Override
+ public void open() throws Exception {
+ super.open();
+ ioManager =
+ new IOManagerImpl(
+ getContainingTask()
+ .getEnvironment()
+ .getIOManager()
+ .getSpillingDirectoriesPaths());
+ write =
+ table.newWrite(
+ commitUser,
+
RuntimeContextUtils.getIndexOfThisSubtask(getRuntimeContext()))
+ .withIOManager(ioManager);
+ write.withWriteRestore(
+ new FileSystemWriteRestore(
+ table.coreOptions(),
+ table.snapshotManager(),
+ table.store().newScan(),
+ table.store().newIndexFileHandler(),
+ snapshotId));
+ rowKeyExtractor = new FixedBucketRowKeyExtractor(table.schema());
+ buckets = new LinkedHashSet<>();
+ endedInputs = new boolean[2];
+ }
+
+ @Override
+ public InputSelection nextSelection() {
+ return InputSelection.ALL;
+ }
+
+ @Override
+ public void processElement1(StreamRecord<InternalRow> element) throws
Exception {
+ InternalRow row = element.getValue();
+ rowKeyExtractor.setRecord(row);
+ buckets.add(Pair.of(rowKeyExtractor.partition().copy(),
rowKeyExtractor.bucket()));
+ write.write(row);
+ }
+
+ @Override
+ public void processElement2(StreamRecord<CompactBucket> element) {
+ CompactBucket bucket = element.getValue();
+ buckets.add(Pair.of(bucket.partition().copy(), bucket.bucket()));
+ }
+
+ @Override
+ public void endInput(int inputId) throws Exception {
+ endedInputs[inputId - 1] = true;
+ if (endedInputs[0] && endedInputs[1]) {
+ for (Pair<BinaryRow, Integer> bucket : buckets) {
+ write.compact(bucket.getLeft(), bucket.getRight(), false);
+ }
+ for (CommitMessage message : write.prepareCommit(true,
Long.MAX_VALUE)) {
+ output.collect(new StreamRecord<>(new
Committable(Long.MAX_VALUE, message)));
+ }
+ }
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ if (write != null) {
+ write.close();
+ }
+ } finally {
+ if (ioManager != null) {
+ ioManager.close();
+ }
+ super.close();
+ }
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSource.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSource.java
index 3e0cd7cf23..5331fd0e9b 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSource.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSource.java
@@ -34,6 +34,7 @@ import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.sink.ChannelComputer;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.snapshot.SnapshotReader;
import org.apache.paimon.utils.Pair;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
@@ -66,11 +67,23 @@ public class PostponeBucketCompactSplitSource extends
AbstractNonCoordinatedSour
private final FileStoreTable table;
private final Map<String, String> partitionSpec;
+ @Nullable private final Long snapshotId;
public PostponeBucketCompactSplitSource(
FileStoreTable table, Map<String, String> partitionSpec) {
+ this(table, partitionSpec, null);
+ }
+
+ public PostponeBucketCompactSplitSource(
+ FileStoreTable table, Map<String, String> partitionSpec, long
snapshotId) {
+ this(table, partitionSpec, Long.valueOf(snapshotId));
+ }
+
+ private PostponeBucketCompactSplitSource(
+ FileStoreTable table, Map<String, String> partitionSpec, @Nullable
Long snapshotId) {
this.table = table;
this.partitionSpec = partitionSpec;
+ this.snapshotId = snapshotId;
}
@Override
@@ -88,8 +101,12 @@ public class PostponeBucketCompactSplitSource extends
AbstractNonCoordinatedSour
@Override
public InputStatus pollNext(ReaderOutput<Split> output) throws
Exception {
+ SnapshotReader snapshotReader = table.newSnapshotReader();
+ if (snapshotId != null) {
+ snapshotReader.withSnapshot(snapshotId);
+ }
List<Split> splits =
- table.newSnapshotReader()
+ snapshotReader
.withPartitionFilter(partitionSpec)
.withBucket(BucketMode.POSTPONE_BUCKET)
.read()
@@ -123,9 +140,32 @@ public class PostponeBucketCompactSplitSource extends
AbstractNonCoordinatedSour
FileStoreTable table,
Map<String, String> partitionSpec,
@Nullable Integer parallelism) {
+ return buildSource(env, table, partitionSpec, null, parallelism);
+ }
+
+ public static Pair<DataStream<RowData>, DataStream<Committable>>
buildSource(
+ StreamExecutionEnvironment env,
+ FileStoreTable table,
+ Map<String, String> partitionSpec,
+ long snapshotId,
+ @Nullable Integer parallelism) {
+ return buildSource(env, table, partitionSpec,
Long.valueOf(snapshotId), parallelism);
+ }
+
+ private static Pair<DataStream<RowData>, DataStream<Committable>>
buildSource(
+ StreamExecutionEnvironment env,
+ FileStoreTable table,
+ Map<String, String> partitionSpec,
+ @Nullable Long snapshotId,
+ @Nullable Integer parallelism) {
+ PostponeBucketCompactSplitSource splitSource =
+ snapshotId == null
+ ? new PostponeBucketCompactSplitSource(table,
partitionSpec)
+ : new PostponeBucketCompactSplitSource(
+ table, partitionSpec, snapshotId.longValue());
DataStream<Split> source =
env.fromSource(
- new PostponeBucketCompactSplitSource(table,
partitionSpec),
+ splitSource,
WatermarkStrategy.noWatermarks(),
String.format(
"Compact split generator: %s - %s",
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSourceTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSourceTest.java
new file mode 100644
index 0000000000..24973e3706
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSourceTest.java
@@ -0,0 +1,57 @@
+/*
+ * 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.paimon.flink.postpone;
+
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.source.snapshot.SnapshotReader;
+
+import org.apache.flink.api.connector.source.ReaderOutput;
+import org.apache.flink.api.connector.source.SourceReader;
+import org.apache.flink.api.connector.source.SourceReaderContext;
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+
+import static org.mockito.Answers.CALLS_REAL_METHODS;
+import static org.mockito.Answers.RETURNS_SELF;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/** Tests for {@link PostponeBucketCompactSplitSource}. */
+class PostponeBucketCompactSplitSourceTest {
+
+ @Test
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ void testReadsPinnedSnapshot() throws Exception {
+ SnapshotReader reader = mock(SnapshotReader.class, RETURNS_SELF);
+ SnapshotReader.Plan plan = mock(SnapshotReader.Plan.class,
CALLS_REAL_METHODS);
+ when(plan.splits()).thenReturn(Collections.emptyList());
+ when(reader.read()).thenReturn(plan);
+ FileStoreTable table = mock(FileStoreTable.class);
+ when(table.newSnapshotReader()).thenReturn(reader);
+
+ PostponeBucketCompactSplitSource source =
+ new PostponeBucketCompactSplitSource(table,
Collections.emptyMap(), 5L);
+ SourceReader sourceReader =
source.createReader(mock(SourceReaderContext.class));
+ sourceReader.pollNext(mock(ReaderOutput.class));
+
+ verify(reader).withSnapshot(5L);
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/procedure/VectorSearchProcedureITCase.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/procedure/VectorSearchProcedureITCase.java
index bc29716356..74ec37d14b 100644
---
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/procedure/VectorSearchProcedureITCase.java
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/procedure/VectorSearchProcedureITCase.java
@@ -32,6 +32,7 @@ import org.apache.paimon.io.CompactIncrement;
import org.apache.paimon.io.DataIncrement;
import org.apache.paimon.options.Options;
import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.PostponeUtils;
import org.apache.paimon.table.sink.BatchTableCommit;
import org.apache.paimon.table.sink.BatchTableWrite;
import org.apache.paimon.table.sink.BatchWriteBuilder;
@@ -40,6 +41,7 @@ import org.apache.paimon.table.sink.CommitMessageImpl;
import org.apache.paimon.types.DataField;
import org.apache.paimon.utils.Range;
+import org.apache.flink.table.api.config.TableConfigOptions;
import org.apache.flink.types.Row;
import org.junit.jupiter.api.Test;
@@ -71,6 +73,65 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
.containsExactlyInAnyOrder("{\"id\":\"2\"}", "{\"id\":\"3\"}");
}
+ @Test
+ public void testPostponeBucketBuildsVectorIndexDuringCompact() throws
Exception {
+ createPostponePrimaryKeyVectorTable("POSTPONE_PK_T");
+
+ sql(
+ "INSERT INTO POSTPONE_PK_T VALUES "
+ + "(1, ARRAY[CAST(3.0 AS FLOAT), CAST(0.0 AS FLOAT)]),
"
+ + "(2, ARRAY[CAST(1.0 AS FLOAT), CAST(0.0 AS FLOAT)]),
"
+ + "(3, ARRAY[CAST(2.0 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+
+ assertThat(sql("SELECT * FROM POSTPONE_PK_T")).isEmpty();
+ assertThat(sql("SELECT bucket FROM `POSTPONE_PK_T$buckets`"))
+ .allMatch(row -> !Integer.valueOf(-2).equals(row.getField(0)));
+ assertThat(sql("SELECT file_path FROM `POSTPONE_PK_T$files` WHERE
level = 0")).isNotEmpty();
+ FileStoreTable table = paimonTable("POSTPONE_PK_T");
+ long snapshotId = table.latestSnapshot().get().id();
+ assertThat(PostponeUtils.getLevel0Buckets(table,
snapshotId)).isNotEmpty();
+
+ tEnv.getConfig().set(TableConfigOptions.TABLE_DML_SYNC, true);
+ sql("CALL sys.compact(`table` => 'default.POSTPONE_PK_T')");
+
+ assertThat(paimonTable("POSTPONE_PK_T").latestSnapshot().get().id())
+ .isGreaterThan(snapshotId);
+ assertThat(sql("SELECT id FROM POSTPONE_PK_T"))
+ .extracting(row -> row.getField(0))
+ .containsExactlyInAnyOrder(1, 2, 3);
+ assertThat(searchPrimaryKeyVectorTable("POSTPONE_PK_T", 2, "id"))
+ .extracting(row -> row.getField(0).toString())
+ .containsExactlyInAnyOrder("{\"id\":\"2\"}", "{\"id\":\"3\"}");
+
+ sql(
+ "INSERT INTO POSTPONE_PK_T VALUES "
+ + "(4, ARRAY[CAST(0.5 AS FLOAT), CAST(0.0 AS FLOAT)]),
"
+ + "(6, ARRAY[CAST(6.0 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+ sql(
+ "INSERT INTO POSTPONE_PK_T /*+
OPTIONS('postpone.batch-write-fixed-bucket' = 'false') */ VALUES "
+ + "(5, ARRAY[CAST(1.5 AS FLOAT), CAST(0.0 AS FLOAT)]),
"
+ + "(7, ARRAY[CAST(7.0 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+ assertThat(sql("SELECT id FROM POSTPONE_PK_T"))
+ .extracting(row -> row.getField(0))
+ .containsExactlyInAnyOrder(1, 2, 3);
+ assertThat(sql("SELECT bucket FROM `POSTPONE_PK_T$buckets`"))
+ .extracting(row -> row.getField(0))
+ .contains(-2);
+ assertThat(sql("SELECT file_path FROM `POSTPONE_PK_T$files` WHERE
level = 0")).isNotEmpty();
+
+ sql("CALL sys.compact(`table` => 'default.POSTPONE_PK_T')");
+
+ assertThat(sql("SELECT id FROM POSTPONE_PK_T"))
+ .extracting(row -> row.getField(0))
+ .containsExactlyInAnyOrder(1, 2, 3, 4, 5, 6, 7);
+ assertThat(sql("SELECT bucket FROM `POSTPONE_PK_T$buckets`"))
+ .allMatch(row -> !Integer.valueOf(-2).equals(row.getField(0)));
+ assertThat(sql("SELECT file_path FROM `POSTPONE_PK_T$files` WHERE
level = 0")).isEmpty();
+ assertThat(searchPrimaryKeyVectorTable("POSTPONE_PK_T", 2, "id"))
+ .extracting(row -> row.getField(0).toString())
+ .containsExactlyInAnyOrder("{\"id\":\"4\"}", "{\"id\":\"2\"}");
+ }
+
@Test
public void testPrimaryKeyVectorSearchAfterUpdateAndDelete() throws
Exception {
createPrimaryKeyVectorTable("PK_UPDATE_T");
@@ -285,6 +346,31 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
tableName, DIMENSION,
TestVectorGlobalIndexerFactory.IDENTIFIER, DIMENSION);
}
+ private void createPostponePrimaryKeyVectorTable(String tableName) {
+ sql(
+ "CREATE TABLE %s ("
+ + "id INT, "
+ + "vec ARRAY<FLOAT>, "
+ + "PRIMARY KEY (id) NOT ENFORCED"
+ + ") WITH ("
+ + "'bucket' = '-2', "
+ + "'write-only' = 'true', "
+ + "'postpone.batch-write-fixed-bucket' = 'true', "
+ + "'file.format' = 'json', "
+ + "'file.compression' = 'deflate', "
+ + "'deletion-vectors.enabled' = 'true', "
+ + "'deletion-vectors.merge-on-read' = 'false', "
+ + "'vector-field' = 'vec', "
+ + "'field.vec.vector-dim' = '%d', "
+ + "'pk-vector.index.columns' = 'vec', "
+ + "'fields.vec.pk-vector.index.type' = '%s', "
+ + "'fields.vec.pk-vector.distance.metric' = 'l2', "
+ + "'test.vector.dimension' = '%d', "
+ + "'test.vector.metric' = 'l2'"
+ + ")",
+ tableName, DIMENSION,
TestVectorGlobalIndexerFactory.IDENTIFIER, DIMENSION);
+ }
+
private List<Row> searchPrimaryKeyVectorTable(String tableName, int topK,
String projection) {
return sql(
"CALL sys.vector_search("
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/procedure/SparkPostponeCompactProcedure.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/procedure/SparkPostponeCompactProcedure.scala
index 4725ce09f9..199702e088 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/procedure/SparkPostponeCompactProcedure.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/procedure/SparkPostponeCompactProcedure.scala
@@ -19,9 +19,9 @@
package org.apache.paimon.spark.procedure
import org.apache.paimon.CoreOptions
-import org.apache.paimon.CoreOptions.BUCKET
import org.apache.paimon.data.BinaryRow
import org.apache.paimon.io.{CompactIncrement, DataFileMeta, DataIncrement}
+import org.apache.paimon.operation.FileSystemWriteRestore
import org.apache.paimon.partition.PartitionPredicate
import org.apache.paimon.postpone.BucketFiles
import org.apache.paimon.spark.PaimonImplicits._
@@ -31,12 +31,13 @@ import org.apache.paimon.spark.util.{ScanPlanHelper,
SparkRowUtils}
import org.apache.paimon.spark.write.{PaimonDataWrite, WriteTaskResult}
import org.apache.paimon.table.{BucketMode, FileStoreTable, PostponeUtils}
import org.apache.paimon.table.sink.{CommitMessage, CommitMessageImpl}
-import org.apache.paimon.utils.BlobDescriptorUtils
+import org.apache.paimon.utils.{BlobDescriptorUtils, SerializationUtils}
+import org.apache.spark.HashPartitioner
+import org.apache.spark.rdd.RDD
import org.apache.spark.sql._
import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
-import org.apache.spark.sql.functions.{col, lit}
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.functions.lit
import org.slf4j.LoggerFactory
import javax.annotation.Nullable
@@ -57,18 +58,13 @@ case class SparkPostponeCompactProcedure(
@transient relation: DataSourceV2Relation) {
private val LOG = LoggerFactory.getLogger(getClass)
- // Unlike `PostponeUtils.tableForFixBucketWrite`, here explicitly set bucket
to 1 without
- // WRITE_ONLY to enable proper compaction logic.
- private lazy val realTable = table.copy(Map(BUCKET.key -> "1").asJava)
-
- // Create bucket computer to determine bucket count for each partition
- private lazy val postponePartitionBucketComputer = {
- val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table)
+ private def createPostponePartitionBucketComputer(snapshotId: Long) = {
+ val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table, snapshotId)
val targetRowNumPerBucket: Option[java.lang.Long] =
table.coreOptions.postponeTargetRowNumPerBucket
val postponeRowCounts =
if (targetRowNumPerBucket.isDefined) {
- PostponeUtils.getPostponeRowCounts(table)
+ PostponeUtils.getPostponeRowCounts(table, snapshotId)
} else {
Collections.emptyMap[BinaryRow, java.lang.Long]()
}
@@ -86,29 +82,11 @@ case class SparkPostponeCompactProcedure(
defaultBucketNum)
}
- private def partitionCols(df: DataFrame): Seq[Column] = {
- val inputSchema = df.schema
- val tableSchema = table.schema
- tableSchema
- .partitionKeys()
- .asScala
- .map(tableSchema.fieldNames().indexOf(_))
- .map(x => col(inputSchema.fieldNames(x)))
- .toSeq
- }
-
- private def repartitionByPartitionsAndBucket(df: DataFrame): DataFrame = {
- df.repartition(partitionCols(df) ++ Seq(col(BUCKET_COL)): _*)
- }
-
- // write data with the bucket processor
- private def writeWithBucketProcessor(
- repartitioned: DataFrame,
- bucketColIdx: Int,
- schema: StructType) = {
- import spark.implicits._
- val rowKindColIdx = SparkRowUtils.getFieldIndex(schema, ROW_KIND_COL)
- val writeBuilder = realTable.newBatchWriteBuilder
+ private def newDataWrite(
+ realTable: FileStoreTable,
+ rowKindColIdx: Int,
+ postponePartitionBucketComputer:
SparkPostponeCompactProcedure.PostponePartitionBucketComputer)
+ : PaimonDataWrite = {
val rowType = table.rowType()
val coreOptions = table.coreOptions()
val catalogContextForBlobDescriptor =
@@ -116,8 +94,8 @@ case class SparkPostponeCompactProcedure(
table.catalogEnvironment().catalogContext(),
coreOptions.toConfiguration)
- def newWrite() = PaimonDataWrite(
- writeBuilder,
+ val dataWrite = PaimonDataWrite(
+ realTable.newBatchWriteBuilder,
rowType,
rowKindColIdx,
writeRowTracking = coreOptions.dataEvolutionEnabled(),
@@ -126,19 +104,7 @@ case class SparkPostponeCompactProcedure(
catalogContextForBlobDescriptor,
Some(postponePartitionBucketComputer)
)
-
- repartitioned.mapPartitions {
- iter =>
- {
- val write = newWrite()
- try {
- iter.foreach(row => write.write(row, row.getInt(bucketColIdx)))
- Iterator.apply(write.commit)
- } finally {
- write.close()
- }
- }
- }
+ dataWrite
}
/** Creates a new BucketFiles instance for tracking file changes */
@@ -156,47 +122,148 @@ case class SparkPostponeCompactProcedure(
partitionPredicate == null,
"Postpone bucket compaction currently does not support specifying
partitions")
+ val snapshot = table.latestSnapshot().orElse(null)
+ if (snapshot == null) {
+ LOG.info("Table has no snapshot, no compact job to execute.")
+ return
+ }
+ val snapshotId = snapshot.id()
+ val postponePartitionBucketComputer =
+ createPostponePartitionBucketComputer(snapshotId)
+ val realTable = PostponeUtils.tableForPostponeCompact(table, 1, snapshotId)
+
// Read data splits from the POSTPONE_BUCKET (-2)
val splits =
table.newSnapshotReader
+ .withSnapshot(snapshotId)
.withBucket(BucketMode.POSTPONE_BUCKET)
.read
.dataSplits
.asScala
+ val compactBuckets = PostponeUtils.getLevel0Buckets(table,
snapshotId).asScala
- if (splits.isEmpty) {
- LOG.info("Partition bucket is empty, no compact job to execute.")
+ if (splits.isEmpty && compactBuckets.isEmpty) {
+ LOG.info("Postpone bucket and real Level-0 buckets are empty, no compact
job to execute.")
return
}
- // Prepare dataset for writing by combining all partitions
- val datasetForWrite: Dataset[Row] = splits
- .groupBy(_.partition)
- .values
- .map(
- split => {
- PaimonUtils
- .createDataset(spark,
ScanPlanHelper.createNewScanPlan(split.toArray, relation))
- })
- .reduce((a, b) => a.union(b))
+ val rowWorkAndKind:
(RDD[SparkPostponeCompactProcedure.PostponeCompactWork], Int) =
+ if (splits.isEmpty) {
+ (spark.sparkContext.emptyRDD, -1)
+ } else {
+ val datasetForWrite: Dataset[Row] = splits
+ .groupBy(_.partition)
+ .values
+ .map(
+ split => {
+ PaimonUtils
+ .createDataset(spark,
ScanPlanHelper.createNewScanPlan(split.toArray, relation))
+ })
+ .reduce((a, b) => a.union(b))
- val withInitBucketCol = datasetForWrite.withColumn(BUCKET_COL, lit(-1))
- val bucketColIdx = SparkRowUtils.getFieldIndex(withInitBucketCol.schema,
BUCKET_COL)
- val encoderGroupWithBucketCol = EncoderSerDeGroup(withInitBucketCol.schema)
+ val withInitBucketCol = datasetForWrite.withColumn(BUCKET_COL, lit(-1))
+ val bucketColIdx =
SparkRowUtils.getFieldIndex(withInitBucketCol.schema, BUCKET_COL)
+ val encoderGroupWithBucketCol =
EncoderSerDeGroup(withInitBucketCol.schema)
+ val processor = PostponeFixBucketProcessor(
+ table,
+ bucketColIdx,
+ encoderGroupWithBucketCol,
+ postponePartitionBucketComputer
+ )
+ val dataFrame = withInitBucketCol
+
.mapPartitions(processor.processPartition)(encoderGroupWithBucketCol.encoder)
+ .toDF()
+ val rowKindColIdx =
SparkRowUtils.getFieldIndex(withInitBucketCol.schema, ROW_KIND_COL)
+ val rowType = table.rowType()
+ val catalogContext = BlobDescriptorUtils.getCatalogContext(
+ table.catalogEnvironment().catalogContext(),
+ table.coreOptions().toConfiguration)
+ val rowWorks = dataFrame.rdd.mapPartitions {
+ rows =>
+ val extractor = realTable.createRowKeyExtractor()
+ val toPaimonRow = SparkRowUtils.toPaimonRow(rowType,
rowKindColIdx, catalogContext)
+ rows.map {
+ row =>
+ extractor.setRecord(toPaimonRow(row))
+ val partition = extractor.partition().copy()
+ SparkPostponeCompactProcedure.PostponeCompactWork(
+ row.copy(),
+ SerializationUtils.serializeBinaryRow(partition),
+ row.getInt(bucketColIdx),
+ compactMarker = false)
+ }
+ }
+ (rowWorks, rowKindColIdx)
+ }
- // Create processor to handle the actual bucket assignment
- val processor = PostponeFixBucketProcessor(
- table,
- bucketColIdx,
- encoderGroupWithBucketCol,
- postponePartitionBucketComputer
- )
+ val markerWorks = compactBuckets.map {
+ bucket =>
+ SparkPostponeCompactProcedure.PostponeCompactWork(
+ null,
+ SerializationUtils.serializeBinaryRow(bucket.partition()),
+ bucket.bucket(),
+ compactMarker = true)
+ }
+ val markerRdd =
+ if (markerWorks.isEmpty) {
+
spark.sparkContext.emptyRDD[SparkPostponeCompactProcedure.PostponeCompactWork]
+ } else {
+ spark.sparkContext.parallelize(
+ markerWorks.toSeq,
+ Math.min(markerWorks.size, Math.max(1,
spark.sparkContext.defaultParallelism)))
+ }
+
+ val partitionedWorks = rowWorkAndKind._1
+ .union(markerRdd)
+ .map(
+ work =>
+ (
+ SparkPostponeCompactProcedure
+ .PostponeCompactKey(work.partition.toIndexedSeq, work.bucket),
+ work))
+ .partitionBy(new HashPartitioner(Math.max(1,
spark.sparkContext.defaultParallelism)))
- val dataFrame = withInitBucketCol
-
.mapPartitions(processor.processPartition)(encoderGroupWithBucketCol.encoder)
- .toDF()
- val repartition = repartitionByPartitionsAndBucket(dataFrame)
- val written = writeWithBucketProcessor(repartition, bucketColIdx,
withInitBucketCol.schema)
+ val written = partitionedWorks.mapPartitions {
+ works =>
+ if (!works.hasNext) {
+ Iterator.empty
+ } else {
+ val dataWrite =
+ newDataWrite(realTable, rowWorkAndKind._2,
postponePartitionBucketComputer)
+ dataWrite.write.withWriteRestore(
+ new FileSystemWriteRestore(
+ realTable.coreOptions(),
+ realTable.snapshotManager(),
+ realTable.store().newScan(),
+ realTable.store().newIndexFileHandler(),
+ snapshotId))
+ var commitInvoked = false
+ try {
+ val pendingBuckets = mutable.LinkedHashMap
+ .empty[SparkPostponeCompactProcedure.PostponeCompactKey,
Array[Byte]]
+ works.foreach {
+ case (key, work) =>
+ pendingBuckets.put(key, work.partition)
+ if (!work.compactMarker) {
+ dataWrite.write(work.row, work.bucket)
+ }
+ }
+ pendingBuckets.foreach {
+ case (key, partition) =>
+ dataWrite.write.compact(
+ SerializationUtils.deserializeBinaryRow(partition),
+ key.bucket,
+ false)
+ }
+ commitInvoked = true
+ Iterator.single(dataWrite.commit)
+ } finally {
+ if (!commitInvoked) {
+ dataWrite.close()
+ }
+ }
+ }
+ }
// Create commit messages for removing old postpone bucket files
val removeMessages = splits.map {
@@ -254,6 +321,16 @@ case class SparkPostponeCompactProcedure(
object SparkPostponeCompactProcedure {
+ private[procedure] case class PostponeCompactWork(
+ row: Row,
+ partition: Array[Byte],
+ bucket: Int,
+ compactMarker: Boolean)
+ extends Serializable
+
+ private[procedure] case class PostponeCompactKey(partition:
IndexedSeq[Byte], bucket: Int)
+ extends Serializable
+
private[procedure] case class PostponePartitionBucketComputer(
knownNumBuckets: java.util.Map[BinaryRow, Integer],
targetRowNumPerBucket: Option[java.lang.Long],
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
index 155f314eba..061fc5238e 100644
---
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
@@ -105,6 +105,75 @@ class PrimaryKeyVectorSearchTest extends
PaimonSparkTestBase {
}
}
+ test("postpone bucket builds primary-key vector index during compact") {
+ withTable("T") {
+ createVectorTable(
+ bucket = -2,
+ extraOptions = Seq(
+ "write-only" -> "true",
+ "postpone.batch-write-fixed-bucket" -> "true",
+ "deletion-vectors.merge-on-read" -> "false"))
+
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, array(3.0f, 0.0f)),
+ | (2, array(1.0f, 0.0f)),
+ | (3, array(2.0f, 0.0f))
+ |""".stripMargin)
+
+ assert(spark.sql("SELECT id FROM T").collect().isEmpty)
+ assert(spark.sql("SELECT file_path FROM `T$files` WHERE level =
0").collect().nonEmpty)
+ assert(!spark.sql("SELECT bucket FROM
`T$buckets`").collect().exists(_.getInt(0) == -2))
+
+ spark.sql("CALL sys.compact(table => 'T')")
+
+ assert(spark.sql("SELECT id FROM T").collect().map(_.getInt(0)).toSet ==
Set(1, 2, 3))
+ val ids = spark
+ .sql("""
+ |SELECT id
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 2)
+ |""".stripMargin)
+ .collect()
+ .map(_.getInt(0))
+ .toSet
+ assert(ids == Set(2, 3))
+
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (4, array(0.5f, 0.0f)),
+ | (6, array(6.0f, 0.0f))
+ |""".stripMargin)
+ withSparkSQLConf("spark.paimon.postpone.batch-write-fixed-bucket" ->
"false") {
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (5, array(1.5f, 0.0f)),
+ | (7, array(7.0f, 0.0f))
+ |""".stripMargin)
+ }
+
+ assert(spark.sql("SELECT id FROM T").collect().map(_.getInt(0)).toSet ==
Set(1, 2, 3))
+ assert(spark.sql("SELECT bucket FROM
`T$buckets`").collect().exists(_.getInt(0) == -2))
+ assert(spark.sql("SELECT file_path FROM `T$files` WHERE level =
0").collect().nonEmpty)
+
+ spark.sql("CALL sys.compact(table => 'T')")
+
+ assert(
+ spark.sql("SELECT id FROM T").collect().map(_.getInt(0)).toSet ==
+ Set(1, 2, 3, 4, 5, 6, 7))
+ assert(!spark.sql("SELECT bucket FROM
`T$buckets`").collect().exists(_.getInt(0) == -2))
+ assert(spark.sql("SELECT file_path FROM `T$files` WHERE level =
0").collect().isEmpty)
+ val mixedIds = spark
+ .sql("""
+ |SELECT id
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 2)
+ |""".stripMargin)
+ .collect()
+ .map(_.getInt(0))
+ .toSet
+ assert(mixedIds == Set(4, 2))
+ }
+ }
+
test("primary-key vector search merges top k across buckets") {
withTable("T") {
createVectorTable(bucket = 4, extraOptions =
Seq("global-index.thread-num" -> "2"))
