This is an automated email from the ASF dual-hosted git repository. asf-gitbox-commits pushed a commit to branch PHOENIX-7998 in repository https://gitbox.apache.org/repos/asf/phoenix.git
commit a6cb4e45aa47e67586cda91e249edd85bab10f0f Author: Andrew Purtell <[email protected]> AuthorDate: Thu Sep 17 15:19:26 2026 -0700 PHOENIX-7998 Phase 3: Exact nearest-neighbor search path Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]> --- phoenix-core-client/src/main/antlr3/PhoenixSQL.g | 4 +- .../phoenix/compile/ExplainPlanAttributes.java | 26 +- .../apache/phoenix/compile/OrderByCompiler.java | 7 +- .../apache/phoenix/compile/SubselectRewriter.java | 10 +- .../org/apache/phoenix/iterate/ExplainTable.java | 17 +- .../iterate/MergeSortTopNResultIterator.java | 34 +- .../apache/phoenix/optimize/DistanceMetric.java | 66 +++ .../phoenix/optimize/VectorSearchDescriptor.java | 89 ++++ .../apache/phoenix/optimize/VectorSearchUtil.java | 213 ++++++++ .../phoenix/parse/DistanceFunctionParseNode.java | 47 ++ .../java/org/apache/phoenix/parse/OrderByNode.java | 62 ++- .../org/apache/phoenix/parse/ParseNodeFactory.java | 16 + .../apache/phoenix/parse/ParseNodeRewriter.java | 21 +- .../phoenix/end2end/VectorCrossRegionMergeIT.java | 371 +++++++++++++ .../phoenix/end2end/VectorDistancePushDownIT.java | 440 ++++++++++++++++ .../DistanceFunctionSerializationTest.java | 159 ++++++ .../expression/function/DistanceFunctionsTest.java | 44 ++ .../phoenix/optimize/VectorSearchUtilTest.java | 571 +++++++++++++++++++++ .../phoenix/query/explain/ExplainPlanTest.java | 1 + .../types/PArrayDataTypeVectorCoercionTest.java | 167 ++++++ 20 files changed, 2331 insertions(+), 34 deletions(-) diff --git a/phoenix-core-client/src/main/antlr3/PhoenixSQL.g b/phoenix-core-client/src/main/antlr3/PhoenixSQL.g index d7e8db78b6..31d274cb4a 100644 --- a/phoenix-core-client/src/main/antlr3/PhoenixSQL.g +++ b/phoenix-core-client/src/main/antlr3/PhoenixSQL.g @@ -1107,11 +1107,11 @@ order_by returns [List<OrderByNode> ret] //parse the individual field for an order by clause parseOrderByField returns [OrderByNode ret] -@init{boolean isAscending = true; boolean nullsLast = false;} +@init{boolean isAscending = true; Boolean nullsLast = null;} : (expr = expression) (ASC {isAscending = true;} | DESC {isAscending = false;})? (NULLS (FIRST {nullsLast = false;} | LAST {nullsLast = true;}))? - { $ret = factory.orderBy(expr, nullsLast, isAscending); } + { $ret = nullsLast == null ? factory.orderBy(expr, isAscending) : factory.orderBy(expr, nullsLast, isAscending); } ; parseFrom returns [TableNode ret] diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/compile/ExplainPlanAttributes.java b/phoenix-core-client/src/main/java/org/apache/phoenix/compile/ExplainPlanAttributes.java index ed71b4c873..ef785bed11 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/compile/ExplainPlanAttributes.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/compile/ExplainPlanAttributes.java @@ -49,12 +49,12 @@ import org.apache.phoenix.schema.PColumn; "serverDistinctFilter", "serverMergeColumns", "serverParsedProjections", "serverProject", "serverFilters", "ignoredHints", "serverFirstKeyOnlyProjection", "serverEmptyColumnOnlyProjection", "serverAggregate", "serverGroupByLimit", "serverSortedBy", - "serverOffset", "serverRowLimit", "clientFilterBy", "clientFilters", "clientAggregate", - "clientDistinctFilter", "clientAfterAggregate", "clientSortAlgo", "clientSortedBy", - "clientOffset", "clientRowLimit", "clientSequenceCount", "clientCursorName", "clientSteps", - "lhsJoinQueryExplainPlan", "rhsJoinQueryExplainPlan", "subPlans", "dynamicServerFilter", - "afterJoinFilter", "joinScannerLimit", "sortMergeSkipMerge", "regionLocations", - "regionLocationsTotalSize", "numRegionLocationLookups" }) + "vectorSearch", "serverOffset", "serverRowLimit", "clientFilterBy", "clientFilters", + "clientAggregate", "clientDistinctFilter", "clientAfterAggregate", "clientSortAlgo", + "clientSortedBy", "clientOffset", "clientRowLimit", "clientSequenceCount", "clientCursorName", + "clientSteps", "lhsJoinQueryExplainPlan", "rhsJoinQueryExplainPlan", "subPlans", + "dynamicServerFilter", "afterJoinFilter", "joinScannerLimit", "sortMergeSkipMerge", + "regionLocations", "regionLocationsTotalSize", "numRegionLocationLookups" }) public class ExplainPlanAttributes { // Top-of-plan disclosures (populated only on the root plan) @@ -111,6 +111,8 @@ public class ExplainPlanAttributes { private final String serverAggregate; private final Integer serverGroupByLimit; private final String serverSortedBy; + // True if the plan is a top-K vector similarity search. + private final boolean isVectorSearch; private final Integer serverOffset; private final Long serverRowLimit; @@ -208,6 +210,7 @@ public class ExplainPlanAttributes { this.serverAggregate = b.serverAggregate; this.serverGroupByLimit = b.serverGroupByLimit; this.serverSortedBy = b.serverSortedBy; + this.isVectorSearch = b.isVectorSearch; this.serverOffset = b.serverOffset; this.serverRowLimit = b.serverRowLimit; this.clientFilterBy = b.clientFilterBy; @@ -441,6 +444,10 @@ public class ExplainPlanAttributes { return serverSortedBy; } + public boolean isVectorSearch() { + return isVectorSearch; + } + public Integer getServerOffset() { return serverOffset; } @@ -642,6 +649,7 @@ public class ExplainPlanAttributes { private String serverAggregate; private Integer serverGroupByLimit; private String serverSortedBy; + private boolean isVectorSearch; private Integer serverOffset; private Long serverRowLimit; private String clientFilterBy; @@ -729,6 +737,7 @@ public class ExplainPlanAttributes { this.serverAggregate = explainPlanAttributes.getServerAggregate(); this.serverGroupByLimit = explainPlanAttributes.getServerGroupByLimit(); this.serverSortedBy = explainPlanAttributes.getServerSortedBy(); + this.isVectorSearch = explainPlanAttributes.isVectorSearch(); this.serverOffset = explainPlanAttributes.getServerOffset(); this.serverRowLimit = explainPlanAttributes.getServerRowLimit(); this.clientFilterBy = explainPlanAttributes.getClientFilterBy(); @@ -1026,6 +1035,11 @@ public class ExplainPlanAttributes { return this; } + public ExplainPlanAttributesBuilder setVectorSearch(boolean isVectorSearch) { + this.isVectorSearch = isVectorSearch; + return this; + } + public ExplainPlanAttributesBuilder setServerOffset(Integer serverOffset) { this.serverOffset = serverOffset; return this; diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/compile/OrderByCompiler.java b/phoenix-core-client/src/main/java/org/apache/phoenix/compile/OrderByCompiler.java index 101eb03510..19f74a1111 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/compile/OrderByCompiler.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/compile/OrderByCompiler.java @@ -29,6 +29,7 @@ import org.apache.phoenix.exception.SQLExceptionCode; import org.apache.phoenix.exception.SQLExceptionInfo; import org.apache.phoenix.expression.Expression; import org.apache.phoenix.expression.OrderByExpression; +import org.apache.phoenix.expression.function.DistanceFunction; import org.apache.phoenix.iterate.OrderedResultIterator; import org.apache.phoenix.parse.HintNode.Hint; import org.apache.phoenix.parse.OrderByNode; @@ -174,7 +175,11 @@ public class OrderByCompiler { } if (!expression.isStateless()) { boolean isAscending = node.isAscending(); - boolean isNullsLast = node.isNullsLast(); + // An ascending distance without a NULLS clause sorts nulls last. OrderByNode#isNullsLast + // applies this rule to a distance parse node. This check applies it to an ordinal after a + // wildcard, which resolves to a distance only at compilation. + boolean isNullsLast = node.isNullsLast() + || node.isNullsDefault() && isAscending && expression instanceof DistanceFunction; if (tracker != null) { tracker.track(expression, isAscending, isNullsLast); } diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/compile/SubselectRewriter.java b/phoenix-core-client/src/main/java/org/apache/phoenix/compile/SubselectRewriter.java index af56113125..d66cdcc846 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/compile/SubselectRewriter.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/compile/SubselectRewriter.java @@ -187,8 +187,8 @@ public class SubselectRewriter extends ParseNodeRewriter { rewrittenOrderByNodes = new ArrayList<OrderByNode>(orderByNodes.size()); for (OrderByNode orderByNode : orderByNodes) { ParseNode parseNode = orderByNode.getNode(); - rewrittenOrderByNodes.add(NODE_FACTORY.orderBy(parseNode.accept(subselectRewriter), - orderByNode.isNullsLast(), orderByNode.isAscending())); + rewrittenOrderByNodes + .add(NODE_FACTORY.orderBy(parseNode.accept(subselectRewriter), orderByNode)); } } @@ -394,8 +394,7 @@ public class SubselectRewriter extends ParseNodeRewriter { for (OrderByNode outerOrderByNode : outerOrderByNodes) { ParseNode outerOrderByParseNode = outerOrderByNode.getNode(); OrderByNode rewrittenOuterOrderByNode = - NODE_FACTORY.orderBy(outerOrderByParseNode.accept(this), outerOrderByNode.isNullsLast(), - outerOrderByNode.isAscending()); + NODE_FACTORY.orderBy(outerOrderByParseNode.accept(this), outerOrderByNode); assert innerOrderByNodeIter.hasNext(); OrderByNode innerOrderByNode = innerOrderByNodeIter.next(); if (!innerOrderByNode.equals(rewrittenOuterOrderByNode)) { @@ -497,8 +496,7 @@ public class SubselectRewriter extends ParseNodeRewriter { orderByRewrite = Lists.newArrayListWithExpectedSize(orderBy.size()); for (OrderByNode orderByNode : orderBy) { ParseNode node = orderByNode.getNode(); - orderByRewrite.add(NODE_FACTORY.orderBy(node.accept(this), orderByNode.isNullsLast(), - orderByNode.isAscending())); + orderByRewrite.add(NODE_FACTORY.orderBy(node.accept(this), orderByNode)); } } diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/iterate/ExplainTable.java b/phoenix-core-client/src/main/java/org/apache/phoenix/iterate/ExplainTable.java index ac6d927dd6..b0dbb8b6a2 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/iterate/ExplainTable.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/iterate/ExplainTable.java @@ -57,6 +57,7 @@ import org.apache.phoenix.filter.EmptyColumnOnlyFilter; import org.apache.phoenix.optimize.OptimizerDecision; import org.apache.phoenix.optimize.OptimizerReasons; import org.apache.phoenix.optimize.RejectedIndexEntry; +import org.apache.phoenix.optimize.VectorSearchUtil; import org.apache.phoenix.parse.HintNode; import org.apache.phoenix.parse.HintNode.Hint; import org.apache.phoenix.query.KeyRange; @@ -508,14 +509,24 @@ public abstract class ExplainTable { } if (!orderBy.getOrderByExpressions().isEmpty() && groupBy.isEmpty()) { // with GROUP BY, sort // happens client-side - String orderByExpressions = - "SERVER" + (limit == null ? "" : " TOP " + limit + " ROW" + (limit == 1 ? "" : "S")) - + " SORTED BY " + orderBy.getOrderByExpressions().toString(); + boolean isVectorSearch = VectorSearchUtil.isVectorSearch(orderBy, limit); + String orderByExpressions; + if (isVectorSearch) { + Expression distanceExpr = orderBy.getOrderByExpressions().get(0).getExpression(); + orderByExpressions = "SERVER TOP-" + limit + " BY " + distanceExpr; + } else { + orderByExpressions = + "SERVER" + (limit == null ? "" : " TOP " + limit + " ROW" + (limit == 1 ? "" : "S")) + + " SORTED BY " + orderBy.getOrderByExpressions().toString(); + } planSteps.add(" " + orderByExpressions); if (explainPlanAttributesBuilder != null) { if (limit != null) { explainPlanAttributesBuilder.setServerRowLimit(limit.longValue()); } + if (isVectorSearch) { + explainPlanAttributesBuilder.setVectorSearch(true); + } explainPlanAttributesBuilder.setServerSortedBy(orderBy.getOrderByExpressions().toString()); } } else { diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/iterate/MergeSortTopNResultIterator.java b/phoenix-core-client/src/main/java/org/apache/phoenix/iterate/MergeSortTopNResultIterator.java index 767656a950..089cc2432f 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/iterate/MergeSortTopNResultIterator.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/iterate/MergeSortTopNResultIterator.java @@ -23,6 +23,7 @@ import org.apache.hadoop.hbase.io.ImmutableBytesWritable; import org.apache.phoenix.compile.ExplainPlanAttributes.ExplainPlanAttributesBuilder; import org.apache.phoenix.expression.Expression; import org.apache.phoenix.expression.OrderByExpression; +import org.apache.phoenix.optimize.VectorSearchUtil; import org.apache.phoenix.schema.tuple.Tuple; /** @@ -98,14 +99,24 @@ public class MergeSortTopNResultIterator extends MergeSortResultIterator { return super.next(); } + /** + * Returns true if EXPLAIN shows the merge and the limit as one TOP-K step. This is true for a + * top-K vector search without an offset. A search with an offset keeps separate merge, offset and + * limit steps. + */ + private boolean isTopKMerge() { + return offset <= 0 && VectorSearchUtil.isVectorSearch(orderByColumns, limit); + } + @Override public void explain(List<String> planSteps) { resultIterators.explain(planSteps); - planSteps.add("CLIENT MERGE SORT"); + boolean topKMerge = isTopKMerge(); + planSteps.add(topKMerge ? "CLIENT MERGE SORT TOP-" + limit : "CLIENT MERGE SORT"); if (offset > 0) { planSteps.add("CLIENT OFFSET " + offset); } - if (limit > 0) { + if (limit > 0 && !topKMerge) { planSteps.add("CLIENT LIMIT " + limit); } } @@ -114,9 +125,14 @@ public class MergeSortTopNResultIterator extends MergeSortResultIterator { public void explain(List<String> planSteps, ExplainPlanAttributesBuilder explainPlanAttributesBuilder) { resultIterators.explain(planSteps, explainPlanAttributesBuilder); - explainPlanAttributesBuilder.setClientSortAlgo("CLIENT MERGE SORT"); - planSteps.add("CLIENT MERGE SORT"); - explainPlanAttributesBuilder.addClientStep("CLIENT MERGE SORT"); + boolean topKMerge = isTopKMerge(); + String mergeStep = topKMerge ? "CLIENT MERGE SORT TOP-" + limit : "CLIENT MERGE SORT"; + if (topKMerge) { + explainPlanAttributesBuilder.setVectorSearch(true); + } + explainPlanAttributesBuilder.setClientSortAlgo(mergeStep); + planSteps.add(mergeStep); + explainPlanAttributesBuilder.addClientStep(mergeStep); if (offset > 0) { explainPlanAttributesBuilder.setClientOffset(offset); String step = "CLIENT OFFSET " + offset; @@ -125,9 +141,11 @@ public class MergeSortTopNResultIterator extends MergeSortResultIterator { } if (limit > 0) { explainPlanAttributesBuilder.setClientRowLimit(limit); - String step = "CLIENT LIMIT " + limit; - planSteps.add(step); - explainPlanAttributesBuilder.addClientStep(step); + if (!topKMerge) { + String step = "CLIENT LIMIT " + limit; + planSteps.add(step); + explainPlanAttributesBuilder.addClientStep(step); + } } } diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/DistanceMetric.java b/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/DistanceMetric.java new file mode 100644 index 0000000000..094ed30256 --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/DistanceMetric.java @@ -0,0 +1,66 @@ +/* + * 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.phoenix.optimize; + +import org.apache.phoenix.expression.function.CosineDistanceFunction; +import org.apache.phoenix.expression.function.InnerProductDistanceFunction; +import org.apache.phoenix.expression.function.L2DistanceFunction; +import org.apache.phoenix.expression.function.L2DistanceSquaredFunction; + +/** The vector distance metrics that Phoenix supports. */ +public enum DistanceMetric { + L2, + COSINE, + INNER_PRODUCT; + + /** + * Returns the metric with the given name, or null if the name is null or unknown. The match + * ignores case and whitespace at the start or end of the name. + */ + public static DistanceMetric fromString(String name) { + if (name == null) { + return null; + } + try { + return valueOf(name.trim().toUpperCase()); + } catch (IllegalArgumentException e) { + return null; + } + } + + /** + * Returns the metric that a distance function ranks by, or null if the name is null or is not a + * distance function. L2_DISTANCE and L2_DISTANCE_SQUARED both rank by L2. + */ + public static DistanceMetric forFunction(String functionName) { + if (functionName == null) { + return null; + } + switch (functionName) { + case L2DistanceFunction.NAME: + case L2DistanceSquaredFunction.NAME: + return L2; + case CosineDistanceFunction.NAME: + return COSINE; + case InnerProductDistanceFunction.NAME: + return INNER_PRODUCT; + default: + return null; + } + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/VectorSearchDescriptor.java b/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/VectorSearchDescriptor.java new file mode 100644 index 0000000000..a01b57c3a2 --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/VectorSearchDescriptor.java @@ -0,0 +1,89 @@ +/* + * 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.phoenix.optimize; + +import java.util.Objects; +import org.apache.phoenix.parse.ColumnParseNode; +import org.apache.phoenix.parse.FunctionParseNode; +import org.apache.phoenix.parse.ParseNode; + +/** + * Describes a vector similarity search in a parsed query. It holds the distance metric, the + * distance node, the source vector expression, the query vector expression and the literal row + * limit. + */ +public class VectorSearchDescriptor { + + private final DistanceMetric metric; + private final FunctionParseNode distanceNode; + private final ParseNode sourceExpression; + private final ParseNode queryExpression; + private final Integer limit; + + public VectorSearchDescriptor(DistanceMetric metric, FunctionParseNode distanceNode, + ParseNode sourceExpression, ParseNode queryExpression, Integer limit) { + this.metric = Objects.requireNonNull(metric, "metric cannot be null"); + this.distanceNode = Objects.requireNonNull(distanceNode, "distanceNode cannot be null"); + this.sourceExpression = + Objects.requireNonNull(sourceExpression, "sourceExpression cannot be null"); + this.queryExpression = + Objects.requireNonNull(queryExpression, "queryExpression cannot be null"); + this.limit = limit; + } + + public DistanceMetric getMetric() { + return metric; + } + + public FunctionParseNode getDistanceNode() { + return distanceNode; + } + + public String getDistanceFunctionName() { + return distanceNode.getName(); + } + + public ParseNode getSourceExpression() { + return sourceExpression; + } + + public boolean isSourceColumn() { + return sourceExpression instanceof ColumnParseNode; + } + + /** Returns the source column name if the source expression is a column reference, or null. */ + public String getSourceColumnName() { + return isSourceColumn() ? ((ColumnParseNode) sourceExpression).getName() : null; + } + + public ParseNode getQueryExpression() { + return queryExpression; + } + + /** Returns the row limit, or null if the LIMIT is not a literal, such as a bind parameter. */ + public Integer getLimit() { + return limit; + } + + @Override + public String toString() { + return "VectorSearchDescriptor{function=" + distanceNode.getName() + ", metric=" + metric + + ", sourceExpression=" + sourceExpression + ", queryExpression=" + queryExpression + + ", limit=" + limit + '}'; + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/VectorSearchUtil.java b/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/VectorSearchUtil.java new file mode 100644 index 0000000000..dbeb77a41f --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/VectorSearchUtil.java @@ -0,0 +1,213 @@ +/* + * 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.phoenix.optimize; + +import java.util.List; +import org.apache.phoenix.compile.OrderByCompiler.OrderBy; +import org.apache.phoenix.expression.OrderByExpression; +import org.apache.phoenix.expression.function.DistanceFunction; +import org.apache.phoenix.parse.BindParseNode; +import org.apache.phoenix.parse.ColumnParseNode; +import org.apache.phoenix.parse.DistanceFunctionParseNode; +import org.apache.phoenix.parse.FunctionParseNode; +import org.apache.phoenix.parse.LimitNode; +import org.apache.phoenix.parse.LiteralParseNode; +import org.apache.phoenix.parse.OrderByNode; +import org.apache.phoenix.parse.ParseNode; +import org.apache.phoenix.parse.SelectStatement; + +/** + * Finds vector similarity search queries and extracts their parameters. + * <p> + * A vector similarity search has a LIMIT and one ascending ORDER BY item. A literal LIMIT must be + * positive. The ORDER BY item is a distance between a source vector expression that references a + * column and a constant or bind parameter query vector. The item can also be an ordinal that refers + * to such a distance in the projection. + */ +public final class VectorSearchUtil { + + private VectorSearchUtil() { + } + + /** + * Returns the expression that an ORDER BY item ranks by. For an ordinal, this is the projected + * expression that the ordinal refers to. If the position of the ordinal is not known before + * compilation, this is the item itself. + */ + static ParseNode getOrderByExpression(SelectStatement statement, OrderByNode orderByNode) { + ParseNode selectNode = orderByNode.getOrdinalSelectNode(statement.getSelect()); + return selectNode != null ? selectNode : orderByNode.getNode(); + } + + /** Returns true if the statement is a vector similarity search. */ + public static boolean isVectorSearchQuery(SelectStatement statement) { + return getVectorSearchDescriptor(statement) != null; + } + + /** + * Returns a {@link VectorSearchDescriptor} for the parsed statement, or null if the statement is + * not a vector similarity search. + */ + public static VectorSearchDescriptor getVectorSearchDescriptor(SelectStatement statement) { + if (statement == null) { + return null; + } + + LimitNode limitNode = statement.getLimit(); + if (limitNode == null) { + return null; + } + Integer limitValue = null; + ParseNode limitParseNode = limitNode.getLimitParseNode(); + if (limitParseNode instanceof LiteralParseNode) { + Object value = ((LiteralParseNode) limitParseNode).getValue(); + if (value instanceof Number) { + int l = ((Number) value).intValue(); + if (l <= 0) { + return null; + } + limitValue = l; + } + } + + List<OrderByNode> orderByNodes = statement.getOrderBy(); + if (orderByNodes == null || orderByNodes.size() != 1) { + return null; + } + OrderByNode orderByNode = orderByNodes.get(0); + + if (!orderByNode.isAscending()) { + return null; + } + + ParseNode exprNode = getOrderByExpression(statement, orderByNode); + if (!isDistanceFunctionNode(exprNode)) { + return null; + } + + List<ParseNode> children = exprNode.getChildren(); + if (children == null || children.size() < 2) { + return null; + } + + ParseNode arg0 = children.get(0); + ParseNode arg1 = children.get(1); + + ParseNode sourceExpr; + ParseNode queryExpr; + + boolean arg0Source = isSourceVectorNode(arg0); + boolean arg0Query = isQueryVectorNode(arg0); + boolean arg1Source = isSourceVectorNode(arg1); + boolean arg1Query = isQueryVectorNode(arg1); + + if (arg0Source && arg1Query && !arg0Query) { + sourceExpr = arg0; + queryExpr = arg1; + } else if (arg1Source && arg0Query && !arg1Query) { + sourceExpr = arg1; + queryExpr = arg0; + } else { + return null; + } + + DistanceMetric metric = getMetric(exprNode); + return new VectorSearchDescriptor(metric, (FunctionParseNode) exprNode, sourceExpr, queryExpr, + limitValue); + } + + /** Returns true if the parse node is a vector distance function or distance operator. */ + public static boolean isDistanceFunctionNode(ParseNode node) { + return node instanceof DistanceFunctionParseNode; + } + + /** Returns the {@link DistanceMetric} of a distance parse node, or null for any other node. */ + public static DistanceMetric getMetric(ParseNode distanceNode) { + return isDistanceFunctionNode(distanceNode) + ? DistanceMetric.forFunction(((FunctionParseNode) distanceNode).getName()) + : null; + } + + /** + * Returns true if the parse node can be the source vector. The node must reference a column and + * must not be stateless. + */ + public static boolean isSourceVectorNode(ParseNode node) { + if (node == null || node.isStateless()) { + return false; + } + return hasColumnParseNode(node); + } + + /** + * Returns true if the parse node can be the query vector. The node must be a literal, a bind + * parameter, or a stateless expression that references no column. + */ + public static boolean isQueryVectorNode(ParseNode node) { + if (node == null) { + return false; + } + if (node instanceof BindParseNode || node instanceof LiteralParseNode) { + return true; + } + return node.isStateless() && !hasColumnParseNode(node); + } + + /** Returns true if the parse node tree contains a column reference. */ + public static boolean hasColumnParseNode(ParseNode node) { + if (node == null) { + return false; + } + if (node instanceof ColumnParseNode) { + return true; + } + List<ParseNode> children = node.getChildren(); + if (children != null) { + for (ParseNode child : children) { + if (hasColumnParseNode(child)) { + return true; + } + } + } + return false; + } + + /** + * Returns true if the compiled {@link OrderBy} and row limit make a vector similarity search. See + * {@link #isVectorSearch(List, Integer)}. + */ + public static boolean isVectorSearch(OrderBy orderBy, Integer limit) { + return orderBy != null && isVectorSearch(orderBy.getOrderByExpressions(), limit); + } + + /** + * Returns true if the compiled ORDER BY expressions and row limit make a vector similarity + * search. The limit must be positive, and the only ORDER BY expression must be an ascending + * {@link DistanceFunction}. + */ + public static boolean isVectorSearch(List<OrderByExpression> orderByExpressions, Integer limit) { + if ( + limit == null || limit <= 0 || orderByExpressions == null || orderByExpressions.size() != 1 + ) { + return false; + } + OrderByExpression orderByExpression = orderByExpressions.get(0); + return orderByExpression.isAscending() + && orderByExpression.getExpression() instanceof DistanceFunction; + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/parse/DistanceFunctionParseNode.java b/phoenix-core-client/src/main/java/org/apache/phoenix/parse/DistanceFunctionParseNode.java index 4437a8c57b..7bffbc6ad9 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/parse/DistanceFunctionParseNode.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/parse/DistanceFunctionParseNode.java @@ -19,9 +19,15 @@ package org.apache.phoenix.parse; import java.sql.SQLException; import java.util.List; +import org.apache.hadoop.hbase.io.ImmutableBytesWritable; import org.apache.phoenix.compile.StatementContext; +import org.apache.phoenix.expression.CoerceExpression; import org.apache.phoenix.expression.Expression; +import org.apache.phoenix.expression.LiteralExpression; import org.apache.phoenix.expression.function.DistanceFunction; +import org.apache.phoenix.schema.types.PDataType; +import org.apache.phoenix.schema.types.PVectorFloat; +import org.apache.phoenix.util.ExpressionUtil; /** * Parse node for the vector distance functions. It converts an array argument to the vector type of @@ -35,6 +41,47 @@ public class DistanceFunctionParseNode extends FunctionParseNode { super(name, children, info); } + @Override + public List<Expression> validate(List<Expression> children, StatementContext context) + throws SQLException { + Expression child0 = children.get(0); + Expression child1 = children.get(1); + PDataType type0 = child0.getDataType(); + PDataType type1 = child1.getDataType(); + PDataType targetType = PVectorFloat.INSTANCE; + Integer targetDim = null; + if (type0 != null && type0.isVectorType()) { + targetType = type0; + targetDim = child0.getMaxLength(); + } else if (type1 != null && type1.isVectorType()) { + targetType = type1; + targetDim = child1.getMaxLength(); + } + if (type0 != null && !type0.isVectorType() && type0.isArrayType()) { + children.set(0, coerceToVector(child0, targetType, targetDim, context)); + } + if (type1 != null && !type1.isVectorType() && type1.isArrayType()) { + children.set(1, coerceToVector(child1, targetType, targetDim, context)); + } + return super.validate(children, context); + } + + private static Expression coerceToVector(Expression expr, PDataType targetType, Integer targetDim, + StatementContext context) throws SQLException { + if (!ExpressionUtil.isConstant(expr)) { + return CoerceExpression.create(expr, targetType, expr.getSortOrder(), targetDim); + } + ImmutableBytesWritable ptr = context.getTempPtr(); + expr.evaluate(null, ptr); + Object array = + expr.getDataType().toObject(ptr, expr.getSortOrder(), expr.getMaxLength(), expr.getScale()); + Object vector = targetType.toObject(array, expr.getDataType()); + // The max length carries the literal dimension, so that validation at compile time can + // compare it with the dimension of the other operand. + return LiteralExpression.newConstant(vector, targetType, targetType.getMaxLength(vector), null, + expr.getSortOrder(), expr.getDeterminism(), true); + } + @Override public Expression create(List<Expression> children, StatementContext context) throws SQLException { diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/parse/OrderByNode.java b/phoenix-core-client/src/main/java/org/apache/phoenix/parse/OrderByNode.java index 340b0d73c9..8e1472234a 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/parse/OrderByNode.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/parse/OrderByNode.java @@ -17,6 +17,7 @@ */ package org.apache.phoenix.parse; +import java.util.List; import org.apache.phoenix.compile.ColumnResolver; import org.apache.phoenix.schema.types.PInteger; @@ -28,15 +29,42 @@ public final class OrderByNode { private final ParseNode child; private final boolean nullsLast; private final boolean orderAscending; + // True if the item has no NULLS FIRST or NULLS LAST clause. + private final boolean nullsDefault; OrderByNode(ParseNode child, boolean nullsLast, boolean orderAscending) { + this(child, nullsLast, orderAscending, false); + } + + OrderByNode(ParseNode child, boolean orderAscending) { + this(child, false, orderAscending, true); + } + + private OrderByNode(ParseNode child, boolean nullsLast, boolean orderAscending, + boolean nullsDefault) { this.child = child; this.nullsLast = nullsLast; this.orderAscending = orderAscending; + this.nullsDefault = nullsDefault; } + /** + * Returns true if nulls sort last. Without a NULLS clause, an ascending vector distance sorts + * nulls last, because a nearest neighbor search ranks rows without a distance after all scored + * rows. All other items without a NULLS clause sort nulls first. + */ public boolean isNullsLast() { - return nullsLast; + return nullsLast + || nullsDefault && orderAscending && child instanceof DistanceFunctionParseNode; + } + + /** + * Returns true if the item has no NULLS FIRST or NULLS LAST clause. Compilation uses this to + * apply the distance default to an item, such as an ordinal, that only compilation resolves to a + * vector distance. + */ + public boolean isNullsDefault() { + return nullsDefault; } public boolean isAscending() { @@ -52,7 +80,7 @@ public final class OrderByNode { final int prime = 31; int result = 1; result = prime * result + ((child == null) ? 0 : child.hashCode()); - result = prime * result + (nullsLast ? 1231 : 1237); + result = prime * result + (isNullsLast() ? 1231 : 1237); result = prime * result + (orderAscending ? 1231 : 1237); return result; } @@ -66,7 +94,7 @@ public final class OrderByNode { if (child == null) { if (other.child != null) return false; } else if (!child.equals(other.child)) return false; - if (nullsLast != other.nullsLast) return false; + if (isNullsLast() != other.isNullsLast()) return false; if (orderAscending != other.orderAscending) return false; return true; } @@ -74,13 +102,15 @@ public final class OrderByNode { @Override public String toString() { return child.toString() + (orderAscending ? " asc" : " desc") + " nulls " - + (nullsLast ? "last" : "first"); + + (isNullsLast() ? "last" : "first"); } public void toSQL(ColumnResolver resolver, StringBuilder buf) { child.toSQL(resolver, buf); if (!orderAscending) buf.append(" DESC"); - if (nullsLast) buf.append(" NULLS LAST "); + if (isNullsLast()) buf.append(" NULLS LAST "); + else if (!nullsDefault && orderAscending && child instanceof DistanceFunctionParseNode) + buf.append(" NULLS FIRST "); } public boolean isIntegerLiteral() { @@ -94,4 +124,26 @@ public final class OrderByNode { } return (Integer) ((LiteralParseNode) child).getValue(); } + + /** + * Returns the projected node that an ordinal item refers to. Returns null if the item is not a + * valid ordinal. Also returns null if a wildcard comes before that position, because the position + * is then not known until compilation. + */ + public ParseNode getOrdinalSelectNode(List<AliasedNode> select) { + Integer ordinal = getValueIfIntegerLiteral(); + if (ordinal == null || ordinal < 1 || ordinal > select.size()) { + return null; + } + for (int i = 0; i < ordinal - 1; i++) { + ParseNode node = select.get(i).getNode(); + if ( + node instanceof WildcardParseNode || node instanceof TableWildcardParseNode + || node instanceof FamilyWildcardParseNode + ) { + return null; + } + } + return select.get(ordinal - 1).getNode(); + } } diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/parse/ParseNodeFactory.java b/phoenix-core-client/src/main/java/org/apache/phoenix/parse/ParseNodeFactory.java index da08d77a76..c4cdb4b4d0 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/parse/ParseNodeFactory.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/parse/ParseNodeFactory.java @@ -921,6 +921,22 @@ public class ParseNodeFactory { return new OrderByNode(expression, nullsLast, orderAscending); } + /** Creates an ORDER BY item without a NULLS FIRST or NULLS LAST clause. */ + public OrderByNode orderBy(ParseNode expression, boolean orderAscending) { + return new OrderByNode(expression, orderAscending); + } + + /** + * Creates an ORDER BY item over a rewritten expression with the ordering of the template. If the + * template resolves its default to NULLS LAST, as for a vector distance, the new item keeps NULLS + * LAST. This is also true if the rewritten expression is not a distance, such as an index column. + */ + public OrderByNode orderBy(ParseNode expression, OrderByNode template) { + return template.isNullsDefault() && !template.isNullsLast() + ? new OrderByNode(expression, template.isAscending()) + : new OrderByNode(expression, template.isNullsLast(), template.isAscending()); + } + public SelectStatement select(TableNode from, HintNode hint, boolean isDistinct, List<AliasedNode> select, ParseNode where, List<ParseNode> groupBy, ParseNode having, List<OrderByNode> orderBy, LimitNode limit, OffsetNode offset, int bindCount, diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/parse/ParseNodeRewriter.java b/phoenix-core-client/src/main/java/org/apache/phoenix/parse/ParseNodeRewriter.java index 59a1349d5e..d004f400eb 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/parse/ParseNodeRewriter.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/parse/ParseNodeRewriter.java @@ -160,7 +160,10 @@ public class ParseNodeRewriter extends TraverseAllParseNodeVisitor<ParseNode> { ParseNode node = orderByNode.getNode(); rewriter.reset(); ParseNode normNode = node.accept(rewriter); - if (node == normNode) { + OrderByNode normOrderByNode = node == normNode + ? resolveDistanceOrdinal(orderByNode, selectNodes) + : NODE_FACTORY.orderBy(normNode, orderByNode); + if (orderByNode == normOrderByNode) { if (orderByNodes != normOrderByNodes) { normOrderByNodes.add(orderByNode); } @@ -169,8 +172,7 @@ public class ParseNodeRewriter extends TraverseAllParseNodeVisitor<ParseNode> { if (orderByNodes == normOrderByNodes) { normOrderByNodes = Lists.newArrayList(orderByNodes.subList(0, i)); } - normOrderByNodes - .add(NODE_FACTORY.orderBy(normNode, orderByNode.isNullsLast(), orderByNode.isAscending())); + normOrderByNodes.add(normOrderByNode); } // Return new SELECT statement with updated WHERE clause @@ -188,6 +190,19 @@ public class ParseNodeRewriter extends TraverseAllParseNodeVisitor<ParseNode> { statement.getSelects(), statement.getUdfParseNodes()); } + /** + * Makes NULLS LAST explicit for an ascending ordinal without a NULLS clause that refers to a + * projected vector distance. Without this step, a rewrite of the projection, such as to an index + * column, loses the distance default. Compilation resolves an ordinal after a wildcard. + */ + private static OrderByNode resolveDistanceOrdinal(OrderByNode orderByNode, + List<AliasedNode> selectNodes) { + return orderByNode.isNullsDefault() && orderByNode.isAscending() + && orderByNode.getOrdinalSelectNode(selectNodes) instanceof DistanceFunctionParseNode + ? NODE_FACTORY.orderBy(orderByNode.getNode(), true, true) + : orderByNode; + } + private Map<String, ParseNode> getAliasMap() { return aliasMap; } diff --git a/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorCrossRegionMergeIT.java b/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorCrossRegionMergeIT.java new file mode 100644 index 0000000000..6cbd7e656b --- /dev/null +++ b/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorCrossRegionMergeIT.java @@ -0,0 +1,371 @@ +/* + * 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.phoenix.end2end; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.phoenix.util.QueryUtil; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +@Category(ParallelStatsDisabledTest.class) +public class VectorCrossRegionMergeIT extends ParallelStatsDisabledIT { + + /** + * Makes sure that the top-K merge across the regions of a salted table gives the exact nearest + * neighbors. + */ + @Test + public void testCrossRegionMergeCorrectness() throws Exception { + String tableName = generateUniqueName(); + int numRows = 200; + int dim = 4; + int limit = 5; + + Random rng = new Random(42); + float[][] vectors = new float[numRows][dim]; + for (int i = 0; i < numRows; i++) { + for (int d = 0; d < dim; d++) { + vectors[i][d] = rng.nextFloat(); + } + } + + float[] queryVec = new float[] { 0.25f, 0.50f, 0.75f, 0.10f }; + + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement().execute("CREATE TABLE " + tableName + + " (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 4)) SALT_BUCKETS=4"); + + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + for (int i = 0; i < numRows; i++) { + ps.setInt(1, i + 1); + Float[] boxed = new Float[dim]; + for (int d = 0; d < dim; d++) { + boxed[d] = vectors[i][d]; + } + ps.setArray(2, conn.createArrayOf("FLOAT", boxed)); + ps.executeUpdate(); + } + } + conn.commit(); + + // The plan must push top-K to the server and merge sort on the client. + String query = "SELECT pk FROM " + tableName + " ORDER BY L2_DISTANCE(v, ?) LIMIT " + limit; + try (PreparedStatement psExplain = conn.prepareStatement("EXPLAIN " + query)) { + Float[] boxedQuery = new Float[dim]; + for (int d = 0; d < dim; d++) { + boxedQuery[d] = queryVec[d]; + } + psExplain.setArray(1, conn.createArrayOf("FLOAT", boxedQuery)); + ResultSet rsExplain = psExplain.executeQuery(); + String explainPlan = QueryUtil.getExplainPlan(rsExplain); + assertNotNull("EXPLAIN plan must not be null", explainPlan); + assertTrue("EXPLAIN plan should indicate client merge sort: " + explainPlan, + explainPlan.contains("CLIENT MERGE SORT")); + assertTrue("EXPLAIN plan should indicate server top-N: " + explainPlan, + explainPlan.contains("SERVER TOP-5")); + } + + // Compute the exact nearest-neighbor order as the reference. + double[] distances = new double[numRows]; + for (int i = 0; i < numRows; i++) { + distances[i] = computeDist("L2_DISTANCE", vectors[i], queryVec); + } + Integer[] indices = new Integer[numRows]; + for (int i = 0; i < numRows; i++) { + indices[i] = i; + } + Arrays.sort(indices, (a, b) -> Double.compare(distances[a], distances[b])); + + List<Integer> expectedPks = new ArrayList<>(); + for (int i = 0; i < limit; i++) { + expectedPks.add(indices[i] + 1); + } + + List<Integer> actualPks = new ArrayList<>(); + try (PreparedStatement ps = conn.prepareStatement(query)) { + Float[] boxedQuery = new Float[dim]; + for (int d = 0; d < dim; d++) { + boxedQuery[d] = queryVec[d]; + } + ps.setArray(1, conn.createArrayOf("FLOAT", boxedQuery)); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + actualPks.add(rs.getInt(1)); + } + } + } + + assertEquals("Should return exactly " + limit + " rows", limit, actualPks.size()); + assertEquals("Global top-" + limit + " across 4 salt buckets must match brute-force", + expectedPks, actualPks); + } + } + + /** Top-K merge across regions with an inline array literal as the query vector. */ + @Test + public void testCrossRegionMergeWithArrayLiteral() throws Exception { + String tableName = generateUniqueName(); + int numRows = 100; + int dim = 4; + int limit = 5; + + Random rng = new Random(123); + float[][] vectors = new float[numRows][dim]; + for (int i = 0; i < numRows; i++) { + for (int d = 0; d < dim; d++) { + vectors[i][d] = rng.nextFloat(); + } + } + float[] queryVec = new float[] { 0.1f, 0.2f, 0.3f, 0.4f }; + + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement().execute("CREATE TABLE " + tableName + + " (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 4)) SALT_BUCKETS=4"); + + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + for (int i = 0; i < numRows; i++) { + ps.setInt(1, i + 1); + Float[] boxed = new Float[dim]; + for (int d = 0; d < dim; d++) { + boxed[d] = vectors[i][d]; + } + ps.setArray(2, conn.createArrayOf("FLOAT", boxed)); + ps.executeUpdate(); + } + } + conn.commit(); + + double[] distances = new double[numRows]; + for (int i = 0; i < numRows; i++) { + distances[i] = computeDist("L2_DISTANCE", vectors[i], queryVec); + } + Integer[] indices = new Integer[numRows]; + for (int i = 0; i < numRows; i++) { + indices[i] = i; + } + Arrays.sort(indices, (a, b) -> Double.compare(distances[a], distances[b])); + + List<Integer> expectedPks = new ArrayList<>(); + for (int i = 0; i < limit; i++) { + expectedPks.add(indices[i] + 1); + } + + String query = "SELECT pk FROM " + tableName + + " ORDER BY L2_DISTANCE(v, ARRAY[0.1, 0.2, 0.3, 0.4]) LIMIT " + limit; + List<Integer> actualPks = new ArrayList<>(); + try (ResultSet rs = conn.createStatement().executeQuery(query)) { + while (rs.next()) { + actualPks.add(rs.getInt(1)); + } + } + + assertEquals("Should return exactly " + limit + " rows", limit, actualPks.size()); + assertEquals("Literal query vector: global top-" + limit + " must match brute-force", + expectedPks, actualPks); + } + } + + /** Top-K merge across regions for each supported distance metric. */ + @Test + public void testCrossRegionMergeAllMetrics() throws Exception { + String tableName = generateUniqueName(); + int numRows = 120; + int dim = 4; + int limit = 5; + + Random rng = new Random(999); + float[][] vectors = new float[numRows][dim]; + for (int i = 0; i < numRows; i++) { + for (int d = 0; d < dim; d++) { + vectors[i][d] = rng.nextFloat(); + } + } + float[] queryVec = new float[] { 0.3f, 0.4f, 0.5f, 0.6f }; + + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement().execute("CREATE TABLE " + tableName + + " (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 4)) SALT_BUCKETS=4"); + + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + for (int i = 0; i < numRows; i++) { + ps.setInt(1, i + 1); + Float[] boxed = new Float[dim]; + for (int d = 0; d < dim; d++) { + boxed[d] = vectors[i][d]; + } + ps.setArray(2, conn.createArrayOf("FLOAT", boxed)); + ps.executeUpdate(); + } + } + conn.commit(); + + String[] metrics = + { "L2_DISTANCE", "L2_DISTANCE_SQUARED", "COSINE_DISTANCE", "INNER_PRODUCT" }; + for (String metric : metrics) { + double[] distances = new double[numRows]; + for (int i = 0; i < numRows; i++) { + distances[i] = computeDist(metric, vectors[i], queryVec); + } + Integer[] indices = new Integer[numRows]; + for (int i = 0; i < numRows; i++) { + indices[i] = i; + } + Arrays.sort(indices, (a, b) -> Double.compare(distances[a], distances[b])); + + List<Integer> expectedPks = new ArrayList<>(); + for (int i = 0; i < limit; i++) { + expectedPks.add(indices[i] + 1); + } + + String query = + "SELECT pk FROM " + tableName + " ORDER BY " + metric + "(v, ?) LIMIT " + limit; + List<Integer> actualPks = new ArrayList<>(); + try (PreparedStatement ps = conn.prepareStatement(query)) { + Float[] boxedQuery = new Float[dim]; + for (int d = 0; d < dim; d++) { + boxedQuery[d] = queryVec[d]; + } + ps.setArray(1, conn.createArrayOf("FLOAT", boxedQuery)); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + actualPks.add(rs.getInt(1)); + } + } + } + + assertEquals(metric + ": should return " + limit + " rows", limit, actualPks.size()); + assertEquals(metric + ": cross-region merge top-" + limit + " must match brute-force", + expectedPks, actualPks); + } + } + } + + /** Top-K merge across regions with double-precision vectors. */ + @Test + public void testCrossRegionMergeDoubleVectors() throws Exception { + String tableName = generateUniqueName(); + int numRows = 100; + int dim = 4; + int limit = 5; + + Random rng = new Random(777); + double[][] vectors = new double[numRows][dim]; + for (int i = 0; i < numRows; i++) { + for (int d = 0; d < dim; d++) { + vectors[i][d] = rng.nextDouble(); + } + } + double[] queryVec = new double[] { 0.2, 0.4, 0.6, 0.8 }; + + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement().execute("CREATE TABLE " + tableName + + " (pk INTEGER PRIMARY KEY, v VECTOR(DOUBLE, 4)) SALT_BUCKETS=4"); + + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + for (int i = 0; i < numRows; i++) { + ps.setInt(1, i + 1); + Double[] boxed = new Double[dim]; + for (int d = 0; d < dim; d++) { + boxed[d] = vectors[i][d]; + } + ps.setArray(2, conn.createArrayOf("DOUBLE", boxed)); + ps.executeUpdate(); + } + } + conn.commit(); + + double[] distances = new double[numRows]; + for (int i = 0; i < numRows; i++) { + double sumSq = 0; + for (int d = 0; d < dim; d++) { + double diff = vectors[i][d] - queryVec[d]; + sumSq += diff * diff; + } + distances[i] = Math.sqrt(sumSq); + } + Integer[] indices = new Integer[numRows]; + for (int i = 0; i < numRows; i++) { + indices[i] = i; + } + Arrays.sort(indices, (a, b) -> Double.compare(distances[a], distances[b])); + + List<Integer> expectedPks = new ArrayList<>(); + for (int i = 0; i < limit; i++) { + expectedPks.add(indices[i] + 1); + } + + String query = "SELECT pk FROM " + tableName + " ORDER BY L2_DISTANCE(v, ?) LIMIT " + limit; + List<Integer> actualPks = new ArrayList<>(); + try (PreparedStatement ps = conn.prepareStatement(query)) { + Double[] boxedQuery = new Double[dim]; + for (int d = 0; d < dim; d++) { + boxedQuery[d] = queryVec[d]; + } + ps.setArray(1, conn.createArrayOf("DOUBLE", boxedQuery)); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + actualPks.add(rs.getInt(1)); + } + } + } + + assertEquals("Should return " + limit + " rows", limit, actualPks.size()); + assertEquals("Double vectors cross-region merge top-" + limit + " must match brute-force", + expectedPks, actualPks); + } + } + + private static double computeDist(String metric, float[] a, float[] b) { + double dot = 0, normA = 0, normB = 0, sumSq = 0; + for (int i = 0; i < a.length; i++) { + double diff = a[i] - b[i]; + sumSq += diff * diff; + dot += a[i] * b[i]; + normA += a[i] * a[i]; + normB += b[i] * b[i]; + } + switch (metric) { + case "L2_DISTANCE": + return Math.sqrt(sumSq); + case "L2_DISTANCE_SQUARED": + return sumSq; + case "COSINE_DISTANCE": + double denom = Math.sqrt(normA) * Math.sqrt(normB); + return (denom == 0) ? 1.0 : 1.0 - dot / denom; + case "INNER_PRODUCT": + return -dot; + default: + throw new IllegalArgumentException(metric); + } + } +} diff --git a/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorDistancePushDownIT.java b/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorDistancePushDownIT.java new file mode 100644 index 0000000000..32158e710b --- /dev/null +++ b/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorDistancePushDownIT.java @@ -0,0 +1,440 @@ +/* + * 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.phoenix.end2end; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.sql.Array; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.phoenix.compile.ExplainPlan; +import org.apache.phoenix.compile.ExplainPlanAttributes; +import org.apache.phoenix.jdbc.PhoenixPreparedStatement; +import org.apache.phoenix.util.QueryUtil; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +/** Tests exact vector similarity search with server-side top-K distance push-down. */ +@Category(ParallelStatsDisabledTest.class) +public class VectorDistancePushDownIT extends ParallelStatsDisabledIT { + + /** Server-side top-K push-down with fixed reference vectors. */ + @Test + public void testServerSidePushDownVerification() throws Exception { + String tableName = generateUniqueName(); + + // Fixed candidate vectors for the query vector [1.0, 0.0, 0.0] + float[][] vectors = new float[50][3]; + for (int i = 0; i < 50; i++) { + // The first five vectors are the nearest to the query vector. + vectors[i][0] = (i < 5) ? (0.9f + 0.02f * i) : ((i % 3) * 0.3f); + vectors[i][1] = (i < 5) ? 0.05f * i : ((i % 5) * 0.2f); + vectors[i][2] = (i < 5) ? 0.02f * i : ((i % 7) * 0.15f); + } + + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement() + .execute("CREATE TABLE " + tableName + " (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 3))"); + + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + for (int i = 0; i < 50; i++) { + ps.setInt(1, i + 1); + Float[] boxed = new Float[] { vectors[i][0], vectors[i][1], vectors[i][2] }; + ps.setArray(2, conn.createArrayOf("FLOAT", boxed)); + ps.executeUpdate(); + } + } + conn.commit(); + + // The plan must push top-K to the server and merge sort on the client. + String query = + "SELECT pk FROM " + tableName + " ORDER BY L2_DISTANCE(v, ARRAY[1.0,0.0,0.0]) LIMIT 5"; + ResultSet explainRs = conn.createStatement().executeQuery("EXPLAIN " + query); + String explainPlan = QueryUtil.getExplainPlan(explainRs); + assertNotNull("EXPLAIN plan must not be null", explainPlan); + assertTrue("EXPLAIN plan should contain 'SERVER TOP-5' but was:\n" + explainPlan, + explainPlan.contains("SERVER TOP-5")); + assertTrue("EXPLAIN plan should contain 'CLIENT MERGE SORT TOP-5' but was:\n" + explainPlan, + explainPlan.contains("CLIENT MERGE SORT TOP-5")); + + // Compute the exact nearest-neighbor order as the reference. + final float[] queryVec = { 1.0f, 0.0f, 0.0f }; + double[] distances = new double[50]; + for (int i = 0; i < 50; i++) { + distances[i] = bruteForceDist("L2_DISTANCE", vectors[i], queryVec); + } + Integer[] indices = new Integer[50]; + for (int i = 0; i < 50; i++) { + indices[i] = i; + } + Arrays.sort(indices, (a, b) -> Double.compare(distances[a], distances[b])); + + List<Integer> expectedPks = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + expectedPks.add(indices[i] + 1); + } + + ResultSet rs = conn.createStatement().executeQuery(query); + List<Integer> resultPks = new ArrayList<>(); + while (rs.next()) { + resultPks.add(rs.getInt(1)); + } + + assertEquals("Should return exactly 5 rows", 5, resultPks.size()); + assertEquals("Returned PKs should match brute-force top-5 nearest neighbors (in order). " + + "Expected: " + expectedPks + ", Got: " + resultPks, expectedPks, resultPks); + } + } + + /** + * Push-down correctness for each supported distance metric, with vectors whose order depends on + * the metric. + */ + @Test + public void testAllDistanceFunctionsPushDownCorrectness() throws Exception { + String tableName = generateUniqueName(); + + // The nearest-neighbor order of these vectors changes with the metric. + float[][] vectors = { { 1.0f, 0.0f, 0.0f }, { 0.0f, 1.0f, 0.0f }, { 0.0f, 0.0f, 1.0f }, + { 0.6f, 0.8f, 0.0f }, { 0.5f, 0.5f, 0.5f } }; + float[] queryVec = { 0.9f, 0.1f, 0.0f }; + + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement() + .execute("CREATE TABLE " + tableName + " (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 3))"); + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + for (int i = 0; i < vectors.length; i++) { + ps.setInt(1, i + 1); + ps.setArray(2, conn.createArrayOf("FLOAT", + new Float[] { vectors[i][0], vectors[i][1], vectors[i][2] })); + ps.executeUpdate(); + } + } + conn.commit(); + + String[] functions = + { "L2_DISTANCE", "L2_DISTANCE_SQUARED", "COSINE_DISTANCE", "INNER_PRODUCT" }; + for (String func : functions) { + String query = + "SELECT pk FROM " + tableName + " ORDER BY " + func + "(v, ARRAY[0.9,0.1,0.0]) LIMIT 2"; + + // Reference order for this metric + double[] dists = new double[vectors.length]; + for (int i = 0; i < vectors.length; i++) { + dists[i] = bruteForceDist(func, vectors[i], queryVec); + } + Integer[] idx = new Integer[vectors.length]; + for (int i = 0; i < idx.length; i++) { + idx[i] = i; + } + Arrays.sort(idx, (a, b) -> Double.compare(dists[a], dists[b])); + List<Integer> expected = new ArrayList<>(); + for (int i = 0; i < 2; i++) { + expected.add(idx[i] + 1); + } + + ResultSet rs = conn.createStatement().executeQuery(query); + List<Integer> actual = new ArrayList<>(); + while (rs.next()) { + actual.add(rs.getInt(1)); + } + assertEquals(func + ": wrong top-2 result", expected, actual); + } + } + } + + private static double bruteForceDist(String func, float[] a, float[] b) { + double dot = 0, normA = 0, normB = 0, sumSq = 0; + for (int i = 0; i < a.length; i++) { + double diff = a[i] - b[i]; + sumSq += diff * diff; + dot += a[i] * b[i]; + normA += a[i] * a[i]; + normB += b[i] * b[i]; + } + switch (func) { + case "L2_DISTANCE": + return Math.sqrt(sumSq); + case "L2_DISTANCE_SQUARED": + return sumSq; + case "COSINE_DISTANCE": + double denom = Math.sqrt(normA) * Math.sqrt(normB); + return (denom == 0) ? 1.0 : 1.0 - dot / denom; + case "INNER_PRODUCT": + return -dot; + default: + throw new IllegalArgumentException(func); + } + } + + /** Push-down with a bind parameter as the query vector. */ + @Test + public void testBindParameterQueryVectorPushDown() throws Exception { + String tableName = generateUniqueName(); + + float[][] vectors = { { 1.0f, 0.0f, 0.0f }, { 0.0f, 1.0f, 0.0f }, { 0.0f, 0.0f, 1.0f }, + { 0.5f, 0.5f, 0.0f }, { 0.7f, 0.2f, 0.1f } }; + + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement() + .execute("CREATE TABLE " + tableName + " (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 3))"); + + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + for (int i = 0; i < vectors.length; i++) { + ps.setInt(1, i + 1); + ps.setArray(2, conn.createArrayOf("FLOAT", + new Float[] { vectors[i][0], vectors[i][1], vectors[i][2] })); + ps.executeUpdate(); + } + } + conn.commit(); + + // The nearest match is pk=1. + String query = "SELECT pk FROM " + tableName + " ORDER BY L2_DISTANCE(v, ?) LIMIT 1"; + try (PreparedStatement ps = conn.prepareStatement(query)) { + Array queryVec = conn.createArrayOf("FLOAT", new Float[] { 1.0f, 0.0f, 0.0f }); + ps.setArray(1, queryVec); + ResultSet rs = ps.executeQuery(); + assertTrue("Should have at least one result", rs.next()); + assertEquals("Row with pk=1 (vector [1,0,0]) should be nearest to query [1,0,0]", 1, + rs.getInt(1)); + } + } + } + + /** Query plan structure and EXPLAIN plan attributes for vector similarity queries. */ + @Test + public void testExactVectorSearchExplainPlan() throws Exception { + String tableName = generateUniqueName(); + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement().execute( + "CREATE TABLE " + tableName + " (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 3), val INTEGER)"); + + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + ps.setInt(1, 1); + ps.setArray(2, conn.createArrayOf("FLOAT", new Float[] { 1.0f, 0.0f, 0.0f })); + ps.setInt(3, 42); + ps.executeUpdate(); + } + conn.commit(); + + // Query plan with a bind parameter query vector + String query = "SELECT * FROM " + tableName + " ORDER BY COSINE_DISTANCE(v, ?) LIMIT 10"; + try (PreparedStatement ps = conn.prepareStatement("EXPLAIN " + query)) { + ps.setArray(1, conn.createArrayOf("FLOAT", new Float[] { 1.0f, 0.0f, 0.0f })); + ResultSet rs = ps.executeQuery(); + String explainPlan = QueryUtil.getExplainPlan(rs); + assertNotNull("EXPLAIN plan must not be null", explainPlan); + assertTrue("EXPLAIN plan should contain 'COSINE_DISTANCE' but was:\n" + explainPlan, + explainPlan.contains("COSINE_DISTANCE")); + assertTrue("EXPLAIN plan should contain 'TOP-10' but was:\n" + explainPlan, + explainPlan.contains("TOP-10")); + assertTrue( + "EXPLAIN plan should contain 'CLIENT MERGE SORT TOP-10' but was:\n" + explainPlan, + explainPlan.contains("CLIENT MERGE SORT TOP-10")); + assertTrue( + "EXPLAIN plan should contain 'SERVER TOP-10 BY COSINE_DISTANCE' but was:\n" + explainPlan, + explainPlan.contains("SERVER TOP-10 BY COSINE_DISTANCE")); + } + + // EXPLAIN plan attributes + try (PreparedStatement ps = conn.prepareStatement(query)) { + ps.setArray(1, conn.createArrayOf("FLOAT", new Float[] { 1.0f, 0.0f, 0.0f })); + ExplainPlan plan = + ps.unwrap(PhoenixPreparedStatement.class).optimizeQuery().getExplainPlan(); + ExplainPlanAttributes attributes = plan.getPlanStepsAsAttributes(); + assertNotNull("ExplainPlanAttributes must not be null", attributes); + assertTrue("isVectorSearch should be true", attributes.isVectorSearch()); + assertEquals("CLIENT MERGE SORT TOP-10", attributes.getClientSortAlgo()); + assertEquals(Integer.valueOf(10), attributes.getClientRowLimit()); + assertEquals(Long.valueOf(10), attributes.getServerRowLimit()); + assertNotNull(attributes.getServerSortedBy()); + assertTrue("serverSortedBy should contain 'COSINE_DISTANCE'", + attributes.getServerSortedBy().contains("COSINE_DISTANCE")); + } + + // Query plan with a literal array query vector + String literalQuery = + "SELECT * FROM " + tableName + " ORDER BY COSINE_DISTANCE(v, ARRAY[1.0,0.0,0.0]) LIMIT 10"; + try (ResultSet rs = conn.createStatement().executeQuery("EXPLAIN " + literalQuery)) { + String explainPlan = QueryUtil.getExplainPlan(rs); + assertTrue("EXPLAIN plan should contain 'COSINE_DISTANCE' but was:\n" + explainPlan, + explainPlan.contains("COSINE_DISTANCE")); + assertTrue("EXPLAIN plan should contain 'TOP-10' but was:\n" + explainPlan, + explainPlan.contains("TOP-10")); + assertTrue( + "EXPLAIN plan should contain 'CLIENT MERGE SORT TOP-10' but was:\n" + explainPlan, + explainPlan.contains("CLIENT MERGE SORT TOP-10")); + } + + // Query plans for the other distance functions + String[] otherFuncs = { "L2_DISTANCE", "L2_DISTANCE_SQUARED", "INNER_PRODUCT" }; + for (String func : otherFuncs) { + String funcQuery = + "SELECT pk FROM " + tableName + " ORDER BY " + func + "(v, ARRAY[1.0,0.0,0.0]) LIMIT 10"; + try (ResultSet rs = conn.createStatement().executeQuery("EXPLAIN " + funcQuery)) { + String explainPlan = QueryUtil.getExplainPlan(rs); + assertTrue("EXPLAIN plan should contain " + func + " but was:\n" + explainPlan, + explainPlan.contains(func)); + assertTrue( + "EXPLAIN plan should contain 'SERVER TOP-10 BY " + func + "' but was:\n" + explainPlan, + explainPlan.contains("SERVER TOP-10 BY " + func)); + assertTrue( + "EXPLAIN plan should contain 'CLIENT MERGE SORT TOP-10' but was:\n" + explainPlan, + explainPlan.contains("CLIENT MERGE SORT TOP-10")); + } + } + + // Query plan for an ORDER BY that is not a vector search, as a baseline + String nonVectorQuery = "SELECT * FROM " + tableName + " ORDER BY val LIMIT 10"; + try (ResultSet rs = conn.createStatement().executeQuery("EXPLAIN " + nonVectorQuery)) { + String explainPlan = QueryUtil.getExplainPlan(rs); + assertTrue( + "Non-vector EXPLAIN plan should contain 'CLIENT MERGE SORT' but was:\n" + explainPlan, + explainPlan.contains("CLIENT MERGE SORT")); + assertTrue( + "Non-vector EXPLAIN plan should contain 'CLIENT LIMIT 10' but was:\n" + explainPlan, + explainPlan.contains("CLIENT LIMIT 10")); + assertFalse("Non-vector EXPLAIN plan must not contain 'TOP-10' but was:\n" + explainPlan, + explainPlan.contains("TOP-10")); + assertTrue( + "Non-vector EXPLAIN plan should contain 'SERVER TOP 10 ROWS SORTED BY' but was:\n" + + explainPlan, + explainPlan.contains("SERVER TOP 10 ROWS SORTED BY")); + } + try (PreparedStatement ps = conn.prepareStatement(nonVectorQuery)) { + ExplainPlan plan = + ps.unwrap(PhoenixPreparedStatement.class).optimizeQuery().getExplainPlan(); + ExplainPlanAttributes attributes = plan.getPlanStepsAsAttributes(); + assertFalse("isVectorSearch should be false for non-vector query", + attributes.isVectorSearch()); + assertEquals("CLIENT MERGE SORT", attributes.getClientSortAlgo()); + assertEquals(Integer.valueOf(10), attributes.getClientRowLimit()); + assertEquals(Long.valueOf(10), attributes.getServerRowLimit()); + assertEquals("[VAL]", attributes.getServerSortedBy()); + } + } + } + + /** Top-K merge across the regions of a salted table. */ + @Test + public void testCrossRegionMergeWithSaltedTable() throws Exception { + String tableName = generateUniqueName(); + final int numRows = 200; + final int dim = 4; + final int topK = 5; + + Random rng = new Random(42); + float[][] vectors = new float[numRows][dim]; + for (int i = 0; i < numRows; i++) { + for (int d = 0; d < dim; d++) { + vectors[i][d] = rng.nextFloat(); + } + } + + float[] queryVec = { 0.5f, 0.5f, 0.5f, 0.5f }; + + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement().execute("CREATE TABLE " + tableName + + " (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, " + dim + ")) SALT_BUCKETS=4"); + + String upsertSql = "UPSERT INTO " + tableName + " VALUES (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(upsertSql)) { + for (int i = 0; i < numRows; i++) { + ps.setInt(1, i + 1); + Float[] boxed = new Float[dim]; + for (int d = 0; d < dim; d++) { + boxed[d] = vectors[i][d]; + } + ps.setArray(2, conn.createArrayOf("FLOAT", boxed)); + ps.executeUpdate(); + } + } + conn.commit(); + + // Compute the exact nearest-neighbor order as the reference. + double[] distances = new double[numRows]; + for (int i = 0; i < numRows; i++) { + distances[i] = bruteForceDist("L2_DISTANCE", vectors[i], queryVec); + } + Integer[] indices = new Integer[numRows]; + for (int i = 0; i < numRows; i++) { + indices[i] = i; + } + Arrays.sort(indices, (a, b) -> Double.compare(distances[a], distances[b])); + List<Integer> expectedPks = new ArrayList<>(); + for (int i = 0; i < topK; i++) { + expectedPks.add(indices[i] + 1); + } + + String query = "SELECT pk FROM " + tableName + + " ORDER BY L2_DISTANCE(v, ARRAY[0.5,0.5,0.5,0.5]) LIMIT " + topK; + + // The plan must show a server top-K, a client top-K merge and a parallel scan. + ResultSet explainRs = conn.createStatement().executeQuery("EXPLAIN " + query); + String explainPlan = QueryUtil.getExplainPlan(explainRs); + assertNotNull("EXPLAIN plan must not be null", explainPlan); + assertTrue("EXPLAIN plan should contain 'SERVER TOP-" + topK + "' but was:\n" + explainPlan, + explainPlan.contains("SERVER TOP-" + topK)); + assertTrue( + "EXPLAIN plan should contain 'CLIENT MERGE SORT TOP-" + topK + "' but was:\n" + explainPlan, + explainPlan.contains("CLIENT MERGE SORT TOP-" + topK)); + assertTrue("EXPLAIN plan should show parallel scan (4-WAY) but was:\n" + explainPlan, + explainPlan.contains("4-WAY")); + + ResultSet rs = conn.createStatement().executeQuery(query); + List<Integer> resultPks = new ArrayList<>(); + while (rs.next()) { + resultPks.add(rs.getInt(1)); + } + + assertEquals("Should return exactly " + topK + " rows", topK, resultPks.size()); + assertEquals( + "Cross-region merge: returned PKs should match brute-force top-" + topK + + " nearest neighbors (in order). Expected: " + expectedPks + ", Got: " + resultPks, + expectedPks, resultPks); + + // A bind parameter query vector must give the same result. + String bindQuery = + "SELECT pk FROM " + tableName + " ORDER BY L2_DISTANCE(v, ?) LIMIT " + topK; + try (PreparedStatement ps = conn.prepareStatement(bindQuery)) { + ps.setArray(1, conn.createArrayOf("FLOAT", new Float[] { 0.5f, 0.5f, 0.5f, 0.5f })); + ResultSet bindRs = ps.executeQuery(); + List<Integer> bindResultPks = new ArrayList<>(); + while (bindRs.next()) { + bindResultPks.add(bindRs.getInt(1)); + } + assertEquals("Bind-param cross-region merge: wrong top-" + topK, expectedPks, + bindResultPks); + } + } + } +} diff --git a/phoenix-core/src/test/java/org/apache/phoenix/expression/function/DistanceFunctionSerializationTest.java b/phoenix-core/src/test/java/org/apache/phoenix/expression/function/DistanceFunctionSerializationTest.java new file mode 100644 index 0000000000..9d1cc114d3 --- /dev/null +++ b/phoenix-core/src/test/java/org/apache/phoenix/expression/function/DistanceFunctionSerializationTest.java @@ -0,0 +1,159 @@ +/* + * 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.phoenix.expression.function; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.util.Arrays; +import java.util.List; +import org.apache.hadoop.hbase.Cell; +import org.apache.hadoop.hbase.KeyValue; +import org.apache.hadoop.hbase.io.ImmutableBytesWritable; +import org.apache.hadoop.hbase.util.Bytes; +import org.apache.hadoop.io.WritableUtils; +import org.apache.phoenix.expression.Expression; +import org.apache.phoenix.expression.ExpressionType; +import org.apache.phoenix.expression.KeyValueColumnExpression; +import org.apache.phoenix.expression.LiteralExpression; +import org.apache.phoenix.schema.PColumn; +import org.apache.phoenix.schema.PColumnImpl; +import org.apache.phoenix.schema.PNameFactory; +import org.apache.phoenix.schema.SortOrder; +import org.apache.phoenix.schema.tuple.MultiKeyValueTuple; +import org.apache.phoenix.schema.tuple.Tuple; +import org.apache.phoenix.schema.types.PDouble; +import org.apache.phoenix.schema.types.PVectorFloat; +import org.junit.Test; + +/** Tests serialization and deserialization of vector distance expressions. */ +public class DistanceFunctionSerializationTest { + + private static final byte[] CF = Bytes.toBytes("0"); + private static final byte[] CQ_V = Bytes.toBytes("V"); + + @Test + public void testL2DistanceFunctionRoundTrip() throws Exception { + verifyRoundTrip(createDistanceFunction(L2DistanceFunction.class), "L2_DISTANCE"); + } + + @Test + public void testL2DistanceSquaredFunctionRoundTrip() throws Exception { + verifyRoundTrip(createDistanceFunction(L2DistanceSquaredFunction.class), "L2_DISTANCE_SQUARED"); + } + + @Test + public void testCosineDistanceFunctionRoundTrip() throws Exception { + verifyRoundTrip(createDistanceFunction(CosineDistanceFunction.class), "COSINE_DISTANCE"); + } + + @Test + public void testInnerProductDistanceFunctionRoundTrip() throws Exception { + verifyRoundTrip(createDistanceFunction(InnerProductDistanceFunction.class), "INNER_PRODUCT"); + } + + @Test + public void testExpressionTypeRegistration() throws Exception { + DistanceFunction[] funcs = { createDistanceFunction(L2DistanceFunction.class), + createDistanceFunction(L2DistanceSquaredFunction.class), + createDistanceFunction(CosineDistanceFunction.class), + createDistanceFunction(InnerProductDistanceFunction.class), }; + + Class<?>[] expectedClasses = { L2DistanceFunction.class, L2DistanceSquaredFunction.class, + CosineDistanceFunction.class, InnerProductDistanceFunction.class, }; + + for (int i = 0; i < funcs.length; i++) { + ExpressionType type = ExpressionType.valueOf(funcs[i]); + assertNotNull(funcs[i].getName() + " must have ExpressionType registration", type); + assertEquals(funcs[i].getName() + " ExpressionType class mismatch", expectedClasses[i], + type.getExpressionClass()); + } + } + + @Test + public void testDistanceUpperBoundNotSerialized() throws Exception { + // The server sets the distance upper bound at execution time. Serialization does not keep the + // bound, so a deserialized expression has no bound. + DistanceFunction original = createDistanceFunction(L2DistanceFunction.class); + original.setDistanceUpperBound(42.5); + + DistanceFunction deserialized = serializeAndDeserialize(original); + assertEquals(Double.MAX_VALUE, deserialized.getDistanceUpperBound(), 0.0); + } + + private DistanceFunction createDistanceFunction(Class<? extends DistanceFunction> clazz) + throws Exception { + int dim = 3; + PColumn col = + new PColumnImpl(PNameFactory.newName("V"), PNameFactory.newName("0"), PVectorFloat.INSTANCE, + dim, null, true, 1, SortOrder.ASC, 0, null, false, null, false, false, null, 0); + KeyValueColumnExpression colExpr = new KeyValueColumnExpression(col); + float[] queryVec = { 0.5f, 0.5f, 0.5f }; + LiteralExpression queryLit = + LiteralExpression.newConstant(queryVec, PVectorFloat.INSTANCE, dim, null); + return clazz.getConstructor(List.class).newInstance(Arrays.asList(colExpr, queryLit)); + } + + private void verifyRoundTrip(DistanceFunction original, String name) throws Exception { + float[] testVec = { 1.0f, 0.0f, 0.0f }; + byte[] rowKey = Bytes.toBytes(1); + byte[] vBytes = PVectorFloat.INSTANCE.toBytes(testVec); + Cell cell = new KeyValue(rowKey, CF, CQ_V, vBytes); + Tuple tuple = new MultiKeyValueTuple(Arrays.asList(cell)); + + ImmutableBytesWritable ptr1 = new ImmutableBytesWritable(); + assertTrue(name + ": original evaluate must succeed", original.evaluate(tuple, ptr1)); + double originalDist = (Double) PDouble.INSTANCE.toObject(ptr1); + + DistanceFunction deserialized = serializeAndDeserialize(original); + + assertEquals(name + ": deserialized class mismatch", original.getClass(), + deserialized.getClass()); + + ImmutableBytesWritable ptr2 = new ImmutableBytesWritable(); + assertTrue(name + ": deserialized evaluate must succeed", deserialized.evaluate(tuple, ptr2)); + double deserializedDist = (Double) PDouble.INSTANCE.toObject(ptr2); + + assertEquals(name + ": distances must match after round-trip", originalDist, deserializedDist, + 0.0); + } + + private DistanceFunction serializeAndDeserialize(DistanceFunction original) throws Exception { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + DataOutputStream out = new DataOutputStream(baos); + int ordinal = ExpressionType.valueOf(original).ordinal(); + WritableUtils.writeVInt(out, ordinal); + original.write(out); + out.flush(); + + byte[] bytes = baos.toByteArray(); + DataInputStream in = new DataInputStream(new ByteArrayInputStream(bytes)); + int readOrdinal = WritableUtils.readVInt(in); + Expression expr = ExpressionType.values()[readOrdinal].newInstance(); + expr.readFields(in); + + assertTrue("Deserialized expression must be a DistanceFunction", + expr instanceof DistanceFunction); + return (DistanceFunction) expr; + } +} diff --git a/phoenix-core/src/test/java/org/apache/phoenix/expression/function/DistanceFunctionsTest.java b/phoenix-core/src/test/java/org/apache/phoenix/expression/function/DistanceFunctionsTest.java index 89aaa3a5c8..0da426058c 100644 --- a/phoenix-core/src/test/java/org/apache/phoenix/expression/function/DistanceFunctionsTest.java +++ b/phoenix-core/src/test/java/org/apache/phoenix/expression/function/DistanceFunctionsTest.java @@ -58,6 +58,7 @@ import org.apache.phoenix.schema.SortOrder; import org.apache.phoenix.schema.types.PDouble; import org.apache.phoenix.schema.types.PVectorDouble; import org.apache.phoenix.schema.types.PVectorFloat; +import org.apache.phoenix.util.QueryUtil; import org.junit.Test; import org.apache.phoenix.thirdparty.com.google.common.collect.Multimap; @@ -494,6 +495,25 @@ public class DistanceFunctionsTest extends BaseConnectionlessQueryTest { } } + @Test + public void testCompilationFailsOnLiteralDimensionMismatch() throws Exception { + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement().execute("CREATE TABLE t_literal_dim_mismatch (" + + "pk INTEGER PRIMARY KEY, " + "embedding VECTOR(FLOAT, 3))"); + + try (PreparedStatement stmt = + conn.prepareStatement("SELECT pk FROM t_literal_dim_mismatch ORDER BY " + + "L2_DISTANCE(embedding, ARRAY[1.0, 2.0]) LIMIT 10")) { + PhoenixPreparedStatement pStmt = stmt.unwrap(PhoenixPreparedStatement.class); + pStmt.compileQuery(); + fail("Expected compilation to fail due to dimension mismatch between the column and " + + "the literal query vector"); + } catch (SQLException e) { + assertEquals(SQLExceptionCode.VECTOR_DIMENSION_MISMATCH.getErrorCode(), e.getErrorCode()); + } + } + } + @Test public void testCompilationFailsOnNonVectorArgument() throws Exception { try (Connection conn = DriverManager.getConnection(getUrl())) { @@ -528,4 +548,28 @@ public class DistanceFunctionsTest extends BaseConnectionlessQueryTest { } } } + + @Test + public void testExplainTopKMerge() throws Exception { + try (Connection conn = DriverManager.getConnection(getUrl())) { + conn.createStatement().execute("CREATE TABLE t_explain_topk (pk INTEGER PRIMARY KEY, " + + "v VECTOR(FLOAT, 3)) SALT_BUCKETS=4"); + String query = "SELECT pk FROM t_explain_topk ORDER BY L2_DISTANCE(v, ARRAY[1.0, 0.0, 0.0])"; + + String plan = QueryUtil + .getExplainPlan(conn.createStatement().executeQuery("EXPLAIN " + query + " LIMIT 5")); + assertTrue(plan, plan.contains("SERVER TOP-5 BY L2_DISTANCE")); + assertTrue(plan, plan.contains("CLIENT MERGE SORT TOP-5")); + assertFalse(plan, plan.contains("CLIENT LIMIT")); + + // With an OFFSET, the plan keeps separate merge, offset and limit steps. + plan = QueryUtil.getExplainPlan( + conn.createStatement().executeQuery("EXPLAIN " + query + " LIMIT 5 OFFSET 2")); + assertFalse(plan, plan.contains("CLIENT MERGE SORT TOP-")); + int merge = plan.indexOf("CLIENT MERGE SORT"); + int offset = plan.indexOf("CLIENT OFFSET 2"); + int limit = plan.indexOf("CLIENT LIMIT 5"); + assertTrue(plan, merge >= 0 && merge < offset && offset < limit); + } + } } diff --git a/phoenix-core/src/test/java/org/apache/phoenix/optimize/VectorSearchUtilTest.java b/phoenix-core/src/test/java/org/apache/phoenix/optimize/VectorSearchUtilTest.java new file mode 100644 index 0000000000..a572e7279d --- /dev/null +++ b/phoenix-core/src/test/java/org/apache/phoenix/optimize/VectorSearchUtilTest.java @@ -0,0 +1,571 @@ +/* + * 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.phoenix.optimize; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +import java.sql.DriverManager; +import org.apache.phoenix.compile.ColumnResolver; +import org.apache.phoenix.compile.FromCompiler; +import org.apache.phoenix.compile.GroupByCompiler.GroupBy; +import org.apache.phoenix.compile.OrderByCompiler; +import org.apache.phoenix.compile.OrderByCompiler.OrderBy; +import org.apache.phoenix.compile.RowProjector; +import org.apache.phoenix.compile.StatementContext; +import org.apache.phoenix.expression.function.CosineDistanceFunction; +import org.apache.phoenix.expression.function.InnerProductDistanceFunction; +import org.apache.phoenix.expression.function.L2DistanceFunction; +import org.apache.phoenix.expression.function.L2DistanceSquaredFunction; +import org.apache.phoenix.jdbc.PhoenixConnection; +import org.apache.phoenix.jdbc.PhoenixStatement; +import org.apache.phoenix.parse.ArrayConstructorNode; +import org.apache.phoenix.parse.BindParseNode; +import org.apache.phoenix.parse.ColumnParseNode; +import org.apache.phoenix.parse.LiteralParseNode; +import org.apache.phoenix.parse.ParseNode; +import org.apache.phoenix.parse.SQLParser; +import org.apache.phoenix.parse.SelectStatement; +import org.apache.phoenix.query.BaseConnectionlessQueryTest; +import org.junit.Test; + +/** Unit tests that find vector search queries and extract their descriptors. */ +public class VectorSearchUtilTest extends BaseConnectionlessQueryTest { + + private static boolean isQueryLiteral(VectorSearchDescriptor d) { + return d.getQueryExpression() instanceof LiteralParseNode + || d.getQueryExpression() instanceof ArrayConstructorNode; + } + + private SelectStatement parse(String sql) throws Exception { + SQLParser parser = new SQLParser(sql); + return (SelectStatement) parser.parseStatement(); + } + + @Test + public void testL2DistanceFunctionLiteralArg() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ARRAY[1.0,0.0,0.0]) ASC LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.L2, d.getMetric()); + assertEquals(L2DistanceFunction.NAME, d.getDistanceFunctionName()); + assertTrue("Source should be a plain column", d.isSourceColumn()); + assertEquals("V", d.getSourceColumnName()); + assertTrue("ARRAY literal should be a query literal", isQueryLiteral(d)); + assertFalse(d.getQueryExpression() instanceof BindParseNode); + assertEquals(Integer.valueOf(10), d.getLimit()); + } + + @Test + public void testL2DistanceSquaredFunctionBindArg() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE_SQUARED(v, ?) ASC LIMIT 5"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.L2, d.getMetric()); + assertEquals(L2DistanceSquaredFunction.NAME, d.getDistanceFunctionName()); + assertTrue(d.isSourceColumn()); + assertEquals("V", d.getSourceColumnName()); + assertTrue(d.getQueryExpression() instanceof BindParseNode); + assertFalse(isQueryLiteral(d)); + assertEquals(Integer.valueOf(5), d.getLimit()); + } + + @Test + public void testCosineDistanceFunctionBindArg() throws Exception { + String sql = "SELECT * FROM t ORDER BY COSINE_DISTANCE(v, ?) ASC LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.COSINE, d.getMetric()); + assertEquals(CosineDistanceFunction.NAME, d.getDistanceFunctionName()); + assertTrue(d.isSourceColumn()); + assertEquals("V", d.getSourceColumnName()); + assertTrue(d.getQueryExpression() instanceof BindParseNode); + assertEquals(Integer.valueOf(10), d.getLimit()); + } + + @Test + public void testInnerProductFunctionBindArg() throws Exception { + String sql = "SELECT * FROM t ORDER BY INNER_PRODUCT(v, ?) ASC LIMIT 25"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.INNER_PRODUCT, d.getMetric()); + assertEquals(InnerProductDistanceFunction.NAME, d.getDistanceFunctionName()); + assertTrue(d.isSourceColumn()); + assertEquals("V", d.getSourceColumnName()); + assertEquals(Integer.valueOf(25), d.getLimit()); + } + + @Test + public void testL2OperatorLiteralArg() throws Exception { + String sql = "SELECT * FROM t ORDER BY v <-> ARRAY[1.0,0.0,0.0] ASC LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.L2, d.getMetric()); + assertTrue(d.isSourceColumn()); + assertEquals("V", d.getSourceColumnName()); + assertEquals(Integer.valueOf(10), d.getLimit()); + } + + @Test + public void testCosineOperatorBindArg() throws Exception { + String sql = "SELECT * FROM t ORDER BY v <=> ? ASC LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.COSINE, d.getMetric()); + assertTrue(d.isSourceColumn()); + assertEquals("V", d.getSourceColumnName()); + assertEquals(Integer.valueOf(10), d.getLimit()); + } + + @Test + public void testInnerProductOperatorBindArg() throws Exception { + String sql = "SELECT * FROM t ORDER BY v <#> ? ASC LIMIT 15"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.INNER_PRODUCT, d.getMetric()); + assertTrue(d.isSourceColumn()); + assertEquals("V", d.getSourceColumnName()); + assertEquals(Integer.valueOf(15), d.getLimit()); + } + + @Test + public void testSwappedArguments_QueryVectorFirst() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(?, v) ASC LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.L2, d.getMetric()); + assertEquals("V", d.getSourceColumnName()); + assertTrue(d.getQueryExpression() instanceof BindParseNode); + } + + @Test + public void testDefaultSortIsAscending() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull("Default ASC should qualify as vector search", d); + assertEquals(DistanceMetric.L2, d.getMetric()); + assertEquals("V", d.getSourceColumnName()); + assertEquals(Integer.valueOf(10), d.getLimit()); + } + + @Test + public void testDescendingSortReturnsFalse() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) DESC LIMIT 10"; + SelectStatement select = parse(sql); + + assertFalse("Descending vector distance order is not a vector search", + VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testOrdinalResolvesProjectedDistance() throws Exception { + VectorSearchDescriptor d = VectorSearchUtil + .getVectorSearchDescriptor(parse("SELECT id, L2_DISTANCE(v, ?) FROM t ORDER BY 2 LIMIT 10")); + assertNotNull("An ordinal referencing a projected distance is a vector search", d); + assertEquals(DistanceMetric.L2, d.getMetric()); + assertEquals("V", d.getSourceColumnName()); + assertEquals(Integer.valueOf(10), d.getLimit()); + assertNull(VectorSearchUtil + .getVectorSearchDescriptor(parse("SELECT id, L2_DISTANCE(v, ?) FROM t ORDER BY 1 LIMIT 10"))); + assertNull(VectorSearchUtil + .getVectorSearchDescriptor(parse("SELECT id, L2_DISTANCE(v, ?) FROM t ORDER BY 3 LIMIT 10"))); + assertNull(VectorSearchUtil.getVectorSearchDescriptor( + parse("SELECT id, L2_DISTANCE(v, ?) FROM t ORDER BY 2 DESC LIMIT 10"))); + // The position of an ordinal after a wildcard is known only after the projection expands. + assertNotNull(VectorSearchUtil.getVectorSearchDescriptor( + parse("SELECT L2_DISTANCE(v, ?), \"0\".* FROM t ORDER BY 1 LIMIT 10"))); + assertNull(VectorSearchUtil.getVectorSearchDescriptor( + parse("SELECT \"0\".*, L2_DISTANCE(v, ?) FROM t ORDER BY 2 LIMIT 10"))); + } + + private static boolean isNullsLast(PhoenixConnection conn, String sql) throws Exception { + OrderBy orderBy = conn.prepareStatement(sql) + .unwrap(org.apache.phoenix.jdbc.PhoenixPreparedStatement.class).optimizeQuery().getOrderBy(); + assertEquals(sql, 1, orderBy.getOrderByExpressions().size()); + return orderBy.getOrderByExpressions().get(0).isNullsLast(); + } + + @Test + public void testAscendingDistanceDefaultsToNullsLast() throws Exception { + try (PhoenixConnection conn = (PhoenixConnection) DriverManager.getConnection(getUrl())) { + conn.createStatement() + .execute("CREATE TABLE t_nulls (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 3), ts INTEGER, " + + "w VECTOR(FLOAT, 3))"); + String q = "ARRAY[1.0, 0.0, 0.0]"; + // Without a NULLS clause, an ascending distance sorts rows without a vector last. This + // applies to a distance written directly, as an operator, or through an alias or an ordinal. + // It also applies to a distance in a flattened subquery, through an alias of the subquery, + // and to a distance group key, directly or through an ordinal. + for (String sql : new String[] { + "SELECT pk FROM t_nulls ORDER BY L2_DISTANCE(v, " + q + ") LIMIT 5", + "SELECT pk FROM t_nulls ORDER BY v <=> " + q + " ASC LIMIT 5", + "SELECT pk, INNER_PRODUCT(v, " + q + ") AS d FROM t_nulls ORDER BY d LIMIT 5", + "SELECT pk, L2_DISTANCE_SQUARED(v, " + q + ") FROM t_nulls ORDER BY 2 LIMIT 5", + "SELECT pk FROM (SELECT pk, v FROM t_nulls) ORDER BY v <-> " + q + " LIMIT 5", + "SELECT * FROM (SELECT pk, L2_DISTANCE(v, " + q + ") d FROM t_nulls) ORDER BY d LIMIT 5", + "SELECT L2_DISTANCE(v, w), COUNT(*) FROM t_nulls GROUP BY L2_DISTANCE(v, w) " + + "ORDER BY L2_DISTANCE(v, w)", + "SELECT L2_DISTANCE(v, w), COUNT(*) FROM t_nulls GROUP BY L2_DISTANCE(v, w) ORDER BY 1", + "SELECT pk FROM t_nulls ORDER BY L2_DISTANCE(v, " + q + ") NULLS LAST LIMIT 5" }) { + assertTrue(sql, isNullsLast(conn, sql)); + } + // An explicit NULLS FIRST, a descending distance and all other orderings keep nulls first. + // This includes an ordinal that a wildcard before it moves to another column, here ts. + for (String sql : new String[] { + "SELECT pk FROM t_nulls ORDER BY L2_DISTANCE(v, " + q + ") NULLS FIRST LIMIT 5", + "SELECT pk, L2_DISTANCE(v, " + q + ") AS d FROM t_nulls ORDER BY d NULLS FIRST LIMIT 5", + "SELECT pk, L2_DISTANCE(v, " + q + ") FROM t_nulls ORDER BY 2 NULLS FIRST LIMIT 5", + "SELECT * FROM (SELECT pk, L2_DISTANCE(v, " + q + + ") d FROM t_nulls) ORDER BY d NULLS FIRST LIMIT 5", + "SELECT pk FROM t_nulls ORDER BY L2_DISTANCE(v, " + q + ") DESC LIMIT 5", + "SELECT pk FROM t_nulls ORDER BY ts LIMIT 5", + "SELECT \"0\".*, L2_DISTANCE(v, " + q + ") FROM t_nulls ORDER BY 2 LIMIT 5", + "SELECT ts, COUNT(*) FROM t_nulls GROUP BY ts ORDER BY 2" }) { + assertFalse(sql, isNullsLast(conn, sql)); + } + // SQL that toString() regenerates from a statement with an explicit NULLS FIRST keeps + // NULLS FIRST. + String regenerated = + parse("SELECT pk FROM t_nulls ORDER BY L2_DISTANCE(v, " + q + ") NULLS FIRST LIMIT 5") + .toString(); + assertFalse(regenerated, isNullsLast(conn, regenerated)); + } + } + + @Test + public void testLimitAsBindVariable() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) ASC LIMIT ?"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertNull("Bind-variable limit has no parse-time integer value", d.getLimit()); + } + + @Test + public void testLargeLimit() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) ASC LIMIT 1000000"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(Integer.valueOf(1_000_000), d.getLimit()); + } + + @Test + public void testMissingLimitReturnsFalse() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) ASC"; + SelectStatement select = parse(sql); + + assertFalse("Unbounded distance sort is not a vector search", + VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testZeroLimitReturnsFalse() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) ASC LIMIT 0"; + SelectStatement select = parse(sql); + + assertFalse("LIMIT 0 returns no rows; not a useful vector search", + VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testTableQualifiedColumn() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(t.v, ?) ASC LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals("V", d.getSourceColumnName()); + assertEquals("T", ((ColumnParseNode) d.getSourceExpression()).getTableName()); + assertNull("No schema qualifier expected", + ((ColumnParseNode) d.getSourceExpression()).getSchemaName()); + } + + @Test + public void testSchemaAndTableQualifiedColumn() throws Exception { + String sql = "SELECT * FROM s.t ORDER BY L2_DISTANCE(s.t.v, ?) ASC LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals("V", d.getSourceColumnName()); + assertEquals("T", ((ColumnParseNode) d.getSourceExpression()).getTableName()); + assertEquals("S", ((ColumnParseNode) d.getSourceExpression()).getSchemaName()); + } + + @Test + public void testBsonVectorPathExpression() throws Exception { + String sql = + "SELECT * FROM t ORDER BY L2_DISTANCE(BSON_VECTOR_VALUE(doc, 'search.embedding', 128), ?)" + + " ASC LIMIT 10"; + VectorSearchDescriptor d = VectorSearchUtil.getVectorSearchDescriptor(parse(sql)); + + assertNotNull(d); + assertEquals(DistanceMetric.L2, d.getMetric()); + assertFalse("BSON expression is not a bare column", d.isSourceColumn()); + assertNotNull(d.getSourceExpression()); + assertTrue(d.getQueryExpression() instanceof BindParseNode); + assertEquals(Integer.valueOf(10), d.getLimit()); + } + + @Test + public void testNoOrderByReturnsFalse() throws Exception { + String sql = "SELECT * FROM t LIMIT 10"; + SelectStatement select = parse(sql); + + assertFalse(VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testNonDistanceOrderByReturnsFalse() throws Exception { + String sql = "SELECT * FROM t ORDER BY created_at ASC LIMIT 10"; + SelectStatement select = parse(sql); + + assertFalse(VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testMultipleDistanceExpressionsReturnsFalse() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v1, ?), L2_DISTANCE(v2, ?) LIMIT 10"; + SelectStatement select = parse(sql); + + assertFalse("Multiple ORDER BY items disqualify the query", + VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testDistanceMixedWithOtherOrderByReturnsFalse() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?), created_at LIMIT 10"; + SelectStatement select = parse(sql); + + assertFalse(VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testTwoColumnsInDistanceReturnsFalse() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v1, v2) ASC LIMIT 10"; + SelectStatement select = parse(sql); + + assertFalse("Column-to-column distance is not a nearest-neighbor search", + VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testTwoLiteralsInDistanceReturnsFalse() throws Exception { + String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(ARRAY[1.0], ARRAY[2.0]) ASC LIMIT 10"; + SelectStatement select = parse(sql); + + assertFalse("Both-literal distance has no indexed column to search", + VectorSearchUtil.isVectorSearchQuery(select)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(select)); + } + + @Test + public void testNullStatementReturnsFalse() { + assertFalse(VectorSearchUtil.isVectorSearchQuery(null)); + assertNull(VectorSearchUtil.getVectorSearchDescriptor(null)); + } + + @Test + public void testDistanceMetricForFunction() { + assertSame(DistanceMetric.L2, DistanceMetric.forFunction(L2DistanceFunction.NAME)); + assertSame(DistanceMetric.L2, DistanceMetric.forFunction(L2DistanceSquaredFunction.NAME)); + assertSame(DistanceMetric.COSINE, DistanceMetric.forFunction(CosineDistanceFunction.NAME)); + assertSame(DistanceMetric.INNER_PRODUCT, + DistanceMetric.forFunction(InnerProductDistanceFunction.NAME)); + assertNull(DistanceMetric.forFunction(null)); + assertNull(DistanceMetric.forFunction("MANHATTAN_DISTANCE")); + } + + @Test + public void testDistanceMetricFromString() { + assertSame(DistanceMetric.L2, DistanceMetric.fromString("L2")); + assertSame(DistanceMetric.L2, DistanceMetric.fromString(" l2 ")); + assertSame(DistanceMetric.COSINE, DistanceMetric.fromString("cosine")); + assertSame(DistanceMetric.INNER_PRODUCT, DistanceMetric.fromString("Inner_Product")); + // Null, blank, function names and other unknown names are not metric names. + for (String other : new String[] { null, " ", "EUCLIDEAN", "DOT_PRODUCT", "IP", "L2_DISTANCE", + "HAMMING" }) { + assertNull("'" + other + "' is not a metric name", DistanceMetric.fromString(other)); + } + } + + @Test + public void testHasColumnParseNode_directColumn() throws Exception { + SelectStatement stmt = parse("SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) ASC LIMIT 1"); + ParseNode distanceNode = stmt.getOrderBy().get(0).getNode(); + ParseNode columnArg = distanceNode.getChildren().get(0); + + assertTrue("Direct ColumnParseNode should be detected", columnArg instanceof ColumnParseNode); + assertTrue(VectorSearchUtil.hasColumnParseNode(columnArg)); + } + + @Test + public void testHasColumnParseNode_deeplyNestedColumn() throws Exception { + SelectStatement stmt = parse( + "SELECT * FROM t ORDER BY L2_DISTANCE(BSON_VECTOR_VALUE(doc, 'emb', 3), ?) ASC LIMIT 1"); + ParseNode distanceNode = stmt.getOrderBy().get(0).getNode(); + ParseNode bsonArg = distanceNode.getChildren().get(0); + + assertTrue("Nested column inside BSON call must be found recursively", + VectorSearchUtil.hasColumnParseNode(bsonArg)); + } + + @Test + public void testHasColumnParseNode_bindAndNullHaveNoColumn() throws Exception { + SelectStatement stmt = parse("SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) ASC LIMIT 1"); + ParseNode distanceNode = stmt.getOrderBy().get(0).getNode(); + ParseNode bindArg = distanceNode.getChildren().get(1); + + assertFalse("Bind parameter has no column child", VectorSearchUtil.hasColumnParseNode(bindArg)); + assertFalse("null input should return false", VectorSearchUtil.hasColumnParseNode(null)); + } + + @Test + public void testIsDistanceFunctionNode_distanceFunctions() throws Exception { + for (String sql : new String[] { "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) ASC LIMIT 1", + "SELECT * FROM t ORDER BY L2_DISTANCE_SQUARED(v, ?) ASC LIMIT 1", + "SELECT * FROM t ORDER BY COSINE_DISTANCE(v, ?) ASC LIMIT 1", + "SELECT * FROM t ORDER BY INNER_PRODUCT(v, ?) ASC LIMIT 1", }) { + ParseNode node = parse(sql).getOrderBy().get(0).getNode(); + assertTrue("Expected distance node for: " + sql, + VectorSearchUtil.isDistanceFunctionNode(node)); + } + } + + @Test + public void testIsDistanceFunctionNode_distanceOperators() throws Exception { + for (String sql : new String[] { "SELECT * FROM t ORDER BY v <-> ? ASC LIMIT 1", + "SELECT * FROM t ORDER BY v <=> ? ASC LIMIT 1", + "SELECT * FROM t ORDER BY v <#> ? ASC LIMIT 1", }) { + ParseNode node = parse(sql).getOrderBy().get(0).getNode(); + assertTrue("Expected distance operator node for: " + sql, + VectorSearchUtil.isDistanceFunctionNode(node)); + } + } + + @Test + public void testIsDistanceFunctionNode_nonDistanceAndNull() throws Exception { + ParseNode colNode = + parse("SELECT * FROM t ORDER BY v ASC LIMIT 1").getOrderBy().get(0).getNode(); + assertFalse("Plain column is not a distance function", + VectorSearchUtil.isDistanceFunctionNode(colNode)); + assertFalse("null should return false", VectorSearchUtil.isDistanceFunctionNode(null)); + } + + @Test + public void testCompiledIsVectorSearch_emptyAndNullOrderByReturnsFalse() { + assertFalse(VectorSearchUtil.isVectorSearch(OrderBy.EMPTY_ORDER_BY, 10)); + assertFalse(VectorSearchUtil.isVectorSearch((OrderBy) null, 10)); + } + + @Test + public void testCompiledIsVectorSearch_nullOrZeroLimitReturnsFalse() throws Exception { + // Two vector columns give both operands a type at compile time without bind values. + String ddl = "CREATE TABLE t_compiled" + + " (pk INTEGER PRIMARY KEY, v1 VECTOR(FLOAT, 3), v2 VECTOR(FLOAT, 3))"; + try (PhoenixConnection conn = (PhoenixConnection) DriverManager.getConnection(getUrl())) { + conn.createStatement().execute(ddl); + + String sql = "SELECT pk FROM t_compiled ORDER BY L2_DISTANCE(v1, v2) ASC LIMIT 5"; + SelectStatement select = parse(sql); + PhoenixStatement stmt = new PhoenixStatement(conn); + ColumnResolver resolver = FromCompiler.getResolverForQuery(select, conn); + StatementContext context = new StatementContext(stmt, resolver); + + OrderBy orderBy = OrderByCompiler.compile(context, select, GroupBy.EMPTY_GROUP_BY, 5, + org.apache.phoenix.compile.CompiledOffset.EMPTY_COMPILED_OFFSET, + RowProjector.EMPTY_PROJECTOR, null, null); + + assertTrue("Compiled OrderBy with LIMIT 5 should be a vector search", + VectorSearchUtil.isVectorSearch(orderBy, 5)); + + assertFalse("null limit must be rejected", VectorSearchUtil.isVectorSearch(orderBy, null)); + assertFalse("zero limit must be rejected", VectorSearchUtil.isVectorSearch(orderBy, 0)); + } + } + + @Test + public void testCompiledIsVectorSearch_allDistanceFunctions() throws Exception { + // Two vector columns let each distance function compile without a literal query vector. + String ddl = "CREATE TABLE t_all_dist" + + " (pk INTEGER PRIMARY KEY, v1 VECTOR(FLOAT, 3), v2 VECTOR(FLOAT, 3))"; + try (PhoenixConnection conn = (PhoenixConnection) DriverManager.getConnection(getUrl())) { + conn.createStatement().execute(ddl); + + String[] sqls = { "SELECT pk FROM t_all_dist ORDER BY L2_DISTANCE(v1, v2) ASC LIMIT 3", + "SELECT pk FROM t_all_dist ORDER BY L2_DISTANCE_SQUARED(v1, v2) ASC LIMIT 3", + "SELECT pk FROM t_all_dist ORDER BY COSINE_DISTANCE(v1, v2) ASC LIMIT 3", + "SELECT pk FROM t_all_dist ORDER BY INNER_PRODUCT(v1, v2) ASC LIMIT 3", }; + + for (String sql : sqls) { + SelectStatement select = parse(sql); + PhoenixStatement stmt = new PhoenixStatement(conn); + ColumnResolver resolver = FromCompiler.getResolverForQuery(select, conn); + StatementContext context = new StatementContext(stmt, resolver); + + OrderBy orderBy = OrderByCompiler.compile(context, select, GroupBy.EMPTY_GROUP_BY, 3, + org.apache.phoenix.compile.CompiledOffset.EMPTY_COMPILED_OFFSET, + RowProjector.EMPTY_PROJECTOR, null, null); + + assertTrue("VectorSearchUtil.isVectorSearch for: " + sql, + VectorSearchUtil.isVectorSearch(orderBy, 3)); + } + } + } + + @Test + public void testCompiledIsVectorSearch_nonDistanceOrderByReturnsFalse() throws Exception { + String ddl = "CREATE TABLE t_no_dist (pk INTEGER PRIMARY KEY, ts INTEGER)"; + try (PhoenixConnection conn = (PhoenixConnection) DriverManager.getConnection(getUrl())) { + conn.createStatement().execute(ddl); + + String sql = "SELECT pk FROM t_no_dist ORDER BY ts ASC LIMIT 10"; + SelectStatement select = parse(sql); + PhoenixStatement stmt = new PhoenixStatement(conn); + ColumnResolver resolver = FromCompiler.getResolverForQuery(select, conn); + StatementContext context = new StatementContext(stmt, resolver); + + OrderBy orderBy = OrderByCompiler.compile(context, select, GroupBy.EMPTY_GROUP_BY, 10, + org.apache.phoenix.compile.CompiledOffset.EMPTY_COMPILED_OFFSET, + RowProjector.EMPTY_PROJECTOR, null, null); + + assertFalse("Non-distance ORDER BY must not be detected as vector search", + VectorSearchUtil.isVectorSearch(orderBy, 10)); + } + } +} diff --git a/phoenix-core/src/test/java/org/apache/phoenix/query/explain/ExplainPlanTest.java b/phoenix-core/src/test/java/org/apache/phoenix/query/explain/ExplainPlanTest.java index 62034de010..700d26c354 100644 --- a/phoenix-core/src/test/java/org/apache/phoenix/query/explain/ExplainPlanTest.java +++ b/phoenix-core/src/test/java/org/apache/phoenix/query/explain/ExplainPlanTest.java @@ -1836,6 +1836,7 @@ public class ExplainPlanTest extends BaseConnectionlessQueryTest { n.putNull("consistency"); n.putNull("hint"); n.putNull("serverSortedBy"); + n.put("vectorSearch", false); n.putNull("explainScanType"); n.putNull("tableName"); n.putNull("keyRanges"); diff --git a/phoenix-core/src/test/java/org/apache/phoenix/schema/types/PArrayDataTypeVectorCoercionTest.java b/phoenix-core/src/test/java/org/apache/phoenix/schema/types/PArrayDataTypeVectorCoercionTest.java new file mode 100644 index 0000000000..a360fb3990 --- /dev/null +++ b/phoenix-core/src/test/java/org/apache/phoenix/schema/types/PArrayDataTypeVectorCoercionTest.java @@ -0,0 +1,167 @@ +/* + * 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.phoenix.schema.types; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; + +/** Unit tests for the coercion and cast rules between array types and vector types. */ +public class PArrayDataTypeVectorCoercionTest { + + @Test + public void testFloatArrayCoercibleToVectorFloat() { + assertTrue("FLOAT ARRAY should be coercible to VECTOR(FLOAT)", + PFloatArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testFloatArrayCoercibleToVectorDouble() { + assertTrue("FLOAT ARRAY should be coercible to VECTOR(DOUBLE)", + PFloatArray.INSTANCE.isCoercibleTo(PVectorDouble.INSTANCE)); + } + + @Test + public void testDoubleArrayCoercibleToVectorFloat() { + assertTrue("DOUBLE ARRAY should be coercible to VECTOR(FLOAT)", + PDoubleArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testDoubleArrayCoercibleToVectorDouble() { + assertTrue("DOUBLE ARRAY should be coercible to VECTOR(DOUBLE)", + PDoubleArray.INSTANCE.isCoercibleTo(PVectorDouble.INSTANCE)); + } + + @Test + public void testIntegerArrayCoercibleToVector() { + assertTrue( + "INTEGER ARRAY should be coercible to VECTOR(FLOAT) " + "(integer is coercible to double)", + PIntegerArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testDecimalArrayCoercibleToVector() { + assertTrue("DECIMAL ARRAY should be coercible to VECTOR(FLOAT)", + PDecimalArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testLongArrayCoercibleToVector() { + assertTrue("BIGINT ARRAY should be coercible to VECTOR(FLOAT)", + PLongArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testSmallintArrayCoercibleToVector() { + assertTrue("SMALLINT ARRAY should be coercible to VECTOR(FLOAT)", + PSmallintArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testTinyintArrayCoercibleToVector() { + assertTrue("TINYINT ARRAY should be coercible to VECTOR(FLOAT)", + PTinyintArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + /** Array types that are not numeric cannot be coerced to vector types. */ + @Test + public void testVarcharArrayNotCoercibleToVector() { + assertFalse("VARCHAR ARRAY must not be coercible to VECTOR(FLOAT)", + PVarcharArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testBooleanArrayNotCoercibleToVector() { + assertFalse("BOOLEAN ARRAY must not be coercible to VECTOR(FLOAT)", + PBooleanArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testBinaryArrayNotCoercibleToVector() { + assertFalse("VARBINARY ARRAY must not be coercible to VECTOR(FLOAT)", + PVarbinaryArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testDateArrayNotCoercibleToVector() { + assertFalse("DATE ARRAY must not be coercible to VECTOR(FLOAT)", + PDateArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testFloatArrayCastableToVectorFloat() { + assertTrue("FLOAT ARRAY should be castable to VECTOR(FLOAT)", + PFloatArray.INSTANCE.isCastableTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testDoubleArrayCastableToVectorDouble() { + assertTrue("DOUBLE ARRAY should be castable to VECTOR(DOUBLE)", + PDoubleArray.INSTANCE.isCastableTo(PVectorDouble.INSTANCE)); + } + + @Test + public void testVarcharArrayNotCastableToVector() { + assertFalse("VARCHAR ARRAY must not be castable to VECTOR(FLOAT)", + PVarcharArray.INSTANCE.isCastableTo(PVectorFloat.INSTANCE)); + } + + @Test + public void testFloatArrayTwoArgCoercibleToVector() { + assertTrue("PFloatArray.isCoercibleTo(PVectorFloat, PFloatArray) should return true", + PFloatArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE, PFloatArray.INSTANCE)); + } + + @Test + public void testDoubleArrayTwoArgCoercibleToVector() { + assertTrue("PDoubleArray.isCoercibleTo(PVectorDouble, PDoubleArray) should return true", + PDoubleArray.INSTANCE.isCoercibleTo(PVectorDouble.INSTANCE, PDoubleArray.INSTANCE)); + } + + @Test + public void testDecimalArrayTwoArgCoercibleToVector() { + assertTrue("PDecimalArray.isCoercibleTo(PVectorFloat, PDecimalArray) should return true", + PDecimalArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE, PDecimalArray.INSTANCE)); + } + + @Test + public void testVarcharArrayTwoArgNotCoercibleToVector() { + assertFalse("PVarcharArray.isCoercibleTo(PVectorFloat, PVarcharArray) must return false", + PVarcharArray.INSTANCE.isCoercibleTo(PVectorFloat.INSTANCE, PVarcharArray.INSTANCE)); + } + + @Test + public void testFloatArrayStillCoercibleToFloatArray() { + assertTrue("FLOAT ARRAY should still be coercible to FLOAT ARRAY", + PFloatArray.INSTANCE.isCoercibleTo(PFloatArray.INSTANCE)); + } + + @Test + public void testFloatArrayStillCoercibleToDoubleArray() { + assertTrue("FLOAT ARRAY should still be coercible to DOUBLE ARRAY", + PFloatArray.INSTANCE.isCoercibleTo(PDoubleArray.INSTANCE)); + } + + @Test + public void testFloatArrayNotCoercibleToVarcharArray() { + assertFalse("FLOAT ARRAY must not be coercible to VARCHAR ARRAY", + PFloatArray.INSTANCE.isCoercibleTo(PVarcharArray.INSTANCE)); + } +}
