Skip to content

[core][spark] Flatten scored global index result unions #8809

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

Open
QuakeWang wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[core][spark] Flatten scored global index result unions #8809

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c70f6f3
[core][spark] Flatten scored global index result unions
QuakeWang Jul 22, 2026
d7eae93
Merge Spark raw vector split results in bulk
QuakeWang Jul 23, 2026
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
27 changes: 27 additions & 0 deletions paimon-common/src/main/java/org/apache/paimon/globalindex/ScoredGlobalIndexResult.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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. */
Expand Down Expand Up @@ -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<ScoredGlobalIndexResult> results) {
if (results.isEmpty()) {
return createEmpty();
}
if (results.size() == 1) {
return results.get(0);
}

RoaringNavigableMap64 mergedRowIds = new RoaringNavigableMap64();
Map<Long, Float> 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() {
Expand Down
16 changes: 7 additions & 9 deletions paimon-common/src/main/java/org/apache/paimon/globalindex/UnionGlobalIndexReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,19 +154,17 @@ public CompletableFuture<Optional<ScoredGlobalIndexResult>> visitVectorSearch(
return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
.thenApply(
v -> {
Optional<ScoredGlobalIndexResult> result = Optional.empty();
List<ScoredGlobalIndexResult> results = new ArrayList<>(futures.size());
for (CompletableFuture<Optional<ScoredGlobalIndexResult>> f : futures) {
Optional<ScoredGlobalIndexResult> 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.<ScoredGlobalIndexResult>empty();
}
return Optional.of(ScoredGlobalIndexResult.merge(results));
})
.whenComplete(
(ignored, throwable) -> {
Expand Down
77 changes: 77 additions & 0 deletions paimon-common/src/test/java/org/apache/paimon/globalindex/ScoredGlobalIndexResultTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<ScoredGlobalIndexResult> 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 =
Expand Down
15 changes: 10 additions & 5 deletions paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionBatchVectorRead.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,19 +126,24 @@ protected ScoredGlobalIndexResult[] readIndexedBatch(

CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();

ScoredGlobalIndexResult[] merged = new ScoredGlobalIndexResult[n];
List<List<ScoredGlobalIndexResult>> resultsByQuery = new ArrayList<>(n);
for (int i = 0; i < n; i++) {
merged[i] = ScoredGlobalIndexResult.createEmpty();
resultsByQuery.add(new ArrayList<>());
}

for (CompletableFuture<List<Optional<ScoredGlobalIndexResult>>> future : futures) {
List<Optional<ScoredGlobalIndexResult>> 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<ScoredGlobalIndexResult> 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);
}

Expand Down
6 changes: 3 additions & 3 deletions paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionFullTextRead.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,15 +171,15 @@ private ScoredGlobalIndexResult evalColumnQuery(

CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();

ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
List<ScoredGlobalIndexResult> results = new ArrayList<>(futures.size());
for (CompletableFuture<Optional<ScoredGlobalIndexResult>> f : futures) {
Optional<ScoredGlobalIndexResult> next = f.join();
if (next.isPresent()) {
result = result.or(next.get());
results.add(next.get());
}
}

return result;
return ScoredGlobalIndexResult.merge(results);
}

@Nullable
Expand Down
7 changes: 4 additions & 3 deletions paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorRead.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,13 +111,14 @@ protected ScoredGlobalIndexResult readIndexed(

CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();

ScoredGlobalIndexResult merged = ScoredGlobalIndexResult.createEmpty();
List<ScoredGlobalIndexResult> results = new ArrayList<>(futures.size());
for (CompletableFuture<Optional<ScoredGlobalIndexResult>> future : futures) {
Optional<ScoredGlobalIndexResult> 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);
}
}
20 changes: 11 additions & 9 deletions ...spark-common/src/main/java/org/apache/paimon/spark/read/SparkDataEvolutionVectorRead.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,14 +149,15 @@ protected ScoredGlobalIndexResult readIndexSplitsInSpark(
executor));
}
CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
List<ScoredGlobalIndexResult> results = new ArrayList<>(futures.size());
for (CompletableFuture<Optional<ScoredGlobalIndexResult>> f : futures) {
Optional<ScoredGlobalIndexResult> 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;
}
Expand Down Expand Up @@ -213,7 +214,7 @@ protected ScoredGlobalIndexResult readRawSplitsInSpark(

SerializableFunction<List<SerializedSplit>, byte[]> task =
group -> {
ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
List<ScoredGlobalIndexResult> results = new ArrayList<>(group.size());
for (SerializedSplit serializedSplit : group) {
List<Range> rowRanges = deserializeRanges(serializedSplit.split);
ScoredGlobalIndexResult splitResult =
Expand All @@ -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;
}
Expand Down Expand Up @@ -302,18 +304,18 @@ private ScoredGlobalIndexResult mergeRemoteResults(List<byte[]> remoteResults) {
}

private ScoredGlobalIndexResult mergeRemoteResults(List<byte[]> remoteResults, int topK) {
ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
List<ScoredGlobalIndexResult> 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<List<Range>> rangeGroups(List<Range> ranges, int parallelism) {
Expand Down
Loading
Loading