From 1bfef2f6899a0c1144d89776f73ff11d746cf08c Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 15 Sep 2026 13:23:11 +0800 Subject: [PATCH] [core] Search fresh level-0 files in PK vector search full modes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PrimaryKeyVectorRead handed only PrimaryKeyIndexSourcePolicy-filtered files (complete COMPACT level>0 outputs) to the bucket search, so Level-0 APPEND files — freshly written rows between compactions — were omitted by every search mode, including full and detail, which are documented to scan unindexed data files exactly. Pass all data files of the split instead: the exact fallback covers the files the ANN segments do not, while FAST keeps ignoring them, and the ANN state builder keeps applying the compact-file policy itself. The synchronous write harness cannot keep a stable level-0 window (the write path compacts eagerly), so the regression coverage pins both halves of the chain: filesToSearch must include an APPEND level-0 file that shouldRead rejects, and the bucket search must exact-scan exactly such a file in full mode while FAST still skips it. Assisted-by: GLM-5.3 --- .../table/source/PrimaryKeyVectorRead.java | 21 +-- .../source/PrimaryKeyVectorReadTest.java | 128 ++++++++++++++++++ .../source/PrimaryKeyVectorSearchTest.java | 9 +- 3 files changed, 148 insertions(+), 10 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java index 6a632ed64607..23d9c137117e 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java @@ -28,7 +28,6 @@ import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.globalindex.VectorSearchMetric; import org.apache.paimon.index.IndexFileHandler; -import org.apache.paimon.index.pk.PrimaryKeyIndexSourcePolicy; import org.apache.paimon.index.pkvector.PkVectorAnnSegmentSearcher; import org.apache.paimon.index.pkvector.PkVectorBucketIndexState; import org.apache.paimon.index.pkvector.PkVectorDataFileReader; @@ -68,7 +67,6 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutorService; -import java.util.stream.Collectors; import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM; import static org.apache.paimon.globalindex.VectorSearchMetric.normalize; @@ -273,6 +271,16 @@ protected SearchResult mergeSearchResults(List results) { topK(indexedCandidates, indexedLimit), topK(exactCandidates, limit)); } + /** + * Files eligible for the bucket search: every data file of the split. Level-0 APPEND files are + * never index sources, but full/detail search modes must exact-scan them instead of silently + * omitting freshly written rows; FAST still ignores uncovered files inside the bucket search. + * The ANN state builder applies the compact-file policy itself. + */ + static List filesToSearch(DataSplit dataSplit) { + return dataSplit.dataFiles(); + } + CompletableFuture searchAsync( BucketVectorSearchSplit split, SearchContext context) throws IOException { return searchBatchAsync(split, context, new float[][] {query}) @@ -283,13 +291,10 @@ CompletableFuture> searchBatchAsync( BucketVectorSearchSplit split, SearchContext context, float[][] queries) throws IOException { DataSplit dataSplit = split.dataSplit(); - List activeFiles = - dataSplit.dataFiles().stream() - .filter(PrimaryKeyIndexSourcePolicy::shouldRead) - .collect(Collectors.toList()); + List dataFiles = filesToSearch(dataSplit); PkVectorBucketIndexState state = PkVectorBucketIndexState.fromActiveDataFiles( - vectorField.id(), indexType, activeFiles, split.payloadFiles()); + vectorField.id(), indexType, dataFiles, split.payloadFiles()); Map deletionVectors = deletionVectors(dataSplit, context.fileIO); PkVectorDataFileReader.Factory readerFactory = new PkVectorDataFileReader.Factory( @@ -317,7 +322,7 @@ CompletableFuture> searchBatchAsync( return bucketSearch .searchBatchAsync( state, - activeFiles, + dataFiles, deletionVectors, rowRangesByFile(split), queries, diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java index 0164c38cbd05..44b840108866 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java @@ -19,7 +19,16 @@ package org.apache.paimon.table.source; import org.apache.paimon.CoreOptions; +import org.apache.paimon.CoreOptions.GlobalIndexSearchMode; import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourcePolicy; +import org.apache.paimon.index.pkvector.PkVectorBucketIndexState; +import org.apache.paimon.index.pkvector.PkVectorDataFileReader; +import org.apache.paimon.index.pkvector.PkVectorSearchResult; +import org.apache.paimon.index.pkvector.PrimaryKeyVectorBucketSearch; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.stats.SimpleStats; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.types.DataField; import org.apache.paimon.types.FloatType; @@ -38,8 +47,11 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -101,6 +113,122 @@ void testMergesGlobalTopKWithDeterministicTies() { org.assertj.core.groups.Tuple.tuple(1, "file-b", 1L)); } + @Test + void testFilesToSearchIncludesAppendLevel0Files() { + // a fresh APPEND level-0 file is not an index source (shouldRead is false) but + // must still reach the bucket search: full/detail modes exact-scan it + DataFileMeta compact = file("data-compact.parquet", 1, FileSource.COMPACT); + DataFileMeta append = file("data-append.parquet", 0, FileSource.APPEND); + + assertThat(PrimaryKeyIndexSourcePolicy.shouldRead(append)).isFalse(); + assertThat(PrimaryKeyIndexSourcePolicy.shouldRead(compact)).isTrue(); + + DataSplit split = + DataSplit.builder() + .withSnapshot(1L) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withDataFiles(Arrays.asList(compact, append)) + .build(); + + assertThat(PrimaryKeyVectorRead.filesToSearch(split)) + .extracting(DataFileMeta::fileName) + .containsExactly("data-compact.parquet", "data-append.parquet"); + } + + @Test + void testFullModeBucketSearchExactScansAppendLevel0Files() throws Exception { + // no ANN payloads: the fresh level-0 APPEND file is only reachable through the + // exact path, which full mode must run over it + DataFileMeta append = file("data-append.parquet", 0, FileSource.APPEND); + PkVectorBucketIndexState state = + PkVectorBucketIndexState.fromActiveDataFiles( + 1, + "test-vector-ann", + Collections.singletonList(append), + Collections.emptyList()); + assertThat(state.annSegments()).isEmpty(); + + PkVectorDataFileReader.Factory factory = mock(PkVectorDataFileReader.Factory.class); + PkVectorDataFileReader reader = mock(PkVectorDataFileReader.class); + when(reader.dimension()).thenReturn(2); + when(reader.rowCount()).thenReturn(2L); + AtomicInteger read = new AtomicInteger(); + doAnswer( + invocation -> { + float[] reuse = invocation.getArgument(0, float[].class); + int index = read.getAndIncrement(); + if (index == 0) { + reuse[0] = 1; + reuse[1] = 0; + return true; + } + if (index == 1) { + reuse[0] = 100; + reuse[1] = 0; + return true; + } + return false; + }) + .when(reader) + .readNextVector(any()); + when(factory.create(append)).thenReturn(reader); + + PrimaryKeyVectorBucketSearch search = + new PrimaryKeyVectorBucketSearch( + factory, null, Collections.emptyMap(), "l2", GlobalIndexSearchMode.FULL); + List results = + search.search( + state, + Collections.singletonList(append), + Collections.emptyMap(), + new float[] {0, 0}, + 1); + + assertThat(results).hasSize(1); + assertThat(results.get(0).dataFileName()).isEqualTo("data-append.parquet"); + assertThat(results.get(0).rowPosition()).isZero(); + + // FAST mode still omits files the index does not cover + PrimaryKeyVectorBucketSearch fast = + new PrimaryKeyVectorBucketSearch( + factory, null, Collections.emptyMap(), "l2", GlobalIndexSearchMode.FAST); + assertThat( + fast.search( + state, + Collections.singletonList(append), + Collections.emptyMap(), + new float[] {0, 0}, + 1)) + .isEmpty(); + } + + private static DataFileMeta file(String fileName, int level, FileSource fileSource) { + // a fresh write: APPEND source at level 0 (never an index source); a complete + // compacted output: COMPACT source at level 1 (an index source) + return DataFileMeta.create( + fileName, + 1024L, + 2L, + BinaryRow.EMPTY_ROW, + BinaryRow.EMPTY_ROW, + SimpleStats.EMPTY_STATS, + SimpleStats.EMPTY_STATS, + 0L, + 1L, + 1L, + level, + Collections.emptyList(), + null, + null, + fileSource, + null, + null, + null, + null); + } + private static PrimaryKeyVectorRead.Candidate candidate( int bucket, String fileName, long position, float distance) { return new PrimaryKeyVectorRead.Candidate( 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..e6b21044ef20 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 @@ -55,7 +55,7 @@ private Schema vectorSchema(String mergeEngine, boolean deletionVectorsEnabled) return vectorSchema(mergeEngine, deletionVectorsEnabled, false); } - private Schema vectorSchema( + private Schema.Builder vectorSchemaBuilder( String mergeEngine, boolean deletionVectorsEnabled, boolean reverseScore) { Schema.Builder builder = Schema.newBuilder() @@ -77,7 +77,12 @@ private Schema vectorSchema( if (reverseScore) { builder.option("test.vector.reverse-score", "true"); } - return builder.build(); + return builder; + } + + private Schema vectorSchema( + String mergeEngine, boolean deletionVectorsEnabled, boolean reverseScore) { + return vectorSchemaBuilder(mergeEngine, deletionVectorsEnabled, reverseScore).build(); } @Test