liangjie3138 commented on code in PR #8834:
URL: https://github.com/apache/paimon/pull/8834#discussion_r3643939483
##########
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:
DataEvolutionPartialWriteOperator uses a single low-level FileStoreScan,
which supports withSnapshot(snapshotId). Vector Search uses several high-level
scan paths: DataEvolutionVectorScan, Raw/refine reads through ReadBuilder, and
the final row lookup. These APIs do not expose a common snapshot parameter.
We could remove the copy after adding explicit snapshot support to both
VectorSearchBuilder and ReadBuilder, but that would require a broader Core API
change.
--
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]
