-
Notifications
You must be signed in to change notification settings - Fork 33
[MOD-17706] finalize SVSIndex::relabelVector #1045
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 12 commits
48fe6db
a68d357
326edc0
06d42e6
f548f3d
4115d72
219cd1d
7066792
743cd20
086a3ba
87d8fc5
b7c0c38
cd09a53
52c5917
98e7edc
88be7eb
492ae88
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -106,12 +106,11 @@ class PyVecSimIndex { | |
| template <typename DataType, typename DistType, typename NPArrayType = DataType> | ||
| inline py::object rawVectorsAsNumpy(labelType label, size_t dim) { | ||
| std::vector<std::vector<DataType>> vectors; | ||
| if (index->basicInfo().algo == VecSimAlgo_BF) { | ||
| dynamic_cast<BruteForceIndex<DataType, DistType> *>(this->index.get()) | ||
| ->getDataByLabel(label, vectors); | ||
| if (auto *tiered = | ||
| dynamic_cast<VecSimTieredIndex<DataType, DistType> *>(this->index.get())) { | ||
| tiered->getDataByLabel(label, vectors); | ||
| } else { | ||
| // index is HNSW | ||
| dynamic_cast<HNSWIndex<DataType, DistType> *>(this->index.get()) | ||
| dynamic_cast<VecSimIndexAbstract<DataType, DistType> *>(this->index.get()) | ||
| ->getDataByLabel(label, vectors); | ||
| } | ||
| size_t n_vectors = vectors.size(); | ||
|
|
@@ -216,6 +215,11 @@ class PyVecSimIndex { | |
|
|
||
| void runGC() { VecSimTieredIndex_GC(index.get()); } | ||
|
|
||
| VecSimRelabelCode relabelVector(labelType old_label, labelType new_label) { | ||
| py::gil_scoped_release py_gil; | ||
| return VecSimIndex_RelabelVector(index.get(), old_label, new_label); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do you think we need to find a way to synchronize the new
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. decided to remove completly the parael tests on plain indexes. we have the concurrent tests on tiered instead |
||
| } | ||
|
|
||
| py::object getVector(labelType label) { | ||
| VecSimIndexBasicInfo info = index->basicInfo(); | ||
| size_t dim = info.dim; | ||
|
|
@@ -713,6 +717,14 @@ PYBIND11_MODULE(VecSim, m) { | |
| .def_readwrite("initialCapacity", &BFParams::initialCapacity) | ||
| .def_readwrite("blockSize", &BFParams::blockSize); | ||
|
|
||
| py::enum_<VecSimRelabelCode>(m, "VecSimRelabelCode") | ||
| .value("VecSimRelabel_OK", VecSimRelabel_OK) | ||
| .value("VecSimRelabel_OldLabelMissing", VecSimRelabel_OldLabelMissing) | ||
| .value("VecSimRelabel_NewLabelTaken", VecSimRelabel_NewLabelTaken) | ||
| .value("VecSimRelabel_SameLabel", VecSimRelabel_SameLabel) | ||
| .value("VecSimRelabel_Unsupported", VecSimRelabel_Unsupported) | ||
| .export_values(); | ||
|
|
||
| py::enum_<VecSimSvsQuantBits>(m, "VecSimSvsQuantBits") | ||
| .value("VecSimSvsQuant_NONE", VecSimSvsQuant_NONE) | ||
| .value("VecSimSvsQuant_Scalar", VecSimSvsQuant_Scalar) | ||
|
|
@@ -799,6 +811,8 @@ PYBIND11_MODULE(VecSim, m) { | |
| .def("create_batch_iterator", &PyVecSimIndex::createBatchIterator, py::arg("query_blob"), | ||
| py::arg("query_param") = nullptr) | ||
| .def("get_vector", &PyVecSimIndex::getVector) | ||
| .def("relabel_vector", &PyVecSimIndex::relabelVector, py::arg("old_label"), | ||
| py::arg("new_label")) | ||
| .def("run_gc", &PyVecSimIndex::runGC); | ||
|
|
||
| py::class_<PyHNSWLibIndex, PyVecSimIndex>(m, "HNSWIndex") | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -621,3 +621,78 @@ def test_multi_range_query(test_logger): | |
| # Expect zero results for radius==0 | ||
| tiered_labels, tiered_distances = index.range_query(query_data, radius=0) | ||
| assert len(tiered_labels[0]) == 0 | ||
|
|
||
|
|
||
| def test_relabel_vector(test_logger): | ||
| dim = 16 | ||
| num_elements = 1000 | ||
| hnsw_params = create_hnsw_params(dim, num_elements, VecSimMetric_L2, VecSimType_FLOAT32) | ||
| # A flat buffer large enough to hold everything, so the relabel below has a real chance of | ||
| # landing while the vector is still buffered with a pending ingest job. | ||
| index = Tiered_HNSWIndex(hnsw_params, create_tiered_hnsw_params(), num_elements) | ||
|
|
||
| data = np.float32(np.random.random((num_elements, dim))) | ||
| for label, vector in enumerate(data): | ||
| index.add_vector(vector, label) | ||
|
|
||
| # Relabel one early and one late label. The workers ingest in insertion order, so by now the | ||
| # early one is most likely already in HNSW while the late one is most likely still buffered | ||
| # with a pending ingest job - between them the two tiers both get covered. The buffered case is | ||
| # the delicate one: a job left holding the old label would either ingest the vector under it or | ||
| # throw out of the worker thread, and neither would survive the assertions below. | ||
| buffered = index.get_curr_bf_size() | ||
| test_logger.info(f"relabeling with {buffered} of {num_elements} vectors still buffered") | ||
| moved = {7: num_elements + 500, num_elements - 1: num_elements + 501} | ||
| for old_label, new_label in moved.items(): | ||
| assert index.relabel_vector(old_label, new_label) == VecSimRelabel_OK | ||
|
|
||
| index.wait_for_index() | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could we also test
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. tested it in C++ now |
||
|
|
||
| # Once ingestion has drained, the vector sits in HNSW under the new label and under no other. | ||
| assert index.index_size() == num_elements | ||
| assert index.hnsw_label_count() == num_elements | ||
| for old_label, new_label in moved.items(): | ||
| assert_allclose(index.get_vector(new_label)[0], data[old_label], rtol=1e-6) | ||
| assert index.get_vector(old_label).shape == (0, dim) | ||
|
|
||
| labels, distances = index.knn_query(data[old_label], 1) | ||
| assert labels[0][0] == new_label | ||
| assert distances[0][0] < 1e-6 | ||
|
|
||
| # Each rejection is reported distinctly, and none of them modifies the index. | ||
| assert index.relabel_vector(num_elements + 1, 0) == VecSimRelabel_OldLabelMissing | ||
| assert index.relabel_vector(0, 1) == VecSimRelabel_NewLabelTaken | ||
| assert index.relabel_vector(0, 0) == VecSimRelabel_SameLabel | ||
| assert index.index_size() == num_elements | ||
| test_logger.info("tiered relabel_vector moved the label across both tiers") | ||
|
|
||
|
|
||
| def test_relabel_vector_multi(test_logger): | ||
| dim = 16 | ||
| num_labels = 200 | ||
| per_label = 5 | ||
| hnsw_params = create_hnsw_params(dim, num_labels * per_label, VecSimMetric_L2, | ||
| VecSimType_FLOAT32, is_multi=True) | ||
| index = Tiered_HNSWIndex(hnsw_params, create_tiered_hnsw_params(), num_labels * per_label) | ||
|
|
||
| data = np.float32(np.random.random((num_labels, per_label, dim))) | ||
| for label in range(num_labels): | ||
| for vector in data[label]: | ||
| index.add_vector(vector, label) | ||
|
|
||
| # In a multi index a label can hold several pending ingest jobs at once, so a late label | ||
| # exercises re-keying all of them together while an early one is most likely already in HNSW. | ||
| buffered = index.get_curr_bf_size() | ||
| test_logger.info(f"relabeling with {buffered} of {num_labels * per_label} vectors buffered") | ||
| moved = {7: num_labels + 500, num_labels - 1: num_labels + 501} | ||
| for old_label, new_label in moved.items(): | ||
| assert index.relabel_vector(old_label, new_label) == VecSimRelabel_OK | ||
|
|
||
| index.wait_for_index() | ||
|
|
||
| assert index.index_size() == num_labels * per_label | ||
| assert index.hnsw_label_count() == num_labels | ||
| for old_label, new_label in moved.items(): | ||
| assert index.get_vector(new_label).shape == (per_label, dim) | ||
| assert index.get_vector(old_label).shape == (0, dim) | ||
| test_logger.info("tiered multi relabel_vector moved every vector under the label") | ||
Uh oh!
There was an error while loading. Please reload this page.