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 06ff91b15d [core][spark][python] Refine vector search read APIs (#8277)
06ff91b15d is described below
commit 06ff91b15dacbec4fd1a139d83a9bdc7323808ec
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Jun 18 15:10:11 2026 +0800
[core][spark][python] Refine vector search read APIs (#8277)
Refine the vector search read APIs so single-vector and batch-vector
reads use separate contracts and implementations. This keeps
`VectorRead` single-vector only, moves batch execution behind
`BatchVectorRead`, and aligns Python with the same base + sibling read
architecture.
---
...VectorReadImpl.java => AbstractVectorRead.java} | 127 +++++-----------
.../{VectorRead.java => BatchVectorRead.java} | 17 +--
.../paimon/table/source/BatchVectorReadImpl.java | 114 ++++++++++++++
.../table/source/BatchVectorSearchBuilder.java | 6 +-
.../table/source/BatchVectorSearchBuilderImpl.java | 4 +-
.../org/apache/paimon/table/source/VectorRead.java | 13 --
.../apache/paimon/table/source/VectorReadImpl.java | 169 ++-------------------
.../table/source/VectorSearchBuilderImpl.java | 3 +-
.../table/source/VectorSearchBuilderTest.java | 8 +-
.../lumina/index/LuminaVectorGlobalIndexTest.java | 8 +-
.../table/source/batch_vector_search_builder.py | 13 +-
.../pypaimon/table/source/vector_search_read.py | 124 ++++++++++-----
.../paimon/spark/read/SparkVectorReadImpl.java | 27 +---
.../spark/read/SparkVectorSearchBuilderImpl.java | 3 +-
.../paimon/vector/index/VectorGlobalIndexTest.java | 12 +-
15 files changed, 292 insertions(+), 356 deletions(-)
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/AbstractVectorRead.java
similarity index 61%
copy from
paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
copy to
paimon-core/src/main/java/org/apache/paimon/table/source/AbstractVectorRead.java
index 14e3bb643d..12bec4fc2f 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/AbstractVectorRead.java
@@ -20,7 +20,6 @@ package org.apache.paimon.table.source;
import org.apache.paimon.fs.FileIO;
import org.apache.paimon.globalindex.GlobalIndexIOMeta;
-import org.apache.paimon.globalindex.GlobalIndexReadThreadPool;
import org.apache.paimon.globalindex.GlobalIndexReader;
import org.apache.paimon.globalindex.GlobalIndexResult;
import org.apache.paimon.globalindex.GlobalIndexScanner;
@@ -34,6 +33,7 @@ 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;
@@ -55,12 +55,10 @@ import java.util.TreeSet;
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.checkNotNull;
-import static org.apache.paimon.utils.Preconditions.checkState;
-/** Implementation for {@link VectorRead}. */
-public class VectorReadImpl implements VectorRead, Serializable {
+/** Base implementation for vector reads. */
+public abstract class AbstractVectorRead implements Serializable {
private static final long serialVersionUID = 1L;
@@ -68,113 +66,37 @@ public class VectorReadImpl implements VectorRead,
Serializable {
private final Predicate filter;
protected final int limit;
protected final DataField vectorColumn;
- protected final float[][] vectors;
protected final Map<String, String> options;
- public VectorReadImpl(
+ protected AbstractVectorRead(
FileStoreTable table,
Predicate filter,
int limit,
DataField vectorColumn,
- float[][] vectors) {
- this(table, filter, limit, vectorColumn, vectors,
Collections.emptyMap());
- }
-
- public VectorReadImpl(
- FileStoreTable table,
- Predicate filter,
- int limit,
- DataField vectorColumn,
- float[][] vectors,
Map<String, String> options) {
this.table = table;
this.filter = filter;
this.limit = limit;
this.vectorColumn = vectorColumn;
- this.vectors = vectors;
this.options =
options == null
? Collections.emptyMap()
: Collections.unmodifiableMap(new HashMap<>(options));
}
- @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()) {
- 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);
-
+ protected GlobalIndexer createGlobalIndexer(List<VectorSearchSplit>
splits) {
IndexFileMeta firstFile = splits.get(0).vectorIndexFiles().get(0);
String indexType = firstFile.indexType();
GlobalIndexMeta firstMeta = checkNotNull(firstFile.globalIndexMeta());
- GlobalIndexer globalIndexer;
if (firstMeta.extraFieldIds() != null) {
- globalIndexer =
- GlobalIndexerFactoryUtils.load(indexType)
- .create(
- firstMeta.getIndexField(table.rowType()),
- firstMeta.getExtraFields(table.rowType()),
- table.coreOptions().toConfiguration());
- } else {
- globalIndexer =
- GlobalIndexerFactoryUtils.load(indexType)
- .create(vectorColumn,
table.coreOptions().toConfiguration());
- }
- IndexPathFactory indexPathFactory =
table.store().pathFactory().globalIndexFileFactory();
-
- int parallelism =
table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM);
- ExecutorService executor =
GlobalIndexReadThreadPool.getExecutorService(parallelism);
-
- List<CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>>
futures =
- new ArrayList<>(splits.size());
- for (VectorSearchSplit split : splits) {
- futures.add(
- evalBatch(
- globalIndexer,
- indexPathFactory,
- split.rowRangeStart(),
- split.rowRangeEnd(),
- split.vectorIndexFiles(),
- preFilter,
- executor));
+ return GlobalIndexerFactoryUtils.load(indexType)
+ .create(
+ firstMeta.getIndexField(table.rowType()),
+ firstMeta.getExtraFields(table.rowType()),
+ table.coreOptions().toConfiguration());
}
-
- CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0])).join();
-
- 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());
- }
- }
- }
-
- List<GlobalIndexResult> results = new ArrayList<>(n);
- for (int i = 0; i < n; i++) {
- results.add(merged[i].topK(limit));
- }
- return results;
+ return GlobalIndexerFactoryUtils.load(indexType)
+ .create(vectorColumn, table.coreOptions().toConfiguration());
}
protected Optional<RoaringNavigableMap64>
preFilter(List<VectorSearchSplit> splits) {
@@ -196,12 +118,37 @@ public class VectorReadImpl implements VectorRead,
Serializable {
}
}
+ protected CompletableFuture<Optional<ScoredGlobalIndexResult>> eval(
+ GlobalIndexer globalIndexer,
+ IndexPathFactory indexPathFactory,
+ long rowRangeStart,
+ long rowRangeEnd,
+ List<IndexFileMeta> vectorIndexFiles,
+ float[] vector,
+ @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);
+ 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));
+ }
+
protected CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>
evalBatch(
GlobalIndexer globalIndexer,
IndexPathFactory indexPathFactory,
long rowRangeStart,
long rowRangeEnd,
List<IndexFileMeta> vectorIndexFiles,
+ float[][] vectors,
@Nullable RoaringNavigableMap64 includeRowIds,
ExecutorService executor) {
List<GlobalIndexIOMeta> indexIOMetaList =
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/BatchVectorRead.java
similarity index 72%
copy from
paimon-core/src/main/java/org/apache/paimon/table/source/VectorRead.java
copy to
paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorRead.java
index 54b696b691..7b2169779d 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/BatchVectorRead.java
@@ -20,17 +20,10 @@ 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. */
-public interface VectorRead {
-
- default GlobalIndexResult read(VectorScan.Plan plan) {
- return read(plan.splits());
- }
-
- GlobalIndexResult read(List<VectorSearchSplit> splits);
+/** Batch vector read to read index files for multiple query vectors. */
+public interface BatchVectorRead {
/** Read batch results; result {@code i} corresponds to input vector
{@code i}. */
default List<GlobalIndexResult> readBatch(VectorScan.Plan plan) {
@@ -38,9 +31,5 @@ public interface VectorRead {
}
/** 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;
- }
+ List<GlobalIndexResult> readBatch(List<VectorSearchSplit> splits);
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorReadImpl.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorReadImpl.java
new file mode 100644
index 0000000000..1da9ad554f
--- /dev/null
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorReadImpl.java
@@ -0,0 +1,114 @@
+/*
+ * 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.GlobalIndexReadThreadPool;
+import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.globalindex.GlobalIndexer;
+import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
+import org.apache.paimon.index.IndexPathFactory;
+import org.apache.paimon.predicate.Predicate;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.types.DataField;
+import org.apache.paimon.utils.RoaringNavigableMap64;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+
+import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM;
+
+/** Implementation for {@link BatchVectorRead}. */
+public class BatchVectorReadImpl extends AbstractVectorRead implements
BatchVectorRead {
+
+ private static final long serialVersionUID = 1L;
+
+ protected final float[][] vectors;
+
+ public BatchVectorReadImpl(
+ FileStoreTable table,
+ Predicate filter,
+ int limit,
+ DataField vectorColumn,
+ float[][] vectors,
+ Map<String, String> options) {
+ super(table, filter, limit, vectorColumn, options);
+ this.vectors = vectors;
+ }
+
+ @Override
+ public List<GlobalIndexResult> readBatch(List<VectorSearchSplit> splits) {
+ int n = vectors.length;
+ if (splits.isEmpty()) {
+ 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);
+
+ GlobalIndexer globalIndexer = createGlobalIndexer(splits);
+ IndexPathFactory indexPathFactory =
table.store().pathFactory().globalIndexFileFactory();
+
+ int parallelism =
table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM);
+ ExecutorService executor =
GlobalIndexReadThreadPool.getExecutorService(parallelism);
+
+ List<CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>>
futures =
+ new ArrayList<>(splits.size());
+ for (VectorSearchSplit split : splits) {
+ futures.add(
+ evalBatch(
+ globalIndexer,
+ indexPathFactory,
+ split.rowRangeStart(),
+ split.rowRangeEnd(),
+ split.vectorIndexFiles(),
+ vectors,
+ preFilter,
+ executor));
+ }
+
+ CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0])).join();
+
+ 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());
+ }
+ }
+ }
+
+ List<GlobalIndexResult> results = new ArrayList<>(n);
+ for (int i = 0; i < n; i++) {
+ results.add(merged[i].topK(limit));
+ }
+ return results;
+ }
+}
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
index a4b5d955b2..f4037249a3 100644
---
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
@@ -59,11 +59,11 @@ public interface BatchVectorSearchBuilder extends
Serializable {
/** Create vector scan to scan index files. */
VectorScan newVectorScan();
- /** Create vector read to read index files. */
- VectorRead newVectorRead();
+ /** Create batch vector read to read index files. */
+ BatchVectorRead newBatchVectorRead();
/** Execute batch vector search locally; result {@code i} corresponds to
{@code vectors[i]}. */
default List<GlobalIndexResult> executeBatchLocal() {
- return newVectorRead().readBatch(newVectorScan().scan());
+ return newBatchVectorRead().readBatch(newVectorScan().scan());
}
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
index 7b78b6369b..3e8dd877df 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java
@@ -106,7 +106,7 @@ public class BatchVectorSearchBuilderImpl implements
BatchVectorSearchBuilder {
}
@Override
- public VectorRead newVectorRead() {
+ public BatchVectorRead newBatchVectorRead() {
checkArgument(limit > 0, "Limit must be positive, set via
withLimit()");
checkNotNull(vectorColumn, "Vector column must be set via
withVectorColumn()");
checkArgument(
@@ -114,6 +114,6 @@ public class BatchVectorSearchBuilderImpl implements
BatchVectorSearchBuilder {
for (float[] vector : vectors) {
checkNotNull(vector, "Search vector element cannot be null");
}
- return new VectorReadImpl(table, filter, limit, vectorColumn, vectors,
options);
+ return new BatchVectorReadImpl(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 54b696b691..74e17e2845 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,7 +20,6 @@ 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. */
@@ -31,16 +30,4 @@ 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 14e3bb643d..eab7ab2737 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
@@ -18,218 +18,81 @@
package org.apache.paimon.table.source;
-import org.apache.paimon.fs.FileIO;
-import org.apache.paimon.globalindex.GlobalIndexIOMeta;
import org.apache.paimon.globalindex.GlobalIndexReadThreadPool;
-import org.apache.paimon.globalindex.GlobalIndexReader;
import org.apache.paimon.globalindex.GlobalIndexResult;
-import org.apache.paimon.globalindex.GlobalIndexScanner;
import org.apache.paimon.globalindex.GlobalIndexer;
-import org.apache.paimon.globalindex.GlobalIndexerFactoryUtils;
-import org.apache.paimon.globalindex.OffsetGlobalIndexReader;
import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
-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.table.FileStoreTable;
import org.apache.paimon.types.DataField;
-import org.apache.paimon.utils.IOUtils;
import org.apache.paimon.utils.RoaringNavigableMap64;
-import javax.annotation.Nullable;
-
-import java.io.IOException;
-import java.io.Serializable;
import java.util.ArrayList;
-import java.util.Collections;
-import java.util.Comparator;
-import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
-import java.util.Set;
-import java.util.TreeSet;
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.checkNotNull;
-import static org.apache.paimon.utils.Preconditions.checkState;
/** Implementation for {@link VectorRead}. */
-public class VectorReadImpl implements VectorRead, Serializable {
+public class VectorReadImpl extends AbstractVectorRead implements VectorRead {
private static final long serialVersionUID = 1L;
- protected final FileStoreTable table;
- private final Predicate filter;
- protected final int limit;
- protected final DataField vectorColumn;
- protected final float[][] vectors;
- protected final Map<String, String> options;
-
- public VectorReadImpl(
- FileStoreTable table,
- Predicate filter,
- int limit,
- DataField vectorColumn,
- float[][] vectors) {
- this(table, filter, limit, vectorColumn, vectors,
Collections.emptyMap());
- }
+ protected final float[] vector;
public VectorReadImpl(
FileStoreTable table,
Predicate filter,
int limit,
DataField vectorColumn,
- float[][] vectors,
+ float[] vector,
Map<String, String> options) {
- this.table = table;
- this.filter = filter;
- this.limit = limit;
- this.vectorColumn = vectorColumn;
- this.vectors = vectors;
- this.options =
- options == null
- ? Collections.emptyMap()
- : Collections.unmodifiableMap(new HashMap<>(options));
+ super(table, filter, limit, vectorColumn, options);
+ this.vector = vector;
}
@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()) {
- List<GlobalIndexResult> empty = new ArrayList<>(n);
- for (int i = 0; i < n; i++) {
- empty.add(GlobalIndexResult.createEmpty());
- }
- return empty;
+ return GlobalIndexResult.createEmpty();
}
RoaringNavigableMap64 preFilter = preFilter(splits).orElse(null);
- IndexFileMeta firstFile = splits.get(0).vectorIndexFiles().get(0);
- String indexType = firstFile.indexType();
- GlobalIndexMeta firstMeta = checkNotNull(firstFile.globalIndexMeta());
- GlobalIndexer globalIndexer;
- if (firstMeta.extraFieldIds() != null) {
- globalIndexer =
- GlobalIndexerFactoryUtils.load(indexType)
- .create(
- firstMeta.getIndexField(table.rowType()),
- firstMeta.getExtraFields(table.rowType()),
- table.coreOptions().toConfiguration());
- } else {
- globalIndexer =
- GlobalIndexerFactoryUtils.load(indexType)
- .create(vectorColumn,
table.coreOptions().toConfiguration());
- }
+ GlobalIndexer globalIndexer = createGlobalIndexer(splits);
IndexPathFactory indexPathFactory =
table.store().pathFactory().globalIndexFileFactory();
int parallelism =
table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM);
ExecutorService executor =
GlobalIndexReadThreadPool.getExecutorService(parallelism);
- List<CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>>
futures =
+ List<CompletableFuture<Optional<ScoredGlobalIndexResult>>> futures =
new ArrayList<>(splits.size());
for (VectorSearchSplit split : splits) {
futures.add(
- evalBatch(
+ eval(
globalIndexer,
indexPathFactory,
split.rowRangeStart(),
split.rowRangeEnd(),
split.vectorIndexFiles(),
+ vector,
preFilter,
executor));
}
CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0])).join();
- 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());
- }
+ ScoredGlobalIndexResult merged = ScoredGlobalIndexResult.createEmpty();
+ for (CompletableFuture<Optional<ScoredGlobalIndexResult>> future :
futures) {
+ Optional<ScoredGlobalIndexResult> splitResult = future.join();
+ if (splitResult.isPresent()) {
+ merged = merged.or(splitResult.get());
}
}
-
- 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) {
- Set<IndexFileMeta> scalarIndexFiles =
- new TreeSet<>(Comparator.comparing(IndexFileMeta::fileName));
- for (VectorSearchSplit split : splits) {
- scalarIndexFiles.addAll(split.scalarIndexFiles());
- }
-
- Optional<GlobalIndexScanner> optionalScanner =
- GlobalIndexScanner.create(table, scalarIndexFiles);
- if (!optionalScanner.isPresent()) {
- return Optional.empty();
- }
- try (GlobalIndexScanner scanner = optionalScanner.get()) {
- return scanner.scan(filter).map(GlobalIndexResult::results);
- } catch (IOException e) {
- throw new RuntimeException(e);
- }
- }
-
- protected CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>
evalBatch(
- GlobalIndexer globalIndexer,
- IndexPathFactory indexPathFactory,
- long rowRangeStart,
- long rowRangeEnd,
- 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());
- indexIOMetaList.add(
- new GlobalIndexIOMeta(
- indexPathFactory.toPath(indexFile),
- indexFile.fileSize(),
- meta.indexMeta()));
- }
- return indexIOMetaList;
+ return merged.topK(limit);
}
}
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 b7652719ad..0ca9612436 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
@@ -107,7 +107,6 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
@Override
public VectorRead newVectorRead() {
checkNotNull(vector, "vector must be set via withVector()");
- return new VectorReadImpl(
- table, filter, limit, vectorColumn, new float[][] {vector},
options);
+ return new VectorReadImpl(table, filter, limit, vectorColumn, 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 ed9705c786..9898f93594 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
@@ -595,12 +595,14 @@ public class VectorSearchBuilderTest extends
TableTestBase {
{0.7f, 0.7f}
};
- List<GlobalIndexResult> results =
+ BatchVectorSearchBuilder searchBuilder =
table.newBatchVectorSearchBuilder()
.withVectors(queryVectors)
.withLimit(2)
- .withVectorColumn(VECTOR_FIELD_NAME)
- .executeBatchLocal();
+ .withVectorColumn(VECTOR_FIELD_NAME);
+
+ List<GlobalIndexResult> results =
+
searchBuilder.newBatchVectorRead().readBatch(searchBuilder.newVectorScan().scan());
assertThat(results).hasSize(3);
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 86bad585b4..6b5fbe3ee6 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
@@ -772,7 +772,7 @@ public class LuminaVectorGlobalIndexTest {
LuminaVectorIndexOptions indexOptions = new
LuminaVectorIndexOptions(options);
LuminaVectorGlobalIndexWriter writer =
new LuminaVectorGlobalIndexWriter(fileWriter, vectorType,
indexOptions);
- Arrays.stream(vectors).forEach(writer::write);
+ writeVectors(writer, vectors);
List<ResultEntry> results = writer.finish();
List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
@@ -823,7 +823,7 @@ public class LuminaVectorGlobalIndexTest {
LuminaVectorIndexOptions indexOptions = new
LuminaVectorIndexOptions(options);
LuminaVectorGlobalIndexWriter writer =
new LuminaVectorGlobalIndexWriter(fileWriter, vectorType,
indexOptions);
- Arrays.stream(vectors).forEach(writer::write);
+ writeVectors(writer, vectors);
List<ResultEntry> results = writer.finish();
List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
@@ -866,7 +866,7 @@ public class LuminaVectorGlobalIndexTest {
new LuminaVectorGlobalIndexWriter(fileWriter, vectorType,
indexOptions);
List<float[]> testVectors = generateRandomVectors(numVectors,
dimension);
- testVectors.forEach(writer::write);
+ writeVectors(writer, testVectors);
List<ResultEntry> results = writer.finish();
List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
@@ -912,7 +912,7 @@ public class LuminaVectorGlobalIndexTest {
new LuminaVectorGlobalIndexWriter(fileWriter, vectorType,
indexOptions);
List<float[]> testVectors = generateRandomVectors(numVectors,
dimension);
- testVectors.forEach(writer::write);
+ writeVectors(writer, testVectors);
List<ResultEntry> results = writer.finish();
List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
diff --git a/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
index 03477ef73b..f987680a36 100644
--- a/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
+++ b/paimon-python/pypaimon/table/source/batch_vector_search_builder.py
@@ -22,7 +22,10 @@ from abc import ABC, abstractmethod
from pypaimon.table.source.vector_search_builder import (
AbstractVectorSearchBuilderImpl,
)
-from pypaimon.table.source.vector_search_read import BatchVectorSearchReadImpl
+from pypaimon.table.source.vector_search_read import (
+ BatchVectorSearchRead,
+ BatchVectorSearchReadImpl,
+)
class BatchVectorSearchBuilder(ABC):
@@ -80,15 +83,15 @@ class BatchVectorSearchBuilder(ABC):
@abstractmethod
def new_batch_vector_search_read(self):
- # type: () -> BatchVectorSearchReadImpl
+ # type: () -> BatchVectorSearchRead
"""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()
+ return self.new_batch_vector_search_read().read_batch_plan(
+ self.new_vector_search_scan().scan()
)
@@ -106,7 +109,7 @@ class
BatchVectorSearchBuilderImpl(AbstractVectorSearchBuilderImpl,
return self
def new_batch_vector_search_read(self):
- # type: () -> BatchVectorSearchReadImpl
+ # type: () -> BatchVectorSearchRead
if self._limit <= 0:
raise ValueError("Limit must be positive, set via with_limit()")
if self._vector_column is None:
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index bd2c6f9287..005a2ce85b 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -40,50 +40,29 @@ class VectorSearchRead(ABC):
pass
-class VectorSearchReadImpl(VectorSearchRead):
- """Implementation for VectorSearchRead."""
+class BatchVectorSearchRead(ABC):
+ """Batch vector search read to read index files."""
- def __init__(self, table, limit, vector_column, query_vector, filter_=None,
- options=None):
+ def read_batch_plan(self, plan):
+ # type: (VectorSearchScanPlan) -> List[GlobalIndexResult]
+ return self.read_batch(plan.splits())
+
+ @abstractmethod
+ def read_batch(self, splits):
+ # type: (List[VectorSearchSplit]) -> List[GlobalIndexResult]
+ pass
+
+
+class AbstractVectorSearchReadImpl:
+ """Base implementation for vector search reads."""
+
+ def __init__(self, table, limit, vector_column, filter_=None,
options=None):
self._table = table
self._limit = limit
self._vector_column = vector_column
- self._query_vector = query_vector
self._filter = filter_
self._options = dict(options or {})
- def read(self, splits):
- # type: (List[VectorSearchSplit]) -> GlobalIndexResult
- if not splits:
- 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, query_vector, pre_filter
- )
- for split in splits
- ]
-
- wait(futures)
-
- merged_scores = {}
- for future in futures:
- split_result = future.result()
- if split_result is not None:
- score_getter = split_result.score_getter()
- for row_id in split_result.results():
- if row_id not in merged_scores:
- merged_scores[row_id] = score_getter(row_id)
-
- return DictBasedScoredIndexResult(merged_scores).top_k(self._limit)
-
def _pre_filter(self, splits):
# type: (list) -> Optional[RoaringBitmap64]
"""Evaluate the scalar filter against scalar global indexes to produce
a row-id bitmap."""
@@ -158,12 +137,50 @@ class VectorSearchReadImpl(VectorSearchRead):
return future
-class BatchVectorSearchReadImpl(VectorSearchReadImpl):
+class VectorSearchReadImpl(AbstractVectorSearchReadImpl, VectorSearchRead):
+ """Implementation for VectorSearchRead."""
+
+ def __init__(self, table, limit, vector_column, query_vector, filter_=None,
+ options=None):
+ super().__init__(table, limit, vector_column,
+ filter_=filter_, options=options)
+ self._query_vector = query_vector
+
+ def read(self, splits):
+ # type: (List[VectorSearchSplit]) -> GlobalIndexResult
+ if not splits:
+ return GlobalIndexResult.create_empty()
+
+ pre_filter = self._pre_filter(splits)
+ futures = [
+ self._eval(
+ split.row_range_start, split.row_range_end,
+ split.vector_index_files, self._query_vector, pre_filter
+ )
+ for split in splits
+ ]
+
+ wait(futures)
+
+ merged_scores = {}
+ for future in futures:
+ split_result = future.result()
+ if split_result is not None:
+ score_getter = split_result.score_getter()
+ for row_id in split_result.results():
+ if row_id not in merged_scores:
+ merged_scores[row_id] = score_getter(row_id)
+
+ return DictBasedScoredIndexResult(merged_scores).top_k(self._limit)
+
+
+class BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
+ BatchVectorSearchRead):
"""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,
+ super().__init__(table, limit, vector_column,
filter_=filter_, options=options)
self._query_vectors = list(query_vectors)
@@ -174,9 +191,32 @@ class BatchVectorSearchReadImpl(VectorSearchReadImpl):
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]
+ futures_by_vector = [
+ [
+ self._eval(
+ split.row_range_start, split.row_range_end,
+ split.vector_index_files, vector, pre_filter
+ )
+ for split in splits
+ ]
+ for vector in self._query_vectors
+ ]
+
+ for futures in futures_by_vector:
+ wait(futures)
+
+ results = []
+ for futures in futures_by_vector:
+ merged_scores = {}
+ for future in futures:
+ split_result = future.result()
+ if split_result is not None:
+ score_getter = split_result.score_getter()
+ for row_id in split_result.results():
+ if row_id not in merged_scores:
+ merged_scores[row_id] = score_getter(row_id)
+
results.append(DictBasedScoredIndexResult(merged_scores).top_k(self._limit))
+ return results
def _create_vector_reader(index_type, file_io, index_path, index_io_meta_list,
options=None):
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 94e45316ce..1ddc9c9ac9 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,7 +45,6 @@ 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
@@ -60,25 +59,13 @@ public class SparkVectorReadImpl extends VectorReadImpl {
Predicate filter,
int limit,
DataField vectorColumn,
- float[][] vectors) {
- super(table, filter, limit, vectorColumn, vectors);
- }
-
- public SparkVectorReadImpl(
- FileStoreTable table,
- Predicate filter,
- int limit,
- DataField vectorColumn,
- float[][] vectors,
+ float[] vector,
Map<String, String> options) {
- super(table, filter, limit, vectorColumn, vectors, options);
+ super(table, filter, limit, vectorColumn, vector, 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();
}
@@ -117,25 +104,25 @@ public class SparkVectorReadImpl extends VectorReadImpl {
ExecutorService executor =
GlobalIndexReadThreadPool.getExecutorService(
Math.min(parallelism, group.size()));
-
List<CompletableFuture<List<Optional<ScoredGlobalIndexResult>>>> futures =
+ List<CompletableFuture<Optional<ScoredGlobalIndexResult>>>
futures =
new ArrayList<>(group.size());
for (byte[] bytes : group) {
VectorSearchSplit split = deserializeSplit(bytes);
futures.add(
- evalBatch(
+ eval(
globalIndexer,
indexPathFactory,
split.rowRangeStart(),
split.rowRangeEnd(),
split.vectorIndexFiles(),
+ vector,
includeRowIds,
executor));
}
CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0])).join();
ScoredGlobalIndexResult result =
ScoredGlobalIndexResult.createEmpty();
- 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);
+ for (CompletableFuture<Optional<ScoredGlobalIndexResult>>
f : futures) {
+ Optional<ScoredGlobalIndexResult> next = f.join();
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 91928b7f84..be8d3d8cad 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
@@ -38,7 +38,6 @@ public class SparkVectorSearchBuilderImpl extends
VectorSearchBuilderImpl {
@Override
public VectorRead newVectorRead() {
- return new SparkVectorReadImpl(
- table, filter, limit, vectorColumn, new float[][] {vector},
options);
+ return new SparkVectorReadImpl(table, filter, limit, vectorColumn,
vector, options);
}
}
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 bd4fc35884..f08c072984 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
@@ -394,7 +394,7 @@ public class VectorGlobalIndexTest {
GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
VectorGlobalIndexWriter writer = createIvfPqWriter(fileWriter,
vectorType, options);
- Arrays.stream(vectors).forEach(writer::write);
+ writeVectors(writer, vectors);
List<ResultEntry> results = writer.finish();
List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
@@ -446,7 +446,7 @@ public class VectorGlobalIndexTest {
GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
VectorGlobalIndexWriter writer = createIvfPqWriter(fileWriter,
vectorType, options);
- Arrays.stream(vectors).forEach(writer::write);
+ writeVectors(writer, vectors);
List<ResultEntry> results = writer.finish();
List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
@@ -501,7 +501,7 @@ public class VectorGlobalIndexTest {
GlobalIndexFileWriter fileWriter = createFileWriter(indexPath);
VectorGlobalIndexWriter writer = createIvfPqWriter(fileWriter,
vectorType, options);
- Arrays.stream(vectors).forEach(writer::write);
+ writeVectors(writer, vectors);
List<ResultEntry> results = writer.finish();
List<GlobalIndexIOMeta> metas = toIOMetas(results, indexPath);
@@ -558,6 +558,12 @@ public class VectorGlobalIndexTest {
return options;
}
+ private void writeVectors(VectorGlobalIndexWriter writer, float[][]
vectors) {
+ for (int i = 0; i < vectors.length; i++) {
+ writer.write(vectors[i], i);
+ }
+ }
+
private GlobalIndexFileWriter createFileWriter(Path path) {
return new GlobalIndexFileWriter() {
@Override
