diff --git a/paimon-common/src/main/java/org/apache/paimon/globalindex/ScoredGlobalIndexResult.java b/paimon-common/src/main/java/org/apache/paimon/globalindex/ScoredGlobalIndexResult.java index 3faeeee7bac5..f256b8c473a7 100644 --- a/paimon-common/src/main/java/org/apache/paimon/globalindex/ScoredGlobalIndexResult.java +++ b/paimon-common/src/main/java/org/apache/paimon/globalindex/ScoredGlobalIndexResult.java @@ -21,6 +21,9 @@ import org.apache.paimon.utils.RoaringNavigableMap64; import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; import java.util.PriorityQueue; /** Vector search global index result for scored index. */ @@ -119,6 +122,30 @@ static ScoredGlobalIndexResult createEmpty() { return create(new RoaringNavigableMap64(), rowId -> 0); } + /** Merges scored results, keeping the score from the first result containing each row ID. */ + static ScoredGlobalIndexResult merge(List results) { + if (results.isEmpty()) { + return createEmpty(); + } + if (results.size() == 1) { + return results.get(0); + } + + RoaringNavigableMap64 mergedRowIds = new RoaringNavigableMap64(); + Map mergedScores = new HashMap<>(); + for (ScoredGlobalIndexResult result : results) { + RoaringNavigableMap64 rowIds = result.results(); + ScoreGetter scoreGetter = result.scoreGetter(); + for (long rowId : rowIds) { + if (!mergedRowIds.contains(rowId)) { + mergedRowIds.add(rowId); + mergedScores.put(rowId, scoreGetter.score(rowId)); + } + } + } + return create(mergedRowIds, mergedScores::get); + } + /** Returns a new {@link ScoredGlobalIndexResult} from bitmap. */ static ScoredGlobalIndexResult create(RoaringNavigableMap64 bitmap, ScoreGetter scoreGetter) { return new ScoredGlobalIndexResult() { diff --git a/paimon-common/src/main/java/org/apache/paimon/globalindex/UnionGlobalIndexReader.java b/paimon-common/src/main/java/org/apache/paimon/globalindex/UnionGlobalIndexReader.java index c901d4905c86..473eb24d0a84 100644 --- a/paimon-common/src/main/java/org/apache/paimon/globalindex/UnionGlobalIndexReader.java +++ b/paimon-common/src/main/java/org/apache/paimon/globalindex/UnionGlobalIndexReader.java @@ -154,19 +154,17 @@ public CompletableFuture> visitVectorSearch( return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])) .thenApply( v -> { - Optional result = Optional.empty(); + List results = new ArrayList<>(futures.size()); for (CompletableFuture> f : futures) { Optional current = f.join(); - if (!current.isPresent()) { - continue; - } - if (!result.isPresent()) { - result = current; - } else { - result = Optional.of(result.get().or(current.get())); + if (current.isPresent()) { + results.add(current.get()); } } - return result; + if (results.isEmpty()) { + return Optional.empty(); + } + return Optional.of(ScoredGlobalIndexResult.merge(results)); }) .whenComplete( (ignored, throwable) -> { diff --git a/paimon-common/src/test/java/org/apache/paimon/globalindex/ScoredGlobalIndexResultTest.java b/paimon-common/src/test/java/org/apache/paimon/globalindex/ScoredGlobalIndexResultTest.java index 756e7733cfef..48f4e120b271 100644 --- a/paimon-common/src/test/java/org/apache/paimon/globalindex/ScoredGlobalIndexResultTest.java +++ b/paimon-common/src/test/java/org/apache/paimon/globalindex/ScoredGlobalIndexResultTest.java @@ -22,14 +22,91 @@ import org.junit.jupiter.api.Test; +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 java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; /** Tests for {@link ScoredGlobalIndexResult}. */ public class ScoredGlobalIndexResultTest { + @Test + public void testMergeEmptyResults() { + ScoredGlobalIndexResult merged = ScoredGlobalIndexResult.merge(Collections.emptyList()); + + assertThat(merged.results()).isEmpty(); + } + + @Test + public void testMergeSingleResult() { + ScoredGlobalIndexResult result = result(new long[] {1}, new float[] {0.1f}); + + assertThat(ScoredGlobalIndexResult.merge(Collections.singletonList(result))) + .isSameAs(result); + } + + @Test + public void testMergeDisjointResults() { + ScoredGlobalIndexResult first = result(new long[] {1, 3}, new float[] {0.1f, 0.3f}); + ScoredGlobalIndexResult second = result(new long[] {2, 4}, new float[] {0.2f, 0.4f}); + + ScoredGlobalIndexResult merged = + ScoredGlobalIndexResult.merge(Arrays.asList(first, second)); + + assertThat(merged.results().getIntCardinality()).isEqualTo(4); + assertThat(merged.results()).contains(1L, 2L, 3L, 4L); + assertThat(merged.scoreGetter().score(1L)).isEqualTo(0.1f); + assertThat(merged.scoreGetter().score(2L)).isEqualTo(0.2f); + assertThat(merged.scoreGetter().score(3L)).isEqualTo(0.3f); + assertThat(merged.scoreGetter().score(4L)).isEqualTo(0.4f); + } + + @Test + public void testMergeKeepsFirstScoreForDuplicateRowId() { + ScoredGlobalIndexResult first = result(new long[] {1, 2}, new float[] {0.1f, 0.2f}); + ScoredGlobalIndexResult second = result(new long[] {2, 3}, new float[] {2.0f, 0.3f}); + + ScoredGlobalIndexResult merged = + ScoredGlobalIndexResult.merge(Arrays.asList(first, second)); + + assertThat(merged.results().getIntCardinality()).isEqualTo(3); + assertThat(merged.scoreGetter().score(2L)).isEqualTo(0.2f); + } + + @Test + public void testMergeManyResultsMaterializesScoresBeforeTopK() { + int resultCount = 5000; + AtomicInteger scoreCalls = new AtomicInteger(); + List results = new ArrayList<>(resultCount); + for (int i = 0; i < resultCount; i++) { + RoaringNavigableMap64 rowIds = new RoaringNavigableMap64(); + rowIds.add(i); + final float score = i; + results.add( + ScoredGlobalIndexResult.create( + rowIds, + ignored -> { + scoreCalls.incrementAndGet(); + return score; + })); + } + + ScoredGlobalIndexResult merged = ScoredGlobalIndexResult.merge(results); + int scoreCallsAfterMerge = scoreCalls.get(); + RoaringNavigableMap64 topK = merged.topK(1).results(); + + assertThat(merged.results().getIntCardinality()).isEqualTo(resultCount); + assertThat(scoreCallsAfterMerge).isEqualTo(resultCount); + assertThat(scoreCalls).hasValue(scoreCallsAfterMerge); + assertThat(topK.getIntCardinality()).isEqualTo(1); + assertThat(topK).contains(resultCount - 1L); + } + @Test public void testTopKBreaksBoundaryTiesByRowId() { ScoredGlobalIndexResult result = diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionBatchVectorRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionBatchVectorRead.java index 9e598b783d4c..9df3e2df0bf3 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionBatchVectorRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionBatchVectorRead.java @@ -126,19 +126,24 @@ protected ScoredGlobalIndexResult[] readIndexedBatch( CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); - ScoredGlobalIndexResult[] merged = new ScoredGlobalIndexResult[n]; + List> resultsByQuery = new ArrayList<>(n); for (int i = 0; i < n; i++) { - merged[i] = ScoredGlobalIndexResult.createEmpty(); + resultsByQuery.add(new ArrayList<>()); } - for (CompletableFuture>> future : futures) { List> splitResults = future.join(); for (int i = 0; i < n; i++) { - if (splitResults.get(i).isPresent()) { - merged[i] = merged[i].or(splitResults.get(i).get()); + Optional splitResult = splitResults.get(i); + if (splitResult.isPresent()) { + resultsByQuery.get(i).add(splitResult.get()); } } } + + ScoredGlobalIndexResult[] merged = new ScoredGlobalIndexResult[n]; + for (int i = 0; i < n; i++) { + merged[i] = ScoredGlobalIndexResult.merge(resultsByQuery.get(i)); + } return maybeRerankIndexedBatchResults(merged, indexType, globalIndexer); } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionFullTextRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionFullTextRead.java index 7e759c648d6c..d9100c2838df 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionFullTextRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionFullTextRead.java @@ -171,15 +171,15 @@ private ScoredGlobalIndexResult evalColumnQuery( CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); - ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty(); + List results = new ArrayList<>(futures.size()); for (CompletableFuture> f : futures) { Optional next = f.join(); if (next.isPresent()) { - result = result.or(next.get()); + results.add(next.get()); } } - return result; + return ScoredGlobalIndexResult.merge(results); } @Nullable diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorRead.java index 8d3236314661..967aecd6bb18 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorRead.java @@ -111,13 +111,14 @@ protected ScoredGlobalIndexResult readIndexed( CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); - ScoredGlobalIndexResult merged = ScoredGlobalIndexResult.createEmpty(); + List results = new ArrayList<>(futures.size()); for (CompletableFuture> future : futures) { Optional splitResult = future.join(); if (splitResult.isPresent()) { - merged = merged.or(splitResult.get()); + results.add(splitResult.get()); } } - return maybeRerankIndexedResult(merged, indexType, globalIndexer, vector); + return maybeRerankIndexedResult( + ScoredGlobalIndexResult.merge(results), indexType, globalIndexer, vector); } } diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorRead.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorRead.java index ddac704b667a..491fa035ef34 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorRead.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorRead.java @@ -149,14 +149,15 @@ protected ScoredGlobalIndexResult readIndexSplitsInSpark( executor)); } CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); - ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty(); + List results = new ArrayList<>(futures.size()); for (CompletableFuture> f : futures) { Optional next = f.join(); if (next.isPresent()) { - result = result.or(next.get()); + results.add(next.get()); } } - result = result.topK(searchLimit); + ScoredGlobalIndexResult result = + ScoredGlobalIndexResult.merge(results).topK(searchLimit); if (result.results().isEmpty()) { return null; } @@ -213,7 +214,7 @@ protected ScoredGlobalIndexResult readRawSplitsInSpark( SerializableFunction, byte[]> task = group -> { - ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty(); + List results = new ArrayList<>(group.size()); for (SerializedSplit serializedSplit : group) { List rowRanges = deserializeRanges(serializedSplit.split); ScoredGlobalIndexResult splitResult = @@ -222,9 +223,10 @@ protected ScoredGlobalIndexResult readRawSplitsInSpark( deserializePreFilter(serializedSplit.preFilter), metric, vector); - result = result.or(splitResult); + results.add(splitResult); } - result = result.topK(limit); + ScoredGlobalIndexResult result = + ScoredGlobalIndexResult.merge(results).topK(limit); if (result.results().isEmpty()) { return null; } @@ -302,18 +304,18 @@ private ScoredGlobalIndexResult mergeRemoteResults(List remoteResults) { } private ScoredGlobalIndexResult mergeRemoteResults(List remoteResults, int topK) { - ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty(); + List results = new ArrayList<>(remoteResults.size()); GlobalIndexResultSerializer serializer = new GlobalIndexResultSerializer(); for (byte[] bytes : remoteResults) { if (bytes != null) { try { - result = result.or(serializer.deserialize(bytes)); + results.add(serializer.deserialize(bytes)); } catch (IOException e) { throw new RuntimeException("Failed to deserialize ScoredGlobalIndexResult", e); } } } - return result.topK(topK); + return ScoredGlobalIndexResult.merge(results).topK(topK); } private List> rangeGroups(List ranges, int parallelism) { diff --git a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorReadTest.java b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorReadTest.java index 756dff2cb698..b89efd1afbf8 100644 --- a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorReadTest.java +++ b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorReadTest.java @@ -18,6 +18,9 @@ package org.apache.paimon.spark.read; +import org.apache.paimon.CoreOptions; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; import org.apache.paimon.globalindex.GlobalIndexIOMeta; import org.apache.paimon.globalindex.GlobalIndexReader; import org.apache.paimon.globalindex.GlobalIndexResult; @@ -30,6 +33,14 @@ import org.apache.paimon.globalindex.io.GlobalIndexFileWriter; import org.apache.paimon.index.GlobalIndexMeta; import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.IndexPathFactory; +import org.apache.paimon.options.Options; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.schema.SchemaManager; +import org.apache.paimon.schema.SchemaUtils; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.FileStoreTableFactory; import org.apache.paimon.table.source.IndexVectorSearchSplit; import org.apache.paimon.table.source.RawVectorSearchSplit; import org.apache.paimon.table.source.VectorScan; @@ -42,6 +53,7 @@ import org.apache.paimon.utils.SerializableFunction; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; import javax.annotation.Nullable; @@ -50,6 +62,8 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; @@ -59,6 +73,8 @@ /** Tests for {@link SparkDataEvolutionVectorRead}. */ public class SparkDataEvolutionVectorReadTest { + @TempDir java.nio.file.Path tempDir; + @Test public void testRawSearchUsesSparkPath() { TestingSparkVectorRead read = new TestingSparkVectorRead(); @@ -103,6 +119,36 @@ public void testDistributedIndexRefinesAfterGlobalMerge() { assertThat(result.results().contains(0L)).isTrue(); } + @Test + public void testDistributedIndexMaterializesManySplitsInSingleTask() throws Exception { + int splitCount = 5000; + SingleTaskSparkVectorRead read = new SingleTaskSparkVectorRead(createTable()); + + ScoredGlobalIndexResult result = + read.readIndexSplitsInSpark( + indexSplits("test-vector-ann", splitCount), new L2Indexer()); + + assertThat(read.sparkParallelism).isEqualTo(1); + assertThat(read.scoreCalls).hasValue(splitCount); + assertThat(result.results().getLongCardinality()).isEqualTo(1); + assertThat(result.results()).contains(splitCount - 1L); + } + + private FileStoreTable createTable() throws Exception { + Path tablePath = new Path(tempDir.toUri()); + Options options = new Options(); + options.set(CoreOptions.PATH, tablePath.toString()); + options.set(CoreOptions.GLOBAL_INDEX_THREAD_NUM, 1); + Schema schema = + Schema.newBuilder() + .column("vec", new ArrayType(DataTypes.FLOAT())) + .options(options.toMap()) + .build(); + TableSchema tableSchema = + SchemaUtils.forceCommit(new SchemaManager(LocalFileIO.create(), tablePath), schema); + return FileStoreTableFactory.create(LocalFileIO.create(), tablePath, tableSchema); + } + private static List indexSplits(String indexType, int count) { List splits = new ArrayList<>(); for (int i = 0; i < count; i++) { @@ -228,6 +274,59 @@ private static ScoredGlobalIndexResult scoredResult(long rowId, float score) { } } + private static class SingleTaskSparkVectorRead extends SparkDataEvolutionVectorRead { + + private final AtomicInteger scoreCalls = new AtomicInteger(); + private int sparkParallelism; + + private SingleTaskSparkVectorRead(FileStoreTable table) { + super( + table, + null, + null, + 1, + new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())), + new float[] {0.0f}, + null); + } + + @Override + protected List preFilters(List splits) { + return Collections.emptyList(); + } + + @Override + protected List mapInSpark( + List data, SerializableFunction func, int parallelism) { + sparkParallelism = parallelism; + assertThat(data).hasSize(1); + return data.stream().map(func::apply).collect(Collectors.toList()); + } + + @Override + protected CompletableFuture> eval( + GlobalIndexer globalIndexer, + IndexPathFactory indexPathFactory, + long rowRangeStart, + long rowRangeEnd, + List vectorIndexFiles, + float[] vector, + int searchLimit, + @Nullable RoaringNavigableMap64 includeRowIds, + ExecutorService executor) { + RoaringNavigableMap64 rows = new RoaringNavigableMap64(); + rows.add(rowRangeStart); + return CompletableFuture.completedFuture( + Optional.of( + ScoredGlobalIndexResult.create( + rows, + rowId -> { + scoreCalls.incrementAndGet(); + return rowId; + }))); + } + } + private static class RecordingSparkVectorRead extends SparkDataEvolutionVectorRead { private final AtomicInteger nextTask = new AtomicInteger();