This is an automated email from the ASF dual-hosted git repository. yiguolei pushed a commit to branch branch-4.2 in repository https://gitbox.apache.org/repos/asf/doris.git
commit b16d11515056c141cd4dbe2b34c63e0dabb984f4 Author: Gabriel <[email protected]> AuthorDate: Sun Sep 20 14:40:51 2026 +0800 [feature](lance) Support prepared vector search on branch-4.1 (#68162) ### What problem does this PR solve? `vector_search("query_vector"=?, "top_k"=?)` currently converts the placeholders into the literal string `?`, so server-side prepared Lance vector queries fail during analysis. Preserve placeholders for `query_vector`, `top_k`, `offset`, and `filter`, and bind fresh property values on every EXECUTE. PREPARE resolves the output schema without requiring a query vector. Table, column, and search configuration remain constant. Each execution performs normal analysis and planning and resolves Lance metadata again; the retained statement never caches a bound vector or snapshot. This supports both indexed and flat searches and reuses the parsed statement. Forwarded executions preserve cached parameter types when the client omits the type table on repeated executions; the forwarding request includes the types needed by a fresh master-side PREPARE. It does not introduce an execution-plan cache or bypass the optimizer. Example with MySQL Connector/J and `useServerPrepStmts=true`: ```java try (PreparedStatement statement = connection.prepareStatement( "SELECT id, _distance FROM vector_search(" + "'table'='catalog.db.items', 'column'='embedding', " + "'query_vector'=?, 'top_k'=?, 'offset'=?, 'filter'=?)")) { for (String vector : vectors) { statement.setString(1, vector); // JSON array matching the vector column dimension. statement.setInt(2, 10); statement.setInt(3, 0); statement.setString(4, "id > 0"); try (ResultSet rows = statement.executeQuery()) { // Consume the result before executing the next vector. } } } ``` ### Release note Support server-side prepared parameters for Lance `vector_search` query vectors, top-k, offsets, and filters. ### Validation - Proxy PREPARE tests reproduce the stale planner context in both search modes before the fix; PREPARE now synchronizes the executor with the parser-created context before parameter decoding. - Parser and analyzer tests cover parameter retention, quoted question marks, case-insensitive properties, constant-only property rejection, repeated execution with different vectors and snapshots, mixed constants/parameters, and invalid or missing parameters. - A JDBC regression reuses one `ServerPreparedStatement` with A/B/A parameter changes and compares results with literal SQL for both indexed and flat searches. On live non-master FEs, it also prepares locally and forces repeated executions to the master. - A protocol test exercises five executions through the follower decoder, RPC request construction, Thrift serialization, and a fresh proxy decoder. It covers omitted type tables, unsigned values, NULL parameters across two bitmap bytes, and later parameter type changes. Local execution retains the original payload without copying it. - PREPARE and EXECUTE share single-vector and multi-vector schema validation. Tests cover schema-only PREPARE, multi-vector binding, and rejection of invalid dimensions and excessive candidate budgets at EXECUTE. The four CI FE UT failures caused by unbound-vector validation were reproduced before the fix. - Local verification: all 154 tests passed via `run-fe-ut.sh` across `MysqlPreparedStatementForwardingTest`, `MysqlConnectProcessorCursorFetchTest`, `FEOpExecutorMysqlProtocolTest`, `LancePreparedSearchTest`, `LancePreparedStatementTest`, `LanceVectorQueryTest`, `VectorSearchTableValuedFunctionTest`, `FullTextSearchTableValuedFunctionTest`, `LanceScanNodeTest`, `PrepareTest`, and `NereidsParserTest`. FE Checkstyle passed. The full external JDBC regression was also executed locally using the current PR FE build, a master FE, an observer FE, one BE, and the MinIO/Lance fixtures: 1 suite passed, 0 failures, 0 skips. It executed 12 real server-prepared queries covering direct and forwarded execution, indexed and flat search, and A/B/A parameter changes. The original script reproduced the CI String.positive() failure before fixing both JDBC URL concatenations. ### Check List (For Author) - Test - [x] Regression test - [x] Unit Test - Behavior changed: - [x] Yes. Supported `vector_search` property values can be prepared parameters. - Does this need documentation? - [x] Yes. Usage and the supported parameter scope are documented above. ### Check List (For Reviewer who merge this PR) - [ ] Confirm the release note - [ ] Confirm test cases - [ ] Confirm document - [ ] Add branch pick label --- .../doris/nereids/analyzer/UnboundTVFRelation.java | 33 ++- .../doris/nereids/parser/LogicalPlanBuilder.java | 29 ++- .../nereids/rules/analysis/BindExpression.java | 40 ++++ .../expressions/functions/table/VectorSearch.java | 9 +- .../trees/plans/commands/PrepareCommand.java | 3 + .../java/org/apache/doris/qe/ConnectContext.java | 5 + .../java/org/apache/doris/qe/FEOpExecutor.java | 25 +- .../org/apache/doris/qe/MysqlConnectProcessor.java | 10 + .../VectorSearchTableValuedFunction.java | 71 +++--- .../nereids/parser/LancePreparedStatementTest.java | 87 +++++++ .../qe/MysqlPreparedStatementForwardingTest.java | 148 ++++++++++++ .../tablefunction/LancePreparedSearchTest.java | 258 +++++++++++++++++++++ .../lance/test_lance_prepared_vector_search.groovy | 112 +++++++++ 13 files changed, 798 insertions(+), 32 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundTVFRelation.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundTVFRelation.java index ee8212b9049..eedb2460cd9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundTVFRelation.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundTVFRelation.java @@ -21,6 +21,8 @@ import org.apache.doris.nereids.exceptions.UnboundException; import org.apache.doris.nereids.memo.GroupExpression; import org.apache.doris.nereids.properties.LogicalProperties; import org.apache.doris.nereids.properties.UnboundLogicalProperties; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.Placeholder; import org.apache.doris.nereids.trees.expressions.Properties; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.functions.table.TableValuedFunction; @@ -33,7 +35,11 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalRelation; import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; import org.apache.doris.nereids.util.Utils; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; + import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -43,6 +49,7 @@ public class UnboundTVFRelation extends LogicalRelation implements TVFRelation, private final String functionName; private final Properties properties; + private final Map<String, Placeholder> propertyParameters; public UnboundTVFRelation(RelationId id, String functionName, Properties properties) { this(id, functionName, properties, Optional.empty(), Optional.empty()); @@ -50,7 +57,19 @@ public class UnboundTVFRelation extends LogicalRelation implements TVFRelation, public UnboundTVFRelation(RelationId id, String functionName, Properties properties, Optional<GroupExpression> groupExpression, Optional<LogicalProperties> logicalProperties) { + this(id, functionName, properties, ImmutableMap.of(), groupExpression, logicalProperties); + } + + public UnboundTVFRelation(RelationId id, String functionName, Properties properties, + Map<String, Placeholder> propertyParameters) { + this(id, functionName, properties, propertyParameters, Optional.empty(), Optional.empty()); + } + + private UnboundTVFRelation(RelationId id, String functionName, Properties properties, + Map<String, Placeholder> propertyParameters, Optional<GroupExpression> groupExpression, + Optional<LogicalProperties> logicalProperties) { super(id, PlanType.LOGICAL_UNBOUND_TVF_RELATION, groupExpression, logicalProperties); + this.propertyParameters = ImmutableMap.copyOf(propertyParameters); this.functionName = Objects.requireNonNull(functionName, "functionName can not be null"); this.properties = Objects.requireNonNull(properties, "properties can not be null"); } @@ -63,6 +82,15 @@ public class UnboundTVFRelation extends LogicalRelation implements TVFRelation, return properties; } + public Map<String, Placeholder> getPropertyParameters() { + return propertyParameters; + } + + @Override + public List<Expression> getExpressions() { + return ImmutableList.copyOf(propertyParameters.values()); + } + @Override public TableValuedFunction getFunction() { throw new UnboundException("getFunction"); @@ -85,14 +113,15 @@ public class UnboundTVFRelation extends LogicalRelation implements TVFRelation, @Override public Plan withGroupExpression(Optional<GroupExpression> groupExpression) { - return new UnboundTVFRelation(relationId, functionName, properties, groupExpression, + return new UnboundTVFRelation(relationId, functionName, properties, propertyParameters, groupExpression, Optional.of(getLogicalProperties())); } @Override public Plan withGroupExprLogicalPropChildren(Optional<GroupExpression> groupExpression, Optional<LogicalProperties> logicalProperties, List<Plan> children) { - return new UnboundTVFRelation(relationId, functionName, properties, groupExpression, logicalProperties); + return new UnboundTVFRelation(relationId, functionName, properties, propertyParameters, + groupExpression, logicalProperties); } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java index 5f6c72d0a5a..8804be4406c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java @@ -2875,9 +2875,34 @@ public class LogicalPlanBuilder extends DorisParserBaseVisitor<Object> { return ParserUtils.withOrigin(ctx, () -> { String functionName = ctx.tvfName.getText(); - Map<String, String> map = visitPropertyItemList(ctx.properties); + Map<String, Placeholder> parameters = new LinkedHashMap<>(); + Map<String, String> map; + if ("vector_search".equalsIgnoreCase(functionName) && ctx.properties != null) { + map = new HashMap<>(); + Set<String> keys = new HashSet<>(); + for (PropertyItemContext argument : ctx.properties.properties) { + if (argument.key.constant() instanceof DorisParser.PlaceholderContext) { + throw new AnalysisException("vector_search property names must be constant"); + } + String key = parsePropertyKey(argument.key).toLowerCase(Locale.ROOT); + if (!keys.add(key)) { + throw new AnalysisException("Duplicate vector_search property: " + key); + } + if (argument.value.constant() instanceof DorisParser.PlaceholderContext) { + if (!ImmutableSet.of("query_vector", "top_k", "offset", "filter").contains(key)) { + throw new AnalysisException("vector_search property '" + key + + "' must be constant in a prepared statement"); + } + parameters.put(key, (Placeholder) visit(argument.value.constant())); + } else { + map.put(key, parsePropertyValue(argument.value)); + } + } + } else { + map = visitPropertyItemList(ctx.properties); + } LogicalPlan relation = new UnboundTVFRelation(StatementScopeIdGenerator.newRelationId(), - functionName, new Properties(map)); + functionName, new Properties(map), parameters); return withTableAlias(relation, ctx.tableAlias()); }); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java index d145e09c875..02737980a3a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java @@ -54,6 +54,7 @@ import org.apache.doris.nereids.trees.expressions.EqualTo; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Placeholder; import org.apache.doris.nereids.trees.expressions.Properties; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; @@ -72,6 +73,8 @@ import org.apache.doris.nereids.trees.expressions.functions.table.FullTextSearch import org.apache.doris.nereids.trees.expressions.functions.table.TableValuedFunction; import org.apache.doris.nereids.trees.expressions.functions.table.VectorSearch; import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; import org.apache.doris.nereids.trees.plans.JoinType; import org.apache.doris.nereids.trees.plans.Plan; @@ -1777,6 +1780,39 @@ public class BindExpression implements AnalysisRuleFactory { String functionName = unboundTVFRelation.getFunctionName(); Properties arguments = unboundTVFRelation.getProperties(); + if (!unboundTVFRelation.getPropertyParameters().isEmpty()) { + // The unbound plan is retained across EXECUTEs. Never overwrite its parameter slots + // or cache a bound TVF, which would retain a previous vector and Lance snapshot. + Map<String, String> boundProperties = new HashMap<>(arguments.getMap()); + for (Map.Entry<String, Placeholder> parameter : unboundTVFRelation.getPropertyParameters().entrySet()) { + String key = parameter.getKey(); + if (statementContext.isPrepareStage()) { + // These values only determine the result schema; PREPARE does not execute a search. + switch (key) { + case "top_k": + boundProperties.put(key, "1"); + break; + case "offset": + boundProperties.put(key, "0"); + break; + case "filter": + boundProperties.put(key, "true"); + break; + default: + break; + } + } else { + Expression value = statementContext.getIdToPlaceholderRealExpr() + .get(parameter.getValue().getPlaceholderId()); + if (!(value instanceof Literal) || value instanceof NullLiteral) { + throw new AnalysisException("vector_search parameter '" + key + + "' must be a non-null literal"); + } + boundProperties.put(key, ((Literal) value).getStringValue()); + } + } + arguments = new Properties(boundProperties); + } FunctionBuilder functionBuilder = functionRegistry.findFunctionBuilder(functionName, arguments); Pair<? extends Expression, ? extends BoundFunction> bindResult = functionBuilder.build(functionName, arguments); @@ -1788,6 +1824,10 @@ public class BindExpression implements AnalysisRuleFactory { sqlCacheContext.get().setCannotProcessExpression(true); } TableValuedFunction tableValuedFunction = (TableValuedFunction) bindResult.first; + if (tableValuedFunction instanceof VectorSearch && statementContext.isPrepareStage() + && unboundTVFRelation.getPropertyParameters().containsKey("query_vector")) { + tableValuedFunction = new VectorSearch(arguments, true); + } LogicalTVFRelation relation = new LogicalTVFRelation( unboundTVFRelation.getRelationId(), tableValuedFunction, ImmutableList.of()); if (!(tableValuedFunction instanceof VectorSearch) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/table/VectorSearch.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/table/VectorSearch.java index 3fc5a1b945a..c94e754c05c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/table/VectorSearch.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/table/VectorSearch.java @@ -28,8 +28,15 @@ import java.util.Map; /** Lance vector_search relation TVF. */ public class VectorSearch extends TableValuedFunction { + private final boolean deferQueryVector; + public VectorSearch(Properties properties) { + this(properties, false); + } + + public VectorSearch(Properties properties, boolean deferQueryVector) { super(VectorSearchTableValuedFunction.NAME, properties); + this.deferQueryVector = deferQueryVector; } @Override @@ -41,7 +48,7 @@ public class VectorSearch extends TableValuedFunction { protected TableValuedFunctionIf toCatalogFunction() { try { Map<String, String> arguments = getTVFProperties().getMap(); - return new VectorSearchTableValuedFunction(arguments); + return new VectorSearchTableValuedFunction(arguments, deferQueryVector); } catch (Throwable t) { throw new AnalysisException("Can not build vector_search(): " + t.getMessage(), t); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PrepareCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PrepareCommand.java index af4230b7ff8..36b52deb8a9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PrepareCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PrepareCommand.java @@ -114,6 +114,9 @@ public class PrepareCommand extends Command { } StatementContext statementContext = ctx.getStatementContext(); statementContext.setPrepareStage(true); + // Forwarded EXECUTE reparses SQL into a new context before reconstructing PREPARE. + // The planner must use that context before the forwarded parameter packet is decoded. + executor.setStatementContext(statementContext); List<Slot> slots; if (logicalPlan instanceof Command) { if (logicalPlan instanceof InsertIntoTableCommand diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java index 620e4427348..0f3e092130c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java @@ -291,6 +291,11 @@ public class ConnectContext { @Setter private ByteBuffer prepareExecuteBuffer; + // Snapshot of cached types omitted by the current COM_STMT_EXECUTE packet. + @Getter + @Setter + private int[] prepareExecuteTypeCodes; + // Whether the current COM_STMT_EXECUTE requested a server-side read-only cursor. private boolean cursorFetchRequested; diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/FEOpExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/FEOpExecutor.java index 43a01af0530..fe0de92857d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/FEOpExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/FEOpExecutor.java @@ -47,6 +47,7 @@ import org.apache.thrift.TException; import org.apache.thrift.transport.TTransportException; import java.nio.ByteBuffer; +import java.nio.ByteOrder; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -214,7 +215,7 @@ public class FEOpExecutor { if (ctx.getCommand() == MysqlCommand.COM_STMT_EXECUTE) { if (null != ctx.getPrepareExecuteBuffer()) { - params.setPrepareExecuteBuffer(ctx.getPrepareExecuteBuffer()); + params.setPrepareExecuteBuffer(buildPrepareExecuteBuffer()); } params.setCursorFetchRequested(ctx.isCursorFetchRequested()); } @@ -232,6 +233,28 @@ public class FEOpExecutor { return params; } + private ByteBuffer buildPrepareExecuteBuffer() { + int[] typeCodes = ctx.getPrepareExecuteTypeCodes(); + if (typeCodes == null || typeCodes.length == 0) { + return ctx.getPrepareExecuteBuffer(); + } + // Every master RPC rebuilds PREPARE without cached types. Expand only for forwarding + // so local executions do not copy potentially large vector or binary parameter values. + ByteBuffer source = ctx.getPrepareExecuteBuffer().duplicate(); + ByteBuffer forwarded = ByteBuffer.allocate(Math.addExact(source.remaining(), typeCodes.length * 2)) + .order(ByteOrder.LITTLE_ENDIAN); + byte[] nullBitmap = new byte[(typeCodes.length + 7) / 8]; + source.get(nullBitmap); + source.get(); // Replace new_params_bind_flag=0 with a complete type table. + forwarded.put(nullBitmap).put((byte) 1); + for (int typeCode : typeCodes) { + forwarded.putChar((char) typeCode); + } + forwarded.put(source); + forwarded.flip(); + return forwarded; + } + public int getStatusCode() { if (result == null || !result.isSetStatusCode()) { return ErrorCode.ERR_UNKNOWN_ERROR.getCode(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java index eb20e67bdac..150bf613356 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java @@ -131,6 +131,7 @@ public class MysqlConnectProcessor extends ConnectProcessor { if (!ctx.isProxy()) { // An empty buffer still identifies a zero-parameter COM_STMT_EXECUTE when forwarding. ctx.setPrepareExecuteBuffer(packetBuf.duplicate()); + ctx.setPrepareExecuteTypeCodes(null); } if (paramCount > 0) { if (LOG.isDebugEnabled()) { @@ -152,6 +153,15 @@ public class MysqlConnectProcessor extends ConnectProcessor { // rewrite with new prepared statment with type info in placeholders prepCtx.command = prepareCommand.withPlaceholders(typedPlaceholders); prepareCommand = (PrepareCommand) prepCtx.command; + } else if (!ctx.isProxy()) { + // A new master proxy cannot reuse the follower's statement-local type cache. + int[] typeCodes = new int[paramCount]; + for (int i = 0; i < paramCount; i++) { + Placeholder parameter = prepareCommand.getPlaceholders().get(i); + typeCodes[i] = parameter.getMysqlColType().getCode() + | (parameter.isUnsigned() ? MysqlColType.UNSIGNED_MASK : 0); + } + ctx.setPrepareExecuteTypeCodes(typeCodes); } // parse param data for (int i = 0; i < paramCount; ++i) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java index 270bbe83d44..c394adbe099 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java @@ -70,10 +70,15 @@ public class VectorSearchTableValuedFunction extends LanceExternalSearchTableVal public VectorSearchTableValuedFunction(Map<String, String> properties) throws AnalysisException { - super(prepare(properties)); + this(properties, false); } - private static PreparedSearch prepare(Map<String, String> properties) + public VectorSearchTableValuedFunction(Map<String, String> properties, boolean deferQueryVector) + throws AnalysisException { + super(prepare(properties, deferQueryVector)); + } + + private static PreparedSearch prepare(Map<String, String> properties, boolean deferQueryVector) throws AnalysisException { Map<String, String> params = normalizeProperties(properties, PROPERTIES, NAME); boolean useIndex = !params.containsKey(USE_INDEX) @@ -85,8 +90,14 @@ public class VectorSearchTableValuedFunction extends LanceExternalSearchTableVal common.metadata().getSchema(), required(params, COLUMN, NAME), "vector"); int vectorFieldId = useIndex ? requireLanceFieldId(common.metadata(), vectorField) : -1; - TSearchVector queryVector = parseAndEncodeQueryVector( - vectorField, required(params, QUERY_VECTOR, NAME)); + // PREPARE needs the table schema, but the vector is supplied only at EXECUTE. + TSearchVector queryVector = null; + if (deferQueryVector) { + analyzeQueryVectorField(vectorField); + } else { + queryVector = parseAndEncodeQueryVector( + vectorField, required(params, QUERY_VECTOR, NAME)); + } TVectorSearchParams vectorParams = new TVectorSearchParams() .setColumn(vectorField.getName()) @@ -95,11 +106,13 @@ public class VectorSearchTableValuedFunction extends LanceExternalSearchTableVal .setOffset(common.offset()); // Pin the planner's default on every split; Lance otherwise inherits an index metric. vectorParams.setMetric(params.containsKey(METRIC) ? parseMetric(params.get(METRIC)) : TVectorMetric.L2); - validateMultiVectorBudget(queryVector, common.topK(), common.offset(), - params.containsKey(REFINE_FACTOR) ? parsePositiveInt(params.get(REFINE_FACTOR), REFINE_FACTOR) : 1); - - if (queryVector.isSetNumVectors() && vectorParams.getMetric() == TVectorMetric.HAMMING) { - throw new AnalysisException("Lance multi-vector search supports l2, cosine, and dot metrics"); + // Query-dependent checks need the bound vector; schema-only PREPARE has no vector yet. + if (!deferQueryVector) { + validateMultiVectorBudget(queryVector, common.topK(), common.offset(), + params.containsKey(REFINE_FACTOR) ? parsePositiveInt(params.get(REFINE_FACTOR), REFINE_FACTOR) : 1); + if (queryVector.isSetNumVectors() && vectorParams.getMetric() == TVectorMetric.HAMMING) { + throw new AnalysisException("Lance multi-vector search supports l2, cosine, and dot metrics"); + } } TExternalSearchRequest searchRequest = new TExternalSearchRequest() .setSchemaVersion(1) @@ -215,23 +228,7 @@ public class VectorSearchTableValuedFunction extends LanceExternalSearchTableVal static TSearchVector parseAndEncodeQueryVector(Field field, String json) throws AnalysisException { boolean multiVector = field.getType().getTypeID() == ArrowType.ArrowTypeID.List; - Field vectorField = field; - if (multiVector) { - if (hasExtension(field) || field.getDictionary() != null || field.getChildren().size() != 1) { - throw unsupportedVectorType(field); - } - vectorField = field.getChildren().get(0); - // Lance's multi-vector distance kernels do not consult inner validity bitmaps. - if (vectorField.isNullable() || vectorField.getChildren().size() != 1) { - throw new AnalysisException("Lance multi-vector columns require non-nullable subvectors"); - } - } - VectorEncodingSpec encodingSpec = analyzeVectorField(vectorField); - if (multiVector && encodingSpec.elementType != TVectorElementType.FLOAT16 - && encodingSpec.elementType != TVectorElementType.FLOAT32 - && encodingSpec.elementType != TVectorElementType.FLOAT64) { - throw unsupportedVectorType(field); - } + VectorEncodingSpec encodingSpec = analyzeQueryVectorField(field); JsonArray values = parseQueryVector(json, field, multiVector ? -1 : encodingSpec.dimension); int numVectors = multiVector ? values.size() : 1; if (multiVector) { @@ -268,6 +265,28 @@ public class VectorSearchTableValuedFunction extends LanceExternalSearchTableVal return query; } + private static VectorEncodingSpec analyzeQueryVectorField(Field field) throws AnalysisException { + boolean multiVector = field.getType().getTypeID() == ArrowType.ArrowTypeID.List; + Field vectorField = field; + if (multiVector) { + if (hasExtension(field) || field.getDictionary() != null || field.getChildren().size() != 1) { + throw unsupportedVectorType(field); + } + vectorField = field.getChildren().get(0); + // Lance's multi-vector distance kernels do not consult inner validity bitmaps. + if (vectorField.isNullable() || vectorField.getChildren().size() != 1) { + throw new AnalysisException("Lance multi-vector columns require non-nullable subvectors"); + } + } + VectorEncodingSpec encodingSpec = analyzeVectorField(vectorField); + if (multiVector && encodingSpec.elementType != TVectorElementType.FLOAT16 + && encodingSpec.elementType != TVectorElementType.FLOAT32 + && encodingSpec.elementType != TVectorElementType.FLOAT64) { + throw unsupportedVectorType(field); + } + return encodingSpec; + } + private static VectorEncodingSpec analyzeVectorField(Field field) throws AnalysisException { if (hasExtension(field) || field.getDictionary() != null || field.getType().getTypeID() != ArrowType.ArrowTypeID.FixedSizeList diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/LancePreparedStatementTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/LancePreparedStatementTest.java new file mode 100644 index 00000000000..bf6f10b47fe --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/LancePreparedStatementTest.java @@ -0,0 +1,87 @@ +// 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.doris.nereids.parser; + +import org.apache.doris.nereids.analyzer.UnboundTVFRelation; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Placeholder; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.qe.ConnectContext; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class LancePreparedStatementTest { + private ConnectContext previous; + + @BeforeEach + public void setUp() { + previous = ConnectContext.get(); + MemoTestUtils.createStatementContext(""); + } + + @AfterEach + public void tearDown() { + ConnectContext.remove(); + if (previous != null) { + previous.setThreadLocalInfo(); + } + } + + @Test + public void testRejectParameterizedTableAndColumn() { + for (String property : new String[] {"table", "column", "use_index"}) { + Assertions.assertThrows(AnalysisException.class, () -> new NereidsParser().parseSingle( + "select * from vector_search('" + property + "'=?)")); + } + } + + @Test + public void testRejectParameterizedKeysAndDuplicateProperties() { + Assertions.assertThrows(AnalysisException.class, () -> new NereidsParser().parseSingle( + "select * from vector_search(?='value')")); + Assertions.assertThrows(AnalysisException.class, () -> new NereidsParser().parseSingle( + "select * from vector_search('query_vector'=?, 'QUERY_VECTOR'='[1,2]')")); + } + + @Test + public void testQuotedQuestionMarkRemainsLiteral() { + LogicalPlan plan = new NereidsParser().parseSingle( + "select * from vector_search('table'='catalog.db.items', " + + "'column'='embedding', 'query_vector'='[1,2]', 'filter'=\"label = '?'\")"); + UnboundTVFRelation relation = plan.<UnboundTVFRelation>collectToList( + UnboundTVFRelation.class::isInstance).get(0); + Assertions.assertEquals("label = '?'", relation.getProperties().getMap().get("filter")); + Assertions.assertTrue(relation.getExpressions().isEmpty()); + } + + @Test + public void testVectorParameterSurvivesParsing() { + LogicalPlan plan = new NereidsParser().parseSingle( + "select * from vector_search('table'='catalog.db.items', " + + "'column'='embedding', 'QUERY_VECTOR'=?)"); + UnboundTVFRelation relation = plan.<UnboundTVFRelation>collectToList( + UnboundTVFRelation.class::isInstance).get(0); + Assertions.assertTrue(relation.getExpressions().stream() + .anyMatch(expression -> expression.anyMatch(Placeholder.class::isInstance)), + "The retained TVF must keep the placeholder instead of the text '?'"); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/MysqlPreparedStatementForwardingTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/MysqlPreparedStatementForwardingTest.java new file mode 100644 index 00000000000..0ab8e1eee6f --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/MysqlPreparedStatementForwardingTest.java @@ -0,0 +1,148 @@ +// 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.doris.qe; + +import org.apache.doris.analysis.UserIdentity; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.MysqlColType; +import org.apache.doris.mysql.MysqlCapability; +import org.apache.doris.mysql.MysqlCommand; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.trees.expressions.Placeholder; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.PlaceholderId; +import org.apache.doris.nereids.trees.plans.commands.PrepareCommand; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.system.SystemInfoService; +import org.apache.doris.thrift.TMasterOpRequest; +import org.apache.doris.thrift.TNetworkAddress; + +import org.apache.thrift.TDeserializer; +import org.apache.thrift.TSerializer; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.MockedConstruction; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +public class MysqlPreparedStatementForwardingTest { + @Test + public void testForwardedExecutionsRetainTypesAndNullBitmap() throws Exception { + ConnectContext follower = new ConnectContext(); + follower.setCommand(MysqlCommand.COM_STMT_EXECUTE); + follower.setCurrentUserIdentity(UserIdentity.ROOT); + follower.setCapability(MysqlCapability.DEFAULT_CAPABILITY); + Env env = Mockito.mock(Env.class); + Mockito.when(env.getSelfNode()).thenReturn(new SystemInfoService.HostInfo("127.0.0.1", 9010)); + PreparedStatementContext prepared = prepare(follower); + try (MockedConstruction<StmtExecutor> executors = Mockito.mockConstruction(StmtExecutor.class); + MockedStatic<AuditLogHelper> audit = Mockito.mockStatic(AuditLogHelper.class); + MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + for (int execution = 0; execution < 5; execution++) { + boolean newTypes = execution == 0 || execution == 3; + boolean wideInteger = execution >= 3; + ByteBuffer packet = packet(execution, newTypes, wideInteger); + byte[] original = Arrays.copyOf(packet.array(), packet.limit()); + new MysqlConnectProcessor(follower).handleExecute( + prepared.command, 7, prepared, packet, null); + Assertions.assertNotEquals(QueryState.MysqlStateType.ERR, follower.getState().getStateType(), + follower.getState().getErrorMessage()); + Assertions.assertEquals(packet.limit(), packet.position()); + Assertions.assertArrayEquals(original, Arrays.copyOf(packet.array(), packet.limit())); + + // Every master RPC reparses PREPARE, so it starts without cached parameter types. + ConnectContext master = new ConnectContext(null, true); + master.setCommand(MysqlCommand.COM_STMT_EXECUTE); + PreparedStatementContext proxyPrepared = prepare(master); + Assertions.assertSame(packet.array(), follower.getPrepareExecuteBuffer().array()); + TMasterOpRequest request = new FEOpExecutor(new TNetworkAddress("127.0.0.1", 9010), + prepared.command.getOriginalStmt(), follower, true).buildStmtForwardParams(); + TMasterOpRequest restored = new TMasterOpRequest(); + new TDeserializer().deserialize(restored, new TSerializer().serialize(request)); + ByteBuffer forwarded = ByteBuffer.wrap(restored.getPrepareExecuteBuffer()).order(ByteOrder.LITTLE_ENDIAN); + new MysqlConnectProcessor(master).handleExecute( + proxyPrepared.command, 7, proxyPrepared, forwarded, null); + Assertions.assertNotEquals(QueryState.MysqlStateType.ERR, master.getState().getStateType(), + master.getState().getErrorMessage()); + Assertions.assertFalse(forwarded.hasRemaining()); + Assertions.assertEquals((execution + 1) * 2, executors.constructed().size()); + for (int i = 0; i < 10; i++) { + PlaceholderId id = new PlaceholderId(i); + Literal expected = (Literal) prepared.statementContext.getIdToPlaceholderRealExpr().get(id); + Literal actual = (Literal) proxyPrepared.statementContext.getIdToPlaceholderRealExpr().get(id); + Assertions.assertEquals(expected.getDataType(), actual.getDataType()); + Assertions.assertEquals(expected.getStringValue(), actual.getStringValue()); + } + Literal integer = (Literal) proxyPrepared.statementContext.getIdToPlaceholderRealExpr() + .get(new PlaceholderId(1)); + Assertions.assertEquals(wideInteger ? "1099511627776" : "4294967295", integer.getStringValue()); + Assertions.assertEquals(!wideInteger, proxyPrepared.command.getPlaceholders().get(1).isUnsigned()); + } + } + } + + private PreparedStatementContext prepare(ConnectContext context) { + List<Placeholder> parameters = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + parameters.add(new Placeholder(new PlaceholderId(i))); + } + OriginStatement origin = new OriginStatement("select ?, ?, ?, ?, ?, ?, ?, ?, ?, ?", 0); + PrepareCommand command = new PrepareCommand("7", Mockito.mock(LogicalPlan.class), parameters, origin); + return new PreparedStatementContext(command, context, new StatementContext(context, origin), "7"); + } + + private ByteBuffer packet(int execution, boolean newTypes, boolean wideInteger) { + ByteBuffer packet = ByteBuffer.allocate(256).order(ByteOrder.LITTLE_ENDIAN); + packet.position(9); // The MySQL dispatcher has already consumed the execute header. + boolean withNulls = execution == 1; + packet.put((byte) (withNulls ? 4 : 0)); + packet.put((byte) (withNulls ? 1 : 0)); // Exercise both bytes of the null bitmap. + packet.put((byte) (newTypes ? 1 : 0)); + if (newTypes) { + packet.putChar((char) MysqlColType.MYSQL_TYPE_VARSTRING.getCode()); + packet.putChar((char) (wideInteger ? MysqlColType.MYSQL_TYPE_LONGLONG.getCode() + : MysqlColType.MYSQL_TYPE_LONG.getCode() | MysqlColType.UNSIGNED_MASK)); + for (int i = 2; i < 10; i++) { + packet.putChar((char) MysqlColType.MYSQL_TYPE_LONG.getCode()); + } + } + byte[] vector = (execution == 1 ? "[3,4]" : "[1,2]").getBytes(StandardCharsets.UTF_8); + packet.put((byte) vector.length).put(vector); + if (wideInteger) { + packet.putLong(1L << 40); + } else { + packet.putInt(-1); + } + for (int i = 2; i < 10; i++) { + if (!withNulls || (i != 2 && i != 8)) { + packet.putInt(execution + i); + } + } + packet.flip(); + packet.position(9); + return packet; + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/tablefunction/LancePreparedSearchTest.java b/fe/fe-core/src/test/java/org/apache/doris/tablefunction/LancePreparedSearchTest.java new file mode 100644 index 00000000000..20e72540a7f --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/tablefunction/LancePreparedSearchTest.java @@ -0,0 +1,258 @@ +// 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.doris.tablefunction; + +import org.apache.doris.analysis.TableName; +import org.apache.doris.analysis.UserIdentity; +import org.apache.doris.catalog.Env; +import org.apache.doris.common.Pair; +import org.apache.doris.common.jmockit.Deencapsulation; +import org.apache.doris.datasource.lance.LanceExternalTable; +import org.apache.doris.datasource.lance.metadata.LanceTableAccess; +import org.apache.doris.datasource.lance.metadata.LanceTableMetadata; +import org.apache.doris.mysql.MysqlCommand; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.glue.LogicalPlanAdapter; +import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.properties.PhysicalProperties; +import org.apache.doris.nereids.trees.expressions.Placeholder; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; +import org.apache.doris.nereids.trees.plans.commands.PrepareCommand; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.OriginStatement; +import org.apache.doris.qe.StmtExecutor; + +import com.google.common.collect.ImmutableMap; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicLong; + +public class LancePreparedSearchTest { + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void testPrepareAndRepeatedBindingUseFreshVectorAndSnapshot(boolean useIndex) { + ConnectContext previous = ConnectContext.get(); + ConnectContext context = MemoTestUtils.createConnectContext(); + AtomicLong version = new AtomicLong(42); + LanceExternalTable table = mockTable(version); + try (MockedStatic<LanceExternalSearchTableValuedFunction> lookup = Mockito.mockStatic( + LanceExternalSearchTableValuedFunction.class, Mockito.CALLS_REAL_METHODS)) { + lookup.when(() -> LanceExternalSearchTableValuedFunction.findLanceExternalTable( + Mockito.any(TableName.class))).thenReturn(table); + Pair<LogicalPlan, StatementContext> parsed = new NereidsParser().parseMultiple( + "select id, _distance from vector_search('table'='catalog.db.items', " + + "'column'='embedding', 'use_index'='" + useIndex + "', " + + "'query_vector'=?, 'top_k'=?, 'offset'=?, 'filter'=?)") + .get(0); + List<Placeholder> parameters = parsed.second.getPlaceholders(); + Assertions.assertEquals(4, parameters.size()); + StatementContext prepare = MemoTestUtils.createStatementContext(context, ""); + prepare.setPrepareStage(true); + VectorSearchTableValuedFunction prepared = analyze(parsed.first, prepare); + Assertions.assertEquals(3, prepared.getTableColumns().size()); + Assertions.assertFalse(prepared.getSearchRequest().getSearchQuery().getVectorSearch().isSetQueryVector()); + + LogicalPlan staticVector = new NereidsParser().parseMultiple( + "select * from vector_search('table'='catalog.db.items', 'column'='embedding', " + + "'query_vector'='[1,2]', 'top_k'=?)").get(0).first; + Assertions.assertTrue(analyze(staticVector, prepare).getSearchRequest() + .getSearchQuery().getVectorSearch().isSetQueryVector()); + + for (int i = 0; i < 2; i++) { + version.set(42 + i); + StatementContext execute = MemoTestUtils.createStatementContext(context, ""); + execute.getIdToPlaceholderRealExpr().put(parameters.get(0).getPlaceholderId(), + new StringLiteral(i == 0 ? "[1,2]" : "[3,4]")); + execute.getIdToPlaceholderRealExpr().put(parameters.get(1).getPlaceholderId(), new IntegerLiteral(3 + i)); + execute.getIdToPlaceholderRealExpr().put(parameters.get(2).getPlaceholderId(), new IntegerLiteral(i)); + execute.getIdToPlaceholderRealExpr().put(parameters.get(3).getPlaceholderId(), + new StringLiteral("id > " + i)); + VectorSearchTableValuedFunction function = analyze(parsed.first, execute); + Assertions.assertEquals(42 + i, function.getMetadata().getVersion()); + Assertions.assertEquals(3 + i, function.getTopK()); + Assertions.assertEquals(i, function.getOffset()); + Assertions.assertEquals("id > " + i, new String( + function.getSearchRequest().getSearchFilter().getPayload(), StandardCharsets.UTF_8)); + byte[] values = function.getSearchRequest().getSearchQuery().getVectorSearch().getQueryVector().getValues(); + Assertions.assertEquals(i == 0 ? 1.0f : 3.0f, + ByteBuffer.wrap(values).order(ByteOrder.LITTLE_ENDIAN).getFloat()); + for (String invalidVector : new String[] {"[1]", "not-json"}) { + execute.getIdToPlaceholderRealExpr().put(parameters.get(0).getPlaceholderId(), + new StringLiteral(invalidVector)); + Assertions.assertThrows(org.apache.doris.nereids.exceptions.AnalysisException.class, + () -> analyze(parsed.first, execute)); + } + execute.getIdToPlaceholderRealExpr().put(parameters.get(0).getPlaceholderId(), NullLiteral.INSTANCE); + Assertions.assertThrows(org.apache.doris.nereids.exceptions.AnalysisException.class, + () -> analyze(parsed.first, execute)); + execute.getIdToPlaceholderRealExpr().put(parameters.get(0).getPlaceholderId(), + new StringLiteral("[1,2]")); + execute.getIdToPlaceholderRealExpr().put(parameters.get(1).getPlaceholderId(), new IntegerLiteral(0)); + Assertions.assertThrows(org.apache.doris.nereids.exceptions.AnalysisException.class, + () -> analyze(parsed.first, execute)); + execute.getIdToPlaceholderRealExpr().put(parameters.get(1).getPlaceholderId(), new IntegerLiteral(3)); + execute.getIdToPlaceholderRealExpr().put(parameters.get(2).getPlaceholderId(), new IntegerLiteral(-1)); + Assertions.assertThrows(org.apache.doris.nereids.exceptions.AnalysisException.class, + () -> analyze(parsed.first, execute)); + execute.getIdToPlaceholderRealExpr().remove(parameters.get(0).getPlaceholderId()); + Assertions.assertThrows(org.apache.doris.nereids.exceptions.AnalysisException.class, + () -> analyze(parsed.first, execute)); + } + } finally { + ConnectContext.remove(); + if (previous != null) { + previous.setThreadLocalInfo(); + } + } + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void testProxyPrepareUsesParsedStatementContext(boolean useIndex) throws Exception { + ConnectContext previous = ConnectContext.get(); + ConnectContext context = new ConnectContext(null, true); + context.setEnv(Env.getCurrentEnv()); + context.setCurrentUserIdentity(UserIdentity.ROOT); + context.setThreadLocalInfo(); + context.setCommand(MysqlCommand.COM_STMT_PREPARE); + LanceExternalTable table = mockTable(new AtomicLong(42)); + try (MockedStatic<LanceExternalSearchTableValuedFunction> lookup = Mockito.mockStatic( + LanceExternalSearchTableValuedFunction.class, Mockito.CALLS_REAL_METHODS)) { + lookup.when(() -> LanceExternalSearchTableValuedFunction.findLanceExternalTable( + Mockito.any(TableName.class))).thenReturn(table); + OriginStatement sql = new OriginStatement( + "select id, _distance from vector_search('table'='catalog.db.items', " + + "'column'='embedding', 'use_index'='" + useIndex + "', " + + "'query_vector'=?, 'top_k'=?, 'offset'=?, 'filter'=?)", 0); + StmtExecutor proxy = new StmtExecutor(context, sql, true); + StatementContext constructorContext = context.getStatementContext(); + // A forwarded EXECUTE reconstructs PREPARE before decoding its parameter packet. + Deencapsulation.invoke(proxy, "parseByNereids"); + StatementContext parsedContext = context.getStatementContext(); + Assertions.assertNotSame(constructorContext, parsedContext); + Assertions.assertTrue(parsedContext.getIdToPlaceholderRealExpr().isEmpty()); + LogicalPlanAdapter parsed = (LogicalPlanAdapter) proxy.getParsedStmt(); + PrepareCommand command = new PrepareCommand("1", parsed.getLogicalPlan(), + parsedContext.getPlaceholders(), sql); + command.run(context, proxy); + Assertions.assertNotNull(context.getPreparedStementContext("1")); + Assertions.assertEquals(4, command.placeholderCount()); + Assertions.assertEquals(2, proxy.planPrepareStatementSlots().size()); + } finally { + ConnectContext.remove(); + if (previous != null) { + previous.setThreadLocalInfo(); + } + } + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void testMultiVectorPrepareDefersValueValidation(boolean useIndex) { + ConnectContext previous = ConnectContext.get(); + ConnectContext context = MemoTestUtils.createConnectContext(); + LanceExternalTable table = mockTable(new AtomicLong(42), true); + try (MockedStatic<LanceExternalSearchTableValuedFunction> lookup = Mockito.mockStatic( + LanceExternalSearchTableValuedFunction.class, Mockito.CALLS_REAL_METHODS)) { + lookup.when(() -> LanceExternalSearchTableValuedFunction.findLanceExternalTable( + Mockito.any(TableName.class))).thenReturn(table); + Pair<LogicalPlan, StatementContext> parsed = new NereidsParser().parseMultiple( + "select id, _distance from vector_search('table'='catalog.db.items', " + + "'column'='embedding', 'use_index'='" + useIndex + "', " + + "'query_vector'=?, 'top_k'=?)").get(0); + StatementContext prepare = MemoTestUtils.createStatementContext(context, ""); + prepare.setPrepareStage(true); + Assertions.assertFalse(analyze(parsed.first, prepare).getSearchRequest() + .getSearchQuery().getVectorSearch().isSetQueryVector()); + StatementContext execute = MemoTestUtils.createStatementContext(context, ""); + List<Placeholder> parameters = parsed.second.getPlaceholders(); + execute.getIdToPlaceholderRealExpr().put(parameters.get(0).getPlaceholderId(), + new StringLiteral("[[1,2],[3,4]]")); + execute.getIdToPlaceholderRealExpr().put(parameters.get(1).getPlaceholderId(), new IntegerLiteral(3)); + Assertions.assertEquals(2, analyze(parsed.first, execute).getSearchRequest() + .getSearchQuery().getVectorSearch().getQueryVector().getNumVectors()); + execute.getIdToPlaceholderRealExpr().put(parameters.get(1).getPlaceholderId(), new IntegerLiteral(100001)); + Exception budget = Assertions.assertThrows(org.apache.doris.nereids.exceptions.AnalysisException.class, + () -> analyze(parsed.first, execute)); + Assertions.assertTrue(budget.getMessage().contains("candidate budget")); + execute.getIdToPlaceholderRealExpr().put(parameters.get(1).getPlaceholderId(), new IntegerLiteral(3)); + execute.getIdToPlaceholderRealExpr().put(parameters.get(0).getPlaceholderId(), + new StringLiteral("[[1],[3,4]]")); + Exception dimension = Assertions.assertThrows(org.apache.doris.nereids.exceptions.AnalysisException.class, + () -> analyze(parsed.first, execute)); + Assertions.assertTrue(dimension.getMessage().contains("dimension")); + } finally { + ConnectContext.remove(); + if (previous != null) { + previous.setThreadLocalInfo(); + } + } + } + + private LanceExternalTable mockTable(AtomicLong version) { + return mockTable(version, false); + } + + private LanceExternalTable mockTable(AtomicLong version, boolean multiVector) { + LanceExternalTable table = Mockito.mock(LanceExternalTable.class); + Field vector = new Field("embedding", org.apache.arrow.vector.types.pojo.FieldType.nullable( + new ArrowType.FixedSizeList(2)), Collections.singletonList( + Field.nullable("item", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)))); + if (multiVector) { + vector = new Field("embedding", org.apache.arrow.vector.types.pojo.FieldType.nullable(ArrowType.List.INSTANCE), + Collections.singletonList(new Field("item", + org.apache.arrow.vector.types.pojo.FieldType.notNullable(new ArrowType.FixedSizeList(2)), + vector.getChildren()))); + } + Schema schema = new Schema(Arrays.asList(Field.nullable("id", new ArrowType.Int(64, true)), vector)); + Mockito.when(table.loadMetadataForSearch()).thenAnswer(invocation -> LanceTableMetadata.createSnapshotWithIndexes( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), + version.get(), schema, Collections.emptyList(), + ImmutableMap.of("id", 0, "embedding", 1), Collections.emptyList())); + Mockito.when(table.loadBasicMetadata()).thenAnswer(invocation -> table.loadMetadataForSearch()); + return table; + } + + private VectorSearchTableValuedFunction analyze(LogicalPlan plan, StatementContext statement) { + CascadesContext cascades = CascadesContext.initContext(statement, plan, PhysicalProperties.ANY); + cascades.newAnalyzer().analyze(); + LogicalTVFRelation relation = cascades.getRewritePlan().<LogicalTVFRelation>collectToList( + LogicalTVFRelation.class::isInstance).get(0); + return (VectorSearchTableValuedFunction) relation.getFunction().getCatalogFunction(); + } +} diff --git a/regression-test/suites/external_table_p0/lance/test_lance_prepared_vector_search.groovy b/regression-test/suites/external_table_p0/lance/test_lance_prepared_vector_search.groovy new file mode 100644 index 00000000000..6d70e55b11c --- /dev/null +++ b/regression-test/suites/external_table_p0/lance/test_lance_prepared_vector_search.groovy @@ -0,0 +1,112 @@ +// 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. + +suite("test_lance_prepared_vector_search", "p0,external") { + if (!"true".equalsIgnoreCase(context.config.otherConfigs.get("enableIcebergTest"))) { + logger.info("Lance prepared search requires the external MinIO fixtures") + return + } + String externalEnvIp = context.config.otherConfigs.get("externalEnvIp") + String minioPort = context.config.otherConfigs.get("iceberg_minio_port") + sql "DROP CATALOG IF EXISTS test_lance_prepared_vector_search" + sql """CREATE CATALOG test_lance_prepared_vector_search PROPERTIES ( + "type" = "lance", + "lance.catalog.type" = "filesystem", + "warehouse" = "s3://warehouse/lance", + "s3.endpoint" = "http://${externalEnvIp}:${minioPort}";, + "s3.access_key" = "admin", + "s3.secret_key" = "password", + "s3.region" = "us-east-1", + "use_path_style" = "true" + )""" + + def checkPreparedSearch = { boolean forwarded -> + def connection = context.getConnection() + connection.createStatement().withCloseable { control -> + // PREPARE must run locally; forwarding is enabled only after the statement exists. + control.execute("SET force_forward_all_queries=false") + if (forwarded) { + control.execute("SYNC") + } + sql "SET enable_file_scanner_v2 = true" + sql "SET enable_sql_cache = false" + sql "SET enable_query_cache = false" + for (String useIndex : ["false", "true"]) { + String fixedProperties = """'table'='test_lance_prepared_vector_search.doris.vs_ivf_pq_f32', + 'column'='embedding', 'metric'='l2', 'use_index'='${useIndex}', + 'nprobes'='4', 'refine_factor'='10'""" + def statement = prepareStatement("""SELECT row_id, _distance + FROM vector_search(${fixedProperties}, "query_vector"=?, "top_k"=?, "offset"=?, "filter"=?) + WHERE row_id > ? ORDER BY _distance, row_id""") + try { + assertTrue(statement instanceof com.mysql.cj.jdbc.ServerPreparedStatement) + // Reuse one server statement, including returning to the first vector after a different execution. + for (int start : [0, 1023, 0]) { + String vector = "[" + (0..<16).collect { it + start }.join(",") + "]" + int topK = start == 0 ? 3 : 2 + int offset = start == 0 ? 0 : 1 + String filter = start == 0 ? "row_id > 0" : "row_id > 1000" + statement.setString(1, vector) + statement.setInt(2, topK) + statement.setInt(3, offset) + statement.setString(4, filter) + statement.setInt(5, 0) + def expected = sql("""SELECT row_id, _distance + FROM vector_search(${fixedProperties}, "query_vector"="${vector}", + "top_k"="${topK}", "offset"="${offset}", "filter"="${filter}") + WHERE row_id > 0 ORDER BY _distance, row_id""") + assertEquals(topK, expected.size()) + if (forwarded) { + control.execute("SET force_forward_all_queries=true") + } + try { + assertEquals(expected, exec(statement)) + } finally { + if (forwarded) { + control.execute("SET force_forward_all_queries=false") + } + } + } + } finally { + statement.close() + } + } + } + } + + // Keep '+' on the preceding line so Groovy continues the expression instead of applying unary plus. + String url = getServerPrepareJdbcUrl(context.config.jdbcUrl, context.dbName) + + "&emulateUnsupportedPstmts=false" + connect(context.config.jdbcUser, context.config.jdbcPassword, url) { + checkPreparedSearch(false) + } + + def followers = sql_return_maparray("SHOW FRONTENDS").findAll { + it.IsMaster == "false" && it.Alive == "true" + } + if (followers.isEmpty()) { + logger.info("Skip Lance prepared forwarding coverage: no live non-master FE") + } + followers.each { fe -> + String followerUrl = getServerPrepareJdbcUrl( + "jdbc:mysql://${fe.Host}:${fe.QueryPort}/", context.dbName, false) + + "&emulateUnsupportedPstmts=false" + connect(context.config.jdbcUser, context.config.jdbcPassword, followerUrl) { + checkPreparedSearch(true) + } + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
