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 f7d77e77d9 [core] Support batch vector search (#7857)
f7d77e77d9 is described below
commit f7d77e77d92eddcb86bbd9b11f14155d1da9f24e
Author: XiaoHongbo <[email protected]>
AuthorDate: Thu Jun 18 13:30:18 2026 +0800
[core] Support batch vector search (#7857)
---
docs/docs/multimodal-table/global-index/vector.mdx | 15 ++
.../paimon/globalindex/GlobalIndexReader.java | 22 +++
.../globalindex/OffsetGlobalIndexReader.java | 17 ++
.../{VectorSearch.java => BatchVectorSearch.java} | 73 +++----
.../org/apache/paimon/predicate/VectorSearch.java | 12 +-
.../apache/paimon/predicate/VectorSearchUtils.java | 24 ++-
.../apache/paimon/predicate/VectorSearchTest.java | 17 ++
.../java/org/apache/paimon/table/FormatTable.java | 6 +
.../java/org/apache/paimon/table/InnerTable.java | 7 +
.../main/java/org/apache/paimon/table/Table.java | 7 +
.../table/source/BatchVectorSearchBuilder.java | 69 +++++++
...Impl.java => BatchVectorSearchBuilderImpl.java} | 35 ++--
.../org/apache/paimon/table/source/VectorRead.java | 13 ++
.../apache/paimon/table/source/VectorReadImpl.java | 85 +++++---
.../table/source/VectorSearchBuilderImpl.java | 5 +-
.../table/source/VectorSearchBuilderTest.java | 140 +++++++++++++
.../index/LuminaVectorGlobalIndexReader.java | 217 ++++++++++++++++-----
.../lumina/index/LuminaVectorGlobalIndexTest.java | 204 +++++++++++++++++++
paimon-python/pypaimon/table/file_store_table.py | 5 +
.../pypaimon/table/format/format_table.py | 3 +
.../pypaimon/table/iceberg/iceberg_table.py | 3 +
.../pypaimon/table/object/object_table.py | 5 +
.../table/source/batch_vector_search_builder.py | 124 ++++++++++++
.../pypaimon/table/source/vector_search_builder.py | 23 ++-
.../pypaimon/table/source/vector_search_read.py | 31 ++-
.../pypaimon/table/system/system_table.py | 3 +
paimon-python/pypaimon/table/table.py | 7 +
.../pypaimon/tests/vector_search_filter_test.py | 83 ++++++++
.../paimon/spark/read/SparkVectorReadImpl.java | 21 +-
.../spark/read/SparkVectorSearchBuilderImpl.java | 5 +-
.../vector/index/VectorGlobalIndexReader.java | 156 ++++++++++++---
.../paimon/vector/index/VectorGlobalIndexTest.java | 168 ++++++++++++++++
32 files changed, 1413 insertions(+), 192 deletions(-)
diff --git a/docs/docs/multimodal-table/global-index/vector.mdx
b/docs/docs/multimodal-table/global-index/vector.mdx
index f4ac4eedcc..5f6cb01691 100644
--- a/docs/docs/multimodal-table/global-index/vector.mdx
+++ b/docs/docs/multimodal-table/global-index/vector.mdx
@@ -182,6 +182,21 @@ try (RecordReader<InternalRow> reader =
readBuilder.newRead().createReader(plan)
}
```
+Batch results keep the same order as input vectors.
+
+```java
+float[][] queryVectors = {
+ {1.0f, 2.0f, 3.0f},
+ {3.0f, 2.0f, 1.0f}
+};
+List<GlobalIndexResult> batchResults = table.newBatchVectorSearchBuilder()
+ .withVectors(queryVectors)
+ .withLimit(5)
+ .withVectorColumn("embedding")
+ .executeBatchLocal();
+// batchResults.get(i) corresponds to queryVectors[i].
+```
+
For Java, use `Table.newVectorSearchBuilder()` to produce a global index
result, then pass
the result to `TableScan.withGlobalIndexResult`.
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
index b16ce888af..5951d40726 100644
---
a/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java
+++
b/paimon-common/src/main/java/org/apache/paimon/globalindex/GlobalIndexReader.java
@@ -18,12 +18,14 @@
package org.apache.paimon.globalindex;
+import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.FullTextSearch;
import org.apache.paimon.predicate.FunctionVisitor;
import org.apache.paimon.predicate.LeafPredicate;
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 +61,24 @@ public interface GlobalIndexReader
FullTextSearch fullTextSearch) {
throw new UnsupportedOperationException();
}
+
+ /** Batch search; result {@code i} matches vector {@code i}. */
+ default CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>
visitBatchVectorSearch(
+ BatchVectorSearch batchVectorSearch) {
+ List<CompletableFuture<Optional<ScoredGlobalIndexResult>>> futures =
new ArrayList<>();
+ for (int i = 0; i < batchVectorSearch.vectorCount(); i++) {
+ futures.add(visitVectorSearch(batchVectorSearch.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
index e2a03bca76..f740223d5a 100644
---
a/paimon-common/src/main/java/org/apache/paimon/globalindex/OffsetGlobalIndexReader.java
+++
b/paimon-common/src/main/java/org/apache/paimon/globalindex/OffsetGlobalIndexReader.java
@@ -18,11 +18,13 @@
package org.apache.paimon.globalindex;
+import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.FieldRef;
import org.apache.paimon.predicate.FullTextSearch;
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 +147,21 @@ public class OffsetGlobalIndexReader implements
GlobalIndexReader {
.thenApply(opt -> opt.map(r -> r.offset(offset)));
}
+ @Override
+ public CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>
visitBatchVectorSearch(
+ BatchVectorSearch batchVectorSearch) {
+ return
wrapped.visitBatchVectorSearch(batchVectorSearch.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/BatchVectorSearch.java
similarity index 57%
copy from
paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
copy to
paimon-common/src/main/java/org/apache/paimon/predicate/BatchVectorSearch.java
index 5f046c6307..f1617f6722 100644
--- a/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
+++
b/paimon-common/src/main/java/org/apache/paimon/predicate/BatchVectorSearch.java
@@ -18,7 +18,6 @@
package org.apache.paimon.predicate;
-import org.apache.paimon.utils.Range;
import org.apache.paimon.utils.RoaringNavigableMap64;
import javax.annotation.Nullable;
@@ -28,30 +27,31 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
-/**
- * VectorSearch to perform vector similarity search.
- *
- * <p>This is an internal pushdown representation. Use {@code
Table.newVectorSearchBuilder()} to
- * configure vector search from Java.
- */
-public class VectorSearch implements Serializable {
+/** Batch vector similarity search. */
+public class BatchVectorSearch implements Serializable {
private static final long serialVersionUID = 1L;
- private final float[] vector;
+ private final float[][] vectors;
private final String fieldName;
private final int limit;
private final Map<String, String> options;
@Nullable private RoaringNavigableMap64 includeRowIds;
- public VectorSearch(float[] vector, int limit, String fieldName) {
- this(vector, limit, fieldName, Collections.emptyMap());
+ public BatchVectorSearch(float[][] vectors, int limit, String fieldName) {
+ this(vectors, limit, fieldName, Collections.emptyMap());
}
- public VectorSearch(float[] vector, int limit, String fieldName,
Map<String, String> options) {
- if (vector == null) {
- throw new IllegalArgumentException("Search cannot be null");
+ public BatchVectorSearch(
+ float[][] vectors, int limit, String fieldName, Map<String,
String> options) {
+ if (vectors == null || vectors.length == 0) {
+ throw new IllegalArgumentException("Search vectors cannot be null
or empty");
+ }
+ for (float[] vector : vectors) {
+ if (vector == null) {
+ throw new IllegalArgumentException("Search vector element
cannot be null");
+ }
}
if (limit <= 0) {
throw new IllegalArgumentException("Limit must be positive, got: "
+ limit);
@@ -59,7 +59,7 @@ public class VectorSearch implements Serializable {
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;
this.options =
@@ -68,8 +68,25 @@ public class VectorSearch implements Serializable {
: Collections.unmodifiableMap(new HashMap<>(options));
}
- public float[] vector() {
- return vector;
+ /** Query vectors in input order. */
+ public float[][] vectors() {
+ return vectors;
+ }
+
+ public int vectorCount() {
+ return vectors.length;
+ }
+
+ public VectorSearch forIndex(int i) {
+ VectorSearch vectorSearch = new VectorSearch(vectors[i], limit,
fieldName, options);
+ if (includeRowIds != null) {
+ vectorSearch.withIncludeRowIds(includeRowIds);
+ }
+ return vectorSearch;
+ }
+
+ public Map<String, String> options() {
+ return options;
}
public int limit() {
@@ -80,30 +97,19 @@ public class VectorSearch implements Serializable {
return fieldName;
}
- public Map<String, String> options() {
- return options == null ? Collections.emptyMap() : options;
- }
-
public RoaringNavigableMap64 includeRowIds() {
return includeRowIds;
}
- public VectorSearch withIncludeRowIds(RoaringNavigableMap64 includeRowIds)
{
+ public BatchVectorSearch withIncludeRowIds(RoaringNavigableMap64
includeRowIds) {
this.includeRowIds = includeRowIds;
return this;
}
- public VectorSearch offsetRange(long from, long to) {
+ public BatchVectorSearch offsetRange(long from, long to) {
if (includeRowIds != null) {
- RoaringNavigableMap64 range = new RoaringNavigableMap64();
- range.addRange(new Range(from, to));
- RoaringNavigableMap64 and64 = RoaringNavigableMap64.and(range,
includeRowIds);
- final RoaringNavigableMap64 roaringNavigableMap64Offset = new
RoaringNavigableMap64();
- for (long rowId : and64) {
- roaringNavigableMap64Offset.add(rowId - from);
- }
- VectorSearch target = new VectorSearch(vector, limit, fieldName,
options());
- target.withIncludeRowIds(roaringNavigableMap64Offset);
+ BatchVectorSearch target = new BatchVectorSearch(vectors, limit,
fieldName, options);
+
target.withIncludeRowIds(VectorSearchUtils.offsetRowIds(includeRowIds, from,
to));
return target;
}
return this;
@@ -111,6 +117,7 @@ public class VectorSearch implements Serializable {
@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-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
index 5f046c6307..1114904bc6 100644
--- a/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
+++ b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
@@ -18,7 +18,6 @@
package org.apache.paimon.predicate;
-import org.apache.paimon.utils.Range;
import org.apache.paimon.utils.RoaringNavigableMap64;
import javax.annotation.Nullable;
@@ -51,7 +50,7 @@ public class VectorSearch implements Serializable {
public VectorSearch(float[] vector, int limit, String fieldName,
Map<String, String> options) {
if (vector == null) {
- throw new IllegalArgumentException("Search cannot be null");
+ throw new IllegalArgumentException("Search vector cannot be null");
}
if (limit <= 0) {
throw new IllegalArgumentException("Limit must be positive, got: "
+ limit);
@@ -95,15 +94,8 @@ public class VectorSearch implements Serializable {
public VectorSearch offsetRange(long from, long to) {
if (includeRowIds != null) {
- RoaringNavigableMap64 range = new RoaringNavigableMap64();
- range.addRange(new Range(from, to));
- RoaringNavigableMap64 and64 = RoaringNavigableMap64.and(range,
includeRowIds);
- final RoaringNavigableMap64 roaringNavigableMap64Offset = new
RoaringNavigableMap64();
- for (long rowId : and64) {
- roaringNavigableMap64Offset.add(rowId - from);
- }
VectorSearch target = new VectorSearch(vector, limit, fieldName,
options());
- target.withIncludeRowIds(roaringNavigableMap64Offset);
+
target.withIncludeRowIds(VectorSearchUtils.offsetRowIds(includeRowIds, from,
to));
return target;
}
return this;
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorRead.java
b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearchUtils.java
similarity index 53%
copy from
paimon-core/src/main/java/org/apache/paimon/table/source/VectorRead.java
copy to
paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearchUtils.java
index 74e17e2845..7d664160b7 100644
--- a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorRead.java
+++
b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearchUtils.java
@@ -16,18 +16,24 @@
* limitations under the License.
*/
-package org.apache.paimon.table.source;
+package org.apache.paimon.predicate;
-import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.utils.Range;
+import org.apache.paimon.utils.RoaringNavigableMap64;
-import java.util.List;
+/** Utilities for vector search. */
+class VectorSearchUtils {
-/** Vector read to read index files. */
-public interface VectorRead {
-
- default GlobalIndexResult read(VectorScan.Plan plan) {
- return read(plan.splits());
+ static RoaringNavigableMap64 offsetRowIds(RoaringNavigableMap64 rowIds,
long from, long to) {
+ RoaringNavigableMap64 range = new RoaringNavigableMap64();
+ range.addRange(new Range(from, to));
+ RoaringNavigableMap64 filtered = RoaringNavigableMap64.and(range,
rowIds);
+ RoaringNavigableMap64 offsetRowIds = new RoaringNavigableMap64();
+ for (long rowId : filtered) {
+ offsetRowIds.add(rowId - from);
+ }
+ return offsetRowIds;
}
- GlobalIndexResult read(List<VectorSearchSplit> splits);
+ private VectorSearchUtils() {}
}
diff --git
a/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java
b/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java
index 0284ff3d37..a1e91d9c46 100644
---
a/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java
+++
b/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java
@@ -73,4 +73,21 @@ public class VectorSearchTest {
.containsEntry("hnsw.ef_search", "64")
.hasSize(2);
}
+
+ @Test
+ public void testBatchVectorSearchOrder() {
+ float[][] vectors = new float[][] {new float[] {1.0f, 0.0f}, new
float[] {0.0f, 1.0f}};
+
+ RoaringNavigableMap64 includeRowIds = new RoaringNavigableMap64();
+ includeRowIds.addRange(new Range(100L, 200L));
+
+ BatchVectorSearch batchSearch =
+ new BatchVectorSearch(vectors, 1,
"test").withIncludeRowIds(includeRowIds);
+ batchSearch = batchSearch.offsetRange(60, 150);
+
+ assertThat(batchSearch.forIndex(0).vector()).isSameAs(vectors[0]);
+ assertThat(batchSearch.forIndex(1).vector()).isSameAs(vectors[1]);
+ assertThat(batchSearch.forIndex(0).includeRowIds().toRangeList())
+ .containsExactly(new Range(40L, 90L));
+ }
}
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/FormatTable.java
b/paimon-core/src/main/java/org/apache/paimon/table/FormatTable.java
index 86441c5dde..9a832176ca 100644
--- a/paimon-core/src/main/java/org/apache/paimon/table/FormatTable.java
+++ b/paimon-core/src/main/java/org/apache/paimon/table/FormatTable.java
@@ -31,6 +31,7 @@ import org.apache.paimon.table.format.FormatBatchWriteBuilder;
import org.apache.paimon.table.format.FormatReadBuilder;
import org.apache.paimon.table.sink.BatchWriteBuilder;
import org.apache.paimon.table.sink.StreamWriteBuilder;
+import org.apache.paimon.table.source.BatchVectorSearchBuilder;
import org.apache.paimon.table.source.FullTextSearchBuilder;
import org.apache.paimon.table.source.HybridSearchBuilder;
import org.apache.paimon.table.source.ReadBuilder;
@@ -286,6 +287,11 @@ public interface FormatTable extends Table {
throw new UnsupportedOperationException("FormatTable does not
support hybrid search.");
}
+ @Override
+ public BatchVectorSearchBuilder newBatchVectorSearchBuilder() {
+ throw new UnsupportedOperationException("FormatTable does not
support vector search.");
+ }
+
@Override
public FullTextSearchBuilder newFullTextSearchBuilder() {
throw new UnsupportedOperationException(
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/InnerTable.java
b/paimon-core/src/main/java/org/apache/paimon/table/InnerTable.java
index 494aae1fe8..2dfc477561 100644
--- a/paimon-core/src/main/java/org/apache/paimon/table/InnerTable.java
+++ b/paimon-core/src/main/java/org/apache/paimon/table/InnerTable.java
@@ -25,6 +25,8 @@ import org.apache.paimon.table.sink.InnerTableWrite;
import org.apache.paimon.table.sink.StreamWriteBuilder;
import org.apache.paimon.table.sink.StreamWriteBuilderImpl;
import org.apache.paimon.table.sink.WriteSelector;
+import org.apache.paimon.table.source.BatchVectorSearchBuilder;
+import org.apache.paimon.table.source.BatchVectorSearchBuilderImpl;
import org.apache.paimon.table.source.FullTextSearchBuilder;
import org.apache.paimon.table.source.FullTextSearchBuilderImpl;
import org.apache.paimon.table.source.HybridSearchBuilder;
@@ -69,6 +71,11 @@ public interface InnerTable extends Table {
return new HybridSearchBuilderImpl(this);
}
+ @Override
+ default BatchVectorSearchBuilder newBatchVectorSearchBuilder() {
+ return new BatchVectorSearchBuilderImpl(this);
+ }
+
@Override
default FullTextSearchBuilder newFullTextSearchBuilder() {
return new FullTextSearchBuilderImpl(this);
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/Table.java
b/paimon-core/src/main/java/org/apache/paimon/table/Table.java
index 3a7b1730b6..b4fe2bc237 100644
--- a/paimon-core/src/main/java/org/apache/paimon/table/Table.java
+++ b/paimon-core/src/main/java/org/apache/paimon/table/Table.java
@@ -28,6 +28,7 @@ import org.apache.paimon.manifest.ManifestFileMeta;
import org.apache.paimon.stats.Statistics;
import org.apache.paimon.table.sink.BatchWriteBuilder;
import org.apache.paimon.table.sink.StreamWriteBuilder;
+import org.apache.paimon.table.source.BatchVectorSearchBuilder;
import org.apache.paimon.table.source.FullTextSearchBuilder;
import org.apache.paimon.table.source.HybridSearchBuilder;
import org.apache.paimon.table.source.ReadBuilder;
@@ -233,6 +234,12 @@ public interface Table extends Serializable {
getClass().getName() + " does not support hybrid search.");
}
+ /** Returns a new batch vector search builder. */
+ default BatchVectorSearchBuilder newBatchVectorSearchBuilder() {
+ throw new UnsupportedOperationException(
+ getClass().getName() + " does not support batch vector
search.");
+ }
+
/** Returns a new full-text search builder. */
FullTextSearchBuilder newFullTextSearchBuilder();
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilder.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilder.java
new file mode 100644
index 0000000000..a4b5d955b2
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilder.java
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.table.source;
+
+import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.partition.PartitionPredicate;
+import org.apache.paimon.predicate.Predicate;
+
+import java.io.Serializable;
+import java.util.List;
+import java.util.Map;
+
+/** Builder to build batch vector search over multiple query vectors. */
+public interface BatchVectorSearchBuilder extends Serializable {
+
+ /** Push partition filters. */
+ BatchVectorSearchBuilder withPartitionFilter(PartitionPredicate
partitionPredicate);
+
+ /** Push pre-filter for vector search. */
+ BatchVectorSearchBuilder withFilter(Predicate predicate);
+
+ /** The top k results to return per query vector. */
+ BatchVectorSearchBuilder withLimit(int limit);
+
+ /** The vector column to search. */
+ BatchVectorSearchBuilder withVectorColumn(String name);
+
+ /** The query vectors; result {@code i} corresponds to {@code vectors[i]}.
*/
+ BatchVectorSearchBuilder withVectors(float[][] vectors);
+
+ /** Option for vector indexes. */
+ default BatchVectorSearchBuilder withOption(String key, String value) {
+ throw new UnsupportedOperationException(
+ getClass().getName() + " does not support vector options.");
+ }
+
+ /** Options for vector indexes. */
+ default BatchVectorSearchBuilder withOptions(Map<String, String> options) {
+ throw new UnsupportedOperationException(
+ getClass().getName() + " does not support vector options.");
+ }
+
+ /** Create vector scan to scan index files. */
+ VectorScan newVectorScan();
+
+ /** Create vector read to read index files. */
+ VectorRead newVectorRead();
+
+ /** Execute batch vector search locally; result {@code i} corresponds to
{@code vectors[i]}. */
+ 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/BatchVectorSearchBuilderImpl.java
similarity index 67%
copy from
paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java
copy to
paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
index d4686b4416..7b78b6369b 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
@@ -29,9 +29,11 @@ import java.util.HashMap;
import java.util.Map;
import static
org.apache.paimon.partition.PartitionPredicate.splitPartitionPredicate;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+import static org.apache.paimon.utils.Preconditions.checkNotNull;
-/** Implementation for {@link VectorSearchBuilder}. */
-public class VectorSearchBuilderImpl implements VectorSearchBuilder {
+/** Implementation for {@link BatchVectorSearchBuilder}. */
+public class BatchVectorSearchBuilderImpl implements BatchVectorSearchBuilder {
private static final long serialVersionUID = 1L;
@@ -41,21 +43,21 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
protected Predicate filter;
protected int limit;
protected DataField vectorColumn;
- protected float[] vector;
+ protected float[][] vectors;
protected Map<String, String> options = new HashMap<>();
- public VectorSearchBuilderImpl(InnerTable table) {
+ public BatchVectorSearchBuilderImpl(InnerTable table) {
this.table = (FileStoreTable) table;
}
@Override
- public VectorSearchBuilder withPartitionFilter(PartitionPredicate
partitionFilter) {
+ public BatchVectorSearchBuilder withPartitionFilter(PartitionPredicate
partitionFilter) {
this.partitionFilter = partitionFilter;
return this;
}
@Override
- public VectorSearchBuilder withFilter(Predicate predicate) {
+ public BatchVectorSearchBuilder withFilter(Predicate predicate) {
if (this.filter == null) {
this.filter = predicate;
} else {
@@ -67,25 +69,25 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
}
@Override
- public VectorSearchBuilder withLimit(int limit) {
+ public BatchVectorSearchBuilder withLimit(int limit) {
this.limit = limit;
return this;
}
@Override
- public VectorSearchBuilder withVectorColumn(String name) {
+ public BatchVectorSearchBuilder withVectorColumn(String name) {
this.vectorColumn = table.rowType().getField(name);
return this;
}
@Override
- public VectorSearchBuilder withVector(float[] vector) {
- this.vector = vector;
+ public BatchVectorSearchBuilder withVectors(float[][] vectors) {
+ this.vectors = vectors;
return this;
}
@Override
- public VectorSearchBuilder withOptions(Map<String, String> options) {
+ public BatchVectorSearchBuilder withOptions(Map<String, String> options) {
if (options != null) {
this.options.putAll(options);
}
@@ -93,7 +95,7 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
}
@Override
- public VectorSearchBuilder withOption(String key, String value) {
+ public BatchVectorSearchBuilder withOption(String key, String value) {
this.options.put(key, value);
return this;
}
@@ -105,6 +107,13 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
@Override
public VectorRead newVectorRead() {
- return new VectorReadImpl(table, filter, limit, vectorColumn, vector,
options);
+ checkArgument(limit > 0, "Limit must be positive, set via
withLimit()");
+ checkNotNull(vectorColumn, "Vector column must be set via
withVectorColumn()");
+ checkArgument(
+ vectors != null && vectors.length > 0, "vectors must be set
via withVectors()");
+ for (float[] vector : vectors) {
+ checkNotNull(vector, "Search vector element cannot be null");
+ }
+ return new VectorReadImpl(table, filter, limit, vectorColumn, vectors,
options);
}
}
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
index 74e17e2845..54b696b691 100644
--- 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 @@ package org.apache.paimon.table.source;
import org.apache.paimon.globalindex.GlobalIndexResult;
+import java.util.ArrayList;
import java.util.List;
/** Vector read to read index files. */
@@ -30,4 +31,16 @@ public interface VectorRead {
}
GlobalIndexResult read(List<VectorSearchSplit> splits);
+
+ /** Read batch results; result {@code i} corresponds to input vector
{@code i}. */
+ default List<GlobalIndexResult> readBatch(VectorScan.Plan plan) {
+ return readBatch(plan.splits());
+ }
+
+ /** Read batch results; result {@code i} corresponds to input vector
{@code i}. */
+ 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
index 1b3601619c..14e3bb643d 100644
---
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
@@ -32,8 +32,8 @@ import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
import org.apache.paimon.index.GlobalIndexMeta;
import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.index.IndexPathFactory;
+import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.Predicate;
-import org.apache.paimon.predicate.VectorSearch;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.types.DataField;
import org.apache.paimon.utils.IOUtils;
@@ -57,6 +57,7 @@ import java.util.concurrent.ExecutorService;
import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM;
import static org.apache.paimon.utils.Preconditions.checkNotNull;
+import static org.apache.paimon.utils.Preconditions.checkState;
/** Implementation for {@link VectorRead}. */
public class VectorReadImpl implements VectorRead, Serializable {
@@ -67,7 +68,7 @@ public class VectorReadImpl implements VectorRead,
Serializable {
private final Predicate filter;
protected final int limit;
protected final DataField vectorColumn;
- protected final float[] vector;
+ protected final float[][] vectors;
protected final Map<String, String> options;
public VectorReadImpl(
@@ -75,8 +76,8 @@ public class VectorReadImpl implements VectorRead,
Serializable {
Predicate filter,
int limit,
DataField vectorColumn,
- float[] vector) {
- this(table, filter, limit, vectorColumn, vector,
Collections.emptyMap());
+ float[][] vectors) {
+ this(table, filter, limit, vectorColumn, vectors,
Collections.emptyMap());
}
public VectorReadImpl(
@@ -84,13 +85,13 @@ public class VectorReadImpl implements VectorRead,
Serializable {
Predicate filter,
int limit,
DataField vectorColumn,
- float[] vector,
+ float[][] vectors,
Map<String, String> options) {
this.table = table;
this.filter = filter;
this.limit = limit;
this.vectorColumn = vectorColumn;
- this.vector = vector;
+ this.vectors = vectors;
this.options =
options == null
? Collections.emptyMap()
@@ -99,8 +100,21 @@ public class VectorReadImpl implements VectorRead,
Serializable {
@Override
public GlobalIndexResult read(List<VectorSearchSplit> splits) {
+ checkState(
+ vectors.length == 1,
+ "read() is 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);
@@ -126,11 +140,11 @@ public class VectorReadImpl implements VectorRead,
Serializable {
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(),
@@ -142,15 +156,25 @@ public class VectorReadImpl implements VectorRead,
Serializable {
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;
}
protected Optional<RoaringNavigableMap64>
preFilter(List<VectorSearchSplit> splits) {
@@ -172,7 +196,7 @@ public class VectorReadImpl implements VectorRead,
Serializable {
}
}
- protected CompletableFuture<Optional<ScoredGlobalIndexResult>> eval(
+ protected CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>
evalBatch(
GlobalIndexer globalIndexer,
IndexPathFactory indexPathFactory,
long rowRangeStart,
@@ -180,6 +204,23 @@ public class VectorReadImpl implements VectorRead,
Serializable {
List<IndexFileMeta> vectorIndexFiles,
@Nullable RoaringNavigableMap64 includeRowIds,
ExecutorService executor) {
+ 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);
+ BatchVectorSearch batchVectorSearch =
+ new BatchVectorSearch(vectors, limit, vectorColumn.name(),
options)
+ .withIncludeRowIds(includeRowIds);
+ return new OffsetGlobalIndexReader(reader, rowRangeStart, rowRangeEnd)
+ .visitBatchVectorSearch(batchVectorSearch)
+ .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());
@@ -189,16 +230,6 @@ public class VectorReadImpl implements VectorRead,
Serializable {
indexFile.fileSize(),
meta.indexMeta()));
}
- @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(), options)
- .withIncludeRowIds(includeRowIds);
- return new OffsetGlobalIndexReader(reader, rowRangeStart, rowRangeEnd)
- .visitVectorSearch(vectorSearch)
- .whenComplete((r, t) -> IOUtils.closeQuietly(reader));
+ return indexIOMetaList;
}
}
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
index d4686b4416..b7652719ad 100644
---
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
@@ -29,6 +29,7 @@ import java.util.HashMap;
import java.util.Map;
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 {
@@ -105,6 +106,8 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
@Override
public VectorRead newVectorRead() {
- return new VectorReadImpl(table, filter, limit, vectorColumn, vector,
options);
+ checkNotNull(vector, "vector must be set via withVector()");
+ return new VectorReadImpl(
+ table, filter, limit, vectorColumn, new float[][] {vector},
options);
}
}
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
index fb5ca42764..ed9705c786 100644
---
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
@@ -572,6 +572,136 @@ public class VectorSearchBuilderTest extends
TableTestBase {
}
}
+ @Test
+ public void testBatchVectorSearch() throws Exception {
+ createTableDefault();
+ FileStoreTable table = getTableDefault();
+
+ float[][] vectors = {
+ {1.0f, 0.0f},
+ {0.95f, 0.1f},
+ {0.1f, 0.95f},
+ {0.98f, 0.05f},
+ {0.0f, 1.0f},
+ {0.05f, 0.98f}
+ };
+
+ writeVectors(table, vectors);
+ buildAndCommitIndex(table, vectors);
+
+ float[][] queryVectors = {
+ {1.0f, 0.0f},
+ {0.0f, 1.0f},
+ {0.7f, 0.7f}
+ };
+
+ List<GlobalIndexResult> results =
+ table.newBatchVectorSearchBuilder()
+ .withVectors(queryVectors)
+ .withLimit(2)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .executeBatchLocal();
+
+ assertThat(results).hasSize(3);
+
+ // Query 0 near (1,0): should find rows 0 (1,0) and 3 (0.98,0.05)
+ assertThat(results.get(0).results().isEmpty()).isFalse();
+ ReadBuilder rb0 = table.newReadBuilder();
+ List<Integer> ids0 = new ArrayList<>();
+ try (RecordReader<InternalRow> reader =
+ rb0.newRead()
+
.createReader(rb0.newScan().withGlobalIndexResult(results.get(0)).plan())) {
+ reader.forEachRemaining(row -> ids0.add(row.getInt(0)));
+ }
+ assertThat(ids0).contains(0);
+
+ // Query 1 near (0,1): should find rows 4 (0,1) and 5 (0.05,0.98)
+ assertThat(results.get(1).results().isEmpty()).isFalse();
+ ReadBuilder rb1 = table.newReadBuilder();
+ List<Integer> ids1 = new ArrayList<>();
+ try (RecordReader<InternalRow> reader =
+ rb1.newRead()
+
.createReader(rb1.newScan().withGlobalIndexResult(results.get(1)).plan())) {
+ reader.forEachRemaining(row -> ids1.add(row.getInt(0)));
+ }
+ assertThat(ids1).contains(4);
+ }
+
+ @Test
+ public void testBatchVectorSearchWithMultipleIndexFiles() throws Exception
{
+ createTableDefault();
+ FileStoreTable table = getTableDefault();
+
+ float[][] allVectors = {
+ {1.0f, 0.0f},
+ {0.95f, 0.1f},
+ {0.1f, 0.95f},
+ {0.98f, 0.05f},
+ {0.0f, 1.0f},
+ {0.05f, 0.98f}
+ };
+
+ writeVectors(table, allVectors);
+ buildAndCommitMultipleIndexFiles(table, allVectors);
+
+ List<GlobalIndexResult> results =
+ table.newBatchVectorSearchBuilder()
+ .withVectors(new float[][] {{0.85f, 0.15f}, {0.0f,
1.0f}})
+ .withLimit(3)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .executeBatchLocal();
+
+ assertThat(results).hasSize(2);
+
+ List<Integer> ids0 = readIds(table, results.get(0));
+ assertThat(ids0.size()).isLessThanOrEqualTo(3);
+ assertThat(ids0).contains(0, 3);
+
+ List<Integer> ids1 = readIds(table, results.get(1));
+ assertThat(ids1.size()).isLessThanOrEqualTo(3);
+ assertThat(ids1).contains(4, 5);
+ }
+
+ @Test
+ public void testBatchSingleVector() throws Exception {
+ createTableDefault();
+ FileStoreTable table = getTableDefault();
+
+ float[][] vectors = {
+ {1.0f, 0.0f},
+ {0.95f, 0.1f},
+ {0.0f, 1.0f},
+ {0.98f, 0.05f}
+ };
+
+ writeVectors(table, vectors);
+ buildAndCommitIndex(table, vectors);
+
+ float[] queryVector = {0.9f, 0.1f};
+
+ GlobalIndexResult singleResult =
+ table.newVectorSearchBuilder()
+ .withVector(queryVector)
+ .withLimit(3)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .executeLocal();
+
+ List<GlobalIndexResult> batchResults =
+ table.newBatchVectorSearchBuilder()
+ .withVectors(new float[][] {queryVector})
+ .withLimit(3)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .executeBatchLocal();
+
+ assertThat(batchResults).hasSize(1);
+ assertThat(batchResults.get(0).results().getIntCardinality())
+ .isEqualTo(singleResult.results().getIntCardinality());
+
+ for (long rowId : singleResult.results()) {
+ assertThat(batchResults.get(0).results().contains(rowId)).isTrue();
+ }
+ }
+
// ====================== Helper methods ======================
private void writeVectors(FileStoreTable table, float[][] vectors) throws
Exception {
@@ -601,6 +731,16 @@ public class VectorSearchBuilderTest extends TableTestBase
{
}
}
+ private List<Integer> readIds(FileStoreTable table, GlobalIndexResult
result) throws Exception {
+ ReadBuilder readBuilder = table.newReadBuilder();
+ TableScan.Plan plan =
readBuilder.newScan().withGlobalIndexResult(result).plan();
+ List<Integer> ids = new ArrayList<>();
+ try (RecordReader<InternalRow> reader =
readBuilder.newRead().createReader(plan)) {
+ reader.forEachRemaining(row -> ids.add(row.getInt(0)));
+ }
+ return ids;
+ }
+
private void buildAndCommitIndex(FileStoreTable table, float[][] vectors)
throws Exception {
buildAndCommitIndex(table, VECTOR_FIELD_NAME, vectors);
}
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
index 20fffa9d11..39aabb7f1a 100644
---
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
@@ -24,6 +24,7 @@ import org.apache.paimon.globalindex.GlobalIndexReader;
import org.apache.paimon.globalindex.GlobalIndexResult;
import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
+import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.FieldRef;
import org.apache.paimon.predicate.VectorSearch;
import org.apache.paimon.types.ArrayType;
@@ -36,6 +37,7 @@ import org.apache.paimon.utils.RoaringNavigableMap64;
import org.aliyun.lumina.LuminaFileInput;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
@@ -107,65 +109,118 @@ public class LuminaVectorGlobalIndexReader implements
GlobalIndexReader {
executor);
}
+ @Override
+ public CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>
visitBatchVectorSearch(
+ BatchVectorSearch batchVectorSearch) {
+ return CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ ensureLoaded();
+ return searchBatch(batchVectorSearch);
+ } catch (IOException e) {
+ throw new RuntimeException(
+ String.format(
+ "Failed to batch search Lumina vector
index with fieldName=%s, limit=%d, vectorCount=%d",
+ batchVectorSearch.fieldName(),
+ batchVectorSearch.limit(),
+ batchVectorSearch.vectorCount()),
+ e);
+ }
+ },
+ executor);
+ }
+
+ private List<Optional<ScoredGlobalIndexResult>>
searchBatch(BatchVectorSearch batchVectorSearch)
+ throws IOException {
+ int n = batchVectorSearch.vectorCount();
+ if (n == 1) {
+ List<Optional<ScoredGlobalIndexResult>> results = new
ArrayList<>(1);
+
results.add(Optional.ofNullable(search(batchVectorSearch.forIndex(0))));
+ return results;
+ }
+
+ float[][] vectors = batchVectorSearch.vectors();
+ int dim = indexMeta.dim();
+ for (float[] v : vectors) {
+ validateSearchVector(v);
+ }
+
+ int limit = batchVectorSearch.limit();
+ int effectiveK = (int) Math.min(limit, index.size());
+ if (effectiveK <= 0) {
+ return emptyResults(n);
+ }
+
+ float[] queryVectors = new float[n * dim];
+ for (int i = 0; i < n; i++) {
+ System.arraycopy(vectors[i], 0, queryVectors, i * dim, dim);
+ }
+
+ RoaringNavigableMap64 includeRowIds =
batchVectorSearch.includeRowIds();
+ long[] scopedIds = toScopedIds(includeRowIds);
+ if (scopedIds != null && scopedIds.length == 0) {
+ return emptyResults(n);
+ }
+ if (scopedIds != null) {
+ effectiveK = Math.min(effectiveK, scopedIds.length);
+ }
+
+ float[] distances = new float[n * effectiveK];
+ long[] labels = new long[n * effectiveK];
+ Map<String, String> searchOptions =
+ buildSearchOptions(scopedIds != null, effectiveK,
batchVectorSearch.options());
+
+ if (scopedIds != null) {
+ index.searchWithFilter(
+ queryVectors, n, effectiveK, distances, labels, scopedIds,
searchOptions);
+ } else {
+ index.search(queryVectors, n, effectiveK, distances, labels,
searchOptions);
+ }
+
+ LuminaVectorMetric indexMetric = indexMeta.metric();
+ List<Optional<ScoredGlobalIndexResult>> results = new ArrayList<>(n);
+ for (int i = 0; i < n; i++) {
+ results.add(
+ buildScoredResult(distances, labels, i * effectiveK,
effectiveK, indexMetric));
+ }
+ return results;
+ }
+
private ScoredGlobalIndexResult search(VectorSearch vectorSearch) throws
IOException {
validateSearchVector(vectorSearch.vector());
float[] queryVector = vectorSearch.vector().clone();
- int limit = vectorSearch.limit();
- LuminaVectorMetric indexMetric = indexMeta.metric();
- int effectiveK = (int) Math.min(limit, index.size());
+ int effectiveK = (int) Math.min(vectorSearch.limit(), index.size());
if (effectiveK <= 0) {
return null;
}
RoaringNavigableMap64 includeRowIds = vectorSearch.includeRowIds();
- float[] distances;
- long[] labels;
-
- if (includeRowIds != null) {
- long cardinality = includeRowIds.getLongCardinality();
- if (cardinality > Integer.MAX_VALUE) {
- throw new IllegalArgumentException(
- "includeRowIds cardinality ("
- + cardinality
- + ") exceeds Integer.MAX_VALUE");
- }
- long[] scopedIds = new long[(int) cardinality];
- Iterator<Long> iter = includeRowIds.iterator();
- for (int i = 0; i < scopedIds.length; i++) {
- scopedIds[i] = iter.next();
- }
- if (scopedIds.length == 0) {
- return null;
- }
+ long[] scopedIds = toScopedIds(includeRowIds);
+ if (scopedIds != null && scopedIds.length == 0) {
+ return null;
+ }
+ if (scopedIds != null) {
effectiveK = Math.min(effectiveK, scopedIds.length);
- distances = new float[effectiveK];
- labels = new long[effectiveK];
- Map<String, String> mergedOptions =
- mergeOptions(
- this.options.toLuminaOptions(),
- indexMeta.options(),
- vectorSearch.options());
- mergedOptions.put("search.thread_safe_filter", "true");
- ensureSearchListSize(mergedOptions, effectiveK);
+ }
+
+ float[] distances = new float[effectiveK];
+ long[] labels = new long[effectiveK];
+ Map<String, String> searchOptions =
+ buildSearchOptions(scopedIds != null, effectiveK,
vectorSearch.options());
+
+ if (scopedIds != null) {
index.searchWithFilter(
- queryVector, 1, effectiveK, distances, labels, scopedIds,
mergedOptions);
+ queryVector, 1, effectiveK, distances, labels, scopedIds,
searchOptions);
} else {
- distances = new float[effectiveK];
- labels = new long[effectiveK];
- Map<String, String> mergedOptions =
- mergeOptions(
- this.options.toLuminaOptions(),
- indexMeta.options(),
- vectorSearch.options());
- ensureSearchListSize(mergedOptions, effectiveK);
- index.search(queryVector, 1, effectiveK, distances, labels,
mergedOptions);
+ index.search(queryVector, 1, effectiveK, distances, labels,
searchOptions);
}
+ LuminaVectorMetric indexMetric = indexMeta.metric();
// Min-heap: smallest score at head, so we can evict the weakest
candidate efficiently.
PriorityQueue<ScoredRow> topK =
new PriorityQueue<>(effectiveK + 1,
Comparator.comparingDouble(s -> s.score));
- collectResults(distances, labels, effectiveK, effectiveK, topK,
indexMetric);
+ collectResults(distances, labels, 0, effectiveK, effectiveK, topK,
indexMetric);
RoaringNavigableMap64 roaringBitmap64 = new RoaringNavigableMap64();
HashMap<Long, Float> id2scores = new HashMap<>(topK.size());
@@ -176,36 +231,96 @@ public class LuminaVectorGlobalIndexReader implements
GlobalIndexReader {
return new LuminaScoredGlobalIndexResult(roaringBitmap64, id2scores);
}
+ private long[] toScopedIds(RoaringNavigableMap64 includeRowIds) {
+ if (includeRowIds == null) {
+ return null;
+ }
+ long cardinality = includeRowIds.getLongCardinality();
+ if (cardinality > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException(
+ "includeRowIds cardinality (" + cardinality + ") exceeds
Integer.MAX_VALUE");
+ }
+ long[] scopedIds = new long[(int) cardinality];
+ Iterator<Long> iter = includeRowIds.iterator();
+ for (int i = 0; i < scopedIds.length; i++) {
+ scopedIds[i] = iter.next();
+ }
+ return scopedIds;
+ }
+
+ private Map<String, String> buildSearchOptions(
+ boolean withFilter, int effectiveK, Map<String, String>
queryOptions) {
+ Map<String, String> searchOptions = options.toLuminaOptions();
+ searchOptions.putAll(indexMeta.options());
+ searchOptions.putAll(queryOptions);
+ if (withFilter) {
+ searchOptions.put("search.thread_safe_filter", "true");
+ }
+ ensureSearchListSize(searchOptions, effectiveK);
+ return searchOptions;
+ }
+
static Map<String, String> mergeOptions(
Map<String, String> baseOptions,
Map<String, String> indexOptions,
Map<String, String> queryOptions) {
- Map<String, String> options = new HashMap<>(baseOptions);
- options.putAll(indexOptions);
- options.putAll(queryOptions);
- return options;
+ Map<String, String> merged = new HashMap<>(baseOptions);
+ merged.putAll(indexOptions);
+ merged.putAll(queryOptions);
+ return merged;
+ }
+
+ private static Optional<ScoredGlobalIndexResult> buildScoredResult(
+ float[] distances,
+ long[] labels,
+ int offset,
+ int effectiveK,
+ LuminaVectorMetric indexMetric) {
+ PriorityQueue<ScoredRow> topK =
+ new PriorityQueue<>(effectiveK + 1,
Comparator.comparingDouble(s -> s.score));
+ collectResults(distances, labels, offset, effectiveK, effectiveK,
topK, indexMetric);
+ if (topK.isEmpty()) {
+ return Optional.empty();
+ }
+ RoaringNavigableMap64 bitmap = new RoaringNavigableMap64();
+ HashMap<Long, Float> id2scores = new HashMap<>(topK.size());
+ for (ScoredRow row : topK) {
+ bitmap.add(row.rowId);
+ id2scores.put(row.rowId, row.score);
+ }
+ return Optional.of(new LuminaScoredGlobalIndexResult(bitmap,
id2scores));
+ }
+
+ private static List<Optional<ScoredGlobalIndexResult>> emptyResults(int n)
{
+ List<Optional<ScoredGlobalIndexResult>> results = new ArrayList<>(n);
+ for (int i = 0; i < n; i++) {
+ results.add(Optional.empty());
+ }
+ return results;
}
- private static void ensureSearchListSize(Map<String, String> options, int
topK) {
- if (!options.containsKey("diskann.search.list_size")) {
+ private static void ensureSearchListSize(Map<String, String>
searchOptions, int topK) {
+ if (!searchOptions.containsKey("diskann.search.list_size")) {
int listSize = Math.max((int) (topK * 1.5), MIN_SEARCH_LIST_SIZE);
- options.put("diskann.search.list_size", String.valueOf(listSize));
+ searchOptions.put("diskann.search.list_size",
String.valueOf(listSize));
}
}
private static void collectResults(
float[] distances,
long[] labels,
+ int offset,
int count,
int limit,
PriorityQueue<ScoredRow> topK,
LuminaVectorMetric metric) {
for (int i = 0; i < count; i++) {
- long rowId = labels[i];
+ int index = offset + i;
+ long rowId = labels[index];
if (rowId < 0) {
continue;
}
- float score = convertDistanceToScore(distances[i], metric);
+ float score = convertDistanceToScore(distances[index], metric);
if (topK.size() < limit) {
topK.offer(new ScoredRow(rowId, score));
} else if (score > topK.peek().score) {
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
index b9ab7e1219..86bad585b4 100644
---
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
@@ -25,9 +25,11 @@ import org.apache.paimon.fs.PositionOutputStream;
import org.apache.paimon.fs.local.LocalFileIO;
import org.apache.paimon.globalindex.GlobalIndexIOMeta;
import org.apache.paimon.globalindex.ResultEntry;
+import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
import org.apache.paimon.options.Options;
+import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.VectorSearch;
import org.apache.paimon.types.ArrayType;
import org.apache.paimon.types.DataType;
@@ -48,6 +50,8 @@ import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
+import java.util.Map;
+import java.util.Optional;
import java.util.Random;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
@@ -749,6 +753,206 @@ public class LuminaVectorGlobalIndexTest {
.hasMessageContaining("-Infinity");
}
+ @Test
+ public void testBatchVectorSearch() throws IOException {
+ int dimension = 2;
+ Options options = createDefaultOptions(dimension);
+
+ float[][] vectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f},
+ new float[] {0.95f, 0.1f},
+ new float[] {0.1f, 0.95f},
+ new float[] {0.98f, 0.05f},
+ new float[] {0.0f, 1.0f},
+ new float[] {0.05f, 0.98f}
+ };
+
+ GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
+ LuminaVectorIndexOptions indexOptions = new
LuminaVectorIndexOptions(options);
+ LuminaVectorGlobalIndexWriter writer =
+ new LuminaVectorGlobalIndexWriter(fileWriter, vectorType,
indexOptions);
+ Arrays.stream(vectors).forEach(writer::write);
+
+ List<ResultEntry> results = writer.finish();
+ List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
+
+ GlobalIndexFileReader fileReader = createFileReader(indexPath);
+ try (LuminaVectorGlobalIndexReader reader =
+ new LuminaVectorGlobalIndexReader(
+ fileReader, metas, vectorType, indexOptions,
executor)) {
+ float[][] queryVectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f},
+ new float[] {0.0f, 1.0f},
+ new float[] {0.7f, 0.7f}
+ };
+ BatchVectorSearch batchSearch = new
BatchVectorSearch(queryVectors, 2, fieldName);
+ List<Optional<ScoredGlobalIndexResult>> batchResults =
+ reader.visitBatchVectorSearch(batchSearch).join();
+
+ assertThat(batchResults).hasSize(3);
+
+ assertThat(batchResults.get(0)).isPresent();
+
assertThat(batchResults.get(0).get().results().contains(0L)).isTrue();
+
assertThat(batchResults.get(0).get().results().contains(3L)).isTrue();
+
+ assertThat(batchResults.get(1)).isPresent();
+
assertThat(batchResults.get(1).get().results().contains(4L)).isTrue();
+
assertThat(batchResults.get(1).get().results().contains(5L)).isTrue();
+
+ assertThat(batchResults.get(2)).isPresent();
+
assertThat(batchResults.get(2).get().results().getLongCardinality()).isEqualTo(2);
+ }
+ }
+
+ @Test
+ public void testBatchVectorSearchWithFilter() throws IOException {
+ int dimension = 2;
+ Options options = createDefaultOptions(dimension);
+
+ float[][] vectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f},
+ new float[] {0.95f, 0.1f},
+ new float[] {0.1f, 0.95f},
+ new float[] {0.0f, 1.0f},
+ };
+
+ GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
+ LuminaVectorIndexOptions indexOptions = new
LuminaVectorIndexOptions(options);
+ LuminaVectorGlobalIndexWriter writer =
+ new LuminaVectorGlobalIndexWriter(fileWriter, vectorType,
indexOptions);
+ Arrays.stream(vectors).forEach(writer::write);
+
+ List<ResultEntry> results = writer.finish();
+ List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
+
+ GlobalIndexFileReader fileReader = createFileReader(indexPath);
+ try (LuminaVectorGlobalIndexReader reader =
+ new LuminaVectorGlobalIndexReader(
+ fileReader, metas, vectorType, indexOptions,
executor)) {
+ float[][] queryVectors =
+ new float[][] {new float[] {1.0f, 0.0f}, new float[]
{0.0f, 1.0f}};
+
+ RoaringNavigableMap64 filter = new RoaringNavigableMap64();
+ filter.add(1L);
+ filter.add(2L);
+
+ BatchVectorSearch batchSearch =
+ new BatchVectorSearch(queryVectors, 2,
fieldName).withIncludeRowIds(filter);
+ List<Optional<ScoredGlobalIndexResult>> batchResults =
+ reader.visitBatchVectorSearch(batchSearch).join();
+
+ assertThat(batchResults).hasSize(2);
+
+ assertThat(batchResults.get(0)).isPresent();
+
assertThat(batchResults.get(0).get().results().contains(1L)).isTrue();
+
+ assertThat(batchResults.get(1)).isPresent();
+
assertThat(batchResults.get(1).get().results().contains(2L)).isTrue();
+ }
+ }
+
+ @Test
+ public void testBatchConsistentWithSingle() throws IOException {
+ int dimension = 32;
+ int numVectors = 100;
+ Options options = createDefaultOptions(dimension);
+
+ GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
+ LuminaVectorIndexOptions indexOptions = new
LuminaVectorIndexOptions(options);
+ LuminaVectorGlobalIndexWriter writer =
+ new LuminaVectorGlobalIndexWriter(fileWriter, vectorType,
indexOptions);
+
+ List<float[]> testVectors = generateRandomVectors(numVectors,
dimension);
+ testVectors.forEach(writer::write);
+
+ List<ResultEntry> results = writer.finish();
+ List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
+
+ GlobalIndexFileReader fileReader = createFileReader(indexPath);
+ try (LuminaVectorGlobalIndexReader reader =
+ new LuminaVectorGlobalIndexReader(
+ fileReader, metas, vectorType, indexOptions,
executor)) {
+ float[][] queryVectors =
+ new float[][] {testVectors.get(10), testVectors.get(50),
testVectors.get(90)};
+ int limit = 5;
+
+ BatchVectorSearch batchSearch = new
BatchVectorSearch(queryVectors, limit, fieldName);
+ List<Optional<ScoredGlobalIndexResult>> batchResults =
+ reader.visitBatchVectorSearch(batchSearch).join();
+
+ for (int i = 0; i < queryVectors.length; i++) {
+ VectorSearch singleSearch = new VectorSearch(queryVectors[i],
limit, fieldName);
+ Optional<ScoredGlobalIndexResult> singleResult =
+ reader.visitVectorSearch(singleSearch).join();
+
+
assertThat(batchResults.get(i).isPresent()).isEqualTo(singleResult.isPresent());
+ if (singleResult.isPresent()) {
+
assertThat(batchResults.get(i).get().results().getIntCardinality())
+
.isEqualTo(singleResult.get().results().getIntCardinality());
+ for (long rowId : singleResult.get().results()) {
+
assertThat(batchResults.get(i).get().results().contains(rowId)).isTrue();
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testBatchVectorSearchAppliesQueryOptions() throws IOException {
+ int dimension = 32;
+ int numVectors = 100;
+ Options options = createDefaultOptions(dimension);
+
+ GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
+ LuminaVectorIndexOptions indexOptions = new
LuminaVectorIndexOptions(options);
+ LuminaVectorGlobalIndexWriter writer =
+ new LuminaVectorGlobalIndexWriter(fileWriter, vectorType,
indexOptions);
+
+ List<float[]> testVectors = generateRandomVectors(numVectors,
dimension);
+ testVectors.forEach(writer::write);
+
+ List<ResultEntry> results = writer.finish();
+ List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
+
+ GlobalIndexFileReader fileReader = createFileReader(indexPath);
+ try (LuminaVectorGlobalIndexReader reader =
+ new LuminaVectorGlobalIndexReader(
+ fileReader, metas, vectorType, indexOptions,
executor)) {
+ float[][] queryVectors =
+ new float[][] {testVectors.get(10), testVectors.get(50),
testVectors.get(90)};
+ int limit = 20;
+
+ // A query option must reach the native batch call as it does for
single search;
+ // if the batch path dropped it, this batch-vs-single equality
would break.
+ Map<String, String> queryOptions =
+ Collections.singletonMap("diskann.search.list_size", "1");
+
+ BatchVectorSearch batchSearch =
+ new BatchVectorSearch(queryVectors, limit, fieldName,
queryOptions);
+ List<Optional<ScoredGlobalIndexResult>> batchResults =
+ reader.visitBatchVectorSearch(batchSearch).join();
+
+ for (int i = 0; i < queryVectors.length; i++) {
+ VectorSearch singleSearch =
+ new VectorSearch(queryVectors[i], limit, fieldName,
queryOptions);
+ Optional<ScoredGlobalIndexResult> singleResult =
+ reader.visitVectorSearch(singleSearch).join();
+
+
assertThat(batchResults.get(i).isPresent()).isEqualTo(singleResult.isPresent());
+ if (singleResult.isPresent()) {
+
assertThat(batchResults.get(i).get().results().getIntCardinality())
+
.isEqualTo(singleResult.get().results().getIntCardinality());
+ for (long rowId : singleResult.get().results()) {
+
assertThat(batchResults.get(i).get().results().contains(rowId)).isTrue();
+ }
+ }
+ }
+ }
+ }
+
private Options createDefaultOptions(int dimension) {
Options options = new Options();
options.setInteger(LuminaVectorIndexOptions.DIMENSION.key(),
dimension);
diff --git a/paimon-python/pypaimon/table/file_store_table.py
b/paimon-python/pypaimon/table/file_store_table.py
index 04f972d9e8..8805806883 100644
--- a/paimon-python/pypaimon/table/file_store_table.py
+++ b/paimon-python/pypaimon/table/file_store_table.py
@@ -434,6 +434,11 @@ class FileStoreTable(Table):
HybridSearchBuilderImpl
return HybridSearchBuilderImpl(self)
+ def new_batch_vector_search_builder(self) -> 'BatchVectorSearchBuilder':
+ from pypaimon.table.source.batch_vector_search_builder import \
+ BatchVectorSearchBuilderImpl
+ return BatchVectorSearchBuilderImpl(self)
+
def create_row_key_extractor(self) -> RowKeyExtractor:
bucket_mode = self.bucket_mode()
if bucket_mode == BucketMode.HASH_FIXED:
diff --git a/paimon-python/pypaimon/table/format/format_table.py
b/paimon-python/pypaimon/table/format/format_table.py
index 795835f47a..1a1218153d 100644
--- a/paimon-python/pypaimon/table/format/format_table.py
+++ b/paimon-python/pypaimon/table/format/format_table.py
@@ -112,3 +112,6 @@ class FormatTable(Table):
def new_hybrid_search_builder(self):
raise NotImplementedError("Format table does not support hybrid
search.")
+
+ def new_batch_vector_search_builder(self):
+ raise NotImplementedError("Format table does not support vector
search.")
diff --git a/paimon-python/pypaimon/table/iceberg/iceberg_table.py
b/paimon-python/pypaimon/table/iceberg/iceberg_table.py
index 51f099ee0f..8bf17e9544 100644
--- a/paimon-python/pypaimon/table/iceberg/iceberg_table.py
+++ b/paimon-python/pypaimon/table/iceberg/iceberg_table.py
@@ -117,3 +117,6 @@ class IcebergTable(Table):
def new_hybrid_search_builder(self):
raise NotImplementedError("IcebergTable does not support hybrid
search.")
+
+ def new_batch_vector_search_builder(self):
+ raise NotImplementedError("IcebergTable does not support vector
search.")
diff --git a/paimon-python/pypaimon/table/object/object_table.py
b/paimon-python/pypaimon/table/object/object_table.py
index d12e00c616..f875b4d3f6 100644
--- a/paimon-python/pypaimon/table/object/object_table.py
+++ b/paimon-python/pypaimon/table/object/object_table.py
@@ -117,3 +117,8 @@ class ObjectTable(Table):
raise NotImplementedError(
"ObjectTable is read-only and does not support hybrid search."
)
+
+ def new_batch_vector_search_builder(self):
+ raise NotImplementedError(
+ "ObjectTable is read-only and does not support vector search."
+ )
diff --git a/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
new file mode 100644
index 0000000000..03477ef73b
--- /dev/null
+++ b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
@@ -0,0 +1,124 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Builder to build batch vector search over multiple query vectors."""
+
+from abc import ABC, abstractmethod
+
+from pypaimon.table.source.vector_search_builder import (
+ AbstractVectorSearchBuilderImpl,
+)
+from pypaimon.table.source.vector_search_read import BatchVectorSearchReadImpl
+
+
+class BatchVectorSearchBuilder(ABC):
+ """Builder to build batch vector search; result ``i`` matches vector
``i``."""
+
+ @abstractmethod
+ def with_limit(self, limit):
+ # type: (int) -> BatchVectorSearchBuilder
+ """The top k results to return per query vector."""
+ pass
+
+ @abstractmethod
+ def with_vector_column(self, name):
+ # type: (str) -> BatchVectorSearchBuilder
+ """The vector column to search."""
+ pass
+
+ @abstractmethod
+ def with_query_vectors(self, vectors):
+ # type: (list) -> BatchVectorSearchBuilder
+ """The query vectors (list of list of floats); result i matches
vectors[i]."""
+ pass
+
+ def with_option(self, key, value):
+ # type: (str, str) -> BatchVectorSearchBuilder
+ """Option for vector indexes."""
+ raise NotImplementedError(
+ "%s does not support vector options."
+ % self.__class__.__name__)
+
+ def with_options(self, options):
+ # type: (dict) -> BatchVectorSearchBuilder
+ """Options for vector indexes."""
+ raise NotImplementedError(
+ "%s does not support vector options."
+ % self.__class__.__name__)
+
+ @abstractmethod
+ def with_filter(self, predicate):
+ # type: (Predicate) -> BatchVectorSearchBuilder
+ """Scalar predicate used to pre-filter rows before vector search."""
+ pass
+
+ @abstractmethod
+ def with_partition_filter(self, partition_filter):
+ # type: (Predicate) -> BatchVectorSearchBuilder
+ """Partition predicate used to prune index manifest entries."""
+ pass
+
+ @abstractmethod
+ def new_vector_search_scan(self):
+ # type: () -> VectorSearchScan
+ """Create vector search scan to scan index files."""
+ pass
+
+ @abstractmethod
+ def new_batch_vector_search_read(self):
+ # type: () -> BatchVectorSearchReadImpl
+ """Create batch vector search read to read index files."""
+ pass
+
+ def execute_batch_local(self):
+ # type: () -> List[GlobalIndexResult]
+ """Execute batch vector search locally; result i matches query vector
i."""
+ return self.new_batch_vector_search_read().read_batch(
+ self.new_vector_search_scan().scan().splits()
+ )
+
+
+class BatchVectorSearchBuilderImpl(AbstractVectorSearchBuilderImpl,
+ BatchVectorSearchBuilder):
+ """Implementation for BatchVectorSearchBuilder."""
+
+ def __init__(self, table):
+ super().__init__(table)
+ self._query_vectors = None
+
+ def with_query_vectors(self, vectors):
+ # type: (list) -> BatchVectorSearchBuilder
+ self._query_vectors = vectors
+ return self
+
+ def new_batch_vector_search_read(self):
+ # type: () -> BatchVectorSearchReadImpl
+ if self._limit <= 0:
+ raise ValueError("Limit must be positive, set via with_limit()")
+ if self._vector_column is None:
+ raise ValueError("Vector column must be set via
with_vector_column()")
+ if not self._query_vectors:
+ raise ValueError(
+ "Query vectors must be set via with_query_vectors()")
+ return BatchVectorSearchReadImpl(
+ self._table,
+ self._limit,
+ self._vector_column,
+ self._query_vectors,
+ filter_=self._filter,
+ options=self._options,
+ )
diff --git a/paimon-python/pypaimon/table/source/vector_search_builder.py
b/paimon-python/pypaimon/table/source/vector_search_builder.py
index 7cb9eb0104..f985471929 100644
--- a/paimon-python/pypaimon/table/source/vector_search_builder.py
+++ b/paimon-python/pypaimon/table/source/vector_search_builder.py
@@ -91,14 +91,13 @@ class VectorSearchBuilder(ABC):
)
-class VectorSearchBuilderImpl(VectorSearchBuilder):
- """Implementation for VectorSearchBuilder."""
+class AbstractVectorSearchBuilderImpl:
+ """Shared state and filter/partition handling for the vector search
builders."""
def __init__(self, table):
self._table = table
self._limit = 0
self._vector_column = None
- self._query_vector = None
self._filter = None
self._partition_filter = None
self._options = {}
@@ -116,11 +115,6 @@ class VectorSearchBuilderImpl(VectorSearchBuilder):
self._vector_column = field_dict[name]
return self
- def with_query_vector(self, vector):
- # type: (list) -> VectorSearchBuilder
- self._query_vector = vector
- return self
-
def with_option(self, key, value):
# type: (str, str) -> VectorSearchBuilder
self._options[key] = value
@@ -224,6 +218,19 @@ class VectorSearchBuilderImpl(VectorSearchBuilder):
partition_filter=self._partition_filter,
)
+
+class VectorSearchBuilderImpl(AbstractVectorSearchBuilderImpl,
VectorSearchBuilder):
+ """Implementation for VectorSearchBuilder."""
+
+ def __init__(self, table):
+ super().__init__(table)
+ self._query_vector = None
+
+ def with_query_vector(self, vector):
+ # type: (list) -> VectorSearchBuilder
+ self._query_vector = vector
+ return self
+
def new_vector_search_read(self):
# type: () -> VectorSearchRead
if self._limit <= 0:
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index 2abac3bed8..bd2c6f9287 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -58,11 +58,15 @@ class VectorSearchReadImpl(VectorSearchRead):
return GlobalIndexResult.create_empty()
pre_filter = self._pre_filter(splits)
+ return self._search_one(self._query_vector, splits, pre_filter)
+ def _search_one(self, query_vector, splits, pre_filter):
+ # type: (list, list, Optional[RoaringBitmap64]) -> GlobalIndexResult
+ """Search one query vector across all splits and merge per-split
results."""
futures = [
self._eval(
split.row_range_start, split.row_range_end,
- split.vector_index_files, pre_filter
+ split.vector_index_files, query_vector, pre_filter
)
for split in splits
]
@@ -112,7 +116,7 @@ class VectorSearchReadImpl(VectorSearchRead):
scanner.close()
def _eval(self, row_range_start, row_range_end, vector_index_files,
- include_row_ids):
+ query_vector, include_row_ids):
from pypaimon.globalindex.global_index_reader import _completed_future
if not vector_index_files:
@@ -136,7 +140,7 @@ class VectorSearchReadImpl(VectorSearchRead):
options = self._table.table_schema.options
vector_search = VectorSearch(
- vector=self._query_vector,
+ vector=query_vector,
limit=self._limit,
field_name=self._vector_column.name,
options=self._options,
@@ -154,6 +158,27 @@ class VectorSearchReadImpl(VectorSearchRead):
return future
+class BatchVectorSearchReadImpl(VectorSearchReadImpl):
+ """Batch vector search read; result ``i`` corresponds to query vector
``i``."""
+
+ def __init__(self, table, limit, vector_column, query_vectors,
+ filter_=None, options=None):
+ super().__init__(table, limit, vector_column, None,
+ filter_=filter_, options=options)
+ self._query_vectors = list(query_vectors)
+
+ def read_batch(self, splits):
+ # type: (List[VectorSearchSplit]) -> List[GlobalIndexResult]
+ n = len(self._query_vectors)
+ if not splits:
+ return [GlobalIndexResult.create_empty() for _ in range(n)]
+
+ pre_filter = self._pre_filter(splits)
+ # result i corresponds to query_vectors[i], in input order.
+ return [self._search_one(vector, splits, pre_filter)
+ for vector in self._query_vectors]
+
+
def _create_vector_reader(index_type, file_io, index_path, index_io_meta_list,
options=None):
"""Create a global index reader for vector search."""
from pypaimon.globalindex.lumina.lumina_vector_global_index_reader import (
diff --git a/paimon-python/pypaimon/table/system/system_table.py
b/paimon-python/pypaimon/table/system/system_table.py
index 8432470076..7ace7f23f3 100644
--- a/paimon-python/pypaimon/table/system/system_table.py
+++ b/paimon-python/pypaimon/table/system/system_table.py
@@ -105,6 +105,9 @@ class SystemTable(Table):
def new_hybrid_search_builder(self):
raise NotImplementedError(_READ_ONLY_MESSAGE)
+ def new_batch_vector_search_builder(self):
+ raise NotImplementedError(_READ_ONLY_MESSAGE)
+
class SystemReadBuilder:
"""ReadBuilder-shaped facade exposing the system table's data.
diff --git a/paimon-python/pypaimon/table/table.py
b/paimon-python/pypaimon/table/table.py
index a89eab6e74..b650e7dbb1 100644
--- a/paimon-python/pypaimon/table/table.py
+++ b/paimon-python/pypaimon/table/table.py
@@ -19,6 +19,7 @@ from abc import ABC, abstractmethod
from pypaimon.read.read_builder import ReadBuilder
from pypaimon.read.stream_read_builder import StreamReadBuilder
+from pypaimon.table.source.batch_vector_search_builder import
BatchVectorSearchBuilder
from pypaimon.table.source.full_text_search_builder import
FullTextSearchBuilder
from pypaimon.table.source.hybrid_search_builder import HybridSearchBuilder
from pypaimon.table.source.vector_search_builder import VectorSearchBuilder
@@ -55,3 +56,9 @@ class Table(ABC):
@abstractmethod
def new_hybrid_search_builder(self) -> HybridSearchBuilder:
"""Returns a new hybrid search builder."""
+
+ def new_batch_vector_search_builder(self) -> BatchVectorSearchBuilder:
+ """Returns a new batch vector search builder."""
+ raise NotImplementedError(
+ "%s does not support batch vector search."
+ % self.__class__.__name__)
diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py
b/paimon-python/pypaimon/tests/vector_search_filter_test.py
index 90cb5bb0e4..5aa886220f 100644
--- a/paimon-python/pypaimon/tests/vector_search_filter_test.py
+++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py
@@ -1595,5 +1595,88 @@ class FullTextSearchManySplitsTest(unittest.TestCase):
mock.patch.stopall()
+class BatchVectorSearchTest(unittest.TestCase):
+ """Batch vector search returns one result per query vector, in input
order."""
+
+ def test_batch_returns_per_query_results_in_order(self):
+ from pypaimon.globalindex.vector_search_result import (
+ DictBasedScoredIndexResult,
+ )
+ from pypaimon.table.source.batch_vector_search_builder import (
+ BatchVectorSearchBuilderImpl,
+ )
+
+ embedding_field = _field(1, "embedding", "FLOAT")
+ entry = _entry(None, field_id=1, index_type="lumina-vector-ann",
+ file_name="vec.index",
+ row_range_start=0, row_range_end=99)
+ table = _StubTable(fields=[embedding_field], entries=[entry])
+ _patch_snapshot(self, [entry])
+
+ # The fake reader routes each query vector to a distinct row id derived
+ # from the vector itself, so result i must reflect query_vectors[i].
+ def _fake_create(index_type, file_io, index_path,
+ index_io_meta_list, options=None):
+ class _FakeReader:
+ def visit_vector_search(self_inner, vs):
+ row_id = int(vs.vector[0])
+ return _completed_future(
+ DictBasedScoredIndexResult({row_id: 1.0}))
+
+ def close(self_inner):
+ pass
+
+ def __enter__(self_inner):
+ return self_inner
+
+ def __exit__(self_inner, *a):
+ return False
+ return _FakeReader()
+
+ query_vectors = [[10.0], [20.0], [30.0]]
+ with mock.patch(
+
"pypaimon.table.source.vector_search_read._create_vector_reader",
+ side_effect=_fake_create):
+ results = (
+ BatchVectorSearchBuilderImpl(table)
+ .with_vector_column("embedding")
+ .with_query_vectors(query_vectors)
+ .with_limit(5)
+ .execute_batch_local()
+ )
+
+ self.assertEqual(len(results), len(query_vectors))
+ for i, query_vector in enumerate(query_vectors):
+ expected_row = int(query_vector[0])
+ self.assertTrue(results[i].results().contains(expected_row))
+ self.assertEqual(results[i].results().cardinality(), 1)
+ # Different query vectors yield different results.
+ self.assertNotEqual(
+ list(results[0].results()), list(results[1].results()))
+
+ def test_batch_empty_splits_returns_empty_per_query(self):
+ from pypaimon.table.source.batch_vector_search_builder import (
+ BatchVectorSearchBuilderImpl,
+ )
+
+ embedding_field = _field(1, "embedding", "FLOAT")
+ table = _StubTable(fields=[embedding_field], entries=[])
+ _patch_snapshot(self, [])
+
+ results = (
+ BatchVectorSearchBuilderImpl(table)
+ .with_vector_column("embedding")
+ .with_query_vectors([[1.0], [2.0]])
+ .with_limit(5)
+ .execute_batch_local()
+ )
+ self.assertEqual(len(results), 2)
+ for result in results:
+ self.assertEqual(result.results().cardinality(), 0)
+
+ def tearDown(self):
+ mock.patch.stopall()
+
+
if __name__ == "__main__":
unittest.main()
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
index 525ed334a8..94e45316ce 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
@@ -45,6 +45,7 @@ import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM;
+import static org.apache.paimon.utils.Preconditions.checkState;
/**
* Spark-aware {@link VectorReadImpl} that distributes grouped vector index
evaluation across the
@@ -59,8 +60,8 @@ public class SparkVectorReadImpl extends VectorReadImpl {
Predicate filter,
int limit,
DataField vectorColumn,
- float[] vector) {
- super(table, filter, limit, vectorColumn, vector);
+ float[][] vectors) {
+ super(table, filter, limit, vectorColumn, vectors);
}
public SparkVectorReadImpl(
@@ -68,13 +69,16 @@ public class SparkVectorReadImpl extends VectorReadImpl {
Predicate filter,
int limit,
DataField vectorColumn,
- float[] vector,
+ float[][] vectors,
Map<String, String> options) {
- super(table, filter, limit, vectorColumn, vector, options);
+ super(table, filter, limit, vectorColumn, vectors, options);
}
@Override
public GlobalIndexResult read(List<VectorSearchSplit> splits) {
+ checkState(
+ vectors.length == 1,
+ "read() is single-vector only; use readBatch() for multiple
vectors");
if (splits.isEmpty()) {
return GlobalIndexResult.createEmpty();
}
@@ -113,12 +117,12 @@ public class SparkVectorReadImpl extends VectorReadImpl {
ExecutorService executor =
GlobalIndexReadThreadPool.getExecutorService(
Math.min(parallelism, group.size()));
- List<CompletableFuture<Optional<ScoredGlobalIndexResult>>>
futures =
+
List<CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>> futures =
new ArrayList<>(group.size());
for (byte[] bytes : group) {
VectorSearchSplit split = deserializeSplit(bytes);
futures.add(
- eval(
+ evalBatch(
globalIndexer,
indexPathFactory,
split.rowRangeStart(),
@@ -129,8 +133,9 @@ public class SparkVectorReadImpl extends VectorReadImpl {
}
CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0])).join();
ScoredGlobalIndexResult result =
ScoredGlobalIndexResult.createEmpty();
- for (CompletableFuture<Optional<ScoredGlobalIndexResult>>
f : futures) {
- Optional<ScoredGlobalIndexResult> next = f.join();
+ for
(CompletableFuture<List<Optional<ScoredGlobalIndexResult>>> f : futures) {
+ // Spark carries a single query vector, so the batch
result has one element.
+ Optional<ScoredGlobalIndexResult> next =
f.join().get(0);
if (next.isPresent()) {
result = result.or(next.get());
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
index bd19a9e565..91928b7f84 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
@@ -25,6 +25,8 @@ import org.apache.paimon.table.source.VectorSearchBuilderImpl;
/**
* Spark-aware {@link VectorSearchBuilderImpl} which produces a {@link
SparkVectorReadImpl} so the
* per-split vector index evaluation is dispatched through Spark instead of
the local thread pool.
+ *
+ * <p>Single-vector only; batch search has no Spark-dispatched path yet (TODO).
*/
public class SparkVectorSearchBuilderImpl extends VectorSearchBuilderImpl {
@@ -36,6 +38,7 @@ public class SparkVectorSearchBuilderImpl extends
VectorSearchBuilderImpl {
@Override
public VectorRead newVectorRead() {
- return new SparkVectorReadImpl(table, filter, limit, vectorColumn,
vector, options);
+ return new SparkVectorReadImpl(
+ table, filter, limit, vectorColumn, new float[][] {vector},
options);
}
}
diff --git
a/paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorGlobalIndexReader.java
b/paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorGlobalIndexReader.java
index 54f8532a22..cde8d2a83d 100644
---
a/paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorGlobalIndexReader.java
+++
b/paimon-vector/paimon-vector-index/src/main/java/org/apache/paimon/vector/index/VectorGlobalIndexReader.java
@@ -30,7 +30,9 @@ import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
import org.apache.paimon.index.vector.VectorIndexInput;
import org.apache.paimon.index.vector.VectorIndexMetadata;
import org.apache.paimon.index.vector.VectorIndexReader;
+import org.apache.paimon.index.vector.VectorSearchBatchResult;
import org.apache.paimon.index.vector.VectorSearchResult;
+import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.FieldRef;
import org.apache.paimon.predicate.VectorSearch;
import org.apache.paimon.types.ArrayType;
@@ -107,6 +109,74 @@ public class VectorGlobalIndexReader implements
GlobalIndexReader {
executor);
}
+ @Override
+ public CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>
visitBatchVectorSearch(
+ BatchVectorSearch batchVectorSearch) {
+ return CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ ensureLoaded();
+ return searchBatch(batchVectorSearch);
+ } catch (IOException e) {
+ throw new RuntimeException(
+ String.format(
+ "Failed batch vector index search:
field=%s, limit=%d, vectorCount=%d",
+ batchVectorSearch.fieldName(),
+ batchVectorSearch.limit(),
+ batchVectorSearch.vectorCount()),
+ e);
+ }
+ },
+ executor);
+ }
+
+ private List<Optional<ScoredGlobalIndexResult>>
searchBatch(BatchVectorSearch batchVectorSearch)
+ throws IOException {
+ int n = batchVectorSearch.vectorCount();
+ // Single vector: reuse the scalar path; no batching benefit.
+ if (n == 1) {
+ List<Optional<ScoredGlobalIndexResult>> results = new
ArrayList<>(1);
+
results.add(Optional.ofNullable(search(batchVectorSearch.forIndex(0))));
+ return results;
+ }
+
+ float[][] vectors = batchVectorSearch.vectors();
+ for (float[] vector : vectors) {
+ validateSearchVector(vector);
+ }
+ int dim = nativeMeta.dimension();
+ int nprobe = nprobe(batchVectorSearch.options());
+ int efSearch = efSearch(batchVectorSearch.options());
+ String metric = nativeMeta.metric();
+
+ SearchScope scope =
+ resolveScope(batchVectorSearch.includeRowIds(),
batchVectorSearch.limit());
+ if (scope == null) {
+ return emptyResults(n);
+ }
+
+ // Flatten query vectors into one contiguous array for a single native
call.
+ float[] queries = new float[n * dim];
+ for (int i = 0; i < n; i++) {
+ System.arraycopy(vectors[i], 0, queries, i * dim, dim);
+ }
+
+ VectorSearchBatchResult batchResult =
+ scope.filterBytes != null
+ ? vectorReader.searchBatch(
+ queries, n, scope.effectiveK, nprobe,
efSearch, scope.filterBytes)
+ : vectorReader.searchBatch(queries, n,
scope.effectiveK, nprobe, efSearch);
+
+ // result i corresponds to vectors[i], matching input order.
+ List<Optional<ScoredGlobalIndexResult>> results = new ArrayList<>(n);
+ for (int i = 0; i < n; i++) {
+ results.add(
+ buildScoredResult(
+ batchResult.idsForQuery(i),
batchResult.distancesForQuery(i), metric));
+ }
+ return results;
+ }
+
private ScoredGlobalIndexResult search(VectorSearch vectorSearch) throws
IOException {
validateSearchVector(vectorSearch.vector());
float[] queryVector = vectorSearch.vector().clone();
@@ -115,26 +185,23 @@ public class VectorGlobalIndexReader implements
GlobalIndexReader {
int efSearch = efSearch(vectorSearch.options());
String metric = nativeMeta.metric();
- RoaringNavigableMap64 includeRowIds = vectorSearch.includeRowIds();
- VectorSearchResult result;
-
- if (includeRowIds != null) {
- long cardinality = includeRowIds.getLongCardinality();
- if (cardinality == 0) {
- return null;
- }
- byte[] filterBytes = includeRowIds.serialize();
- int effectiveK = (int) Math.min(limit, cardinality);
- result = vectorReader.search(queryVector, effectiveK, nprobe,
efSearch, filterBytes);
- } else {
- result = vectorReader.search(queryVector, limit, nprobe, efSearch);
+ SearchScope scope = resolveScope(vectorSearch.includeRowIds(), limit);
+ if (scope == null) {
+ return null;
}
+ VectorSearchResult result =
+ scope.filterBytes != null
+ ? vectorReader.search(
+ queryVector, scope.effectiveK, nprobe,
efSearch, scope.filterBytes)
+ : vectorReader.search(queryVector, scope.effectiveK,
nprobe, efSearch);
- long[] ids = result.ids();
- float[] distances = result.distances();
+ return buildScoredResult(result.ids(), result.distances(),
metric).orElse(null);
+ }
+ private static Optional<ScoredGlobalIndexResult> buildScoredResult(
+ long[] ids, float[] distances, String metric) {
if (ids.length == 0) {
- return null;
+ return Optional.empty();
}
RoaringNavigableMap64 resultBitmap = new RoaringNavigableMap64();
@@ -151,21 +218,54 @@ public class VectorGlobalIndexReader implements
GlobalIndexReader {
}
if (resultBitmap.isEmpty()) {
+ return Optional.empty();
+ }
+
+ return Optional.of(
+ ScoredGlobalIndexResult.create(
+ resultBitmap,
+ rowId -> {
+ Float score = id2scores.get(rowId);
+ if (score == null) {
+ throw new IllegalArgumentException(
+ "No score found for rowId: "
+ + rowId
+ + ". Only rowIds present in
results() are valid.");
+ }
+ return score;
+ }));
+ }
+
+ private static List<Optional<ScoredGlobalIndexResult>> emptyResults(int n)
{
+ List<Optional<ScoredGlobalIndexResult>> results = new ArrayList<>(n);
+ for (int i = 0; i < n; i++) {
+ results.add(Optional.empty());
+ }
+ return results;
+ }
+
+ /** Resolves filter bytes and effective top-K; returns null when the
filter selects no rows. */
+ private static SearchScope resolveScope(RoaringNavigableMap64
includeRowIds, int limit)
+ throws IOException {
+ if (includeRowIds == null) {
+ return new SearchScope(null, limit);
+ }
+ long cardinality = includeRowIds.getLongCardinality();
+ if (cardinality == 0) {
return null;
}
+ return new SearchScope(includeRowIds.serialize(), (int)
Math.min(limit, cardinality));
+ }
- return ScoredGlobalIndexResult.create(
- resultBitmap,
- rowId -> {
- Float score = id2scores.get(rowId);
- if (score == null) {
- throw new IllegalArgumentException(
- "No score found for rowId: "
- + rowId
- + ". Only rowIds present in results()
are valid.");
- }
- return score;
- });
+ /** Resolved filter state for a query. {@code filterBytes} is null when no
filter is set. */
+ private static final class SearchScope {
+ private final byte[] filterBytes;
+ private final int effectiveK;
+
+ private SearchScope(byte[] filterBytes, int effectiveK) {
+ this.filterBytes = filterBytes;
+ this.effectiveK = effectiveK;
+ }
}
private static float convertDistanceToScore(float distance, String metric)
{
diff --git
a/paimon-vector/paimon-vector-index/src/test/java/org/apache/paimon/vector/index/VectorGlobalIndexTest.java
b/paimon-vector/paimon-vector-index/src/test/java/org/apache/paimon/vector/index/VectorGlobalIndexTest.java
index e82664f806..bd4fc35884 100644
---
a/paimon-vector/paimon-vector-index/src/test/java/org/apache/paimon/vector/index/VectorGlobalIndexTest.java
+++
b/paimon-vector/paimon-vector-index/src/test/java/org/apache/paimon/vector/index/VectorGlobalIndexTest.java
@@ -29,6 +29,7 @@ import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
import org.apache.paimon.index.vector.NativeLoader;
import org.apache.paimon.options.Options;
+import org.apache.paimon.predicate.BatchVectorSearch;
import org.apache.paimon.predicate.VectorSearch;
import org.apache.paimon.types.ArrayType;
import org.apache.paimon.types.DataType;
@@ -49,6 +50,7 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -371,6 +373,172 @@ public class VectorGlobalIndexTest {
}
}
+ @Test
+ public void testBatchVectorSearch() throws IOException {
+ Assumptions.assumeTrue(isNativeAvailable(), "Vector index native
library not available");
+
+ int dimension = 2;
+ Options options = createDefaultOptions(dimension);
+ options.setInteger("ivf-pq.nlist", 2);
+ options.setInteger("ivf-pq.pq.m", 1);
+
+ float[][] vectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f},
+ new float[] {0.95f, 0.1f},
+ new float[] {0.1f, 0.95f},
+ new float[] {0.98f, 0.05f},
+ new float[] {0.0f, 1.0f},
+ new float[] {0.05f, 0.98f}
+ };
+
+ GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
+ VectorGlobalIndexWriter writer = createIvfPqWriter(fileWriter,
vectorType, options);
+ Arrays.stream(vectors).forEach(writer::write);
+ List<ResultEntry> results = writer.finish();
+ List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
+
+ GlobalIndexFileReader fileReader = createFileReader(indexPath);
+ try (VectorGlobalIndexReader reader =
+ new VectorGlobalIndexReader(fileReader, metas, vectorType,
executor)) {
+ float[][] queryVectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f},
+ new float[] {0.0f, 1.0f},
+ new float[] {0.7f, 0.7f}
+ };
+ BatchVectorSearch batchSearch = new
BatchVectorSearch(queryVectors, 3, fieldName);
+ List<Optional<ScoredGlobalIndexResult>> batchResults =
+ reader.visitBatchVectorSearch(batchSearch).join();
+
+ // result i corresponds to queryVectors[i], in input order.
+ assertThat(batchResults).hasSize(3);
+
+ assertThat(batchResults.get(0)).isPresent();
+
assertThat(batchResults.get(0).get().results().contains(0L)).isTrue();
+
+ assertThat(batchResults.get(1)).isPresent();
+
assertThat(batchResults.get(1).get().results().contains(4L)).isTrue();
+
+ assertThat(batchResults.get(2)).isPresent();
+
assertThat(batchResults.get(2).get().results().getLongCardinality()).isEqualTo(3);
+ }
+ }
+
+ @Test
+ public void testBatchVectorSearchWithFilter() throws IOException {
+ Assumptions.assumeTrue(isNativeAvailable(), "Vector index native
library not available");
+
+ int dimension = 2;
+ Options options = createDefaultOptions(dimension);
+ options.setInteger("ivf-pq.nlist", 2);
+ options.setInteger("ivf-pq.pq.m", 1);
+
+ float[][] vectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f},
+ new float[] {0.95f, 0.1f},
+ new float[] {0.9f, 0.2f},
+ new float[] {-1.0f, 0.0f},
+ new float[] {-0.95f, 0.1f},
+ new float[] {-0.9f, 0.2f}
+ };
+
+ GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
+ VectorGlobalIndexWriter writer = createIvfPqWriter(fileWriter,
vectorType, options);
+ Arrays.stream(vectors).forEach(writer::write);
+ List<ResultEntry> results = writer.finish();
+ List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
+
+ GlobalIndexFileReader fileReader = createFileReader(indexPath);
+ try (VectorGlobalIndexReader reader =
+ new VectorGlobalIndexReader(fileReader, metas, vectorType,
executor)) {
+ float[][] queryVectors =
+ new float[][] {new float[] {1.0f, 0.0f}, new float[]
{-1.0f, 0.0f}};
+
+ // Both queries are scoped to rows {1, 4}.
+ RoaringNavigableMap64 filter = new RoaringNavigableMap64();
+ filter.add(1L);
+ filter.add(4L);
+
+ BatchVectorSearch batchSearch =
+ new BatchVectorSearch(queryVectors, 6,
fieldName).withIncludeRowIds(filter);
+ List<Optional<ScoredGlobalIndexResult>> batchResults =
+ reader.visitBatchVectorSearch(batchSearch).join();
+
+ assertThat(batchResults).hasSize(2);
+
+ assertThat(batchResults.get(0)).isPresent();
+
assertThat(batchResults.get(0).get().results().getLongCardinality()).isEqualTo(2);
+
assertThat(batchResults.get(0).get().results().contains(1L)).isTrue();
+
assertThat(batchResults.get(0).get().results().contains(4L)).isTrue();
+
+ assertThat(batchResults.get(1)).isPresent();
+
assertThat(batchResults.get(1).get().results().getLongCardinality()).isEqualTo(2);
+
assertThat(batchResults.get(1).get().results().contains(1L)).isTrue();
+
assertThat(batchResults.get(1).get().results().contains(4L)).isTrue();
+ }
+ }
+
+ @Test
+ public void testBatchConsistentWithSingle() throws IOException {
+ Assumptions.assumeTrue(isNativeAvailable(), "Vector index native
library not available");
+
+ int dimension = 2;
+ Options options = createDefaultOptions(dimension);
+ options.setInteger("ivf-pq.nlist", 2);
+ options.setInteger("ivf-pq.pq.m", 1);
+
+ float[][] vectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f},
+ new float[] {0.95f, 0.1f},
+ new float[] {0.1f, 0.95f},
+ new float[] {0.98f, 0.05f},
+ new float[] {0.0f, 1.0f},
+ new float[] {0.05f, 0.98f}
+ };
+
+ GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
+ VectorGlobalIndexWriter writer = createIvfPqWriter(fileWriter,
vectorType, options);
+ Arrays.stream(vectors).forEach(writer::write);
+ List<ResultEntry> results = writer.finish();
+ List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
+
+ GlobalIndexFileReader fileReader = createFileReader(indexPath);
+ try (VectorGlobalIndexReader reader =
+ new VectorGlobalIndexReader(fileReader, metas, vectorType,
executor)) {
+ float[][] queryVectors =
+ new float[][] {
+ new float[] {1.0f, 0.0f},
+ new float[] {0.0f, 1.0f},
+ new float[] {0.7f, 0.7f}
+ };
+ int limit = 3;
+
+ // The batch path must return exactly what looping the single path
returns, in order.
+ BatchVectorSearch batchSearch = new
BatchVectorSearch(queryVectors, limit, fieldName);
+ List<Optional<ScoredGlobalIndexResult>> batchResults =
+ reader.visitBatchVectorSearch(batchSearch).join();
+
+ assertThat(batchResults).hasSize(queryVectors.length);
+ for (int i = 0; i < queryVectors.length; i++) {
+ VectorSearch singleSearch = new VectorSearch(queryVectors[i],
limit, fieldName);
+ Optional<ScoredGlobalIndexResult> singleResult =
+ reader.visitVectorSearch(singleSearch).join();
+
+
assertThat(batchResults.get(i).isPresent()).isEqualTo(singleResult.isPresent());
+ if (singleResult.isPresent()) {
+
assertThat(batchResults.get(i).get().results().getLongCardinality())
+
.isEqualTo(singleResult.get().results().getLongCardinality());
+ for (long rowId : singleResult.get().results()) {
+
assertThat(batchResults.get(i).get().results().contains(rowId)).isTrue();
+ }
+ }
+ }
+ }
+ }
+
// =================== Helpers =====================
private VectorGlobalIndexWriter createIvfPqWriter(
