[core] Support batch vector search

XiaoHongbo-Hope · XiaoHongbo-Hope · commit 6f736638ff64 · 2026-06-02T13:52:11.000+08:00
diff --git a/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java b/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java
@@ -24,6 +24,7 @@
 import org.apache.paimon.predicate.VectorSearch;
 
 import java.io.Closeable;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -59,4 +60,23 @@ default CompletableFuture<Optional<ScoredGlobalIndexResult>> visitFullTextSearch
             FullTextSearch fullTextSearch) {
         throw new UnsupportedOperationException();
     }
+
+    default CompletableFuture<List<Optional<ScoredGlobalIndexResult>>> visitBatchVectorSearch(
+            VectorSearch vectorSearch) {
+        List<CompletableFuture<Optional<ScoredGlobalIndexResult>>> futures = new ArrayList<>();
+        for (int i = 0; i < vectorSearch.vectorCount(); i++) {
+            futures.add(visitVectorSearch(vectorSearch.forIndex(i)));
+        }
+        return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+                .thenApply(
+                        ignored -> {
+                            List<Optional<ScoredGlobalIndexResult>> results =
+                                    new ArrayList<>(futures.size());
+                            for (CompletableFuture<Optional<ScoredGlobalIndexResult>> future :
+                                    futures) {
+                                results.add(future.join());
+                            }
+                            return results;
+                        });
+    }
 }
diff --git a/paimon-common/src/main/java/org/apache/paimon/globalindex/OffsetGlobalIndexReader.java b/paimon-common/src/main/java/org/apache/paimon/globalindex/OffsetGlobalIndexReader.java
@@ -23,6 +23,7 @@
 import org.apache.paimon.predicate.VectorSearch;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -145,6 +146,21 @@ public CompletableFuture<Optional<ScoredGlobalIndexResult>> visitFullTextSearch(
                 .thenApply(opt -> opt.map(r -> r.offset(offset)));
     }
 
+    @Override
+    public CompletableFuture<List<Optional<ScoredGlobalIndexResult>>> visitBatchVectorSearch(
+            VectorSearch vectorSearch) {
+        return wrapped.visitBatchVectorSearch(vectorSearch.offsetRange(this.offset, this.to))
+                .thenApply(
+                        results -> {
+                            List<Optional<ScoredGlobalIndexResult>> offsetResults =
+                                    new ArrayList<>(results.size());
+                            for (Optional<ScoredGlobalIndexResult> result : results) {
+                                offsetResults.add(result.map(r -> r.offset(offset)));
+                            }
+                            return offsetResults;
+                        });
+    }
+
     private Optional<GlobalIndexResult> applyOffset(Optional<GlobalIndexResult> result) {
         return result.map(r -> r.offset(offset));
     }
diff --git a/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
@@ -30,29 +30,54 @@ public class VectorSearch implements Serializable {
 
     private static final long serialVersionUID = 1L;
 
-    private final float[] vector;
+    private final float[][] vectors;
     private final String fieldName;
     private final int limit;
 
     @Nullable private RoaringNavigableMap64 includeRowIds;
 
     public VectorSearch(float[] vector, int limit, String fieldName) {
-        if (vector == null) {
-            throw new IllegalArgumentException("Search cannot be null");
+        this(new float[][] {vector}, limit, fieldName);
+    }
+
+    public VectorSearch(float[][] vectors, int limit, String fieldName) {
+        if (vectors == null || vectors.length == 0) {
+            throw new IllegalArgumentException("Search vectors cannot be null or empty");
+        }
+        for (float[] v : vectors) {
+            if (v == null) {
+                throw new IllegalArgumentException("Search vector element cannot be null");
+            }
         }
         if (limit <= 0) {
             throw new IllegalArgumentException("Limit must be positive, got: " + limit);
         }
         if (fieldName == null || fieldName.isEmpty()) {
             throw new IllegalArgumentException("Field name cannot be null or empty");
         }
-        this.vector = vector;
+        this.vectors = vectors;
         this.limit = limit;
         this.fieldName = fieldName;
     }
 
     public float[] vector() {
-        return vector;
+        return vectors[0];
+    }
+
+    public float[][] vectors() {
+        return vectors;
+    }
+
+    public int vectorCount() {
+        return vectors.length;
+    }
+
+    public VectorSearch forIndex(int i) {
+        VectorSearch single = new VectorSearch(vectors[i], limit, fieldName);
+        if (includeRowIds != null) {
+            single.withIncludeRowIds(includeRowIds);
+        }
+        return single;
     }
 
     public int limit() {
@@ -81,7 +106,7 @@ public VectorSearch offsetRange(long from, long to) {
             for (long rowId : and64) {
                 roaringNavigableMap64Offset.add(rowId - from);
             }
-            VectorSearch target = new VectorSearch(vector, limit, fieldName);
+            VectorSearch target = new VectorSearch(vectors, limit, fieldName);
             target.withIncludeRowIds(roaringNavigableMap64Offset);
             return target;
         }
@@ -90,6 +115,7 @@ public VectorSearch offsetRange(long from, long to) {
 
     @Override
     public String toString() {
-        return String.format("FieldName(%s), Limit(%s)", fieldName, limit);
+        return String.format(
+                "FieldName(%s), Limit(%s), VectorCount(%s)", fieldName, limit, vectors.length);
     }
 }
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorRead.java
@@ -20,6 +20,7 @@
 
 import org.apache.paimon.globalindex.GlobalIndexResult;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /** Vector read to read index files. */
@@ -30,4 +31,14 @@ default GlobalIndexResult read(VectorScan.Plan plan) {
     }
 
     GlobalIndexResult read(List<VectorSearchSplit> splits);
+
+    default List<GlobalIndexResult> readBatch(VectorScan.Plan plan) {
+        return readBatch(plan.splits());
+    }
+
+    default List<GlobalIndexResult> readBatch(List<VectorSearchSplit> splits) {
+        List<GlobalIndexResult> results = new ArrayList<>(1);
+        results.add(read(splits));
+        return results;
+    }
 }
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
@@ -61,25 +61,39 @@ public class VectorReadImpl implements VectorRead {
     private final Predicate filter;
     private final int limit;
     private final DataField vectorColumn;
-    private final float[] vector;
+    private final float[][] vectors;
 
     public VectorReadImpl(
             FileStoreTable table,
             Predicate filter,
             int limit,
             DataField vectorColumn,
-            float[] vector) {
+            float[][] vectors) {
         this.table = table;
         this.filter = filter;
         this.limit = limit;
         this.vectorColumn = vectorColumn;
-        this.vector = vector;
+        this.vectors = vectors;
     }
 
     @Override
     public GlobalIndexResult read(List<VectorSearchSplit> splits) {
+        if (vectors.length > 1) {
+            throw new IllegalStateException(
+                    "read() supports single vector only; use readBatch() for multiple vectors");
+        }
+        return readBatch(splits).get(0);
+    }
+
+    @Override
+    public List<GlobalIndexResult> readBatch(List<VectorSearchSplit> splits) {
+        int n = vectors.length;
         if (splits.isEmpty()) {
-            return GlobalIndexResult.createEmpty();
+            List<GlobalIndexResult> empty = new ArrayList<>(n);
+            for (int i = 0; i < n; i++) {
+                empty.add(GlobalIndexResult.createEmpty());
+            }
+            return empty;
         }
 
         RoaringNavigableMap64 preFilter = preFilter(splits).orElse(null);
@@ -93,11 +107,11 @@ public GlobalIndexResult read(List<VectorSearchSplit> splits) {
         int parallelism = table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM);
         ExecutorService executor = GlobalIndexReadThreadPool.getExecutorService(parallelism);
 
-        List<CompletableFuture<Optional<ScoredGlobalIndexResult>>> futures =
+        List<CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>> futures =
                 new ArrayList<>(splits.size());
         for (VectorSearchSplit split : splits) {
             futures.add(
-                    eval(
+                    evalBatch(
                             globalIndexer,
                             indexPathFactory,
                             split.rowRangeStart(),
@@ -109,15 +123,25 @@ public GlobalIndexResult read(List<VectorSearchSplit> splits) {
 
         CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
 
-        ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
-        for (CompletableFuture<Optional<ScoredGlobalIndexResult>> f : futures) {
-            Optional<ScoredGlobalIndexResult> next = f.join();
-            if (next.isPresent()) {
-                result = result.or(next.get());
+        ScoredGlobalIndexResult[] merged = new ScoredGlobalIndexResult[n];
+        for (int i = 0; i < n; i++) {
+            merged[i] = ScoredGlobalIndexResult.createEmpty();
+        }
+
+        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());
+                }
             }
         }
 
-        return result.topK(limit);
+        List<GlobalIndexResult> results = new ArrayList<>(n);
+        for (int i = 0; i < n; i++) {
+            results.add(merged[i].topK(limit));
+        }
+        return results;
     }
 
     private Optional<RoaringNavigableMap64> preFilter(List<VectorSearchSplit> splits) {
@@ -139,33 +163,40 @@ private Optional<RoaringNavigableMap64> preFilter(List<VectorSearchSplit> splits
         }
     }
 
-    private CompletableFuture<Optional<ScoredGlobalIndexResult>> eval(
+    private CompletableFuture<List<Optional<ScoredGlobalIndexResult>>> evalBatch(
             GlobalIndexer globalIndexer,
             IndexPathFactory indexPathFactory,
             long rowRangeStart,
             long rowRangeEnd,
             List<IndexFileMeta> vectorIndexFiles,
             @Nullable RoaringNavigableMap64 includeRowIds,
             ExecutorService executor) {
-        List<GlobalIndexIOMeta> indexIOMetaList = new ArrayList<>();
-        for (IndexFileMeta indexFile : vectorIndexFiles) {
-            GlobalIndexMeta meta = checkNotNull(indexFile.globalIndexMeta());
-            indexIOMetaList.add(
-                    new GlobalIndexIOMeta(
-                            indexPathFactory.toPath(indexFile),
-                            indexFile.fileSize(),
-                            meta.indexMeta()));
-        }
+        List<GlobalIndexIOMeta> indexIOMetaList =
+                buildIOMetaList(indexPathFactory, vectorIndexFiles);
         @SuppressWarnings("resource")
         FileIO fileIO = table.fileIO();
         GlobalIndexFileReader indexFileReader = m -> fileIO.newInputStream(m.filePath());
         GlobalIndexReader reader =
                 globalIndexer.createReader(indexFileReader, indexIOMetaList, executor);
         VectorSearch vectorSearch =
-                new VectorSearch(vector, limit, vectorColumn.name())
+                new VectorSearch(vectors, limit, vectorColumn.name())
                         .withIncludeRowIds(includeRowIds);
         return new OffsetGlobalIndexReader(reader, rowRangeStart, rowRangeEnd)
-                .visitVectorSearch(vectorSearch)
+                .visitBatchVectorSearch(vectorSearch)
                 .whenComplete((r, t) -> IOUtils.closeQuietly(reader));
     }
+
+    private List<GlobalIndexIOMeta> buildIOMetaList(
+            IndexPathFactory indexPathFactory, List<IndexFileMeta> vectorIndexFiles) {
+        List<GlobalIndexIOMeta> indexIOMetaList = new ArrayList<>();
+        for (IndexFileMeta indexFile : vectorIndexFiles) {
+            GlobalIndexMeta meta = checkNotNull(indexFile.globalIndexMeta());
+            indexIOMetaList.add(
+                    new GlobalIndexIOMeta(
+                            indexPathFactory.toPath(indexFile),
+                            indexFile.fileSize(),
+                            meta.indexMeta()));
+        }
+        return indexIOMetaList;
+    }
 }
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilder.java b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilder.java
@@ -23,6 +23,7 @@
 import org.apache.paimon.predicate.Predicate;
 
 import java.io.Serializable;
+import java.util.List;
 
 /** Builder to build vector search. */
 public interface VectorSearchBuilder extends Serializable {
@@ -42,6 +43,9 @@ public interface VectorSearchBuilder extends Serializable {
     /** The vector to search. */
     VectorSearchBuilder withVector(float[] vector);
 
+    /** The vectors to batch search. */
+    VectorSearchBuilder withVectors(float[][] vectors);
+
     /** Create vector scan to scan index files. */
     VectorScan newVectorScan();
 
@@ -52,4 +56,9 @@ public interface VectorSearchBuilder extends Serializable {
     default GlobalIndexResult executeLocal() {
         return newVectorRead().read(newVectorScan().scan());
     }
+
+    /** Execute batch vector index search in local. */
+    default List<GlobalIndexResult> executeBatchLocal() {
+        return newVectorRead().readBatch(newVectorScan().scan());
+    }
 }
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java
@@ -26,6 +26,7 @@
 import org.apache.paimon.types.DataField;
 
 import static org.apache.paimon.partition.PartitionPredicate.splitPartitionPredicate;
+import static org.apache.paimon.utils.Preconditions.checkNotNull;
 
 /** Implementation for {@link VectorSearchBuilder}. */
 public class VectorSearchBuilderImpl implements VectorSearchBuilder {
@@ -38,7 +39,7 @@ public class VectorSearchBuilderImpl implements VectorSearchBuilder {
     private Predicate filter;
     private int limit;
     private DataField vectorColumn;
-    private float[] vector;
+    private float[][] vectors;
 
     public VectorSearchBuilderImpl(InnerTable table) {
         this.table = (FileStoreTable) table;
@@ -76,7 +77,13 @@ public VectorSearchBuilder withVectorColumn(String name) {
 
     @Override
     public VectorSearchBuilder withVector(float[] vector) {
-        this.vector = vector;
+        this.vectors = new float[][] {vector};
+        return this;
+    }
+
+    @Override
+    public VectorSearchBuilder withVectors(float[][] vectors) {
+        this.vectors = vectors;
         return this;
     }
 
@@ -87,6 +94,7 @@ public VectorScan newVectorScan() {
 
     @Override
     public VectorRead newVectorRead() {
-        return new VectorReadImpl(table, filter, limit, vectorColumn, vector);
+        checkNotNull(vectors, "vectors must be set via withVector() or withVectors()");
+        return new VectorReadImpl(table, filter, limit, vectorColumn, vectors);
     }
 }
diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
diff --git a/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaVectorGlobalIndexReader.java b/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaVectorGlobalIndexReader.java
diff --git a/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/LuminaVectorGlobalIndexTest.java b/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/LuminaVectorGlobalIndexTest.java
