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 5f6876c025 [flink] Support distributed vector search (#8834)
5f6876c025 is described below
commit 5f6876c02523e2ccdd56767cf4f14f5950047d29
Author: liangjie <[email protected]>
AuthorDate: Tue Jul 28 08:24:46 2026 +0800
[flink] Support distributed vector search (#8834)
---
docs/docs/flink/procedures.md | 12 +-
docs/docs/multimodal-table/global-index/vector.mdx | 4 +-
docs/docs/primary-key-table/global-index.mdx | 5 +-
.../table/source/VectorSearchBuilderImpl.java | 2 +-
.../flink/procedure/VectorSearchProcedure.java | 429 +++++++++++--
.../vectorsearch/FlinkDataEvolutionVectorRead.java | 590 ++++++++++++++++++
.../vectorsearch/FlinkPrimaryKeyVectorRead.java | 193 ++++++
.../vectorsearch/FlinkVectorSearchBuilderImpl.java | 51 ++
.../procedure/VectorSearchProcedureITCase.java | 674 ++++++++++++++++++++-
.../FlinkDataEvolutionVectorReadTest.java | 555 +++++++++++++++++
.../FlinkPrimaryKeyVectorReadTest.java | 163 +++++
11 files changed, 2602 insertions(+), 76 deletions(-)
diff --git a/docs/docs/flink/procedures.md b/docs/docs/flink/procedures.md
index eb91819554..927bfe231b 100644
--- a/docs/docs/flink/procedures.md
+++ b/docs/docs/flink/procedures.md
@@ -1177,8 +1177,10 @@ All available procedures are listed below.
vector_column => 'columnName',<br/>
query_vector => 'v1,v2,...',<br/>
top_k => topK,<br/>
- projection => 'col1,col2',<br/>
- options => 'key1=value1;key2=value2')<br/>
+ projection => 'col1,col2,__paimon_search_score',<br/>
+ options => 'key1=value1;key2=value2',<br/>
+ `where` => 'predicate',<br/>
+ partitions => 'pt1=v1,pt2=v2;pt1=v3,pt2=v4')<br/>
</td>
<td>
To perform vector similarity search on a table with a global vector
index. Returns JSON-serialized rows. Arguments:
@@ -1188,6 +1190,8 @@ All available procedures are listed below.
<li>top_k(required): the number of nearest neighbors to
return.</li>
<li>projection(optional): comma-separated column names to include
in the result. If omitted, all columns are returned.</li>
<li>options(optional): additional dynamic options of the
table.</li>
+ <li>where(optional): a SQL predicate applied before Top-K.</li>
+ <li>partitions(optional): semicolon-separated specs for a
partitioned table.</li>
</td>
<td>
CALL sys.vector_search(<br/>
@@ -1200,7 +1204,9 @@ All available procedures are listed below.
vector_column => 'embedding',<br/>
query_vector => '1.0,2.0,3.0',<br/>
top_k => 5,<br/>
- projection => 'id,name')
+ projection => 'id,name,__paimon_search_score',<br/>
+ options =>
'vector-search.distribute.enabled=true;ivf.nprobe=32',<br/>
+ `where` => 'status = ''active''')
</td>
</tr>
</tbody>
diff --git a/docs/docs/multimodal-table/global-index/vector.mdx
b/docs/docs/multimodal-table/global-index/vector.mdx
index 6399d3c258..d29d076769 100644
--- a/docs/docs/multimodal-table/global-index/vector.mdx
+++ b/docs/docs/multimodal-table/global-index/vector.mdx
@@ -335,7 +335,9 @@ CALL sys.vector_search(
vector_column => 'embedding',
query_vector => '1.0,2.0,3.0',
top_k => 5,
- projection => 'id,name'
+ projection => 'id,name,__paimon_search_score',
+ options => 'ivf.nprobe=32',
+ `where` => 'id > 100'
);
```
diff --git a/docs/docs/primary-key-table/global-index.mdx
b/docs/docs/primary-key-table/global-index.mdx
index e76845e6c3..04008c2db5 100644
--- a/docs/docs/primary-key-table/global-index.mdx
+++ b/docs/docs/primary-key-table/global-index.mdx
@@ -349,8 +349,9 @@ CALL sys.vector_search(
vector_column => 'embedding',
query_vector => '0.1,0.2,0.3',
top_k => 10,
- projection => 'id,status',
- options => 'ivf.nprobe=32'
+ projection => 'id,status,__paimon_search_score',
+ options => 'ivf.nprobe=32',
+ `where` => 'status = ''active'''
);
```
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 4b59d43812..fb0c972365 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
@@ -155,7 +155,7 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
&&
table.coreOptions().primaryKeyVectorIndexColumns().contains(vectorColumn.name());
}
- VectorSearchBuilderImpl withSnapshot(Snapshot snapshot) {
+ public VectorSearchBuilderImpl withSnapshot(Snapshot snapshot) {
this.pinnedSnapshot = snapshot;
return this;
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/VectorSearchProcedure.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/VectorSearchProcedure.java
index 8a75b8f9d4..d49810c9dd 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/VectorSearchProcedure.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/VectorSearchProcedure.java
@@ -18,17 +18,34 @@
package org.apache.paimon.flink.procedure;
+import org.apache.paimon.CoreOptions;
+import org.apache.paimon.Snapshot;
+import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.flink.predicate.SimpleSqlPredicateConvertor;
+import org.apache.paimon.flink.vectorsearch.FlinkVectorSearchBuilderImpl;
import org.apache.paimon.format.json.JsonFormatWriter;
import org.apache.paimon.format.json.JsonOptions;
import org.apache.paimon.fs.PositionOutputStream;
import org.apache.paimon.globalindex.GlobalIndexResult;
import org.apache.paimon.options.Options;
+import org.apache.paimon.partition.PartitionPredicate;
+import org.apache.paimon.predicate.Predicate;
import org.apache.paimon.reader.RecordReader;
+import org.apache.paimon.reader.ScoreRecordIterator;
+import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.Table;
import org.apache.paimon.table.source.ReadBuilder;
import org.apache.paimon.table.source.TableScan;
+import org.apache.paimon.table.source.VectorScan;
+import org.apache.paimon.table.source.VectorSearchBuilder;
+import org.apache.paimon.table.source.VectorSearchBuilderImpl;
+import org.apache.paimon.table.source.snapshot.TimeTravelUtil;
+import org.apache.paimon.types.DataField;
+import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
+import org.apache.paimon.utils.InternalRowUtils;
+import org.apache.paimon.utils.Pair;
import org.apache.paimon.utils.StringUtils;
import org.apache.flink.table.annotation.ArgumentHint;
@@ -36,40 +53,32 @@ import org.apache.flink.table.annotation.DataTypeHint;
import org.apache.flink.table.annotation.ProcedureHint;
import org.apache.flink.table.procedure.ProcedureContext;
+import javax.annotation.Nullable;
+
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
-import java.util.stream.Collectors;
+import java.util.Optional;
+import java.util.Set;
-/**
- * Vector search procedure. This procedure takes one vector and searches for
topK nearest vectors.
- * Usage:
- *
- * <pre><code>
- * CALL sys.vector_search(
- * `table` => 'tableId',
- * vector_column => 'v',
- * query_vector => '1.0,2.0,3.0',
- * top_k => 5
- * )
- *
- * -- with projection and options
- * CALL sys.vector_search(
- * `table` => 'tableId',
- * vector_column => 'v',
- * query_vector => '1.0,2.0,3.0',
- * top_k => 5,
- * projection => 'id,name',
- * options => 'k1=v1;k2=v2'
- * )
- * </code></pre>
- */
+import static
org.apache.paimon.partition.PartitionPredicate.splitPartitionPredicatesAndDataPredicates;
+import static org.apache.paimon.predicate.PredicateVisitor.collectFieldNames;
+import static org.apache.paimon.utils.ParameterUtils.getPartitions;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+
+/** Procedure for local or distributed vector search with optional filtering
and scores. */
public class VectorSearchProcedure extends ProcedureBase {
public static final String IDENTIFIER = "vector_search";
+ public static final String SEARCH_SCORE = "__paimon_search_score";
+
+ private static final int MAX_TOP_K = 10_000;
+ private static final DataField SEARCH_SCORE_FIELD =
+ new DataField(Integer.MAX_VALUE, SEARCH_SCORE, DataTypes.FLOAT());
@ProcedureHint(
argument = {
@@ -81,7 +90,12 @@ public class VectorSearchProcedure extends ProcedureBase {
name = "projection",
type = @DataTypeHint("STRING"),
isOptional = true),
- @ArgumentHint(name = "options", type =
@DataTypeHint("STRING"), isOptional = true)
+ @ArgumentHint(name = "options", type =
@DataTypeHint("STRING"), isOptional = true),
+ @ArgumentHint(name = "where", type = @DataTypeHint("STRING"),
isOptional = true),
+ @ArgumentHint(
+ name = "partitions",
+ type = @DataTypeHint("STRING"),
+ isOptional = true)
})
public String[] call(
ProcedureContext procedureContext,
@@ -90,89 +104,370 @@ public class VectorSearchProcedure extends ProcedureBase {
String queryVectorStr,
Integer topK,
String projection,
- String options)
+ String options,
+ String where,
+ String partitions)
throws Exception {
- Table table = table(tableId);
+ validateSearch(vectorColumn, queryVectorStr, topK);
+ Table table = table(tableId);
Map<String, String> optionsMap = optionalConfigMap(options);
+ String queryAuthOption = CoreOptions.QUERY_AUTH_ENABLED.key();
+ checkArgument(
+ !optionsMap.containsKey(queryAuthOption),
+ "Option '%s' is not allowed",
+ queryAuthOption);
if (!optionsMap.isEmpty()) {
table = table.copy(optionsMap);
}
+ checkArgument(
+ table instanceof FileStoreTable, "Vector search requires a
file store table.");
+
+ FileStoreTable fileStoreTable = (FileStoreTable) table;
+ checkArgument(
+ !fileStoreTable.coreOptions().queryAuthEnabled(),
+ "Vector search does not support tables with query auth
enabled.");
+ RowType tableType = fileStoreTable.rowType();
+ checkArgument(
+ tableType.containsField(vectorColumn),
+ "Vector column '%s' does not exist in table '%s'.",
+ vectorColumn,
+ tableId);
+ checkArgument(
+ tableType.notContainsField(SEARCH_SCORE),
+ "Table column '%s' conflicts with vector-search metadata.",
+ SEARCH_SCORE);
float[] queryVector = parseVector(queryVectorStr);
+ Predicate filter = parseFilter(where, tableType);
+ PartitionPredicate partitionFilter = parsePartitions(partitions,
fileStoreTable);
+ Projection parsedProjection = Projection.parse(projection, tableType,
filter);
+ FilterParts filterParts = FilterParts.from(filter, fileStoreTable);
+ validatePrimaryKeyFilter(fileStoreTable, vectorColumn, filterParts);
- GlobalIndexResult result =
- table.newVectorSearchBuilder()
+ Snapshot snapshot = TimeTravelUtil.tryTravelOrLatest(fileStoreTable);
+ if (snapshot == null) {
+ return new String[0];
+ }
+
+ VectorSearchBuilder builder =
+ newVectorSearchBuilder(procedureContext, fileStoreTable)
+ .withSnapshot(snapshot)
.withVector(queryVector)
.withVectorColumn(vectorColumn)
.withLimit(topK)
- .withOptions(optionsMap)
- .executeLocal();
+ .withOptions(optionsMap);
+ if (filter != null) {
+ builder.withFilter(filter);
+ }
+ if (partitionFilter != null) {
+ builder.withPartitionFilter(partitionFilter);
+ }
- RowType tableRowType = table.rowType();
- int[] projectionIndices = parseProjection(projection, tableRowType);
+ VectorScan.Plan vectorPlan = builder.newVectorScan().scan();
+ GlobalIndexResult result = builder.newVectorRead().read(vectorPlan);
- ReadBuilder readBuilder = table.newReadBuilder();
- if (projectionIndices != null) {
- readBuilder.withProjection(projectionIndices);
+ ReadBuilder readBuilder = fileStoreTable.newReadBuilder();
+ if (filter != null) {
+ readBuilder.withFilter(filter);
+ }
+ PartitionPredicate effectivePartitionFilter =
+ filterParts.mergePartitionFilter(partitionFilter);
+ if (effectivePartitionFilter != null) {
+ readBuilder.withPartitionFilter(effectivePartitionFilter);
+ }
+ if (parsedProjection.readProjection != null) {
+ readBuilder.withProjection(parsedProjection.readProjection);
}
+ TableScan.Plan readPlan =
readBuilder.newScan().withGlobalIndexResult(result).plan();
+ return readRows(readBuilder, readPlan, parsedProjection);
+ }
- TableScan.Plan plan =
readBuilder.newScan().withGlobalIndexResult(result).plan();
+ private static VectorSearchBuilderImpl newVectorSearchBuilder(
+ ProcedureContext context, FileStoreTable table) {
+ if (!table.coreOptions().vectorSearchDistributeEnabled()) {
+ return (VectorSearchBuilderImpl) table.newVectorSearchBuilder();
+ }
+ return new FlinkVectorSearchBuilderImpl(table,
context.getExecutionEnvironment());
+ }
+
+ private static void validateSearch(
+ String vectorColumn, String queryVector, @Nullable Integer topK) {
+ checkArgument(
+ !StringUtils.isNullOrWhitespaceOnly(vectorColumn),
+ "vector_column must not be blank");
+ checkArgument(
+ !StringUtils.isNullOrWhitespaceOnly(queryVector),
"query_vector must not be blank");
+ checkArgument(topK != null && topK > 0, "top_k must be positive");
+ checkArgument(topK <= MAX_TOP_K, "top_k must not exceed %s",
MAX_TOP_K);
+ }
- RowType readType =
- projectionIndices != null ?
tableRowType.project(projectionIndices) : tableRowType;
+ private static float[] parseVector(String vectorStr) {
+ String[] parts = vectorStr.split(",", -1);
+ float[] vector = new float[parts.length];
+ for (int i = 0; i < parts.length; i++) {
+ String part = parts[i].trim();
+ checkArgument(!part.isEmpty(), "query_vector contains an empty
value");
+ try {
+ vector[i] = Float.parseFloat(part);
+ } catch (NumberFormatException e) {
+ throw new IllegalArgumentException(
+ String.format("Invalid query vector value: '%s'",
part), e);
+ }
+ checkArgument(
+ !Float.isNaN(vector[i]) && !Float.isInfinite(vector[i]),
+ "query_vector values must be finite");
+ }
+ return vector;
+ }
+
+ @Nullable
+ private static Predicate parseFilter(@Nullable String where, RowType
tableType)
+ throws Exception {
+ if (StringUtils.isNullOrWhitespaceOnly(where)) {
+ return null;
+ }
+ return new
SimpleSqlPredicateConvertor(tableType).convertSqlToPredicate(where);
+ }
+
+ @Nullable
+ private static PartitionPredicate parsePartitions(
+ @Nullable String partitions, FileStoreTable table) {
+ if (StringUtils.isNullOrWhitespaceOnly(partitions)) {
+ return null;
+ }
+
+ List<String> partitionKeys = table.partitionKeys();
+ checkArgument(
+ !partitionKeys.isEmpty(),
+ "partitions can only be used with a partitioned table '%s'.",
+ table.name());
+ List<Map<String, String>> specs = getPartitions(partitions.split(";"));
+ Set<String> expectedKeys = new HashSet<>(partitionKeys);
+ for (Map<String, String> spec : specs) {
+ checkArgument(
+ spec.keySet().equals(expectedKeys),
+ "Partition spec for table '%s' must exactly match
partition keys %s, but was %s.",
+ table.name(),
+ partitionKeys,
+ spec);
+ }
+ return PartitionPredicate.fromMaps(
+ table.schema().logicalPartitionType(),
+ specs,
+ table.coreOptions().partitionDefaultName());
+ }
+
+ private static void validatePrimaryKeyFilter(
+ FileStoreTable table, String vectorColumn, FilterParts
filterParts) {
+ if
(!table.coreOptions().primaryKeyVectorIndexColumns().contains(vectorColumn)
+ || !filterParts.hasDataFilter()) {
+ return;
+ }
+ checkArgument(
+ table.coreOptions().deletionVectorsEnabled()
+ && !table.coreOptions().deletionVectorsMergeOnRead(),
+ "Primary-key vector where filter on non-partition columns
requires "
+ + "deletion-vectors.enabled = true and "
+ + "deletion-vectors.merge-on-read = false.");
+ }
+ private static String[] readRows(
+ ReadBuilder readBuilder, TableScan.Plan plan, Projection
projection)
+ throws IOException {
ByteArrayOutputStream byteOut = new ByteArrayOutputStream(1024);
JsonOptions jsonOptions = new JsonOptions(new Options());
+ List<Float> scores = new ArrayList<>();
try (JsonFormatWriter jsonWriter =
new JsonFormatWriter(
new ByteArrayPositionOutputStream(byteOut),
- readType,
+ projection.outputType,
jsonOptions,
"none");
- RecordReader<InternalRow> reader =
readBuilder.newRead().createReader(plan)) {
- reader.forEachRemaining(
- row -> {
- try {
- jsonWriter.addElement(row);
- } catch (Exception e) {
- throw new RuntimeException("Failed to convert row
to JSON string", e);
- }
- });
+ RecordReader<InternalRow> reader =
+
readBuilder.newRead().executeFilter().createReader(plan)) {
+ RecordReader.RecordIterator<InternalRow> batch;
+ while ((batch = reader.readBatch()) != null) {
+ try {
+ checkArgument(
+ batch instanceof ScoreRecordIterator,
+ "Vector-search reader did not expose search
scores.");
+ ScoreRecordIterator<InternalRow> scoredBatch =
+ (ScoreRecordIterator<InternalRow>) batch;
+ InternalRow row;
+ while ((row = scoredBatch.next()) != null) {
+ float score = scoredBatch.returnedScore();
+ checkArgument(
+ !Float.isNaN(score),
+ "Vector-search reader returned a missing
score.");
+ jsonWriter.addElement(projection.output(row, score));
+ scores.add(score);
+ }
+ } finally {
+ batch.releaseBatch();
+ }
+ }
}
String[] lines =
StringUtils.split(byteOut.toString("UTF-8"),
jsonOptions.getLineDelimiter());
- List<String> rows = new ArrayList<>(lines.length);
+ List<ScoredJson> rows = new ArrayList<>(lines.length);
+ int scoreIndex = 0;
for (String line : lines) {
String trimmed = line.trim();
if (!trimmed.isEmpty()) {
- rows.add(trimmed);
+ checkArgument(
+ scoreIndex < scores.size(),
+ "Vector-search JSON rows and scores are misaligned.");
+ rows.add(new ScoredJson(trimmed, scores.get(scoreIndex++)));
}
}
- return rows.toArray(new String[0]);
+ checkArgument(
+ scoreIndex == scores.size(), "Vector-search JSON rows and
scores are misaligned.");
+ Collections.sort(
+ rows,
+ (left, right) -> {
+ int scoreOrder = Float.compare(right.score, left.score);
+ return scoreOrder != 0 ? scoreOrder :
left.json.compareTo(right.json);
+ });
+ return rows.stream().map(row -> row.json).toArray(String[]::new);
}
- private static float[] parseVector(String vectorStr) {
- String[] parts = StringUtils.split(vectorStr, ",");
- float[] vector = new float[parts.length];
- for (int i = 0; i < parts.length; i++) {
- vector[i] = Float.parseFloat(parts[i].trim());
+ @Override
+ public String identifier() {
+ return IDENTIFIER;
+ }
+
+ private static class Projection {
+
+ private final RowType readType;
+ private final RowType outputType;
+ @Nullable private final int[] readProjection;
+ @Nullable private final int[] outputToRead;
+
+ private Projection(
+ RowType readType,
+ RowType outputType,
+ @Nullable int[] readProjection,
+ @Nullable int[] outputToRead) {
+ this.readType = readType;
+ this.outputType = outputType;
+ this.readProjection = readProjection;
+ this.outputToRead = outputToRead;
+ }
+
+ private static Projection parse(
+ @Nullable String projection, RowType tableType, @Nullable
Predicate filter) {
+ if (StringUtils.isNullOrWhitespaceOnly(projection)) {
+ return new Projection(tableType, tableType, null, null);
+ }
+
+ String[] names = projection.split(",", -1);
+ List<Integer> readIndices = new ArrayList<>();
+ List<DataField> outputFields = new ArrayList<>();
+ int[] outputToRead = new int[names.length];
+ Set<String> seen = new HashSet<>();
+ for (int i = 0; i < names.length; i++) {
+ String name = names[i].trim();
+ checkArgument(!name.isEmpty(), "Projection column must not be
blank");
+ checkArgument(seen.add(name), "Duplicate projection column:
%s", name);
+ if (SEARCH_SCORE.equals(name)) {
+ outputFields.add(SEARCH_SCORE_FIELD);
+ outputToRead[i] = -1;
+ continue;
+ }
+
+ int tableIndex = tableType.getFieldIndex(name);
+ checkArgument(tableIndex >= 0, "Unknown projection column:
%s", name);
+ outputFields.add(tableType.getFields().get(tableIndex));
+ outputToRead[i] = readIndices.size();
+ readIndices.add(tableIndex);
+ }
+
+ if (filter != null) {
+ Set<String> filterFields = collectFieldNames(filter);
+ Set<Integer> included = new HashSet<>(readIndices);
+ List<String> tableFields = tableType.getFieldNames();
+ for (int i = 0; i < tableFields.size(); i++) {
+ if (filterFields.contains(tableFields.get(i)) &&
included.add(i)) {
+ readIndices.add(i);
+ }
+ }
+ }
+
+ int[] readProjection =
readIndices.stream().mapToInt(Integer::intValue).toArray();
+ return new Projection(
+ tableType.project(readProjection),
+ new RowType(outputFields),
+ readProjection,
+ outputToRead);
+ }
+
+ private InternalRow output(InternalRow row, float score) {
+ if (outputToRead == null) {
+ return row;
+ }
+ GenericRow output = new GenericRow(outputToRead.length);
+ for (int i = 0; i < outputToRead.length; i++) {
+ int readIndex = outputToRead[i];
+ output.setField(
+ i,
+ readIndex < 0
+ ? score
+ : InternalRowUtils.get(
+ row, readIndex,
readType.getTypeAt(readIndex)));
+ }
+ return output;
}
- return vector;
}
- private static int[] parseProjection(String projection, RowType rowType) {
- if (StringUtils.isNullOrWhitespaceOnly(projection)) {
- return null;
+ private static class FilterParts {
+
+ @Nullable private final PartitionPredicate partitionFilter;
+ private final boolean hasDataFilter;
+
+ private FilterParts(@Nullable PartitionPredicate partitionFilter,
boolean hasDataFilter) {
+ this.partitionFilter = partitionFilter;
+ this.hasDataFilter = hasDataFilter;
+ }
+
+ private static FilterParts from(@Nullable Predicate filter,
FileStoreTable table) {
+ if (filter == null) {
+ return new FilterParts(null, false);
+ }
+ Pair<Optional<PartitionPredicate>, List<Predicate>> parts =
+ splitPartitionPredicatesAndDataPredicates(
+ filter, table.rowType(), table.partitionKeys());
+ return new FilterParts(parts.getLeft().orElse(null),
!parts.getRight().isEmpty());
+ }
+
+ private boolean hasDataFilter() {
+ return hasDataFilter;
+ }
+
+ @Nullable
+ private PartitionPredicate mergePartitionFilter(
+ @Nullable PartitionPredicate explicitPartitionFilter) {
+ List<PartitionPredicate> filters = new ArrayList<>(2);
+ if (partitionFilter != null) {
+ filters.add(partitionFilter);
+ }
+ if (explicitPartitionFilter != null) {
+ filters.add(explicitPartitionFilter);
+ }
+ return PartitionPredicate.and(filters);
}
- String[] projectionNames = StringUtils.split(projection, ",");
- return
rowType.getFieldIndices(Arrays.stream(projectionNames).collect(Collectors.toList()));
}
- @Override
- public String identifier() {
- return IDENTIFIER;
+ private static class ScoredJson {
+
+ private final String json;
+ private final float score;
+
+ private ScoredJson(String json, float score) {
+ this.json = json;
+ this.score = score;
+ }
}
/** A {@link PositionOutputStream} wrapping a {@link
ByteArrayOutputStream}. */
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkDataEvolutionVectorRead.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkDataEvolutionVectorRead.java
new file mode 100644
index 0000000000..e1ea1e0d6a
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkDataEvolutionVectorRead.java
@@ -0,0 +1,590 @@
+/*
+ * 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.flink.vectorsearch;
+
+import org.apache.paimon.flink.utils.StreamExecutionEnvironmentUtils;
+import org.apache.paimon.globalindex.GlobalIndexReadThreadPool;
+import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.globalindex.GlobalIndexResultSerializer;
+import org.apache.paimon.globalindex.GlobalIndexer;
+import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
+import org.apache.paimon.index.IndexPathFactory;
+import org.apache.paimon.partition.PartitionPredicate;
+import org.apache.paimon.predicate.Predicate;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.source.DataEvolutionVectorRead;
+import org.apache.paimon.table.source.IndexVectorSearchSplit;
+import org.apache.paimon.table.source.RawVectorSearchSplit;
+import org.apache.paimon.table.source.VectorScan;
+import org.apache.paimon.types.DataField;
+import org.apache.paimon.utils.InstantiationUtil;
+import org.apache.paimon.utils.Range;
+import org.apache.paimon.utils.RoaringNavigableMap64;
+
+import org.apache.flink.api.common.functions.MapFunction;
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.common.typeinfo.PrimitiveArrayTypeInfo;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.api.java.typeutils.TupleTypeInfo;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.util.CloseableIterator;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+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;
+import static org.apache.paimon.utils.Preconditions.checkNotNull;
+
+/** Flink-aware {@link DataEvolutionVectorRead}. */
+public class FlinkDataEvolutionVectorRead extends DataEvolutionVectorRead {
+
+ private static final long serialVersionUID = 1L;
+ private static final byte INDEX_RESULT = 0;
+ private static final byte RAW_RESULT = 1;
+ private static final TupleTypeInfo<Tuple2<Byte, byte[]>> RESULT_TYPE_INFO =
+ new TupleTypeInfo<Tuple2<Byte, byte[]>>(
+ BasicTypeInfo.BYTE_TYPE_INFO,
+ PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO);
+
+ private final transient StreamExecutionEnvironment env;
+
+ public FlinkDataEvolutionVectorRead(
+ FileStoreTable table,
+ @Nullable PartitionPredicate partitionFilter,
+ @Nullable Predicate filter,
+ int limit,
+ DataField vectorColumn,
+ float[] vector,
+ @Nullable Map<String, String> options,
+ StreamExecutionEnvironment env) {
+ super(table, partitionFilter, filter, limit, vectorColumn, vector,
options);
+ this.env = checkNotNull(env);
+ }
+
+ @Override
+ public GlobalIndexResult read(VectorScan.Plan plan) {
+ List<IndexVectorSearchSplit> indexSplits = new ArrayList<>();
+ List<RawVectorSearchSplit> rawSplits = new ArrayList<>();
+ splitSearchSplits(plan.splits(), indexSplits, rawSplits);
+ if (indexSplits.isEmpty() && rawSplits.isEmpty()) {
+ return GlobalIndexResult.createEmpty();
+ }
+
+ GlobalIndexer globalIndexer =
+ !indexSplits.isEmpty() && !rawSplits.isEmpty()
+ ? createGlobalIndexer(indexSplits)
+ : null;
+ if (!indexSplits.isEmpty() && !rawSplits.isEmpty()) {
+ int parallelism = flinkParallelism();
+ List<Range> rawRowRanges = rawRowRanges(rawSplits);
+ if (indexSplits.size() >= parallelism * 2L
+ && rawRowCount(rawRowRanges) >= parallelism * 2L) {
+ return readIndexAndRawSplitsInFlink(
+ indexSplits,
+ rawSplits,
+ rawRowRanges,
+ globalIndexer,
+ rawPreFilter(rawSplits),
+ parallelism);
+ }
+ }
+
+ ScoredGlobalIndexResult indexed =
+ indexSplits.isEmpty()
+ ? ScoredGlobalIndexResult.createEmpty()
+ : readIndexSplitsInFlink(indexSplits, globalIndexer);
+ ScoredGlobalIndexResult raw =
+ readRawSplitsInFlink(rawSplits, globalIndexer,
rawPreFilter(rawSplits));
+ return indexed.or(raw).topK(limit);
+ }
+
+ protected ScoredGlobalIndexResult readIndexSplitsInFlink(
+ List<IndexVectorSearchSplit> splits, @Nullable GlobalIndexer
globalIndexer) {
+ if (splits.isEmpty()) {
+ return ScoredGlobalIndexResult.createEmpty();
+ }
+
+ int parallelism = flinkParallelism();
+ if (splits.size() < parallelism * 2L) {
+ return readIndexed(
+ splits, globalIndexer == null ?
createGlobalIndexer(splits) : globalIndexer);
+ }
+
+ List<RoaringNavigableMap64> preFilters = preFilters(splits);
+ String indexType = vectorIndexType(splits);
+ int searchLimit = indexedSearchLimit(indexType);
+ List<List<SerializedSplit>> splitGroups = indexSplitGroups(splits,
preFilters, parallelism);
+ List<byte[]> remoteResults =
+ executeIndexSearchGroups(splitGroups, searchLimit,
parallelism);
+ GlobalIndexer rerankGlobalIndexer =
+ globalIndexer == null ? createGlobalIndexer(splits) :
globalIndexer;
+ return maybeRerankIndexedResult(
+ mergeRemoteResults(remoteResults, searchLimit),
+ indexType,
+ rerankGlobalIndexer,
+ vector);
+ }
+
+ private List<List<SerializedSplit>> indexSplitGroups(
+ List<IndexVectorSearchSplit> splits,
+ List<RoaringNavigableMap64> preFilters,
+ int parallelism) {
+ List<SerializedSplit> serializedSplits = new
ArrayList<>(splits.size());
+ for (int i = 0; i < splits.size(); i++) {
+ try {
+ IndexVectorSearchSplit split = splits.get(i);
+ RoaringNavigableMap64 preFilter = preFilters.isEmpty() ? null
: preFilters.get(i);
+ serializedSplits.add(
+ new SerializedSplit(
+ InstantiationUtil.serializeObject(split),
+ preFilter == null
+ ? null
+ :
InstantiationUtil.serializeObject(preFilter)));
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to serialize vector-search
split.", e);
+ }
+ }
+ return splitGroups(serializedSplits, parallelism);
+ }
+
+ protected ScoredGlobalIndexResult readRawSplitsInFlink(
+ List<RawVectorSearchSplit> splits,
+ @Nullable GlobalIndexer globalIndexer,
+ @Nullable RoaringNavigableMap64 preFilter) {
+ List<Range> rawRowRanges = rawRowRanges(splits);
+ if (rawRowRanges.isEmpty()) {
+ return ScoredGlobalIndexResult.createEmpty();
+ }
+
+ int parallelism = flinkParallelism();
+ if (rawRowCount(rawRowRanges) < parallelism * 2L) {
+ return readRawSearch(
+ rawRowRanges, preFilter, rawSearchIndexer(splits,
globalIndexer), vector);
+ }
+
+ String metric = rawSearchMetric(rawSearchIndexer(splits,
globalIndexer));
+ List<List<SerializedSplit>> splitGroups =
+ rawSplitGroups(rawRowRanges, preFilter, parallelism);
+ List<byte[]> remoteResults = executeRawSearchGroups(splitGroups,
metric, parallelism);
+ return mergeRemoteResults(remoteResults, limit);
+ }
+
+ private List<List<SerializedSplit>> rawSplitGroups(
+ List<Range> rawRowRanges, @Nullable RoaringNavigableMap64
preFilter, int parallelism) {
+ List<List<Range>> rangeGroups = rangeGroups(rawRowRanges, parallelism);
+ List<SerializedSplit> serializedSplits = new
ArrayList<>(rangeGroups.size());
+ for (List<Range> rangeGroup : rangeGroups) {
+ try {
+ RoaringNavigableMap64 groupPreFilter =
groupPreFilter(rangeGroup, preFilter);
+ serializedSplits.add(
+ new SerializedSplit(
+ InstantiationUtil.serializeObject(rangeGroup),
+ groupPreFilter == null
+ ? null
+ :
InstantiationUtil.serializeObject(groupPreFilter)));
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to serialize raw vector row
ranges.", e);
+ }
+ }
+ return splitGroups(serializedSplits, parallelism);
+ }
+
+ @Nullable
+ private RoaringNavigableMap64 groupPreFilter(
+ List<Range> rangeGroup, @Nullable RoaringNavigableMap64 preFilter)
{
+ if (preFilter == null) {
+ return null;
+ }
+
+ RoaringNavigableMap64 groupRows = new RoaringNavigableMap64();
+ for (Range range : rangeGroup) {
+ groupRows.addRange(range);
+ }
+ groupRows.and(preFilter);
+ if (groupRows.getLongCardinality() == rawRowCount(rangeGroup)) {
+ return null;
+ }
+ groupRows.runOptimize();
+ return groupRows;
+ }
+
+ private ScoredGlobalIndexResult readIndexAndRawSplitsInFlink(
+ List<IndexVectorSearchSplit> indexSplits,
+ List<RawVectorSearchSplit> rawSplits,
+ List<Range> rawRowRanges,
+ GlobalIndexer globalIndexer,
+ @Nullable RoaringNavigableMap64 rawPreFilter,
+ int parallelism) {
+ String indexType = vectorIndexType(indexSplits);
+ int searchLimit = indexedSearchLimit(indexType);
+ List<List<SerializedSplit>> indexGroups =
+ indexSplitGroups(indexSplits, preFilters(indexSplits),
parallelism);
+
+ String rawMetric = rawSearchMetric(rawSearchIndexer(rawSplits,
globalIndexer));
+ List<List<SerializedSplit>> rawGroups =
+ rawSplitGroups(rawRowRanges, rawPreFilter, parallelism);
+
+ List<Tuple2<Byte, byte[]>> results =
+ executeIndexAndRawSearchGroups(
+ indexGroups, searchLimit, rawGroups, rawMetric,
parallelism);
+ List<byte[]> indexedResults = new ArrayList<>();
+ List<byte[]> rawResults = new ArrayList<>();
+ for (Tuple2<Byte, byte[]> result : results) {
+ byte resultType = checkNotNull(result.f0, "Missing vector-search
result type.");
+ byte[] resultBytes = checkNotNull(result.f1, "Missing
vector-search result payload.");
+ if (resultType == INDEX_RESULT) {
+ indexedResults.add(resultBytes);
+ } else if (resultType == RAW_RESULT) {
+ rawResults.add(resultBytes);
+ } else {
+ throw new IllegalArgumentException(
+ "Unknown vector-search result type: " + resultType);
+ }
+ }
+
+ ScoredGlobalIndexResult indexed =
+ maybeRerankIndexedResult(
+ mergeRemoteResults(indexedResults, searchLimit),
+ indexType,
+ globalIndexer,
+ vector);
+ ScoredGlobalIndexResult raw = mergeRemoteResults(rawResults, limit);
+ return indexed.or(raw).topK(limit);
+ }
+
+ protected int flinkParallelism() {
+ int maxParallelism =
+ Math.max(1,
table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM));
+ int envParallelism = checkNotNull(env).getParallelism();
+ return envParallelism > 0
+ ? Math.max(1, Math.min(envParallelism, maxParallelism))
+ : maxParallelism;
+ }
+
+ protected List<byte[]> executeIndexSearchGroups(
+ List<List<SerializedSplit>> groups, int searchLimit, int
parallelism) {
+ String operatorName = "Vector Index Search";
+ return collectResults(
+ indexSearchStream(groups, searchLimit, parallelism,
operatorName), operatorName);
+ }
+
+ protected List<byte[]> executeRawSearchGroups(
+ List<List<SerializedSplit>> groups, String metric, int
parallelism) {
+ String operatorName = "Raw Vector Search";
+ return collectResults(
+ rawSearchStream(groups, metric, parallelism, operatorName),
operatorName);
+ }
+
+ protected List<Tuple2<Byte, byte[]>> executeIndexAndRawSearchGroups(
+ List<List<SerializedSplit>> indexGroups,
+ int searchLimit,
+ List<List<SerializedSplit>> rawGroups,
+ String rawMetric,
+ int parallelism) {
+ DataStream<Tuple2<Byte, byte[]>> indexResults =
+ addResultType(
+ indexSearchStream(
+ indexGroups, searchLimit, parallelism, "Vector
Index Search"),
+ INDEX_RESULT,
+ "Add Vector Index Result Type");
+ DataStream<Tuple2<Byte, byte[]>> rawResults =
+ addResultType(
+ rawSearchStream(rawGroups, rawMetric, parallelism,
"Raw Vector Search"),
+ RAW_RESULT,
+ "Add Raw Result Type");
+ return collectResults(indexResults.union(rawResults), "Vector Index
and Raw Search");
+ }
+
+ private DataStream<byte[]> indexSearchStream(
+ List<List<SerializedSplit>> groups,
+ int searchLimit,
+ int parallelism,
+ String operatorName) {
+ DataStream<byte[]> results =
+ taskSource(groups, operatorName)
+ .map(
+ (MapFunction<byte[], byte[]>)
+ bytes ->
+ executeIndexSearchGroup(
+
deserializeSplitGroup(bytes),
+ searchLimit,
+ parallelism))
+
.returns(PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO)
+ .name(operatorName)
+ .setParallelism(Math.min(parallelism, groups.size()));
+ return results;
+ }
+
+ private DataStream<byte[]> rawSearchStream(
+ List<List<SerializedSplit>> groups,
+ String metric,
+ int parallelism,
+ String operatorName) {
+ return taskSource(groups, operatorName)
+ .map(
+ (MapFunction<byte[], byte[]>)
+ bytes ->
+
executeRawSearchGroup(deserializeSplitGroup(bytes), metric))
+ .returns(PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO)
+ .name(operatorName)
+ .setParallelism(Math.min(parallelism, groups.size()));
+ }
+
+ static DataStream<Tuple2<Byte, byte[]>> addResultType(
+ DataStream<byte[]> results, byte resultType, String operatorName) {
+ return results.map(
+ (MapFunction<byte[], Tuple2<Byte, byte[]>>)
+ bytes -> Tuple2.of(resultType, bytes))
+ .returns(RESULT_TYPE_INFO)
+ .name(operatorName)
+ .setParallelism(results.getParallelism());
+ }
+
+ private DataStream<byte[]> taskSource(List<List<SerializedSplit>> groups,
String operatorName) {
+ List<byte[]> serializedGroups = new ArrayList<>(groups.size());
+ try {
+ for (List<SerializedSplit> group : groups) {
+ serializedGroups.add(InstantiationUtil.serializeObject(group));
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to serialize vector-search task
groups.", e);
+ }
+
+ return StreamExecutionEnvironmentUtils.fromData(
+ checkNotNull(env),
+ serializedGroups,
+ PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO)
+ .name(operatorName + " Source")
+ .setParallelism(1)
+ .rebalance();
+ }
+
+ private byte[] executeIndexSearchGroup(
+ List<SerializedSplit> group, int searchLimit, int parallelism) {
+ List<IndexVectorSearchSplit> splits = new ArrayList<>(group.size());
+ for (SerializedSplit serializedSplit : group) {
+ splits.add(deserializeSplit(serializedSplit.split));
+ }
+
+ GlobalIndexer taskGlobalIndexer = createGlobalIndexer(splits);
+ IndexPathFactory indexPathFactory =
table.store().pathFactory().globalIndexFileFactory();
+ ExecutorService executor =
+
GlobalIndexReadThreadPool.getExecutorService(Math.min(parallelism,
group.size()));
+
+ List<CompletableFuture<Optional<ScoredGlobalIndexResult>>> futures =
+ new ArrayList<>(group.size());
+ for (int i = 0; i < group.size(); i++) {
+ SerializedSplit serializedSplit = group.get(i);
+ IndexVectorSearchSplit split = splits.get(i);
+ futures.add(
+ eval(
+ taskGlobalIndexer,
+ indexPathFactory,
+ split.rowRangeStart(),
+ split.rowRangeEnd(),
+ split.vectorIndexFiles(),
+ vector,
+ searchLimit,
+ deserializePreFilter(serializedSplit.preFilter),
+ executor));
+ }
+ CompletableFuture.allOf(futures.toArray(new
CompletableFuture[0])).join();
+
+ ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
+ for (CompletableFuture<Optional<ScoredGlobalIndexResult>> future :
futures) {
+ Optional<ScoredGlobalIndexResult> next = future.join();
+ if (next.isPresent()) {
+ result = result.or(next.get());
+ }
+ }
+ return serializeResult(result.topK(searchLimit));
+ }
+
+ protected byte[] executeRawSearchGroup(List<SerializedSplit> group, String
metric) {
+ ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
+ for (SerializedSplit serializedSplit : group) {
+ ScoredGlobalIndexResult splitResult =
+ readRawSearch(
+ deserializeRanges(serializedSplit.split),
+ deserializePreFilter(serializedSplit.preFilter),
+ metric,
+ vector);
+ result = result.or(splitResult);
+ }
+ return serializeResult(result.topK(limit));
+ }
+
+ private List<SerializedSplit> deserializeSplitGroup(byte[] bytes) {
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException("Failed to deserialize vector-search
task group.", e);
+ }
+ }
+
+ private <T> List<T> collectResults(DataStream<T> results, String
operatorName) {
+ List<T> output = new ArrayList<>();
+ try (CloseableIterator<T> iterator =
+ results.executeAndCollect(flinkJobName(operatorName))) {
+ iterator.forEachRemaining(output::add);
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to execute " + operatorName +
".", e);
+ }
+ return output;
+ }
+
+ protected String flinkJobName(String operatorName) {
+ return "Vector Search - " + operatorName + " : " + table.fullName();
+ }
+
+ private byte[] serializeResult(ScoredGlobalIndexResult result) {
+ if (result.results().isEmpty()) {
+ return new byte[0];
+ }
+ try {
+ return new GlobalIndexResultSerializer().serialize(result);
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to serialize scored
vector-search result.", e);
+ }
+ }
+
+ private IndexVectorSearchSplit deserializeSplit(byte[] bytes) {
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException("Failed to deserialize vector-search
split.", e);
+ }
+ }
+
+ private List<Range> deserializeRanges(byte[] bytes) {
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException("Failed to deserialize raw vector row
ranges.", e);
+ }
+ }
+
+ @Nullable
+ private RoaringNavigableMap64 deserializePreFilter(@Nullable byte[] bytes)
{
+ if (bytes == null) {
+ return null;
+ }
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException("Failed to deserialize vector
pre-filter.", e);
+ }
+ }
+
+ private List<List<SerializedSplit>> splitGroups(
+ List<SerializedSplit> serializedSplits, int parallelism) {
+ List<List<SerializedSplit>> groups = new ArrayList<>(parallelism);
+ int groupSize = (serializedSplits.size() + parallelism - 1) /
parallelism;
+ for (int start = 0; start < serializedSplits.size(); start +=
groupSize) {
+ groups.add(
+ new ArrayList<>(
+ serializedSplits.subList(
+ start, Math.min(start + groupSize,
serializedSplits.size()))));
+ }
+ return groups;
+ }
+
+ private ScoredGlobalIndexResult mergeRemoteResults(List<byte[]>
remoteResults, int topK) {
+ ScoredGlobalIndexResult result = ScoredGlobalIndexResult.createEmpty();
+ GlobalIndexResultSerializer serializer = new
GlobalIndexResultSerializer();
+ for (byte[] bytes : remoteResults) {
+ if (bytes != null && bytes.length > 0) {
+ try {
+ result = result.or(serializer.deserialize(bytes));
+ } catch (IOException e) {
+ throw new RuntimeException(
+ "Failed to deserialize scored vector-search
result.", e);
+ }
+ }
+ }
+ return result.topK(topK);
+ }
+
+ private List<List<Range>> rangeGroups(List<Range> ranges, int parallelism)
{
+ long rowCount = rawRowCount(ranges);
+ int groupCount = (int) Math.min(parallelism, rowCount);
+ long targetRowsPerGroup = (rowCount - 1) / groupCount + 1;
+
+ List<List<Range>> groups = new ArrayList<>(groupCount);
+ List<Range> currentGroup = new ArrayList<>();
+ long currentRows = 0;
+ for (Range range : ranges) {
+ long from = range.from;
+ while (from <= range.to) {
+ if (currentRows == targetRowsPerGroup) {
+ groups.add(currentGroup);
+ currentGroup = new ArrayList<>();
+ currentRows = 0;
+ }
+ long remainingGroupRows = targetRowsPerGroup - currentRows;
+ long to = Math.min(range.to, from + remainingGroupRows - 1);
+ Range next = new Range(from, to);
+ currentGroup.add(next);
+ currentRows += next.count();
+ from = to + 1;
+ }
+ }
+ if (!currentGroup.isEmpty()) {
+ groups.add(currentGroup);
+ }
+ return groups;
+ }
+
+ private long rawRowCount(List<Range> ranges) {
+ long rowCount = 0;
+ for (Range range : ranges) {
+ long count = range.count();
+ if (Long.MAX_VALUE - rowCount < count) {
+ return Long.MAX_VALUE;
+ }
+ rowCount += count;
+ }
+ return rowCount;
+ }
+
+ static class SerializedSplit implements java.io.Serializable {
+
+ private static final long serialVersionUID = 1L;
+
+ private final byte[] split;
+ @Nullable private final byte[] preFilter;
+
+ private SerializedSplit(@Nullable byte[] split, @Nullable byte[]
preFilter) {
+ this.split = split;
+ this.preFilter = preFilter;
+ }
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkPrimaryKeyVectorRead.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkPrimaryKeyVectorRead.java
new file mode 100644
index 0000000000..613c59465b
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkPrimaryKeyVectorRead.java
@@ -0,0 +1,193 @@
+/*
+ * 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.flink.vectorsearch;
+
+import org.apache.paimon.flink.utils.StreamExecutionEnvironmentUtils;
+import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.predicate.Predicate;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.source.BucketVectorSearchSplit;
+import org.apache.paimon.table.source.PrimaryKeyVectorRead;
+import org.apache.paimon.table.source.PrimaryKeyVectorScan;
+import org.apache.paimon.table.source.VectorScan;
+import org.apache.paimon.types.DataField;
+import org.apache.paimon.utils.InstantiationUtil;
+
+import org.apache.flink.api.common.functions.MapFunction;
+import org.apache.flink.api.common.typeinfo.PrimitiveArrayTypeInfo;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.util.CloseableIterator;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM;
+import static org.apache.paimon.utils.Preconditions.checkNotNull;
+
+/** Flink-aware {@link PrimaryKeyVectorRead}. */
+public class FlinkPrimaryKeyVectorRead extends PrimaryKeyVectorRead {
+
+ private static final long serialVersionUID = 1L;
+
+ private final transient StreamExecutionEnvironment env;
+
+ public FlinkPrimaryKeyVectorRead(
+ FileStoreTable table,
+ DataField vectorField,
+ float[] query,
+ int limit,
+ Map<String, String> searchOptions,
+ @Nullable Predicate filter,
+ StreamExecutionEnvironment env) {
+ super(table, vectorField, query, limit, searchOptions, filter);
+ this.env = checkNotNull(env);
+ }
+
+ @Override
+ public GlobalIndexResult read(VectorScan.Plan plan) {
+ PrimaryKeyVectorScan.Plan primaryKeyPlan = primaryKeyPlan(plan);
+ List<BucketVectorSearchSplit> splits = bucketSplits(primaryKeyPlan);
+ int parallelism = flinkParallelism();
+ if (splits.size() < parallelism * 2L) {
+ return super.read(plan);
+ }
+
+ List<byte[]> serializedSplits = new ArrayList<>(splits.size());
+ for (BucketVectorSearchSplit split : splits) {
+ try {
+ serializedSplits.add(InstantiationUtil.serializeObject(split));
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to serialize primary-key
vector split.", e);
+ }
+ }
+
+ List<List<byte[]>> groups = splitGroups(serializedSplits, parallelism);
+ List<byte[]> groupResults = executeBucketSearchGroups(groups,
parallelism);
+ List<SearchResult> searchResults = new
ArrayList<>(groupResults.size());
+ for (byte[] groupResult : groupResults) {
+ searchResults.add(deserializeSearchResult(groupResult));
+ }
+ return createResult(primaryKeyPlan, mergeSearchResults(searchResults));
+ }
+
+ protected int flinkParallelism() {
+ int maxParallelism =
+ Math.max(1,
table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM));
+ int envParallelism = checkNotNull(env).getParallelism();
+ return envParallelism > 0
+ ? Math.max(1, Math.min(envParallelism, maxParallelism))
+ : maxParallelism;
+ }
+
+ protected List<byte[]> executeBucketSearchGroups(List<List<byte[]>>
groups, int parallelism) {
+ String operatorName = "Primary-Key Vector Search";
+ List<byte[]> serializedGroups = new ArrayList<>(groups.size());
+ try {
+ for (List<byte[]> group : groups) {
+ serializedGroups.add(InstantiationUtil.serializeObject(group));
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to serialize primary-key vector
task groups.", e);
+ }
+
+ DataStream<byte[]> results =
+ StreamExecutionEnvironmentUtils.fromData(
+ checkNotNull(env),
+ serializedGroups,
+
PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO)
+ .name(operatorName + " Source")
+ .setParallelism(1)
+ .rebalance()
+ .map(
+ (MapFunction<byte[], byte[]>)
+ bytes ->
executeBucketSearchGroup(deserializeGroup(bytes)))
+
.returns(PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO)
+ .name(operatorName)
+ .setParallelism(Math.min(parallelism, groups.size()));
+
+ List<byte[]> output = new ArrayList<>(groups.size());
+ try (CloseableIterator<byte[]> iterator =
+ results.executeAndCollect(flinkJobName(operatorName))) {
+ iterator.forEachRemaining(output::add);
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to execute " + operatorName +
".", e);
+ }
+ return output;
+ }
+
+ protected String flinkJobName(String operatorName) {
+ return "Vector Search - " + operatorName + " : " + table.fullName();
+ }
+
+ protected byte[] executeBucketSearchGroup(List<byte[]> group) {
+ List<BucketVectorSearchSplit> taskSplits = new
ArrayList<>(group.size());
+ for (byte[] bytes : group) {
+ taskSplits.add(deserializeSplit(bytes));
+ }
+ try {
+ return
InstantiationUtil.serializeObject(searchBuckets(taskSplits));
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to serialize primary-key vector
candidates.", e);
+ }
+ }
+
+ private List<byte[]> deserializeGroup(byte[] bytes) {
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException("Failed to deserialize primary-key
vector task group.", e);
+ }
+ }
+
+ private List<List<byte[]>> splitGroups(List<byte[]> splits, int
parallelism) {
+ List<List<byte[]>> groups = new ArrayList<>(parallelism);
+ int groupSize = (splits.size() + parallelism - 1) / parallelism;
+ for (int start = 0; start < splits.size(); start += groupSize) {
+ groups.add(
+ new ArrayList<>(
+ splits.subList(start, Math.min(start + groupSize,
splits.size()))));
+ }
+ return groups;
+ }
+
+ private BucketVectorSearchSplit deserializeSplit(byte[] bytes) {
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException("Failed to deserialize primary-key
vector split.", e);
+ }
+ }
+
+ private SearchResult deserializeSearchResult(byte[] bytes) {
+ try {
+ return InstantiationUtil.deserializeObject(
+ bytes, Thread.currentThread().getContextClassLoader());
+ } catch (IOException | ClassNotFoundException e) {
+ throw new RuntimeException(
+ "Failed to deserialize primary-key vector search result.",
e);
+ }
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkVectorSearchBuilderImpl.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkVectorSearchBuilderImpl.java
new file mode 100644
index 0000000000..cad0f86835
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/vectorsearch/FlinkVectorSearchBuilderImpl.java
@@ -0,0 +1,51 @@
+/*
+ * 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.flink.vectorsearch;
+
+import org.apache.paimon.table.InnerTable;
+import org.apache.paimon.table.source.VectorRead;
+import org.apache.paimon.table.source.VectorSearchBuilderImpl;
+
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+
+import static org.apache.paimon.utils.Preconditions.checkNotNull;
+
+/** Flink-aware vector-search builder which creates distributed Flink readers.
*/
+public class FlinkVectorSearchBuilderImpl extends VectorSearchBuilderImpl {
+
+ private static final long serialVersionUID = 1L;
+
+ private final transient StreamExecutionEnvironment env;
+
+ public FlinkVectorSearchBuilderImpl(InnerTable table,
StreamExecutionEnvironment env) {
+ super(table);
+ this.env = checkNotNull(env);
+ }
+
+ @Override
+ public VectorRead newVectorRead() {
+ checkNotNull(vector, "vector must be set via withVector()");
+ if (isPrimaryKeyVectorSearch()) {
+ return new FlinkPrimaryKeyVectorRead(
+ table, vectorColumn, vector, limit, options, filter, env);
+ }
+ return new FlinkDataEvolutionVectorRead(
+ table, partitionFilter, filter, limit, vectorColumn, vector,
options, env);
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/procedure/VectorSearchProcedureITCase.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/procedure/VectorSearchProcedureITCase.java
index 74ec37d14b..9fb817ba10 100644
---
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/procedure/VectorSearchProcedureITCase.java
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/procedure/VectorSearchProcedureITCase.java
@@ -19,10 +19,12 @@
package org.apache.paimon.flink.procedure;
import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.GenericArray;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.flink.CatalogITCaseBase;
import org.apache.paimon.globalindex.GlobalIndexBuilderUtils;
+import org.apache.paimon.globalindex.GlobalIndexMultiColumnWriter;
import org.apache.paimon.globalindex.GlobalIndexSingleColumnWriter;
import org.apache.paimon.globalindex.ResultEntry;
import org.apache.paimon.globalindex.testvector.TestVectorGlobalIndexer;
@@ -31,6 +33,8 @@ import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.io.CompactIncrement;
import org.apache.paimon.io.DataIncrement;
import org.apache.paimon.options.Options;
+import org.apache.paimon.predicate.Predicate;
+import org.apache.paimon.predicate.PredicateBuilder;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.PostponeUtils;
import org.apache.paimon.table.sink.BatchTableCommit;
@@ -38,6 +42,9 @@ import org.apache.paimon.table.sink.BatchTableWrite;
import org.apache.paimon.table.sink.BatchWriteBuilder;
import org.apache.paimon.table.sink.CommitMessage;
import org.apache.paimon.table.sink.CommitMessageImpl;
+import org.apache.paimon.table.source.IndexVectorSearchSplit;
+import org.apache.paimon.table.source.RawVectorSearchSplit;
+import org.apache.paimon.table.source.VectorScan;
import org.apache.paimon.types.DataField;
import org.apache.paimon.utils.Range;
@@ -45,10 +52,14 @@ import org.apache.flink.table.api.config.TableConfigOptions;
import org.apache.flink.types.Row;
import org.junit.jupiter.api.Test;
+import java.util.Arrays;
import java.util.Collections;
import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** IT cases for {@link VectorSearchProcedure}. */
public class VectorSearchProcedureITCase extends CatalogITCaseBase {
@@ -73,6 +84,155 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
.containsExactlyInAnyOrder("{\"id\":\"2\"}", "{\"id\":\"3\"}");
}
+ @Test
+ public void testPrimaryKeyVectorSearchWithSnapshot() throws Exception {
+ createPrimaryKeyVectorTable("PK_SNAPSHOT_T");
+ sql(
+ "INSERT INTO PK_SNAPSHOT_T VALUES "
+ + "(1, ARRAY[CAST(1.0 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+ long snapshotId =
paimonTable("PK_SNAPSHOT_T").latestSnapshot().get().id();
+ sql(
+ "INSERT INTO PK_SNAPSHOT_T VALUES "
+ + "(2, ARRAY[CAST(0.0 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+
+ assertThat(search("PK_SNAPSHOT_T", "0.0,0.0", 2, "id",
"scan.snapshot-id=" + snapshotId))
+ .extracting(row -> row.getField(0).toString())
+ .containsExactly("{\"id\":\"1\"}");
+ }
+
+ @Test
+ public void testEmptyTableVectorSearch() {
+ createVectorTable("EMPTY_T");
+ createPrimaryKeyVectorTable("EMPTY_PK_T");
+
+ assertThat(search("EMPTY_T", "0.0,0.0", 1, "id", null)).isEmpty();
+ assertThat(searchPrimaryKeyVectorTable("EMPTY_PK_T", 1,
"id")).isEmpty();
+ }
+
+ @Test
+ public void testVectorSearchKeepsResolvedSchema() throws Exception {
+ createVectorTable("SCHEMA_EVOLUTION_T", "'global-index.search-mode' =
'full'");
+ FileStoreTable table = paimonTable("SCHEMA_EVOLUTION_T");
+ float[][] vectors = {{0.0f, 0.0f}, {1.0f, 0.0f}};
+ writeVectors(table, vectors);
+ buildAndCommitVectorIndex(table, vectors);
+
+ sql("ALTER TABLE SCHEMA_EVOLUTION_T ADD payload STRING FIRST");
+ FileStoreTable evolvedTable = paimonTable("SCHEMA_EVOLUTION_T");
+ assertThat(evolvedTable.schema().id())
+ .isNotEqualTo(evolvedTable.latestSnapshot().get().schemaId());
+
+ List<String> defaultProjection =
+ sql(
+ "CALL sys.vector_search("
+ + "`table` =>
'default.SCHEMA_EVOLUTION_T', "
+ + "vector_column => 'vec', "
+ + "query_vector => '0.0,0.0', "
+ + "top_k => 1)")
+ .stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+ assertThat(defaultProjection).hasSize(1);
+ assertThat(defaultProjection.get(0))
+ .contains("\"payload\":null", "\"id\":\"0\"", "\"vec\"");
+
+ assertThat(search("SCHEMA_EVOLUTION_T", "0.0,0.0", 1, "payload,id",
null))
+ .extracting(row -> row.getField(0).toString())
+ .containsExactly("{\"payload\":null,\"id\":\"0\"}");
+
+ List<String> filtered =
+ sql(
+ "CALL sys.vector_search("
+ + "`table` =>
'default.SCHEMA_EVOLUTION_T', "
+ + "vector_column => 'vec', "
+ + "query_vector => '0.0,0.0', "
+ + "top_k => 1, "
+ + "projection => 'payload,id', "
+ + "`where` => 'payload IS NULL AND id
= 1')")
+ .stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+
assertThat(filtered).containsExactly("{\"payload\":null,\"id\":\"1\"}");
+ }
+
+ @Test
+ public void testDistributedPrimaryKeyVectorSearch() {
+ createPrimaryKeyVectorTable(
+ "DISTRIBUTED_PK_T",
+ 4,
+ "'vector-search.distribute.enabled' = 'true', "
+ + "'global-index.thread-num' = '2'");
+ String values =
+ IntStream.range(0, 32)
+ .mapToObj(
+ i ->
+ String.format(
+ "(%d, ARRAY[CAST(%d.0 AS
FLOAT), CAST(0.0 AS FLOAT)])",
+ i, i))
+ .collect(Collectors.joining(","));
+ sql("INSERT INTO DISTRIBUTED_PK_T VALUES " + values);
+ assertThat(sql("SELECT DISTINCT bucket FROM
`DISTRIBUTED_PK_T$files`")).hasSize(4);
+
+ List<String> rows =
+ searchPrimaryKeyVectorTable("DISTRIBUTED_PK_T", 3,
"id,__paimon_search_score")
+ .stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+
+ assertThat(rows).hasSize(3);
+ assertThat(rows.get(0)).contains("\"id\":\"0\"");
+ assertThat(rows.get(1)).contains("\"id\":\"1\"");
+ assertThat(rows.get(2)).contains("\"id\":\"2\"");
+ assertThat(rows).allMatch(row ->
row.contains("\"__paimon_search_score\""));
+
+ List<String> scoreOnlyRows =
+ searchPrimaryKeyVectorTable("DISTRIBUTED_PK_T", 2,
"__paimon_search_score").stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+ assertThat(scoreOnlyRows).hasSize(2);
+ assertThat(scoreOnlyRows)
+ .allMatch(
+ row ->
+ row.startsWith("{\"__paimon_search_score\":")
+ && row.endsWith("}")
+ && !row.contains("\"id\"")
+ && !row.contains("\"vec\""));
+ }
+
+ @Test
+ public void testDistributedPrimaryKeyVectorSearchWithResidualFilter() {
+ createPrimaryKeyVectorTableWithPayload(
+ "DISTRIBUTED_FILTERED_PK_T",
+ 4,
+ "'vector-search.distribute.enabled' = 'true', "
+ + "'global-index.thread-num' = '2'");
+ String values =
+ IntStream.rangeClosed(1, 16)
+ .mapToObj(
+ i ->
+ String.format(
+ "(%d, '%s', ARRAY[CAST(%d.0 AS
FLOAT), CAST(0.0 AS FLOAT)])",
+ i, i == 1 ? "drop" : "keep",
i))
+ .collect(Collectors.joining(","));
+ sql("INSERT INTO DISTRIBUTED_FILTERED_PK_T VALUES " + values);
+ assertThat(sql("SELECT DISTINCT bucket FROM
`DISTRIBUTED_FILTERED_PK_T$files`")).hasSize(4);
+
+ List<String> rows =
+ sql(
+ "CALL sys.vector_search("
+ + "`table` =>
'default.DISTRIBUTED_FILTERED_PK_T', "
+ + "vector_column => 'vec', "
+ + "query_vector => '0.0,0.0', "
+ + "top_k => 2, "
+ + "projection => 'id', "
+ + "`where` => 'payload = ''keep''')")
+ .stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+
+ assertThat(rows).containsExactly("{\"id\":\"2\"}", "{\"id\":\"3\"}");
+ }
+
@Test
public void testPostponeBucketBuildsVectorIndexDuringCompact() throws
Exception {
createPostponePrimaryKeyVectorTable("POSTPONE_PK_T");
@@ -303,6 +463,336 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
assertThat(result.size()).isLessThanOrEqualTo(2);
}
+ @Test
+ public void testDistributedVectorSearchWithFilterAndScore() throws
Exception {
+ createVectorTable(
+ "DISTRIBUTED_T",
+ "'vector-search.distribute.enabled' = 'true', "
+ + "'global-index.thread-num' = '2', "
+ + "'global-index.search-mode' = 'full'");
+ FileStoreTable table = paimonTable("DISTRIBUTED_T");
+ float[][] vectors = {
+ {0.0f, 0.0f},
+ {1.0f, 0.0f},
+ {2.0f, 0.0f},
+ {3.0f, 0.0f},
+ {4.0f, 0.0f},
+ {5.0f, 0.0f}
+ };
+ writeVectors(table, vectors);
+
+ List<String> rows =
+ sql(
+ "CALL sys.vector_search("
+ + "`table` => 'default.DISTRIBUTED_T',
"
+ + "vector_column => 'vec', "
+ + "query_vector => '0.0,0.0', "
+ + "top_k => 2, "
+ + "projection =>
'vec,__paimon_search_score', "
+ + "`where` => 'id >= 3')")
+ .stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+
+ assertThat(rows).hasSize(2);
+ assertThat(rows.get(0)).contains("3.0");
+ assertThat(rows.get(1)).contains("4.0");
+ assertThat(rows)
+ .allMatch(
+ row ->
+ row.contains("\"vec\"")
+ &&
row.contains("\"__paimon_search_score\"")
+ && !row.contains("\"id\""));
+
+ List<String> scoreOnlyRows =
+ search("DISTRIBUTED_T", "0.0,0.0", 2, "__paimon_search_score",
null).stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+ assertThat(scoreOnlyRows).hasSize(2);
+ assertThat(scoreOnlyRows)
+ .allMatch(
+ row ->
+ row.startsWith("{\"__paimon_search_score\":")
+ && row.endsWith("}")
+ && !row.contains("\"id\"")
+ && !row.contains("\"vec\""));
+ }
+
+ @Test
+ public void testDistributedMultiFieldVectorIndex() throws Exception {
+ createVectorTable(
+ "DISTRIBUTED_MULTI_FIELD_T",
+ "'vector-search.distribute.enabled' = 'true', "
+ + "'global-index.thread-num' = '2', "
+ + "'global-index.search-mode' = 'full'");
+ FileStoreTable table = paimonTable("DISTRIBUTED_MULTI_FIELD_T");
+ float[][] vectors = {
+ {0.0f, 0.0f},
+ {1.0f, 0.0f},
+ {2.0f, 0.0f},
+ {3.0f, 0.0f},
+ {4.0f, 0.0f},
+ {5.0f, 0.0f},
+ {6.0f, 0.0f},
+ {7.0f, 0.0f}
+ };
+ writeVectors(table, vectors);
+ for (int start = 0; start < vectors.length; start += 2) {
+ buildAndCommitMultiFieldVectorIndex(
+ table,
+ new float[][] {vectors[start], vectors[start + 1]},
+ new Range(start, start + 1));
+ }
+
+ VectorScan.Plan plan =
+ table.newVectorSearchBuilder()
+ .withVector(new float[] {0.0f, 0.0f})
+ .withLimit(2)
+ .withVectorColumn(VECTOR_FIELD)
+ .newVectorScan()
+ .scan();
+ List<IndexVectorSearchSplit> indexSplits =
+ plan.splits().stream()
+ .filter(IndexVectorSearchSplit.class::isInstance)
+ .map(IndexVectorSearchSplit.class::cast)
+ .collect(Collectors.toList());
+ assertThat(indexSplits).hasSize(4);
+ assertThat(indexSplits).allMatch(split ->
split.vectorIndexFiles().size() == 2);
+
+ List<String> rows =
+ sql(
+ "CALL sys.vector_search("
+ + "`table` =>
'default.DISTRIBUTED_MULTI_FIELD_T', "
+ + "vector_column => 'vec', "
+ + "query_vector => '0.0,0.0', "
+ + "top_k => 2, "
+ + "projection => 'id', "
+ + "`where` => 'id >= 5')")
+ .stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+
+ assertThat(rows).containsExactly("{\"id\":\"5\"}", "{\"id\":\"6\"}");
+ }
+
+ @Test
+ public void testDistributedMixedIndexedAndRawVectorSearch() throws
Exception {
+ createVectorTable(
+ "DISTRIBUTED_MIXED_T",
+ "'vector-search.distribute.enabled' = 'true', "
+ + "'global-index.thread-num' = '2', "
+ + "'global-index.search-mode' = 'full'");
+ FileStoreTable table = paimonTable("DISTRIBUTED_MIXED_T");
+ float[][] vectors = {
+ {20.0f, 0.0f},
+ {21.0f, 0.0f},
+ {22.0f, 0.0f},
+ {23.0f, 0.0f},
+ {24.0f, 0.0f},
+ {25.0f, 0.0f},
+ {26.0f, 0.0f},
+ {27.0f, 0.0f},
+ {0.0f, 0.0f},
+ {1.0f, 0.0f},
+ {2.0f, 0.0f},
+ {3.0f, 0.0f}
+ };
+ writeVectors(table, vectors);
+ for (int start = 0; start < 8; start += 2) {
+ buildAndCommitVectorIndex(
+ table,
+ new float[][] {vectors[start], vectors[start + 1]},
+ new Range(start, start + 1));
+ }
+
+ VectorScan.Plan plan =
+ table.newVectorSearchBuilder()
+ .withVector(new float[] {0.0f, 0.0f})
+ .withLimit(3)
+ .withVectorColumn(VECTOR_FIELD)
+ .newVectorScan()
+ .scan();
+
assertThat(plan.splits().stream().filter(IndexVectorSearchSplit.class::isInstance))
+ .hasSize(4);
+
assertThat(plan.splits().stream().filter(RawVectorSearchSplit.class::isInstance))
+ .hasSize(1);
+
+ List<String> rows =
+ search("DISTRIBUTED_MIXED_T", "0.0,0.0", 3, "id",
null).stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+
+ assertThat(rows).containsExactly("{\"id\":\"8\"}", "{\"id\":\"9\"}",
"{\"id\":\"10\"}");
+ }
+
+ @Test
+ public void testDistributedMixedIndexedAndRawVectorSearchWithFilter()
throws Exception {
+ createVectorTable(
+ "DISTRIBUTED_FILTERED_MIXED_T",
+ "'vector-search.distribute.enabled' = 'true', "
+ + "'global-index.thread-num' = '2', "
+ + "'global-index.search-mode' = 'full'");
+ FileStoreTable table = paimonTable("DISTRIBUTED_FILTERED_MIXED_T");
+ float[][] vectors = {
+ {0.0f, 0.0f},
+ {0.1f, 0.0f},
+ {0.2f, 0.0f},
+ {0.3f, 0.0f},
+ {0.4f, 0.0f},
+ {1.0f, 0.0f},
+ {4.0f, 0.0f},
+ {5.0f, 0.0f},
+ {0.5f, 0.0f},
+ {2.0f, 0.0f},
+ {0.05f, 0.0f},
+ {0.15f, 0.0f}
+ };
+ writeVectors(table, vectors);
+ for (int start = 0; start < 8; start += 2) {
+ buildAndCommitMultiFieldVectorIndex(
+ table,
+ new float[][] {vectors[start], vectors[start + 1]},
+ new Range(start, start + 1));
+ }
+
+ PredicateBuilder predicateBuilder = new
PredicateBuilder(table.rowType());
+ Predicate filter =
+ PredicateBuilder.and(
+ predicateBuilder.greaterOrEqual(0, 5),
predicateBuilder.lessThan(0, 10));
+ VectorScan.Plan plan =
+ table.newVectorSearchBuilder()
+ .withVector(new float[] {0.0f, 0.0f})
+ .withLimit(3)
+ .withVectorColumn(VECTOR_FIELD)
+ .withFilter(filter)
+ .newVectorScan()
+ .scan();
+ List<IndexVectorSearchSplit> indexSplits =
+ plan.splits().stream()
+ .filter(IndexVectorSearchSplit.class::isInstance)
+ .map(IndexVectorSearchSplit.class::cast)
+ .collect(Collectors.toList());
+ assertThat(indexSplits).hasSize(4);
+ assertThat(indexSplits)
+ .allSatisfy(
+ split -> {
+ assertThat(split.vectorIndexFiles()).hasSize(2);
+ assertThat(split.scalarIndexFiles())
+
.containsExactlyElementsOf(split.vectorIndexFiles());
+ });
+ List<RawVectorSearchSplit> rawSplits =
+ plan.splits().stream()
+ .filter(RawVectorSearchSplit.class::isInstance)
+ .map(RawVectorSearchSplit.class::cast)
+ .collect(Collectors.toList());
+ assertThat(rawSplits).hasSize(1);
+ assertThat(rawSplits.get(0).rowRanges()).containsExactly(new Range(8,
11));
+ assertThat(rawSplits.get(0).scalarIndexFiles()).isEmpty();
+
+ List<String> rows =
+ sql(
+ "CALL sys.vector_search("
+ + "`table` =>
'default.DISTRIBUTED_FILTERED_MIXED_T', "
+ + "vector_column => 'vec', "
+ + "query_vector => '0.0,0.0', "
+ + "top_k => 3, "
+ + "projection => 'id', "
+ + "`where` => 'id >= 5 AND id < 10')")
+ .stream()
+ .map(row -> row.getField(0).toString())
+ .collect(Collectors.toList());
+
+ assertThat(rows).containsExactly("{\"id\":\"8\"}", "{\"id\":\"5\"}",
"{\"id\":\"9\"}");
+ }
+
+ @Test
+ public void testVectorSearchWithPartitions() throws Exception {
+ createPartitionedVectorTable("PARTITIONED_T");
+ FileStoreTable table = paimonTable("PARTITIONED_T");
+ writePartitionedVectors(table);
+
+ List<Row> result = searchWithFilters("PARTITIONED_T", 1, null, "pt=b");
+
+ assertThat(result)
+ .extracting(row -> row.getField(0).toString())
+ .containsExactly("{\"id\":\"2\"}");
+
+ assertThat(searchWithFilters("PARTITIONED_T", 1, "pt = 'a'",
"pt=b")).isEmpty();
+ assertThat(searchWithFilters("PARTITIONED_T", 1, "pt = 'b'", "pt=b"))
+ .extracting(row -> row.getField(0).toString())
+ .containsExactly("{\"id\":\"2\"}");
+ }
+
+ @Test
+ public void testVectorSearchRejectsInvalidPartitions() {
+ createPartitionedVectorTable("INVALID_PARTITIONS_T");
+
+ assertThat(searchWithFilters("INVALID_PARTITIONS_T", 1, null, "
")).isEmpty();
+ assertThat(searchWithFilters("INVALID_PARTITIONS_T", 1, null,
"pt=")).isEmpty();
+ String separator = ";";
+ for (String partitions :
+ Arrays.asList(separator + "pt=a", "pt=a" + separator +
separator + "pt=b")) {
+ assertThatThrownBy(() -> searchWithFilters("INVALID_PARTITIONS_T",
1, null, partitions))
+ .hasStackTraceContaining("must exactly match partition
keys");
+ }
+ for (String partitions :
+ Arrays.asList(separator, separator + separator, "pt=a" +
separator)) {
+ assertThat(searchWithFilters("INVALID_PARTITIONS_T", 1, null,
partitions)).isEmpty();
+ }
+ assertThat(searchWithFilters("INVALID_PARTITIONS_T", 1, null,
"pt=a,pt=b")).isEmpty();
+ }
+
+ @Test
+ public void testFirstRowPrimaryKeyVectorSearchWhereBoundary() {
+ createFirstRowPartitionedPrimaryKeyVectorTable("FIRST_ROW_PK_T");
+ sql(
+ "INSERT INTO FIRST_ROW_PK_T VALUES "
+ + "(1, 'a', ARRAY[CAST(1.0 AS FLOAT), CAST(0.0 AS
FLOAT)]), "
+ + "(2, 'b', ARRAY[CAST(2.0 AS FLOAT), CAST(0.0 AS
FLOAT)]), "
+ + "(3, 'b', ARRAY[CAST(3.0 AS FLOAT), CAST(0.0 AS
FLOAT)])");
+
+ assertThat(searchWithFilters("FIRST_ROW_PK_T", 1, "pt = 'b'", null))
+ .extracting(row -> row.getField(0).toString())
+ .containsExactly("{\"id\":\"2\"}");
+ assertThatThrownBy(() -> searchWithFilters("FIRST_ROW_PK_T", 1, "id >=
2", null))
+ .hasStackTraceContaining(
+ "Primary-key vector where filter on non-partition
columns requires")
+ .hasStackTraceContaining("deletion-vectors.enabled = true");
+ }
+
+ @Test
+ public void testVectorSearchValidation() {
+ createVectorTable("VALIDATION_T");
+ createVectorTable("QUERY_AUTH_T", "'query-auth.enabled' = 'true'");
+ createVectorTable(
+ "DISTRIBUTED_QUERY_AUTH_T",
+ "'query-auth.enabled' = 'true', " +
"'vector-search.distribute.enabled' = 'true'");
+
+ assertThatThrownBy(() -> search("VALIDATION_T", "0.0,0.0", 0, "id",
null))
+ .hasStackTraceContaining("top_k must be positive");
+ assertThatThrownBy(() -> search("VALIDATION_T", "NaN,0.0", 1, "id",
null))
+ .hasStackTraceContaining("query_vector values must be finite");
+ assertThatThrownBy(() -> search("VALIDATION_T", "0.0,0.0", 1,
"missing", null))
+ .hasStackTraceContaining("Unknown projection column");
+ assertThatThrownBy(() -> search("VALIDATION_T", "0.0,0.0", 1, "id,id",
null))
+ .hasStackTraceContaining("Duplicate projection column");
+ assertThatThrownBy(
+ () ->
+ search(
+ "VALIDATION_T",
+ "0.0,0.0",
+ 1,
+ "id",
+ "query-auth.enabled=false"))
+ .hasStackTraceContaining("Option 'query-auth.enabled' is not
allowed");
+ assertThatThrownBy(() -> search("QUERY_AUTH_T", "0.0,0.0", 1, "id",
null))
+ .hasStackTraceContaining(
+ "Vector search does not support tables with query auth
enabled");
+ assertThatThrownBy(() -> search("DISTRIBUTED_QUERY_AUTH_T", "0.0,0.0",
1, "id", null))
+ .hasStackTraceContaining(
+ "Vector search does not support tables with query auth
enabled");
+ }
+
private void createVectorTable(String tableName) {
createVectorTable(tableName, "");
}
@@ -325,13 +815,18 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
}
private void createPrimaryKeyVectorTable(String tableName) {
+ createPrimaryKeyVectorTable(tableName, 2, "");
+ }
+
+ private void createPrimaryKeyVectorTable(String tableName, int bucket,
String extraOptions) {
+ String formattedExtraOptions = extraOptions.isEmpty() ? "" : ", " +
extraOptions;
sql(
"CREATE TABLE %s ("
+ "id INT, "
+ "vec ARRAY<FLOAT>, "
+ "PRIMARY KEY (id) NOT ENFORCED"
+ ") WITH ("
- + "'bucket' = '2', "
+ + "'bucket' = '%d', "
+ "'file.format' = 'json', "
+ "'file.compression' = 'none', "
+ "'deletion-vectors.enabled' = 'true', "
@@ -342,10 +837,87 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
+ "'fields.vec.pk-vector.distance.metric' = 'l2', "
+ "'test.vector.dimension' = '%d', "
+ "'test.vector.metric' = 'l2'"
+ + "%s"
+ + ")",
+ tableName,
+ bucket,
+ DIMENSION,
+ TestVectorGlobalIndexerFactory.IDENTIFIER,
+ DIMENSION,
+ formattedExtraOptions);
+ }
+
+ private void createPrimaryKeyVectorTableWithPayload(
+ String tableName, int bucket, String extraOptions) {
+ String formattedExtraOptions = extraOptions.isEmpty() ? "" : ", " +
extraOptions;
+ sql(
+ "CREATE TABLE %s ("
+ + "id INT, "
+ + "payload STRING, "
+ + "vec ARRAY<FLOAT>, "
+ + "PRIMARY KEY (id) NOT ENFORCED"
+ + ") WITH ("
+ + "'bucket' = '%d', "
+ + "'file.format' = 'json', "
+ + "'file.compression' = 'none', "
+ + "'deletion-vectors.enabled' = 'true', "
+ + "'vector-field' = 'vec', "
+ + "'field.vec.vector-dim' = '%d', "
+ + "'pk-vector.index.columns' = 'vec', "
+ + "'fields.vec.pk-vector.index.type' = '%s', "
+ + "'fields.vec.pk-vector.distance.metric' = 'l2', "
+ + "'test.vector.dimension' = '%d', "
+ + "'test.vector.metric' = 'l2'"
+ + "%s"
+ + ")",
+ tableName,
+ bucket,
+ DIMENSION,
+ TestVectorGlobalIndexerFactory.IDENTIFIER,
+ DIMENSION,
+ formattedExtraOptions);
+ }
+
+ private void createFirstRowPartitionedPrimaryKeyVectorTable(String
tableName) {
+ sql(
+ "CREATE TABLE %s ("
+ + "id INT, "
+ + "pt STRING, "
+ + "vec ARRAY<FLOAT>, "
+ + "PRIMARY KEY (id, pt) NOT ENFORCED"
+ + ") PARTITIONED BY (pt) WITH ("
+ + "'bucket' = '2', "
+ + "'merge-engine' = 'first-row', "
+ + "'file.format' = 'json', "
+ + "'file.compression' = 'none', "
+ + "'vector-field' = 'vec', "
+ + "'field.vec.vector-dim' = '%d', "
+ + "'pk-vector.index.columns' = 'vec', "
+ + "'fields.vec.pk-vector.index.type' = '%s', "
+ + "'fields.vec.pk-vector.distance.metric' = 'l2', "
+ + "'test.vector.dimension' = '%d', "
+ + "'test.vector.metric' = 'l2'"
+ ")",
tableName, DIMENSION,
TestVectorGlobalIndexerFactory.IDENTIFIER, DIMENSION);
}
+ private void createPartitionedVectorTable(String tableName) {
+ sql(
+ "CREATE TABLE %s ("
+ + "id INT, "
+ + "pt STRING, "
+ + "vec ARRAY<FLOAT>"
+ + ") PARTITIONED BY (pt) WITH ("
+ + "'bucket' = '-1', "
+ + "'row-tracking.enabled' = 'true', "
+ + "'data-evolution.enabled' = 'true', "
+ + "'global-index.search-mode' = 'full', "
+ + "'test.vector.dimension' = '%d', "
+ + "'test.vector.metric' = 'l2'"
+ + ")",
+ tableName, DIMENSION);
+ }
+
private void createPostponePrimaryKeyVectorTable(String tableName) {
sql(
"CREATE TABLE %s ("
@@ -382,6 +954,35 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
tableName, topK, projection);
}
+ private List<Row> search(
+ String tableName, String queryVector, int topK, String projection,
String options) {
+ String optionsArgument = options == null ? "" : ", options => '" +
options + "'";
+ return sql(
+ "CALL sys.vector_search("
+ + "`table` => 'default.%s', "
+ + "vector_column => 'vec', "
+ + "query_vector => '%s', "
+ + "top_k => %d, "
+ + "projection => '%s'%s)",
+ tableName, queryVector, topK, projection, optionsArgument);
+ }
+
+ private List<Row> searchWithFilters(
+ String tableName, int topK, String where, String partitions) {
+ String whereArgument =
+ where == null ? "" : ", `where` => '" + where.replace("'",
"''") + "'";
+ String partitionsArgument =
+ partitions == null ? "" : ", partitions => '" +
partitions.replace("'", "''") + "'";
+ return sql(
+ "CALL sys.vector_search("
+ + "`table` => 'default.%s', "
+ + "vector_column => 'vec', "
+ + "query_vector => '0.0,0.0', "
+ + "top_k => %d, "
+ + "projection => 'id'%s%s)",
+ tableName, topK, whereArgument, partitionsArgument);
+ }
+
private void createPartialUpdatePrimaryKeyVectorTable(String tableName) {
sql(
"CREATE TABLE %s ("
@@ -418,8 +1019,36 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
}
}
+ private void writePartitionedVectors(FileStoreTable table) throws
Exception {
+ BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder();
+ try (BatchTableWrite write = writeBuilder.newWrite();
+ BatchTableCommit commit = writeBuilder.newCommit()) {
+ write.write(
+ GenericRow.of(
+ 1,
+ BinaryString.fromString("a"),
+ new GenericArray(new float[] {0.0f, 0.0f})));
+ write.write(
+ GenericRow.of(
+ 2,
+ BinaryString.fromString("b"),
+ new GenericArray(new float[] {1.0f, 0.0f})));
+ write.write(
+ GenericRow.of(
+ 3,
+ BinaryString.fromString("b"),
+ new GenericArray(new float[] {2.0f, 0.0f})));
+ commit.commit(write.prepareCommit());
+ }
+ }
+
private void buildAndCommitVectorIndex(FileStoreTable table, float[][]
vectors)
throws Exception {
+ buildAndCommitVectorIndex(table, vectors, new Range(0, vectors.length
- 1));
+ }
+
+ private void buildAndCommitVectorIndex(FileStoreTable table, float[][]
vectors, Range rowRange)
+ throws Exception {
Options options = table.coreOptions().toConfiguration();
DataField vectorField = table.rowType().getField(VECTOR_FIELD);
@@ -435,7 +1064,6 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
}
List<ResultEntry> entries = writer.finish();
- Range rowRange = new Range(0, vectors.length - 1);
List<IndexFileMeta> indexFiles =
GlobalIndexBuilderUtils.toIndexFileMetas(
table.fileIO(),
@@ -458,4 +1086,46 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
commit.commit(Collections.singletonList(message));
}
}
+
+ private void buildAndCommitMultiFieldVectorIndex(
+ FileStoreTable table, float[][] vectors, Range rowRange) throws
Exception {
+ Options options = table.coreOptions().toConfiguration();
+ DataField vectorField = table.rowType().getField(VECTOR_FIELD);
+ DataField idField = table.rowType().getField("id");
+
+ GlobalIndexMultiColumnWriter writer =
+ (GlobalIndexMultiColumnWriter)
+ GlobalIndexBuilderUtils.createIndexWriter(
+ table,
+ TestVectorGlobalIndexerFactory.IDENTIFIER,
+ vectorField,
+ Collections.singletonList(idField),
+ options);
+ for (int i = 0; i < vectors.length; i++) {
+ writer.write(i, GenericRow.of(new GenericArray(vectors[i]), (int)
(rowRange.from + i)));
+ }
+ List<ResultEntry> entries = writer.finish();
+
+ List<IndexFileMeta> indexFiles =
+ GlobalIndexBuilderUtils.toIndexFileMetas(
+ table.fileIO(),
+ table.store().pathFactory().globalIndexFileFactory(),
+ table.coreOptions(),
+ rowRange,
+ Arrays.asList(vectorField, idField),
+ TestVectorGlobalIndexerFactory.IDENTIFIER,
+ entries);
+
+ DataIncrement dataIncrement = DataIncrement.indexIncrement(indexFiles);
+ CommitMessage message =
+ new CommitMessageImpl(
+ BinaryRow.EMPTY_ROW,
+ 0,
+ null,
+ dataIncrement,
+ CompactIncrement.emptyIncrement());
+ try (BatchTableCommit commit =
table.newBatchWriteBuilder().newCommit()) {
+ commit.commit(Collections.singletonList(message));
+ }
+ }
}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/vectorsearch/FlinkDataEvolutionVectorReadTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/vectorsearch/FlinkDataEvolutionVectorReadTest.java
new file mode 100644
index 0000000000..9bd45a7776
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/vectorsearch/FlinkDataEvolutionVectorReadTest.java
@@ -0,0 +1,555 @@
+/*
+ * 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.flink.vectorsearch;
+
+import org.apache.paimon.flink.utils.StreamExecutionEnvironmentUtils;
+import org.apache.paimon.globalindex.GlobalIndexIOMeta;
+import org.apache.paimon.globalindex.GlobalIndexReader;
+import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.globalindex.GlobalIndexResultSerializer;
+import org.apache.paimon.globalindex.GlobalIndexWriter;
+import org.apache.paimon.globalindex.GlobalIndexer;
+import org.apache.paimon.globalindex.ScoredGlobalIndexResult;
+import org.apache.paimon.globalindex.VectorGlobalIndexer;
+import org.apache.paimon.globalindex.io.GlobalIndexFileReader;
+import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
+import org.apache.paimon.index.GlobalIndexMeta;
+import org.apache.paimon.index.IndexFileMeta;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.source.IndexVectorSearchSplit;
+import org.apache.paimon.table.source.RawVectorSearchSplit;
+import org.apache.paimon.table.source.VectorScan;
+import org.apache.paimon.table.source.VectorSearchSplit;
+import org.apache.paimon.types.ArrayType;
+import org.apache.paimon.types.DataField;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.utils.Range;
+import org.apache.paimon.utils.RoaringNavigableMap64;
+
+import org.apache.flink.api.common.functions.MapFunction;
+import org.apache.flink.api.common.typeinfo.PrimitiveArrayTypeInfo;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.junit.jupiter.api.Test;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/** Tests for {@link FlinkDataEvolutionVectorRead}. */
+public class FlinkDataEvolutionVectorReadTest {
+
+ @Test
+ public void testExecutesSerializedRawSearchTasks() {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(2);
+ ExecutingFlinkVectorRead read = new ExecutingFlinkVectorRead(env);
+
+ ScoredGlobalIndexResult result =
+
read.readRawSplitsInFlink(Collections.singletonList(rawSplit(0, 63)), null,
null);
+ ScoredGlobalIndexResult secondResult =
+ read.readRawSplitsInFlink(
+ Collections.singletonList(rawSplit(100, 103)), null,
null);
+
+ assertThat(result.results().getLongCardinality()).isEqualTo(64);
+ assertThat(secondResult.results().getLongCardinality()).isEqualTo(4);
+ }
+
+ @Test
+ public void testRawSearchUsesFlinkPath() {
+ TestingFlinkVectorRead read = new TestingFlinkVectorRead();
+ RawVectorSearchSplit rawSplit =
+ new RawVectorSearchSplit(
+ Collections.singletonList(new Range(42, 42)),
+ Collections.emptyList(),
+ null);
+ VectorScan.Plan plan = () -> Collections.singletonList(rawSplit);
+
+ GlobalIndexResult result = read.read(plan);
+
+ assertThat(read.rawFlinkPathUsed).isTrue();
+ assertThat(result.results().contains(42L)).isTrue();
+ }
+
+ @Test
+ public void testRawSearchSplitsRangesAcrossFlinkTasks() {
+ RecordingFlinkVectorRead read = new RecordingFlinkVectorRead();
+ RawVectorSearchSplit rawSplit =
+ new RawVectorSearchSplit(
+ Collections.singletonList(new Range(0, 63)),
Collections.emptyList(), null);
+
+ ScoredGlobalIndexResult result =
+ read.readRawSplitsInFlink(Collections.singletonList(rawSplit),
null, null);
+
+ assertThat(read.rawSearchRanges).containsExactly(new Range(0, 31), new
Range(32, 63));
+ assertThat(read.flinkParallelism).isEqualTo(2);
+ assertThat(result.results().getLongCardinality()).isEqualTo(64);
+ }
+
+ @Test
+ public void testRawSearchScopesPreFilterToTaskRanges() {
+ PreFilterRecordingFlinkVectorRead read = new
PreFilterRecordingFlinkVectorRead();
+ RoaringNavigableMap64 preFilter = new RoaringNavigableMap64();
+ preFilter.addRange(new Range(0, 29));
+ preFilter.add(35L);
+ preFilter.add(39L);
+
+ read.readRawSplitsInFlink(Collections.singletonList(rawSplit(0, 89)),
null, preFilter);
+
+ assertThat(read.rawSearchRanges)
+ .containsExactly(new Range(0, 29), new Range(30, 59), new
Range(60, 89));
+ assertThat(read.rawSearchPreFilters)
+ .containsExactly(null, Arrays.asList(35L, 39L),
Collections.emptyList());
+ }
+
+ @Test
+ public void testDistributedIndexRefinesAfterGlobalMerge() {
+ DistributedRefineFlinkVectorRead read = new
DistributedRefineFlinkVectorRead();
+
+ ScoredGlobalIndexResult result =
+ read.readIndexSplitsInFlink(indexSplits("test-vector-ann", 4),
new L2Indexer());
+
+ assertThat(read.flinkParallelism).isEqualTo(2);
+ assertThat(read.rawSearchCandidateRows).containsExactly(0L, 2L);
+ assertThat(result.results().getLongCardinality()).isEqualTo(1);
+ assertThat(result.results().contains(0L)).isTrue();
+ }
+
+ @Test
+ public void testMixedSearchUsesCombinedFlinkJob() {
+ MixedFlinkVectorRead read = new MixedFlinkVectorRead();
+ List<VectorSearchSplit> splits = new
ArrayList<>(indexSplits("test-vector-ann", 4));
+ splits.add(rawSplit(42, 105));
+
+ GlobalIndexResult result = read.read(() -> splits);
+
+ assertThat(read.combinedExecution).isTrue();
+ assertThat(result.results().contains(0L)).isTrue();
+ assertThat(result.results().contains(42L)).isTrue();
+ }
+
+ @Test
+ public void testMixedResultTypeKeepsSearchParallelism() {
+ StreamExecutionEnvironment env = environment();
+ env.setParallelism(8);
+ DataStream<byte[]> searchResults =
+ StreamExecutionEnvironmentUtils.fromData(
+ env,
+ Collections.singletonList(new byte[] {1}),
+
PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO)
+ .map((MapFunction<byte[], byte[]>) bytes -> bytes)
+
.returns(PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO)
+ .setParallelism(2);
+
+ DataStream<Tuple2<Byte, byte[]>> typedResults =
+ FlinkDataEvolutionVectorRead.addResultType(
+ searchResults, (byte) 0, "Add Vector Search Result
Type");
+
+ assertThat(typedResults.getParallelism()).isEqualTo(2);
+ }
+
+ @Test
+ public void testFlinkJobNameContainsFullTableName() {
+ FileStoreTable table = mock(FileStoreTable.class);
+ when(table.fullName()).thenReturn("default.T");
+ FlinkDataEvolutionVectorRead read =
+ new FlinkDataEvolutionVectorRead(
+ table,
+ null,
+ null,
+ 1,
+ new DataField(0, "vec", new
ArrayType(DataTypes.FLOAT())),
+ new float[] {1.0f},
+ null,
+ environment());
+
+ assertThat(read.flinkJobName("Vector Index Search"))
+ .isEqualTo("Vector Search - Vector Index Search : default.T");
+ }
+
+ private static StreamExecutionEnvironment environment() {
+ return StreamExecutionEnvironment.getExecutionEnvironment();
+ }
+
+ private static RawVectorSearchSplit rawSplit(long from, long to) {
+ return new RawVectorSearchSplit(
+ Collections.singletonList(new Range(from, to)),
Collections.emptyList(), null);
+ }
+
+ private static List<IndexVectorSearchSplit> indexSplits(String indexType,
int count) {
+ List<IndexVectorSearchSplit> splits = new ArrayList<>();
+ for (int i = 0; i < count; i++) {
+ GlobalIndexMeta globalIndexMeta = new GlobalIndexMeta(i, i, 0,
null, new byte[0]);
+ IndexFileMeta indexFile =
+ new IndexFileMeta(indexType, "index-" + i, 1L, 1L,
globalIndexMeta, null);
+ splits.add(
+ new IndexVectorSearchSplit(
+ i, i, Collections.singletonList(indexFile),
Collections.emptyList()));
+ }
+ return splits;
+ }
+
+ private static ScoredGlobalIndexResult scoredResult(long rowId, float
score) {
+ RoaringNavigableMap64 rows = new RoaringNavigableMap64();
+ rows.add(rowId);
+ return ScoredGlobalIndexResult.create(rows, candidate -> score);
+ }
+
+ private static class TestingFlinkVectorRead extends
FlinkDataEvolutionVectorRead {
+
+ private boolean rawFlinkPathUsed;
+
+ private TestingFlinkVectorRead() {
+ super(
+ null,
+ null,
+ null,
+ 10,
+ new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())),
+ new float[] {1.0f},
+ null,
+ environment());
+ }
+
+ @Override
+ protected GlobalIndexResult readSplits(List<? extends
VectorSearchSplit> splits) {
+ throw new AssertionError("Raw search should not fall back to local
vector read.");
+ }
+
+ @Override
+ protected ScoredGlobalIndexResult readIndexSplitsInFlink(
+ List<IndexVectorSearchSplit> splits, @Nullable GlobalIndexer
globalIndexer) {
+ throw new AssertionError("Index search is not part of this test.");
+ }
+
+ @Override
+ protected ScoredGlobalIndexResult readRawSplitsInFlink(
+ List<RawVectorSearchSplit> splits,
+ @Nullable GlobalIndexer globalIndexer,
+ @Nullable RoaringNavigableMap64 preFilter) {
+ rawFlinkPathUsed = true;
+ assertThat(splits).hasSize(1);
+ assertThat(globalIndexer).isNull();
+ assertThat(preFilter).isNull();
+
+ RoaringNavigableMap64 rows = new RoaringNavigableMap64();
+ rows.add(42L);
+ return ScoredGlobalIndexResult.create(rows, rowId -> 1.0f);
+ }
+ }
+
+ private static class DistributedRefineFlinkVectorRead extends
FlinkDataEvolutionVectorRead {
+
+ private int flinkParallelism;
+ private List<Long> rawSearchCandidateRows = Collections.emptyList();
+
+ private DistributedRefineFlinkVectorRead() {
+ super(
+ null,
+ null,
+ null,
+ 1,
+ new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())),
+ new float[] {0.0f},
+ Collections.singletonMap("refine_factor", "2"),
+ environment());
+ }
+
+ @Override
+ protected int flinkParallelism() {
+ return 2;
+ }
+
+ @Override
+ protected List<byte[]> executeIndexSearchGroups(
+ List<List<SerializedSplit>> groups, int searchLimit, int
parallelism) {
+ flinkParallelism = parallelism;
+ assertThat(groups).hasSize(2);
+ assertThat(searchLimit).isEqualTo(2);
+ try {
+ GlobalIndexResultSerializer serializer = new
GlobalIndexResultSerializer();
+ return Arrays.asList(
+ serializer.serialize(scoredResult(2L, 100.0f)),
+ serializer.serialize(scoredResult(0L, 1.0f)));
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
+ protected ScoredGlobalIndexResult readRawRefineSearch(
+ RoaringNavigableMap64 candidates,
+ @Nullable GlobalIndexer globalIndexer,
+ float[] queryVector) {
+ assertThat(globalIndexer).isInstanceOf(VectorGlobalIndexer.class);
+ assertThat(((VectorGlobalIndexer)
globalIndexer).metric()).isEqualTo("l2");
+ assertThat(queryVector).containsExactly(0.0f);
+ rawSearchCandidateRows = new ArrayList<>();
+ for (long rowId : candidates) {
+ rawSearchCandidateRows.add(rowId);
+ }
+
+ RoaringNavigableMap64 rows = new RoaringNavigableMap64();
+ for (long rowId : candidates) {
+ rows.add(rowId);
+ }
+ return ScoredGlobalIndexResult.create(
+ rows, rowId -> rowId == 0L ? 1.0f : 1.0f / (1.0f +
rowId * rowId))
+ .topK(1);
+ }
+ }
+
+ private static class MixedFlinkVectorRead extends
FlinkDataEvolutionVectorRead {
+
+ private boolean combinedExecution;
+
+ private MixedFlinkVectorRead() {
+ super(
+ null,
+ null,
+ null,
+ 2,
+ new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())),
+ new float[] {0.0f},
+ Collections.singletonMap("test.vector.metric", "l2"),
+ environment());
+ }
+
+ @Override
+ protected GlobalIndexer
createGlobalIndexer(List<IndexVectorSearchSplit> splits) {
+ return new L2Indexer();
+ }
+
+ @Override
+ protected int flinkParallelism() {
+ return 2;
+ }
+
+ @Override
+ protected List<Tuple2<Byte, byte[]>> executeIndexAndRawSearchGroups(
+ List<List<SerializedSplit>> indexGroups,
+ int searchLimit,
+ List<List<SerializedSplit>> rawGroups,
+ String rawMetric,
+ int parallelism) {
+ combinedExecution = true;
+ assertThat(indexGroups).hasSize(2);
+ assertThat(rawGroups).hasSize(2);
+ assertThat(searchLimit).isEqualTo(2);
+ assertThat(rawMetric).isEqualTo("l2");
+ assertThat(parallelism).isEqualTo(2);
+ try {
+ GlobalIndexResultSerializer serializer = new
GlobalIndexResultSerializer();
+ return Arrays.asList(
+ Tuple2.of((byte) 0,
serializer.serialize(scoredResult(0L, 1.0f))),
+ Tuple2.of((byte) 1,
serializer.serialize(scoredResult(42L, 0.9f))));
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
+ protected List<byte[]> executeIndexSearchGroups(
+ List<List<SerializedSplit>> groups, int searchLimit, int
parallelism) {
+ throw new AssertionError("Mixed search must not submit an
indexed-only Flink job.");
+ }
+
+ @Override
+ protected List<byte[]> executeRawSearchGroups(
+ List<List<SerializedSplit>> groups, String metric, int
parallelism) {
+ throw new AssertionError("Mixed search must not submit a raw-only
Flink job.");
+ }
+ }
+
+ private static class RecordingFlinkVectorRead extends
FlinkDataEvolutionVectorRead {
+
+ private final AtomicInteger nextTask = new AtomicInteger();
+ private final List<Range> rawSearchRanges =
Collections.synchronizedList(new ArrayList<>());
+ private int flinkParallelism;
+
+ private RecordingFlinkVectorRead() {
+ super(
+ null,
+ null,
+ null,
+ 100,
+ new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())),
+ new float[] {1.0f},
+ Collections.singletonMap("test.vector.metric", "l2"),
+ environment());
+ }
+
+ @Override
+ protected int flinkParallelism() {
+ return 2;
+ }
+
+ @Override
+ protected List<byte[]> executeRawSearchGroups(
+ List<List<SerializedSplit>> groups, String metric, int
parallelism) {
+ flinkParallelism = parallelism;
+ return groups.stream()
+ .map(group -> executeRawSearchGroup(group, metric))
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ protected ScoredGlobalIndexResult readRawSearch(
+ List<Range> rawRowRanges,
+ @Nullable RoaringNavigableMap64 preFilter,
+ String metric,
+ float[] queryVector) {
+ assertThat(preFilter).isNull();
+ assertThat(metric).isEqualTo("l2");
+ assertThat(queryVector).containsExactly(1.0f);
+ rawSearchRanges.addAll(rawRowRanges);
+
+ RoaringNavigableMap64 rows = new RoaringNavigableMap64();
+ int scoreBase = nextTask.getAndIncrement();
+ for (Range range : rawRowRanges) {
+ rows.addRange(range);
+ }
+ return ScoredGlobalIndexResult.create(rows, rowId -> scoreBase +
(float) rowId);
+ }
+ }
+
+ private static class PreFilterRecordingFlinkVectorRead extends
FlinkDataEvolutionVectorRead {
+
+ private final List<Range> rawSearchRanges = new ArrayList<>();
+ private final List<List<Long>> rawSearchPreFilters = new ArrayList<>();
+
+ private PreFilterRecordingFlinkVectorRead() {
+ super(
+ null,
+ null,
+ null,
+ 100,
+ new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())),
+ new float[] {1.0f},
+ Collections.singletonMap("test.vector.metric", "l2"),
+ environment());
+ }
+
+ @Override
+ protected int flinkParallelism() {
+ return 3;
+ }
+
+ @Override
+ protected List<byte[]> executeRawSearchGroups(
+ List<List<SerializedSplit>> groups, String metric, int
parallelism) {
+ return groups.stream()
+ .map(group -> executeRawSearchGroup(group, metric))
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ protected ScoredGlobalIndexResult readRawSearch(
+ List<Range> rawRowRanges,
+ @Nullable RoaringNavigableMap64 preFilter,
+ String metric,
+ float[] queryVector) {
+ assertThat(rawRowRanges).hasSize(1);
+ assertThat(metric).isEqualTo("l2");
+ assertThat(queryVector).containsExactly(1.0f);
+ rawSearchRanges.add(rawRowRanges.get(0));
+ if (preFilter == null) {
+ rawSearchPreFilters.add(null);
+ } else {
+ List<Long> rowIds = new ArrayList<>();
+ for (long rowId : preFilter) {
+ rowIds.add(rowId);
+ }
+ rawSearchPreFilters.add(rowIds);
+ }
+ return ScoredGlobalIndexResult.createEmpty();
+ }
+ }
+
+ private static class ExecutingFlinkVectorRead extends
FlinkDataEvolutionVectorRead {
+
+ private ExecutingFlinkVectorRead(StreamExecutionEnvironment env) {
+ super(
+ null,
+ null,
+ null,
+ 100,
+ new DataField(0, "vec", new ArrayType(DataTypes.FLOAT())),
+ new float[] {1.0f},
+ Collections.singletonMap("test.vector.metric", "l2"),
+ env);
+ }
+
+ @Override
+ protected int flinkParallelism() {
+ return 2;
+ }
+
+ @Override
+ protected String flinkJobName(String operatorName) {
+ return "Vector Search Test - " + operatorName;
+ }
+
+ @Override
+ protected ScoredGlobalIndexResult readRawSearch(
+ List<Range> rawRowRanges,
+ @Nullable RoaringNavigableMap64 preFilter,
+ String metric,
+ float[] queryVector) {
+ RoaringNavigableMap64 rows = new RoaringNavigableMap64();
+ for (Range range : rawRowRanges) {
+ rows.addRange(range);
+ }
+ return ScoredGlobalIndexResult.create(rows, rowId -> (float)
rowId);
+ }
+ }
+
+ private static class L2Indexer implements VectorGlobalIndexer {
+
+ @Override
+ public GlobalIndexWriter createWriter(GlobalIndexFileWriter
fileWriter) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public GlobalIndexReader createReader(
+ GlobalIndexFileReader fileReader,
+ List<GlobalIndexIOMeta> files,
+ ExecutorService executor) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public String metric() {
+ return "l2";
+ }
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/vectorsearch/FlinkPrimaryKeyVectorReadTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/vectorsearch/FlinkPrimaryKeyVectorReadTest.java
new file mode 100644
index 0000000000..e09dfdaec3
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/vectorsearch/FlinkPrimaryKeyVectorReadTest.java
@@ -0,0 +1,163 @@
+/*
+ * 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.flink.vectorsearch;
+
+import org.apache.paimon.CoreOptions;
+import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.globalindex.GlobalIndexResult;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.source.BucketVectorSearchSplit;
+import org.apache.paimon.table.source.DataSplit;
+import org.apache.paimon.table.source.PrimaryKeyVectorRead;
+import org.apache.paimon.table.source.PrimaryKeyVectorScan;
+import org.apache.paimon.table.source.VectorScan;
+import org.apache.paimon.types.DataField;
+import org.apache.paimon.types.FloatType;
+import org.apache.paimon.types.VectorType;
+
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/** Tests for {@link FlinkPrimaryKeyVectorRead}. */
+public class FlinkPrimaryKeyVectorReadTest {
+
+ @Test
+ public void testDistributesSerializedBucketGroups() {
+ List<BucketVectorSearchSplit> splits = new ArrayList<>();
+ for (int bucket = 0; bucket < 4; bucket++) {
+ DataSplit dataSplit =
+ DataSplit.builder()
+ .withSnapshot(1L)
+ .withPartition(BinaryRow.EMPTY_ROW)
+ .withBucket(bucket)
+ .withBucketPath("bucket-" + bucket)
+ .withDataFiles(Collections.emptyList())
+ .isStreaming(false)
+ .rawConvertible(false)
+ .build();
+ splits.add(new BucketVectorSearchSplit(dataSplit,
Collections.emptyList()));
+ }
+
+ TestingFlinkPrimaryKeyVectorRead read = new
TestingFlinkPrimaryKeyVectorRead(splits);
+ GlobalIndexResult result = read.read(() -> Collections.emptyList());
+
+ assertThat(result.results()).isEmpty();
+ assertThat(read.taskGroupSizes).containsExactly(2, 2);
+ assertThat(read.dispatchedBuckets).containsExactlyInAnyOrder(0, 1, 2,
3);
+ assertThat(read.operatorParallelism).isEqualTo(2);
+ assertThat(read.createResultCalled).isTrue();
+ }
+
+ @Test
+ public void testFlinkJobNameContainsFullTableName() {
+ TestingFlinkPrimaryKeyVectorRead read =
+ new TestingFlinkPrimaryKeyVectorRead(Collections.emptyList());
+
+ assertThat(read.flinkJobName("Primary-Key Vector Search"))
+ .isEqualTo("Vector Search - Primary-Key Vector Search :
default.PK_T");
+ }
+
+ private static FileStoreTable table() {
+ Map<String, String> options = new HashMap<>();
+ options.put("fields.vector.pk-vector.index.type", "test-vector-ann");
+ options.put("fields.vector.pk-vector.distance.metric", "l2");
+ FileStoreTable table = mock(FileStoreTable.class);
+ when(table.coreOptions()).thenReturn(new CoreOptions(options));
+ when(table.options()).thenReturn(options);
+ when(table.fullName()).thenReturn("default.PK_T");
+ return table;
+ }
+
+ private static class TestingFlinkPrimaryKeyVectorRead extends
FlinkPrimaryKeyVectorRead {
+
+ private final PrimaryKeyVectorScan.Plan primaryKeyPlan;
+ private final List<BucketVectorSearchSplit> splits;
+ private final List<Integer> taskGroupSizes = new ArrayList<>();
+ private final List<Integer> dispatchedBuckets = new ArrayList<>();
+ private int operatorParallelism;
+ private boolean createResultCalled;
+
+ private TestingFlinkPrimaryKeyVectorRead(List<BucketVectorSearchSplit>
splits) {
+ super(
+ table(),
+ new DataField(1, "vector", new VectorType(2, new
FloatType())),
+ new float[] {0.0f, 0.0f},
+ 1,
+ Collections.emptyMap(),
+ null,
+ mock(StreamExecutionEnvironment.class));
+ this.primaryKeyPlan = mock(PrimaryKeyVectorScan.Plan.class);
+ this.splits = splits;
+ }
+
+ @Override
+ protected PrimaryKeyVectorScan.Plan primaryKeyPlan(VectorScan.Plan
plan) {
+ return primaryKeyPlan;
+ }
+
+ @Override
+ protected List<BucketVectorSearchSplit>
bucketSplits(PrimaryKeyVectorScan.Plan plan) {
+ assertThat(plan).isSameAs(primaryKeyPlan);
+ return splits;
+ }
+
+ @Override
+ protected int flinkParallelism() {
+ return 2;
+ }
+
+ @Override
+ protected List<byte[]> executeBucketSearchGroups(
+ List<List<byte[]>> groups, int parallelism) {
+ operatorParallelism = parallelism;
+
taskGroupSizes.addAll(groups.stream().map(List::size).collect(Collectors.toList()));
+ return
groups.stream().map(this::executeBucketSearchGroup).collect(Collectors.toList());
+ }
+
+ @Override
+ protected SearchResult searchBuckets(List<BucketVectorSearchSplit>
taskSplits) {
+ dispatchedBuckets.addAll(
+ taskSplits.stream()
+ .map(split -> split.dataSplit().bucket())
+ .collect(Collectors.toList()));
+ return new SearchResult(Collections.emptyList(),
Collections.emptyList());
+ }
+
+ @Override
+ protected GlobalIndexResult createResult(
+ PrimaryKeyVectorScan.Plan plan,
PrimaryKeyVectorRead.SearchResult searchResult) {
+ assertThat(plan).isSameAs(primaryKeyPlan);
+ assertThat(searchResult.indexedCandidates()).isEmpty();
+ assertThat(searchResult.exactCandidates()).isEmpty();
+ createResultCalled = true;
+ return GlobalIndexResult.createEmpty();
+ }
+ }
+}
