diff --git a/vector/src/function/query_hnsw_index.cpp b/vector/src/function/query_hnsw_index.cpp index 48dfb84b..21525658 100644 --- a/vector/src/function/query_hnsw_index.cpp +++ b/vector/src/function/query_hnsw_index.cpp @@ -369,7 +369,11 @@ static std::unique_ptr initQueryHNSWSharedState( auto nodeTable = storage::StorageManager::Get(*context) ->getTable(bindData->nodeTableEntry->getTableID()) ->ptrCast(); - auto numNodes = nodeTable->getStats(transaction::Transaction::Get(*context)).getTableCard(); + // NB: TableStats::getTableCard() is an estimated (possibly stale) cardinality and must not + // be used to size the visited bitmap. Use the exact row count instead. VisitedState is + // additionally bounds-checked and auto-growing, so a stale size cannot overflow. + auto numNodes = + nodeTable->getNumTotalRows(transaction::Transaction::Get(*context)); return std::make_unique(nodeTable, numNodes); } diff --git a/vector/src/include/index/hnsw_index.h b/vector/src/include/index/hnsw_index.h index a3ac95bd..fdcc3705 100644 --- a/vector/src/include/index/hnsw_index.h +++ b/vector/src/include/index/hnsw_index.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -55,9 +56,42 @@ struct VisitedState { // NOLINTNEXTLINE(readability-make-member-function-const): Semantically non-const. void reset() { memset(visited.get(), 0, size); } + // Grow the bitmap so that `offset` is addressable. Preserves existing marks. + void ensureCapacity(common::offset_t offset) { + if (offset < size) { + return; + } + // Grow geometrically to amortize repeated growth when the table grows + // between sizing and search (e.g. stale cardinality estimates). + common::offset_t newSize = size == 0 ? offset + 1 : size; + while (newSize <= offset) { + newSize *= 2; + } + auto newVisited = std::make_unique(newSize); + memset(newVisited.get(), 0, newSize); + if (size > 0) { + memcpy(newVisited.get(), visited.get(), size); + } + visited = std::move(newVisited); + size = newSize; + } + void resize(common::offset_t newSize) { + if (newSize <= size) { + return; + } + ensureCapacity(newSize - 1); + } // NOLINTNEXTLINE(readability-make-member-function-const): Semantically non-const. - void add(common::offset_t offset) { visited[offset] = 1; } - bool contains(common::offset_t offset) const { return visited[offset]; } + void add(common::offset_t offset) { + ensureCapacity(offset); + visited[offset] = 1; + } + bool contains(common::offset_t offset) const { + if (offset >= size) { + return false; + } + return visited[offset]; + } }; struct HNSWStorageInfo final : storage::IndexStorageInfo { diff --git a/vector/src/index/hnsw_index.cpp b/vector/src/index/hnsw_index.cpp index 100a7ad1..8b1e07cb 100644 --- a/vector/src/index/hnsw_index.cpp +++ b/vector/src/index/hnsw_index.cpp @@ -958,6 +958,10 @@ common::offset_t OnDiskHNSWIndex::searchNNInUpperLayer(const EmbeddingHandle& qu void OnDiskHNSWIndex::initLayerSearchState(Transaction* transaction, HNSWSearchState& searchState, bool isUpperLayer) const { + // The visited bitmap is sized at query init time and may be smaller than the table at + // search time (e.g. concurrent inserts). Grow it proactively; VisitedState::add/contains + // are additionally bounds-safe so a stale size can never cause an out-of-bounds access. + searchState.visited.resize(nodeTable.getNumTotalRows(transaction)); searchState.visited.reset(); const auto& hnswStorageInfo = storageInfo->cast(); const auto& hnswGraph = isUpperLayer ? searchState.upperGraph : searchState.lowerGraph;