This is an automated email from the ASF dual-hosted git repository. JingsongLi pushed a commit to branch codex/pk-vector-07-spark-search in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 1ca7d3070edff2f9b8b81f88a2e6b9fcc8c0c8d4 Author: JingsongLi <[email protected]> AuthorDate: Sun Jul 12 21:44:58 2026 +0800 [spark] Support primary-key vector search --- .../apache/paimon/reader/ScoreRecordReader.java | 31 +++ .../globalindex/IndexedSplitRecordReader.java | 3 +- .../source/PrimaryKeyVectorPositionReader.java | 4 +- .../table/source/VectorSearchBuilderImpl.java | 2 +- .../spark/read/SparkVectorSearchBuilderImpl.java | 8 +- .../paimon/spark/PaimonRecordReaderIterator.scala | 5 +- .../apache/paimon/spark/PaimonScanBuilder.scala | 5 + .../spark/sql/PrimaryKeyVectorSearchTest.scala | 218 +++++++++++++++++++++ .../paimon/spark/sql/VectorSearchOptionsTest.scala | 9 + 9 files changed, 276 insertions(+), 9 deletions(-) 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/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/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-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..889a522835 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.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 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. * * <p>Single-vector only; batch search has no Spark-dispatched path yet (TODO). */ @@ -38,6 +39,9 @@ public class SparkVectorSearchBuilderImpl extends VectorSearchBuilderImpl { @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 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-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..ed4557139e --- /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<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)) } } }
