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 4a2d409f96 [core] Add asynchronous PK vector ANN compaction (#8584)
4a2d409f96 is described below
commit 4a2d409f96069dc7ccfe126bf53b5aa7a724181c
Author: Jingsong Lee <[email protected]>
AuthorDate: Mon Jul 13 12:09:20 2026 +0800
[core] Add asynchronous PK vector ANN compaction (#8584)
- Build primary-key vector ANN segments asynchronously on a dedicated
executor and honor `prepareCommit(waitCompaction)`.
- Derive LSM-like ANN levels from active segments and rebuild
high-fanout or stale segments without persisting level metadata.
- Search stale-source segments with active row filtering and apply
exact-search fallback according to `GlobalIndexSearchMode`.
- Add independent ANN compaction target-size, level-fanout, and
stale-ratio options with validation and documentation.
Synchronous per-compaction ANN construction increases commit latency,
while accumulating many small or stale ANN segments increases query
fan-out. The new maintenance path decouples ANN work from data
compaction while preserving restart correctness and snapshot-safe search
coverage.
---
docs/docs/primary-key-table/vector-index.md | 6 +
docs/generated/core_configuration.html | 12 +
.../main/java/org/apache/paimon/CoreOptions.java | 22 +
.../pkvector/BucketedVectorIndexMaintainer.java | 336 +++++++++++---
.../paimon/index/pkvector/PkVectorAnnLevels.java | 159 +++++++
.../index/pkvector/PkVectorAnnSegmentSearcher.java | 50 +-
.../pkvector/PrimaryKeyVectorBucketSearch.java | 46 +-
.../paimon/operation/AbstractFileStoreWrite.java | 105 +++--
.../org/apache/paimon/schema/SchemaValidation.java | 10 +
.../paimon/table/source/PrimaryKeyVectorRead.java | 7 +-
.../BucketedVectorIndexMaintainerTest.java | 501 ++++++++++++++++++++-
.../index/pkvector/PkVectorAnnLevelsTest.java | 139 ++++++
.../index/pkvector/PkVectorAnnSegmentFileTest.java | 41 ++
.../pkvector/PrimaryKeyVectorBucketSearchTest.java | 106 ++++-
.../operation/PrimaryKeyVectorIndexWriteTest.java | 68 ++-
.../PrimaryKeyVectorIndexValidationTest.java | 20 +
16 files changed, 1500 insertions(+), 128 deletions(-)
diff --git a/docs/docs/primary-key-table/vector-index.md
b/docs/docs/primary-key-table/vector-index.md
index 8ac875b29d..9604df843e 100644
--- a/docs/docs/primary-key-table/vector-index.md
+++ b/docs/docs/primary-key-table/vector-index.md
@@ -102,6 +102,8 @@ The vector comment directive converts the SQL
`ARRAY<FLOAT>` column to Paimon's
| `fields.<column>.pk-vector.index.type` | Yes | ANN implementation, such as
`ivf-flat`, `ivf-pq`, `ivf-hnsw-flat`, `ivf-hnsw-sq`, or `lumina`. |
| `fields.<column>.pk-vector.distance.metric` | No | `l2`, `cosine`, or
`inner_product`. The default is `inner_product`. |
| `fields.<column>.pk-vector.index.options` | No | JSON object containing
build options for the selected ANN implementation. Unqualified keys are scoped
to that implementation. |
+| `pk-vector.index.compaction.level-fanout` | No | Number of similarly sized
ANN segments which triggers a rebuild and maximum row-count ratio within one
size tier. Default: `5`. |
+| `pk-vector.index.compaction.stale-ratio-threshold` | No | Ratio of rows from
inactive source files which triggers an ANN rebuild. Default: `0.2`. |
For algorithm-specific build and search options, see
[Vector Index](../multimodal-table/global-index/vector).
@@ -113,6 +115,10 @@ 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.
+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`.
+
The maintenance behavior depends on the merge engine:
- `deduplicate`: an update indexes the latest row and the deletion vector
hides the replaced
diff --git a/docs/generated/core_configuration.html
b/docs/generated/core_configuration.html
index d22fb24cdf..5d46fda800 100644
--- a/docs/generated/core_configuration.html
+++ b/docs/generated/core_configuration.html
@@ -1211,6 +1211,18 @@ This config option does not affect the default
filesystem metastore.</td>
<td>String</td>
<td>Comma-separated VECTOR columns indexed by primary-key vector
indexes. Each column owns one index and must define
fields.<column>.pk-vector.index.type. Index options and distance metric
are also field-scoped. The first release supports exactly one column.</td>
</tr>
+ <tr>
+ <td><h5>pk-vector.index.compaction.level-fanout</h5></td>
+ <td style="word-wrap: break-word;">5</td>
+ <td>Integer</td>
+ <td>Number of similarly sized ANN segments that triggers a rebuild
and the maximum row-count ratio within one size tier.</td>
+ </tr>
+ <tr>
+ <td><h5>pk-vector.index.compaction.stale-ratio-threshold</h5></td>
+ <td style="word-wrap: break-word;">0.2</td>
+ <td>Double</td>
+ <td>Ratio of rows belonging to inactive source files that triggers
an ANN segment rebuild.</td>
+ </tr>
<tr>
<td><h5>postpone.batch-write-fixed-bucket</h5></td>
<td style="word-wrap: break-word;">true</td>
diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
index 9b7798948f..e4f43582a1 100644
--- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
+++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
@@ -2741,6 +2741,20 @@ public class CoreOptions implements Serializable {
+ "metric are also field-scoped. The first
release supports exactly "
+ "one column.");
+ public static final ConfigOption<Integer>
PK_VECTOR_INDEX_COMPACTION_LEVEL_FANOUT =
+ key("pk-vector.index.compaction.level-fanout")
+ .intType()
+ .defaultValue(5)
+ .withDescription(
+ "Number of similarly sized ANN segments that
triggers a rebuild and the maximum row-count ratio within one size tier.");
+
+ public static final ConfigOption<Double>
PK_VECTOR_INDEX_COMPACTION_STALE_RATIO_THRESHOLD =
+ key("pk-vector.index.compaction.stale-ratio-threshold")
+ .doubleType()
+ .defaultValue(0.2)
+ .withDescription(
+ "Ratio of rows belonging to inactive source files
that triggers an ANN segment rebuild.");
+
@Immutable
public static final ConfigOption<Boolean> PK_CLUSTERING_OVERRIDE =
key("pk-clustering-override")
@@ -4273,6 +4287,14 @@ public class CoreOptions implements Serializable {
return options.getOptional(PK_VECTOR_INDEX_COLUMNS).isPresent();
}
+ public int primaryKeyVectorIndexCompactionLevelFanout() {
+ return options.get(PK_VECTOR_INDEX_COMPACTION_LEVEL_FANOUT);
+ }
+
+ public double primaryKeyVectorIndexCompactionStaleRatioThreshold() {
+ return options.get(PK_VECTOR_INDEX_COMPACTION_STALE_RATIO_THRESHOLD);
+ }
+
public List<String> primaryKeyVectorIndexColumns() {
String columns = options.get(PK_VECTOR_INDEX_COLUMNS);
if (columns == null) {
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java
index 56ba5f9cfc..e1fefc5479 100644
---
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java
@@ -18,6 +18,7 @@
package org.apache.paimon.index.pkvector;
+import org.apache.paimon.CoreOptions;
import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.index.IndexFileHandler;
import org.apache.paimon.index.IndexFileMeta;
@@ -42,6 +43,10 @@ import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
import static org.apache.paimon.utils.Preconditions.checkArgument;
@@ -55,8 +60,11 @@ public class BucketedVectorIndexMaintainer {
private final String metric;
private final String algorithm;
private final PkVectorDataFileReader.Factory vectorReaderFactory;
+ private final PkVectorAnnLevels annLevels;
+ private ExecutorService executor;
private final List<IndexFileMeta> annSegments;
private final Map<String, DataFileMeta> activeSourceFiles;
+ @Nullable private PendingBuild pendingBuild;
BucketedVectorIndexMaintainer(
int vectorFieldId,
@@ -67,7 +75,8 @@ public class BucketedVectorIndexMaintainer {
String algorithm,
PkVectorDataFileReader.Factory vectorReaderFactory,
List<DataFileMeta> restoredDataFiles,
- List<IndexFileMeta> restoredPayloads) {
+ List<IndexFileMeta> restoredPayloads,
+ ExecutorService executor) {
this.vectorFieldId = vectorFieldId;
this.annSegmentFile = annSegmentFile;
this.vectorField = vectorField;
@@ -75,6 +84,12 @@ public class BucketedVectorIndexMaintainer {
this.metric = metric;
this.algorithm = algorithm;
this.vectorReaderFactory = vectorReaderFactory;
+ CoreOptions coreOptions = new CoreOptions(indexOptions);
+ this.annLevels =
+ new PkVectorAnnLevels(
+
coreOptions.primaryKeyVectorIndexCompactionLevelFanout(),
+
coreOptions.primaryKeyVectorIndexCompactionStaleRatioThreshold());
+ this.executor = executor;
PkVectorBucketIndexState restoredState =
new PkVectorBucketIndexState(vectorFieldId, algorithm,
restoredPayloads);
this.annSegments = new ArrayList<>(restoredState.annSegments());
@@ -87,67 +102,223 @@ public class BucketedVectorIndexMaintainer {
validateCoverage(annSegments, activeSourceFiles);
}
- /** Produces ANN changes in the same compact increment as their source
data-file transition. */
+ /** Produces ANN changes while optionally waiting for asynchronous ANN
construction. */
public synchronized VectorIndexCommit prepareCommit(
- DataIncrement appendIncrement, CompactIncrement compactIncrement) {
+ DataIncrement appendIncrement,
+ CompactIncrement compactIncrement,
+ boolean waitCompaction)
+ throws Exception {
checkArgument(
eligibleFiles(appendIncrement.newFiles()).isEmpty(),
"Append files must not be primary-key vector index sources.");
- Map<String, DataFileMeta> nextSources = new
LinkedHashMap<>(activeSourceFiles);
- Set<String> deletedFiles = new HashSet<>();
- for (DataFileMeta file : compactIncrement.compactBefore()) {
- if (!containsFile(compactIncrement.compactAfter(),
file.fileName())) {
- deletedFiles.add(file.fileName());
- nextSources.remove(file.fileName());
+ List<IndexFileMeta> originalSegments = new ArrayList<>(annSegments);
+ Map<String, DataFileMeta> originalSourceFiles = new
LinkedHashMap<>(activeSourceFiles);
+ List<IndexFileMeta> generated = new ArrayList<>();
+ try {
+ Map<String, DataFileMeta> nextSources = new
LinkedHashMap<>(activeSourceFiles);
+ for (DataFileMeta file : compactIncrement.compactBefore()) {
+ if (!containsFile(compactIncrement.compactAfter(),
file.fileName())) {
+ nextSources.remove(file.fileName());
+ }
+ }
+ for (DataFileMeta file : compactIncrement.compactAfter()) {
+ if (PkVectorSourcePolicy.shouldRead(file)) {
+ nextSources.put(file.fileName(), file);
+ }
+ }
+
+ List<IndexFileMeta> removed = new ArrayList<>();
+
+ List<IndexFileMeta> created = new ArrayList<>();
+ activeSourceFiles.clear();
+ activeSourceFiles.putAll(nextSources);
+
+ while (true) {
+ Optional<CompletedBuild> completed =
finishPendingBuild(waitCompaction);
+ if (completed.isPresent()) {
+ CompletedBuild build = completed.get();
+ generated.add(build.segment);
+ if (canAccept(build)) {
+ replaceSegments(build, created, removed);
+ } else {
+ annSegmentFile.delete(build.segment);
+ }
+ }
+
+ if (pendingBuild == null) {
+ List<DataFileMeta> uncovered = uncoveredFiles();
+ if (!uncovered.isEmpty()) {
+ startPendingBuild(uncovered,
Collections.<IndexFileMeta>emptyList());
+ } else {
+ Optional<PkVectorAnnLevels.Plan> plan =
+ annLevels.pick(annSegments, activeSourceFiles);
+ if (plan.isPresent()) {
+ if (plan.get().sourceFiles().isEmpty()) {
+ removeSegments(plan.get().inputSegments(),
created, removed);
+ continue;
+ }
+ startPendingBuild(plan.get().sourceFiles(),
plan.get().inputSegments());
+ }
+ }
+ }
+ if (!waitCompaction || pendingBuild == null) {
+ break;
+ }
+ }
+
+ validateCoverage(annSegments, activeSourceFiles);
+ boolean hasCompactDataTransition =
+ !compactIncrement.compactBefore().isEmpty()
+ || !compactIncrement.compactAfter().isEmpty();
+ boolean indexChanged = !created.isEmpty() || !removed.isEmpty();
+ Optional<VectorIndexIncrement> appendChange =
+ !indexChanged || hasCompactDataTransition
+ ? Optional.empty()
+ : Optional.of(new VectorIndexIncrement(created,
removed));
+ Optional<VectorIndexIncrement> compactChange =
+ !indexChanged || !hasCompactDataTransition
+ ? Optional.empty()
+ : Optional.of(new VectorIndexIncrement(created,
removed));
+ return new VectorIndexCommit(appendChange, compactChange);
+ } catch (Throwable failure) {
+ rollbackPrepareCommit(originalSegments, originalSourceFiles,
generated, failure);
+ if (failure instanceof Exception) {
+ throw (Exception) failure;
}
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ throw new RuntimeException(failure);
}
- for (DataFileMeta file : compactIncrement.compactAfter()) {
- if (PkVectorSourcePolicy.shouldRead(file)) {
- nextSources.put(file.fileName(), file);
+ }
+
+ private void startPendingBuild(
+ List<DataFileMeta> sourceFiles, List<IndexFileMeta> inputSegments)
{
+ PendingBuild build = new PendingBuild(sourceFiles, inputSegments);
+ build.start();
+ pendingBuild = build;
+ }
+
+ private void rollbackPrepareCommit(
+ List<IndexFileMeta> originalSegments,
+ Map<String, DataFileMeta> originalSourceFiles,
+ List<IndexFileMeta> generated,
+ Throwable failure) {
+ annSegments.clear();
+ annSegments.addAll(originalSegments);
+ activeSourceFiles.clear();
+ activeSourceFiles.putAll(originalSourceFiles);
+
+ PendingBuild build = pendingBuild;
+ pendingBuild = null;
+ if (build != null) {
+ try {
+ build.cancel();
+ } catch (Throwable cleanupFailure) {
+ failure.addSuppressed(cleanupFailure);
}
}
+ for (IndexFileMeta segment : generated) {
+ try {
+ annSegmentFile.delete(segment);
+ } catch (Throwable cleanupFailure) {
+ failure.addSuppressed(cleanupFailure);
+ }
+ }
+ }
+
+ private void replaceSegments(
+ CompletedBuild build, List<IndexFileMeta> created,
List<IndexFileMeta> removed) {
+ removeSegments(build.inputSegments, created, removed);
+ annSegments.add(build.segment);
+ created.add(build.segment);
+ }
- List<IndexFileMeta> retained = new ArrayList<>();
- List<IndexFileMeta> removed = new ArrayList<>();
- for (IndexFileMeta ann : annSegments) {
- if (referencesAny(ann, deletedFiles)) {
- removed.add(ann);
+ private void removeSegments(
+ List<IndexFileMeta> segments,
+ List<IndexFileMeta> created,
+ List<IndexFileMeta> removed) {
+ for (IndexFileMeta segment : segments) {
+ checkArgument(annSegments.remove(segment), "ANN rebuild input is
no longer active.");
+ if (created.remove(segment)) {
+ annSegmentFile.delete(segment);
} else {
- retained.add(ann);
+ removed.add(segment);
}
}
+ }
- Set<String> covered = coveredSources(retained);
+ private List<DataFileMeta> uncoveredFiles() {
+ Set<String> covered = coveredSources(annSegments);
List<DataFileMeta> uncovered = new ArrayList<>();
- for (DataFileMeta file : nextSources.values()) {
+ for (DataFileMeta file : activeSourceFiles.values()) {
if (!covered.contains(file.fileName())) {
uncovered.add(file);
}
}
uncovered.sort(Comparator.comparing(DataFileMeta::fileName));
+ return uncovered;
+ }
+
+ private Optional<CompletedBuild> finishPendingBuild(boolean blocking)
throws Exception {
+ if (pendingBuild == null || (!blocking && !pendingBuild.isDone())) {
+ return Optional.empty();
+ }
- List<IndexFileMeta> created = new ArrayList<>();
+ PendingBuild completed = pendingBuild;
try {
- if (!uncovered.isEmpty()) {
- created.add(buildAnnSegment(uncovered));
+ IndexFileMeta segment = completed.get();
+ pendingBuild = null;
+ return Optional.of(new CompletedBuild(segment,
completed.inputSegments));
+ } catch (CancellationException e) {
+ pendingBuild = null;
+ return Optional.empty();
+ } catch (ExecutionException e) {
+ pendingBuild = null;
+ Throwable cause = e.getCause();
+ if (cause instanceof Exception) {
+ throw (Exception) cause;
+ }
+ if (cause instanceof Error) {
+ throw (Error) cause;
}
- List<IndexFileMeta> nextAnn = new ArrayList<>(retained);
- nextAnn.addAll(created);
- validateCoverage(nextAnn, nextSources);
+ throw new RuntimeException(cause);
+ }
+ }
- activeSourceFiles.clear();
- activeSourceFiles.putAll(nextSources);
- annSegments.clear();
- annSegments.addAll(nextAnn);
- Optional<VectorIndexIncrement> compactChange =
- created.isEmpty() && removed.isEmpty()
- ? Optional.empty()
- : Optional.of(new VectorIndexIncrement(created,
removed));
- return new VectorIndexCommit(Optional.empty(), compactChange);
- } catch (RuntimeException e) {
- deleteAnnSegments(created);
- throw e;
+ private boolean canAccept(CompletedBuild build) {
+ if (!annSegments.containsAll(build.inputSegments)) {
+ return false;
+ }
+ List<IndexFileMeta> retained = new ArrayList<>(annSegments);
+ retained.removeAll(build.inputSegments);
+ Set<String> covered = coveredSources(retained);
+ for (PkVectorSourceFile source :
sourceMeta(build.segment).sourceFiles()) {
+ DataFileMeta activeFile = activeSourceFiles.get(source.fileName());
+ if (activeFile == null
+ || activeFile.rowCount() != source.rowCount()
+ || covered.contains(source.fileName())) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ public synchronized boolean buildNotCompleted() {
+ return pendingBuild != null;
+ }
+
+ public synchronized void withExecutor(ExecutorService executor) {
+ checkArgument(pendingBuild == null, "Cannot replace executor during an
ANN build.");
+ this.executor = executor;
+ }
+
+ public synchronized void close() {
+ PendingBuild build = pendingBuild;
+ pendingBuild = null;
+ if (build != null) {
+ build.cancel();
}
}
@@ -173,11 +344,9 @@ public class BucketedVectorIndexMaintainer {
new PkVectorBucketIndexState(vectorFieldId, algorithm,
candidateAnn);
for (Map.Entry<String, IndexFileMeta> entry :
state.sourceFileToAnnSegment().entrySet()) {
DataFileMeta file = sourceFiles.get(entry.getKey());
- checkArgument(
- file != null,
- "ANN segment %s references inactive data file %s.",
- entry.getValue().fileName(),
- entry.getKey());
+ if (file == null) {
+ continue;
+ }
PkVectorSourceMeta sourceMeta = sourceMeta(entry.getValue());
for (PkVectorSourceFile source : sourceMeta.sourceFiles()) {
if (source.fileName().equals(entry.getKey())) {
@@ -198,9 +367,69 @@ public class BucketedVectorIndexMaintainer {
return new PkVectorBucketIndexState(vectorFieldId, algorithm,
annSegments);
}
- private void deleteAnnSegments(List<IndexFileMeta> segments) {
- for (IndexFileMeta segment : segments) {
- annSegmentFile.delete(segment);
+ private class PendingBuild {
+
+ private final List<DataFileMeta> sourceFiles;
+ private final List<IndexFileMeta> inputSegments;
+ @Nullable private IndexFileMeta result;
+ @Nullable private Future<IndexFileMeta> future;
+ private boolean cancelled;
+
+ private PendingBuild(List<DataFileMeta> sourceFiles,
List<IndexFileMeta> inputSegments) {
+ this.sourceFiles = new ArrayList<>(sourceFiles);
+ this.inputSegments = new ArrayList<>(inputSegments);
+ }
+
+ private void start() {
+ future =
+ executor.submit(
+ () -> {
+ IndexFileMeta segment =
buildAnnSegment(sourceFiles);
+ synchronized (PendingBuild.this) {
+ if (!cancelled) {
+ result = segment;
+ return segment;
+ }
+ }
+ annSegmentFile.delete(segment);
+ throw new CancellationException();
+ });
+ }
+
+ private boolean isDone() {
+ return future.isDone();
+ }
+
+ private IndexFileMeta get() throws InterruptedException,
ExecutionException {
+ return future.get();
+ }
+
+ private void cancel() {
+ Future<IndexFileMeta> buildFuture;
+ IndexFileMeta segment;
+ synchronized (this) {
+ cancelled = true;
+ buildFuture = future;
+ segment = result;
+ result = null;
+ }
+ if (buildFuture != null) {
+ buildFuture.cancel(true);
+ }
+ if (segment != null) {
+ annSegmentFile.delete(segment);
+ }
+ }
+ }
+
+ private static class CompletedBuild {
+
+ private final IndexFileMeta segment;
+ private final List<IndexFileMeta> inputSegments;
+
+ private CompletedBuild(IndexFileMeta segment, List<IndexFileMeta>
inputSegments) {
+ this.segment = segment;
+ this.inputSegments = new ArrayList<>(inputSegments);
}
}
@@ -223,15 +452,6 @@ public class BucketedVectorIndexMaintainer {
return false;
}
- private static boolean referencesAny(IndexFileMeta ann, Set<String>
dataFiles) {
- for (PkVectorSourceFile source : sourceMeta(ann).sourceFiles()) {
- if (dataFiles.contains(source.fileName())) {
- return true;
- }
- }
- return false;
- }
-
private static Set<String> coveredSources(List<IndexFileMeta> segments) {
Set<String> sources = new HashSet<>();
for (IndexFileMeta ann : segments) {
@@ -331,7 +551,8 @@ public class BucketedVectorIndexMaintainer {
BinaryRow partition,
int bucket,
@Nullable List<DataFileMeta> restoredDataFiles,
- @Nullable List<IndexFileMeta> restoredPayloads) {
+ @Nullable List<IndexFileMeta> restoredPayloads,
+ ExecutorService executor) {
checkArgument(indexOptions != null, "ANN index options are not
configured.");
List<DataFileMeta> dataFiles =
restoredDataFiles == null ? Collections.emptyList() :
restoredDataFiles;
@@ -347,7 +568,8 @@ public class BucketedVectorIndexMaintainer {
new PkVectorDataFileReader.Factory(
readerFactoryBuilder, partition, bucket,
vectorField, vectorDimension),
dataFiles,
- payloads);
+ payloads,
+ executor);
}
}
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnLevels.java
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnLevels.java
new file mode 100644
index 0000000000..04c3dd7209
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnLevels.java
@@ -0,0 +1,159 @@
+/*
+ * 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.index.pkvector;
+
+import org.apache.paimon.index.IndexFileMeta;
+import org.apache.paimon.io.DataFileMeta;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.TreeMap;
+
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+
+/** Derives logical ANN compaction levels from immutable segment source
metadata. */
+final class PkVectorAnnLevels {
+
+ private final int fanout;
+ private final double staleRatioThreshold;
+
+ PkVectorAnnLevels(int fanout, double staleRatioThreshold) {
+ checkArgument(fanout > 1, "ANN level fanout must be greater than
one.");
+ checkArgument(
+ staleRatioThreshold > 0 && staleRatioThreshold <= 1,
+ "ANN stale ratio threshold must be in (0, 1].");
+ this.fanout = fanout;
+ this.staleRatioThreshold = staleRatioThreshold;
+ }
+
+ Optional<Plan> pick(List<IndexFileMeta> segments, Map<String,
DataFileMeta> activeSourceFiles) {
+ IndexFileMeta staleCandidate = null;
+ double highestStaleRatio = -1;
+ for (IndexFileMeta segment : segments) {
+ double staleRatio = staleRatio(segment, activeSourceFiles);
+ if (staleRatio >= staleRatioThreshold
+ && (staleRatio > highestStaleRatio
+ || (staleRatio == highestStaleRatio
+ && (staleCandidate == null
+ || segment.fileName()
+
.compareTo(staleCandidate.fileName())
+ < 0)))) {
+ staleCandidate = segment;
+ highestStaleRatio = staleRatio;
+ }
+ }
+ if (staleCandidate != null) {
+ return Optional.of(
+ createPlan(Collections.singletonList(staleCandidate),
activeSourceFiles));
+ }
+
+ List<IndexFileMeta> candidates = new ArrayList<>(segments);
+ candidates.sort(
+ Comparator.comparingLong(PkVectorAnnLevels::buildRowCount)
+ .thenComparing(IndexFileMeta::fileName));
+ for (int start = 0; start + fanout <= candidates.size(); start++) {
+ long smallest = buildRowCount(candidates.get(start));
+ long largest = buildRowCount(candidates.get(start + fanout - 1));
+ if (largest <= saturatedMultiply(smallest, fanout)) {
+ return Optional.of(
+ createPlan(
+ new ArrayList<>(candidates.subList(start,
start + fanout)),
+ activeSourceFiles));
+ }
+ }
+ return Optional.empty();
+ }
+
+ private static double staleRatio(
+ IndexFileMeta segment, Map<String, DataFileMeta>
activeSourceFiles) {
+ long totalRows = 0;
+ long staleRows = 0;
+ for (PkVectorSourceFile source :
PkVectorSourceMeta.fromIndexFile(segment).sourceFiles()) {
+ totalRows = Math.addExact(totalRows, source.rowCount());
+ DataFileMeta active = activeSourceFiles.get(source.fileName());
+ if (active == null) {
+ staleRows = Math.addExact(staleRows, source.rowCount());
+ } else {
+ checkArgument(
+ active.rowCount() == source.rowCount(),
+ "ANN source %s row count does not match active data
file.",
+ source.fileName());
+ }
+ }
+ return totalRows == 0 ? 0 : ((double) staleRows) / totalRows;
+ }
+
+ private static Plan createPlan(
+ List<IndexFileMeta> inputSegments, Map<String, DataFileMeta>
activeSourceFiles) {
+ Map<String, DataFileMeta> selectedSources = new TreeMap<>();
+ for (IndexFileMeta segment : inputSegments) {
+ for (PkVectorSourceFile source :
+ PkVectorSourceMeta.fromIndexFile(segment).sourceFiles()) {
+ DataFileMeta active = activeSourceFiles.get(source.fileName());
+ if (active != null) {
+ checkArgument(
+ active.rowCount() == source.rowCount(),
+ "ANN source %s row count does not match active
data file.",
+ source.fileName());
+ selectedSources.put(active.fileName(), active);
+ }
+ }
+ }
+ return new Plan(inputSegments, new
ArrayList<>(selectedSources.values()));
+ }
+
+ private static long buildRowCount(IndexFileMeta segment) {
+ long rowCount = 0;
+ for (PkVectorSourceFile source :
PkVectorSourceMeta.fromIndexFile(segment).sourceFiles()) {
+ rowCount = Math.addExact(rowCount, source.rowCount());
+ }
+ return rowCount;
+ }
+
+ private static long saturatedMultiply(long value, int multiplier) {
+ if (value > Long.MAX_VALUE / multiplier) {
+ return Long.MAX_VALUE;
+ }
+ return value * multiplier;
+ }
+
+ /** A deterministic ANN rebuild selection. */
+ static final class Plan {
+
+ private final List<IndexFileMeta> inputSegments;
+ private final List<DataFileMeta> sourceFiles;
+
+ private Plan(List<IndexFileMeta> inputSegments, List<DataFileMeta>
sourceFiles) {
+ this.inputSegments = Collections.unmodifiableList(new
ArrayList<>(inputSegments));
+ this.sourceFiles = Collections.unmodifiableList(new
ArrayList<>(sourceFiles));
+ }
+
+ List<IndexFileMeta> inputSegments() {
+ return inputSegments;
+ }
+
+ List<DataFileMeta> sourceFiles() {
+ return sourceFiles;
+ }
+ }
+}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java
index 5722995225..b59d6a6ae5 100644
---
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java
@@ -41,9 +41,11 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
+import java.util.Set;
import java.util.concurrent.ExecutorService;
import static org.apache.paimon.utils.Preconditions.checkArgument;
@@ -102,6 +104,28 @@ public class PkVectorAnnSegmentSearcher {
int limit,
Map<String, DeletionVector> deletionVectors,
Map<String, String> searchOptions) {
+ Set<String> activeSourceFiles = new HashSet<>();
+ for (PkVectorSourceFile sourceFile : sourceMeta.sourceFiles()) {
+ activeSourceFiles.add(sourceFile.fileName());
+ }
+ return search(
+ segment,
+ sourceMeta,
+ query,
+ limit,
+ deletionVectors,
+ activeSourceFiles,
+ searchOptions);
+ }
+
+ public List<PkVectorSearchResult> search(
+ IndexFileMeta segment,
+ PkVectorSourceMeta sourceMeta,
+ float[] query,
+ int limit,
+ Map<String, DeletionVector> deletionVectors,
+ Set<String> activeSourceFiles,
+ Map<String, String> searchOptions) {
checkArgument(limit > 0, "Vector search limit must be positive: %s.",
limit);
GlobalIndexMeta globalIndexMeta = segment.globalIndexMeta();
checkArgument(
@@ -135,7 +159,7 @@ public class PkVectorAnnSegmentSearcher {
try {
VectorSearch search = new VectorSearch(query, limit,
vectorField.name(), searchOptions);
RoaringNavigableMap64 liveRows =
- liveRowPositions(sourceMeta.sourceFiles(),
deletionVectors);
+ liveRowPositions(sourceMeta.sourceFiles(),
activeSourceFiles, deletionVectors);
if (liveRows != null) {
search.withIncludeRowIds(liveRows);
}
@@ -155,6 +179,11 @@ public class PkVectorAnnSegmentSearcher {
ordinal,
sourceRowCount);
FilePosition filePosition =
filePosition(sourceMeta.sourceFiles(), ordinal);
+ checkArgument(
+ activeSourceFiles.contains(filePosition.dataFileName),
+ "ANN segment %s returned inactive source %s.",
+ segment.fileName(),
+ filePosition.dataFileName);
DeletionVector deletionVector =
deletionVectors.get(filePosition.dataFileName);
checkArgument(
deletionVector == null
@@ -178,18 +207,29 @@ public class PkVectorAnnSegmentSearcher {
@Nullable
private static RoaringNavigableMap64 liveRowPositions(
- List<PkVectorSourceFile> sourceFiles, Map<String, DeletionVector>
deletionVectors) {
- if (deletionVectors.isEmpty()) {
+ List<PkVectorSourceFile> sourceFiles,
+ Set<String> activeSourceFiles,
+ Map<String, DeletionVector> deletionVectors) {
+ boolean allSourcesActive = true;
+ for (PkVectorSourceFile sourceFile : sourceFiles) {
+ if (!activeSourceFiles.contains(sourceFile.fileName())) {
+ allSourcesActive = false;
+ break;
+ }
+ }
+ if (allSourcesActive && deletionVectors.isEmpty()) {
return null;
}
RoaringNavigableMap64 live = new RoaringNavigableMap64();
RoaringNavigableMap64 deleted = new RoaringNavigableMap64();
long fileOffset = 0;
for (PkVectorSourceFile sourceFile : sourceFiles) {
- if (sourceFile.rowCount() > 0) {
+ boolean active = activeSourceFiles.contains(sourceFile.fileName());
+ if (active && sourceFile.rowCount() > 0) {
live.addRange(new Range(fileOffset, fileOffset +
sourceFile.rowCount() - 1));
}
- DeletionVector deletionVector =
deletionVectors.get(sourceFile.fileName());
+ DeletionVector deletionVector =
+ active ? deletionVectors.get(sourceFile.fileName()) : null;
if (deletionVector != null) {
final long offset = fileOffset;
deletionVector.forEachDeletedPosition(position ->
deleted.add(offset + position));
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
index f7b358cf3b..8c2a77a135 100644
---
a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
+++
b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java
@@ -18,6 +18,7 @@
package org.apache.paimon.index.pkvector;
+import org.apache.paimon.CoreOptions.GlobalIndexSearchMode;
import org.apache.paimon.deletionvectors.DeletionVector;
import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.io.DataFileMeta;
@@ -38,7 +39,7 @@ import java.util.function.LongPredicate;
import static org.apache.paimon.utils.Preconditions.checkArgument;
-/** ANN plus exact data-file fallback search for one snapshot bucket. */
+/** ANN search plus search-mode-controlled exact fallback for one snapshot
bucket. */
public class PrimaryKeyVectorBucketSearch {
private static final Comparator<PkVectorSearchResult> BEST_FIRST =
@@ -50,16 +51,19 @@ public class PrimaryKeyVectorBucketSearch {
@Nullable private final PkVectorAnnSegmentSearcher annSearcher;
private final Map<String, String> searchOptions;
private final String metric;
+ private final GlobalIndexSearchMode searchMode;
public PrimaryKeyVectorBucketSearch(
PkVectorDataFileReader.Factory vectorReaderFactory,
@Nullable PkVectorAnnSegmentSearcher annSearcher,
Map<String, String> searchOptions,
- String metric) {
+ String metric,
+ GlobalIndexSearchMode searchMode) {
this.vectorReaderFactory = vectorReaderFactory;
this.annSearcher = annSearcher;
this.searchOptions = Collections.unmodifiableMap(new
HashMap<>(searchOptions));
this.metric = metric;
+ this.searchMode = searchMode;
}
public List<PkVectorSearchResult> search(
@@ -76,13 +80,17 @@ public class PrimaryKeyVectorBucketSearch {
}
PriorityQueue<PkVectorSearchResult> nearest =
new PriorityQueue<>(limit, BEST_FIRST.reversed());
+ Set<String> activeSourceFiles = new HashSet<>(filesByName.keySet());
Set<String> covered = new HashSet<>();
for (IndexFileMeta ann : state.annSegments()) {
PkVectorSourceMeta sourceMeta =
PkVectorSourceMeta.fromIndexFile(ann);
for (PkVectorSourceFile source : sourceMeta.sourceFiles()) {
DataFileMeta file = filesByName.get(source.fileName());
+ if (file == null) {
+ continue;
+ }
checkArgument(
- file != null && file.rowCount() == source.rowCount(),
+ file.rowCount() == source.rowCount(),
"ANN source %s does not match the active data file.",
source.fileName());
covered.add(source.fileName());
@@ -90,22 +98,30 @@ public class PrimaryKeyVectorBucketSearch {
checkArgument(annSearcher != null, "ANN search is not
configured.");
for (PkVectorSearchResult result :
annSearcher.search(
- ann, sourceMeta, query, limit, deletionVectors,
searchOptions)) {
+ ann,
+ sourceMeta,
+ query,
+ limit,
+ deletionVectors,
+ activeSourceFiles,
+ searchOptions)) {
add(nearest, result, limit);
}
}
- for (DataFileMeta file : activeFiles) {
- if (covered.contains(file.fileName())) {
- continue;
- }
- DeletionVector dv = deletionVectors.get(file.fileName());
- LongPredicate excluded = dv == null ? position -> false :
dv::isDeleted;
- try (PkVectorReader reader = vectorReaderFactory.create(file)) {
- for (PkVectorSearchResult result :
- PkVectorExactSearcher.search(
- file.fileName(), reader, query, metric, limit,
excluded)) {
- add(nearest, result, limit);
+ if (searchMode != GlobalIndexSearchMode.FAST) {
+ for (DataFileMeta file : activeFiles) {
+ if (covered.contains(file.fileName())) {
+ continue;
+ }
+ DeletionVector dv = deletionVectors.get(file.fileName());
+ LongPredicate excluded = dv == null ? position -> false :
dv::isDeleted;
+ try (PkVectorReader reader = vectorReaderFactory.create(file))
{
+ for (PkVectorSearchResult result :
+ PkVectorExactSearcher.search(
+ file.fileName(), reader, query, metric,
limit, excluded)) {
+ add(nearest, result, limit);
+ }
}
}
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java
b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java
index 1014ddd118..0409f9492a 100644
---
a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java
+++
b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java
@@ -96,6 +96,7 @@ public abstract class AbstractFileStoreWrite<T> implements
FileStoreWrite<T> {
protected WriteRestore restore;
private ExecutorService lazyCompactExecutor;
+ private ExecutorService lazyVectorIndexExecutor;
private boolean closeCompactExecutorWhenLeaving = true;
private boolean ignorePreviousFiles = false;
private boolean ignoreNumBucketCheck = false;
@@ -248,33 +249,11 @@ public abstract class AbstractFileStoreWrite<T>
implements FileStoreWrite<T> {
.newIndexFiles()
.addAll(writerContainer.dynamicBucketMaintainer.prepareCommit());
}
- if (writerContainer.vectorIndexMaintainer != null) {
- BucketedVectorIndexMaintainer.VectorIndexCommit
vectorCommit =
-
writerContainer.vectorIndexMaintainer.prepareCommit(
- newFilesIncrement, compactIncrement);
- vectorCommit
- .appendIncrement()
- .ifPresent(
- vectorIncrement -> {
- newFilesIncrement
- .newIndexFiles()
-
.addAll(vectorIncrement.newIndexFiles());
- newFilesIncrement
- .deletedIndexFiles()
-
.addAll(vectorIncrement.deletedIndexFiles());
- });
- vectorCommit
- .compactIncrement()
- .ifPresent(
- vectorIncrement -> {
- compactIncrement
- .newIndexFiles()
-
.addAll(vectorIncrement.newIndexFiles());
- compactIncrement
- .deletedIndexFiles()
-
.addAll(vectorIncrement.deletedIndexFiles());
- });
- }
+ applyVectorIndexCommit(
+ writerContainer.vectorIndexMaintainer,
+ newFilesIncrement,
+ compactIncrement,
+ waitCompaction);
CompactDeletionFile compactDeletionFile =
increment.compactDeletionFile();
if (compactDeletionFile != null) {
compactDeletionFile
@@ -291,7 +270,10 @@ public abstract class AbstractFileStoreWrite<T> implements
FileStoreWrite<T> {
result.add(committable);
if (committable.isEmpty()) {
- if (writerCleanChecker.apply(writerContainer)) {
+ if (writerCleanChecker.apply(writerContainer)
+ && (writerContainer.vectorIndexMaintainer == null
+ || !writerContainer.vectorIndexMaintainer
+ .buildNotCompleted())) {
// Clear writer if no update, and if its latest
modification has committed.
//
// We need a mechanism to clear writers, otherwise
there will be more and
@@ -307,6 +289,9 @@ public abstract class AbstractFileStoreWrite<T> implements
FileStoreWrite<T> {
commitIdentifier);
}
writerContainer.writer.close();
+ if (writerContainer.vectorIndexMaintainer != null) {
+ writerContainer.vectorIndexMaintainer.close();
+ }
bucketIter.remove();
}
} else {
@@ -322,6 +307,43 @@ public abstract class AbstractFileStoreWrite<T> implements
FileStoreWrite<T> {
return result;
}
+ private void applyVectorIndexCommit(
+ @Nullable BucketedVectorIndexMaintainer vectorIndexMaintainer,
+ DataIncrement newFilesIncrement,
+ CompactIncrement compactIncrement,
+ boolean waitCompaction)
+ throws Exception {
+ if (vectorIndexMaintainer == null) {
+ return;
+ }
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit vectorCommit =
+ vectorIndexMaintainer.prepareCommit(
+ newFilesIncrement, compactIncrement, waitCompaction);
+ vectorCommit
+ .appendIncrement()
+ .ifPresent(
+ vectorIncrement -> {
+ newFilesIncrement
+ .newIndexFiles()
+ .addAll(vectorIncrement.newIndexFiles());
+ newFilesIncrement
+ .deletedIndexFiles()
+
.addAll(vectorIncrement.deletedIndexFiles());
+ });
+ vectorCommit
+ .compactIncrement()
+ .ifPresent(
+ vectorIncrement -> {
+ compactIncrement
+ .newIndexFiles()
+ .addAll(vectorIncrement.newIndexFiles());
+ compactIncrement
+ .deletedIndexFiles()
+
.addAll(vectorIncrement.deletedIndexFiles());
+ });
+ }
+
// This abstract function returns a whole function (instead of just a
boolean value),
// because we do not want to introduce `commitUser` into this base class.
//
@@ -360,12 +382,18 @@ public abstract class AbstractFileStoreWrite<T>
implements FileStoreWrite<T> {
for (Map<Integer, WriterContainer<T>> bucketWriters :
writers.values()) {
for (WriterContainer<T> writerContainer : bucketWriters.values()) {
writerContainer.writer.close();
+ if (writerContainer.vectorIndexMaintainer != null) {
+ writerContainer.vectorIndexMaintainer.close();
+ }
}
}
writers.clear();
if (lazyCompactExecutor != null && closeCompactExecutorWhenLeaving) {
lazyCompactExecutor.shutdownNow();
}
+ if (lazyVectorIndexExecutor != null) {
+ lazyVectorIndexExecutor.shutdownNow();
+ }
if (compactionMetrics != null) {
compactionMetrics.close();
}
@@ -386,6 +414,11 @@ public abstract class AbstractFileStoreWrite<T> implements
FileStoreWrite<T> {
CommitIncrement increment;
try {
increment = writerContainer.writer.prepareCommit(false);
+ applyVectorIndexCommit(
+ writerContainer.vectorIndexMaintainer,
+ increment.newFilesIncrement(),
+ increment.compactIncrement(),
+ true);
} catch (Exception e) {
throw new RuntimeException(
"Failed to extract state from writer of partition "
@@ -435,6 +468,9 @@ public abstract class AbstractFileStoreWrite<T> implements
FileStoreWrite<T> {
// not ignore them.
false);
notifyNewWriter(writer);
+ if (state.vectorIndexMaintainer != null) {
+
state.vectorIndexMaintainer.withExecutor(vectorIndexExecutor());
+ }
WriterContainer<T> writerContainer =
new WriterContainer<>(
writer,
@@ -515,7 +551,8 @@ public abstract class AbstractFileStoreWrite<T> implements
FileStoreWrite<T> {
partition,
bucket,
restored.dataFiles(),
- restored.vectorIndexPayloads());
+ restored.vectorIndexPayloads(),
+ vectorIndexExecutor());
List<DataFileMeta> restoreFiles = restored.dataFiles();
if (restoreFiles == null) {
@@ -636,6 +673,16 @@ public abstract class AbstractFileStoreWrite<T> implements
FileStoreWrite<T> {
return lazyCompactExecutor;
}
+ private ExecutorService vectorIndexExecutor() {
+ if (lazyVectorIndexExecutor == null) {
+ lazyVectorIndexExecutor =
+ Executors.newSingleThreadExecutor(
+ new ExecutorThreadFactory(
+ Thread.currentThread().getName() +
"-vector-index"));
+ }
+ return lazyVectorIndexExecutor;
+ }
+
@VisibleForTesting
public ExecutorService getCompactExecutor() {
return lazyCompactExecutor;
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 09cb61c0bf..4f68c0b996 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
@@ -890,6 +890,16 @@ public class SchemaValidation {
return;
}
+ checkArgument(
+ options.primaryKeyVectorIndexCompactionLevelFanout() > 1,
+ "%s must be greater than 1.",
+ CoreOptions.PK_VECTOR_INDEX_COMPACTION_LEVEL_FANOUT.key());
+ double staleRatio =
options.primaryKeyVectorIndexCompactionStaleRatioThreshold();
+ checkArgument(
+ staleRatio > 0 && staleRatio <= 1,
+ "%s must be in (0, 1].",
+
CoreOptions.PK_VECTOR_INDEX_COMPACTION_STALE_RATIO_THRESHOLD.key());
+
List<String> indexColumns = options.primaryKeyVectorIndexColumns();
checkArgument(
new HashSet<>(indexColumns).size() == indexColumns.size(),
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
index b70f19eb73..5f1868f0ec 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
@@ -198,7 +198,12 @@ public class PrimaryKeyVectorRead implements VectorRead,
Serializable {
metric,
context.executor);
PrimaryKeyVectorBucketSearch bucketSearch =
- new PrimaryKeyVectorBucketSearch(readerFactory, annSearcher,
searchOptions, metric);
+ new PrimaryKeyVectorBucketSearch(
+ readerFactory,
+ annSearcher,
+ searchOptions,
+ metric,
+ table.coreOptions().globalIndexSearchMode());
List<Candidate> candidates = new ArrayList<>();
for (PkVectorSearchResult result :
bucketSearch.search(state, activeFiles, deletionVectors,
query, limit)) {
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java
index 31b4437151..d86e412366 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java
@@ -31,15 +31,28 @@ import org.apache.paimon.stats.SimpleStats;
import org.apache.paimon.types.DataField;
import org.apache.paimon.types.DataTypes;
+import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
+import java.util.List;
import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.Assertions.catchThrowable;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -47,14 +60,152 @@ import static org.mockito.Mockito.when;
class BucketedVectorIndexMaintainerTest {
@TempDir java.nio.file.Path tempPath;
+ private final ExecutorService executor =
Executors.newSingleThreadExecutor();
+
+ @AfterEach
+ void shutdownExecutor() {
+ executor.shutdownNow();
+ }
+
+ @Test
+ void testNonBlockingPrepareCommitPublishesCompletedAnnLater() throws
Exception {
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO,
pathFactory());
+ DataField vectorField =
+ new DataField(7, "embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()));
+ DataFileMeta data = dataFile("data");
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ CountDownLatch buildStarted = new CountDownLatch(1);
+ CountDownLatch allowBuild = new CountDownLatch(1);
+ AtomicReference<Thread> buildThread = new AtomicReference<>();
+ Thread callerThread = Thread.currentThread();
+ when(readerFactory.create(data))
+ .thenAnswer(
+ ignored -> {
+ buildThread.set(Thread.currentThread());
+ buildStarted.countDown();
+ assertThat(allowBuild.await(30,
TimeUnit.SECONDS)).isTrue();
+ return reader(new float[][] {{1, 0}});
+ });
+ try {
+ BucketedVectorIndexMaintainer maintainer =
+ new BucketedVectorIndexMaintainer(
+ 7,
+ annFile,
+ vectorField,
+ indexOptions(),
+ "l2",
+ "test-vector-ann",
+ readerFactory,
+ Collections.emptyList(),
+ Collections.emptyList(),
+ executor);
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit firstCommit =
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
+ new CompactIncrement(
+ Collections.emptyList(),
+ Collections.singletonList(data),
+ Collections.emptyList()),
+ false);
+
+ assertThat(firstCommit.compactIncrement()).isEmpty();
+ assertThat(buildStarted.await(30, TimeUnit.SECONDS)).isTrue();
+ assertThat(buildThread.get()).isNotSameAs(callerThread);
+ allowBuild.countDown();
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit secondCommit =
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
+ CompactIncrement.emptyIncrement(),
+ true);
+
+ assertThat(secondCommit.appendIncrement()).isPresent();
+
assertThat(secondCommit.appendIncrement().get().newIndexFiles()).hasSize(1);
+ assertThat(secondCommit.compactIncrement()).isEmpty();
+ } finally {
+ allowBuild.countDown();
+ }
+ }
@Test
- void testPartialCompactionRebuildsMultiSourceAnnFromDataFiles() throws
Exception {
+ void testDiscardsStaleBuildAndIndexesLatestSources() throws Exception {
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO,
pathFactory());
+ DataField vectorField =
+ new DataField(7, "embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()));
+ DataFileMeta data1 = dataFile("data-1");
+ DataFileMeta data2 = dataFile("data-2");
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ CountDownLatch buildStarted = new CountDownLatch(1);
+ CountDownLatch allowBuild = new CountDownLatch(1);
+ when(readerFactory.create(data1))
+ .thenAnswer(
+ ignored -> {
+ buildStarted.countDown();
+ assertThat(allowBuild.await(30,
TimeUnit.SECONDS)).isTrue();
+ return reader(new float[][] {{1, 0}});
+ });
+ PkVectorDataFileReader reader2 = reader(new float[][] {{2, 0}});
+ when(readerFactory.create(data2)).thenReturn(reader2);
+
+ try {
+ BucketedVectorIndexMaintainer maintainer =
+ new BucketedVectorIndexMaintainer(
+ 7,
+ annFile,
+ vectorField,
+ indexOptions(),
+ "l2",
+ "test-vector-ann",
+ readerFactory,
+ Collections.emptyList(),
+ Collections.emptyList(),
+ executor);
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
+ new CompactIncrement(
+ Collections.emptyList(),
+ Collections.singletonList(data1),
+ Collections.emptyList()),
+ false);
+ assertThat(buildStarted.await(30, TimeUnit.SECONDS)).isTrue();
+
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
+ new CompactIncrement(
+ Collections.singletonList(data1),
+ Collections.singletonList(data2),
+ Collections.emptyList()),
+ false);
+ allowBuild.countDown();
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit commit =
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
+ CompactIncrement.emptyIncrement(),
+ true);
+
+ assertThat(commit.appendIncrement()).isPresent();
+ IndexFileMeta segment =
commit.appendIncrement().get().newIndexFiles().get(0);
+ assertThat(PkVectorSourceMeta.fromIndexFile(segment).sourceFiles())
+ .extracting(PkVectorSourceFile::fileName)
+ .containsExactly("data-2");
+ assertThat(annFile.exists(segment)).isTrue();
+ } finally {
+ allowBuild.countDown();
+ }
+ }
+
+ @Test
+ void testPartialCompactionKeepsOldAnnAndIndexesNewSource() throws
Exception {
LocalFileIO fileIO = LocalFileIO.create();
PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO,
pathFactory());
DataField vectorField =
new DataField(7, "embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()));
Options options = indexOptions();
+ options.setString("pk-vector.index.compaction.stale-ratio-threshold",
"1.0");
DataFileMeta data1 = dataFile("data-1");
DataFileMeta data2 = dataFile("data-2");
DataFileMeta data3 = dataFile("data-3");
@@ -84,7 +235,8 @@ class BucketedVectorIndexMaintainerTest {
"test-vector-ann",
readerFactory,
Arrays.asList(data1, data2),
- Collections.singletonList(initialAnn));
+ Collections.singletonList(initialAnn),
+ executor);
BucketedVectorIndexMaintainer.VectorIndexCommit commit =
maintainer.prepareCommit(
@@ -92,17 +244,19 @@ class BucketedVectorIndexMaintainerTest {
new CompactIncrement(
Collections.singletonList(data1),
Collections.singletonList(data3),
- Collections.emptyList()));
+ Collections.emptyList()),
+ true);
BucketedVectorIndexMaintainer.VectorIndexIncrement increment =
commit.compactIncrement().get();
- assertThat(increment.deletedIndexFiles()).containsExactly(initialAnn);
+ assertThat(increment.deletedIndexFiles()).isEmpty();
assertThat(increment.newIndexFiles()).hasSize(1);
- IndexFileMeta repaired = increment.newIndexFiles().get(0);
- assertThat(repaired.indexType()).isEqualTo("test-vector-ann");
- assertThat(PkVectorSourceMeta.fromIndexFile(repaired).sourceFiles())
+ IndexFileMeta delta = increment.newIndexFiles().get(0);
+ assertThat(delta.indexType()).isEqualTo("test-vector-ann");
+ assertThat(PkVectorSourceMeta.fromIndexFile(delta).sourceFiles())
.extracting(PkVectorSourceFile::fileName)
- .containsExactly("data-2", "data-3");
+ .containsExactly("data-3");
+ assertThat(maintainer.segments()).containsExactly(initialAnn, delta);
}
@Test
@@ -125,7 +279,8 @@ class BucketedVectorIndexMaintainerTest {
"test-vector-ann",
readerFactory,
Collections.emptyList(),
- Collections.emptyList());
+ Collections.emptyList(),
+ executor);
BucketedVectorIndexMaintainer.VectorIndexCommit commit =
maintainer.prepareCommit(
@@ -133,12 +288,331 @@ class BucketedVectorIndexMaintainerTest {
new CompactIncrement(
Collections.emptyList(),
Collections.singletonList(data),
- Collections.emptyList()));
+ Collections.emptyList()),
+ true);
assertThat(commit.compactIncrement()).isPresent();
assertThat(commit.compactIncrement().get().newIndexFiles()).hasSize(1);
}
+ @Test
+ void testRebuildsDerivedLevelAndAtomicallyReplacesInputs() throws
Exception {
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO,
pathFactory());
+ DataField vectorField =
+ new DataField(7, "embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()));
+ Options options = indexOptions();
+ options.setString("pk-vector.index.compaction.level-fanout", "3");
+ DataFileMeta data1 = dataFile("data-1");
+ DataFileMeta data2 = dataFile("data-2");
+ DataFileMeta data3 = dataFile("data-3");
+ IndexFileMeta ann1 =
+ annFile.build(
+ Collections.singletonList(
+ new PkVectorAnnSegmentFile.Source(
+ data1, new ArrayReader(new float[][]
{{1, 0}}))),
+ vectorField,
+ options,
+ "l2",
+ "test-vector-ann");
+ IndexFileMeta ann2 =
+ annFile.build(
+ Collections.singletonList(
+ new PkVectorAnnSegmentFile.Source(
+ data2, new ArrayReader(new float[][]
{{2, 0}}))),
+ vectorField,
+ options,
+ "l2",
+ "test-vector-ann");
+ IndexFileMeta ann3 =
+ annFile.build(
+ Collections.singletonList(
+ new PkVectorAnnSegmentFile.Source(
+ data3, new ArrayReader(new float[][]
{{3, 0}}))),
+ vectorField,
+ options,
+ "l2",
+ "test-vector-ann");
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ PkVectorDataFileReader reader1 = reader(new float[][] {{1, 0}});
+ PkVectorDataFileReader reader2 = reader(new float[][] {{2, 0}});
+ PkVectorDataFileReader reader3 = reader(new float[][] {{3, 0}});
+ when(readerFactory.create(data1)).thenReturn(reader1);
+ when(readerFactory.create(data2)).thenReturn(reader2);
+ when(readerFactory.create(data3)).thenReturn(reader3);
+ BucketedVectorIndexMaintainer maintainer =
+ new BucketedVectorIndexMaintainer(
+ 7,
+ annFile,
+ vectorField,
+ options,
+ "l2",
+ "test-vector-ann",
+ readerFactory,
+ Arrays.asList(data1, data2, data3),
+ Arrays.asList(ann3, ann1, ann2),
+ executor);
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit commit =
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
CompactIncrement.emptyIncrement(), true);
+
+ BucketedVectorIndexMaintainer.VectorIndexIncrement increment =
+ commit.appendIncrement().get();
+
assertThat(increment.deletedIndexFiles()).containsExactlyInAnyOrder(ann1, ann2,
ann3);
+ assertThat(increment.newIndexFiles()).hasSize(1);
+
assertThat(PkVectorSourceMeta.fromIndexFile(increment.newIndexFiles().get(0)).sourceFiles())
+ .extracting(PkVectorSourceFile::fileName)
+ .containsExactly("data-1", "data-2", "data-3");
+ }
+
+ @Test
+ void testPrepareCommitRollsBackMultipleRebuildsOnFailure() throws
Exception {
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO,
pathFactory());
+ DataField vectorField =
+ new DataField(7, "embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()));
+ Options options = indexOptions();
+ options.setString("pk-vector.index.compaction.level-fanout", "3");
+ List<DataFileMeta> dataFiles = new ArrayList<>();
+ List<IndexFileMeta> initialSegments = new ArrayList<>();
+ for (int i = 1; i <= 6; i++) {
+ DataFileMeta data = dataFile("data-" + i);
+ dataFiles.add(data);
+ initialSegments.add(
+ annFile.build(
+ Collections.singletonList(
+ new PkVectorAnnSegmentFile.Source(
+ data, new ArrayReader(new
float[][] {{(float) i, 0}}))),
+ vectorField,
+ options,
+ "l2",
+ "test-vector-ann"));
+ }
+
+ AtomicInteger readerCount = new AtomicInteger();
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+
when(readerFactory.create(org.mockito.ArgumentMatchers.any(DataFileMeta.class)))
+ .thenAnswer(
+ ignored -> {
+ if (readerCount.incrementAndGet() == 4) {
+ throw new IOException("Expected second ANN
rebuild failure.");
+ }
+ return reader(new float[][] {{1, 0}});
+ });
+ BucketedVectorIndexMaintainer maintainer =
+ new BucketedVectorIndexMaintainer(
+ 7,
+ annFile,
+ vectorField,
+ options,
+ "l2",
+ "test-vector-ann",
+ readerFactory,
+ dataFiles,
+ initialSegments,
+ executor);
+
+ assertThatThrownBy(
+ () ->
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
+ CompactIncrement.emptyIncrement(),
+ true))
+ .isInstanceOf(UncheckedIOException.class);
+
assertThat(maintainer.segments()).containsExactlyInAnyOrderElementsOf(initialSegments);
+ assertThat(fileCount()).isEqualTo(6);
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit retry =
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
CompactIncrement.emptyIncrement(), true);
+
+ BucketedVectorIndexMaintainer.VectorIndexIncrement increment =
+ retry.appendIncrement().get();
+ assertThat(increment.deletedIndexFiles())
+ .containsExactlyInAnyOrderElementsOf(initialSegments);
+ assertThat(increment.newIndexFiles()).hasSize(2);
+ }
+
+ @Test
+ void testInterruptedPrepareCommitCleansPendingBuildOutput() throws
Exception {
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO,
pathFactory());
+ DataField vectorField =
+ new DataField(7, "embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()));
+ DataFileMeta data = dataFile("data");
+ CountDownLatch buildStarted = new CountDownLatch(1);
+ CountDownLatch allowBuild = new CountDownLatch(1);
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ when(readerFactory.create(data))
+ .thenAnswer(
+ ignored -> {
+ buildStarted.countDown();
+ boolean released = false;
+ while (!released) {
+ try {
+ released = allowBuild.await(30,
TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ // Simulate a native ANN build which does
not stop immediately.
+ }
+ }
+ return reader(new float[][] {{1, 0}});
+ });
+ BucketedVectorIndexMaintainer maintainer =
+ new BucketedVectorIndexMaintainer(
+ 7,
+ annFile,
+ vectorField,
+ indexOptions(),
+ "l2",
+ "test-vector-ann",
+ readerFactory,
+ Collections.emptyList(),
+ Collections.emptyList(),
+ executor);
+ AtomicReference<Throwable> failure = new AtomicReference<>();
+ Thread caller =
+ new Thread(
+ () -> {
+ try {
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
+ new CompactIncrement(
+ Collections.emptyList(),
+
Collections.singletonList(data),
+ Collections.emptyList()),
+ true);
+ } catch (Throwable t) {
+ failure.set(t);
+ }
+ });
+
+ try {
+ caller.start();
+ assertThat(buildStarted.await(30, TimeUnit.SECONDS)).isTrue();
+ caller.interrupt();
+ caller.join(TimeUnit.SECONDS.toMillis(30));
+ assertThat(caller.isAlive()).isFalse();
+ assertThat(failure.get()).isInstanceOf(InterruptedException.class);
+ maintainer.close();
+ } finally {
+ allowBuild.countDown();
+ }
+ executor.submit(() -> {}).get(30, TimeUnit.SECONDS);
+
+ assertThat(fileCount()).isZero();
+ assertThat(maintainer.segments()).isEmpty();
+ assertThat(maintainer.buildNotCompleted()).isFalse();
+ }
+
+ @Test
+ void testPrepareCommitRollsBackSourceFilesOnFailure() throws Exception {
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO,
pathFactory());
+ DataField vectorField =
+ new DataField(7, "embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()));
+ Options options = indexOptions();
+ DataFileMeta oldData = dataFile("old-data");
+ DataFileMeta newData = dataFile("new-data");
+ IndexFileMeta oldAnn =
+ annFile.build(
+ Collections.singletonList(
+ new PkVectorAnnSegmentFile.Source(
+ oldData, new ArrayReader(new float[][]
{{1, 0}}))),
+ vectorField,
+ options,
+ "l2",
+ "test-vector-ann");
+ AtomicInteger readerCount = new AtomicInteger();
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ when(readerFactory.create(newData))
+ .thenAnswer(
+ ignored -> {
+ if (readerCount.incrementAndGet() == 1) {
+ throw new IOException("Expected ANN build
failure.");
+ }
+ return reader(new float[][] {{2, 0}});
+ });
+ BucketedVectorIndexMaintainer maintainer =
+ new BucketedVectorIndexMaintainer(
+ 7,
+ annFile,
+ vectorField,
+ options,
+ "l2",
+ "test-vector-ann",
+ readerFactory,
+ Collections.singletonList(oldData),
+ Collections.singletonList(oldAnn),
+ executor);
+
+ assertThatThrownBy(
+ () ->
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
+ new CompactIncrement(
+
Collections.singletonList(oldData),
+
Collections.singletonList(newData),
+ Collections.emptyList()),
+ true))
+ .isInstanceOf(UncheckedIOException.class);
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit retry =
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
CompactIncrement.emptyIncrement(), true);
+
+ assertThat(retry.appendIncrement()).isEmpty();
+ assertThat(retry.compactIncrement()).isEmpty();
+ assertThat(maintainer.segments()).containsExactly(oldAnn);
+ }
+
+ @Test
+ void testRejectedBuildSubmissionDoesNotPoisonMaintainer() throws Exception
{
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO,
pathFactory());
+ DataField vectorField =
+ new DataField(7, "embedding", DataTypes.VECTOR(2,
DataTypes.FLOAT()));
+ DataFileMeta data = dataFile("data");
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+ PkVectorDataFileReader dataReader = reader(new float[][] {{1, 0}});
+ when(readerFactory.create(data)).thenReturn(dataReader);
+ ExecutorService rejectingExecutor =
Executors.newSingleThreadExecutor();
+ rejectingExecutor.shutdownNow();
+ BucketedVectorIndexMaintainer maintainer =
+ new BucketedVectorIndexMaintainer(
+ 7,
+ annFile,
+ vectorField,
+ indexOptions(),
+ "l2",
+ "test-vector-ann",
+ readerFactory,
+ Collections.emptyList(),
+ Collections.emptyList(),
+ rejectingExecutor);
+ CompactIncrement increment =
+ new CompactIncrement(
+ Collections.emptyList(),
+ Collections.singletonList(data),
+ Collections.emptyList());
+
+ Throwable failure =
+ catchThrowable(
+ () ->
+ maintainer.prepareCommit(
+ DataIncrement.emptyIncrement(),
increment, true));
+
+ assertThat(failure).isInstanceOf(RejectedExecutionException.class);
+ assertThat(failure.getSuppressed()).isEmpty();
+ assertThat(maintainer.buildNotCompleted()).isFalse();
+ maintainer.withExecutor(executor);
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit retry =
+ maintainer.prepareCommit(DataIncrement.emptyIncrement(),
increment, true);
+ assertThat(retry.compactIncrement()).isPresent();
+ assertThat(retry.compactIncrement().get().newIndexFiles()).hasSize(1);
+ }
+
private static Options indexOptions() {
Options options = new Options();
options.setString("test.vector.dimension", "2");
@@ -201,6 +675,13 @@ class BucketedVectorIndexMaintainerTest {
};
}
+ private long fileCount() throws IOException {
+ try (java.util.stream.Stream<java.nio.file.Path> files =
+ java.nio.file.Files.list(tempPath)) {
+ return files.count();
+ }
+ }
+
private static class ArrayReader implements PkVectorReader {
private final float[][] vectors;
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnLevelsTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnLevelsTest.java
new file mode 100644
index 0000000000..fb7eb22d3f
--- /dev/null
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnLevelsTest.java
@@ -0,0 +1,139 @@
+/*
+ * 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.index.pkvector;
+
+import org.apache.paimon.index.GlobalIndexMeta;
+import org.apache.paimon.index.IndexFileMeta;
+import org.apache.paimon.io.DataFileMeta;
+import org.apache.paimon.manifest.FileSource;
+import org.apache.paimon.stats.SimpleStats;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for derived primary-key vector ANN levels. */
+class PkVectorAnnLevelsTest {
+
+ @Test
+ void testPicksSimilarSegmentsAcrossAbsoluteBoundary() {
+ PkVectorAnnLevels levels = new PkVectorAnnLevels(3, 0.2);
+ DataFileMeta dataA = dataFile("data-a", 99);
+ DataFileMeta dataB = dataFile("data-b", 100);
+ DataFileMeta dataC = dataFile("data-c", 101);
+ IndexFileMeta annA = segment("ann-a", dataA);
+ IndexFileMeta annB = segment("ann-b", dataB);
+ IndexFileMeta annC = segment("ann-c", dataC);
+ Map<String, DataFileMeta> active = new LinkedHashMap<>();
+ active.put(dataA.fileName(), dataA);
+ active.put(dataB.fileName(), dataB);
+ active.put(dataC.fileName(), dataC);
+
+ PkVectorAnnLevels.Plan plan = levels.pick(Arrays.asList(annC, annA,
annB), active).get();
+
+ assertThat(plan.inputSegments()).containsExactly(annA, annB, annC);
+ }
+
+ @Test
+ void testPicksLevelZeroSegmentsAtFanout() {
+ PkVectorAnnLevels levels = new PkVectorAnnLevels(3, 0.2);
+ DataFileMeta dataA = dataFile("data-a", 30);
+ DataFileMeta dataB = dataFile("data-b", 40);
+ DataFileMeta dataC = dataFile("data-c", 50);
+ IndexFileMeta annA = segment("ann-a", dataA);
+ IndexFileMeta annB = segment("ann-b", dataB);
+ IndexFileMeta annC = segment("ann-c", dataC);
+ Map<String, DataFileMeta> active = new LinkedHashMap<>();
+ active.put(dataA.fileName(), dataA);
+ active.put(dataB.fileName(), dataB);
+ active.put(dataC.fileName(), dataC);
+
+ PkVectorAnnLevels.Plan plan = levels.pick(Arrays.asList(annC, annA,
annB), active).get();
+
+ assertThat(plan.inputSegments()).containsExactly(annA, annB, annC);
+ assertThat(plan.sourceFiles()).containsExactly(dataA, dataB, dataC);
+ }
+
+ @Test
+ void testPicksSegmentAboveStaleRatio() {
+ PkVectorAnnLevels levels = new PkVectorAnnLevels(5, 0.2);
+ DataFileMeta retired = dataFile("retired", 60);
+ DataFileMeta activeData = dataFile("active", 60);
+ IndexFileMeta ann = segment("ann", retired, activeData);
+ Map<String, DataFileMeta> active =
+ Collections.singletonMap(activeData.fileName(), activeData);
+
+ PkVectorAnnLevels.Plan plan =
levels.pick(Collections.singletonList(ann), active).get();
+
+ assertThat(plan.inputSegments()).containsExactly(ann);
+ assertThat(plan.sourceFiles()).containsExactly(activeData);
+ }
+
+ private static IndexFileMeta segment(String fileName, long rowCount) {
+ return segment(fileName, dataFile(fileName + "-data", rowCount));
+ }
+
+ private static IndexFileMeta segment(String fileName, DataFileMeta source)
{
+ return segment(fileName, new DataFileMeta[] {source});
+ }
+
+ private static IndexFileMeta segment(String fileName, DataFileMeta...
sources) {
+ long rowCount =
Arrays.stream(sources).mapToLong(DataFileMeta::rowCount).sum();
+ byte[] sourceMeta =
+ new PkVectorSourceMeta(
+ Arrays.stream(sources)
+ .map(
+ source ->
+ new PkVectorSourceFile(
+
source.fileName(),
+
source.rowCount()))
+
.collect(java.util.stream.Collectors.toList()))
+ .serialize();
+ return new IndexFileMeta(
+ "test-vector-ann",
+ fileName,
+ 100,
+ rowCount,
+ new GlobalIndexMeta(0, rowCount - 1, 7, null, new byte[] {1},
sourceMeta),
+ null);
+ }
+
+ private static DataFileMeta dataFile(String fileName, long rowCount) {
+ return DataFileMeta.forAppend(
+ fileName,
+ 100,
+ rowCount,
+ SimpleStats.EMPTY_STATS,
+ 0,
+ 0,
+ 1,
+ Collections.emptyList(),
+ null,
+ FileSource.COMPACT,
+ null,
+ null,
+ null,
+ null);
+ }
+}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java
index 0f7557c848..3178636075 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java
@@ -127,6 +127,47 @@ class PkVectorAnnSegmentFileTest {
org.assertj.core.groups.Tuple.tuple("data-1", 1L));
}
+ @Test
+ void testSearchFiltersInactiveSources() throws Exception {
+ LocalFileIO fileIO = LocalFileIO.create();
+ PkVectorAnnSegmentFile annFile = annFile(fileIO);
+ IndexFileMeta segment =
+ annFile.build(
+ Arrays.asList(
+ new PkVectorAnnSegmentFile.Source(
+ dataFile("retired", 2),
+ new ArrayReader(new float[][] {{0, 0},
{1, 0}})),
+ new PkVectorAnnSegmentFile.Source(
+ dataFile("active", 2),
+ new ArrayReader(new float[][] {{2, 0},
{3, 0}}))),
+ vectorField(),
+ indexOptions(),
+ "l2",
+ "test-vector-ann");
+
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ List<PkVectorSearchResult> candidates;
+ try {
+ candidates =
+ new PkVectorAnnSegmentSearcher(
+ fileIO, annFile, vectorField(),
indexOptions(), "l2", executor)
+ .search(
+ segment,
+ PkVectorSourceMeta.fromIndexFile(segment),
+ new float[] {0, 0},
+ 3,
+ Collections.emptyMap(),
+ Collections.singleton("active"),
+ Collections.emptyMap());
+ } finally {
+ executor.shutdownNow();
+ }
+
+ assertThat(candidates)
+ .extracting(PkVectorSearchResult::dataFileName)
+ .containsExactly("active", "active");
+ }
+
@Test
void testRejectsSourcesWithoutRows() {
LocalFileIO fileIO = LocalFileIO.create();
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
index 5e99f96deb..fbc40071f0 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java
@@ -18,6 +18,7 @@
package org.apache.paimon.index.pkvector;
+import org.apache.paimon.CoreOptions.GlobalIndexSearchMode;
import org.apache.paimon.deletionvectors.BitmapDeletionVector;
import org.apache.paimon.deletionvectors.DeletionVector;
import org.apache.paimon.index.GlobalIndexMeta;
@@ -45,11 +46,40 @@ import static org.mockito.Mockito.when;
/** Tests for {@link PrimaryKeyVectorBucketSearch}. */
class PrimaryKeyVectorBucketSearchTest {
+ @Test
+ void testFastModeSkipsExactFallback() throws Exception {
+ DataFileMeta data = dataFile("data");
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+
+ List<PkVectorSearchResult> results =
+ new PrimaryKeyVectorBucketSearch(
+ readerFactory,
+ null,
+ Collections.emptyMap(),
+ "l2",
+ GlobalIndexSearchMode.FAST)
+ .search(
+ new PkVectorBucketIndexState(
+ 7, "test-vector-ann",
Collections.emptyList()),
+ Collections.singletonList(data),
+ Collections.emptyMap(),
+ new float[] {0, 0},
+ 1);
+
+ assertThat(results).isEmpty();
+ verify(readerFactory, never()).create(data);
+ }
+
@Test
void testRejectsNonPositiveLimitForEmptyBucket() {
PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
PrimaryKeyVectorBucketSearch search =
- new PrimaryKeyVectorBucketSearch(readerFactory, null,
Collections.emptyMap(), "l2");
+ new PrimaryKeyVectorBucketSearch(
+ readerFactory,
+ null,
+ Collections.emptyMap(),
+ "l2",
+ GlobalIndexSearchMode.FULL);
assertThatIllegalArgumentException()
.isThrownBy(
@@ -85,11 +115,18 @@ class PrimaryKeyVectorBucketSearchTest {
org.mockito.ArgumentMatchers.any(float[].class),
org.mockito.ArgumentMatchers.eq(2),
org.mockito.ArgumentMatchers.eq(deletionVectors),
+ org.mockito.ArgumentMatchers.eq(
+ new
java.util.HashSet<>(Arrays.asList("data-1", "data-2"))),
org.mockito.ArgumentMatchers.eq(searchOptions)))
.thenReturn(Collections.singletonList(new
PkVectorSearchResult("data-1", 1, 0.5F)));
List<PkVectorSearchResult> results =
- new PrimaryKeyVectorBucketSearch(readerFactory, annSearcher,
searchOptions, "l2")
+ new PrimaryKeyVectorBucketSearch(
+ readerFactory,
+ annSearcher,
+ searchOptions,
+ "l2",
+ GlobalIndexSearchMode.FULL)
.search(
state,
Arrays.asList(data1, data2),
@@ -108,6 +145,45 @@ class PrimaryKeyVectorBucketSearchTest {
verify(readerFactory, never()).create(data1);
}
+ @Test
+ void testSearchesActivePartOfAnnWithInactiveSource() throws Exception {
+ DataFileMeta retired = dataFile("retired");
+ DataFileMeta active = dataFile("active");
+ IndexFileMeta ann = segment("ann", Arrays.asList(retired, active));
+ PkVectorAnnSegmentSearcher annSearcher =
mock(PkVectorAnnSegmentSearcher.class);
+ Map<String, DeletionVector> deletionVectors = Collections.emptyMap();
+ when(annSearcher.search(
+ org.mockito.ArgumentMatchers.eq(ann),
+
org.mockito.ArgumentMatchers.any(PkVectorSourceMeta.class),
+ org.mockito.ArgumentMatchers.any(float[].class),
+ org.mockito.ArgumentMatchers.eq(1),
+ org.mockito.ArgumentMatchers.eq(deletionVectors),
+
org.mockito.ArgumentMatchers.eq(Collections.singleton("active")),
+
org.mockito.ArgumentMatchers.eq(Collections.emptyMap())))
+ .thenReturn(Collections.singletonList(new
PkVectorSearchResult("active", 0, 1F)));
+ PkVectorDataFileReader.Factory readerFactory =
mock(PkVectorDataFileReader.Factory.class);
+
+ List<PkVectorSearchResult> results =
+ new PrimaryKeyVectorBucketSearch(
+ readerFactory,
+ annSearcher,
+ Collections.emptyMap(),
+ "l2",
+ GlobalIndexSearchMode.FULL)
+ .search(
+ new PkVectorBucketIndexState(
+ 7, "test-vector-ann",
Collections.singletonList(ann)),
+ Collections.singletonList(active),
+ deletionVectors,
+ new float[] {0, 0},
+ 1);
+
+ assertThat(results)
+ .extracting(PkVectorSearchResult::dataFileName)
+ .containsExactly("active");
+ verify(readerFactory, never()).create(active);
+ }
+
@Test
void testExactFallbackMergesFilesAndAppliesDeletionVectors() throws
Exception {
DataFileMeta data1 = dataFile("data-1");
@@ -123,7 +199,12 @@ class PrimaryKeyVectorBucketSearchTest {
deletionVectors.put("data-1", data1Deletes);
List<PkVectorSearchResult> results =
- new PrimaryKeyVectorBucketSearch(readerFactory, null,
Collections.emptyMap(), "l2")
+ new PrimaryKeyVectorBucketSearch(
+ readerFactory,
+ null,
+ Collections.emptyMap(),
+ "l2",
+ GlobalIndexSearchMode.DETAIL)
.search(
new PkVectorBucketIndexState(
7, "test-vector-ann",
Collections.emptyList()),
@@ -180,18 +261,27 @@ class PrimaryKeyVectorBucketSearchTest {
}
private static IndexFileMeta segment(String fileName, DataFileMeta source)
{
+ return segment(fileName, Collections.singletonList(source));
+ }
+
+ private static IndexFileMeta segment(String fileName, List<DataFileMeta>
sources) {
+ long rowCount =
sources.stream().mapToLong(DataFileMeta::rowCount).sum();
byte[] sourceMeta =
new PkVectorSourceMeta(
- Collections.singletonList(
- new PkVectorSourceFile(
- source.fileName(),
source.rowCount())))
+ sources.stream()
+ .map(
+ source ->
+ new PkVectorSourceFile(
+
source.fileName(),
+
source.rowCount()))
+
.collect(java.util.stream.Collectors.toList()))
.serialize();
return new IndexFileMeta(
"test-vector-ann",
fileName,
100,
- source.rowCount(),
- new GlobalIndexMeta(0, source.rowCount() - 1, 7, null, new
byte[] {1}, sourceMeta),
+ rowCount,
+ new GlobalIndexMeta(0, rowCount - 1, 7, null, new byte[] {1},
sourceMeta),
null);
}
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyVectorIndexWriteTest.java
b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyVectorIndexWriteTest.java
index d2feef4c25..994c51445f 100644
---
a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyVectorIndexWriteTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyVectorIndexWriteTest.java
@@ -45,6 +45,8 @@ import java.util.function.Function;
import static org.apache.paimon.data.BinaryRow.EMPTY_ROW;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
@@ -53,6 +55,56 @@ import static org.mockito.Mockito.when;
/** Tests primary-key vector index publication through {@link
AbstractFileStoreWrite}. */
class PrimaryKeyVectorIndexWriteTest {
+ @Test
+ void testRestoreRebindsVectorIndexExecutor() throws Exception {
+ DataIncrement dataIncrement = DataIncrement.emptyIncrement();
+ CompactIncrement compactIncrement = CompactIncrement.emptyIncrement();
+ RecordWriter<String> writer = mock(RecordWriter.class);
+ when(writer.prepareCommit(false))
+ .thenReturn(new CommitIncrement(dataIncrement,
compactIncrement, null));
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit vectorCommit =
+ mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class);
+ when(vectorCommit.appendIncrement()).thenReturn(Optional.empty());
+ when(vectorCommit.compactIncrement()).thenReturn(Optional.empty());
+ BucketedVectorIndexMaintainer maintainer =
mock(BucketedVectorIndexMaintainer.class);
+ when(maintainer.prepareCommit(same(dataIncrement),
same(compactIncrement), eq(true)))
+ .thenReturn(vectorCommit);
+
+ TestingFileStoreWrite write = new TestingFileStoreWrite();
+ write.install(writer, maintainer);
+ List<FileStoreWrite.State<String>> states = write.checkpoint();
+
+ TestingFileStoreWrite restored = new
TestingFileStoreWrite(mock(RecordWriter.class));
+ restored.restore(states);
+
+ verify(maintainer).withExecutor(any(ExecutorService.class));
+ }
+
+ @Test
+ void testCheckpointWaitsForPendingVectorIndexBuild() throws Exception {
+ DataIncrement dataIncrement = DataIncrement.emptyIncrement();
+ CompactIncrement compactIncrement = CompactIncrement.emptyIncrement();
+ RecordWriter<String> writer = mock(RecordWriter.class);
+ when(writer.prepareCommit(false))
+ .thenReturn(new CommitIncrement(dataIncrement,
compactIncrement, null));
+
+ BucketedVectorIndexMaintainer.VectorIndexCommit vectorCommit =
+ mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class);
+ when(vectorCommit.appendIncrement()).thenReturn(Optional.empty());
+ when(vectorCommit.compactIncrement()).thenReturn(Optional.empty());
+ BucketedVectorIndexMaintainer maintainer =
mock(BucketedVectorIndexMaintainer.class);
+ when(maintainer.prepareCommit(same(dataIncrement),
same(compactIncrement), eq(true)))
+ .thenReturn(vectorCommit);
+
+ TestingFileStoreWrite write = new TestingFileStoreWrite();
+ write.install(writer, maintainer);
+
+ write.checkpoint();
+
+ verify(maintainer).prepareCommit(same(dataIncrement),
same(compactIncrement), eq(true));
+ }
+
@Test
void testPublishesVectorChangesWithCompactDataTransition() throws
Exception {
DataIncrement dataIncrement = DataIncrement.emptyIncrement();
@@ -72,7 +124,7 @@ class PrimaryKeyVectorIndexWriteTest {
when(vectorCommit.appendIncrement()).thenReturn(Optional.empty());
when(vectorCommit.compactIncrement()).thenReturn(Optional.of(vectorIncrement));
BucketedVectorIndexMaintainer maintainer =
mock(BucketedVectorIndexMaintainer.class);
- when(maintainer.prepareCommit(same(dataIncrement),
same(compactIncrement)))
+ when(maintainer.prepareCommit(same(dataIncrement),
same(compactIncrement), eq(false)))
.thenReturn(vectorCommit);
TestingFileStoreWrite write = new TestingFileStoreWrite();
@@ -82,12 +134,18 @@ class PrimaryKeyVectorIndexWriteTest {
assertThat(message.compactIncrement().newIndexFiles()).containsExactly(added);
assertThat(message.compactIncrement().deletedIndexFiles()).containsExactly(deleted);
- verify(maintainer).prepareCommit(same(dataIncrement),
same(compactIncrement));
+ verify(maintainer).prepareCommit(same(dataIncrement),
same(compactIncrement), eq(false));
}
private static class TestingFileStoreWrite extends
AbstractFileStoreWrite<String> {
+ @Nullable private final RecordWriter<String> restoredWriter;
+
private TestingFileStoreWrite() {
+ this(null);
+ }
+
+ private TestingFileStoreWrite(@Nullable RecordWriter<String>
restoredWriter) {
super(
mock(SnapshotManager.class),
mock(FileStoreScan.class),
@@ -97,6 +155,7 @@ class PrimaryKeyVectorIndexWriteTest {
"test-table",
new CoreOptions(new HashMap<>()),
RowType.of());
+ this.restoredWriter = restoredWriter;
}
private void install(
@@ -120,7 +179,10 @@ class PrimaryKeyVectorIndexWriteTest {
ExecutorService compactExecutor,
@Nullable BucketedDvMaintainer deletionVectorsMaintainer,
boolean ignorePreviousFiles) {
- throw new UnsupportedOperationException();
+ if (restoredWriter == null) {
+ throw new UnsupportedOperationException();
+ }
+ return restoredWriter;
}
}
}
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 6f69b2a70c..8341d6bc4e 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
@@ -234,6 +234,26 @@ class PrimaryKeyVectorIndexValidationTest {
.hasMessageContaining("l2, cosine, inner_product");
}
+ @Test
+ void testRejectsInvalidAnnCompactionLevelFanout() {
+ Map<String, String> options = enabledOptions();
+ options.put(CoreOptions.PK_VECTOR_INDEX_COMPACTION_LEVEL_FANOUT.key(),
"1");
+
+ assertThatThrownBy(() -> validateTableSchema(schema(options)))
+
.hasMessageContaining("pk-vector.index.compaction.level-fanout")
+ .hasMessageContaining("greater than 1");
+ }
+
+ @Test
+ void testRejectsInvalidAnnCompactionStaleRatio() {
+ Map<String, String> options = enabledOptions();
+
options.put(CoreOptions.PK_VECTOR_INDEX_COMPACTION_STALE_RATIO_THRESHOLD.key(),
"1.1");
+
+ assertThatThrownBy(() -> validateTableSchema(schema(options)))
+
.hasMessageContaining("pk-vector.index.compaction.stale-ratio-threshold")
+ .hasMessageContaining("(0, 1]");
+ }
+
private static Map<String, String> enabledOptions() {
Map<String, String> options = new HashMap<>();
options.put(CoreOptions.BUCKET.key(), "1");
