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/3] [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/3] [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( From e787eb7f964037758535aece846373865e5b9dfb Mon Sep 17 00:00:00 2001 From: tzphh <742429651@qq.com> Date: Thu, 22 Oct 2026 21:34:02 +0800 Subject: [PATCH 3/3] [core] Support bounded parallel builds for primary-key indexes --- docs/docs/primary-key-table/global-index.mdx | 13 + docs/generated/core_configuration.html | 6 + .../java/org/apache/paimon/CoreOptions.java | 20 ++ .../pk/BucketedPrimaryKeyIndexMaintainer.java | 123 +++----- .../BucketedVectorIndexMaintainer.java | 12 + .../operation/AbstractFileStoreWrite.java | 3 +- .../paimon/schema/SchemaValidation.java | 1 + ...BucketedPrimaryKeyIndexMaintainerTest.java | 283 ++++++++++++++++-- .../PrimaryKeySortedIndexMaintenanceTest.java | 23 +- .../BucketedVectorIndexMaintainerTest.java | 29 +- .../operation/PrimaryKeyIndexWriteTest.java | 17 ++ .../PrimaryKeyVectorIndexWriteTest.java | 93 ++++-- .../PrimaryKeyVectorIndexValidationTest.java | 16 + .../source/PrimaryKeyVectorSearchTest.java | 34 ++- 14 files changed, 514 insertions(+), 159 deletions(-) diff --git a/docs/docs/primary-key-table/global-index.mdx b/docs/docs/primary-key-table/global-index.mdx index c0eeb9f7b214..7cd7ac4652fd 100644 --- a/docs/docs/primary-key-table/global-index.mdx +++ b/docs/docs/primary-key-table/global-index.mdx @@ -299,6 +299,7 @@ not cover a query on `text_embedding`. | Option | Default | Description | |---|---|---| +| `pk-index.build.parallelism` | `1` | Maximum concurrent primary-key index build tasks per writer instance. Must be positive. | | `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`. | @@ -330,6 +331,18 @@ in the [`paimon-full-text` README](https://github.com/apache/paimon/blob/master/ ## Maintenance and Coverage +Set `pk-index.build.parallelism` to `2`, for example, to allow two primary-key index build tasks +to run concurrently within one writer instance. The default is `1`. Vector, Full Text, BTree, +Bitmap, Multivalue, and FM indexes share this limit across that writer's columns and buckets. +It is not the Flink or Spark job parallelism. + +Blocking commits start all configured indexes before waiting for their results. Non-blocking +commits start builds and collect finished results without waiting; outstanding builds are +collected by later commits. If an index fails during preparation, the coordinator rolls back +that preparation, cancels outstanding builds, and cleans up uncommitted outputs. Previously +committed snapshots remain valid. More concurrent builds can increase memory and I/O usage; +each index implementation may also use its own internal threads. + Paimon builds primary-key indexes from complete `COMPACT` data files above Level 0. An index group stores the source file names and row counts required to map index results back to physical rows. Data-file and index-file changes are committed in the same snapshot. diff --git a/docs/generated/core_configuration.html b/docs/generated/core_configuration.html index 069c4d15cdb3..fcdb1a3412b4 100644 --- a/docs/generated/core_configuration.html +++ b/docs/generated/core_configuration.html @@ -1392,6 +1392,12 @@ String Comma-separated character columns indexed by primary-key full-text indexes. The first release supports exactly one column. + +
pk-index.build.parallelism
+ 1 + Integer + The maximum number of concurrent primary-key index build tasks per writer instance, shared by its index columns and buckets. Must be greater than zero. Increasing this value may increase memory and I/O usage. +
pk-multivalue.index.columns
(none) 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 2f14a68b19a1..8f5c47660315 100644 --- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java +++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java @@ -3163,6 +3163,16 @@ public String toString() { "The batch size for lateral vector search. Each batch executes vector " + "topK search and table lookup for multiple query vectors."); + public static final ConfigOption PK_INDEX_BUILD_PARALLELISM = + key("pk-index.build.parallelism") + .intType() + .defaultValue(1) + .withDescription( + "The maximum number of concurrent primary-key index build tasks per " + + "writer instance, shared by its index columns and buckets. " + + "Must be greater than zero. Increasing this value may increase " + + "memory and I/O usage."); + public static final ConfigOption PK_VECTOR_INDEX_COLUMNS = key("pk-vector.index.columns") .stringType() @@ -4990,6 +5000,16 @@ public boolean primaryKeyFMIndexEnabled() { return options.getOptional(PK_FM_INDEX_COLUMNS).isPresent(); } + public int primaryKeyIndexBuildParallelism() { + int parallelism = options.get(PK_INDEX_BUILD_PARALLELISM); + checkArgument( + parallelism > 0, + "%s must be greater than zero, but is %s.", + PK_INDEX_BUILD_PARALLELISM.key(), + parallelism); + return parallelism; + } + public List primaryKeyVectorIndexColumns() { return primaryKeyIndexColumns(PK_VECTOR_INDEX_COLUMNS); } 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 44ff1c4be569..dee743618a9b 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 @@ -51,6 +51,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ExecutorService; +import java.util.function.Consumer; import static org.apache.paimon.utils.Preconditions.checkArgument; @@ -60,7 +61,6 @@ public final class BucketedPrimaryKeyIndexMaintainer { private final List vectorMaintainers; @Nullable private final BucketedFullTextIndexMaintainer fullTextMaintainer; private final List sortedMaintainers; - private int nextSortedMaintainerIndex; private BucketedPrimaryKeyIndexMaintainer( List vectorMaintainers, @@ -122,45 +122,25 @@ public synchronized void prepareCommit( CompactIncrement compactIncrement, boolean waitCompaction) throws Exception { - List vectorCommits = new ArrayList<>(); - BucketedFullTextIndexMaintainer.FullTextIndexCommit fullTextCommit = null; - List sortedCommits = new ArrayList<>(); + List commits = new ArrayList<>(); + List> aborts = new ArrayList<>(); try { - for (BucketedVectorIndexMaintainer maintainer : vectorMaintainers) { - vectorCommits.add( - maintainer.prepareCommit( - appendIncrement, compactIncrement, waitCompaction)); + // Start all indexes before waiting. Retain both rounds for reverse-order rollback, + // and merge changes only after every index has prepared successfully. + int indexCount = + vectorMaintainers.size() + + sortedMaintainers.size() + + (fullTextMaintainer == null ? 0 : 1); + if (waitCompaction && indexCount > 1) { + prepareIndexes(appendIncrement, compactIncrement, false, commits, aborts); } - if (fullTextMaintainer != null) { - fullTextCommit = - fullTextMaintainer.prepareCommit( - appendIncrement, compactIncrement, waitCompaction); - } - - if (waitCompaction) { - prepareSortedBlocking(appendIncrement, compactIncrement, sortedCommits); - } else { - prepareSortedNonBlocking(appendIncrement, compactIncrement, sortedCommits); - } - - for (BucketedVectorIndexMaintainer.VectorIndexCommit commit : vectorCommits) { - mergeVectorCommit(appendIncrement, compactIncrement, commit); - } - if (fullTextCommit != null) { - mergeFullTextCommit(appendIncrement, compactIncrement, fullTextCommit); - } - for (BucketedSortedIndexMaintainer.SortedIndexCommit sortedCommit : sortedCommits) { - mergeSortedCommit(appendIncrement, compactIncrement, sortedCommit); + prepareIndexes(appendIncrement, compactIncrement, waitCompaction, commits, aborts); + for (Runnable commit : commits) { + commit.run(); } } catch (Throwable failure) { - for (int i = sortedCommits.size() - 1; i >= 0; i--) { - sortedCommits.get(i).abort(failure); - } - if (fullTextCommit != null) { - fullTextCommit.abort(failure); - } - for (int i = vectorCommits.size() - 1; i >= 0; i--) { - vectorCommits.get(i).abort(failure); + for (int i = aborts.size() - 1; i >= 0; i--) { + aborts.get(i).accept(failure); } if (failure instanceof Exception) { throw (Exception) failure; @@ -172,58 +152,32 @@ public synchronized void prepareCommit( } } - private void prepareSortedBlocking( + private void prepareIndexes( DataIncrement appendIncrement, CompactIncrement compactIncrement, - List commits) + boolean waitCompaction, + List commits, + List> aborts) throws Exception { - BucketedSortedIndexMaintainer active = activeSortedMaintainer(); - for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { - commits.add( - maintainer.prepareCommit( - appendIncrement, compactIncrement, true, maintainer == active)); - } - for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { - if (maintainer != active) { - commits.add( - maintainer.prepareCommit(appendIncrement, compactIncrement, true, true)); - } - } - } - - private void prepareSortedNonBlocking( - DataIncrement appendIncrement, - CompactIncrement compactIncrement, - List commits) - throws Exception { - for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { - commits.add(maintainer.prepareCommit(appendIncrement, compactIncrement, false, false)); - } - if (activeSortedMaintainer() != null) { - return; + for (BucketedVectorIndexMaintainer maintainer : vectorMaintainers) { + BucketedVectorIndexMaintainer.VectorIndexCommit commit = + maintainer.prepareCommit(appendIncrement, compactIncrement, waitCompaction); + aborts.add(commit::abort); + commits.add(() -> mergeVectorCommit(appendIncrement, compactIncrement, commit)); } - - int count = sortedMaintainers.size(); - for (int offset = 0; offset < count; offset++) { - int index = (nextSortedMaintainerIndex + offset) % count; - BucketedSortedIndexMaintainer maintainer = sortedMaintainers.get(index); - if (maintainer.hasPendingMaintenance()) { - commits.add( - maintainer.prepareCommit(appendIncrement, compactIncrement, false, true)); - nextSortedMaintainerIndex = (index + 1) % count; - break; - } + if (fullTextMaintainer != null) { + BucketedFullTextIndexMaintainer.FullTextIndexCommit commit = + fullTextMaintainer.prepareCommit( + appendIncrement, compactIncrement, waitCompaction); + aborts.add(commit::abort); + commits.add(() -> mergeFullTextCommit(appendIncrement, compactIncrement, commit)); } - } - - @Nullable - private BucketedSortedIndexMaintainer activeSortedMaintainer() { for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { - if (maintainer.buildNotCompleted()) { - return maintainer; - } + BucketedSortedIndexMaintainer.SortedIndexCommit commit = + maintainer.prepareCommit(appendIncrement, compactIncrement, waitCompaction); + aborts.add(commit::abort); + commits.add(() -> mergeSortedCommit(appendIncrement, compactIncrement, commit)); } - return null; } private static void mergeSortedCommit( @@ -310,7 +264,12 @@ public boolean buildNotCompleted() { if (fullTextMaintainer != null && fullTextMaintainer.buildNotCompleted()) { return true; } - return activeSortedMaintainer() != null; + for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { + if (maintainer.buildNotCompleted()) { + return true; + } + } + return false; } public void withExecutor(ExecutorService executor) { diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java index b923e052b978..85ad6dea127c 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java @@ -183,6 +183,18 @@ public synchronized VectorIndexCommit prepareCommit( activeSourceFiles.clear(); activeSourceFiles.putAll(nextSources); + // A source transition retires its old coverage before starting replacement builds. + // Non-blocking preparation can then publish the data transition without stale indexes. + removeSegments( + PkVectorBucketIndexState.fromActiveDataFiles( + vectorFieldId, + algorithm, + new ArrayList<>(activeSourceFiles.values()), + annSegments) + .staleSegments(), + created, + removed); + while (true) { Optional completed = finishPendingBuild(waitCompaction); if (completed.isPresent()) { diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java index 0c0738c6173e..f08b54f8188c 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java @@ -716,7 +716,8 @@ private ExecutorService compactExecutor() { private ExecutorService primaryKeyIndexExecutor() { if (lazyPrimaryKeyIndexExecutor == null) { lazyPrimaryKeyIndexExecutor = - Executors.newSingleThreadExecutor( + Executors.newFixedThreadPool( + options.primaryKeyIndexBuildParallelism(), new ExecutorThreadFactory( Thread.currentThread().getName() + "-primary-key-index")); } 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 e30c8e39f330..7879b15c783a 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 @@ -402,6 +402,7 @@ public static void validateTableSchema(TableSchema schema, Set dynamicOp fieldNamesSpecifiedAsVector.isEmpty(), "Some of the columns specified as vector-field are unknown."); + options.primaryKeyIndexBuildParallelism(); validatePrimaryKeyIndexColumns(options); validatePrimaryKeySortedIndexes(schema, options); validatePrimaryKeyVectorIndex(schema, options); 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 b029b42c492f..f49201320893 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 @@ -27,6 +27,7 @@ import org.apache.paimon.index.IndexFileMeta; import org.apache.paimon.index.IndexPathFactory; import org.apache.paimon.index.pkfulltext.BucketedFullTextIndexMaintainer; +import org.apache.paimon.index.pkfulltext.PkFullTextIndexBuilder; import org.apache.paimon.index.pkfulltext.PkFullTextIndexFile; import org.apache.paimon.index.pksorted.BucketedSortedIndexMaintainer; import org.apache.paimon.index.pksorted.PkSortedIndexFile; @@ -46,6 +47,9 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.InOrder; import java.util.ArrayList; @@ -59,6 +63,7 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -174,6 +179,158 @@ void testVectorFailureRollsBackEarlierIndexAndAllowsRetry() throws Exception { maintainer.close(); } + @ParameterizedTest + @CsvSource({"false,1", "false,2", "true,1", "true,2"}) + void testIndexBuildsOverlapWithinExecutorLimit(boolean waitCompaction, int parallelism) + throws Exception { + DataFileMeta source = dataFile("data-1", 3); + CountDownLatch started = new CountDownLatch(parallelism); + CountDownLatch release = new CountDownLatch(1); + AtomicInteger active = new AtomicInteger(); + AtomicInteger peak = new AtomicInteger(); + List payloads = new ArrayList<>(); + IndexFileMeta vectorPayload = payload("vector", source, 3, 1, "vector"); + payloads.add(vectorPayload); + PkVectorAnnSegmentFile vectorFile = mock(PkVectorAnnSegmentFile.class); + when(vectorFile.build(anyList(), any(), any(), eq("l2"), eq("vector"))) + .thenAnswer(ignored -> blockedBuild(vectorPayload, active, peak, started, release)); + IndexFileMeta textPayload = payload("text", source, 3, 2, "full-text"); + payloads.add(textPayload); + PkFullTextIndexBuilder textBuilder = mock(PkFullTextIndexBuilder.class); + when(textBuilder.build(any(DataFileMeta.class))) + .thenAnswer(ignored -> blockedBuild(textPayload, active, peak, started, release)); + BucketedFullTextIndexMaintainer fullText = fullTextMaintainer(textBuilder); + List sorted = new ArrayList<>(); + for (String type : Arrays.asList("btree", "bitmap", "multivalue", "fm")) { + int fieldId = sorted.size() + 3; + IndexFileMeta result = payload(type, source, 3, fieldId, type); + payloads.add(result); + sorted.add( + sortedMaintainer( + fieldId, + type, + source, + files -> blockedBuild(result, active, peak, started, release))); + } + BucketedPrimaryKeyIndexMaintainer maintainer = + BucketedPrimaryKeyIndexMaintainer.of( + vectorMaintainer(1, vectorFile), fullText, sorted); + ExecutorService executor = Executors.newFixedThreadPool(parallelism); + maintainer.withExecutor(executor); + CompactIncrement compact = compactAfter(source); + ExecutorService committer = Executors.newSingleThreadExecutor(); + try { + Future commit = + committer.submit( + () -> { + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), compact, waitCompaction); + return null; + }); + assertThat(started.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(active).hasValue(parallelism); + if (waitCompaction) { + assertThat(commit.isDone()).isFalse(); + } else { + commit.get(5, TimeUnit.SECONDS); + assertThat(compact.newIndexFiles()).isEmpty(); + assertThat(maintainer.buildNotCompleted()).isTrue(); + } + release.countDown(); + commit.get(5, TimeUnit.SECONDS); + if (!waitCompaction) { + maintainer.prepareCommit(DataIncrement.emptyIncrement(), compact, true); + } + assertThat(peak).hasValue(parallelism); + assertThat(compact.newIndexFiles()).containsExactlyInAnyOrderElementsOf(payloads); + assertThat(maintainer.buildNotCompleted()).isFalse(); + } finally { + release.countDown(); + committer.shutdownNow(); + assertThat(committer.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); + maintainer.close(); + executor.shutdownNow(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testFirstVectorFailureCancelsAnotherRunningBuildAndAllowsRetry(boolean useFullText) + throws Exception { + DataFileMeta source = dataFile("data-1", 3); + IndexFileMeta firstPayload = payload("first", source, 3, 1, "vector"); + IndexFileMeta secondPayload = + payload("second", source, 3, 2, useFullText ? "full-text" : "vector"); + CountDownLatch secondStarted = new CountDownLatch(1); + CountDownLatch secondStopped = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + PkVectorAnnSegmentFile firstFile = mock(PkVectorAnnSegmentFile.class); + PkVectorAnnSegmentFile secondFile = mock(PkVectorAnnSegmentFile.class); + when(firstFile.build(anyList(), any(), any(), eq("l2"), eq("vector"))) + .thenAnswer( + ignored -> { + assertThat(secondStarted.await(5, TimeUnit.SECONDS)).isTrue(); + throw new IllegalStateException("first vector failed"); + }) + .thenReturn(firstPayload); + when(secondFile.build(anyList(), any(), any(), eq("l2"), eq("vector"))) + .thenAnswer( + ignored -> { + secondStarted.countDown(); + try { + assertThat(release.await(10, TimeUnit.SECONDS)).isTrue(); + return secondPayload; + } finally { + secondStopped.countDown(); + } + }) + .thenReturn(secondPayload); + BucketedVectorIndexMaintainer first = vectorMaintainer(1, firstFile); + BucketedVectorIndexMaintainer second = vectorMaintainer(2, secondFile); + PkFullTextIndexBuilder textBuilder = mock(PkFullTextIndexBuilder.class); + when(textBuilder.build(source)) + .thenAnswer( + ignored -> { + secondStarted.countDown(); + try { + assertThat(release.await(10, TimeUnit.SECONDS)).isTrue(); + return secondPayload; + } finally { + secondStopped.countDown(); + } + }) + .thenReturn(secondPayload); + BucketedFullTextIndexMaintainer fullText = fullTextMaintainer(textBuilder); + BucketedPrimaryKeyIndexMaintainer maintainer = + useFullText + ? BucketedPrimaryKeyIndexMaintainer.of( + first, fullText, Collections.emptyList()) + : BucketedPrimaryKeyIndexMaintainer.ofVectors( + Arrays.asList(first, second), null, Collections.emptyList()); + CompactIncrement failed = compactAfter(source); + try { + assertThatThrownBy( + () -> + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), failed, true)) + .hasMessage("first vector failed"); + assertThat(secondStopped.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(failed.newIndexFiles()).isEmpty(); + assertThat(failed.deletedIndexFiles()).isEmpty(); + assertThat(first.segments()).isEmpty(); + assertThat(second.segments()).isEmpty(); + assertThat(fullText.payloads()).isEmpty(); + assertThat(maintainer.buildNotCompleted()).isFalse(); + + CompactIncrement retry = compactAfter(source); + maintainer.prepareCommit(DataIncrement.emptyIncrement(), retry, true); + assertThat(retry.newIndexFiles()).containsExactly(firstPayload, secondPayload); + } finally { + release.countDown(); + maintainer.close(); + } + } + @Test void testFullTextFailureAbortsAllPreparedVectorsInReverseOrder() throws Exception { BucketedVectorIndexMaintainer first = mock(BucketedVectorIndexMaintainer.class); @@ -182,9 +339,18 @@ void testFullTextFailureAbortsAllPreparedVectorsInReverseOrder() throws Exceptio mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class); BucketedVectorIndexMaintainer.VectorIndexCommit secondCommit = mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class); + BucketedVectorIndexMaintainer.VectorIndexCommit firstStarted = + mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class); + BucketedVectorIndexMaintainer.VectorIndexCommit secondStarted = + mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class); + when(first.prepareCommit(any(), any(), eq(false))).thenReturn(firstStarted); + when(second.prepareCommit(any(), any(), eq(false))).thenReturn(secondStarted); when(first.prepareCommit(any(), any(), eq(true))).thenReturn(firstCommit); when(second.prepareCommit(any(), any(), eq(true))).thenReturn(secondCommit); BucketedFullTextIndexMaintainer fullText = mock(BucketedFullTextIndexMaintainer.class); + BucketedFullTextIndexMaintainer.FullTextIndexCommit textStarted = + mock(BucketedFullTextIndexMaintainer.FullTextIndexCommit.class); + when(fullText.prepareCommit(any(), any(), eq(false))).thenReturn(textStarted); IllegalStateException failure = new IllegalStateException("full-text failed"); when(fullText.prepareCommit(any(), any(), eq(true))).thenThrow(failure); BucketedPrimaryKeyIndexMaintainer maintainer = @@ -194,12 +360,27 @@ void testFullTextFailureAbortsAllPreparedVectorsInReverseOrder() throws Exceptio CompactIncrement compact = CompactIncrement.emptyIncrement(); assertThatThrownBy(() -> maintainer.prepareCommit(append, compact, true)).isSameAs(failure); - InOrder order = inOrder(first, second, fullText, firstCommit, secondCommit); + InOrder order = + inOrder( + first, + second, + fullText, + firstCommit, + secondCommit, + firstStarted, + secondStarted, + textStarted); + order.verify(first).prepareCommit(append, compact, false); + order.verify(second).prepareCommit(append, compact, false); + order.verify(fullText).prepareCommit(append, compact, false); 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); + order.verify(textStarted).abort(failure); + order.verify(secondStarted).abort(failure); + order.verify(firstStarted).abort(failure); verify(firstCommit, never()).appendIncrement(); verify(secondCommit, never()).appendIncrement(); assertThat(append.newIndexFiles()).isEmpty(); @@ -279,6 +460,8 @@ void testLaterDefinitionFailureAbortsPreparedFullTextCommit() throws Exception { BucketedFullTextIndexMaintainer fullText = mock(BucketedFullTextIndexMaintainer.class); BucketedFullTextIndexMaintainer.FullTextIndexCommit fullTextCommit = mock(BucketedFullTextIndexMaintainer.FullTextIndexCommit.class); + when(fullText.prepareCommit(any(), any(), eq(false))) + .thenReturn(mock(BucketedFullTextIndexMaintainer.FullTextIndexCommit.class)); when(fullText.prepareCommit(any(), any(), eq(true))).thenReturn(fullTextCommit); when(fullTextCommit.appendIncrement()).thenReturn(Optional.empty()); when(fullTextCommit.compactIncrement()).thenReturn(Optional.empty()); @@ -313,6 +496,8 @@ void testScalarFailureAbortsOtherDefinitionsAndVector() throws Exception { mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class); BucketedVectorIndexMaintainer.VectorIndexIncrement vectorIncrement = mock(BucketedVectorIndexMaintainer.VectorIndexIncrement.class); + when(vector.prepareCommit(any(), any(), eq(false))) + .thenReturn(mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class)); when(vector.prepareCommit(any(), any(), eq(true))).thenReturn(vectorCommit); when(vectorCommit.appendIncrement()).thenReturn(Optional.empty()); when(vectorCommit.compactIncrement()).thenReturn(Optional.of(vectorIncrement)); @@ -320,6 +505,7 @@ void testScalarFailureAbortsOtherDefinitionsAndVector() throws Exception { when(vectorIncrement.deletedIndexFiles()).thenReturn(Collections.emptyList()); List buildOrder = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch bitmapStarted = new CountDownLatch(1); BucketedSortedIndexMaintainer btree = sortedMaintainer( 7, @@ -327,6 +513,7 @@ void testScalarFailureAbortsOtherDefinitionsAndVector() throws Exception { source, dataFile -> { buildOrder.add("btree"); + assertThat(bitmapStarted.await(5, TimeUnit.SECONDS)).isTrue(); throw new IllegalStateException("expected BTree failure"); }); IndexFileMeta bitmapPayload = payload("bitmap", source, 3, 8, "bitmap"); @@ -337,6 +524,7 @@ void testScalarFailureAbortsOtherDefinitionsAndVector() throws Exception { source, dataFile -> { buildOrder.add("bitmap"); + bitmapStarted.countDown(); return bitmapPayload; }); BucketedPrimaryKeyIndexMaintainer maintainer = @@ -350,7 +538,9 @@ void testScalarFailureAbortsOtherDefinitionsAndVector() throws Exception { .isInstanceOf(IllegalStateException.class) .hasMessage("expected BTree failure"); - assertThat(buildOrder).containsExactly("btree", "btree", "btree"); + assertThat(buildOrder).containsExactlyInAnyOrder("btree", "btree", "btree", "bitmap"); + assertThat(btree.state().groups()).isEmpty(); + assertThat(bitmap.state().groups()).isEmpty(); assertThat(compactIncrement.newIndexFiles()).isEmpty(); assertThat(compactIncrement.deletedIndexFiles()).isEmpty(); assertThat(maintainer.buildNotCompleted()).isFalse(); @@ -358,7 +548,7 @@ void testScalarFailureAbortsOtherDefinitionsAndVector() throws Exception { } @Test - void testScalarDefinitionsNeverBuildConcurrently() throws Exception { + void testScalarDefinitionsBuildConcurrently() throws Exception { DataFileMeta source = dataFile("data-1", 3); AtomicInteger activeBuilds = new AtomicInteger(); AtomicInteger peakBuilds = new AtomicInteger(); @@ -383,6 +573,7 @@ void testScalarDefinitionsNeverBuildConcurrently() throws Exception { "bitmap", source, dataFile -> { + assertThat(firstStarted.await(5, TimeUnit.SECONDS)).isTrue(); enterBuild(activeBuilds, peakBuilds); secondStarted.countDown(); activeBuilds.decrementAndGet(); @@ -393,21 +584,14 @@ void testScalarDefinitionsNeverBuildConcurrently() throws Exception { maintainer.prepareCommit(DataIncrement.emptyIncrement(), compactAfter(source), false); assertThat(firstStarted.await(5, TimeUnit.SECONDS)).isTrue(); - assertThat(secondStarted.getCount()).isEqualTo(1); + assertThat(secondStarted.await(5, TimeUnit.SECONDS)).isTrue(); assertThat(maintainer.buildNotCompleted()).isTrue(); releaseFirst.countDown(); - long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); - while (secondStarted.getCount() > 0 && System.nanoTime() < deadline) { - maintainer.prepareCommit( - DataIncrement.emptyIncrement(), CompactIncrement.emptyIncrement(), false); - secondStarted.await(10, TimeUnit.MILLISECONDS); - } - assertThat(secondStarted.getCount()).isZero(); maintainer.prepareCommit( DataIncrement.emptyIncrement(), CompactIncrement.emptyIncrement(), true); - assertThat(peakBuilds).hasValue(1); + assertThat(peakBuilds).hasValue(2); assertThat(maintainer.buildNotCompleted()).isFalse(); } @@ -449,6 +633,7 @@ void testBlockingPrepareRollsBackAllDefinitionsAfterInterruption() throws Except DataFileMeta source = dataFile("data-1", 3); AtomicInteger firstBuilds = new AtomicInteger(); AtomicInteger secondBuilds = new AtomicInteger(); + CountDownLatch firstStarted = new CountDownLatch(1); CountDownLatch secondStarted = new CountDownLatch(1); CountDownLatch blockSecond = new CountDownLatch(1); BucketedSortedIndexMaintainer first = @@ -458,6 +643,7 @@ void testBlockingPrepareRollsBackAllDefinitionsAfterInterruption() throws Except source, dataFile -> { int attempt = firstBuilds.incrementAndGet(); + firstStarted.countDown(); return payload("btree-" + attempt, source, 3, 7, "btree"); }); BucketedSortedIndexMaintainer second = @@ -489,25 +675,33 @@ void testBlockingPrepareRollsBackAllDefinitionsAfterInterruption() throws Except }); commitThread.start(); - assertThat(secondStarted.await(5, TimeUnit.SECONDS)).isTrue(); - commitThread.interrupt(); - commitThread.join(TimeUnit.SECONDS.toMillis(5)); - - assertThat(commitThread.isAlive()).isFalse(); - assertThat(failure.get()).isInstanceOf(InterruptedException.class); - assertThat(failedIncrement.newIndexFiles()).isEmpty(); - assertThat(failedIncrement.deletedIndexFiles()).isEmpty(); - assertThat(first.state().groups()).isEmpty(); - assertThat(second.state().groups()).isEmpty(); - - CompactIncrement retryIncrement = compactAfter(source); - maintainer.prepareCommit(DataIncrement.emptyIncrement(), retryIncrement, true); - - assertThat(firstBuilds).hasValue(2); - assertThat(secondBuilds).hasValue(2); - assertThat(retryIncrement.newIndexFiles()) - .extracting(IndexFileMeta::fileName) - .containsExactly("btree-2", "bitmap-2"); + try { + assertThat(firstStarted.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(secondStarted.await(5, TimeUnit.SECONDS)).isTrue(); + commitThread.interrupt(); + commitThread.join(TimeUnit.SECONDS.toMillis(5)); + + assertThat(commitThread.isAlive()).isFalse(); + assertThat(failure.get()).isInstanceOf(InterruptedException.class); + assertThat(failedIncrement.newIndexFiles()).isEmpty(); + assertThat(failedIncrement.deletedIndexFiles()).isEmpty(); + assertThat(first.state().groups()).isEmpty(); + assertThat(second.state().groups()).isEmpty(); + + CompactIncrement retryIncrement = compactAfter(source); + maintainer.prepareCommit(DataIncrement.emptyIncrement(), retryIncrement, true); + + assertThat(firstBuilds).hasValue(2); + assertThat(secondBuilds).hasValue(2); + assertThat(retryIncrement.newIndexFiles()) + .extracting(IndexFileMeta::fileName) + .containsExactly("btree-2", "bitmap-2"); + } finally { + blockSecond.countDown(); + commitThread.interrupt(); + commitThread.join(TimeUnit.SECONDS.toMillis(5)); + maintainer.close(); + } } private BucketedVectorIndexMaintainer vectorMaintainer( @@ -547,6 +741,33 @@ private BucketedSortedIndexMaintainer sortedMaintainer( buildExecutor); } + private BucketedFullTextIndexMaintainer fullTextMaintainer(PkFullTextIndexBuilder builder) { + return new BucketedFullTextIndexMaintainer( + 2, + mock(PkFullTextIndexFile.class), + builder, + Collections.emptyList(), + Collections.emptyList(), + buildExecutor); + } + + private static IndexFileMeta blockedBuild( + IndexFileMeta payload, + AtomicInteger active, + AtomicInteger peak, + CountDownLatch started, + CountDownLatch release) + throws InterruptedException { + enterBuild(active, peak); + started.countDown(); + try { + assertThat(release.await(10, TimeUnit.SECONDS)).isTrue(); + return payload; + } finally { + active.decrementAndGet(); + } + } + private static void enterBuild(AtomicInteger activeBuilds, AtomicInteger peakBuilds) { int active = activeBuilds.incrementAndGet(); while (true) { diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PrimaryKeySortedIndexMaintenanceTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PrimaryKeySortedIndexMaintenanceTest.java index 796485d7f622..087b7eb3d7f7 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PrimaryKeySortedIndexMaintenanceTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PrimaryKeySortedIndexMaintenanceTest.java @@ -58,8 +58,9 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import java.io.IOException; import java.util.ArrayList; @@ -90,9 +91,10 @@ void after() throws Exception { ioManager.close(); } - @Test - void testBuildRestoreAndReplaceSourceGroups() throws Exception { - TestFileStore store = createStore(); + @ParameterizedTest + @ValueSource(ints = {1, 2}) + void testBuildRestoreAndReplaceSourceGroups(int parallelism) throws Exception { + TestFileStore store = createStore(parallelism); TestKeyValueGenerator generator = new TestKeyValueGenerator(); List firstBatch = Arrays.asList( @@ -142,9 +144,10 @@ void testBuildRestoreAndReplaceSourceGroups() throws Exception { restored.close(); } - @Test - void testRestoreMixedVectorAndScalarPayloads() throws Exception { - TestFileStore store = createMixedStore(); + @ParameterizedTest + @ValueSource(ints = {1, 2}) + void testRestoreMixedVectorAndScalarPayloads(int parallelism) throws Exception { + TestFileStore store = createMixedStore(parallelism); TestKeyValueGenerator generator = new TestKeyValueGenerator(); KeyValue first = withVector( @@ -317,8 +320,9 @@ private static KeyValue withVector(KeyValue record, float[] vector) { return record.replaceValue(extended); } - private TestFileStore createStore() throws Exception { + private TestFileStore createStore(int parallelism) throws Exception { Map options = new HashMap<>(); + options.put(CoreOptions.PK_INDEX_BUILD_PARALLELISM.key(), String.valueOf(parallelism)); options.put(CoreOptions.BUCKET.key(), "1"); options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); options.put(CoreOptions.COMPACTION_FORCE_REWRITE_ALL_FILES.key(), "true"); @@ -352,8 +356,9 @@ private TestFileStore createStore() throws Exception { .build(); } - private TestFileStore createMixedStore() throws Exception { + private TestFileStore createMixedStore(int parallelism) throws Exception { Map options = new HashMap<>(); + options.put(CoreOptions.PK_INDEX_BUILD_PARALLELISM.key(), String.valueOf(parallelism)); options.put(CoreOptions.BUCKET.key(), "1"); options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); options.put(CoreOptions.COMPACTION_FORCE_REWRITE_ALL_FILES.key(), "true"); diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java index 3553f858ed50..e6b1986a070e 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java @@ -38,6 +38,8 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import java.io.IOException; import java.io.UncheckedIOException; @@ -138,13 +140,19 @@ void testPreparedCommitCanBeAbortedByCoordinator() throws Exception { assertThat(fileCount()).isZero(); } - @Test - void testNonBlockingPrepareCommitPublishesCompletedAnnLater() throws Exception { + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testNonBlockingPrepareCommitPublishesCompletedAnnLater(boolean replaceExisting) + throws Exception { LocalFileIO fileIO = LocalFileIO.create(); PkVectorAnnSegmentFile annFile = new PkVectorAnnSegmentFile(fileIO, pathFactory()); DataField vectorField = new DataField(7, "embedding", DataTypes.VECTOR(2, DataTypes.FLOAT())); DataFileMeta data = dataFile("data"); + DataFileMeta oldData = dataFile("old-data"); + IndexFileMeta oldAnn = payload("old-ann", oldData, 7, "test-vector-ann"); + List previousSources = + replaceExisting ? Collections.singletonList(oldData) : Collections.emptyList(); PkVectorDataFileReader.Factory readerFactory = mock(PkVectorDataFileReader.Factory.class); CountDownLatch buildStarted = new CountDownLatch(1); CountDownLatch allowBuild = new CountDownLatch(1); @@ -168,20 +176,29 @@ void testNonBlockingPrepareCommitPublishesCompletedAnnLater() throws Exception { "l2", "test-vector-ann", readerFactory, - Collections.emptyList(), - Collections.emptyList(), + previousSources, + replaceExisting + ? Collections.singletonList(oldAnn) + : Collections.emptyList(), executor); BucketedVectorIndexMaintainer.VectorIndexCommit firstCommit = maintainer.prepareCommit( DataIncrement.emptyIncrement(), new CompactIncrement( - Collections.emptyList(), + previousSources, Collections.singletonList(data), Collections.emptyList()), false); - assertThat(firstCommit.compactIncrement()).isEmpty(); + if (replaceExisting) { + assertThat(firstCommit.compactIncrement().get().deletedIndexFiles()) + .containsExactly(oldAnn); + assertThat(firstCommit.compactIncrement().get().newIndexFiles()).isEmpty(); + assertThat(maintainer.segments()).isEmpty(); + } else { + assertThat(firstCommit.compactIncrement()).isEmpty(); + } assertThat(buildStarted.await(30, TimeUnit.SECONDS)).isTrue(); assertThat(buildThread.get()).isNotSameAs(callerThread); allowBuild.countDown(); 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 994d59ef817f..3fd63c950b4e 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 @@ -123,12 +123,15 @@ void testCreatesCoordinatorForMultipleScalarDefinitionsAndVector() throws Except options.put(CoreOptions.PK_BITMAP_INDEX_COLUMNS.key(), "comment,state"); options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding"); options.put("fields.embedding.pk-vector.index.type", "hnsw"); + options.put(CoreOptions.PK_FULL_TEXT_INDEX_COLUMNS.key(), "full_text"); + options.put(CoreOptions.PK_INDEX_BUILD_PARALLELISM.key(), "2"); List fields = new ArrayList<>(TestKeyValueGenerator.DEFAULT_ROW_TYPE.getFields()); fields.add(new DataField(7, "score", DataTypes.BIGINT())); fields.add(new DataField(8, "state", DataTypes.INT())); fields.add(new DataField(9, "embedding", DataTypes.VECTOR(3, DataTypes.FLOAT()))); + fields.add(new DataField(10, "full_text", DataTypes.STRING())); TestFileStore store = createStore(options, new RowType(fields)); KeyValueFileStoreWrite write = (KeyValueFileStoreWrite) store.newWrite(); write.withIOManager(ioManager); @@ -145,6 +148,20 @@ void testCreatesCoordinatorForMultipleScalarDefinitionsAndVector() throws Except assertThat((List) readField(container.primaryKeyIndexMaintainer, "sortedMaintainers")) .hasSize(4); + Object executor = readField(vectorMaintainer, "executor"); + Object fullText = readField(container.primaryKeyIndexMaintainer, "fullTextMaintainer"); + assertThat(readField(fullText, "executor")).isSameAs(executor); + for (Object sorted : + (List) readField(container.primaryKeyIndexMaintainer, "sortedMaintainers")) { + assertThat(readField(sorted, "executor")).isSameAs(executor); + } + AbstractFileStoreWrite.WriterContainer otherBucket = + write.createWriterContainer(generator.getPartition(record), 2); + Object otherVector = + ((List) readField(otherBucket.primaryKeyIndexMaintainer, "vectorMaintainers")) + .get(0); + assertThat(readField(otherVector, "executor")).isSameAs(executor); + BitmapDeletionVector deletionVector = new BitmapDeletionVector(); deletionVector.delete(1); container.deletionVectorsMaintainer.deletionVectors().put("data-file", deletionVector); 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 34e6cbad340e..ca58f124cb0c 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 @@ -23,6 +23,8 @@ import org.apache.paimon.deletionvectors.BucketedDvMaintainer; import org.apache.paimon.index.IndexFileMeta; import org.apache.paimon.index.pk.BucketedPrimaryKeyIndexMaintainer; +import org.apache.paimon.index.pkfulltext.BucketedFullTextIndexMaintainer; +import org.apache.paimon.index.pksorted.BucketedSortedIndexMaintainer; import org.apache.paimon.index.pkvector.BucketedVectorIndexMaintainer; import org.apache.paimon.io.CompactIncrement; import org.apache.paimon.io.DataFileMeta; @@ -34,6 +36,7 @@ import org.apache.paimon.utils.SnapshotManager; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import javax.annotation.Nullable; @@ -42,12 +45,16 @@ import java.util.HashMap; import java.util.List; import java.util.Optional; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.function.Function; import static org.apache.paimon.data.BinaryRow.EMPTY_ROW; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.same; import static org.mockito.Mockito.mock; @@ -58,7 +65,7 @@ class PrimaryKeyVectorIndexWriteTest { @Test - void testRestoreRebindsVectorIndexExecutor() throws Exception { + void testRestoreRebindsSharedIndexExecutor() throws Exception { DataIncrement dataIncrement = DataIncrement.emptyIncrement(); CompactIncrement compactIncrement = CompactIncrement.emptyIncrement(); RecordWriter writer = mock(RecordWriter.class); @@ -70,21 +77,66 @@ void testRestoreRebindsVectorIndexExecutor() throws Exception { when(vectorCommit.appendIncrement()).thenReturn(Optional.empty()); when(vectorCommit.compactIncrement()).thenReturn(Optional.empty()); BucketedVectorIndexMaintainer maintainer = mock(BucketedVectorIndexMaintainer.class); + when(maintainer.prepareCommit(same(dataIncrement), same(compactIncrement), eq(false))) + .thenReturn(vectorCommit); when(maintainer.prepareCommit(same(dataIncrement), same(compactIncrement), eq(true))) .thenReturn(vectorCommit); BucketedVectorIndexMaintainer secondMaintainer = mock(BucketedVectorIndexMaintainer.class); + when(secondMaintainer.prepareCommit(same(dataIncrement), same(compactIncrement), eq(false))) + .thenReturn(vectorCommit); when(secondMaintainer.prepareCommit(same(dataIncrement), same(compactIncrement), eq(true))) .thenReturn(vectorCommit); TestingFileStoreWrite write = new TestingFileStoreWrite(); - write.install(writer, maintainer, secondMaintainer); + BucketedFullTextIndexMaintainer fullText = mock(BucketedFullTextIndexMaintainer.class); + when(fullText.prepareCommit(any(), any(), anyBoolean())) + .thenReturn(mock(BucketedFullTextIndexMaintainer.FullTextIndexCommit.class)); + BucketedSortedIndexMaintainer sorted = mock(BucketedSortedIndexMaintainer.class); + when(sorted.prepareCommit(any(), any(), anyBoolean())) + .thenReturn(mock(BucketedSortedIndexMaintainer.SortedIndexCommit.class)); + write.install( + writer, + BucketedPrimaryKeyIndexMaintainer.ofVectors( + Arrays.asList(maintainer, secondMaintainer), + fullText, + Collections.singletonList(sorted))); List> states = write.checkpoint(); - TestingFileStoreWrite restored = new TestingFileStoreWrite(mock(RecordWriter.class)); + TestingFileStoreWrite restored = new TestingFileStoreWrite(mock(RecordWriter.class), 2); restored.restore(states); - verify(maintainer).withExecutor(any(ExecutorService.class)); - verify(secondMaintainer).withExecutor(any(ExecutorService.class)); + ArgumentCaptor executor = ArgumentCaptor.forClass(ExecutorService.class); + verify(maintainer).withExecutor(executor.capture()); + verify(secondMaintainer).withExecutor(same(executor.getValue())); + verify(fullText).withExecutor(same(executor.getValue())); + verify(sorted).withExecutor(same(executor.getValue())); + CountDownLatch started = new CountDownLatch(2); + CountDownLatch release = new CountDownLatch(1); + try { + Future first = + executor.getValue() + .submit( + () -> { + started.countDown(); + return release.await(10, TimeUnit.SECONDS); + }); + Future second = + executor.getValue() + .submit( + () -> { + started.countDown(); + return release.await(10, TimeUnit.SECONDS); + }); + assertThat(started.await(5, TimeUnit.SECONDS)).isTrue(); + release.countDown(); + assertThat(first.get(5, TimeUnit.SECONDS)).isEqualTo(true); + assertThat(second.get(5, TimeUnit.SECONDS)).isEqualTo(true); + } finally { + release.countDown(); + restored.close(); + write.close(); + } + assertThat(executor.getValue().isShutdown()).isTrue(); } @Test @@ -148,10 +200,11 @@ private static class TestingFileStoreWrite extends AbstractFileStoreWrite restoredWriter; private TestingFileStoreWrite() { - this(null); + this(null, 1); } - private TestingFileStoreWrite(@Nullable RecordWriter restoredWriter) { + private TestingFileStoreWrite( + @Nullable RecordWriter restoredWriter, int buildParallelism) { super( mock(SnapshotManager.class), mock(FileStoreScan.class), @@ -159,26 +212,26 @@ private TestingFileStoreWrite(@Nullable RecordWriter restoredWriter) { null, null, "test-table", - new CoreOptions(new HashMap<>()), + new CoreOptions( + Collections.singletonMap( + CoreOptions.PK_INDEX_BUILD_PARALLELISM.key(), + String.valueOf(buildParallelism))), RowType.of()); this.restoredWriter = restoredWriter; } private void install( RecordWriter writer, BucketedVectorIndexMaintainer... vectorMaintainers) { + install( + writer, + BucketedPrimaryKeyIndexMaintainer.ofVectors( + Arrays.asList(vectorMaintainers), null, Collections.emptyList())); + } + + private void install( + RecordWriter writer, BucketedPrimaryKeyIndexMaintainer maintainer) { writers.computeIfAbsent(EMPTY_ROW, ignored -> new HashMap<>()) - .put( - 0, - new WriterContainer<>( - writer, - 1, - null, - null, - BucketedPrimaryKeyIndexMaintainer.ofVectors( - Arrays.asList(vectorMaintainers), - null, - Collections.emptyList()), - null)); + .put(0, new WriterContainer<>(writer, 1, null, null, maintainer, null)); } @Override 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 495e934d422b..bb0738bbd09f 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 @@ -24,6 +24,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.ValueSource; import java.util.ArrayList; @@ -54,6 +55,7 @@ void testValidPluralPrimaryKeyVectorIndexConfiguration() { @Test void testSupportsMultiplePrimaryKeyVectorIndexColumns() { Map options = multipleIndexOptions(); + options.put(CoreOptions.PK_INDEX_BUILD_PARALLELISM.key(), "2"); 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}"); @@ -63,6 +65,20 @@ void testSupportsMultiplePrimaryKeyVectorIndexColumns() { assertThatCode(() -> validateTableSchema(schema(options))).doesNotThrowAnyException(); } + @ParameterizedTest + @CsvSource({"0,true", "-1,true", "0,false", "-1,false"}) + void testRejectsNonPositiveIndexBuildParallelism(int parallelism, boolean vectorIndex) { + Map options = multipleIndexOptions(); + if (!vectorIndex) { + options.remove(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key()); + options.put(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "payload"); + } + options.put(CoreOptions.PK_INDEX_BUILD_PARALLELISM.key(), String.valueOf(parallelism)); + + assertThatThrownBy(() -> validateTableSchema(schema(options))) + .hasMessageContaining("pk-index.build.parallelism must be greater than zero"); + } + @Test void testRequiresIndexTypeForSecondVectorColumn() { Map options = multipleIndexOptions(); 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 5594f4161d8a..e033a9430bca 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 @@ -42,6 +42,8 @@ import org.apache.paimon.types.RowKind; 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; @@ -87,9 +89,12 @@ private Schema vectorSchema( return builder.build(); } - @Test - void testIndependentVectorIndexesSurviveReopenAndCompaction() throws Exception { - catalog.createTable(identifier(), multipleVectorSchema(2, "l2", "l2"), false); + @ParameterizedTest + @ValueSource(ints = {1, 2}) + void testIndependentVectorIndexesSurviveReopenAndCompaction(int buildParallelism) + throws Exception { + catalog.createTable( + identifier(), multipleVectorSchema(2, "l2", "l2", buildParallelism), false); FileStoreTable table = getTableDefault(); write( table, @@ -131,9 +136,12 @@ void testIndependentVectorIndexesSurviveReopenAndCompaction() throws Exception { .containsExactlyInAnyOrder(1, 2); } - @Test - void testIndependentVectorIndexesWithDifferentDimensionsAndMetrics() throws Exception { - catalog.createTable(identifier(), multipleVectorSchema(3, "l2", "cosine"), false); + @ParameterizedTest + @ValueSource(ints = {1, 2}) + void testIndependentVectorIndexesWithDifferentDimensionsAndMetrics(int buildParallelism) + throws Exception { + catalog.createTable( + identifier(), multipleVectorSchema(3, "l2", "cosine", buildParallelism), 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))); @@ -158,9 +166,11 @@ void testIndependentVectorIndexesWithDifferentDimensionsAndMetrics() throws Exce assertThat(readIds(reopened, results.get(1))).containsExactly(1); } - @Test - void testCompactionBackfillsAddedVectorIndexWithoutNewData() throws Exception { - Schema schema = multipleVectorSchema(2, "l2", "l2"); + @ParameterizedTest + @ValueSource(ints = {1, 2}) + void testCompactionBackfillsAddedVectorIndexWithoutNewData(int buildParallelism) + throws Exception { + Schema schema = multipleVectorSchema(2, "l2", "l2", buildParallelism); 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); @@ -225,12 +235,16 @@ void testCompactionBackfillsAddedVectorIndexWithoutNewData() throws Exception { completedPayloads.stream().map(IndexFileMeta::fileName).collect(toList())); } - private Schema multipleVectorSchema(int textDimension, String imageMetric, String textMetric) { + private Schema multipleVectorSchema( + int textDimension, String imageMetric, String textMetric, int buildParallelism) { 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.PK_INDEX_BUILD_PARALLELISM.key(), + String.valueOf(buildParallelism)) .option(CoreOptions.BUCKET.key(), "1") .option(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true") .option(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "image_embedding,text_embedding")