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 96cd9a0ace [vector] Support vector search options (#8203)
96cd9a0ace is described below
commit 96cd9a0acee6449d901aa9504de4fa8c7c500b4e
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Jun 11 17:19:59 2026 +0800
[vector] Support vector search options (#8203)
Expose query-time vector search options through the Java and Python
vector search APIs, and thread them into Flink, Spark, and Lumina search
execution. This lets callers configure index-specific query parameters
such as `ivf.nprobe` and `hnsw.ef_search` at search time.
---
.../org/apache/paimon/predicate/VectorSearch.java | 18 ++++-
.../testvector/TestVectorGlobalIndexReader.java | 25 +++++++
.../testvector/TestVectorGlobalIndexer.java | 11 ++-
.../apache/paimon/predicate/VectorSearchTest.java | 24 +++++++
.../apache/paimon/table/source/VectorReadImpl.java | 20 +++++-
.../paimon/table/source/VectorSearchBuilder.java | 13 ++++
.../table/source/VectorSearchBuilderImpl.java | 20 +++++-
.../table/source/VectorSearchBuilderTest.java | 33 +++++++++
.../flink/procedure/VectorSearchProcedure.java | 1 +
.../procedure/VectorSearchProcedureITCase.java | 42 ++++++++++-
.../apache/paimon/lumina/index/LuminaIndex.java | 23 +++---
.../index/LuminaVectorGlobalIndexReader.java | 40 +++++++----
.../lumina/index/LuminaVectorOptionsTest.java | 52 ++++++++++++++
.../lumina/lumina_vector_global_index_reader.py | 34 ++++++---
.../pypaimon/globalindex/vector_search.py | 14 +++-
.../pypaimon/table/source/vector_search_builder.py | 27 +++++++
.../pypaimon/table/source/vector_search_read.py | 7 +-
.../pypaimon/tests/vector_search_filter_test.py | 75 +++++++++++++++++++
.../paimon/spark/PaimonScanBuilderTest.scala | 4 ++
.../paimon/spark/read/SparkVectorReadImpl.java | 11 +++
.../spark/read/SparkVectorSearchBuilderImpl.java | 2 +-
.../org/apache/paimon/spark/PaimonBaseScan.scala | 1 +
.../plans/logical/PaimonTableValuedFunctions.scala | 83 ++++++++++++++++++++--
.../plans/logical/VectorSearchQueryTest.scala | 77 ++++++++++++++++++++
.../paimon/spark/sql/VectorSearchOptionsTest.scala | 70 ++++++++++++++++++
25 files changed, 673 insertions(+), 54 deletions(-)
diff --git
a/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
index 5e660ed17f..c2b608346b 100644
--- a/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
+++ b/paimon-common/src/main/java/org/apache/paimon/predicate/VectorSearch.java
@@ -24,6 +24,9 @@ import org.apache.paimon.utils.RoaringNavigableMap64;
import javax.annotation.Nullable;
import java.io.Serializable;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
/** VectorSearch to perform vector similarity search. * */
public class VectorSearch implements Serializable {
@@ -33,10 +36,15 @@ public class VectorSearch implements Serializable {
private final float[] vector;
private final String fieldName;
private final int limit;
+ private final Map<String, String> options;
@Nullable private RoaringNavigableMap64 includeRowIds;
public VectorSearch(float[] vector, int limit, String fieldName) {
+ this(vector, limit, fieldName, Collections.emptyMap());
+ }
+
+ public VectorSearch(float[] vector, int limit, String fieldName,
Map<String, String> options) {
if (vector == null) {
throw new IllegalArgumentException("Search cannot be null");
}
@@ -49,6 +57,10 @@ public class VectorSearch implements Serializable {
this.vector = vector;
this.limit = limit;
this.fieldName = fieldName;
+ this.options =
+ options == null
+ ? Collections.emptyMap()
+ : Collections.unmodifiableMap(new HashMap<>(options));
}
public float[] vector() {
@@ -63,6 +75,10 @@ public class VectorSearch implements Serializable {
return fieldName;
}
+ public Map<String, String> options() {
+ return options == null ? Collections.emptyMap() : options;
+ }
+
public RoaringNavigableMap64 includeRowIds() {
return includeRowIds;
}
@@ -81,7 +97,7 @@ public class VectorSearch implements Serializable {
for (long rowId : and64) {
roaringNavigableMap64Offset.add(rowId - from);
}
- VectorSearch target = new VectorSearch(vector, limit, fieldName);
+ VectorSearch target = new VectorSearch(vector, limit, fieldName,
options());
target.withIncludeRowIds(roaringNavigableMap64Offset);
return target;
}
diff --git
a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexReader.java
b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexReader.java
index da7f533f8c..fd233ba06b 100644
---
a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexReader.java
+++
b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexReader.java
@@ -56,6 +56,8 @@ public class TestVectorGlobalIndexReader implements
GlobalIndexReader {
private final GlobalIndexFileReader fileReader;
private final GlobalIndexIOMeta ioMeta;
private final String metric;
+ private final String requiredOptionKey;
+ private final String requiredOptionValue;
private float[][] vectors;
private int dimension;
@@ -63,9 +65,20 @@ public class TestVectorGlobalIndexReader implements
GlobalIndexReader {
public TestVectorGlobalIndexReader(
GlobalIndexFileReader fileReader, GlobalIndexIOMeta ioMeta, String
metric) {
+ this(fileReader, ioMeta, metric, null, null);
+ }
+
+ public TestVectorGlobalIndexReader(
+ GlobalIndexFileReader fileReader,
+ GlobalIndexIOMeta ioMeta,
+ String metric,
+ String requiredOptionKey,
+ String requiredOptionValue) {
this.fileReader = fileReader;
this.ioMeta = ioMeta;
this.metric = metric;
+ this.requiredOptionKey = requiredOptionKey;
+ this.requiredOptionValue = requiredOptionValue;
}
@Override
@@ -78,6 +91,18 @@ public class TestVectorGlobalIndexReader implements
GlobalIndexReader {
}
float[] queryVector = vectorSearch.vector();
+ if (requiredOptionKey != null) {
+ String actual = vectorSearch.options().get(requiredOptionKey);
+ if (!requiredOptionValue.equals(actual)) {
+ throw new IllegalArgumentException(
+ "Required option "
+ + requiredOptionKey
+ + " expected "
+ + requiredOptionValue
+ + " but got "
+ + actual);
+ }
+ }
if (queryVector.length != dimension) {
throw new IllegalArgumentException(
String.format(
diff --git
a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexer.java
b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexer.java
index d12da8ffca..cd3a140681 100644
---
a/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexer.java
+++
b/paimon-common/src/test/java/org/apache/paimon/globalindex/testvector/TestVectorGlobalIndexer.java
@@ -55,9 +55,15 @@ public class TestVectorGlobalIndexer implements
GlobalIndexer {
/** Option key for distance metric. */
public static final String OPT_METRIC = "test.vector.metric";
+ public static final String OPT_REQUIRED_OPTION_KEY =
"test.vector.required-option.key";
+
+ public static final String OPT_REQUIRED_OPTION_VALUE =
"test.vector.required-option.value";
+
private final DataType fieldType;
private final int dimension;
private final String metric;
+ private final String requiredOptionKey;
+ private final String requiredOptionValue;
public TestVectorGlobalIndexer(DataType fieldType, Options options) {
checkArgument(
@@ -67,6 +73,8 @@ public class TestVectorGlobalIndexer implements GlobalIndexer
{
this.fieldType = fieldType;
this.dimension = options.getInteger(OPT_DIMENSION, 0);
this.metric = options.getString(OPT_METRIC, "l2");
+ this.requiredOptionKey = options.getString(OPT_REQUIRED_OPTION_KEY,
null);
+ this.requiredOptionValue =
options.getString(OPT_REQUIRED_OPTION_VALUE, null);
}
@Override
@@ -80,7 +88,8 @@ public class TestVectorGlobalIndexer implements GlobalIndexer
{
List<GlobalIndexIOMeta> files,
ExecutorService executor) {
checkArgument(files.size() == 1, "Expected exactly one index file per
shard");
- return new TestVectorGlobalIndexReader(fileReader, files.get(0),
metric);
+ return new TestVectorGlobalIndexReader(
+ fileReader, files.get(0), metric, requiredOptionKey,
requiredOptionValue);
}
public int dimension() {
diff --git
a/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java
b/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java
index 735874ce84..0284ff3d37 100644
---
a/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java
+++
b/paimon-common/src/test/java/org/apache/paimon/predicate/VectorSearchTest.java
@@ -23,7 +23,9 @@ import org.apache.paimon.utils.RoaringNavigableMap64;
import org.junit.jupiter.api.Test;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
@@ -49,4 +51,26 @@ public class VectorSearchTest {
List<Range> ranges = vectorSearch.includeRowIds().toRangeList();
assertThat(ranges.get(0)).isEqualTo(new Range(40L, 90L));
}
+
+ @Test
+ public void testVectorSearchOffsetKeepsOptions() {
+ Map<String, String> options = new HashMap<>();
+ options.put("ivf.nprobe", "16");
+ options.put("hnsw.ef_search", "64");
+
+ VectorSearch vectorSearch = new VectorSearch(new float[] {1.0f, 0.0f},
1, "test", options);
+
+ RoaringNavigableMap64 includeRowIds = new RoaringNavigableMap64();
+ includeRowIds.addRange(new Range(100L, 200L));
+ vectorSearch.withIncludeRowIds(includeRowIds);
+
+ VectorSearch offset = vectorSearch.offsetRange(60, 150);
+
+ assertThat(offset.options()).isEqualTo(options);
+ options.put("ivf.nprobe", "32");
+ assertThat(offset.options())
+ .containsEntry("ivf.nprobe", "16")
+ .containsEntry("hnsw.ef_search", "64")
+ .hasSize(2);
+ }
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
index 2eae2d4877..a4ef24637d 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java
@@ -44,8 +44,11 @@ import javax.annotation.Nullable;
import java.io.IOException;
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.Comparator;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.TreeSet;
@@ -65,6 +68,7 @@ public class VectorReadImpl implements VectorRead,
Serializable {
protected final int limit;
protected final DataField vectorColumn;
protected final float[] vector;
+ protected final Map<String, String> options;
public VectorReadImpl(
FileStoreTable table,
@@ -72,11 +76,25 @@ public class VectorReadImpl implements VectorRead,
Serializable {
int limit,
DataField vectorColumn,
float[] vector) {
+ this(table, filter, limit, vectorColumn, vector,
Collections.emptyMap());
+ }
+
+ public VectorReadImpl(
+ FileStoreTable table,
+ Predicate filter,
+ int limit,
+ DataField vectorColumn,
+ float[] vector,
+ Map<String, String> options) {
this.table = table;
this.filter = filter;
this.limit = limit;
this.vectorColumn = vectorColumn;
this.vector = vector;
+ this.options =
+ options == null
+ ? Collections.emptyMap()
+ : Collections.unmodifiableMap(new HashMap<>(options));
}
@Override
@@ -165,7 +183,7 @@ public class VectorReadImpl implements VectorRead,
Serializable {
GlobalIndexReader reader =
globalIndexer.createReader(indexFileReader, indexIOMetaList,
executor);
VectorSearch vectorSearch =
- new VectorSearch(vector, limit, vectorColumn.name())
+ new VectorSearch(vector, limit, vectorColumn.name(), options)
.withIncludeRowIds(includeRowIds);
return new OffsetGlobalIndexReader(reader, rowRangeStart, rowRangeEnd)
.visitVectorSearch(vectorSearch)
diff --git
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilder.java
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilder.java
index ae7e7bf48e..e33f102866 100644
---
a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilder.java
+++
b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilder.java
@@ -23,6 +23,7 @@ import org.apache.paimon.partition.PartitionPredicate;
import org.apache.paimon.predicate.Predicate;
import java.io.Serializable;
+import java.util.Map;
/** Builder to build vector search. */
public interface VectorSearchBuilder extends Serializable {
@@ -42,6 +43,18 @@ public interface VectorSearchBuilder extends Serializable {
/** The vector to search. */
VectorSearchBuilder withVector(float[] vector);
+ /** Option for vector indexes. */
+ default VectorSearchBuilder withOption(String key, String value) {
+ throw new UnsupportedOperationException(
+ getClass().getName() + " does not support vector options.");
+ }
+
+ /** Options for vector indexes. */
+ default VectorSearchBuilder withOptions(Map<String, String> options) {
+ throw new UnsupportedOperationException(
+ getClass().getName() + " does not support vector options.");
+ }
+
/** Create vector scan to scan index files. */
VectorScan newVectorScan();
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 a0d11ff21f..d4686b4416 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
@@ -25,6 +25,9 @@ import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.InnerTable;
import org.apache.paimon.types.DataField;
+import java.util.HashMap;
+import java.util.Map;
+
import static
org.apache.paimon.partition.PartitionPredicate.splitPartitionPredicate;
/** Implementation for {@link VectorSearchBuilder}. */
@@ -39,6 +42,7 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
protected int limit;
protected DataField vectorColumn;
protected float[] vector;
+ protected Map<String, String> options = new HashMap<>();
public VectorSearchBuilderImpl(InnerTable table) {
this.table = (FileStoreTable) table;
@@ -80,6 +84,20 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
return this;
}
+ @Override
+ public VectorSearchBuilder withOptions(Map<String, String> options) {
+ if (options != null) {
+ this.options.putAll(options);
+ }
+ return this;
+ }
+
+ @Override
+ public VectorSearchBuilder withOption(String key, String value) {
+ this.options.put(key, value);
+ return this;
+ }
+
@Override
public VectorScan newVectorScan() {
return new VectorScanImpl(table, partitionFilter, filter,
vectorColumn);
@@ -87,6 +105,6 @@ public class VectorSearchBuilderImpl implements
VectorSearchBuilder {
@Override
public VectorRead newVectorRead() {
- return new VectorReadImpl(table, filter, limit, vectorColumn, vector);
+ return new VectorReadImpl(table, filter, limit, vectorColumn, vector,
options);
}
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
index 79928bacca..17ff21bb61 100644
---
a/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java
@@ -227,6 +227,39 @@ public class VectorSearchBuilderTest extends TableTestBase
{
assertThat(ids.size()).isLessThanOrEqualTo(5);
}
+ @Test
+ public void testVectorSearchThreadsOptions() throws Exception {
+ catalog.createTable(
+ identifier("options_table"),
+ Schema.newBuilder()
+ .column("id", DataTypes.INT())
+ .column(VECTOR_FIELD_NAME, new
ArrayType(DataTypes.FLOAT()))
+ .option(CoreOptions.BUCKET.key(), "-1")
+ .option(CoreOptions.ROW_TRACKING_ENABLED.key(), "true")
+ .option(CoreOptions.DATA_EVOLUTION_ENABLED.key(),
"true")
+ .option("test.vector.dimension",
String.valueOf(DIMENSION))
+ .option("test.vector.metric", "l2")
+ .option("test.vector.required-option.key",
"ivf.nprobe")
+ .option("test.vector.required-option.value", "16")
+ .build(),
+ false);
+ FileStoreTable table = getTable(identifier("options_table"));
+
+ float[][] vectors = {{1.0f, 0.0f}, {0.0f, 1.0f}};
+ writeVectors(table, vectors);
+ buildAndCommitIndex(table, vectors);
+
+ GlobalIndexResult result =
+ table.newVectorSearchBuilder()
+ .withVector(new float[] {1.0f, 0.0f})
+ .withLimit(1)
+ .withVectorColumn(VECTOR_FIELD_NAME)
+ .withOption("ivf.nprobe", "16")
+ .executeLocal();
+
+ assertThat(result.results().isEmpty()).isFalse();
+ }
+
@Test
public void testVectorSearchWithMultipleIndexFiles() throws Exception {
createTableDefault();
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 2b5df413f8..8a75b8f9d4 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
@@ -106,6 +106,7 @@ public class VectorSearchProcedure extends ProcedureBase {
.withVector(queryVector)
.withVectorColumn(vectorColumn)
.withLimit(topK)
+ .withOptions(optionsMap)
.executeLocal();
RowType tableRowType = table.rowType();
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 935da25330..7beab44fdc 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
@@ -25,6 +25,7 @@ import org.apache.paimon.flink.CatalogITCaseBase;
import org.apache.paimon.globalindex.GlobalIndexBuilderUtils;
import org.apache.paimon.globalindex.GlobalIndexSingletonWriter;
import org.apache.paimon.globalindex.ResultEntry;
+import org.apache.paimon.globalindex.testvector.TestVectorGlobalIndexer;
import org.apache.paimon.globalindex.testvector.TestVectorGlobalIndexerFactory;
import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.io.CompactIncrement;
@@ -147,7 +148,45 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
assertThat(result.size()).isLessThanOrEqualTo(3);
}
+ @Test
+ public void testVectorSearchWithOptions() throws Exception {
+ createVectorTable(
+ "T4",
+ "'"
+ + TestVectorGlobalIndexer.OPT_REQUIRED_OPTION_KEY
+ + "' = 'ivf.nprobe', "
+ + "'"
+ + TestVectorGlobalIndexer.OPT_REQUIRED_OPTION_VALUE
+ + "' = '16'");
+ FileStoreTable table = paimonTable("T4");
+
+ float[][] vectors = {
+ {1.0f, 0.0f}, // row 0
+ {0.0f, 1.0f}, // row 1
+ };
+
+ writeVectors(table, vectors);
+ buildAndCommitVectorIndex(table, vectors);
+
+ List<Row> result =
+ sql(
+ "CALL sys.vector_search("
+ + "`table` => 'default.T4', "
+ + "vector_column => 'vec', "
+ + "query_vector => '1.0,0.0', "
+ + "top_k => 2, "
+ + "options => 'ivf.nprobe=16')");
+
+ assertThat(result).isNotEmpty();
+ assertThat(result.size()).isLessThanOrEqualTo(2);
+ }
+
private void createVectorTable(String tableName) {
+ createVectorTable(tableName, "");
+ }
+
+ private void createVectorTable(String tableName, String extraOptions) {
+ String formattedExtraOptions = extraOptions.isEmpty() ? "" : ", " +
extraOptions;
sql(
"CREATE TABLE %s ("
+ "id INT, "
@@ -158,8 +197,9 @@ public class VectorSearchProcedureITCase extends
CatalogITCaseBase {
+ "'data-evolution.enabled' = 'true', "
+ "'test.vector.dimension' = '%d', "
+ "'test.vector.metric' = 'l2'"
+ + "%s"
+ ")",
- tableName, DIMENSION);
+ tableName, DIMENSION, formattedExtraOptions);
}
private void writeVectors(FileStoreTable table, float[][] vectors) throws
Exception {
diff --git
a/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaIndex.java
b/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaIndex.java
index 1850f53ea8..daa8c552df 100644
---
a/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaIndex.java
+++
b/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaIndex.java
@@ -106,10 +106,10 @@ public class LuminaIndex implements Closeable {
int k,
float[] distances,
long[] labels,
- Map<String, String> searchOptions) {
+ Map<String, String> options) {
ensureOpen();
ensureSearcher();
- searcher.search(n, queryVectors, k, distances, labels,
filterSearchOptions(searchOptions));
+ searcher.search(n, queryVectors, k, distances, labels,
filterOptions(options));
}
/** Search for k nearest neighbors with native pre-filtering on vector
IDs. */
@@ -120,17 +120,11 @@ public class LuminaIndex implements Closeable {
float[] distances,
long[] labels,
long[] filterIds,
- Map<String, String> searchOptions) {
+ Map<String, String> options) {
ensureOpen();
ensureSearcher();
searcher.searchWithFilter(
- n,
- queryVectors,
- k,
- distances,
- labels,
- filterIds,
- filterSearchOptions(searchOptions));
+ n, queryVectors, k, distances, labels, filterIds,
filterOptions(options));
}
/** Get the number of vectors (searcher mode). */
@@ -149,13 +143,12 @@ public class LuminaIndex implements Closeable {
}
/**
- * Filters an options map to only include keys valid for Lumina
SearchOptions. This mirrors
- * paimon-cpp's {@code NormalizeSearchOptions} which extracts only
search-relevant keys.
+ * Filters an options map to only include keys accepted by Lumina at query
time.
*
- * <p>Valid search option prefixes: {@code search.*} (core search options)
and {@code
- * diskann.search.*} (DiskANN-specific search options).
+ * <p>Valid query-time prefixes: {@code search.*} (core query options) and
{@code
+ * diskann.search.*} (DiskANN-specific query options).
*/
- private static Map<String, String> filterSearchOptions(Map<String, String>
options) {
+ private static Map<String, String> filterOptions(Map<String, String>
options) {
Map<String, String> searchOpts = new LinkedHashMap<>();
for (Map.Entry<String, String> entry : options.entrySet()) {
String key = entry.getKey();
diff --git
a/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaVectorGlobalIndexReader.java
b/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaVectorGlobalIndexReader.java
index 68dd9b43fb..20fffa9d11 100644
---
a/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaVectorGlobalIndexReader.java
+++
b/paimon-lumina/src/main/java/org/apache/paimon/lumina/index/LuminaVectorGlobalIndexReader.java
@@ -141,19 +141,25 @@ public class LuminaVectorGlobalIndexReader implements
GlobalIndexReader {
effectiveK = Math.min(effectiveK, scopedIds.length);
distances = new float[effectiveK];
labels = new long[effectiveK];
- Map<String, String> searchOptions = options.toLuminaOptions();
- searchOptions.putAll(indexMeta.options());
- searchOptions.put("search.thread_safe_filter", "true");
- ensureSearchListSize(searchOptions, effectiveK);
+ Map<String, String> mergedOptions =
+ mergeOptions(
+ this.options.toLuminaOptions(),
+ indexMeta.options(),
+ vectorSearch.options());
+ mergedOptions.put("search.thread_safe_filter", "true");
+ ensureSearchListSize(mergedOptions, effectiveK);
index.searchWithFilter(
- queryVector, 1, effectiveK, distances, labels, scopedIds,
searchOptions);
+ queryVector, 1, effectiveK, distances, labels, scopedIds,
mergedOptions);
} else {
distances = new float[effectiveK];
labels = new long[effectiveK];
- Map<String, String> searchOptions = options.toLuminaOptions();
- searchOptions.putAll(indexMeta.options());
- ensureSearchListSize(searchOptions, effectiveK);
- index.search(queryVector, 1, effectiveK, distances, labels,
searchOptions);
+ Map<String, String> mergedOptions =
+ mergeOptions(
+ this.options.toLuminaOptions(),
+ indexMeta.options(),
+ vectorSearch.options());
+ ensureSearchListSize(mergedOptions, effectiveK);
+ index.search(queryVector, 1, effectiveK, distances, labels,
mergedOptions);
}
// Min-heap: smallest score at head, so we can evict the weakest
candidate efficiently.
@@ -170,10 +176,20 @@ public class LuminaVectorGlobalIndexReader implements
GlobalIndexReader {
return new LuminaScoredGlobalIndexResult(roaringBitmap64, id2scores);
}
- private static void ensureSearchListSize(Map<String, String>
searchOptions, int topK) {
- if (!searchOptions.containsKey("diskann.search.list_size")) {
+ static Map<String, String> mergeOptions(
+ Map<String, String> baseOptions,
+ Map<String, String> indexOptions,
+ Map<String, String> queryOptions) {
+ Map<String, String> options = new HashMap<>(baseOptions);
+ options.putAll(indexOptions);
+ options.putAll(queryOptions);
+ return options;
+ }
+
+ private static void ensureSearchListSize(Map<String, String> options, int
topK) {
+ if (!options.containsKey("diskann.search.list_size")) {
int listSize = Math.max((int) (topK * 1.5), MIN_SEARCH_LIST_SIZE);
- searchOptions.put("diskann.search.list_size",
String.valueOf(listSize));
+ options.put("diskann.search.list_size", String.valueOf(listSize));
}
}
diff --git
a/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/LuminaVectorOptionsTest.java
b/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/LuminaVectorOptionsTest.java
new file mode 100644
index 0000000000..7ea5088780
--- /dev/null
+++
b/paimon-lumina/src/test/java/org/apache/paimon/lumina/index/LuminaVectorOptionsTest.java
@@ -0,0 +1,52 @@
+/*
+ * 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.lumina.index;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for Lumina vector options. */
+public class LuminaVectorOptionsTest {
+
+ @Test
+ public void testQueryOptionsOverrideIndexOptions() {
+ Map<String, String> baseOptions = new HashMap<>();
+ baseOptions.put("diskann.search.list_size", "16");
+ baseOptions.put("search.parallel_number", "2");
+ Map<String, String> indexOptions = new HashMap<>();
+ indexOptions.put("diskann.search.list_size", "32");
+ indexOptions.put("index.dimension", "4");
+ Map<String, String> queryOptions = new HashMap<>();
+ queryOptions.put("diskann.search.list_size", "64");
+ queryOptions.put("hnsw.ef_search", "128");
+
+ Map<String, String> merged =
+ LuminaVectorGlobalIndexReader.mergeOptions(baseOptions,
indexOptions, queryOptions);
+
+ assertThat(merged)
+ .containsEntry("diskann.search.list_size", "64")
+ .containsEntry("search.parallel_number", "2")
+ .containsEntry("index.dimension", "4")
+ .containsEntry("hnsw.ef_search", "128");
+ }
+}
diff --git
a/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py
b/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py
index fd425b3a13..cabd491911 100644
---
a/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py
+++
b/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py
@@ -36,11 +36,18 @@ LUMINA_IDENTIFIERS = (LUMINA_IDENTIFIER,
LUMINA_VECTOR_ANN_IDENTIFIER)
MIN_SEARCH_LIST_SIZE = 16
-def _ensure_search_list_size(search_options, top_k):
+def _ensure_search_list_size(options, top_k):
"""Set diskann.search.list_size when not explicitly configured."""
- if "diskann.search.list_size" not in search_options:
+ if "diskann.search.list_size" not in options:
list_size = max(int(top_k * 1.5), MIN_SEARCH_LIST_SIZE)
- search_options["diskann.search.list_size"] = str(list_size)
+ options["diskann.search.list_size"] = str(list_size)
+
+
+def _merge_options(base_options, index_options, query_options):
+ options = dict(base_options)
+ options.update(index_options)
+ options.update(query_options or {})
+ return options
class LuminaVectorGlobalIndexReader(GlobalIndexReader):
@@ -51,10 +58,10 @@ class LuminaVectorGlobalIndexReader(GlobalIndexReader):
self._file_io = file_io
self._index_path = index_path
self._io_meta = io_metas[0]
- self._options = options or {}
+ self._table_options = dict(options or {})
+ self._options = {}
self._searcher = None
self._index_meta = None
- self._search_options = None
self._stream = None
self._load_lock = threading.Lock()
@@ -78,19 +85,22 @@ class LuminaVectorGlobalIndexReader(GlobalIndexReader):
return _completed_future(None)
include_row_ids = vector_search.include_row_ids
+ query_options = vector_search.options
if include_row_ids is not None:
filter_id_list = list(include_row_ids)
if len(filter_id_list) == 0:
return _completed_future(None)
effective_k = min(effective_k, len(filter_id_list))
- search_opts = dict(self._search_options)
+ search_opts = _merge_options(
+ self._options, {}, query_options)
search_opts["search.thread_safe_filter"] = "true"
_ensure_search_list_size(search_opts, effective_k)
distances, labels = self._searcher.search_with_filter_list(
query_flat, 1, effective_k, filter_id_list, search_opts)
else:
- search_opts = dict(self._search_options)
+ search_opts = _merge_options(
+ self._options, {}, query_options)
_ensure_search_list_size(search_opts, effective_k)
distances, labels = self._searcher.search_list(
query_flat, 1, effective_k, search_opts)
@@ -123,16 +133,18 @@ class LuminaVectorGlobalIndexReader(GlobalIndexReader):
)
self._index_meta =
LuminaIndexMeta.deserialize(self._io_meta.metadata)
- searcher_options = strip_lumina_options(self._options)
- searcher_options.update(self._index_meta.options)
- self._search_options = searcher_options
+ self._options = _merge_options(
+ strip_lumina_options(self._table_options),
+ self._index_meta.options,
+ {},
+ )
file_path = (self._io_meta.external_path
if self._io_meta.external_path
else os.path.join(self._index_path,
self._io_meta.file_name))
stream = self._file_io.new_input_stream(file_path)
try:
- self._searcher = LuminaSearcher(searcher_options)
+ self._searcher = LuminaSearcher(self._options)
self._searcher.open_stream(stream, self._io_meta.file_size)
self._stream = stream
except Exception:
diff --git a/paimon-python/pypaimon/globalindex/vector_search.py
b/paimon-python/pypaimon/globalindex/vector_search.py
index 3cd8d2c88a..a5e5f709ec 100644
--- a/paimon-python/pypaimon/globalindex/vector_search.py
+++ b/paimon-python/pypaimon/globalindex/vector_search.py
@@ -19,7 +19,7 @@
from concurrent.futures import Future
from dataclasses import dataclass, field
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
import numpy as np
@@ -33,12 +33,14 @@ class VectorSearch:
limit: Maximum number of results to return
field_name: Name of the vector field to search
include_row_ids: Optional bitmap of row IDs to include in search
+ options: Query-time options for vector indexes
"""
vector: Union[List[float], np.ndarray]
limit: int
field_name: str
include_row_ids: Optional['RoaringBitmap64'] = field(default=None)
+ options: Optional[Dict[str, str]] = field(default=None)
def __post_init__(self):
if self.vector is None:
@@ -51,6 +53,10 @@ class VectorSearch:
# Convert list to numpy array if needed
if isinstance(self.vector, list):
self.vector = np.array(self.vector, dtype=np.float32)
+ if self.options is None:
+ self.options = {}
+ else:
+ self.options = dict(self.options)
def with_include_row_ids(self, include_row_ids: 'RoaringBitmap64') ->
'VectorSearch':
"""Return a new VectorSearch with the specified include_row_ids."""
@@ -58,7 +64,8 @@ class VectorSearch:
vector=self.vector,
limit=self.limit,
field_name=self.field_name,
- include_row_ids=include_row_ids
+ include_row_ids=include_row_ids,
+ options=self.options
)
def offset_range(self, from_: int, to: int) -> 'VectorSearch':
@@ -80,7 +87,8 @@ class VectorSearch:
vector=self.vector,
limit=self.limit,
field_name=self.field_name,
- include_row_ids=offset_bitmap
+ include_row_ids=offset_bitmap,
+ options=self.options
)
return self
diff --git a/paimon-python/pypaimon/table/source/vector_search_builder.py
b/paimon-python/pypaimon/table/source/vector_search_builder.py
index 4622d5901c..7cb9eb0104 100644
--- a/paimon-python/pypaimon/table/source/vector_search_builder.py
+++ b/paimon-python/pypaimon/table/source/vector_search_builder.py
@@ -45,6 +45,20 @@ class VectorSearchBuilder(ABC):
"""The query vector (list of floats)."""
pass
+ def with_option(self, key, value):
+ # type: (str, str) -> VectorSearchBuilder
+ """Option for vector indexes."""
+ raise NotImplementedError(
+ "%s does not support vector options."
+ % self.__class__.__name__)
+
+ def with_options(self, options):
+ # type: (dict) -> VectorSearchBuilder
+ """Options for vector indexes."""
+ raise NotImplementedError(
+ "%s does not support vector options."
+ % self.__class__.__name__)
+
@abstractmethod
def with_filter(self, predicate):
# type: (Predicate) -> VectorSearchBuilder
@@ -87,6 +101,7 @@ class VectorSearchBuilderImpl(VectorSearchBuilder):
self._query_vector = None
self._filter = None
self._partition_filter = None
+ self._options = {}
def with_limit(self, limit):
# type: (int) -> VectorSearchBuilder
@@ -106,6 +121,17 @@ class VectorSearchBuilderImpl(VectorSearchBuilder):
self._query_vector = vector
return self
+ def with_option(self, key, value):
+ # type: (str, str) -> VectorSearchBuilder
+ self._options[key] = value
+ return self
+
+ def with_options(self, options):
+ # type: (dict) -> VectorSearchBuilder
+ if options is not None:
+ self._options.update(options)
+ return self
+
def with_filter(self, predicate):
# type: (Predicate) -> VectorSearchBuilder
if predicate is None:
@@ -212,4 +238,5 @@ class VectorSearchBuilderImpl(VectorSearchBuilder):
self._vector_column,
self._query_vector,
filter_=self._filter,
+ options=self._options,
)
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index e6839ebe10..2abac3bed8 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -43,12 +43,14 @@ class VectorSearchRead(ABC):
class VectorSearchReadImpl(VectorSearchRead):
"""Implementation for VectorSearchRead."""
- def __init__(self, table, limit, vector_column, query_vector,
filter_=None):
+ def __init__(self, table, limit, vector_column, query_vector, filter_=None,
+ options=None):
self._table = table
self._limit = limit
self._vector_column = vector_column
self._query_vector = query_vector
self._filter = filter_
+ self._options = dict(options or {})
def read(self, splits):
# type: (List[VectorSearchSplit]) -> GlobalIndexResult
@@ -136,7 +138,8 @@ class VectorSearchReadImpl(VectorSearchRead):
vector_search = VectorSearch(
vector=self._query_vector,
limit=self._limit,
- field_name=self._vector_column.name
+ field_name=self._vector_column.name,
+ options=self._options,
)
if include_row_ids is not None:
vector_search = vector_search.with_include_row_ids(include_row_ids)
diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py
b/paimon-python/pypaimon/tests/vector_search_filter_test.py
index 2932ea8dac..1605d51a6c 100644
--- a/paimon-python/pypaimon/tests/vector_search_filter_test.py
+++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py
@@ -36,6 +36,7 @@ from pypaimon.globalindex.btree.btree_index_meta import
BTreeIndexMeta
from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta,
GlobalIndexMeta
from pypaimon.globalindex.global_index_reader import _completed_future
from pypaimon.globalindex.global_index_result import GlobalIndexResult
+from pypaimon.globalindex.vector_search import VectorSearch
from pypaimon.globalindex.vector_search_result import ScoredGlobalIndexResult
from pypaimon.index.index_file_meta import IndexFileMeta
from pypaimon.manifest.index_manifest_entry import IndexManifestEntry
@@ -355,6 +356,47 @@ class VectorReaderFactoryTest(unittest.TestCase):
reader.close()
+class VectorOptionsTest(unittest.TestCase):
+ """VectorSearch options compatibility."""
+
+ def test_offset_range_preserves_options(self):
+ search = VectorSearch(
+ vector=[1.0, 0.0],
+ limit=1,
+ field_name="embedding",
+ options={"ivf.nprobe": "16", "hnsw.ef_search": "64"},
+ )
+ include_row_ids = RoaringBitmap64()
+ include_row_ids.add_range(100, 200)
+
+ offset = search.with_include_row_ids(include_row_ids).offset_range(60,
150)
+
+ self.assertEqual(
+ {"ivf.nprobe": "16", "hnsw.ef_search": "64"},
+ offset.options,
+ )
+
+
+class LuminaOptionsTest(unittest.TestCase):
+ """Lumina query-time option compatibility."""
+
+ def test_query_options_override_index_options(self):
+ from pypaimon.globalindex.lumina.lumina_vector_global_index_reader
import (
+ _merge_options,
+ )
+
+ merged = _merge_options(
+ {"diskann.search.list_size": "16", "search.parallel_number": "2"},
+ {"diskann.search.list_size": "32", "index.dimension": "4"},
+ {"diskann.search.list_size": "64", "hnsw.ef_search": "128"},
+ )
+
+ self.assertEqual("64", merged["diskann.search.list_size"])
+ self.assertEqual("2", merged["search.parallel_number"])
+ self.assertEqual("4", merged["index.dimension"])
+ self.assertEqual("128", merged["hnsw.ef_search"])
+
+
class TantivyFullTextIndexOptionsTest(unittest.TestCase):
"""Tantivy full-text tokenizer metadata compatibility."""
@@ -883,6 +925,39 @@ class VectorSearchFilterTest(unittest.TestCase):
{"oss://bucket/vec-0.index", "oss://bucket/vec-1.index"},
seen_paths)
+ def test_read_threads_options_to_vector_search(self):
+ scan_plan = self._builder().new_vector_search_scan().scan()
+
+ captured_searches = []
+
+ def _capture_create(index_type, file_io, index_path,
+ index_io_meta_list, options=None):
+ class _FakeReader:
+ def visit_vector_search(self_inner, vs):
+ captured_searches.append(vs)
+ return
_completed_future(ScoredGlobalIndexResult.create_empty())
+
+ def close(self_inner):
+ pass
+
+ return _FakeReader()
+
+ with mock.patch(
+
"pypaimon.table.source.vector_search_read._create_vector_reader",
+ side_effect=_capture_create):
+ (self._builder()
+ .with_option("ivf.nprobe", "16")
+ .with_options({"hnsw.ef_search": "64"})
+ .new_vector_search_read()
+ .read_plan(scan_plan))
+
+ self.assertEqual(2, len(captured_searches))
+ for search in captured_searches:
+ self.assertEqual(
+ {"ivf.nprobe": "16", "hnsw.ef_search": "64"},
+ search.options,
+ )
+
def test_scanner_threads_external_path_to_btree_reader(self):
"""GlobalIndexScanner (backing _pre_filter) must thread external_path
onto the GlobalIndexIOMeta handed to the btree reader factory."""
diff --git
a/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/PaimonScanBuilderTest.scala
b/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/PaimonScanBuilderTest.scala
index bbd79c8d78..d55bb440aa 100755
---
a/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/PaimonScanBuilderTest.scala
+++
b/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/PaimonScanBuilderTest.scala
@@ -39,6 +39,10 @@ class PaimonScanBuilderTest extends PaimonSparkTestBase {
rows =
spark.sql("select id, embs from vector_search('T', 'embs', array(1.0f,
2.0f, 3.0f), 5)")
assert(rows.isEmpty)
+ rows = spark.sql(
+ "select id, embs from vector_search(" +
+ "'T', 'embs', array(1.0f, 2.0f, 3.0f), 5, map('ivf.nprobe', '16'))")
+ assert(rows.isEmpty)
}
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
index df60afb489..525ed334a8 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java
@@ -39,6 +39,7 @@ import org.apache.spark.broadcast.Broadcast;
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;
@@ -62,6 +63,16 @@ public class SparkVectorReadImpl extends VectorReadImpl {
super(table, filter, limit, vectorColumn, vector);
}
+ public SparkVectorReadImpl(
+ FileStoreTable table,
+ Predicate filter,
+ int limit,
+ DataField vectorColumn,
+ float[] vector,
+ Map<String, String> options) {
+ super(table, filter, limit, vectorColumn, vector, options);
+ }
+
@Override
public GlobalIndexResult read(List<VectorSearchSplit> splits) {
if (splits.isEmpty()) {
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
index 0e4e347f64..bd19a9e565 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorSearchBuilderImpl.java
@@ -36,6 +36,6 @@ public class SparkVectorSearchBuilderImpl extends
VectorSearchBuilderImpl {
@Override
public VectorRead newVectorRead() {
- return new SparkVectorReadImpl(table, filter, limit, vectorColumn,
vector);
+ return new SparkVectorReadImpl(table, filter, limit, vectorColumn,
vector, options);
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
index abff5ba3dc..d7f6cd23d3 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
@@ -89,6 +89,7 @@ abstract class PaimonBaseScan(table: InnerTable)
.withVector(vectorSearch.vector())
.withVectorColumn(vectorSearch.fieldName())
.withLimit(vectorSearch.limit())
+ .withOptions(vectorSearch.options())
if (pushedPartitionFilters.nonEmpty) {
vectorBuilder.withPartitionFilter(PartitionPredicate.and(pushedPartitionFilters.asJava))
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala
index fcc6738733..6dc0c77ced 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala
@@ -30,11 +30,13 @@ import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.FunctionIdentifier
import org.apache.spark.sql.catalyst.analysis.FunctionRegistryBase
import
org.apache.spark.sql.catalyst.analysis.TableFunctionRegistry.TableFunctionBuilder
-import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateArray,
Expression, ExpressionInfo, Literal}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateArray,
CreateMap, Expression, ExpressionInfo, Literal}
import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, LogicalPlan}
+import org.apache.spark.sql.catalyst.util.MapData
import org.apache.spark.sql.connector.catalog.{Identifier, Table, TableCatalog}
import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.unsafe.types.UTF8String
import scala.collection.JavaConverters._
@@ -292,11 +294,12 @@ case class IncrementalToAutoTag(override val args:
Seq[Expression])
/**
* Plan for the [[VECTOR_SEARCH]] table-valued function.
*
- * Usage: vector_search(table_name, column_name, query_vector, limit)
+ * Usage: vector_search(table_name, column_name, query_vector, limit[,
options])
* - table_name: the Paimon table to search
* - column_name: the vector column name
* - query_vector: array of floats representing the query vector
* - limit: the number of top results to return
+ * - options: optional options as a map or semicolon-separated key-value
string
*
* Example: SELECT * FROM vector_search('T', 'v', array(50.0f, 51.0f, 52.0f),
5)
*/
@@ -311,9 +314,10 @@ case class VectorSearchQuery(override val args:
Seq[Expression])
def createVectorSearch(
innerTable: InnerTable,
argsWithoutTable: Seq[Expression]): VectorSearch = {
- if (argsWithoutTable.size != 3) {
+ if (argsWithoutTable.size != 3 && argsWithoutTable.size != 4) {
throw new RuntimeException(
- s"$VECTOR_SEARCH needs three parameters after table_name: column_name,
query_vector, limit. " +
+ s"$VECTOR_SEARCH needs three or four parameters after table_name: " +
+ s"column_name, query_vector, limit[, options]. " +
s"Got ${argsWithoutTable.size} parameters after table_name."
)
}
@@ -325,7 +329,13 @@ case class VectorSearchQuery(override val args:
Seq[Expression])
}
val queryVector = extractQueryVector(argsWithoutTable(1))
val limit = parsePositiveLimit(argsWithoutTable(2).eval())
- new VectorSearch(queryVector, limit, columnName)
+ val options: Map[String, String] =
+ if (argsWithoutTable.size == 4) {
+ extractOptions(argsWithoutTable(3))
+ } else {
+ Map.empty[String, String]
+ }
+ new VectorSearch(queryVector, limit, columnName, options.asJava)
}
private def extractQueryVector(expr: Expression): Array[Float] = {
@@ -345,6 +355,69 @@ case class VectorSearchQuery(override val args:
Seq[Expression])
throw new RuntimeException(s"Cannot extract query vector from
expression: $expr")
}
}
+
+ private def extractOptions(expr: Expression): Map[String, String] = {
+ expr match {
+ case CreateMap(children, _) if children != null =>
+ children
+ .grouped(2)
+ .map {
+ case Seq(keyExpr, valueExpr) => (extractString(keyExpr),
extractString(valueExpr))
+ case other =>
+ throw new RuntimeException(s"Invalid options map entries:
$other")
+ }
+ .toMap
+ case _ =>
+ expr.eval() match {
+ case null => Map.empty
+ case options: MapData => mapDataToStringMap(options)
+ case options: java.util.Map[_, _] =>
+ options.asScala.map {
+ case (key, value) => (stringValue(key), stringValue(value))
+ }.toMap
+ case options: String => parseOptionsString(options)
+ case options: UTF8String => parseOptionsString(options.toString)
+ case other =>
+ throw new RuntimeException(
+ s"Invalid options type: ${other.getClass.getName}. " +
+ "Expected a map or semicolon-separated key-value string.")
+ }
+ }
+ }
+
+ private def mapDataToStringMap(mapData: MapData): Map[String, String] = {
+ val keys = mapData.keyArray().array
+ val values = mapData.valueArray().array
+ keys.indices.map(i => (stringValue(keys(i)), stringValue(values(i)))).toMap
+ }
+
+ private def parseOptionsString(options: String): Map[String, String] = {
+ if (options == null || options.trim.isEmpty) {
+ Map.empty
+ } else {
+ options
+ .split(";")
+ .map {
+ kvString =>
+ val kv = kvString.split("=", 2)
+ if (kv.length != 2) {
+ throw new IllegalArgumentException(
+ s"Invalid option '$kvString'. Please use format 'key=value'.")
+ }
+ (kv(0).trim, kv(1).trim)
+ }
+ .toMap
+ }
+ }
+
+ private def extractString(expr: Expression): String =
stringValue(expr.eval())
+
+ private def stringValue(value: Any): String = {
+ if (value == null) {
+ throw new IllegalArgumentException("Option key and value cannot be
null.")
+ }
+ value.toString
+ }
}
/**
diff --git
a/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/catalyst/plans/logical/VectorSearchQueryTest.scala
b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/catalyst/plans/logical/VectorSearchQueryTest.scala
new file mode 100644
index 0000000000..d3f1b8d265
--- /dev/null
+++
b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/catalyst/plans/logical/VectorSearchQueryTest.scala
@@ -0,0 +1,77 @@
+/*
+ * 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.spark.catalyst.plans.logical
+
+import org.apache.paimon.table.InnerTable
+import org.apache.paimon.types.{ArrayType, DataType, DataTypes, RowType}
+
+import org.apache.spark.sql.catalyst.expressions.{CreateArray, CreateMap,
Expression, Literal}
+import org.scalatest.funsuite.AnyFunSuite
+
+import java.lang.reflect.{InvocationHandler, Method, Proxy}
+
+/** Tests for [[VectorSearchQuery]]. */
+class VectorSearchQueryTest extends AnyFunSuite {
+
+ test("create vector search with string options") {
+ val vectorSearch = createVectorSearch(
+ Literal("v"),
+ CreateArray(Seq(Literal(1.0f), Literal(2.0f))),
+ Literal(5),
+ Literal("ivf.nprobe=16;hnsw.ef_search=64"))
+
+ assert(vectorSearch.options().get("ivf.nprobe") == "16")
+ assert(vectorSearch.options().get("hnsw.ef_search") == "64")
+ }
+
+ test("create vector search with map options") {
+ val vectorSearch = createVectorSearch(
+ Literal("v"),
+ CreateArray(Seq(Literal(1.0f), Literal(2.0f))),
+ Literal(5),
+ CreateMap(Seq(Literal("ivf.nprobe"), Literal("16"),
Literal("hnsw.ef_search"), Literal("64")))
+ )
+
+ assert(vectorSearch.options().get("ivf.nprobe") == "16")
+ assert(vectorSearch.options().get("hnsw.ef_search") == "64")
+ }
+
+ private def createVectorSearch(args: Expression*) =
+ VectorSearchQuery(Seq.empty).createVectorSearch(innerTable, args)
+
+ private val innerTable =
+ Proxy
+ .newProxyInstance(
+ classOf[InnerTable].getClassLoader,
+ Array(classOf[InnerTable]),
+ new InvocationHandler {
+ private val rowType =
+ RowType.of(Array[DataType](new ArrayType(DataTypes.FLOAT())),
Array[String]("v"))
+
+ override def invoke(proxy: Any, method: Method, args:
Array[AnyRef]): AnyRef = {
+ method.getName match {
+ case "name" => "T"
+ case "rowType" => rowType
+ case other => throw new UnsupportedOperationException(other)
+ }
+ }
+ }
+ )
+ .asInstanceOf[InnerTable]
+}
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/VectorSearchOptionsTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/VectorSearchOptionsTest.scala
new file mode 100644
index 0000000000..79b78a6146
--- /dev/null
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/VectorSearchOptionsTest.scala
@@ -0,0 +1,70 @@
+/*
+ * 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.spark.sql
+
+import org.apache.paimon.globalindex.testvector.TestVectorGlobalIndexerFactory
+import org.apache.paimon.spark.PaimonSparkTestBase
+
+/** Tests for vector search query-time options. */
+class VectorSearchOptionsTest extends PaimonSparkTestBase {
+
+ test("vector search forwards query options to reader") {
+ withTable("T") {
+ spark.sql("""
+ |CREATE TABLE T (id INT, v ARRAY<FLOAT>)
+ |TBLPROPERTIES (
+ | 'bucket' = '-1',
+ | 'global-index.row-count-per-shard' = '10000',
+ | 'row-tracking.enabled' = 'true',
+ | 'data-evolution.enabled' = 'true',
+ | 'test.vector.dimension' = '2',
+ | 'test.vector.required-option.key' = 'ivf.nprobe',
+ | 'test.vector.required-option.value' = '16')
+ |""".stripMargin)
+
+ spark.sql("""
+ |INSERT INTO T VALUES
+ | (0, array(1.0f, 0.0f)),
+ | (1, array(0.0f, 1.0f))
+ |""".stripMargin)
+
+ spark
+ .sql(s"CALL sys.create_global_index(table => 'test.T', index_column =>
'v', " +
+ s"index_type => '${TestVectorGlobalIndexerFactory.IDENTIFIER}')")
+ .collect()
+
+ intercept[Exception] {
+ spark
+ .sql("""
+ |SELECT id FROM vector_search('T', 'v', array(1.0f, 0.0f), 1)
+ |""".stripMargin)
+ .collect()
+ }
+
+ val result = spark
+ .sql("""
+ |SELECT id FROM vector_search(
+ | 'T', 'v', array(1.0f, 0.0f), 1, map('ivf.nprobe', '16'))
+ |""".stripMargin)
+ .collect()
+
+ assert(result.length == 1)
+ }
+ }
+}
