From 0479ef5dcc6147a601af9d70f701309dbde8637e Mon Sep 17 00:00:00 2001 From: tzphh <742429651@qq.com> Date: Thu, 22 Oct 2026 17:33:00 +0800 Subject: [PATCH 1/2] [core] Prepare PK vector index lifecycle for multiple columns --- .../pk/BucketedPrimaryKeyIndexMaintainer.java | 125 +++++++++------ ...BucketedPrimaryKeyIndexMaintainerTest.java | 143 +++++++++++++++++- .../operation/PrimaryKeyIndexWriteTest.java | 7 +- .../PrimaryKeyVectorIndexWriteTest.java | 14 +- 4 files changed, 230 insertions(+), 59 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainer.java b/paimon-core/src/main/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainer.java index 585c9a7b6c8a..44ff1c4be569 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainer.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainer.java @@ -40,6 +40,8 @@ import org.apache.paimon.io.KeyValueFileReaderFactory; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.types.DataField; +import org.apache.paimon.utils.ExceptionUtils; +import org.apache.paimon.utils.IOUtils; import javax.annotation.Nullable; @@ -55,16 +57,16 @@ /** Coordinates all source-backed primary-key indexes for one bucket. */ public final class BucketedPrimaryKeyIndexMaintainer { - @Nullable private final BucketedVectorIndexMaintainer vectorMaintainer; + private final List vectorMaintainers; @Nullable private final BucketedFullTextIndexMaintainer fullTextMaintainer; private final List sortedMaintainers; private int nextSortedMaintainerIndex; private BucketedPrimaryKeyIndexMaintainer( - @Nullable BucketedVectorIndexMaintainer vectorMaintainer, + List vectorMaintainers, @Nullable BucketedFullTextIndexMaintainer fullTextMaintainer, List sortedMaintainers) { - this.vectorMaintainer = vectorMaintainer; + this.vectorMaintainers = Collections.unmodifiableList(new ArrayList<>(vectorMaintainers)); this.fullTextMaintainer = fullTextMaintainer; List sorted = new ArrayList<>(sortedMaintainers); sorted.sort(Comparator.comparingInt(BucketedSortedIndexMaintainer::fieldId)); @@ -74,19 +76,19 @@ private BucketedPrimaryKeyIndexMaintainer( public static BucketedPrimaryKeyIndexMaintainer ofVector( BucketedVectorIndexMaintainer vectorMaintainer) { return new BucketedPrimaryKeyIndexMaintainer( - vectorMaintainer, null, Collections.emptyList()); + Collections.singletonList(vectorMaintainer), null, Collections.emptyList()); } public static BucketedPrimaryKeyIndexMaintainer ofFullText( BucketedFullTextIndexMaintainer fullTextMaintainer) { return new BucketedPrimaryKeyIndexMaintainer( - null, fullTextMaintainer, Collections.emptyList()); + Collections.emptyList(), fullTextMaintainer, Collections.emptyList()); } public static BucketedPrimaryKeyIndexMaintainer of( BucketedVectorIndexMaintainer vectorMaintainer, List sortedMaintainers) { - return new BucketedPrimaryKeyIndexMaintainer(vectorMaintainer, null, sortedMaintainers); + return of(vectorMaintainer, null, sortedMaintainers); } public static BucketedPrimaryKeyIndexMaintainer of( @@ -94,12 +96,25 @@ public static BucketedPrimaryKeyIndexMaintainer of( @Nullable BucketedFullTextIndexMaintainer fullTextMaintainer, List sortedMaintainers) { return new BucketedPrimaryKeyIndexMaintainer( - vectorMaintainer, fullTextMaintainer, sortedMaintainers); + vectorMaintainer == null + ? Collections.emptyList() + : Collections.singletonList(vectorMaintainer), + fullTextMaintainer, + sortedMaintainers); + } + + public static BucketedPrimaryKeyIndexMaintainer ofVectors( + List vectorMaintainers, + @Nullable BucketedFullTextIndexMaintainer fullTextMaintainer, + List sortedMaintainers) { + return new BucketedPrimaryKeyIndexMaintainer( + vectorMaintainers, fullTextMaintainer, sortedMaintainers); } public static BucketedPrimaryKeyIndexMaintainer ofSorted( List sortedMaintainers) { - return new BucketedPrimaryKeyIndexMaintainer(null, null, sortedMaintainers); + return new BucketedPrimaryKeyIndexMaintainer( + Collections.emptyList(), null, sortedMaintainers); } public synchronized void prepareCommit( @@ -107,14 +122,14 @@ public synchronized void prepareCommit( CompactIncrement compactIncrement, boolean waitCompaction) throws Exception { - BucketedVectorIndexMaintainer.VectorIndexCommit vectorCommit = null; + List vectorCommits = new ArrayList<>(); BucketedFullTextIndexMaintainer.FullTextIndexCommit fullTextCommit = null; List sortedCommits = new ArrayList<>(); try { - if (vectorMaintainer != null) { - vectorCommit = - vectorMaintainer.prepareCommit( - appendIncrement, compactIncrement, waitCompaction); + for (BucketedVectorIndexMaintainer maintainer : vectorMaintainers) { + vectorCommits.add( + maintainer.prepareCommit( + appendIncrement, compactIncrement, waitCompaction)); } if (fullTextMaintainer != null) { fullTextCommit = @@ -128,8 +143,8 @@ public synchronized void prepareCommit( prepareSortedNonBlocking(appendIncrement, compactIncrement, sortedCommits); } - if (vectorCommit != null) { - mergeVectorCommit(appendIncrement, compactIncrement, vectorCommit); + for (BucketedVectorIndexMaintainer.VectorIndexCommit commit : vectorCommits) { + mergeVectorCommit(appendIncrement, compactIncrement, commit); } if (fullTextCommit != null) { mergeFullTextCommit(appendIncrement, compactIncrement, fullTextCommit); @@ -144,8 +159,8 @@ public synchronized void prepareCommit( if (fullTextCommit != null) { fullTextCommit.abort(failure); } - if (vectorCommit != null) { - vectorCommit.abort(failure); + for (int i = vectorCommits.size() - 1; i >= 0; i--) { + vectorCommits.get(i).abort(failure); } if (failure instanceof Exception) { throw (Exception) failure; @@ -287,8 +302,10 @@ private static void applyIndexIncrement( } public boolean buildNotCompleted() { - if (vectorMaintainer != null && vectorMaintainer.buildNotCompleted()) { - return true; + for (BucketedVectorIndexMaintainer maintainer : vectorMaintainers) { + if (maintainer.buildNotCompleted()) { + return true; + } } if (fullTextMaintainer != null && fullTextMaintainer.buildNotCompleted()) { return true; @@ -297,8 +314,11 @@ public boolean buildNotCompleted() { } public void withExecutor(ExecutorService executor) { - if (vectorMaintainer != null) { - vectorMaintainer.withExecutor(executor); + checkArgument( + !buildNotCompleted(), + "Cannot replace executor during primary-key index maintenance."); + for (BucketedVectorIndexMaintainer maintainer : vectorMaintainers) { + maintainer.withExecutor(executor); } if (fullTextMaintainer != null) { fullTextMaintainer.withExecutor(executor); @@ -309,14 +329,20 @@ public void withExecutor(ExecutorService executor) { } public void close() { - if (vectorMaintainer != null) { - vectorMaintainer.close(); + List closeables = new ArrayList<>(); + for (BucketedVectorIndexMaintainer maintainer : vectorMaintainers) { + closeables.add(maintainer::close); } if (fullTextMaintainer != null) { - fullTextMaintainer.close(); + closeables.add(fullTextMaintainer::close); } for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { - maintainer.close(); + closeables.add(maintainer::close); + } + try { + IOUtils.closeAll(closeables); + } catch (Exception e) { + ExceptionUtils.rethrow(e); } } @@ -324,24 +350,27 @@ public void close() { public static final class Factory { private final IndexFileHandler handler; - @Nullable private final BucketedVectorIndexMaintainer.Factory vectorFactory; + private final List vectorFactories; @Nullable private final FullTextDefinitionFactory fullTextFactory; private final List sortedFactories; private Factory( IndexFileHandler handler, - @Nullable BucketedVectorIndexMaintainer.Factory vectorFactory, + List vectorFactories, @Nullable FullTextDefinitionFactory fullTextFactory, List sortedFactories) { this.handler = handler; - this.vectorFactory = vectorFactory; + this.vectorFactories = Collections.unmodifiableList(new ArrayList<>(vectorFactories)); this.fullTextFactory = fullTextFactory; this.sortedFactories = Collections.unmodifiableList(new ArrayList<>(sortedFactories)); } public static Factory ofVector(BucketedVectorIndexMaintainer.Factory vectorFactory) { return new Factory( - vectorFactory.indexFileHandler(), vectorFactory, null, Collections.emptyList()); + vectorFactory.indexFileHandler(), + Collections.singletonList(vectorFactory), + null, + Collections.emptyList()); } public static Factory create( @@ -350,7 +379,7 @@ public static Factory create( TableSchema schema) { CoreOptions coreOptions = new CoreOptions(schema.options()); Map fields = schema.nameToFieldMap(); - BucketedVectorIndexMaintainer.Factory vectorFactory = null; + List vectorFactories = new ArrayList<>(); FullTextDefinitionFactory fullTextFactory = null; List sortedFactories = new ArrayList<>(); for (PrimaryKeyIndexDefinition definition : @@ -362,10 +391,7 @@ public static Factory create( definition.column()); switch (definition.family()) { case VECTOR: - checkArgument( - vectorFactory == null, - "Only one primary-key vector index is supported."); - vectorFactory = + vectorFactories.add( new BucketedVectorIndexMaintainer.Factory( handler, readerFactoryBuilder, @@ -373,7 +399,7 @@ public static Factory create( coreOptions.primaryKeyVectorDistanceMetric( definition.column()), definition.indexType()) - .withIndexOptions(definition.options()); + .withIndexOptions(definition.options())); break; case BTREE: case BITMAP: @@ -409,9 +435,11 @@ public static Factory create( } } checkArgument( - vectorFactory != null || fullTextFactory != null || !sortedFactories.isEmpty(), + !vectorFactories.isEmpty() + || fullTextFactory != null + || !sortedFactories.isEmpty(), "No primary-key index definition is configured."); - return new Factory(handler, vectorFactory, fullTextFactory, sortedFactories); + return new Factory(handler, vectorFactories, fullTextFactory, sortedFactories); } public IndexFileHandler indexFileHandler() { @@ -456,16 +484,17 @@ public BucketedPrimaryKeyIndexMaintainer create( restoredDataFiles == null ? Collections.emptyList() : restoredDataFiles; List payloads = restoredPayloads == null ? Collections.emptyList() : restoredPayloads; - BucketedVectorIndexMaintainer vector = - vectorFactory == null - ? null - : vectorFactory.create( - partition, - bucket, - dataFiles, - payloads, - executor, - deletionVectorFactory); + List vectors = new ArrayList<>(); + for (BucketedVectorIndexMaintainer.Factory factory : vectorFactories) { + vectors.add( + factory.create( + partition, + bucket, + dataFiles, + payloads, + executor, + deletionVectorFactory)); + } BucketedFullTextIndexMaintainer fullText = fullTextFactory == null ? null @@ -478,7 +507,7 @@ public BucketedPrimaryKeyIndexMaintainer create( handler, partition, bucket, dataFiles, payloads, executor, ioManager)); } - return new BucketedPrimaryKeyIndexMaintainer(vector, fullText, sorted); + return new BucketedPrimaryKeyIndexMaintainer(vectors, fullText, sorted); } private static final class FullTextDefinitionFactory { diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainerTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainerTest.java index 531f440930c8..b029b42c492f 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainerTest.java @@ -31,11 +31,13 @@ import org.apache.paimon.index.pksorted.BucketedSortedIndexMaintainer; import org.apache.paimon.index.pksorted.PkSortedIndexFile; import org.apache.paimon.index.pkvector.BucketedVectorIndexMaintainer; +import org.apache.paimon.index.pkvector.PkVectorAnnSegmentFile; import org.apache.paimon.io.CompactIncrement; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.io.DataIncrement; import org.apache.paimon.io.KeyValueFileReaderFactory; import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.options.Options; import org.apache.paimon.schema.TableSchema; import org.apache.paimon.stats.SimpleStats; import org.apache.paimon.types.DataField; @@ -44,6 +46,7 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.mockito.InOrder; import java.util.ArrayList; import java.util.Arrays; @@ -64,8 +67,12 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Answers.RETURNS_SELF; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -82,18 +89,121 @@ void shutdownExecutor() { @Test void testDelegatesVectorLifecycle() { - BucketedVectorIndexMaintainer vector = mock(BucketedVectorIndexMaintainer.class); + BucketedVectorIndexMaintainer first = mock(BucketedVectorIndexMaintainer.class); + BucketedVectorIndexMaintainer second = mock(BucketedVectorIndexMaintainer.class); ExecutorService executor = mock(ExecutorService.class); - when(vector.buildNotCompleted()).thenReturn(true); + when(second.buildNotCompleted()).thenReturn(true); BucketedPrimaryKeyIndexMaintainer maintainer = - BucketedPrimaryKeyIndexMaintainer.ofVector(vector); + BucketedPrimaryKeyIndexMaintainer.ofVectors( + Arrays.asList(first, second), null, Collections.emptyList()); assertThat(maintainer.buildNotCompleted()).isTrue(); + assertThatThrownBy(() -> maintainer.withExecutor(executor)) + .hasMessageContaining("Cannot replace executor"); + verify(first, never()).withExecutor(any()); + verify(second, never()).withExecutor(any()); + + when(second.buildNotCompleted()).thenReturn(false); + assertThat(maintainer.buildNotCompleted()).isFalse(); maintainer.withExecutor(executor); maintainer.close(); - verify(vector).withExecutor(executor); - verify(vector).close(); + verify(first).withExecutor(executor); + verify(second).withExecutor(executor); + verify(first).close(); + verify(second).close(); + } + + @Test + void testClosesAllIndexesWhenVectorCloseFails() { + BucketedVectorIndexMaintainer first = mock(BucketedVectorIndexMaintainer.class); + BucketedVectorIndexMaintainer second = mock(BucketedVectorIndexMaintainer.class); + RuntimeException failure = new RuntimeException("first close failed"); + RuntimeException suppressed = new RuntimeException("second close failed"); + doThrow(failure).when(first).close(); + doThrow(suppressed).when(second).close(); + BucketedFullTextIndexMaintainer fullText = mock(BucketedFullTextIndexMaintainer.class); + BucketedSortedIndexMaintainer sorted = mock(BucketedSortedIndexMaintainer.class); + BucketedPrimaryKeyIndexMaintainer maintainer = + BucketedPrimaryKeyIndexMaintainer.ofVectors( + Arrays.asList(first, second), fullText, Collections.singletonList(sorted)); + + assertThatThrownBy(maintainer::close).isSameAs(failure).hasSuppressedException(suppressed); + verify(second).close(); + verify(fullText).close(); + verify(sorted).close(); + } + + @Test + void testVectorFailureRollsBackEarlierIndexAndAllowsRetry() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + IndexFileMeta firstPayload = payload("first", source, 3, 2, "vector"); + IndexFileMeta secondPayload = payload("second", source, 3, 1, "vector"); + PkVectorAnnSegmentFile firstFile = mock(PkVectorAnnSegmentFile.class); + PkVectorAnnSegmentFile secondFile = mock(PkVectorAnnSegmentFile.class); + when(firstFile.build(anyList(), any(), any(), eq("l2"), eq("vector"))) + .thenReturn(firstPayload); + when(secondFile.build(anyList(), any(), any(), eq("l2"), eq("vector"))) + .thenThrow(new IllegalStateException("second vector failed")) + .thenReturn(secondPayload); + BucketedVectorIndexMaintainer first = vectorMaintainer(2, firstFile); + BucketedVectorIndexMaintainer second = vectorMaintainer(1, secondFile); + BucketedPrimaryKeyIndexMaintainer maintainer = + BucketedPrimaryKeyIndexMaintainer.ofVectors( + Arrays.asList(first, second), null, Collections.emptyList()); + DataIncrement append = DataIncrement.emptyIncrement(); + CompactIncrement failed = compactAfter(source); + + assertThatThrownBy(() -> maintainer.prepareCommit(append, failed, true)) + .hasMessage("second vector failed"); + assertThat(append.newIndexFiles()).isEmpty(); + assertThat(append.deletedIndexFiles()).isEmpty(); + assertThat(failed.newIndexFiles()).isEmpty(); + assertThat(failed.deletedIndexFiles()).isEmpty(); + assertThat(first.segments()).isEmpty(); + assertThat(second.segments()).isEmpty(); + assertThat(maintainer.buildNotCompleted()).isFalse(); + verify(firstFile).delete(firstPayload); + + CompactIncrement retry = compactAfter(source); + maintainer.prepareCommit(DataIncrement.emptyIncrement(), retry, true); + + assertThat(retry.newIndexFiles()).containsExactly(firstPayload, secondPayload); + assertThat(first.segments()).containsExactly(firstPayload); + assertThat(second.segments()).containsExactly(secondPayload); + maintainer.close(); + } + + @Test + void testFullTextFailureAbortsAllPreparedVectorsInReverseOrder() throws Exception { + BucketedVectorIndexMaintainer first = mock(BucketedVectorIndexMaintainer.class); + BucketedVectorIndexMaintainer second = mock(BucketedVectorIndexMaintainer.class); + BucketedVectorIndexMaintainer.VectorIndexCommit firstCommit = + mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class); + BucketedVectorIndexMaintainer.VectorIndexCommit secondCommit = + mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class); + when(first.prepareCommit(any(), any(), eq(true))).thenReturn(firstCommit); + when(second.prepareCommit(any(), any(), eq(true))).thenReturn(secondCommit); + BucketedFullTextIndexMaintainer fullText = mock(BucketedFullTextIndexMaintainer.class); + IllegalStateException failure = new IllegalStateException("full-text failed"); + when(fullText.prepareCommit(any(), any(), eq(true))).thenThrow(failure); + BucketedPrimaryKeyIndexMaintainer maintainer = + BucketedPrimaryKeyIndexMaintainer.ofVectors( + Arrays.asList(first, second), fullText, Collections.emptyList()); + DataIncrement append = DataIncrement.emptyIncrement(); + CompactIncrement compact = CompactIncrement.emptyIncrement(); + + assertThatThrownBy(() -> maintainer.prepareCommit(append, compact, true)).isSameAs(failure); + InOrder order = inOrder(first, second, fullText, firstCommit, secondCommit); + order.verify(first).prepareCommit(append, compact, true); + order.verify(second).prepareCommit(append, compact, true); + order.verify(fullText).prepareCommit(append, compact, true); + order.verify(secondCommit).abort(failure); + order.verify(firstCommit).abort(failure); + verify(firstCommit, never()).appendIncrement(); + verify(secondCommit, never()).appendIncrement(); + assertThat(append.newIndexFiles()).isEmpty(); + assertThat(compact.newIndexFiles()).isEmpty(); } @Test @@ -119,6 +229,7 @@ void testDelegatesFullTextLifecycleAndMergesCommit() throws Exception { assertThat(compactIncrement.newIndexFiles()).containsExactly(payload); assertThat(maintainer.buildNotCompleted()).isTrue(); + when(fullText.buildNotCompleted()).thenReturn(false); maintainer.withExecutor(buildExecutor); maintainer.close(); verify(fullText).withExecutor(buildExecutor); @@ -399,6 +510,28 @@ void testBlockingPrepareRollsBackAllDefinitionsAfterInterruption() throws Except .containsExactly("btree-2", "bitmap-2"); } + private BucketedVectorIndexMaintainer vectorMaintainer( + int fieldId, PkVectorAnnSegmentFile indexFile) { + IndexFileHandler handler = mock(IndexFileHandler.class); + when(handler.pkVectorAnnSegment(any(), eq(0))).thenReturn(indexFile); + return new BucketedVectorIndexMaintainer.Factory( + handler, + mock(KeyValueFileReaderFactory.Builder.class, RETURNS_SELF), + new DataField( + fieldId, + "vector_" + fieldId, + DataTypes.VECTOR(2, DataTypes.FLOAT())), + "l2", + "vector") + .withIndexOptions(new Options()) + .create( + BinaryRow.EMPTY_ROW, + 0, + Collections.emptyList(), + Collections.emptyList(), + buildExecutor); + } + private BucketedSortedIndexMaintainer sortedMaintainer( int fieldId, String indexType, diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexWriteTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexWriteTest.java index 1c5e4587ffc4..994d59ef817f 100644 --- a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexWriteTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexWriteTest.java @@ -138,9 +138,10 @@ void testCreatesCoordinatorForMultipleScalarDefinitionsAndVector() throws Except AbstractFileStoreWrite.WriterContainer container = write.createWriterContainer(generator.getPartition(record), 1); - Object vectorMaintainer = - readField(container.primaryKeyIndexMaintainer, "vectorMaintainer"); - assertThat(vectorMaintainer).isNotNull(); + List vectorMaintainers = + (List) readField(container.primaryKeyIndexMaintainer, "vectorMaintainers"); + assertThat(vectorMaintainers).hasSize(1); + Object vectorMaintainer = vectorMaintainers.get(0); assertThat((List) readField(container.primaryKeyIndexMaintainer, "sortedMaintainers")) .hasSize(4); diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyVectorIndexWriteTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyVectorIndexWriteTest.java index 9b2aa3dcb1eb..34e6cbad340e 100644 --- a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyVectorIndexWriteTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyVectorIndexWriteTest.java @@ -37,6 +37,7 @@ import javax.annotation.Nullable; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -71,15 +72,19 @@ void testRestoreRebindsVectorIndexExecutor() throws Exception { BucketedVectorIndexMaintainer maintainer = mock(BucketedVectorIndexMaintainer.class); when(maintainer.prepareCommit(same(dataIncrement), same(compactIncrement), eq(true))) .thenReturn(vectorCommit); + BucketedVectorIndexMaintainer secondMaintainer = mock(BucketedVectorIndexMaintainer.class); + when(secondMaintainer.prepareCommit(same(dataIncrement), same(compactIncrement), eq(true))) + .thenReturn(vectorCommit); TestingFileStoreWrite write = new TestingFileStoreWrite(); - write.install(writer, maintainer); + write.install(writer, maintainer, secondMaintainer); List> states = write.checkpoint(); TestingFileStoreWrite restored = new TestingFileStoreWrite(mock(RecordWriter.class)); restored.restore(states); verify(maintainer).withExecutor(any(ExecutorService.class)); + verify(secondMaintainer).withExecutor(any(ExecutorService.class)); } @Test @@ -160,7 +165,7 @@ private TestingFileStoreWrite(@Nullable RecordWriter restoredWriter) { } private void install( - RecordWriter writer, BucketedVectorIndexMaintainer vectorMaintainer) { + RecordWriter writer, BucketedVectorIndexMaintainer... vectorMaintainers) { writers.computeIfAbsent(EMPTY_ROW, ignored -> new HashMap<>()) .put( 0, @@ -169,7 +174,10 @@ private void install( 1, null, null, - BucketedPrimaryKeyIndexMaintainer.ofVector(vectorMaintainer), + BucketedPrimaryKeyIndexMaintainer.ofVectors( + Arrays.asList(vectorMaintainers), + null, + Collections.emptyList()), null)); } From 7136275025055125fd6bfcd8437c03e08c881b0d Mon Sep 17 00:00:00 2001 From: tzphh <742429651@qq.com> Date: Thu, 22 Oct 2026 18:02:02 +0800 Subject: [PATCH 2/2] [core] Enable multiple vector index columns on primary-key tables --- docs/docs/primary-key-table/global-index.mdx | 36 +++- docs/generated/core_configuration.html | 2 +- .../java/org/apache/paimon/CoreOptions.java | 12 +- .../TestVectorGlobalIndexerFactory.java | 5 + .../paimon/schema/SchemaValidation.java | 66 +++--- .../PrimaryKeyVectorIndexOptionsTest.java | 22 +- .../PrimaryKeyVectorIndexValidationTest.java | 118 ++++++++-- .../source/PrimaryKeyVectorSearchTest.java | 202 ++++++++++++++++++ .../VectorSearchProcedureITCase.java | 83 +++++++ .../sql/PrimaryKeyVectorSearchTest.scala | 70 +++++- 10 files changed, 550 insertions(+), 66 deletions(-) diff --git a/docs/docs/primary-key-table/global-index.mdx b/docs/docs/primary-key-table/global-index.mdx index a5c2a8d986d3..c0eeb9f7b214 100644 --- a/docs/docs/primary-key-table/global-index.mdx +++ b/docs/docs/primary-key-table/global-index.mdx @@ -104,7 +104,8 @@ Vector additionally requires: - `deletion-vectors.enabled = true`, except for `first-row`, where it must be `false`. - The configured ANN implementation on every writer and reader classpath. -Exactly one vector column is currently supported per table. +Multiple vector columns are supported. Each column owns an independent index with its own +dimension, ANN implementation, distance metric, and build options. @@ -265,11 +266,40 @@ schema validation. +### Index Multiple Vector Columns + +List the vector columns in `pk-vector.index.columns` and configure each column separately. +For example, this Flink table indexes image and text embeddings independently: + +```sql +CREATE TABLE multimodal_items ( + id BIGINT, + image_embedding ARRAY COMMENT '__VECTOR_FIELD;512', + text_embedding ARRAY COMMENT '__VECTOR_FIELD;768', + PRIMARY KEY (id) NOT ENFORCED +) WITH ( + 'bucket' = '8', + 'deletion-vectors.enabled' = 'true', + 'pk-vector.index.columns' = 'image_embedding,text_embedding', + 'fields.image_embedding.pk-vector.index.type' = 'ivf-flat', + 'fields.image_embedding.pk-vector.distance.metric' = 'cosine', + 'fields.image_embedding.pk-vector.index.options' = '{"nlist":"256"}', + 'fields.text_embedding.pk-vector.index.type' = 'ivf-flat', + 'fields.text_embedding.pk-vector.distance.metric' = 'inner_product', + 'fields.text_embedding.pk-vector.index.options' = '{"nlist":"512"}' +); +``` + +Each vector search selects one column and uses that column's index, dimension, and distance +metric. Listing multiple columns does not combine their vectors or define a fused ranking. +Index coverage is tracked separately for each column, so an index on `image_embedding` does +not cover a query on `text_embedding`. + ### Options | Option | Default | Description | |---|---|---| -| `pk-vector.index.columns` | Not set | Vector column to index. Exactly one vector column is currently supported. | +| `pk-vector.index.columns` | Not set | Comma-separated vector columns which own independent vector indexes. | | `fields..pk-vector.index.type` | Required | ANN implementation, such as `ivf-flat`, `ivf-pq`, `ivf-sq`, `ivf-rq`, or `diskann`. | | `fields..pk-vector.distance.metric` | `inner_product` | Distance metric: `l2`, `cosine`, or `inner_product`. | | `fields..pk-vector.index.options` | Not set | JSON object containing build options for the selected ANN implementation. | @@ -610,7 +640,7 @@ create a table with the desired definition and migrate the data. - Multivalue supports arrays of scalar element types accepted by the global-index key serializer. Null elements are not indexed. Higher array and element cardinality increases the expanded-sort I/O and index-file size. -- Exactly one vector index column is currently supported per table. +- Vector indexes are single-column indexes; each configured vector column has an independent index. - Exactly one full-text index column is currently supported per table. - Only `FLOAT` vectors are supported. - Indexes are built from eligible compact output, not directly from Level-0 appends. diff --git a/docs/generated/core_configuration.html b/docs/generated/core_configuration.html index 601ed2a35884..069c4d15cdb3 100644 --- a/docs/generated/core_configuration.html +++ b/docs/generated/core_configuration.html @@ -1402,7 +1402,7 @@
pk-vector.index.columns
(none) String - Comma-separated VECTOR columns indexed by primary-key vector indexes. Each column owns one index and must define fields.<column>.pk-vector.index.type. Index options and distance metric are also field-scoped. The first release supports exactly one column. + Comma-separated VECTOR columns indexed by primary-key vector indexes. Each column owns one index and must define fields.<column>.pk-vector.index.type. Index options and distance metric are also field-scoped.
postpone.batch-write-fixed-bucket
diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java index 4e8a1cec469f..2f14a68b19a1 100644 --- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java +++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java @@ -3171,8 +3171,7 @@ public String toString() { "Comma-separated VECTOR columns indexed by primary-key vector indexes. " + "Each column owns one index and must define " + "fields..pk-vector.index.type. Index options and distance " - + "metric are also field-scoped. The first release supports exactly " - + "one column."); + + "metric are also field-scoped."); public static final ConfigOption PK_BTREE_INDEX_COLUMNS = key("pk-btree.index.columns") @@ -5122,15 +5121,6 @@ private Options primaryKeySortedIndexOptions( return resolved; } - public String primaryKeyVectorIndexColumn() { - List columns = primaryKeyVectorIndexColumns(); - checkArgument( - columns.size() == 1, - "pk-vector.index.columns must contain exactly one column in the first release, but is %s.", - columns); - return columns.get(0); - } - @Nullable public String primaryKeyVectorIndexType(String column) { return options.get("fields." + column + ".pk-vector.index.type"); diff --git a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexerFactory.java b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexerFactory.java index e87a8eec3246..67fa339cb962 100644 --- a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexerFactory.java +++ b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexerFactory.java @@ -40,6 +40,11 @@ public String identifier() { @Override public GlobalIndexer create(List indexFields, Options options) { + String metric = options.getString(IDENTIFIER + ".metric", null); + if (metric != null) { + options = new Options(options.toMap()); + options.setString(TestVectorGlobalIndexer.OPT_METRIC, metric); + } if (indexFields.size() == 1) { return new TestVectorGlobalIndexer(indexFields.get(0).type(), options); } 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 9aaca5f6d9d2..e30c8e39f330 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 @@ -1235,19 +1235,6 @@ private static void validatePrimaryKeyVectorIndex(TableSchema schema, CoreOption new HashSet<>(indexColumns).size() == indexColumns.size(), "pk-vector.index.columns must not contain duplicate columns, but is %s.", indexColumns); - checkArgument( - indexColumns.size() == 1, - "pk-vector.index.columns must contain exactly one column in the first release, but is %s.", - indexColumns); - String indexColumn = indexColumns.get(0); - String indexType = options.primaryKeyVectorIndexType(indexColumn); - checkArgument( - !StringUtils.isNullOrWhitespaceOnly(indexColumn), - "pk-vector.index.columns must contain a non-empty column."); - checkArgument( - !StringUtils.isNullOrWhitespaceOnly(indexType), - "fields.%s.pk-vector.index.type must be configured when a primary-key vector index is defined.", - indexColumn); checkArgument( !schema.primaryKeys().isEmpty(), "Primary-key vector index requires a primary-key table."); @@ -1266,28 +1253,39 @@ private static void validatePrimaryKeyVectorIndex(TableSchema schema, CoreOption checkArgument( !options.pkClusteringOverride(), "Primary-key vector index does not support pk-clustering-override."); - options.primaryKeyVectorIndexOptions(indexColumn); - DataField vectorField = - schema.fields().stream() - .filter(field -> field.name().equals(indexColumn)) - .findFirst() - .orElse(null); - checkArgument( - vectorField != null && vectorField.type().getTypeRoot() == VECTOR, - "pk-vector.index.columns entry '%s' must reference a VECTOR column.", - indexColumn); - checkArgument( - ((VectorType) vectorField.type()).getElementType().getTypeRoot() - == DataTypeRoot.FLOAT, - "pk-vector.index.columns entry '%s' must use FLOAT elements.", - indexColumn); - checkArgument( - Arrays.asList("l2", "cosine", "inner_product") - .contains(options.primaryKeyVectorDistanceMetric(indexColumn)), - "fields.%s.pk-vector.distance.metric must be one of l2, cosine, inner_product, but is %s.", - indexColumn, - options.primaryKeyVectorDistanceMetric(indexColumn)); + for (String indexColumn : indexColumns) { + checkArgument( + !StringUtils.isNullOrWhitespaceOnly(indexColumn), + "pk-vector.index.columns must contain a non-empty column."); + String indexType = options.primaryKeyVectorIndexType(indexColumn); + checkArgument( + !StringUtils.isNullOrWhitespaceOnly(indexType), + "fields.%s.pk-vector.index.type must be configured when a primary-key vector index is defined.", + indexColumn); + options.primaryKeyVectorIndexOptions(indexColumn); + + DataField vectorField = + schema.fields().stream() + .filter(field -> field.name().equals(indexColumn)) + .findFirst() + .orElse(null); + checkArgument( + vectorField != null && vectorField.type().getTypeRoot() == VECTOR, + "pk-vector.index.columns entry '%s' must reference a VECTOR column.", + indexColumn); + checkArgument( + ((VectorType) vectorField.type()).getElementType().getTypeRoot() + == DataTypeRoot.FLOAT, + "pk-vector.index.columns entry '%s' must use FLOAT elements.", + indexColumn); + checkArgument( + Arrays.asList("l2", "cosine", "inner_product") + .contains(options.primaryKeyVectorDistanceMetric(indexColumn)), + "fields.%s.pk-vector.distance.metric must be one of l2, cosine, inner_product, but is %s.", + indexColumn, + options.primaryKeyVectorDistanceMetric(indexColumn)); + } } private static void validatePrimaryKeyFullTextIndex(TableSchema schema, CoreOptions options) { diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorIndexOptionsTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorIndexOptionsTest.java index 449d640aa3f6..9c634c1595f3 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorIndexOptionsTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorIndexOptionsTest.java @@ -41,11 +41,25 @@ void testPluralFieldRegistryEnablesIndex() { } @Test - void testResolvesSingleIndexColumn() { + void testResolvesIndependentOptionsForSameAlgorithm() { Map options = new HashMap<>(); - options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding"); - - assertThat(new CoreOptions(options).primaryKeyVectorIndexColumn()).isEqualTo("embedding"); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), " embedding , other_embedding "); + options.put("fields.embedding.pk-vector.index.type", "ivf-flat"); + options.put("fields.other_embedding.pk-vector.index.type", "ivf-flat"); + options.put("fields.embedding.pk-vector.distance.metric", "cosine"); + options.put("fields.other_embedding.pk-vector.distance.metric", "l2"); + options.put("fields.embedding.pk-vector.index.options", "{\"nlist\":64}"); + options.put("fields.other_embedding.pk-vector.index.options", "{\"nlist\":128}"); + CoreOptions coreOptions = new CoreOptions(options); + + assertThat(coreOptions.primaryKeyVectorIndexColumns()) + .containsExactly("embedding", "other_embedding"); + Options first = coreOptions.primaryKeyVectorIndexOptions("embedding"); + Options second = coreOptions.primaryKeyVectorIndexOptions("other_embedding"); + assertThat(first.get("ivf-flat.nlist")).isEqualTo("64"); + assertThat(first.get("ivf-flat.metric")).isEqualTo("cosine"); + assertThat(second.get("ivf-flat.nlist")).isEqualTo("128"); + assertThat(second.get("ivf-flat.metric")).isEqualTo("l2"); } @Test 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 748b697a2641..495e934d422b 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 @@ -23,10 +23,14 @@ import org.apache.paimon.types.DataTypes; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import static org.apache.paimon.schema.SchemaValidation.validateTableSchema; @@ -48,16 +52,94 @@ void testValidPluralPrimaryKeyVectorIndexConfiguration() { } @Test - void testRejectsMultiplePrimaryKeyVectorIndexColumnsForFirstRelease() { - Map options = new HashMap<>(); - options.put(CoreOptions.BUCKET.key(), "1"); - options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); - options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding,other_embedding"); - options.put("fields.embedding.pk-vector.index.type", "ivf-pq"); + void testSupportsMultiplePrimaryKeyVectorIndexColumns() { + Map options = multipleIndexOptions(); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), " embedding , other_embedding "); + options.put("fields.embedding.pk-vector.distance.metric", "cosine"); + options.put("fields.embedding.pk-vector.index.options", "{\"nlist\":64}"); + options.put("fields.other_embedding.pk-vector.distance.metric", "l2"); + options.put("fields.other_embedding.pk-vector.index.options", "{\"nlist\":128}"); + + assertThatCode(() -> validateTableSchema(schema(options))).doesNotThrowAnyException(); + } + + @Test + void testRequiresIndexTypeForSecondVectorColumn() { + Map options = multipleIndexOptions(); + options.remove("fields.other_embedding.pk-vector.index.type"); + + assertThatThrownBy(() -> validateTableSchema(schema(options))) + .hasMessageContaining("fields.other_embedding.pk-vector.index.type"); + } + + @Test + void testRequiresSecondColumnToReferenceVector() { + Map options = enabledOptions(); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding,payload"); + options.put("fields.payload.pk-vector.index.type", "ivf-flat"); + + assertThatThrownBy(() -> validateTableSchema(schema(options))) + .hasMessageContaining("entry 'payload' must reference a VECTOR column"); + } + + @Test + void testRejectsUnknownSecondVectorColumn() { + Map options = enabledOptions(); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding,missing"); + options.put("fields.missing.pk-vector.index.type", "ivf-flat"); assertThatThrownBy(() -> validateTableSchema(schema(options))) - .hasMessageContaining("exactly one column") - .hasMessageContaining("pk-vector.index.columns"); + .hasMessageContaining("Field missing can not be found in table schema"); + } + + @ParameterizedTest + @ValueSource(strings = {"", "embedding,"}) + void testRejectsEmptyVectorIndexColumn(String columns) { + Map options = enabledOptions(); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), columns); + + assertThatThrownBy(() -> validateTableSchema(schema(options))) + .hasMessageContaining("pk-vector.index.columns") + .hasMessageContaining("non-empty"); + } + + @Test + void testRequiresFloatElementsForSecondVectorColumn() { + List fields = new ArrayList<>(fields()); + fields.set(3, new DataField(3, "other_embedding", DataTypes.VECTOR(3, DataTypes.DOUBLE()))); + + assertThatThrownBy(() -> validateTableSchema(schema(multipleIndexOptions(), fields))) + .hasMessageContaining("entry 'other_embedding' must use FLOAT elements"); + } + + @Test + void testRejectsUnsupportedDistanceMetricForSecondVectorColumn() { + Map options = multipleIndexOptions(); + options.put("fields.other_embedding.pk-vector.distance.metric", "manhattan"); + + assertThatThrownBy(() -> validateTableSchema(schema(options))) + .hasMessageContaining("fields.other_embedding.pk-vector.distance.metric") + .hasMessageContaining("l2, cosine, inner_product"); + } + + @Test + void testRejectsInvalidIndexOptionsForSecondVectorColumn() { + Map options = multipleIndexOptions(); + options.put("fields.other_embedding.pk-vector.index.options", "[1,2]"); + + assertThatThrownBy(() -> validateTableSchema(schema(options))) + .hasMessageContaining("fields.other_embedding.pk-vector.index.options") + .hasMessageContaining("JSON object"); + } + + @Test + void testRejectsSecondVectorColumnInAnotherIndexFamily() { + Map options = multipleIndexOptions(); + options.put(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "other_embedding"); + + assertThatThrownBy(() -> validateTableSchema(schema(options))) + .hasMessageContaining("other_embedding") + .hasMessageContaining("at most one primary-key index"); } @Test @@ -65,7 +147,7 @@ void testRejectsDuplicatePrimaryKeyVectorIndexColumns() { Map options = new HashMap<>(); options.put(CoreOptions.BUCKET.key(), "1"); options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); - options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding,embedding"); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding, embedding "); assertThatThrownBy(() -> validateTableSchema(schema(options))) .hasMessageContaining("pk-vector.index.columns") @@ -253,9 +335,20 @@ private static Map enabledOptions() { } private static TableSchema schema(Map options) { + return schema(options, fields()); + } + + private static Map multipleIndexOptions() { + Map options = enabledOptions(); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding,other_embedding"); + options.put("fields.other_embedding.pk-vector.index.type", "ivf-flat"); + return options; + } + + private static TableSchema schema(Map options, List fields) { return new TableSchema( 0, - fields(), + fields, 0, Collections.emptyList(), Collections.singletonList("id"), @@ -263,10 +356,11 @@ private static TableSchema schema(Map options) { ""); } - private static java.util.List fields() { + private static List fields() { return Arrays.asList( new DataField(0, "id", DataTypes.INT().notNull()), new DataField(1, "embedding", DataTypes.VECTOR(8, DataTypes.FLOAT())), - new DataField(2, "payload", DataTypes.STRING())); + new DataField(2, "payload", DataTypes.STRING()), + new DataField(3, "other_embedding", DataTypes.VECTOR(3, DataTypes.FLOAT()))); } } 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 06d376be2626..5594f4161d8a 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 @@ -26,21 +26,28 @@ import org.apache.paimon.globalindex.GlobalIndexResult; import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.globalindex.testvector.TestVectorGlobalIndexerFactory; +import org.apache.paimon.index.IndexFileMeta; import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.Schema; +import org.apache.paimon.schema.SchemaChange; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.TableTestBase; import org.apache.paimon.table.sink.BatchTableCommit; import org.apache.paimon.table.sink.BatchTableWrite; import org.apache.paimon.table.sink.BatchWriteBuilder; +import org.apache.paimon.table.sink.CommitMessage; +import org.apache.paimon.table.sink.CommitMessageImpl; import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowKind; import org.junit.jupiter.api.Test; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; +import static java.util.stream.Collectors.toList; import static org.assertj.core.api.Assertions.assertThat; /** End-to-end tests for bucket-local primary-key vector search. */ @@ -80,6 +87,201 @@ private Schema vectorSchema( return builder.build(); } + @Test + void testIndependentVectorIndexesSurviveReopenAndCompaction() throws Exception { + catalog.createTable(identifier(), multipleVectorSchema(2, "l2", "l2"), false); + FileStoreTable table = getTableDefault(); + write( + table, + ioManager, + GenericRow.of(1, vector(1, 0), vector(9, 0)), + GenericRow.of(2, vector(2, 0), vector(8, 0))); + write( + table, + ioManager, + GenericRow.of(3, vector(8, 0), vector(2, 0)), + GenericRow.of(4, vector(9, 0), vector(1, 0))); + compact(table, BinaryRow.EMPTY_ROW, 0, ioManager, true); + assertMultipleVectorPayloads(table); + + // Identical algorithms and dimensions must not cause one field's payload to serve another. + FileStoreTable reopened = getTableDefault(); + assertThat(searchIds(reopened, "image_embedding", new float[] {0, 0}, 1)) + .containsExactly(1); + assertThat(searchIds(reopened, "text_embedding", new float[] {0, 0}, 1)).containsExactly(4); + + // A new writer restores both indexes before replacing their source files. + write( + reopened, + ioManager, + GenericRow.of(1, vector(10, 0), vector(0.5f, 0)), + GenericRow.of(2, null, vector(7, 0)), + GenericRow.of(3, vector(8, 0), null), + GenericRow.ofKind(RowKind.DELETE, 4, vector(9, 0), vector(1, 0))); + compact(reopened, BinaryRow.EMPTY_ROW, 0, ioManager, true); + assertMultipleVectorPayloads(reopened); + + FileStoreTable restored = getTableDefault(); + assertThat(searchIds(restored, "image_embedding", new float[] {0, 0}, 1)) + .containsExactly(3); + assertThat(searchIds(restored, "text_embedding", new float[] {0, 0}, 1)).containsExactly(1); + assertThat(searchIds(restored, "image_embedding", new float[] {0, 0}, 10)) + .containsExactlyInAnyOrder(1, 3); + assertThat(searchIds(restored, "text_embedding", new float[] {0, 0}, 10)) + .containsExactlyInAnyOrder(1, 2); + } + + @Test + void testIndependentVectorIndexesWithDifferentDimensionsAndMetrics() throws Exception { + catalog.createTable(identifier(), multipleVectorSchema(3, "l2", "cosine"), false); + FileStoreTable table = getTableDefault(); + write(table, ioManager, GenericRow.of(1, vector(1, 0), vector(1, 1, 0))); + write(table, ioManager, GenericRow.of(2, vector(2, 0), vector(10, 0, 0))); + compact(table, BinaryRow.EMPTY_ROW, 0, ioManager, true); + assertMultipleVectorPayloads(table); + + FileStoreTable reopened = getTableDefault(); + assertThat(searchIds(reopened, "image_embedding", new float[] {0, 0}, 1)) + .containsExactly(1); + // Cosine ranks id 2 first, whereas L2 would rank id 1 first for this field. + assertThat(searchIds(reopened, "text_embedding", new float[] {1, 0, 0}, 1)) + .containsExactly(2); + + List results = + reopened.newBatchVectorSearchBuilder() + .withVectorColumn("text_embedding") + .withVectors(new float[][] {{1, 0, 0}, {0, 1, 0}}) + .withLimit(1) + .executeBatchLocal(); + assertThat(results).hasSize(2); + assertThat(readIds(reopened, results.get(0))).containsExactly(2); + assertThat(readIds(reopened, results.get(1))).containsExactly(1); + } + + @Test + void testCompactionBackfillsAddedVectorIndexWithoutNewData() throws Exception { + Schema schema = multipleVectorSchema(2, "l2", "l2"); + schema.options().put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "image_embedding"); + schema.options().remove("fields.text_embedding.pk-vector.index.type"); + catalog.createTable(identifier(), schema, false); + FileStoreTable original = getTableDefault(); + write(original, ioManager, GenericRow.of(1, vector(1, 0), vector(9, 0))); + write(original, ioManager, GenericRow.of(2, vector(9, 0), vector(1, 0))); + compact(original, BinaryRow.EMPTY_ROW, 0, ioManager, true); + List originalPayloads = vectorPayloads(original); + assertThat(originalPayloads).hasSize(1); + assertThat(originalPayloads.get(0).globalIndexMeta().indexFieldId()) + .isEqualTo(original.rowType().getField("image_embedding").id()); + + catalog.alterTable( + identifier(), + Arrays.asList( + SchemaChange.setOption( + CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), + "image_embedding,text_embedding"), + SchemaChange.setOption( + "fields.text_embedding.pk-vector.index.type", + TestVectorGlobalIndexerFactory.IDENTIFIER)), + false); + FileStoreTable reopened = getTableDefault(); + assertThat(searchIds(reopened, "image_embedding", new float[] {0, 0}, 1)) + .containsExactly(1); + assertThat(searchIds(reopened, "text_embedding", new float[] {0, 0}, 1)).isEmpty(); + + // A fresh compactor must discover the coverage gap even without a data-file transition. + BatchWriteBuilder writeBuilder = reopened.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + write.withIOManager(ioManager); + write.compact(BinaryRow.EMPTY_ROW, 0, true); + List messages = write.prepareCommit(); + assertThat(messages).hasSize(1); + CommitMessageImpl message = (CommitMessageImpl) messages.get(0); + assertThat(message.newFilesIncrement().newFiles()).isEmpty(); + assertThat(message.compactIncrement().compactBefore()).isEmpty(); + assertThat(message.compactIncrement().compactAfter()).isEmpty(); + assertThat(message.newFilesIncrement().newIndexFiles()) + .extracting(payload -> payload.globalIndexMeta().indexFieldId()) + .containsExactly(reopened.rowType().getField("text_embedding").id()); + assertThat(message.newFilesIncrement().deletedIndexFiles()).isEmpty(); + assertThat(message.compactIncrement().deletedIndexFiles()).isEmpty(); + commit.commit(messages); + } + + FileStoreTable restored = getTableDefault(); + assertMultipleVectorPayloads(restored); + List completedPayloads = vectorPayloads(restored); + assertThat(completedPayloads) + .extracting(IndexFileMeta::fileName) + .contains(originalPayloads.get(0).fileName()); + assertThat(searchIds(restored, "image_embedding", new float[] {0, 0}, 1)) + .containsExactly(1); + assertThat(searchIds(restored, "text_embedding", new float[] {0, 0}, 1)).containsExactly(2); + + compact(restored, BinaryRow.EMPTY_ROW, 0, ioManager, true); + assertThat(vectorPayloads(getTableDefault())) + .extracting(IndexFileMeta::fileName) + .containsExactlyInAnyOrderElementsOf( + completedPayloads.stream().map(IndexFileMeta::fileName).collect(toList())); + } + + private Schema multipleVectorSchema(int textDimension, String imageMetric, String textMetric) { + return Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("image_embedding", DataTypes.VECTOR(2, DataTypes.FLOAT())) + .column("text_embedding", DataTypes.VECTOR(textDimension, DataTypes.FLOAT())) + .primaryKey("id") + .option(CoreOptions.BUCKET.key(), "1") + .option(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true") + .option(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "image_embedding,text_embedding") + .option( + "fields.image_embedding.pk-vector.index.type", + TestVectorGlobalIndexerFactory.IDENTIFIER) + .option( + "fields.text_embedding.pk-vector.index.type", + TestVectorGlobalIndexerFactory.IDENTIFIER) + .option("fields.image_embedding.pk-vector.distance.metric", imageMetric) + .option("fields.text_embedding.pk-vector.distance.metric", textMetric) + .option("fields.image_embedding.vector-index.search-mode", "fast") + .option("fields.text_embedding.vector-index.search-mode", "fast") + .build(); + } + + private static BinaryVector vector(float... values) { + return BinaryVector.fromPrimitiveArray(values); + } + + private static void assertMultipleVectorPayloads(FileStoreTable table) { + List payloads = vectorPayloads(table); + assertThat(payloads) + .extracting(payload -> payload.globalIndexMeta().indexFieldId()) + .containsExactlyInAnyOrder( + table.rowType().getField("image_embedding").id(), + table.rowType().getField("text_embedding").id()); + assertThat(payloads) + .extracting(IndexFileMeta::indexType) + .containsOnly(TestVectorGlobalIndexerFactory.IDENTIFIER); + assertThat(payloads).extracting(IndexFileMeta::fileName).doesNotHaveDuplicates(); + } + + private static List vectorPayloads(FileStoreTable table) { + return table.store() + .newIndexFileHandler() + .scanSourceIndexes( + table.store().snapshotManager().latestSnapshot(), BinaryRow.EMPTY_ROW, 0); + } + + private static List searchIds( + FileStoreTable table, String field, float[] query, int limit) throws Exception { + return readIds( + table, + table.newVectorSearchBuilder() + .withVectorColumn(field) + .withVector(query) + .withLimit(limit) + .executeLocal()); + } + @Test void testRefineFactorReranksAnnCandidates() throws Exception { catalog.createTable(identifier(), vectorSchema("deduplicate", true, true), false); 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 544836ee88df..8a5f171ab658 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 @@ -48,6 +48,8 @@ import org.apache.paimon.types.DataField; import org.apache.paimon.utils.Range; +import org.apache.flink.client.program.ClusterClient; +import org.apache.flink.runtime.client.JobStatusMessage; import org.apache.flink.table.api.config.TableConfigOptions; import org.apache.flink.types.Row; import org.junit.jupiter.api.Test; @@ -55,6 +57,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -199,6 +202,86 @@ public void testDistributedPrimaryKeyVectorSearch() { && !row.contains("\"vec\"")); } + @Test + public void testDistributedPrimaryKeySearchWithMultipleVectorIndexes() throws Exception { + sql( + "CREATE TABLE MULTI_VECTOR_PK_T (" + + "id INT, image_vec ARRAY, text_vec ARRAY, " + + "PRIMARY KEY (id) NOT ENFORCED) WITH (" + + "'bucket' = '4', " + + "'file.format' = 'json', " + + "'file.compression' = 'none', " + + "'deletion-vectors.enabled' = 'true', " + + "'vector-field' = 'image_vec,text_vec', " + + "'field.image_vec.vector-dim' = '2', " + + "'field.text_vec.vector-dim' = '2', " + + "'pk-vector.index.columns' = 'image_vec,text_vec', " + + "'fields.image_vec.pk-vector.index.type' = '%s', " + + "'fields.text_vec.pk-vector.index.type' = '%s', " + + "'fields.image_vec.pk-vector.distance.metric' = 'l2', " + + "'fields.text_vec.pk-vector.distance.metric' = 'l2', " + + "'vector-index.search-mode' = 'fast', " + + "'vector-search.distribute.enabled' = 'true', " + + "'global-index.thread-num' = '2')", + TestVectorGlobalIndexerFactory.IDENTIFIER, + TestVectorGlobalIndexerFactory.IDENTIFIER); + String values = + IntStream.range(0, 32) + .mapToObj( + i -> + String.format( + "(%d, ARRAY[CAST(%d AS FLOAT), CAST(0 AS FLOAT)], " + + "ARRAY[CAST(%d AS FLOAT), CAST(0 AS FLOAT)])", + i, i, 31 - i)) + .collect(Collectors.joining(",")); + sql("INSERT INTO MULTI_VECTOR_PK_T VALUES " + values); + sql( + "CALL sys.compact(`table` => 'default.MULTI_VECTOR_PK_T', " + + "compact_strategy => 'full', " + + "options => 'compaction.force-rewrite-all-files=true')"); + assertThat(sql("SELECT DISTINCT bucket FROM `MULTI_VECTOR_PK_T$files`")).hasSize(4); + FileStoreTable table = paimonTable("MULTI_VECTOR_PK_T"); + for (int bucket = 0; bucket < 4; bucket++) { + assertThat( + table.store() + .newIndexFileHandler() + .scanSourceIndexes( + table.latestSnapshot().get(), + BinaryRow.EMPTY_ROW, + bucket)) + .extracting(payload -> payload.globalIndexMeta().indexFieldId()) + .containsExactlyInAnyOrder( + table.rowType().getField("image_vec").id(), + table.rowType().getField("text_vec").id()); + } + + for (String field : Arrays.asList("image_vec", "text_vec")) { + assertThat( + sql( + "CALL sys.vector_search(" + + "`table` => 'default.MULTI_VECTOR_PK_T', " + + "vector_column => '%s', " + + "query_vector => '0.0,0.0', " + + "top_k => 1, projection => 'id')", + field)) + .extracting(row -> row.getField(0).toString()) + .containsExactly( + field.equals("image_vec") ? "{\"id\":\"0\"}" : "{\"id\":\"31\"}"); + } + + // Both fields must run the distributed search operator, not the local fallback. + try (ClusterClient client = MINI_CLUSTER_EXTENSION.createRestClusterClient()) { + assertThat(client.listJobs().get(30, TimeUnit.SECONDS)) + .filteredOn( + job -> + job.getJobName().contains("Primary-Key Vector Search") + && job.getJobName().contains("MULTI_VECTOR_PK_T")) + .extracting(JobStatusMessage::getJobId) + .hasSize(2) + .doesNotHaveDuplicates(); + } + } + @Test public void testDistributedPrimaryKeyVectorSearchWithResidualFilter() { createPrimaryKeyVectorTableWithPayload( 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 80806858289c..5e2c6b038cb8 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 @@ -21,7 +21,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 org.apache.paimon.table.source.{BucketVectorSearchSplit, DataSplit, PrimaryKeyVectorResult} import scala.collection.JavaConverters._ @@ -105,6 +105,74 @@ class PrimaryKeyVectorSearchTest extends PaimonSparkTestBase { } } + test("distributed primary-key vector search isolates multiple indexed columns") { + withTable("T") { + val indexType = TestVectorGlobalIndexerFactory.IDENTIFIER + spark.sql(s""" + |CREATE TABLE T ( + | id INT, image_embedding ARRAY, text_embedding ARRAY + |) TBLPROPERTIES ( + | 'primary-key' = 'id', 'bucket' = '2', + | 'deletion-vectors.enabled' = 'true', + | 'vector-field' = 'image_embedding,text_embedding', + | 'field.image_embedding.vector-dim' = '2', + | 'field.text_embedding.vector-dim' = '2', + | 'pk-vector.index.columns' = 'image_embedding,text_embedding', + | 'fields.image_embedding.pk-vector.index.type' = '$indexType', + | 'fields.text_embedding.pk-vector.index.type' = '$indexType', + | 'fields.image_embedding.pk-vector.distance.metric' = 'l2', + | 'fields.text_embedding.pk-vector.distance.metric' = 'l2', + | 'global-index.search-mode' = 'fast', + | 'global-index.thread-num' = '1', + | 'test.vector.dimension' = '2', 'test.vector.metric' = 'l2' + |) + |""".stripMargin) + val values = (1 to 16) + .map(id => s"($id, array($id.0f, 0.0f), array(${17 - id}.0f, 0.0f))") + .mkString(",") + spark.sql(s"INSERT INTO T VALUES $values") + spark.sql("CALL sys.compact(table => 'T')") + + Seq("image_embedding" -> Set(1, 2), "text_embedding" -> Set(15, 16)).foreach { + case (column, expectedIds) => + val table = loadTable("T") + val builder = new SparkVectorSearchBuilderImpl(table) + builder.withVectorColumn(column).withVector(Array(0.0f, 0.0f)).withLimit(2) + val plan = builder.newVectorScan().scan() + val splits = plan.splits().asScala.map(_.asInstanceOf[BucketVectorSearchSplit]) + assert(splits.map(_.dataSplit().bucket()).toSet == Set(0, 1)) + assert(splits.forall(!_.payloadFiles().isEmpty)) + val fieldId = table.rowType().getField(column).id() + assert( + splits + .flatMap(_.payloadFiles().asScala) + .forall(_.globalIndexMeta().indexFieldId() == fieldId)) + + val jobGroup = s"primary-key-multiple-vectors-$column-${System.nanoTime()}" + spark.sparkContext.setJobGroup(jobGroup, jobGroup) + try { + val result = builder.newVectorRead().read(plan).asInstanceOf[PrimaryKeyVectorResult] + assert(result.positions().size() == 2) + } finally { + spark.sparkContext.clearJobGroup() + } + assert(spark.sparkContext.statusTracker.getJobIdsForGroup(jobGroup).nonEmpty) + + withSparkSQLConf("spark.paimon.vector-search.distribute.enabled" -> "true") { + val ids = spark + .sql(s""" + |SELECT id + |FROM vector_search('T', '$column', array(0.0f, 0.0f), 2) + |""".stripMargin) + .collect() + .map(_.getInt(0)) + .toSet + assert(ids == expectedIds) + } + } + } + } + test("postpone bucket builds primary-key vector index during compact") { withTable("T") { createVectorTable(