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 3d463ba392 [spark][flink] Support primary-key vector search (#8581)
3d463ba392 is described below
commit 3d463ba392dabc28ce09676ac87144e80ddb74aa
Author: Jingsong Lee <[email protected]>
AuthorDate: Mon Jul 13 08:29:18 2026 +0800
[spark][flink] Support primary-key vector search (#8581)
Add end-to-end Spark SQL and Flink procedure support for primary-key
vector search backed by the bucket-local indexes introduced in #8579.
The PR supports every primary-key merge engine, distributes sufficiently
large Spark searches by bucket group, and documents table creation,
index maintenance, compaction freshness, query APIs, and current
limitations.
---
docs/docs/primary-key-table/vector-index.md | 210 ++++++++++++++++
docs/sidebars.js | 1 +
.../apache/paimon/reader/ScoreRecordReader.java | 31 +++
.../globalindex/IndexedSplitRecordReader.java | 3 +-
.../org/apache/paimon/schema/SchemaValidation.java | 7 +-
.../source/PrimaryKeyVectorPositionReader.java | 4 +-
.../paimon/table/source/PrimaryKeyVectorRead.java | 99 +++++---
.../table/source/VectorSearchBuilderImpl.java | 2 +-
.../PrimaryKeyVectorIndexValidationTest.java | 27 +++
.../table/source/PrimaryKeyVectorSearchTest.java | 71 +++++-
.../procedure/VectorSearchProcedureITCase.java | 119 +++++++++
...Impl.java => SparkDataEvolutionVectorRead.java} | 4 +-
.../spark/read/SparkPrimaryKeyVectorRead.java | 133 ++++++++++
.../spark/read/SparkVectorSearchBuilderImpl.java | 9 +-
.../paimon/spark/PaimonRecordReaderIterator.scala | 5 +-
.../apache/paimon/spark/PaimonScanBuilder.scala | 5 +
....java => SparkDataEvolutionVectorReadTest.java} | 10 +-
.../spark/sql/PrimaryKeyVectorSearchTest.scala | 270 +++++++++++++++++++++
.../paimon/spark/sql/VectorSearchOptionsTest.scala | 9 +
19 files changed, 966 insertions(+), 53 deletions(-)
diff --git a/docs/docs/primary-key-table/vector-index.md
b/docs/docs/primary-key-table/vector-index.md
new file mode 100644
index 0000000000..8ac875b29d
--- /dev/null
+++ b/docs/docs/primary-key-table/vector-index.md
@@ -0,0 +1,210 @@
+---
+title: "Vector Index"
+sidebar_position: 9
+---
+
+<!--
+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.
+-->
+
+# Vector Index
+
+Primary key tables can maintain a bucket-local approximate nearest neighbor
(ANN) index together
+with their data. Unlike a global vector index created by
`create_global_index`, a primary-key
+vector index is part of the normal write and compaction lifecycle. Paimon
builds it synchronously
+when complete compact-output files are produced and commits the index changes
together with those
+files.
+
+Use a primary-key vector index when vectors are frequently updated and the ANN
index should follow
+the primary-key table's compaction lifecycle. For append-only or Data
Evolution tables whose index
+is built separately from writes, see
+[Global Vector Index](../multimodal-table/global-index/vector).
+
+## Requirements
+
+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`).
+- `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`.
+- `pk-clustering-override` is disabled.
+- The configured vector index implementation is available on every writer and
reader classpath.
+
+The first release supports exactly one indexed vector column per table. The
option layout is
+field-scoped so that more independently indexed vector columns can be
supported in a future
+release.
+
+## Create Table
+
+The following Flink SQL example creates a three-dimensional vector column and
maintains an
+IVF-Flat index for it. Use the dimension produced by your embedding model in
production.
+
+```sql
+CREATE TABLE item_embeddings (
+ id BIGINT,
+ payload STRING,
+ embedding ARRAY<FLOAT> COMMENT '__VECTOR_FIELD;3',
+ PRIMARY KEY (id) NOT ENFORCED
+) WITH (
+ 'bucket' = '16',
+ 'deletion-vectors.enabled' = 'true',
+ 'pk-vector.index.columns' = 'embedding',
+ 'fields.embedding.pk-vector.index.type' = 'ivf-flat',
+ 'fields.embedding.pk-vector.distance.metric' = 'cosine',
+ 'fields.embedding.pk-vector.index.options' = '{"nlist":"256"}'
+);
+```
+
+Use the same properties in Spark SQL:
+
+```sql
+CREATE TABLE item_embeddings (
+ id BIGINT,
+ payload STRING,
+ embedding ARRAY<FLOAT> COMMENT '__VECTOR_FIELD;3'
+) USING paimon
+TBLPROPERTIES (
+ 'primary-key' = 'id',
+ 'bucket' = '16',
+ 'deletion-vectors.enabled' = 'true',
+ 'pk-vector.index.columns' = 'embedding',
+ 'fields.embedding.pk-vector.index.type' = 'ivf-flat',
+ 'fields.embedding.pk-vector.distance.metric' = 'cosine',
+ 'fields.embedding.pk-vector.index.options' = '{"nlist":"256"}'
+);
+```
+
+The vector comment directive converts the SQL `ARRAY<FLOAT>` column to
Paimon's fixed-length
+`VECTOR<FLOAT>` type. Java API users can define the column directly with
+`DataTypes.VECTOR(3, DataTypes.FLOAT())`.
+
+### Options
+
+| Option | Required | Description |
+|---|---|---|
+| `pk-vector.index.columns` | Yes | Indexed vector column. Exactly one column
is supported in the first release. |
+| `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. |
+
+For algorithm-specific build and search options, see
+[Vector Index](../multimodal-table/global-index/vector).
+
+## Index Maintenance
+
+Paimon builds immutable ANN segments from complete compact-output data files
inside each bucket.
+The index segment records the source data files and maps ANN ordinals back to
their physical row
+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.
+
+The maintenance behavior depends on the merge engine:
+
+- `deduplicate`: an update indexes the latest row and the deletion vector
hides the replaced
+ physical row. A delete removes the old row from search results through the
deletion vector.
+- `partial-update`: Paimon builds the vector index from the lookup-completed
Level-1
+ compact-output row.
+- `aggregation`: Paimon builds the vector index from the aggregated Level-1
compact-output row.
+- `first-row`: Paimon indexes the retained first row. Deletion vectors must be
disabled because
+ later rows with the same primary key are ignored rather than deleting the
retained row.
+
+When compaction replaces source data files, Paimon removes ANN segments that
reference those files
+and creates replacement segments for the new compact-output files. Small
outputs are indexed as
+well; there is no minimum-row threshold before a new segment can be built.
+
+The index follows compaction freshness. Newly appended level-0 files are not
ANN sources, so a
+streaming write may not be searchable until compaction has produced and
committed its complete
+level-1 output. Wait for that compaction when read-after-write vector-search
visibility is
+required. Batch writes which wait for compaction can publish the data and its
index together.
+
+## Search
+
+### Spark SQL
+
+Use the `vector_search` table-valued function. Spark exposes the ANN score
through the
+`__paimon_search_score` metadata column.
+
+```sql
+SELECT id, payload, __paimon_search_score
+FROM vector_search(
+ 'item_embeddings',
+ 'embedding',
+ array(0.1f, 0.2f, 0.3f),
+ 10,
+ map('ivf.nprobe', '32')
+);
+```
+
+The query vector dimension must match the indexed column dimension. For
partitioned tables, Spark
+applies a partition predicate before running ANN and merging the global Top-K.
+When `spark.paimon.vector-search.distribute.enabled` is `true`, Spark
distributes sufficiently
+large groups of bucket-local ANN searches across executors and merges their
task-local Top-K
+results on the driver. Small plans stay local to avoid Spark job startup
overhead.
+
+### Flink SQL
+
+Flink exposes vector search as a procedure and returns JSON-serialized rows.
Use `projection` to
+avoid reading columns that are not needed.
+
+```sql
+CALL sys.vector_search(
+ `table` => 'default.item_embeddings',
+ vector_column => 'embedding',
+ query_vector => '0.1,0.2,0.3',
+ top_k => 10,
+ projection => 'id,payload',
+ options => 'ivf.nprobe=32'
+);
+```
+
+### Java API
+
+```java
+GlobalIndexResult result = table.newVectorSearchBuilder()
+ .withVectorColumn("embedding")
+ .withVector(queryVector)
+ .withLimit(10)
+ .withOption("ivf.nprobe", "32")
+ .executeLocal();
+
+ReadBuilder readBuilder = table.newReadBuilder();
+TableScan.Plan plan =
readBuilder.newScan().withGlobalIndexResult(result).plan();
+try (RecordReader<InternalRow> reader =
readBuilder.newRead().createReader(plan)) {
+ reader.forEachRemaining(row -> consume(row));
+}
+```
+
+## Query Planning
+
+A search captures one table snapshot, plans the active ANN segments for every
selected bucket,
+searches those segments, and merges their candidates into one global Top-K.
The returned candidates
+are materialized from the source data files by physical row position. Deletion
vectors are applied
+while searching and reading, so stale versions and deleted rows are not
returned.
+
+For low latency on object storage, cache data files and ANN payloads with a
caching file system.
+The first query may still need to download index files; subsequent queries can
search the local
+cached payloads and fetch only the selected data-file positions.
+
+## Limitations
+
+- Exactly one vector index column is supported per table in the first release.
+- Only `FLOAT` vectors are supported.
+- Dynamic-bucket and `pk-clustering-override` tables are not supported.
+- Flink's procedure returns rows but does not expose the ANN score as a
separate column.
+- Vector search is snapshot-scoped batch reading; streaming search and lateral
vector search for
+ primary-key tables are not supported.
diff --git a/docs/sidebars.js b/docs/sidebars.js
index 4223c1347f..7e3967dade 100644
--- a/docs/sidebars.js
+++ b/docs/sidebars.js
@@ -84,6 +84,7 @@ const sidebars = {
"primary-key-table/sequence-rowkind",
"primary-key-table/compaction",
"primary-key-table/query-performance",
+ "primary-key-table/vector-index",
"primary-key-table/chain-table",
"primary-key-table/pk-clustering-override",
{
diff --git
a/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordReader.java
b/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordReader.java
new file mode 100644
index 0000000000..69b9463325
--- /dev/null
+++
b/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordReader.java
@@ -0,0 +1,31 @@
+/*
+ * 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.reader;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+
+/** A {@link RecordReader} whose records expose vector-search scores and row
identifiers. */
+public interface ScoreRecordReader<T> extends RecordReader<T> {
+
+ @Nullable
+ @Override
+ ScoreRecordIterator<T> readBatch() throws IOException;
+}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplitRecordReader.java
b/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplitRecordReader.java
index 3778f37cb9..3d5bc89e00 100644
---
a/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplitRecordReader.java
+++
b/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplitRecordReader.java
@@ -21,6 +21,7 @@ package org.apache.paimon.globalindex;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.reader.RecordReader;
import org.apache.paimon.reader.ScoreRecordIterator;
+import org.apache.paimon.reader.ScoreRecordReader;
import org.apache.paimon.table.SpecialFields;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.ProjectedRow;
@@ -35,7 +36,7 @@ import java.util.Map;
import static org.apache.paimon.utils.Preconditions.checkArgument;
/** Return value with score. */
-public class IndexedSplitRecordReader implements RecordReader<InternalRow> {
+public class IndexedSplitRecordReader implements
ScoreRecordReader<InternalRow> {
private final RecordReader<InternalRow> reader;
@Nullable private final Map<Long, Float> rowIdToScore;
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 2122c04ed5..09cb61c0bf 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
@@ -912,13 +912,8 @@ public class SchemaValidation {
!schema.primaryKeys().isEmpty(),
"Primary-key vector index requires a primary-key table.");
checkArgument(
- options.deletionVectorsEnabled(),
+ options.mergeEngine() == MergeEngine.FIRST_ROW ||
options.deletionVectorsEnabled(),
"Primary-key vector index requires deletion-vectors.enabled =
true.");
- checkArgument(
- options.mergeEngine() == MergeEngine.DEDUPLICATE
- || options.mergeEngine() == MergeEngine.PARTIAL_UPDATE,
- "Primary-key vector index only supports merge-engine =
deduplicate or partial-update, but is %s.",
- options.mergeEngine());
checkArgument(
!options.deletionVectorsMergeOnRead(),
"Primary-key vector index with merge-engine = %s requires
deletion-vectors.merge-on-read = false.",
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java
index 6470fc26eb..52f7423a92 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorPositionReader.java
@@ -21,8 +21,8 @@ package org.apache.paimon.table.source;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.reader.FileRecordIterator;
import org.apache.paimon.reader.FileRecordReader;
-import org.apache.paimon.reader.RecordReader;
import org.apache.paimon.reader.ScoreRecordIterator;
+import org.apache.paimon.reader.ScoreRecordReader;
import org.apache.paimon.utils.RoaringBitmap32;
import javax.annotation.Nullable;
@@ -33,7 +33,7 @@ import java.util.function.IntFunction;
import static org.apache.paimon.utils.Preconditions.checkArgument;
/** Reads selected physical file positions and exposes their vector-search
scores. */
-public class PrimaryKeyVectorPositionReader implements
RecordReader<InternalRow> {
+public class PrimaryKeyVectorPositionReader implements
ScoreRecordReader<InternalRow> {
private final FileRecordReader<InternalRow> reader;
private final RoaringBitmap32 rowPositions;
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 68350ffb83..b70f19eb73 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
@@ -40,6 +40,7 @@ import org.apache.paimon.types.DataField;
import org.apache.paimon.types.VectorType;
import java.io.IOException;
+import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
@@ -56,7 +57,9 @@ import static
org.apache.paimon.globalindex.VectorSearchMetric.normalize;
import static org.apache.paimon.utils.Preconditions.checkArgument;
/** Executes bucket-local primary-key vector search and merges physical
candidates globally. */
-public class PrimaryKeyVectorRead implements VectorRead {
+public class PrimaryKeyVectorRead implements VectorRead, Serializable {
+
+ private static final long serialVersionUID = 1L;
private static final Comparator<Candidate> BEST_FIRST =
(left, right) -> {
@@ -76,16 +79,13 @@ public class PrimaryKeyVectorRead implements VectorRead {
return fileName != 0 ? fileName :
Long.compare(left.rowPosition, right.rowPosition);
};
- private final FileIO fileIO;
- private final IndexFileHandler indexFileHandler;
- private final KeyValueFileReaderFactory.Builder readerFactoryBuilder;
- private final DataField vectorField;
+ protected final FileStoreTable table;
+ protected final DataField vectorField;
private final String indexType;
private final Options indexOptions;
- private final ExecutorService executor;
private final Map<String, String> searchOptions;
private final float[] query;
- private final int limit;
+ protected final int limit;
private final String metric;
public PrimaryKeyVectorRead(
@@ -102,15 +102,10 @@ public class PrimaryKeyVectorRead implements VectorRead {
query.length,
((VectorType) vectorField.type()).getLength());
checkArgument(limit > 0, "Vector search limit must be positive: %s.",
limit);
- this.fileIO = table.fileIO();
- this.indexFileHandler = table.store().newIndexFileHandler();
- this.readerFactoryBuilder =
keyValueStore(table).newReaderFactoryBuilder();
+ this.table = table;
this.vectorField = vectorField;
this.indexType =
table.coreOptions().primaryKeyVectorIndexType(vectorField.name());
this.indexOptions =
table.coreOptions().primaryKeyVectorIndexOptions(vectorField.name());
- this.executor =
- GlobalIndexReadThreadPool.getExecutorService(
-
table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM));
this.searchOptions = Collections.unmodifiableMap(new
HashMap<>(searchOptions));
this.query = query.clone();
this.limit = limit;
@@ -131,26 +126,52 @@ public class PrimaryKeyVectorRead implements VectorRead {
@Override
public GlobalIndexResult read(VectorScan.Plan plan) {
+ PrimaryKeyVectorScan.Plan primaryKeyPlan = primaryKeyPlan(plan);
+ return createResult(primaryKeyPlan,
searchBuckets(bucketSplits(primaryKeyPlan)));
+ }
+
+ protected PrimaryKeyVectorScan.Plan primaryKeyPlan(VectorScan.Plan plan) {
checkArgument(
plan instanceof PrimaryKeyVectorScan.Plan,
"Primary-key vector read requires a PrimaryKeyVectorScan
plan.");
PrimaryKeyVectorScan.Plan primaryKeyPlan = (PrimaryKeyVectorScan.Plan)
plan;
+ for (VectorSearchSplit searchSplit : primaryKeyPlan.splits()) {
+ BucketVectorSearchSplit split = (BucketVectorSearchSplit)
searchSplit;
+ checkArgument(
+ split.dataSplit().snapshotId() ==
primaryKeyPlan.snapshotId(),
+ "Vector bucket split snapshot does not match its plan.");
+ }
+ return primaryKeyPlan;
+ }
+
+ protected List<BucketVectorSearchSplit>
bucketSplits(PrimaryKeyVectorScan.Plan plan) {
+ List<BucketVectorSearchSplit> splits = new
ArrayList<>(plan.splits().size());
+ for (VectorSearchSplit split : plan.splits()) {
+ splits.add((BucketVectorSearchSplit) split);
+ }
+ return splits;
+ }
+
+ protected List<Candidate> searchBuckets(List<BucketVectorSearchSplit>
splits) {
try {
+ SearchContext context = new SearchContext(table);
List<Candidate> candidates = new ArrayList<>();
- for (VectorSearchSplit searchSplit : primaryKeyPlan.splits()) {
- BucketVectorSearchSplit split = (BucketVectorSearchSplit)
searchSplit;
- checkArgument(
- split.dataSplit().snapshotId() ==
primaryKeyPlan.snapshotId(),
- "Vector bucket split snapshot does not match its
plan.");
- candidates.addAll(search(split));
+ for (BucketVectorSearchSplit split : splits) {
+ candidates.addAll(search(split, context));
}
- return new PrimaryKeyVectorResult(primaryKeyPlan, topK(candidates,
limit), metric);
+ return topK(candidates, limit);
} catch (IOException e) {
throw new RuntimeException("Failed to search primary-key vector
index.", e);
}
}
- private List<Candidate> search(BucketVectorSearchSplit split) throws
IOException {
+ protected GlobalIndexResult createResult(
+ PrimaryKeyVectorScan.Plan plan, List<Candidate> candidates) {
+ return new PrimaryKeyVectorResult(plan, topK(candidates, limit),
metric);
+ }
+
+ private List<Candidate> search(BucketVectorSearchSplit split,
SearchContext context)
+ throws IOException {
DataSplit dataSplit = split.dataSplit();
List<DataFileMeta> activeFiles =
dataSplit.dataFiles().stream()
@@ -159,23 +180,23 @@ public class PrimaryKeyVectorRead implements VectorRead {
PkVectorBucketIndexState state =
PkVectorBucketIndexState.fromActivePayloads(
vectorField.id(), indexType, split.payloadFiles());
- Map<String, DeletionVector> deletionVectors =
deletionVectors(dataSplit);
+ Map<String, DeletionVector> deletionVectors =
deletionVectors(dataSplit, context.fileIO);
PkVectorDataFileReader.Factory readerFactory =
new PkVectorDataFileReader.Factory(
- readerFactoryBuilder,
+ context.readerFactoryBuilder,
dataSplit.partition(),
dataSplit.bucket(),
vectorField,
((VectorType) vectorField.type()).getLength());
PkVectorAnnSegmentSearcher annSearcher =
new PkVectorAnnSegmentSearcher(
- fileIO,
- indexFileHandler.pkVectorAnnSegment(
+ context.fileIO,
+ context.indexFileHandler.pkVectorAnnSegment(
dataSplit.partition(), dataSplit.bucket()),
vectorField,
indexOptions,
metric,
- executor);
+ context.executor);
PrimaryKeyVectorBucketSearch bucketSearch =
new PrimaryKeyVectorBucketSearch(readerFactory, annSearcher,
searchOptions, metric);
List<Candidate> candidates = new ArrayList<>();
@@ -192,7 +213,8 @@ public class PrimaryKeyVectorRead implements VectorRead {
return candidates;
}
- private Map<String, DeletionVector> deletionVectors(DataSplit split)
throws IOException {
+ private Map<String, DeletionVector> deletionVectors(DataSplit split,
FileIO fileIO)
+ throws IOException {
DeletionVector.Factory factory =
DeletionVector.factory(
fileIO, split.dataFiles(),
split.deletionFiles().orElse(null));
@@ -206,7 +228,7 @@ public class PrimaryKeyVectorRead implements VectorRead {
return result;
}
- static List<Candidate> topK(List<Candidate> candidates, int limit) {
+ protected static List<Candidate> topK(List<Candidate> candidates, int
limit) {
checkArgument(limit > 0, "Vector search limit must be positive: %s.",
limit);
PriorityQueue<Candidate> nearest = new PriorityQueue<>(limit,
BEST_FIRST.reversed());
for (Candidate candidate : candidates) {
@@ -234,7 +256,9 @@ public class PrimaryKeyVectorRead implements VectorRead {
}
/** Snapshot-scoped physical row candidate. */
- public static class Candidate {
+ public static class Candidate implements Serializable {
+
+ private static final long serialVersionUID = 1L;
private final BinaryRow partition;
private final int bucket;
@@ -275,4 +299,21 @@ public class PrimaryKeyVectorRead implements VectorRead {
return distance;
}
}
+
+ private static class SearchContext {
+
+ private final FileIO fileIO;
+ private final IndexFileHandler indexFileHandler;
+ private final KeyValueFileReaderFactory.Builder readerFactoryBuilder;
+ private final ExecutorService executor;
+
+ private SearchContext(FileStoreTable table) {
+ this.fileIO = table.fileIO();
+ this.indexFileHandler = table.store().newIndexFileHandler();
+ this.readerFactoryBuilder =
keyValueStore(table).newReaderFactoryBuilder();
+ this.executor =
+ GlobalIndexReadThreadPool.getExecutorService(
+
table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM));
+ }
+ }
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java
index f2bcb7af65..e706e40ba6 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java
@@ -151,7 +151,7 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
table, partitionFilter, filter, limit, vectorColumn, vector,
options);
}
- private boolean isPrimaryKeyVectorSearch() {
+ protected boolean isPrimaryKeyVectorSearch() {
return vectorColumn != null
&&
table.coreOptions().primaryKeyVectorIndexColumns().contains(vectorColumn.name());
}
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 0daa7bc5c4..6f69b2a70c 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
@@ -123,6 +123,33 @@ class PrimaryKeyVectorIndexValidationTest {
assertThatCode(() ->
validateTableSchema(schema(options))).doesNotThrowAnyException();
}
+ @Test
+ void testSupportsAggregationMergeEngine() {
+ Map<String, String> options = enabledOptions();
+ options.put(CoreOptions.MERGE_ENGINE.key(), "aggregation");
+
+ assertThatCode(() ->
validateTableSchema(schema(options))).doesNotThrowAnyException();
+ }
+
+ @Test
+ void testSupportsFirstRowWithoutDeletionVectors() {
+ Map<String, String> options = enabledOptions();
+ options.put(CoreOptions.MERGE_ENGINE.key(), "first-row");
+ options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "false");
+
+ assertThatCode(() ->
validateTableSchema(schema(options))).doesNotThrowAnyException();
+ }
+
+ @Test
+ void testRejectsFirstRowWithDeletionVectors() {
+ Map<String, String> options = enabledOptions();
+ options.put(CoreOptions.MERGE_ENGINE.key(), "first-row");
+
+ assertThatThrownBy(() -> validateTableSchema(schema(options)))
+ .hasMessageContaining(
+ "First row merge engine does not need deletion vectors
because there is no deletion of old data in this merge engine");
+ }
+
@Test
void testPartialUpdateRejectsDeletionVectorMergeOnRead() {
Map<String, String> options = enabledOptions();
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
index 090488de9e..5c6558fdfe 100644
---
a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
@@ -46,12 +46,19 @@ class PrimaryKeyVectorSearchTest extends TableTestBase {
@Override
protected Schema schemaDefault() {
+ return vectorSchema("deduplicate", true);
+ }
+
+ private Schema vectorSchema(String mergeEngine, boolean
deletionVectorsEnabled) {
return Schema.newBuilder()
.column("id", DataTypes.INT())
.column("embedding", DataTypes.VECTOR(2, DataTypes.FLOAT()))
.primaryKey("id")
.option(CoreOptions.BUCKET.key(), "1")
- .option(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true")
+ .option(CoreOptions.MERGE_ENGINE.key(), mergeEngine)
+ .option(
+ CoreOptions.DELETION_VECTORS_ENABLED.key(),
+ Boolean.toString(deletionVectorsEnabled))
.option(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding")
.option(
"fields.embedding.pk-vector.index.type",
@@ -94,4 +101,66 @@ class PrimaryKeyVectorSearchTest extends TableTestBase {
assertThat(ids).containsExactly(2, 3);
}
+
+ @Test
+ void testFirstRowVectorSearch() throws Exception {
+ catalog.createTable(identifier(), vectorSchema("first-row", false),
false);
+ FileStoreTable table = getTableDefault();
+
+ write(
+ table,
+ ioManager,
+ GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[]
{3, 0})),
+ GenericRow.of(2, BinaryVector.fromPrimitiveArray(new float[]
{1, 0})));
+ write(
+ table,
+ ioManager,
+ GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[]
{0.5f, 0})));
+
+ GlobalIndexResult result =
+ table.newVectorSearchBuilder()
+ .withVectorColumn("embedding")
+ .withVector(new float[] {0, 0})
+ .withLimit(1)
+ .executeLocal();
+ ReadBuilder readBuilder = table.newReadBuilder();
+ TableScan.Plan plan =
readBuilder.newScan().withGlobalIndexResult(result).plan();
+ List<Integer> ids = new ArrayList<>();
+ try (RecordReader<InternalRow> reader =
readBuilder.newRead().createReader(plan)) {
+ reader.forEachRemaining(row -> ids.add(row.getInt(0)));
+ }
+
+ assertThat(ids).containsExactly(2);
+ }
+
+ @Test
+ void testAggregationVectorSearch() throws Exception {
+ catalog.createTable(identifier(), vectorSchema("aggregation", true),
false);
+ FileStoreTable table = getTableDefault();
+
+ write(
+ table,
+ ioManager,
+ GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[]
{3, 0})),
+ GenericRow.of(2, BinaryVector.fromPrimitiveArray(new float[]
{1, 0})));
+ write(
+ table,
+ ioManager,
+ GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[]
{0.5f, 0})));
+
+ GlobalIndexResult result =
+ table.newVectorSearchBuilder()
+ .withVectorColumn("embedding")
+ .withVector(new float[] {0, 0})
+ .withLimit(1)
+ .executeLocal();
+ ReadBuilder readBuilder = table.newReadBuilder();
+ TableScan.Plan plan =
readBuilder.newScan().withGlobalIndexResult(result).plan();
+ List<Integer> ids = new ArrayList<>();
+ try (RecordReader<InternalRow> reader =
readBuilder.newRead().createReader(plan)) {
+ reader.forEachRemaining(row -> ids.add(row.getInt(0)));
+ }
+
+ assertThat(ids).containsExactly(1);
+ }
}
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 d71de22d2a..bc29716356 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
@@ -54,6 +54,67 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
private static final String VECTOR_FIELD = "vec";
private static final int DIMENSION = 2;
+ @Test
+ public void testPrimaryKeyVectorSearch() throws Exception {
+ createPrimaryKeyVectorTable("PK_T");
+
+ sql(
+ "INSERT INTO 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)])");
+
+ List<Row> result = searchPrimaryKeyVectorTable("PK_T", 2, "id");
+
+ assertThat(result)
+ .extracting(row -> row.getField(0).toString())
+ .containsExactlyInAnyOrder("{\"id\":\"2\"}", "{\"id\":\"3\"}");
+ }
+
+ @Test
+ public void testPrimaryKeyVectorSearchAfterUpdateAndDelete() throws
Exception {
+ createPrimaryKeyVectorTable("PK_UPDATE_T");
+
+ sql(
+ "INSERT INTO PK_UPDATE_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)])");
+ sql(
+ "INSERT INTO PK_UPDATE_T VALUES "
+ + "(1, ARRAY[CAST(0.5 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+
+ List<Row> updated = searchPrimaryKeyVectorTable("PK_UPDATE_T", 1,
"id");
+ assertThat(updated)
+ .extracting(row -> row.getField(0).toString())
+ .containsExactly("{\"id\":\"1\"}");
+
+ sql("DELETE FROM PK_UPDATE_T WHERE id = 1");
+
+ List<Row> afterDelete = searchPrimaryKeyVectorTable("PK_UPDATE_T", 1,
"id");
+ assertThat(afterDelete)
+ .extracting(row -> row.getField(0).toString())
+ .containsExactly("{\"id\":\"2\"}");
+ }
+
+ @Test
+ public void testPartialUpdatePrimaryKeyVectorSearch() throws Exception {
+ createPartialUpdatePrimaryKeyVectorTable("PK_PARTIAL_T");
+
+ sql(
+ "INSERT INTO PK_PARTIAL_T VALUES "
+ + "(1, 'keep', ARRAY[CAST(3.0 AS FLOAT), CAST(0.0 AS
FLOAT)]), "
+ + "(2, 'other', ARRAY[CAST(1.0 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+ sql(
+ "INSERT INTO PK_PARTIAL_T (id, vec) VALUES "
+ + "(1, ARRAY[CAST(0.5 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+
+ List<Row> result = searchPrimaryKeyVectorTable("PK_PARTIAL_T", 1,
"id,payload");
+
+ assertThat(result)
+ .extracting(row -> row.getField(0).toString())
+ .containsExactly("{\"id\":\"1\",\"payload\":\"keep\"}");
+ }
+
@Test
public void testVectorSearchBasic() throws Exception {
createVectorTable("T");
@@ -202,6 +263,64 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
tableName, DIMENSION, formattedExtraOptions);
}
+ private void createPrimaryKeyVectorTable(String tableName) {
+ sql(
+ "CREATE TABLE %s ("
+ + "id INT, "
+ + "vec ARRAY<FLOAT>, "
+ + "PRIMARY KEY (id) NOT ENFORCED"
+ + ") WITH ("
+ + "'bucket' = '2', "
+ + "'file.format' = 'json', "
+ + "'file.compression' = 'none', "
+ + "'deletion-vectors.enabled' = 'true', "
+ + "'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("
+ + "`table` => 'default.%s', "
+ + "vector_column => 'vec', "
+ + "query_vector => '0.0,0.0', "
+ + "top_k => %d, "
+ + "projection => '%s')",
+ tableName, topK, projection);
+ }
+
+ private void createPartialUpdatePrimaryKeyVectorTable(String tableName) {
+ sql(
+ "CREATE TABLE %s ("
+ + "id INT, "
+ + "payload STRING, "
+ + "vec ARRAY<FLOAT>, "
+ + "PRIMARY KEY (id) NOT ENFORCED"
+ + ") WITH ("
+ + "'bucket' = '1', "
+ + "'file.format' = 'json', "
+ + "'file.compression' = 'none', "
+ + "'merge-engine' = 'partial-update', "
+ + "'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 void writeVectors(FileStoreTable table, float[][] vectors) throws
Exception {
BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder();
try (BatchTableWrite write = writeBuilder.newWrite();
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorRead.java
similarity index 99%
rename from
paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
rename to
paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorRead.java
index 3c13fb4c0c..ddac704b66 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorRead.java
@@ -54,11 +54,11 @@ import static
org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM;
* Spark-aware {@link DataEvolutionVectorRead} that distributes grouped vector
index evaluation
* across the Spark cluster instead of evaluating them with the local thread
pool.
*/
-public class SparkVectorReadImpl extends DataEvolutionVectorRead {
+public class SparkDataEvolutionVectorRead extends DataEvolutionVectorRead {
private static final long serialVersionUID = 1L;
- public SparkVectorReadImpl(
+ public SparkDataEvolutionVectorRead(
FileStoreTable table,
@Nullable PartitionPredicate partitionFilter,
@Nullable Predicate filter,
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkPrimaryKeyVectorRead.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkPrimaryKeyVectorRead.java
new file mode 100644
index 0000000000..053abc9677
--- /dev/null
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkPrimaryKeyVectorRead.java
@@ -0,0 +1,133 @@
+/*
+ * 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.spark.read;
+
+import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.source.BucketVectorSearchSplit;
+import org.apache.paimon.table.source.PrimaryKeyVectorRead;
+import org.apache.paimon.table.source.PrimaryKeyVectorScan;
+import org.apache.paimon.table.source.VectorScan;
+import org.apache.paimon.types.DataField;
+import org.apache.paimon.utils.InstantiationUtil;
+import org.apache.paimon.utils.SerializableFunction;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM;
+
+/** Spark-aware {@link PrimaryKeyVectorRead}. */
+public class SparkPrimaryKeyVectorRead extends PrimaryKeyVectorRead {
+
+ private static final long serialVersionUID = 1L;
+
+ public SparkPrimaryKeyVectorRead(
+ FileStoreTable table,
+ DataField vectorField,
+ float[] query,
+ int limit,
+ Map<String, String> searchOptions) {
+ super(table, vectorField, query, limit, searchOptions);
+ }
+
+ @Override
+ public GlobalIndexResult read(VectorScan.Plan plan) {
+ PrimaryKeyVectorScan.Plan primaryKeyPlan = primaryKeyPlan(plan);
+ List<BucketVectorSearchSplit> splits = bucketSplits(primaryKeyPlan);
+ int parallelism = sparkParallelism();
+ if (splits.size() < parallelism * 2) {
+ return super.read(plan);
+ }
+
+ List<byte[]> serializedSplits = new ArrayList<>(splits.size());
+ for (BucketVectorSearchSplit split : splits) {
+ try {
+ serializedSplits.add(InstantiationUtil.serializeObject(split));
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to serialize primary-key
vector split.", e);
+ }
+ }
+ List<List<byte[]>> groups = splitGroups(serializedSplits, parallelism);
+ SerializableFunction<List<byte[]>, byte[]> task =
+ group -> {
+ List<BucketVectorSearchSplit> taskSplits = new
ArrayList<>(group.size());
+ for (byte[] bytes : group) {
+ taskSplits.add(deserializeSplit(bytes));
+ }
+ try {
+ return
InstantiationUtil.serializeObject(searchBuckets(taskSplits));
+ } catch (IOException e) {
+ throw new RuntimeException(
+ "Failed to serialize primary-key vector
candidates.", e);
+ }
+ };
+ List<byte[]> groupResults = mapInSpark(groups, task, groups.size());
+ List<Candidate> candidates = new ArrayList<>();
+ for (byte[] groupResult : groupResults) {
+ candidates.addAll(deserializeCandidates(groupResult));
+ }
+ return createResult(primaryKeyPlan, candidates);
+ }
+
+ protected int sparkParallelism() {
+ return Math.max(1,
table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM));
+ }
+
+ protected SparkEngineContext createEngineContext() {
+ return new SparkEngineContext();
+ }
+
+ protected <I, O> List<O> mapInSpark(
+ List<I> data, SerializableFunction<I, O> function, int
parallelism) {
+ return createEngineContext().map(data, function, parallelism);
+ }
+
+ private List<List<byte[]>> splitGroups(List<byte[]> splits, int
parallelism) {
+ List<List<byte[]>> groups = new ArrayList<>(parallelism);
+ int groupSize = (splits.size() + parallelism - 1) / parallelism;
+ for (int start = 0; start < splits.size(); start += groupSize) {
+ groups.add(
+ new ArrayList<>(
+ splits.subList(start, Math.min(start + groupSize,
splits.size()))));
+ }
+ return groups;
+ }
+
+ private BucketVectorSearchSplit deserializeSplit(byte[] bytes) {
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException("Failed to deserialize primary-key
vector split.", e);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private List<Candidate> deserializeCandidates(byte[] bytes) {
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException("Failed to deserialize primary-key
vector candidates.", e);
+ }
+ }
+}
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
index 8704486258..7638f57269 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
@@ -23,8 +23,8 @@ import org.apache.paimon.table.source.VectorRead;
import org.apache.paimon.table.source.VectorSearchBuilderImpl;
/**
- * Spark-aware {@link VectorSearchBuilderImpl} which produces a {@link
SparkVectorReadImpl} so the
- * per-split vector index evaluation is dispatched through Spark instead of
the local thread pool.
+ * Spark-aware {@link VectorSearchBuilderImpl} which produces Spark-specific
vector readers so
+ * data-evolution splits and primary-key bucket groups can be evaluated across
the Spark cluster.
*
* <p>Single-vector only; batch search has no Spark-dispatched path yet (TODO).
*/
@@ -38,7 +38,10 @@ public class SparkVectorSearchBuilderImpl extends
VectorSearchBuilderImpl {
@Override
public VectorRead newVectorRead() {
- return new SparkVectorReadImpl(
+ if (isPrimaryKeyVectorSearch()) {
+ return new SparkPrimaryKeyVectorRead(table, vectorColumn, vector,
limit, options);
+ }
+ return new SparkDataEvolutionVectorRead(
table, partitionFilter, filter, limit, vectorColumn, vector,
options);
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
index 64b0b5166c..93256df2c0 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
@@ -20,8 +20,7 @@ package org.apache.paimon.spark
import org.apache.paimon.data.{BinaryString, GenericRow, InternalRow =>
PaimonInternalRow, JoinedRow}
import org.apache.paimon.fs.Path
-import org.apache.paimon.globalindex.IndexedSplitRecordReader
-import org.apache.paimon.reader.{FileRecordIterator, RecordReader,
ScoreRecordIterator}
+import org.apache.paimon.reader.{FileRecordIterator, RecordReader,
ScoreRecordIterator, ScoreRecordReader}
import org.apache.paimon.spark.schema.PaimonMetadataColumn
import
org.apache.paimon.spark.schema.PaimonMetadataColumn.{PARTITION_AND_BUCKET_META_COLUMNS,
PATH_AND_INDEX_META_COLUMNS, VECTOR_SEARCH_META_COLUMN_NAMES}
import org.apache.paimon.table.source.{DataSplit, Split}
@@ -49,7 +48,7 @@ case class PaimonRecordReaderIterator(
private val needMetadata = metadataColumns.nonEmpty
private val needPathAndIndexMetadata =
metadataColumns.exists(c => PATH_AND_INDEX_META_COLUMNS.contains(c.name))
- private val needVectorSearchMetadata =
reader.isInstanceOf[IndexedSplitRecordReader] &&
+ private val needVectorSearchMetadata =
reader.isInstanceOf[ScoreRecordReader[_]] &&
metadataColumns.exists(c =>
VECTOR_SEARCH_META_COLUMN_NAMES.contains(c.name))
Preconditions.checkArgument(
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
index 6d7f894bea..844c4de52b 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
@@ -18,6 +18,7 @@
package org.apache.paimon.spark
+import org.apache.paimon.CoreOptions
import org.apache.paimon.partition.PartitionPredicate
import org.apache.paimon.predicate._
import org.apache.paimon.predicate.SortValue.{NullOrdering, SortDirection}
@@ -141,6 +142,10 @@ class PaimonScanBuilder(val table: InnerTable)
if (
vectorSearch.isDefined &&
+ !CoreOptions
+ .fromMap(actualTable.options)
+ .primaryKeyVectorIndexColumns()
+ .contains(vectorSearch.get.fieldName()) &&
VectorSearchResultUtils.isVectorSearchMetaOnly(requiredSchema.fieldNames.toSeq)
) {
val result = PaimonBaseScan.evalVectorSearch(
diff --git
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkVectorReadImplTest.java
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorReadTest.java
similarity index 97%
rename from
paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkVectorReadImplTest.java
rename to
paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorReadTest.java
index a88d2d9cbf..756dff2cb6 100644
---
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkVectorReadImplTest.java
+++
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorReadTest.java
@@ -56,8 +56,8 @@ import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
-/** Tests for {@link SparkVectorReadImpl}. */
-public class SparkVectorReadImplTest {
+/** Tests for {@link SparkDataEvolutionVectorRead}. */
+public class SparkDataEvolutionVectorReadTest {
@Test
public void testRawSearchUsesSparkPath() {
@@ -116,7 +116,7 @@ public class SparkVectorReadImplTest {
return splits;
}
- private static class TestingSparkVectorRead extends SparkVectorReadImpl {
+ private static class TestingSparkVectorRead extends
SparkDataEvolutionVectorRead {
private boolean rawSparkPathUsed;
@@ -158,7 +158,7 @@ public class SparkVectorReadImplTest {
}
}
- private static class DistributedRefineSparkVectorRead extends
SparkVectorReadImpl {
+ private static class DistributedRefineSparkVectorRead extends
SparkDataEvolutionVectorRead {
private int sparkParallelism;
private List<Long> rawSearchCandidateRows = Collections.emptyList();
@@ -228,7 +228,7 @@ public class SparkVectorReadImplTest {
}
}
- private static class RecordingSparkVectorRead extends SparkVectorReadImpl {
+ private static class RecordingSparkVectorRead extends
SparkDataEvolutionVectorRead {
private final AtomicInteger nextTask = new AtomicInteger();
private final List<Range> rawSearchRanges =
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
new file mode 100644
index 0000000000..84e896ae50
--- /dev/null
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala
@@ -0,0 +1,270 @@
+/*
+ * 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.spark.sql
+
+import org.apache.paimon.globalindex.testvector.TestVectorGlobalIndexerFactory
+import org.apache.paimon.spark.PaimonSparkTestBase
+import org.apache.paimon.spark.read.{SparkPrimaryKeyVectorRead,
SparkVectorSearchBuilderImpl}
+import org.apache.paimon.table.source.DataSplit
+
+import scala.collection.JavaConverters._
+
+/** End-to-end tests for primary-key vector search through Spark SQL. */
+class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase {
+
+ test("distributed primary-key vector search selects Spark reader") {
+ withTable("T") {
+ createVectorTable()
+
+ val builder = new SparkVectorSearchBuilderImpl(loadTable("T"))
+ builder
+ .withVectorColumn("embedding")
+ .withVector(Array(0.0f, 0.0f))
+ .withLimit(1)
+
+ assert(builder.newVectorRead().isInstanceOf[SparkPrimaryKeyVectorRead])
+ }
+ }
+
+ test("distributed primary-key vector search evaluates buckets in Spark") {
+ withTable("T") {
+ createVectorTable(bucket = 2, extraOptions =
Seq("global-index.thread-num" -> "1"))
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, array(1.0f, 0.0f)),
+ | (2, array(2.0f, 0.0f)),
+ | (3, array(3.0f, 0.0f)),
+ | (4, array(4.0f, 0.0f))
+ |""".stripMargin)
+
+ val builder = new SparkVectorSearchBuilderImpl(loadTable("T"))
+ builder
+ .withVectorColumn("embedding")
+ .withVector(Array(0.0f, 0.0f))
+ .withLimit(2)
+
+ val jobGroup = s"primary-key-vector-${System.nanoTime()}"
+ spark.sparkContext.setJobGroup(jobGroup, jobGroup)
+ try {
+ builder.newVectorRead().read(builder.newVectorScan().scan())
+ } finally {
+ spark.sparkContext.clearJobGroup()
+ }
+
+
assert(spark.sparkContext.statusTracker.getJobIdsForGroup(jobGroup).nonEmpty)
+ }
+ }
+
+ test("primary-key vector search uses bucket-local indexes") {
+ withTable("T") {
+ createVectorTable()
+
+ 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)
+
+ withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" ->
"true") {
+ val rows = spark
+ .sql("""
+ |SELECT id, __paimon_search_score
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 2)
+ |""".stripMargin)
+ .collect()
+
+ assert(rows.map(_.getInt(0)).toSet == Set(2, 3))
+ assert(rows.forall(!_.isNullAt(1)))
+
+ val scores = spark
+ .sql("""
+ |SELECT __paimon_search_score
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 2)
+ |""".stripMargin)
+ .collect()
+ assert(scores.length == 2)
+ assert(scores.forall(!_.isNullAt(0)))
+ }
+ }
+ }
+
+ test("primary-key vector search merges top k across buckets") {
+ withTable("T") {
+ createVectorTable(bucket = 4, extraOptions =
Seq("global-index.thread-num" -> "2"))
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, array(1.0f, 0.0f)),
+ | (2, array(2.0f, 0.0f)),
+ | (3, array(3.0f, 0.0f)),
+ | (4, array(4.0f, 0.0f)),
+ | (5, array(5.0f, 0.0f)),
+ | (6, array(6.0f, 0.0f)),
+ | (7, array(7.0f, 0.0f)),
+ | (8, array(8.0f, 0.0f)),
+ | (9, array(9.0f, 0.0f)),
+ | (10, array(10.0f, 0.0f)),
+ | (11, array(11.0f, 0.0f)),
+ | (12, array(12.0f, 0.0f)),
+ | (13, array(13.0f, 0.0f)),
+ | (14, array(14.0f, 0.0f)),
+ | (15, array(15.0f, 0.0f)),
+ | (16, array(16.0f, 0.0f))
+ |""".stripMargin)
+
+ val buckets = loadTable("T")
+ .newReadBuilder()
+ .newScan()
+ .plan()
+ .splits()
+ .asScala
+ .map(_.asInstanceOf[DataSplit].bucket())
+ .toSet
+ assert(buckets == Set(0, 1, 2, 3))
+
+ withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" ->
"true") {
+ val ids = spark
+ .sql("""
+ |SELECT id
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 3)
+ |""".stripMargin)
+ .collect()
+ .map(_.getInt(0))
+ .toSet
+ assert(ids == Set(1, 2, 3))
+ }
+ }
+ }
+
+ test("primary-key vector search prunes partitions before top k") {
+ withTable("T") {
+ createVectorTable(
+ columns = "id INT, embedding ARRAY<FLOAT>, dt STRING",
+ primaryKey = "id,dt",
+ partitionedBy = Some("dt"))
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, array(1.0f, 0.0f), 'A'),
+ | (2, array(2.0f, 0.0f), 'A'),
+ | (3, array(0.1f, 0.0f), 'B')
+ |""".stripMargin)
+
+ withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" ->
"true") {
+ val ids = spark
+ .sql("""
+ |SELECT id
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 2)
+ |WHERE dt = 'A'
+ |""".stripMargin)
+ .collect()
+ .map(_.getInt(0))
+ .toSet
+ assert(ids == Set(1, 2))
+ }
+ }
+ }
+
+ test("deduplicate updates and deletes primary-key vector results") {
+ withTable("T") {
+ createVectorTable()
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, array(3.0f, 0.0f)),
+ | (2, array(1.0f, 0.0f))
+ |""".stripMargin)
+ spark.sql("INSERT INTO T VALUES (1, array(0.5f, 0.0f))")
+
+ withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" ->
"true") {
+ val updated = spark
+ .sql("""
+ |SELECT id
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 1)
+ |""".stripMargin)
+ .collect()
+ assert(updated.map(_.getInt(0)).toSeq == Seq(1))
+
+ spark.sql("DELETE FROM T WHERE id = 1")
+
+ val afterDelete = spark
+ .sql("""
+ |SELECT id
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 1)
+ |""".stripMargin)
+ .collect()
+ assert(afterDelete.map(_.getInt(0)).toSeq == Seq(2))
+ }
+ }
+ }
+
+ test("partial update completes rows before publishing vector results") {
+ withTable("T") {
+ createVectorTable(
+ columns = "id INT, payload STRING, embedding ARRAY<FLOAT>",
+ extraOptions =
+ Seq("merge-engine" -> "partial-update",
"deletion-vectors.merge-on-read" -> "false")
+ )
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (1, 'keep', array(3.0f, 0.0f)),
+ | (2, 'other', array(1.0f, 0.0f))
+ |""".stripMargin)
+ spark.sql("INSERT INTO T (id, embedding) VALUES (1, array(0.5f, 0.0f))")
+
+ withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" ->
"true") {
+ val rows = spark
+ .sql("""
+ |SELECT id, payload
+ |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 1)
+ |""".stripMargin)
+ .collect()
+ assert(rows.length == 1)
+ assert(rows.head.getInt(0) == 1)
+ assert(rows.head.getString(1) == "keep")
+ }
+ }
+ }
+
+ private def createVectorTable(
+ columns: String = "id INT, embedding ARRAY<FLOAT>",
+ primaryKey: String = "id",
+ bucket: Int = 1,
+ extraOptions: Seq[(String, String)] = Seq.empty,
+ partitionedBy: Option[String] = None): Unit = {
+ val properties = (Seq(
+ "primary-key" -> primaryKey,
+ "bucket" -> bucket.toString,
+ "deletion-vectors.enabled" -> "true",
+ "vector-field" -> "embedding",
+ "field.embedding.vector-dim" -> "2",
+ "pk-vector.index.columns" -> "embedding",
+ "fields.embedding.pk-vector.index.type" ->
TestVectorGlobalIndexerFactory.IDENTIFIER,
+ "fields.embedding.pk-vector.distance.metric" -> "l2",
+ "test.vector.dimension" -> "2",
+ "test.vector.metric" -> "l2"
+ ) ++ extraOptions)
+ .map { case (key, value) => s"'$key' = '$value'" }
+ .mkString(",\n")
+ val partitioning = partitionedBy.map(column => s"PARTITIONED BY
($column)").getOrElse("")
+ spark.sql(s"""
+ |CREATE TABLE T ($columns)
+ |$partitioning
+ |TBLPROPERTIES ($properties)
+ |""".stripMargin)
+ }
+}
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/VectorSearchOptionsTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/VectorSearchOptionsTest.scala
index 79b78a6146..2531fffad7 100644
---
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/VectorSearchOptionsTest.scala
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/VectorSearchOptionsTest.scala
@@ -65,6 +65,15 @@ class VectorSearchOptionsTest extends PaimonSparkTestBase {
.collect()
assert(result.length == 1)
+
+ val scores = spark
+ .sql("""
+ |SELECT __paimon_search_score FROM vector_search(
+ | 'T', 'v', array(1.0f, 0.0f), 1, map('ivf.nprobe', '16'))
+ |""".stripMargin)
+ .collect()
+ assert(scores.length == 1)
+ assert(!scores.head.isNullAt(0))
}
}
}
