Skip to content

[core] Search fresh level-0 files in PK vector search full modes #9896

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
LuciferYang wants to merge 1 commit into from
+148 −10
Closed

[core] Search fresh level-0 files in PK vector search full modes #9896

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 13 additions & 8 deletions paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -273,6 +271,16 @@ protected SearchResult mergeSearchResults(List<SearchResult> 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<DataFileMeta> filesToSearch(DataSplit dataSplit) {
return dataSplit.dataFiles();
}

CompletableFuture<SearchResult> searchAsync(
BucketVectorSearchSplit split, SearchContext context) throws IOException {
return searchBatchAsync(split, context, new float[][] {query})
Expand All @@ -283,13 +291,10 @@ CompletableFuture<List<SearchResult>> searchBatchAsync(
BucketVectorSearchSplit split, SearchContext context, float[][] queries)
throws IOException {
DataSplit dataSplit = split.dataSplit();
List<DataFileMeta> activeFiles =
dataSplit.dataFiles().stream()
.filter(PrimaryKeyIndexSourcePolicy::shouldRead)
.collect(Collectors.toList());
List<DataFileMeta> dataFiles = filesToSearch(dataSplit);
PkVectorBucketIndexState state =
PkVectorBucketIndexState.fromActiveDataFiles(
vectorField.id(), indexType, activeFiles, split.payloadFiles());
vectorField.id(), indexType, dataFiles, split.payloadFiles());
Map<String, DeletionVector> deletionVectors = deletionVectors(dataSplit, context.fileIO);
PkVectorDataFileReader.Factory readerFactory =
new PkVectorDataFileReader.Factory(
Expand Down Expand Up @@ -317,7 +322,7 @@ CompletableFuture<List<SearchResult>> searchBatchAsync(
return bucketSearch
.searchBatchAsync(
state,
activeFiles,
dataFiles,
deletionVectors,
rowRangesByFile(split),
queries,
Expand Down
128 changes: 128 additions & 0 deletions paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

Expand Down Expand Up @@ -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<PkVectorSearchResult> 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.<String>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(
Expand Down
9 changes: 7 additions & 2 deletions paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
Expand Down
Loading