This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 0a8d042604 [core][spark] Flatten scored global index result unions
(#8809)
0a8d042604 is described below
commit 0a8d042604c1f96f0ce9b552d29c38f005005326
Author: QuakeWang <[email protected]>
AuthorDate: Sun Jul 26 15:56:55 2026 +0800
[core][spark] Flatten scored global index result unions (#8809)
---
.../globalindex/ScoredGlobalIndexResult.java | 27 ++++++
.../paimon/globalindex/UnionGlobalIndexReader.java | 16 ++--
.../globalindex/ScoredGlobalIndexResultTest.java | 77 +++++++++++++++++
.../table/source/DataEvolutionBatchVectorRead.java | 15 ++--
.../table/source/DataEvolutionFullTextRead.java | 6 +-
.../table/source/DataEvolutionVectorRead.java | 7 +-
.../spark/read/SparkDataEvolutionVectorRead.java | 20 +++--
.../read/SparkDataEvolutionVectorReadTest.java | 99 ++++++++++++++++++++++
8 files changed, 238 insertions(+), 29 deletions(-)
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 3faeeee7ba..f256b8c473 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 @@ package org.apache.paimon.globalindex;
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 @@ public interface ScoredGlobalIndexResult extends
GlobalIndexResult {
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() {
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 c901d4905c..473eb24d0a 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 class UnionGlobalIndexReader implements
GlobalIndexReader {
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) -> {
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 756e7733cf..48f4e120b2 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.apache.paimon.utils.RoaringNavigableMap64;
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 =
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 9e598b783d..9df3e2df0b 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 @@ public class DataEvolutionBatchVectorRead extends
AbstractDataEvolutionVectorRea
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);
}
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 7e759c648d..d9100c2838 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 @@ public class DataEvolutionFullTextRead implements
FullTextRead {
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
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 8d32363146..967aecd6bb 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 @@ public class DataEvolutionVectorRead extends
AbstractDataEvolutionVectorRead imp
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);
}
}
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 ddac704b66..491fa035ef 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 @@ public class SparkDataEvolutionVectorRead extends
DataEvolutionVectorRead {
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;
}
@@ -213,7 +214,7 @@ public class SparkDataEvolutionVectorRead extends
DataEvolutionVectorRead {
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 =
@@ -222,9 +223,10 @@ public class SparkDataEvolutionVectorRead extends
DataEvolutionVectorRead {
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 @@ public class SparkDataEvolutionVectorRead extends
DataEvolutionVectorRead {
}
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)
{
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 756dff2cb6..b89efd1afb 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.GlobalIndexFileReader;
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.RoaringNavigableMap64;
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.ArrayList;
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 @@ import static org.assertj.core.api.Assertions.assertThat;
/** 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 class SparkDataEvolutionVectorReadTest {
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<IndexVectorSearchSplit> indexSplits(String indexType,
int count) {
List<IndexVectorSearchSplit> splits = new ArrayList<>();
for (int i = 0; i < count; i++) {
@@ -228,6 +274,59 @@ public class SparkDataEvolutionVectorReadTest {
}
}
+ 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<RoaringNavigableMap64>
preFilters(List<IndexVectorSearchSplit> splits) {
+ return Collections.emptyList();
+ }
+
+ @Override
+ protected <I, O> List<O> mapInSpark(
+ List<I> data, SerializableFunction<I, O> func, int
parallelism) {
+ sparkParallelism = parallelism;
+ assertThat(data).hasSize(1);
+ return data.stream().map(func::apply).collect(Collectors.toList());
+ }
+
+ @Override
+ protected CompletableFuture<Optional<ScoredGlobalIndexResult>> eval(
+ GlobalIndexer globalIndexer,
+ IndexPathFactory indexPathFactory,
+ long rowRangeStart,
+ long rowRangeEnd,
+ List<IndexFileMeta> 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();
