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..8ac875b29d03 --- /dev/null +++ b/docs/docs/primary-key-table/vector-index.md @@ -0,0 +1,210 @@ +--- +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`, except for `first-row`, where it must be `false`. +- Its merge engine is `deduplicate`, `partial-update`, `aggregation`, or `first-row`. +- The indexed column is a `VECTOR` whose element type is `FLOAT`. +- `pk-clustering-override` is disabled. +- The configured vector index implementation is available on every writer and reader classpath. + +The first release supports exactly one indexed vector column per table. The option layout is +field-scoped so that more independently indexed vector columns can be supported in a future +release. + +## Create Table + +The following Flink SQL example creates a three-dimensional vector column and maintains an +IVF-Flat index for it. Use the dimension produced by your embedding model in production. + +```sql +CREATE TABLE item_embeddings ( + id BIGINT, + payload STRING, + embedding ARRAY COMMENT '__VECTOR_FIELD;3', + PRIMARY KEY (id) NOT ENFORCED +) WITH ( + 'bucket' = '16', + 'deletion-vectors.enabled' = 'true', + 'pk-vector.index.columns' = 'embedding', + 'fields.embedding.pk-vector.index.type' = 'ivf-flat', + 'fields.embedding.pk-vector.distance.metric' = 'cosine', + 'fields.embedding.pk-vector.index.options' = '{"nlist":"256"}' +); +``` + +Use the same properties in Spark SQL: + +```sql +CREATE TABLE item_embeddings ( + id BIGINT, + payload STRING, + embedding ARRAY COMMENT '__VECTOR_FIELD;3' +) USING paimon +TBLPROPERTIES ( + 'primary-key' = 'id', + 'bucket' = '16', + 'deletion-vectors.enabled' = 'true', + 'pk-vector.index.columns' = 'embedding', + 'fields.embedding.pk-vector.index.type' = 'ivf-flat', + 'fields.embedding.pk-vector.distance.metric' = 'cosine', + 'fields.embedding.pk-vector.index.options' = '{"nlist":"256"}' +); +``` + +The vector comment directive converts the SQL `ARRAY` 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 builds the vector index from the lookup-completed Level-1 + compact-output row. +- `aggregation`: Paimon builds the vector index from the aggregated Level-1 compact-output row. +- `first-row`: Paimon indexes the retained first row. Deletion vectors must be disabled because + later rows with the same primary key are ignored rather than deleting the retained row. + +When compaction replaces source data files, Paimon removes ANN segments that reference those files +and creates replacement segments for the new compact-output files. Small outputs are indexed as +well; there is no minimum-row threshold before a new segment can be built. + +The index follows compaction freshness. Newly appended level-0 files are not ANN sources, so a +streaming write may not be searchable until compaction has produced and committed its complete +level-1 output. Wait for that compaction when read-after-write vector-search visibility is +required. Batch writes which wait for compaction can publish the data and its index together. + +## Search + +### Spark SQL + +Use the `vector_search` table-valued function. Spark exposes the ANN score through the +`__paimon_search_score` metadata column. + +```sql +SELECT id, payload, __paimon_search_score +FROM vector_search( + 'item_embeddings', + 'embedding', + array(0.1f, 0.2f, 0.3f), + 10, + map('ivf.nprobe', '32') +); +``` + +The query vector dimension must match the indexed column dimension. For partitioned tables, Spark +applies a partition predicate before running ANN and merging the global Top-K. +When `spark.paimon.vector-search.distribute.enabled` is `true`, Spark distributes sufficiently +large groups of bucket-local ANN searches across executors and merges their task-local Top-K +results on the driver. Small plans stay local to avoid Spark job startup overhead. + +### Flink SQL + +Flink exposes vector search as a procedure and returns JSON-serialized rows. Use `projection` to +avoid reading columns that are not needed. + +```sql +CALL sys.vector_search( + `table` => 'default.item_embeddings', + vector_column => 'embedding', + query_vector => '0.1,0.2,0.3', + top_k => 10, + projection => 'id,payload', + options => 'ivf.nprobe=32' +); +``` + +### Java API + +```java +GlobalIndexResult result = table.newVectorSearchBuilder() + .withVectorColumn("embedding") + .withVector(queryVector) + .withLimit(10) + .withOption("ivf.nprobe", "32") + .executeLocal(); + +ReadBuilder readBuilder = table.newReadBuilder(); +TableScan.Plan plan = readBuilder.newScan().withGlobalIndexResult(result).plan(); +try (RecordReader reader = readBuilder.newRead().createReader(plan)) { + reader.forEachRemaining(row -> consume(row)); +} +``` + +## Query Planning + +A search captures one table snapshot, plans the active ANN segments for every selected bucket, +searches those segments, and merges their candidates into one global Top-K. The returned candidates +are materialized from the source data files by physical row position. Deletion vectors are applied +while searching and reading, so stale versions and deleted rows are not returned. + +For low latency on object storage, cache data files and ANN payloads with a caching file system. +The first query may still need to download index files; subsequent queries can search the local +cached payloads and fetch only the selected data-file positions. + +## Limitations + +- Exactly one vector index column is supported per table in the first release. +- Only `FLOAT` vectors are supported. +- Dynamic-bucket and `pk-clustering-override` tables are not supported. +- Flink's procedure returns rows but does not expose the ANN score as a separate column. +- Vector search is snapshot-scoped batch reading; streaming search and lateral vector search for + primary-key tables are not supported. diff --git a/docs/sidebars.js b/docs/sidebars.js index 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-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/schema/SchemaValidation.java b/paimon-core/src/main/java/org/apache/paimon/schema/SchemaValidation.java index 2122c04ed515..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 @@ -912,13 +912,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/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/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-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-core/src/test/java/org/apache/paimon/schema/PrimaryKeyVectorIndexValidationTest.java b/paimon-core/src/test/java/org/apache/paimon/schema/PrimaryKeyVectorIndexValidationTest.java index 0daa7bc5c473..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 @@ -123,6 +123,33 @@ 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 testRejectsFirstRowWithDeletionVectors() { + Map options = enabledOptions(); + options.put(CoreOptions.MERGE_ENGINE.key(), "first-row"); + + assertThatThrownBy(() -> validateTableSchema(schema(options))) + .hasMessageContaining( + "First row merge engine does not need deletion vectors because there is no deletion of old data in this merge engine"); + } + @Test void testPartialUpdateRejectsDeletionVectorMergeOnRead() { Map 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..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 @@ -46,12 +46,19 @@ class PrimaryKeyVectorSearchTest extends TableTestBase { @Override protected Schema schemaDefault() { + return vectorSchema("deduplicate", true); + } + + private Schema vectorSchema(String mergeEngine, boolean deletionVectorsEnabled) { return Schema.newBuilder() .column("id", DataTypes.INT()) .column("embedding", DataTypes.VECTOR(2, DataTypes.FLOAT())) .primaryKey("id") .option(CoreOptions.BUCKET.key(), "1") - .option(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true") + .option(CoreOptions.MERGE_ENGINE.key(), mergeEngine) + .option( + CoreOptions.DELETION_VECTORS_ENABLED.key(), + Boolean.toString(deletionVectorsEnabled)) .option(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding") .option( "fields.embedding.pk-vector.index.type", @@ -94,4 +101,66 @@ void testVectorSearchMaterializesPhysicalRows() throws Exception { assertThat(ids).containsExactly(2, 3); } + + @Test + void testFirstRowVectorSearch() throws Exception { + catalog.createTable(identifier(), vectorSchema("first-row", false), false); + FileStoreTable table = getTableDefault(); + + write( + table, + ioManager, + GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[] {3, 0})), + GenericRow.of(2, BinaryVector.fromPrimitiveArray(new float[] {1, 0}))); + write( + table, + ioManager, + GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[] {0.5f, 0}))); + + GlobalIndexResult result = + table.newVectorSearchBuilder() + .withVectorColumn("embedding") + .withVector(new float[] {0, 0}) + .withLimit(1) + .executeLocal(); + ReadBuilder readBuilder = table.newReadBuilder(); + TableScan.Plan plan = readBuilder.newScan().withGlobalIndexResult(result).plan(); + List 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); + } } 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(); 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 87044862582e..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,8 +23,8 @@ import org.apache.paimon.table.source.VectorSearchBuilderImpl; /** - * Spark-aware {@link VectorSearchBuilderImpl} which produces a {@link SparkVectorReadImpl} so the - * per-split vector index evaluation is dispatched through Spark instead of the local thread pool. + * Spark-aware {@link VectorSearchBuilderImpl} which produces Spark-specific vector readers so + * data-evolution splits and primary-key bucket groups can be evaluated across the Spark cluster. * *

Single-vector only; batch search has no Spark-dispatched path yet (TODO). */ @@ -38,7 +38,10 @@ public SparkVectorSearchBuilderImpl(InnerTable table) { @Override public VectorRead newVectorRead() { - return new SparkVectorReadImpl( + if (isPrimaryKeyVectorSearch()) { + return new SparkPrimaryKeyVectorRead(table, vectorColumn, vector, limit, options); + } + return new SparkDataEvolutionVectorRead( table, partitionFilter, filter, limit, vectorColumn, vector, options); } } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala index 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-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 new file mode 100644 index 000000000000..84e896ae5006 --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PrimaryKeyVectorSearchTest.scala @@ -0,0 +1,270 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.spark.sql + +import org.apache.paimon.globalindex.testvector.TestVectorGlobalIndexerFactory +import org.apache.paimon.spark.PaimonSparkTestBase +import org.apache.paimon.spark.read.{SparkPrimaryKeyVectorRead, SparkVectorSearchBuilderImpl} +import org.apache.paimon.table.source.DataSplit + +import scala.collection.JavaConverters._ + +/** End-to-end tests for primary-key vector search through Spark SQL. */ +class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase { + + test("distributed primary-key vector search selects Spark reader") { + withTable("T") { + createVectorTable() + + val builder = new SparkVectorSearchBuilderImpl(loadTable("T")) + builder + .withVectorColumn("embedding") + .withVector(Array(0.0f, 0.0f)) + .withLimit(1) + + assert(builder.newVectorRead().isInstanceOf[SparkPrimaryKeyVectorRead]) + } + } + + test("distributed primary-key vector search evaluates buckets in Spark") { + withTable("T") { + createVectorTable(bucket = 2, extraOptions = Seq("global-index.thread-num" -> "1")) + spark.sql(""" + |INSERT INTO T VALUES + | (1, array(1.0f, 0.0f)), + | (2, array(2.0f, 0.0f)), + | (3, array(3.0f, 0.0f)), + | (4, array(4.0f, 0.0f)) + |""".stripMargin) + + val builder = new SparkVectorSearchBuilderImpl(loadTable("T")) + builder + .withVectorColumn("embedding") + .withVector(Array(0.0f, 0.0f)) + .withLimit(2) + + val jobGroup = s"primary-key-vector-${System.nanoTime()}" + spark.sparkContext.setJobGroup(jobGroup, jobGroup) + try { + builder.newVectorRead().read(builder.newVectorScan().scan()) + } finally { + spark.sparkContext.clearJobGroup() + } + + assert(spark.sparkContext.statusTracker.getJobIdsForGroup(jobGroup).nonEmpty) + } + } + + test("primary-key vector search uses bucket-local indexes") { + withTable("T") { + createVectorTable() + + spark.sql(""" + |INSERT INTO T VALUES + | (1, array(3.0f, 0.0f)), + | (2, array(1.0f, 0.0f)), + | (3, array(2.0f, 0.0f)) + |""".stripMargin) + + withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" -> "true") { + val rows = spark + .sql(""" + |SELECT id, __paimon_search_score + |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 2) + |""".stripMargin) + .collect() + + assert(rows.map(_.getInt(0)).toSet == Set(2, 3)) + assert(rows.forall(!_.isNullAt(1))) + + val scores = spark + .sql(""" + |SELECT __paimon_search_score + |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 2) + |""".stripMargin) + .collect() + assert(scores.length == 2) + assert(scores.forall(!_.isNullAt(0))) + } + } + } + + test("primary-key vector search merges top k across buckets") { + withTable("T") { + createVectorTable(bucket = 4, extraOptions = Seq("global-index.thread-num" -> "2")) + spark.sql(""" + |INSERT INTO T VALUES + | (1, array(1.0f, 0.0f)), + | (2, array(2.0f, 0.0f)), + | (3, array(3.0f, 0.0f)), + | (4, array(4.0f, 0.0f)), + | (5, array(5.0f, 0.0f)), + | (6, array(6.0f, 0.0f)), + | (7, array(7.0f, 0.0f)), + | (8, array(8.0f, 0.0f)), + | (9, array(9.0f, 0.0f)), + | (10, array(10.0f, 0.0f)), + | (11, array(11.0f, 0.0f)), + | (12, array(12.0f, 0.0f)), + | (13, array(13.0f, 0.0f)), + | (14, array(14.0f, 0.0f)), + | (15, array(15.0f, 0.0f)), + | (16, array(16.0f, 0.0f)) + |""".stripMargin) + + val buckets = loadTable("T") + .newReadBuilder() + .newScan() + .plan() + .splits() + .asScala + .map(_.asInstanceOf[DataSplit].bucket()) + .toSet + assert(buckets == Set(0, 1, 2, 3)) + + withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" -> "true") { + val ids = spark + .sql(""" + |SELECT id + |FROM vector_search('T', 'embedding', array(0.0f, 0.0f), 3) + |""".stripMargin) + .collect() + .map(_.getInt(0)) + .toSet + assert(ids == Set(1, 2, 3)) + } + } + } + + test("primary-key vector search prunes partitions before top k") { + withTable("T") { + createVectorTable( + columns = "id INT, embedding ARRAY, 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)) } } }