From 1ca7d3070edff2f9b8b81f88a2e6b9fcc8c0c8d4 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 21:44:58 +0800 Subject: [PATCH 1/8] [spark] Support primary-key vector search --- .../paimon/reader/ScoreRecordReader.java | 31 +++ .../globalindex/IndexedSplitRecordReader.java | 3 +- .../PrimaryKeyVectorPositionReader.java | 4 +- .../table/source/VectorSearchBuilderImpl.java | 2 +- .../read/SparkVectorSearchBuilderImpl.java | 8 +- .../spark/PaimonRecordReaderIterator.scala | 5 +- .../paimon/spark/PaimonScanBuilder.scala | 5 + .../sql/PrimaryKeyVectorSearchTest.scala | 218 ++++++++++++++++++ .../spark/sql/VectorSearchOptionsTest.scala | 9 + 9 files changed, 276 insertions(+), 9 deletions(-) create mode 100644 paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordReader.java create mode 100644 paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala 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 000000000000..69b946332595 --- /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 extends RecordReader { + + @Nullable + @Override + ScoreRecordIterator 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 3778f37cb90a..3d5bc89e00e9 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 @@ 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 static org.apache.paimon.utils.Preconditions.checkArgument; /** Return value with score. */ -public class IndexedSplitRecordReader implements RecordReader { +public class IndexedSplitRecordReader implements ScoreRecordReader { private final RecordReader reader; @Nullable private final Map rowIdToScore; 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 6470fc26ebd7..52f7423a9260 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 @@ 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 static org.apache.paimon.utils.Preconditions.checkArgument; /** Reads selected physical file positions and exposes their vector-search scores. */ -public class PrimaryKeyVectorPositionReader implements RecordReader { +public class PrimaryKeyVectorPositionReader implements ScoreRecordReader { private final FileRecordReader reader; private final RoaringBitmap32 rowPositions; 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 f2bcb7af65cb..e706e40ba6ca 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 VectorRead newVectorRead() { 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-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 87044862582e..889a522835db 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,9 @@ 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 a {@link SparkVectorReadImpl} for + * data-evolution tables so the per-split vector index evaluation is dispatched through Spark + * instead of the local thread pool. Primary-key vector search keeps its bucket-local Core reader. * *

Single-vector only; batch search has no Spark-dispatched path yet (TODO). */ @@ -38,6 +39,9 @@ public SparkVectorSearchBuilderImpl(InnerTable table) { @Override public VectorRead newVectorRead() { + if (isPrimaryKeyVectorSearch()) { + return super.newVectorRead(); + } return new SparkVectorReadImpl( 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 64b0b5166c94..93256df2c00f 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 6d7f894beaeb..844c4de52b1d 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-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 000000000000..ed4557139ed5 --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala @@ -0,0 +1,218 @@ +/* + * 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.table.source.DataSplit + +import scala.collection.JavaConverters._ + +/** End-to-end tests for primary-key vector search through Spark SQL. */ +class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase { + + 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 = 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)) + |""".stripMargin) + + val buckets = loadTable("T") + .newReadBuilder() + .newScan() + .plan() + .splits() + .asScala + .map(_.asInstanceOf[DataSplit].bucket()) + .toSet + assert(buckets == Set(0, 1)) + + 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, 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", + 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", + 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 79b78a614624..2531fffad78b 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)) } } } From 8c2c17c9b6e396254200cfb3ac2839bb98b08718 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 22:18:19 +0800 Subject: [PATCH 2/8] [flink][docs] Test and document primary-key vector search --- docs/docs/primary-key-table/vector-index.md | 208 ++++++++++++++++++ docs/sidebars.js | 1 + .../VectorSearchProcedureITCase.java | 119 ++++++++++ 3 files changed, 328 insertions(+) create mode 100644 docs/docs/primary-key-table/vector-index.md 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 000000000000..b497b7449a61 --- /dev/null +++ b/docs/docs/primary-key-table/vector-index.md @@ -0,0 +1,208 @@ +--- +title: "Vector Index" +sidebar_position: 9 +--- + + + +# 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` and `deletion-vectors.merge-on-read` is `false`. +- Its merge engine is `deduplicate` (the default) or `partial-update`. +- 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 COMMENT '__VECTOR_FIELD;3', + PRIMARY KEY (id) NOT ENFORCED +) WITH ( + 'bucket' = '16', + 'deletion-vectors.enabled' = 'true', + 'deletion-vectors.merge-on-read' = 'false', + '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 COMMENT '__VECTOR_FIELD;3' +) USING paimon +TBLPROPERTIES ( + 'primary-key' = 'id', + 'bucket' = '16', + 'deletion-vectors.enabled' = 'true', + 'deletion-vectors.merge-on-read' = 'false', + '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` column to Paimon's fixed-length +`VECTOR` 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..pk-vector.index.type` | Yes | ANN implementation, such as `ivf-flat`, `ivf-pq`, `ivf-hnsw-flat`, `ivf-hnsw-sq`, or `lumina`. | +| `fields..pk-vector.distance.metric` | No | `l2`, `cosine`, or `inner_product`. The default is `inner_product`. | +| `fields..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 first completes the row through lookup and compaction, then builds the + vector index from the complete compact-output row. The index therefore does not require + `lookup-wait = true` and never indexes an incomplete input 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. + +### 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 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. +- `aggregation` and `first-row` merge engines 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 4223c1347faf..7e3967dadea2 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-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 d71de22d2a64..bc29716356c5 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 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 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 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 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 @@ private void createVectorTable(String tableName, String extraOptions) { tableName, DIMENSION, formattedExtraOptions); } + private void createPrimaryKeyVectorTable(String tableName) { + sql( + "CREATE TABLE %s (" + + "id INT, " + + "vec ARRAY, " + + "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 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, " + + "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(); From a899dc1d0e649cabc0db7b0efe24cef9b36ef026 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 22:28:35 +0800 Subject: [PATCH 3/8] [docs] Simplify primary-key vector index configuration --- docs/docs/primary-key-table/vector-index.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/docs/docs/primary-key-table/vector-index.md b/docs/docs/primary-key-table/vector-index.md index b497b7449a61..cf421e5c4b07 100644 --- a/docs/docs/primary-key-table/vector-index.md +++ b/docs/docs/primary-key-table/vector-index.md @@ -40,7 +40,7 @@ is built separately from writes, see A table with a primary-key vector index must satisfy all of the following: - It is a primary-key table in fixed-bucket mode (`bucket > 0`). -- `deletion-vectors.enabled` is `true` and `deletion-vectors.merge-on-read` is `false`. +- `deletion-vectors.enabled` is `true`. - Its merge engine is `deduplicate` (the default) or `partial-update`. - The indexed column is a `VECTOR` whose element type is `FLOAT`. - `pk-clustering-override` is disabled. @@ -64,7 +64,6 @@ CREATE TABLE item_embeddings ( ) WITH ( 'bucket' = '16', 'deletion-vectors.enabled' = 'true', - 'deletion-vectors.merge-on-read' = 'false', 'pk-vector.index.columns' = 'embedding', 'fields.embedding.pk-vector.index.type' = 'ivf-flat', 'fields.embedding.pk-vector.distance.metric' = 'cosine', @@ -84,7 +83,6 @@ TBLPROPERTIES ( 'primary-key' = 'id', 'bucket' = '16', 'deletion-vectors.enabled' = 'true', - 'deletion-vectors.merge-on-read' = 'false', 'pk-vector.index.columns' = 'embedding', 'fields.embedding.pk-vector.index.type' = 'ivf-flat', 'fields.embedding.pk-vector.distance.metric' = 'cosine', From f4afc2837197acf7af33b1034bd45b9de6789b69 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 22:30:50 +0800 Subject: [PATCH 4/8] [docs] Clarify vector index compaction visibility --- docs/docs/primary-key-table/vector-index.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/docs/primary-key-table/vector-index.md b/docs/docs/primary-key-table/vector-index.md index cf421e5c4b07..8348f080c773 100644 --- a/docs/docs/primary-key-table/vector-index.md +++ b/docs/docs/primary-key-table/vector-index.md @@ -117,9 +117,8 @@ 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 first completes the row through lookup and compaction, then builds the - vector index from the complete compact-output row. The index therefore does not require - `lookup-wait = true` and never indexes an incomplete input row. +- `partial-update`: Paimon builds the vector index from the lookup-completed Level-1 + compact-output 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 From 22dee4dc1e1707e2b0f46d82686d4fc55a922b5e Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 22:40:10 +0800 Subject: [PATCH 5/8] [core] Support all merge engines for primary-key vector index --- docs/docs/primary-key-table/vector-index.md | 8 +- .../MergeTreeCompactManagerFactory.java | 2 +- .../paimon/schema/SchemaValidation.java | 12 +-- .../PrimaryKeyVectorIndexValidationTest.java | 25 ++++++ .../source/PrimaryKeyVectorSearchTest.java | 80 ++++++++++++++++++- 5 files changed, 114 insertions(+), 13 deletions(-) diff --git a/docs/docs/primary-key-table/vector-index.md b/docs/docs/primary-key-table/vector-index.md index 8348f080c773..841b4bf7cf2a 100644 --- a/docs/docs/primary-key-table/vector-index.md +++ b/docs/docs/primary-key-table/vector-index.md @@ -40,8 +40,8 @@ is built separately from writes, see A table with a primary-key vector index must satisfy all of the following: - It is a primary-key table in fixed-bucket mode (`bucket > 0`). -- `deletion-vectors.enabled` is `true`. -- Its merge engine is `deduplicate` (the default) or `partial-update`. +- `deletion-vectors.enabled` is `true`, except for `first-row`, where it is optional. +- 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. @@ -119,6 +119,9 @@ The maintenance behavior depends on the merge engine: 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 are optional because later + rows with the same primary key are ignored. 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 @@ -199,7 +202,6 @@ cached payloads and fetch only the selected data-file positions. - 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. -- `aggregation` and `first-row` merge engines 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/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java index d90fd76c9f2e..ce82b2f251d9 100644 --- a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java +++ b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java @@ -288,7 +288,7 @@ private MergeTreeCompactRewriter createRewriter( LookupMergeTreeCompactRewriter.MergeFunctionWrapperFactory wrapperFactory; FileReaderFactory lookupReaderFactory = readerFactory; if (lookupStrategy.isFirstRow) { - if (options.deletionVectorsEnabled()) { + if (options.deletionVectorsEnabled() && !options.primaryKeyVectorIndexEnabled()) { throw new UnsupportedOperationException( "First row merge engine does not need deletion vectors because there is no deletion of old data in this merge engine."); } 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 2122c04ed515..fd9b66754709 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 @@ -877,8 +877,9 @@ private static void validateForDeletionVectors(CoreOptions options) { || options.changelogProducer() == ChangelogProducer.LOOKUP, "Deletion vectors mode is only supported for NONE/INPUT/LOOKUP changelog producer now."); - // pk-clustering-override mode allows deletion vectors for first-row - if (!options.pkClusteringOverride()) { + // pk-clustering-override and primary-key vector index modes allow deletion vectors for + // first-row. + if (!options.pkClusteringOverride() && !options.primaryKeyVectorIndexEnabled()) { checkArgument( !options.mergeEngine().equals(MergeEngine.FIRST_ROW), "First row merge engine does not need deletion vectors because there is no deletion of old data in this merge engine."); @@ -912,13 +913,8 @@ private static void validatePrimaryKeyVectorIndex(TableSchema schema, CoreOption !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/test/java/org/apache/paimon/schema/PrimaryKeyVectorIndexValidationTest.java b/paimon-core/src/test/java/org/apache/paimon/schema/PrimaryKeyVectorIndexValidationTest.java index 0daa7bc5c473..17982096c1a4 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,31 @@ void testSupportsPartialUpdateMergeEngine() { assertThatCode(() -> validateTableSchema(schema(options))).doesNotThrowAnyException(); } + @Test + void testSupportsAggregationMergeEngine() { + Map options = enabledOptions(); + options.put(CoreOptions.MERGE_ENGINE.key(), "aggregation"); + + assertThatCode(() -> validateTableSchema(schema(options))).doesNotThrowAnyException(); + } + + @Test + void testSupportsFirstRowWithoutDeletionVectors() { + Map 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 testSupportsFirstRowWithDeletionVectors() { + Map options = enabledOptions(); + options.put(CoreOptions.MERGE_ENGINE.key(), "first-row"); + + assertThatCode(() -> validateTableSchema(schema(options))).doesNotThrowAnyException(); + } + @Test void testPartialUpdateRejectsDeletionVectorMergeOnRead() { Map 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 090488de9e53..494d4ae19379 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,75 @@ void testVectorSearchMaterializesPhysicalRows() throws Exception { assertThat(ids).containsExactly(2, 3); } + + @Test + void testFirstRowVectorSearchWithDeletionVectors() throws Exception { + assertFirstRowVectorSearch(true); + } + + @Test + void testFirstRowVectorSearchWithoutDeletionVectors() throws Exception { + assertFirstRowVectorSearch(false); + } + + private void assertFirstRowVectorSearch(boolean deletionVectorsEnabled) throws Exception { + catalog.createTable(identifier(), vectorSchema("first-row", deletionVectorsEnabled), 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 ids = new ArrayList<>(); + try (RecordReader 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 ids = new ArrayList<>(); + try (RecordReader reader = readBuilder.newRead().createReader(plan)) { + reader.forEachRemaining(row -> ids.add(row.getInt(0))); + } + + assertThat(ids).containsExactly(1); + } } From b27bb354ab4b0f9a8fb0d1a9b283cd26794d38bf Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 22:48:20 +0800 Subject: [PATCH 6/8] [core] Reject deletion vectors for first-row vector index --- docs/docs/primary-key-table/vector-index.md | 6 +++--- .../compact/MergeTreeCompactManagerFactory.java | 2 +- .../org/apache/paimon/schema/SchemaValidation.java | 5 ++--- .../schema/PrimaryKeyVectorIndexValidationTest.java | 6 ++++-- .../table/source/PrimaryKeyVectorSearchTest.java | 13 ++----------- 5 files changed, 12 insertions(+), 20 deletions(-) diff --git a/docs/docs/primary-key-table/vector-index.md b/docs/docs/primary-key-table/vector-index.md index 841b4bf7cf2a..9a9abb4d9413 100644 --- a/docs/docs/primary-key-table/vector-index.md +++ b/docs/docs/primary-key-table/vector-index.md @@ -40,7 +40,7 @@ is built separately from writes, see A table with a primary-key vector index must satisfy all of the following: - It is a primary-key table in fixed-bucket mode (`bucket > 0`). -- `deletion-vectors.enabled` is `true`, except for `first-row`, where it is optional. +- `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. @@ -120,8 +120,8 @@ The maintenance behavior depends on the merge engine: - `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 are optional because later - rows with the same primary key are ignored. +- `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 diff --git a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java index ce82b2f251d9..d90fd76c9f2e 100644 --- a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java +++ b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/MergeTreeCompactManagerFactory.java @@ -288,7 +288,7 @@ private MergeTreeCompactRewriter createRewriter( LookupMergeTreeCompactRewriter.MergeFunctionWrapperFactory wrapperFactory; FileReaderFactory lookupReaderFactory = readerFactory; if (lookupStrategy.isFirstRow) { - if (options.deletionVectorsEnabled() && !options.primaryKeyVectorIndexEnabled()) { + if (options.deletionVectorsEnabled()) { throw new UnsupportedOperationException( "First row merge engine does not need deletion vectors because there is no deletion of old data in this merge engine."); } 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 fd9b66754709..09cb61c0bf70 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 @@ -877,9 +877,8 @@ private static void validateForDeletionVectors(CoreOptions options) { || options.changelogProducer() == ChangelogProducer.LOOKUP, "Deletion vectors mode is only supported for NONE/INPUT/LOOKUP changelog producer now."); - // pk-clustering-override and primary-key vector index modes allow deletion vectors for - // first-row. - if (!options.pkClusteringOverride() && !options.primaryKeyVectorIndexEnabled()) { + // pk-clustering-override mode allows deletion vectors for first-row + if (!options.pkClusteringOverride()) { checkArgument( !options.mergeEngine().equals(MergeEngine.FIRST_ROW), "First row merge engine does not need deletion vectors because there is no deletion of old data in this merge engine."); 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 17982096c1a4..6f69b2a70c36 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 @@ -141,11 +141,13 @@ void testSupportsFirstRowWithoutDeletionVectors() { } @Test - void testSupportsFirstRowWithDeletionVectors() { + void testRejectsFirstRowWithDeletionVectors() { Map options = enabledOptions(); options.put(CoreOptions.MERGE_ENGINE.key(), "first-row"); - assertThatCode(() -> validateTableSchema(schema(options))).doesNotThrowAnyException(); + 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 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 494d4ae19379..5c6558fdfe0c 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 @@ -103,17 +103,8 @@ void testVectorSearchMaterializesPhysicalRows() throws Exception { } @Test - void testFirstRowVectorSearchWithDeletionVectors() throws Exception { - assertFirstRowVectorSearch(true); - } - - @Test - void testFirstRowVectorSearchWithoutDeletionVectors() throws Exception { - assertFirstRowVectorSearch(false); - } - - private void assertFirstRowVectorSearch(boolean deletionVectorsEnabled) throws Exception { - catalog.createTable(identifier(), vectorSchema("first-row", deletionVectorsEnabled), false); + void testFirstRowVectorSearch() throws Exception { + catalog.createTable(identifier(), vectorSchema("first-row", false), false); FileStoreTable table = getTableDefault(); write( From 7d13db75ebf366ab59dcd4667ea2f10195930eed Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 12 Jul 2026 23:51:04 +0800 Subject: [PATCH 7/8] [spark] Distribute primary-key vector search --- docs/docs/primary-key-table/vector-index.md | 3 + .../table/source/PrimaryKeyVectorRead.java | 99 +++++++++---- ...java => SparkDataEvolutionVectorRead.java} | 4 +- .../spark/read/SparkPrimaryKeyVectorRead.java | 133 ++++++++++++++++++ .../read/SparkVectorSearchBuilderImpl.java | 9 +- ... => SparkDataEvolutionVectorReadTest.java} | 10 +- .../sql/PrimaryKeyVectorSearchTest.scala | 62 +++++++- 7 files changed, 276 insertions(+), 44 deletions(-) rename paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/{SparkVectorReadImpl.java => SparkDataEvolutionVectorRead.java} (99%) create mode 100644 paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkPrimaryKeyVectorRead.java rename paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/{SparkVectorReadImplTest.java => SparkDataEvolutionVectorReadTest.java} (97%) diff --git a/docs/docs/primary-key-table/vector-index.md b/docs/docs/primary-key-table/vector-index.md index 9a9abb4d9413..8ac875b29d03 100644 --- a/docs/docs/primary-key-table/vector-index.md +++ b/docs/docs/primary-key-table/vector-index.md @@ -152,6 +152,9 @@ FROM vector_search( 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 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 68350ffb83b3..b70f19eb738e 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.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.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 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 searchOptions; private final float[] query; - private final int limit; + protected final int limit; private final String metric; public PrimaryKeyVectorRead( @@ -102,15 +102,10 @@ public PrimaryKeyVectorRead( 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 @@ private static KeyValueFileStore keyValueStore(FileStoreTable table) { @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 bucketSplits(PrimaryKeyVectorScan.Plan plan) { + List splits = new ArrayList<>(plan.splits().size()); + for (VectorSearchSplit split : plan.splits()) { + splits.add((BucketVectorSearchSplit) split); + } + return splits; + } + + protected List searchBuckets(List splits) { try { + SearchContext context = new SearchContext(table); List 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 search(BucketVectorSearchSplit split) throws IOException { + protected GlobalIndexResult createResult( + PrimaryKeyVectorScan.Plan plan, List candidates) { + return new PrimaryKeyVectorResult(plan, topK(candidates, limit), metric); + } + + private List search(BucketVectorSearchSplit split, SearchContext context) + throws IOException { DataSplit dataSplit = split.dataSplit(); List activeFiles = dataSplit.dataFiles().stream() @@ -159,23 +180,23 @@ private List search(BucketVectorSearchSplit split) throws IOException PkVectorBucketIndexState state = PkVectorBucketIndexState.fromActivePayloads( vectorField.id(), indexType, split.payloadFiles()); - Map deletionVectors = deletionVectors(dataSplit); + Map 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 candidates = new ArrayList<>(); @@ -192,7 +213,8 @@ private List search(BucketVectorSearchSplit split) throws IOException return candidates; } - private Map deletionVectors(DataSplit split) throws IOException { + private Map deletionVectors(DataSplit split, FileIO fileIO) + throws IOException { DeletionVector.Factory factory = DeletionVector.factory( fileIO, split.dataFiles(), split.deletionFiles().orElse(null)); @@ -206,7 +228,7 @@ private Map deletionVectors(DataSplit split) throws IOEx return result; } - static List topK(List candidates, int limit) { + protected static List topK(List candidates, int limit) { checkArgument(limit > 0, "Vector search limit must be positive: %s.", limit); PriorityQueue nearest = new PriorityQueue<>(limit, BEST_FIRST.reversed()); for (Candidate candidate : candidates) { @@ -234,7 +256,9 @@ private static int compareBytes(byte[] left, byte[] right) { } /** 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 float distance() { 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-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 3c13fb4c0c65..ddac704b667a 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 @@ * 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 000000000000..053abc9677e9 --- /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 searchOptions) { + super(table, vectorField, query, limit, searchOptions); + } + + @Override + public GlobalIndexResult read(VectorScan.Plan plan) { + PrimaryKeyVectorScan.Plan primaryKeyPlan = primaryKeyPlan(plan); + List splits = bucketSplits(primaryKeyPlan); + int parallelism = sparkParallelism(); + if (splits.size() < parallelism * 2) { + return super.read(plan); + } + + List 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> groups = splitGroups(serializedSplits, parallelism); + SerializableFunction, byte[]> task = + group -> { + List 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 groupResults = mapInSpark(groups, task, groups.size()); + List 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 List mapInSpark( + List data, SerializableFunction function, int parallelism) { + return createEngineContext().map(data, function, parallelism); + } + + private List> splitGroups(List splits, int parallelism) { + List> 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 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 889a522835db..7638f5726965 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,9 +23,8 @@ import org.apache.paimon.table.source.VectorSearchBuilderImpl; /** - * Spark-aware {@link VectorSearchBuilderImpl} which produces a {@link SparkVectorReadImpl} for - * data-evolution tables so the per-split vector index evaluation is dispatched through Spark - * instead of the local thread pool. Primary-key vector search keeps its bucket-local Core reader. + * 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. * *

Single-vector only; batch search has no Spark-dispatched path yet (TODO). */ @@ -40,9 +39,9 @@ public SparkVectorSearchBuilderImpl(InnerTable table) { @Override public VectorRead newVectorRead() { if (isPrimaryKeyVectorSearch()) { - return super.newVectorRead(); + return new SparkPrimaryKeyVectorRead(table, vectorColumn, vector, limit, options); } - return new SparkVectorReadImpl( + return new SparkDataEvolutionVectorRead( table, partitionFilter, filter, limit, vectorColumn, vector, options); } } 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 a88d2d9cbfc7..756dff2cb698 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 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 @@ private static List indexSplits(String indexType, int co return splits; } - private static class TestingSparkVectorRead extends SparkVectorReadImpl { + private static class TestingSparkVectorRead extends SparkDataEvolutionVectorRead { private boolean rawSparkPathUsed; @@ -158,7 +158,7 @@ protected ScoredGlobalIndexResult readRawSplitsInSpark( } } - private static class DistributedRefineSparkVectorRead extends SparkVectorReadImpl { + private static class DistributedRefineSparkVectorRead extends SparkDataEvolutionVectorRead { private int sparkParallelism; private List rawSearchCandidateRows = Collections.emptyList(); @@ -228,7 +228,7 @@ private static ScoredGlobalIndexResult scoredResult(long rowId, float score) { } } - private static class RecordingSparkVectorRead extends SparkVectorReadImpl { + private static class RecordingSparkVectorRead extends SparkDataEvolutionVectorRead { private final AtomicInteger nextTask = new AtomicInteger(); private final List 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 index ed4557139ed5..11c56f40b94a 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala @@ -20,6 +20,7 @@ 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._ @@ -27,6 +28,51 @@ 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() @@ -63,7 +109,9 @@ class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase { test("primary-key vector search merges top k across buckets") { withTable("T") { - createVectorTable(bucket = 2) + createVectorTable( + bucket = 4, + extraOptions = Seq("global-index.thread-num" -> "2")) spark.sql(""" |INSERT INTO T VALUES | (1, array(1.0f, 0.0f)), @@ -73,7 +121,15 @@ class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase { | (5, array(5.0f, 0.0f)), | (6, array(6.0f, 0.0f)), | (7, array(7.0f, 0.0f)), - | (8, array(8.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") @@ -84,7 +140,7 @@ class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase { .asScala .map(_.asInstanceOf[DataSplit].bucket()) .toSet - assert(buckets == Set(0, 1)) + assert(buckets == Set(0, 1, 2, 3)) withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" -> "true") { val ids = spark From 3d4cfe7e6502f4bcbb5afc73b80955cd00f46eab Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 13 Jul 2026 08:17:00 +0800 Subject: [PATCH 8/8] checkstyle --- .../paimon/spark/sql/PrimaryKeyVectorSearchTest.scala | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala index 11c56f40b94a..84e896ae5006 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala @@ -44,9 +44,7 @@ class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase { test("distributed primary-key vector search evaluates buckets in Spark") { withTable("T") { - createVectorTable( - bucket = 2, - extraOptions = Seq("global-index.thread-num" -> "1")) + createVectorTable(bucket = 2, extraOptions = Seq("global-index.thread-num" -> "1")) spark.sql(""" |INSERT INTO T VALUES | (1, array(1.0f, 0.0f)), @@ -109,9 +107,7 @@ class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase { test("primary-key vector search merges top k across buckets") { withTable("T") { - createVectorTable( - bucket = 4, - extraOptions = Seq("global-index.thread-num" -> "2")) + createVectorTable(bucket = 4, extraOptions = Seq("global-index.thread-num" -> "2")) spark.sql(""" |INSERT INTO T VALUES | (1, array(1.0f, 0.0f)),