steFaiz commented on code in PR #8834:
URL: https://github.com/apache/paimon/pull/8834#discussion_r3654062324
##########
paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/procedure/VectorSearchProcedure.java:
##########
@@ -90,89 +110,418 @@ public String[] call(
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);
+
+ fileStoreTable = resolveAndPinSnapshot(fileStoreTable);
+ if (fileStoreTable == null) {
+ return new String[0];
+ }
- GlobalIndexResult result =
- table.newVectorSearchBuilder()
+ VectorSearchBuilder builder =
+ newVectorSearchBuilder(procedureContext, fileStoreTable)
.withVector(queryVector)
.withVectorColumn(vectorColumn)
.withLimit(topK)
- .withOptions(optionsMap)
- .executeLocal();
+ .withOptions(optionsMap);
+ if (filter != null) {
+ builder.withFilter(filter);
+ }
+ if (partitionFilter != null) {
+ builder.withPartitionFilter(partitionFilter);
+ }
+
+ VectorScan.Plan vectorPlan = builder.newVectorScan().scan();
+ GlobalIndexResult result = builder.newVectorRead().read(vectorPlan);
+
+ 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);
+ }
+
+ private static VectorSearchBuilder newVectorSearchBuilder(
+ ProcedureContext context, FileStoreTable table) {
+ if (!table.coreOptions().vectorSearchDistributeEnabled()) {
+ return table.newVectorSearchBuilder();
+ }
+ return new FlinkVectorSearchBuilderImpl(table,
context.getExecutionEnvironment());
+ }
- RowType tableRowType = table.rowType();
- int[] projectionIndices = parseProjection(projection, tableRowType);
+ @Nullable
+ static FileStoreTable resolveAndPinSnapshot(FileStoreTable table) {
+ Snapshot snapshot = TimeTravelUtil.tryTravelOrLatest(table);
+ if (snapshot == null) {
+ return null;
+ }
+
+ Map<String, String> snapshotOptions = new LinkedHashMap<>();
+ snapshotOptions.put(SCAN_VERSION.key(), null);
+ snapshotOptions.put(SCAN_TAG_NAME.key(), null);
+ snapshotOptions.put(SCAN_WATERMARK.key(), null);
+ snapshotOptions.put(SCAN_TIMESTAMP.key(), null);
+ snapshotOptions.put(SCAN_TIMESTAMP_MILLIS.key(), null);
+ snapshotOptions.put(SCAN_MODE.key(),
CoreOptions.StartupMode.FROM_SNAPSHOT.toString());
+ snapshotOptions.put(SCAN_SNAPSHOT_ID.key(),
String.valueOf(snapshot.id()));
+ return table.copyWithoutTimeTravel(snapshotOptions);
Review Comment:
Thanks! Could you refer to FTS search to pin snapshot?
Just like:
org.apache.paimon.table.source.FullTextSearchBuilderImpl#withSnapshot
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
