liangjie3138 commented on code in PR #8834:
URL: https://github.com/apache/paimon/pull/8834#discussion_r3654389481
##########
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 for the suggestion. I removed the Procedure-level table copy and now
pass the resolved snapshot through VectorSearchBuilderImpl.withSnapshot,
following FullTextSearchBuilderImpl.
--
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]
