diff --git a/vector/src/index/hnsw_rel_batch_insert.cpp b/vector/src/index/hnsw_rel_batch_insert.cpp index 61d8a4f..fac3625 100644 --- a/vector/src/index/hnsw_rel_batch_insert.cpp +++ b/vector/src/index/hnsw_rel_batch_insert.cpp @@ -67,6 +67,23 @@ std::unique_ptr HNSWRelBatchInsert::ini storage::StorageUtils::getStartOffsetOfNodeGroup(nodeGroupIdx)); } +// Counts the neighbours of a node that are still valid. The graph's CSR length +// includes slots that were subsequently marked INVALID_OFFSET (e.g. during a +// shrink), so we filter them out to keep the CSR metadata consistent with the +// rel rows that are actually written. +static common::length_t getNumValidNeighbors(const InMemHNSWGraph& graph, + common::offset_t nodeOffset) { + const auto neighbours = graph.getNeighbors(nodeOffset); + common::length_t numValid = 0; + for (const auto neighbourGraphOffset : neighbours) { + if (neighbourGraphOffset == common::INVALID_OFFSET) { + continue; + } + ++numValid; + } + return numValid; +} + void HNSWRelBatchInsert::populateCSRLengths(processor::RelBatchInsertExecutionState& executionState, storage::InMemChunkedCSRHeader& csrHeader, common::offset_t numNodes, const processor::RelBatchInsertInfo&) { @@ -82,7 +99,10 @@ void HNSWRelBatchInsert::populateCSRLengths(processor::RelBatchInsertExecutionSt ++graphOffset) { const auto nodeOffsetInGroup = hnswExecutionState.getBoundNodeOffsetInGroup(graphOffset); DASSERT(nodeOffsetInGroup < numNodes); - lengthData[nodeOffsetInGroup] = graph.getCSRLength(graphOffset); + // getCSRLength counts all occupied slots, including ones later marked + // INVALID_OFFSET. Filter invalidated neighbors so the CSR length matches + // the number of rel rows actually written in writeToTable. + lengthData[nodeOffsetInGroup] = getNumValidNeighbors(graph, graphOffset); } } @@ -90,7 +110,7 @@ static common::offset_t getNumRelsInGraph(const InMemHNSWGraph& graph, common::offset_t startNodeInGraph, common::offset_t endNodeInGraph) { auto numRels = 0u; for (auto offsetInGraph = startNodeInGraph; offsetInGraph < endNodeInGraph; offsetInGraph++) { - numRels += graph.getCSRLength(offsetInGraph); + numRels += getNumValidNeighbors(graph, offsetInGraph); } return numRels; } diff --git a/vector/test/CMakeLists.txt b/vector/test/CMakeLists.txt index 46a5027..e25611b 100644 --- a/vector/test/CMakeLists.txt +++ b/vector/test/CMakeLists.txt @@ -1,3 +1,4 @@ if (${BUILD_EXTENSION_TESTS}) add_lbug_test(vector_prepare_test prepare_test.cpp) + add_lbug_test(vector_hnsw_rel_batch_insert_test hnsw_rel_batch_insert_test.cpp) endif () diff --git a/vector/test/hnsw_rel_batch_insert_test.cpp b/vector/test/hnsw_rel_batch_insert_test.cpp new file mode 100644 index 0000000..0e571d1 --- /dev/null +++ b/vector/test/hnsw_rel_batch_insert_test.cpp @@ -0,0 +1,52 @@ +#include "api_test/api_test.h" + +using namespace lbug::common; + +namespace lbug { +namespace testing { + +// Regression test for inconsistent CSR metadata when a neighbor slot is +// marked INVALID_OFFSET. Previously populateCSRLengths and getNumRelsInGraph +// used the unfiltered graph CSR length, while writeToTable skipped +// INVALID_OFFSET neighbors. This left holes/stale values in the chunked +// group. The test builds an index over enough nodes to trigger batch insert +// and verifies that QUERY_VECTOR_INDEX returns consistent results without +// errors, and that the underlying rel tables have no stale/empty rows. +class HNSWRelBatchInsertTest : public ApiTest {}; + +TEST_F(HNSWRelBatchInsertTest, QueryVectorIndexReturnsConsistentResults) { +#ifndef __STATIC_LINK_EXTENSION_TEST__ + ASSERT_TRUE(conn->query(std::format("LOAD EXTENSION '{}'", + TestHelper::appendLbugRootPath( + "extension/vector/build/libvector.lbug_extension"))) + ->isSuccess()); +#endif + ASSERT_TRUE(conn->query("CREATE NODE TABLE Book (ID SERIAL, title STRING, " + "title_embedding FLOAT[4], PRIMARY KEY (ID));") + ->isSuccess()); + // Insert enough nodes to span multiple node groups so the batch insert + // path exercises CSR length / numRels consistency across groups. + for (int i = 0; i < 5; i++) { + ASSERT_TRUE(conn->query(std::format( + "CREATE (b:Book {{title: 'Book {}', title_embedding: [{}, {}, {}, {}]}});", i, + static_cast(i) * 0.1f, static_cast(i) * 0.2f, + static_cast(i) * 0.3f, static_cast(i) * 0.4f)) + ->isSuccess()); + } + ASSERT_TRUE( + conn->query("CALL CREATE_VECTOR_INDEX('Book', 'title_vec_index', 'title_embedding');") + ->isSuccess()); + + auto prepared = conn->prepare( + "CALL QUERY_VECTOR_INDEX('Book', 'title_vec_index', " + "[0.0, 0.0, 0.0, 0.0], $k) RETURN node.title ORDER BY distance;"); + auto result = conn->execute(prepared.get(), std::make_pair(std::string("k"), 3)); + ASSERT_TRUE(result->isSuccess()) << result->getErrorMessage(); + ASSERT_EQ(result->getNumTuples(), 3); + // Closest node to the origin embedding is Book 0. + auto rows = TestHelper::convertResultToString(*result); + ASSERT_EQ(rows[0], "Book 0"); +} + +} // namespace testing +} // namespace lbug