Skip to content

Commit ee36b1f

Browse files
jbajicjbajic
authored
[COR-45] Add vector index hints (#22141)
* Add vector index hints * Rearrange index hinting logic * Add tests * Add changelog * Remove unused * Remove unintended files (feed and clang-format temp file) * Remove changes on feed * Remove * Remove redundant
1 parent 954344b commit ee36b1f

File tree

3 files changed

+251
-5
lines changed

3 files changed

+251
-5
lines changed

‎CHANGELOG‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
devel
22
-----
33

4+
* COR-45: Added support for index hints in vector search queries.
5+
46
* COR-13: New ArangoSearch consolidation policy
57
- new algorithm that uses maxSkewThreshold and minDeletionRatio
68
- Bumped ArangoSearch
79

8-
* The query explainer now displays the vector index used by EnumerateNearVectoNode.
10+
* The query explainer now displays the vector index used by EnumerateNearVectorNode.
911

1012
* FE-636: bump webpack-dev-middleware from 5.3.3 to 5.3.4.
1113

‎arangod/Aql/Optimizer/Rule/OptimizerRuleVectorIndex.cpp‎

Lines changed: 39 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -233,14 +233,46 @@ AstNode* getApproxNearAttributeExpression(
233233
}
234234

235235
std::vector<std::shared_ptr<Index>> getVectorIndexes(
236-
auto& enumerateCollectionNode) {
236+
auto& enumerateCollectionNode, IndexHint const& indexHint) {
237+
auto const& indexes = enumerateCollectionNode->collection()->indexes();
238+
if (indexHint.isSimple() && indexHint.isForced()) {
239+
auto& idxNames = indexHint.candidateIndexes();
240+
for (auto const& idxName : idxNames) {
241+
if (auto foundHintedIndexIt = std::ranges::find_if(
242+
indexes,
243+
[&idxName](const auto& elem) { return elem->name() == idxName; });
244+
foundHintedIndexIt != indexes.end()) {
245+
return {*foundHintedIndexIt};
246+
} else {
247+
THROW_ARANGO_EXCEPTION_MESSAGE(
248+
TRI_ERROR_QUERY_FORCED_INDEX_HINT_UNUSABLE,
249+
std::format("Could not use hinted vector index `{}`", idxName));
250+
}
251+
}
252+
}
253+
237254
std::vector<std::shared_ptr<Index>> vectorIndexes;
238-
auto indexes = enumerateCollectionNode->collection()->indexes();
239255
std::ranges::copy_if(
240256
indexes, std::back_inserter(vectorIndexes), [](auto const& elem) {
241257
return elem->type() == Index::IndexType::TRI_IDX_TYPE_VECTOR_INDEX;
242258
});
243259

260+
// Reorder indexes
261+
if (indexHint.isSimple()) {
262+
auto& idxNames = indexHint.candidateIndexes();
263+
auto currentIt = vectorIndexes.begin();
264+
for (auto const& idxName : idxNames) {
265+
if (auto foundHintedIndexIt = std::ranges::find_if(
266+
vectorIndexes,
267+
[&idxName](const auto& elem) { return elem->name() == idxName; });
268+
foundHintedIndexIt != vectorIndexes.end()) {
269+
auto oldIt = currentIt;
270+
std::iter_swap(oldIt, foundHintedIndexIt);
271+
currentIt++;
272+
}
273+
}
274+
}
275+
244276
return vectorIndexes;
245277
}
246278

@@ -261,7 +293,7 @@ std::vector<std::shared_ptr<Index>> getVectorIndexes(
261293
//
262294
// Filter Pushdown:
263295
// When a filter expression is detected, the rule triggers
264-
// pushFilterIntoEnumerateNear which will handle the filter expressing usage
296+
// pushFilterIntoEnumerateNear which will handle the filter expression usage
265297
// along with vector index.
266298
void useVectorIndexRule(Optimizer* opt, std::unique_ptr<ExecutionPlan> plan,
267299
OptimizerRule const& rule) {
@@ -275,7 +307,10 @@ void useVectorIndexRule(Optimizer* opt, std::unique_ptr<ExecutionPlan> plan,
275307
ExecutionNode::castTo<EnumerateCollectionNode*>(node);
276308

277309
// check if there are vector indexes on collection
278-
auto const vectorIndexes = getVectorIndexes(enumerateCollectionNode);
310+
auto const& colNodeHints = enumerateCollectionNode->hint();
311+
// We will prioritize the hinted indexes
312+
auto const vectorIndexes =
313+
getVectorIndexes(enumerateCollectionNode, colNodeHints);
279314
if (vectorIndexes.empty()) {
280315
continue;
281316
}

‎tests/js/client/aql/vector/aql-vector-index-hints.js‎

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
/*jshint globalstrict:false, strict:false, maxlen: 500 */
2+
/*global fail, assertEqual, assertNotEqual, assertTrue, assertFalse */
3+
4+
// //////////////////////////////////////////////////////////////////////////////
5+
// / DISCLAIMER
6+
// /
7+
// / Copyright 2014-2024 ArangoDB GmbH, Cologne, Germany
8+
// / Copyright 2004-2014 triAGENS GmbH, Cologne, Germany
9+
// /
10+
// / Licensed under the Business Source License 1.1 (the "License");
11+
// / you may not use this file except in compliance with the License.
12+
// / You may obtain a copy of the License at
13+
// /
14+
// / https://github.com/arangodb/arangodb/blob/devel/LICENSE
15+
// /
16+
// / Unless required by applicable law or agreed to in writing, software
17+
// / distributed under the License is distributed on an "AS IS" BASIS,
18+
// / WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19+
// / See the License for the specific language governing permissions and
20+
// / limitations under the License.
21+
// /
22+
// / Copyright holder is ArangoDB GmbH, Cologne, Germany
23+
// /
24+
/// @author Jure Bajic
25+
// //////////////////////////////////////////////////////////////////////////////
26+
27+
const internal = require("internal");
28+
const jsunity = require("jsunity");
29+
const arangodb = require("@arangodb");
30+
const helper = require("@arangodb/aql-helper");
31+
const errors = internal.errors;
32+
const db = internal.db;
33+
const {
34+
randomNumberGeneratorFloat,
35+
randomInteger,
36+
} = require("@arangodb/testutils/seededRandom");
37+
38+
const dbName = "vectorIndexHintDb";
39+
const collName = "vectorIndexHintColl";
40+
41+
////////////////////////////////////////////////////////////////////////////////
42+
/// @brief helper function to get the used vector index from a query plan
43+
////////////////////////////////////////////////////////////////////////////////
44+
function getVectorIndexName(query, bindVars = {}) {
45+
const explain = db._createStatement({ query, bindVars }).explain();
46+
const vectorNodes = explain.plan.nodes.filter(
47+
(node) => node.type === "EnumerateNearVectorNode"
48+
);
49+
if (vectorNodes.length === 0) {
50+
return null;
51+
}
52+
assertEqual(1, vectorNodes.length, "Expected exactly one EnumerateNearVectorNode");
53+
return vectorNodes[0].index ? vectorNodes[0].index.name : null;
54+
}
55+
56+
////////////////////////////////////////////////////////////////////////////////
57+
/// @brief Test suite for vector index hints
58+
////////////////////////////////////////////////////////////////////////////////
59+
function VectorIndexHintsSuite() {
60+
let collection;
61+
let randomPoint;
62+
const dimension = 128;
63+
const numberOfDocs = 100;
64+
const seed = randomInteger();
65+
66+
return {
67+
setUpAll: function () {
68+
db._createDatabase(dbName);
69+
db._useDatabase(dbName);
70+
71+
collection = db._create(collName, { numberOfShards: 3 });
72+
73+
// Generate random vectors
74+
let docs = [];
75+
let gen = randomNumberGeneratorFloat(seed);
76+
for (let i = 0; i < numberOfDocs; ++i) {
77+
const vector = Array.from({ length: dimension }, () => gen());
78+
const vectorCosine = Array.from({ length: dimension }, () => gen());
79+
const vectorInnerProduct = Array.from({ length: dimension }, () => gen());
80+
81+
if (i === Math.floor(numberOfDocs / 2)) {
82+
randomPoint = vector;
83+
}
84+
85+
docs.push({
86+
_key: `doc${i}`,
87+
vector: vector,
88+
vectorCosine: vectorCosine,
89+
vectorInnerProduct: vectorInnerProduct,
90+
value: i,
91+
});
92+
}
93+
collection.insert(docs);
94+
95+
collection.ensureIndex({
96+
name: "vector_l2",
97+
type: "vector",
98+
fields: ["vector"],
99+
params: {
100+
metric: "l2",
101+
dimension: dimension,
102+
nLists: 5,
103+
},
104+
});
105+
106+
collection.ensureIndex({
107+
name: "vector_l2_secondary",
108+
type: "vector",
109+
fields: ["vector"],
110+
params: {
111+
metric: "l2",
112+
dimension: dimension,
113+
nLists: 3,
114+
},
115+
});
116+
},
117+
118+
tearDownAll: function () {
119+
db._useDatabase("_system");
120+
db._dropDatabase(dbName);
121+
},
122+
123+
testVectorL2HintFirstIndex: function () {
124+
const query = `
125+
FOR doc IN ${collName} OPTIONS { indexHint: "vector_l2" }
126+
SORT APPROX_NEAR_L2(doc.vector, @qp)
127+
LIMIT 5
128+
RETURN doc
129+
`;
130+
const bindVars = { qp: randomPoint };
131+
const indexName = getVectorIndexName(query, bindVars);
132+
133+
assertEqual("vector_l2", indexName);
134+
},
135+
136+
testVectorL2HintSecondIndex: function () {
137+
const query = `
138+
FOR doc IN ${collName} OPTIONS { indexHint: "vector_l2_secondary" }
139+
SORT APPROX_NEAR_L2(doc.vector, @qp)
140+
LIMIT 5
141+
RETURN doc
142+
`;
143+
const bindVars = { qp: randomPoint };
144+
const indexName = getVectorIndexName(query, bindVars);
145+
146+
assertEqual("vector_l2_secondary", indexName);
147+
},
148+
149+
testVectorL2HintNonExistingIndex: function () {
150+
const query = `
151+
FOR doc IN ${collName} OPTIONS { indexHint: "nonexistent_index" }
152+
SORT APPROX_NEAR_L2(doc.vector, @qp)
153+
LIMIT 5
154+
RETURN doc
155+
`;
156+
const bindVars = { qp: randomPoint };
157+
const indexName = getVectorIndexName(query, bindVars);
158+
159+
assertTrue(["vector_l2_secondary", "vector_l2"].includes(indexName));
160+
},
161+
162+
testVectorL2WithHintForced: function () {
163+
const query = `
164+
FOR doc IN ${collName} OPTIONS { indexHint: "vector_l2_secondary", forceIndexHint: true }
165+
SORT APPROX_NEAR_L2(doc.vector, @qp)
166+
LIMIT 5
167+
RETURN doc
168+
`;
169+
const bindVars = { qp: randomPoint };
170+
const indexName = getVectorIndexName(query, bindVars);
171+
172+
assertEqual("vector_l2_secondary", indexName);
173+
},
174+
175+
testVectorL2WithHintForcedNonExistingIndex: function () {
176+
const query = `
177+
FOR doc IN ${collName} OPTIONS { indexHint: "nonexistent_index", forceIndexHint: true }
178+
SORT APPROX_NEAR_L2(doc.vector, @qp)
179+
LIMIT 5
180+
RETURN doc
181+
`;
182+
const bindVars = { qp: randomPoint };
183+
184+
try {
185+
db._createStatement({ query, bindVars }).explain();
186+
fail();
187+
} catch (err) {
188+
assertEqual(errors.ERROR_QUERY_FORCED_INDEX_HINT_UNUSABLE.code, err.errorNum);
189+
}
190+
},
191+
192+
testVectorL2HintMultipleIndexes: function () {
193+
const query = `
194+
FOR doc IN ${collName} OPTIONS { indexHint: ["non_existing_index", "vector_l2_secondary"] }
195+
SORT APPROX_NEAR_L2(doc.vector, @qp)
196+
LIMIT 5
197+
RETURN doc
198+
`;
199+
const bindVars = { qp: randomPoint };
200+
const indexName = getVectorIndexName(query, bindVars);
201+
202+
assertEqual("vector_l2_secondary", indexName);
203+
},
204+
};
205+
}
206+
207+
jsunity.run(VectorIndexHintsSuite);
208+
209+
return jsunity.done();

0 commit comments

Comments
 (0)