Parallel CreateHNSWIndex finalize (#4938)

Author: ray6080

src/function/function_collection.cpp

@@ -231,6 +231,7 @@ FunctionCollection* FunctionCollection::getFunctions() {
         // TODO(Guodong): Move this from builtin to extension and also move _CreateHNSWIndexFunction
         // and _DropHNSWIndexFunction to private functions.
         STANDALONE_TABLE_FUNCTION(InternalCreateHNSWIndexFunction),
+        STANDALONE_TABLE_FUNCTION(InternalFinalizeHNSWIndexFunction),
         STANDALONE_TABLE_FUNCTION(CreateHNSWIndexFunction),
         STANDALONE_TABLE_FUNCTION(InternalDropHNSWIndexFunction),
         STANDALONE_TABLE_FUNCTION(DropHNSWIndexFunction),

src/function/table/hnsw/create_hnsw_index.cpp

@@ -29,27 +29,40 @@ struct CreateHNSWIndexBindData final : TableFuncBindData {
     }
 };
 
-struct CreateHNSWSharedState final : TableFuncSharedState {
+struct CreateInMemHNSWSharedState final : TableFuncSharedState {
     std::string name;
-    std::unique_ptr<storage::InMemHNSWIndex> hnswIndex;
+    std::shared_ptr<storage::InMemHNSWIndex> hnswIndex;
     storage::NodeTable& nodeTable;
     common::offset_t numNodes;
     std::atomic<common::offset_t> numNodesInserted = 0;
 
     const CreateHNSWIndexBindData* bindData;
-    std::shared_ptr<storage::HNSWIndexPartitionerSharedState> partitionerSharedState;
 
-    explicit CreateHNSWSharedState(const CreateHNSWIndexBindData& bindData);
+    explicit CreateInMemHNSWSharedState(const CreateHNSWIndexBindData& bindData);
 };
 
-struct CreateHNSWLocalState final : TableFuncLocalState {
+struct CreateInMemHNSWLocalState final : TableFuncLocalState {
     storage::VisitedState upperVisited;
     storage::VisitedState lowerVisited;
 
-    explicit CreateHNSWLocalState(common::offset_t numNodes)
+    explicit CreateInMemHNSWLocalState(common::offset_t numNodes)
         : upperVisited{numNodes}, lowerVisited{numNodes} {}
 };
 
+struct FinalizeHNSWSharedState final : TableFuncSharedState {
+    std::shared_ptr<storage::InMemHNSWIndex> hnswIndex;
+    std::shared_ptr<storage::HNSWIndexPartitionerSharedState> partitionerSharedState;
+    std::unique_ptr<TableFuncBindData> bindData;
+
+    std::atomic<common::node_group_idx_t> numNodeGroupsFinalized = 0;
+
+    explicit FinalizeHNSWSharedState(storage::MemoryManager& mm) {
+        partitionerSharedState = std::make_shared<storage::HNSWIndexPartitionerSharedState>(mm);
+    }
+
+    TableFuncMorsel getMorsel() override;
+};
+
 struct BoundQueryHNSWIndexInput {
     catalog::NodeTableCatalogEntry* nodeTableEntry;
     catalog::IndexCatalogEntry* indexEntry;
@@ -105,6 +118,12 @@ struct InternalCreateHNSWIndexFunction final {
     static function_set getFunctionSet();
 };
 
+struct InternalFinalizeHNSWIndexFunction final {
+    static constexpr const char* name = "_FINALIZE_HNSW_INDEX";
+
+    static function_set getFunctionSet();
+};
+
 struct CreateHNSWIndexFunction final {
     static constexpr const char* name = "CREATE_HNSW_INDEX";
 

src/include/function/table/hnsw/hnsw_index_functions.h

@@ -34,7 +34,6 @@ class InMemEmbeddings final : public EmbeddingColumn {
         const common::LogicalType& columnType);
 
     float* getEmbedding(common::offset_t offset) const;
-    ColumnChunkData& getData() const { return *data; }
 
     void initialize(main::ClientContext* context, NodeTable& table,
         common::column_id_t columnID) override;
@@ -97,10 +96,8 @@ class InMemHNSWGraph {
         resetCSRLengthAndDstNodes();
     }
 
-    common::offset_vec_t getNeighbors(transaction::Transaction* transaction,
-        common::offset_t nodeOffset) const;
-    common::offset_vec_t getNeighbors(transaction::Transaction* transaction,
-        common::offset_t nodeOffset, common::length_t numNbrs) const;
+    common::offset_vec_t getNeighbors(common::offset_t nodeOffset) const;
+    common::offset_vec_t getNeighbors(common::offset_t nodeOffset, common::length_t numNbrs) const;
 
     common::length_t getMaxDegree() const { return maxDegree; }
 
@@ -120,14 +117,14 @@ class InMemHNSWGraph {
         dstNodes[csrOffset].store(dstNode, std::memory_order_relaxed);
     }
 
-    void finalize(MemoryManager& mm,
+    void finalize(MemoryManager& mm, common::node_group_idx_t nodeGroupIdx,
         const processor::PartitionerSharedState& partitionerSharedState);
 
 private:
     void resetCSRLengthAndDstNodes();
 
     void finalizeNodeGroup(MemoryManager& mm, common::node_group_idx_t nodeGroupIdx,
-        common::table_id_t srcNodeTableID, common::table_id_t dstNodeTableID,
+        uint64_t numRels, common::table_id_t srcNodeTableID, common::table_id_t dstNodeTableID,
         common::table_id_t relTableID, InMemChunkedNodeGroupCollection& partition) const;
 
     common::offset_t getDstNode(common::offset_t csrOffset) const {
@@ -142,8 +139,6 @@ class InMemHNSWGraph {
     std::atomic<common::offset_t>* dstNodes;
     // Max allowed degree of a node in the graph before shrinking.
     common::length_t maxDegree;
-
-    std::vector<common::length_t> numRelsPerNodeGroup;
 };
 
 } // namespace storage

src/include/storage/index/hnsw_graph.h

@@ -107,23 +107,18 @@ class InMemHNSWLayer {
     void setEntryPoint(common::offset_t offset) { entryPoint.store(offset); }
     common::offset_t getEntryPoint() const { return entryPoint.load(); }
 
-    void insert(transaction::Transaction* transaction, common::offset_t offset,
-        common::offset_t entryPoint_, VisitedState& visited);
-    common::offset_t searchNN(transaction::Transaction* transaction, common::offset_t node,
-        common::offset_t entryNode) const;
-    void shrink(transaction::Transaction* transaction);
-    void finalize(MemoryManager& mm,
+    void insert(common::offset_t offset, common::offset_t entryPoint_, VisitedState& visited);
+    common::offset_t searchNN(common::offset_t node, common::offset_t entryNode) const;
+    void finalize(MemoryManager& mm, common::node_group_idx_t nodeGroupIdx,
         const processor::PartitionerSharedState& partitionerSharedState) const;
 
 private:
-    std::vector<NodeWithDistance> searchKNN(transaction::Transaction* transaction,
-        const float* queryVector, common::offset_t entryNode, common::length_t k,
-        uint64_t configuredEf, VisitedState& visited) const;
-    static void shrinkForNode(transaction::Transaction* transaction, const InMemHNSWLayerInfo& info,
-        InMemHNSWGraph* graph, common::offset_t nodeOffset, common::length_t numNbrs);
+    std::vector<NodeWithDistance> searchKNN(const float* queryVector, common::offset_t entryNode,
+        common::length_t k, uint64_t configuredEf, VisitedState& visited) const;
+    static void shrinkForNode(const InMemHNSWLayerInfo& info, InMemHNSWGraph* graph,
+        common::offset_t nodeOffset, common::length_t numNbrs);
 
-    void insertRel(transaction::Transaction* transaction, common::offset_t srcNode,
-        common::offset_t dstNode);
+    void insertRel(common::offset_t srcNode, common::offset_t dstNode);
 
 private:
     std::atomic<common::offset_t> entryPoint;
@@ -145,10 +140,11 @@ class InMemHNSWIndex final : public HNSWIndex {
     common::offset_t getLowerEntryPoint() const override { return lowerLayer->getEntryPoint(); }
 
     // Note that the input is only `offset`, as we assume embeddings are already cached in memory.
-    void insert(common::offset_t offset, transaction::Transaction* transaction,
-        VisitedState& upperVisited, VisitedState& lowerVisited);
-    void shrink(transaction::Transaction* transaction);
-    void finalize(MemoryManager& mm, const HNSWIndexPartitionerSharedState& partitionerSharedState);
+    void insert(common::offset_t offset, VisitedState& upperVisited, VisitedState& lowerVisited);
+    void finalize(MemoryManager& mm, common::node_group_idx_t nodeGroupIdx,
+        const HNSWIndexPartitionerSharedState& partitionerSharedState);
+
+    void resetEmbeddings() { embeddings.reset(); }
 
 private:
     static constexpr int64_t INSERT_TO_UPPER_LAYER_RAND_UPPER_BOUND = 100;

src/include/storage/index/hnsw_index.h

@@ -9,9 +9,12 @@ class TableCatalogEntry;
 } // namespace catalog
 
 namespace storage {
-class EmbeddingColumn;
 
 struct HNSWIndexUtils {
+
+    static double computeDistance(DistFuncType funcType, const float* left, const float* right,
+        uint32_t dimension);
+
     static void validateColumnType(const catalog::TableCatalogEntry& tableEntry,
         const std::string& columnName);
 
@@ -29,9 +32,6 @@ struct HNSWIndexUtils {
         return &listChunk.getDataColumnChunk()->getData<T>()[offset * dimension];
     }
 
-    static double computeDistance(DistFuncType funcType, const float* left, const float* right,
-        uint32_t dimension);
-
 private:
     static void validateColumnType(const common::LogicalType& type);
 };

src/include/storage/index/hnsw_index_utils.h

@@ -87,32 +87,24 @@ float* OnDiskEmbeddings::getEmbedding(transaction::Transaction* transaction,
     return reinterpret_cast<float*>(dataVector->getData()) + value.offset;
 }
 
-void InMemHNSWGraph::finalize(MemoryManager& mm,
+// NOLINTNEXTLINE(readability-make-member-function-const): Semantically non-const function.
+void InMemHNSWGraph::finalize(MemoryManager& mm, common::node_group_idx_t nodeGroupIdx,
     const processor::PartitionerSharedState& partitionerSharedState) {
     const auto& partitionBuffers = partitionerSharedState.partitioningBuffers[0]->partitions;
-    const auto numNodeGroups = (numNodes + common::StorageConfig::NODE_GROUP_SIZE - 1) /
-                               common::StorageConfig::NODE_GROUP_SIZE;
-    KU_ASSERT(numNodeGroups == partitionerSharedState.numPartitions[0]);
-    numRelsPerNodeGroup.resize(numNodeGroups);
-    for (auto nodeGroupIdx = 0u; nodeGroupIdx < numNodeGroups; nodeGroupIdx++) {
-        auto numRels = 0u;
-        const auto startNodeOffset = StorageUtils::getStartOffsetOfNodeGroup(nodeGroupIdx);
-        const auto numNodesInGroup =
-            std::min(common::StorageConfig::NODE_GROUP_SIZE, numNodes - startNodeOffset);
-        for (auto i = 0u; i < numNodesInGroup; i++) {
-            numRels += getCSRLength(startNodeOffset + i);
-        }
-        numRelsPerNodeGroup[nodeGroupIdx] = numRels;
-    }
-    for (auto nodeGroupIdx = 0u; nodeGroupIdx < numNodeGroups; nodeGroupIdx++) {
-        finalizeNodeGroup(mm, nodeGroupIdx, partitionerSharedState.srcNodeTable->getTableID(),
-            partitionerSharedState.dstNodeTable->getTableID(),
-            partitionerSharedState.relTable->getTableID(), *partitionBuffers[nodeGroupIdx]);
+    auto numRels = 0u;
+    const auto startNodeOffset = StorageUtils::getStartOffsetOfNodeGroup(nodeGroupIdx);
+    const auto numNodesInGroup =
+        std::min(common::StorageConfig::NODE_GROUP_SIZE, numNodes - startNodeOffset);
+    for (auto i = 0u; i < numNodesInGroup; i++) {
+        numRels += getCSRLength(startNodeOffset + i);
     }
+    finalizeNodeGroup(mm, nodeGroupIdx, numRels, partitionerSharedState.srcNodeTable->getTableID(),
+        partitionerSharedState.dstNodeTable->getTableID(),
+        partitionerSharedState.relTable->getTableID(), *partitionBuffers[nodeGroupIdx]);
 }
 
 void InMemHNSWGraph::finalizeNodeGroup(MemoryManager& mm, common::node_group_idx_t nodeGroupIdx,
-    common::table_id_t srcNodeTableID, common::table_id_t dstNodeTableID,
+    uint64_t numRels, common::table_id_t srcNodeTableID, common::table_id_t dstNodeTableID,
     common::table_id_t relTableID, InMemChunkedNodeGroupCollection& partition) const {
     const auto startNodeOffset = StorageUtils::getStartOffsetOfNodeGroup(nodeGroupIdx);
     const auto numNodesInGroup =
@@ -122,10 +114,8 @@ void InMemHNSWGraph::finalizeNodeGroup(MemoryManager& mm, common::node_group_idx
     columnTypes.push_back(common::LogicalType::INTERNAL_ID());
     columnTypes.push_back(common::LogicalType::INTERNAL_ID());
     columnTypes.push_back(common::LogicalType::INTERNAL_ID());
-    auto numRelsInGroup = numRelsPerNodeGroup[nodeGroupIdx];
-    auto chunkedNodeGroup =
-        std::make_unique<ChunkedNodeGroup>(mm, columnTypes, false /* enableCompression */,
-            numRelsInGroup, 0 /* startRowIdx */, ResidencyState::IN_MEMORY);
+    auto chunkedNodeGroup = std::make_unique<ChunkedNodeGroup>(mm, columnTypes,
+        false /* enableCompression */, numRels, 0 /* startRowIdx */, ResidencyState::IN_MEMORY);
 
     auto currNumRels = 0u;
     auto& boundColumnChunk = chunkedNodeGroup->getColumnChunk(0).getData();
@@ -158,23 +148,22 @@ void InMemHNSWGraph::finalizeNodeGroup(MemoryManager& mm, common::node_group_idx
     partition.merge(std::move(chunkedNodeGroup));
 }
 
-common::offset_vec_t InMemHNSWGraph::getNeighbors(transaction::Transaction* transaction,
-    common::offset_t nodeOffset) const {
+common::offset_vec_t InMemHNSWGraph::getNeighbors(common::offset_t nodeOffset) const {
     const auto numNbrs = getCSRLength(nodeOffset);
-    return getNeighbors(transaction, nodeOffset, numNbrs);
+    return getNeighbors(nodeOffset, numNbrs);
 }
 
-common::offset_vec_t InMemHNSWGraph::getNeighbors(transaction::Transaction*,
-    common::offset_t nodeOffset, common::length_t numNbrs) const {
+common::offset_vec_t InMemHNSWGraph::getNeighbors(common::offset_t nodeOffset,
+    common::length_t numNbrs) const {
     common::offset_vec_t neighbors;
     neighbors.reserve(numNbrs);
     for (common::offset_t i = 0; i < numNbrs; i++) {
         auto nbr = getDstNode(nodeOffset * maxDegree + i);
         // Note: we might have INVALID_OFFSET at the end of the array of neighbor nodes. This is due
-        // to that when we append a new neighbor node to node x, we don't exclusively lock the x,
-        // instead, we increment the csrLength first, then set the dstNode. This design eases lock
-        // contentions. However, if this function (`getNeighbors`) is called before the dstNode is
-        // set, we will get INVALID_OFFSET. As csrLength is always synchorized, this design
+        // to that when we append a new neighbor node to node x, we don't exclusively lock the
+        // x, instead, we increment the csrLength first, then set the dstNode. This design eases
+        // lock contentions. However, if this function (`getNeighbors`) is called before the dstNode
+        // is set, we will get INVALID_OFFSET. As csrLength is always synchorized, this design
         // shouldn't have correctness issue.
         if (nbr == common::INVALID_OFFSET) {
             continue;