Add HNSW parallel loading support with thread-safe lifecycle management

extension/CMakeLists.txt

@@ -8,7 +8,7 @@ include(extension_config.cmake)
 set(STATICALLY_LINKED_EXTENSIONS "${STATICALLY_LINKED_EXTENSIONS}" PARENT_SCOPE)
 
 function(set_extension_properties target_name output_name extension_name)
-    set_target_properties(${target_name}
+set_target_properties(${target_name}
             PROPERTIES
             ARCHIVE_OUTPUT_DIRECTORY "${PROJECT_SOURCE_DIR}/extension/${extension_name}/build"
             LIBRARY_OUTPUT_DIRECTORY "${PROJECT_SOURCE_DIR}/extension/${extension_name}/build"
@@ -20,7 +20,9 @@ function(set_extension_properties target_name output_name extension_name)
 endfunction()
 
 function(set_apple_dynamic_lookup target_name)
-    set_target_properties(${target_name} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
+    if (APPLE)
+        set_target_properties(${target_name} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
+    endif ()
 endfunction()
 
 function(build_extension_lib build_static ext_name)
@@ -42,9 +44,10 @@ function(build_extension_lib build_static ext_name)
     if (WIN32 OR build_static)
         # See comments in extension/httpfs/CMakeLists.txt.
         target_link_libraries(kuzu_${EXTENSION_LIB_NAME}_extension PRIVATE kuzu)
-    endif ()
-    if (APPLE AND NOT build_static)
-        set_apple_dynamic_lookup(kuzu_${EXTENSION_LIB_NAME}_extension)
+    elseif (APPLE AND NOT build_static)
+        # Two-level namespace on macOS prevents dlopen'ed libraries from seeing executable symbols.
+        # Link against the shared Kuzu library so extensions can resolve their dependencies.
+        target_link_libraries(kuzu_${EXTENSION_LIB_NAME}_extension PRIVATE kuzu_shared)
     endif ()
 endfunction()
 

extension/fts/src/include/index/fts_index.h

@@ -31,8 +31,8 @@ class FTSIndex final : public storage::Index {
         FTSConfig ftsConfig, main::ClientContext* context);
 
     static std::unique_ptr<Index> load(main::ClientContext* context,
-        storage::StorageManager* storageManager, storage::IndexInfo indexInfo,
-        std::span<uint8_t> storageInfoBuffer);
+        storage::StorageManager* storageManager, const catalog::IndexCatalogEntry* catalogEntry,
+        storage::IndexInfo indexInfo, std::span<uint8_t> storageInfoBuffer);
 
     std::unique_ptr<InsertState> initInsertState(main::ClientContext*,
         storage::visible_func isVisible) override;

extension/fts/src/index/fts_index.cpp

@@ -3,6 +3,7 @@
 #include "catalog/catalog.h"
 #include "catalog/fts_index_catalog_entry.h"
 #include "index/fts_update_state.h"
+#include "extension/extension.h"
 #include "re2.h"
 #include "utils/fts_utils.h"
 
@@ -20,14 +21,13 @@ FTSIndex::FTSIndex(IndexInfo indexInfo, std::unique_ptr<IndexStorageInfo> storag
       config{std::move(config)} {}
 
 std::unique_ptr<Index> FTSIndex::load(main::ClientContext* context, StorageManager*,
-    IndexInfo indexInfo, std::span<uint8_t> storageInfoBuffer) {
-    auto catalog = catalog::Catalog::Get(*context);
+    const catalog::IndexCatalogEntry* catalogEntry, IndexInfo indexInfo,
+    std::span<uint8_t> storageInfoBuffer) {
     auto reader =
         std::make_unique<BufferReader>(storageInfoBuffer.data(), storageInfoBuffer.size());
     auto storageInfo = FTSStorageInfo::deserialize(std::move(reader));
-    auto indexEntry = catalog->getIndex(transaction::Transaction::Get(*context), indexInfo.tableID,
-        indexInfo.name);
-    auto ftsConfig = indexEntry->getAuxInfo().cast<FTSIndexAuxInfo>().config;
+    KU_ASSERT(catalogEntry != nullptr);
+    auto ftsConfig = catalogEntry->getAuxInfo().cast<FTSIndexAuxInfo>().config;
     return std::make_unique<FTSIndex>(std::move(indexInfo), std::move(storageInfo),
         std::move(ftsConfig), context);
 }
@@ -201,8 +201,8 @@ void FTSIndex::delete_(Transaction* transaction, const ValueVector& nodeIDVector
 
 void FTSIndex::finalize(main::ClientContext* context) {
     auto& ftsStorageInfo = storageInfo->cast<FTSStorageInfo>();
-    const auto numTotalRows =
-        internalTableInfo.table->getNumTotalRows(&DUMMY_CHECKPOINT_TRANSACTION);
+    const auto numTotalRows = internalTableInfo.table->getNumTotalRows(
+        extension::getExtensionCheckpointTransaction());
     if (numTotalRows == ftsStorageInfo.numCheckpointedNodes) {
         return;
     }
@@ -228,17 +228,17 @@ void FTSIndex::checkpoint(main::ClientContext* context, storage::PageAllocator&
     KU_ASSERT(!context->isInMemory());
     auto catalog = catalog::Catalog::Get(*context);
     internalTableInfo.docTable->checkpoint(context,
-        catalog->getTableCatalogEntry(&DUMMY_CHECKPOINT_TRANSACTION,
+        catalog->getTableCatalogEntry(extension::getExtensionCheckpointTransaction(),
             internalTableInfo.docTable->getTableID()),
         pageAllocator);
     internalTableInfo.termsTable->checkpoint(context,
-        catalog->getTableCatalogEntry(&DUMMY_CHECKPOINT_TRANSACTION,
+        catalog->getTableCatalogEntry(extension::getExtensionCheckpointTransaction(),
             internalTableInfo.termsTable->getTableID()),
         pageAllocator);
     auto appearsInTableName =
         FTSUtils::getAppearsInTableName(internalTableInfo.table->getTableID(), indexInfo.name);
-    auto appearsInTableEntry =
-        catalog->getTableCatalogEntry(&DUMMY_CHECKPOINT_TRANSACTION, appearsInTableName);
+    auto appearsInTableEntry = catalog->getTableCatalogEntry(
+        extension::getExtensionCheckpointTransaction(), appearsInTableName);
     internalTableInfo.appearsInfoTable->checkpoint(context, appearsInTableEntry, pageAllocator);
 }
 

extension/fts/src/main/fts_extension.cpp

@@ -28,7 +28,7 @@ static void initFTSEntries(main::ClientContext* context, catalog::Catalog& catal
             KU_ASSERT_UNCONDITIONAL(
                 optionalIndex.has_value() && !optionalIndex.value().get().isLoaded());
             auto& unloadedIndex = optionalIndex.value().get();
-            unloadedIndex.load(context, storageManager);
+            unloadedIndex.load(context, storageManager, indexEntry);
         }
     }
 }

extension/vector/src/include/index/hnsw_index.h

@@ -308,8 +308,8 @@ class OnDiskHNSWIndex final : public HNSWIndex {
         const EmbeddingHandle& queryVector, HNSWSearchState& searchState) const;
 
     static std::unique_ptr<Index> load(main::ClientContext* context,
-        storage::StorageManager* storageManager, storage::IndexInfo indexInfo,
-        std::span<uint8_t> storageInfoBuffer);
+        storage::StorageManager* storageManager, const catalog::IndexCatalogEntry* catalogEntry,
+        storage::IndexInfo indexInfo, std::span<uint8_t> storageInfoBuffer);
     std::unique_ptr<InsertState> initInsertState(main::ClientContext* context,
         storage::visible_func) override;
     bool needCommitInsert() const override { return true; }

extension/vector/src/index/hnsw_index.cpp

@@ -4,6 +4,7 @@
 #include "catalog/hnsw_index_catalog_entry.h"
 #include "function/hnsw_index_functions.h"
 #include "index/hnsw_rel_batch_insert.h"
+#include "extension/extension.h"
 #include "storage/storage_manager.h"
 #include "storage/table/node_table.h"
 #include "storage/table/rel_table.h"
@@ -468,16 +469,18 @@ OnDiskHNSWIndex::OnDiskHNSWIndex(const main::ClientContext* context, IndexInfo i
 }
 
 std::unique_ptr<Index> OnDiskHNSWIndex::load(main::ClientContext* context, StorageManager*,
-    IndexInfo indexInfo, std::span<uint8_t> storageInfoBuffer) {
+    const catalog::IndexCatalogEntry* catalogEntry, IndexInfo indexInfo,
+    std::span<uint8_t> storageInfoBuffer) {
     auto reader =
         std::make_unique<common::BufferReader>(storageInfoBuffer.data(), storageInfoBuffer.size());
     auto storageInfo = HNSWStorageInfo::deserialize(std::move(reader));
-    const auto catalog = catalog::Catalog::Get(*context);
-    const auto transaction = Transaction::Get(*context);
-    const auto indexEntry = catalog->getIndex(transaction, indexInfo.tableID, indexInfo.name);
-    const auto auxInfo = indexEntry->getAuxInfo().cast<HNSWIndexAuxInfo>();
-    return std::make_unique<OnDiskHNSWIndex>(context, std::move(indexInfo), std::move(storageInfo),
+
+    KU_ASSERT(catalogEntry != nullptr);
+    const auto auxInfo = catalogEntry->getAuxInfo().cast<HNSWIndexAuxInfo>();
+
+    auto result = std::make_unique<OnDiskHNSWIndex>(context, std::move(indexInfo), std::move(storageInfo),
         auxInfo.config.copy());
+    return result;
 }
 
 std::vector<NodeWithDistance> OnDiskHNSWIndex::search(Transaction* transaction,
@@ -605,7 +608,8 @@ void OnDiskHNSWIndex::commitInsert(Transaction* transaction,
 
 void OnDiskHNSWIndex::finalize(main::ClientContext* context) {
     auto& hnswStorageInfo = storageInfo->cast<HNSWStorageInfo>();
-    const auto numTotalRows = nodeTable.getNumTotalRows(&DUMMY_CHECKPOINT_TRANSACTION);
+    const auto numTotalRows =
+        nodeTable.getNumTotalRows(extension::getExtensionCheckpointTransaction());
     if (numTotalRows == hnswStorageInfo.numCheckpointedNodes) {
         return;
     }
@@ -637,7 +641,7 @@ void OnDiskHNSWIndex::finalize(main::ClientContext* context) {
 void OnDiskHNSWIndex::checkpoint(main::ClientContext* context,
     storage::PageAllocator& pageAllocator) {
     auto [nodeTableEntry, upperRelTableEntry, lowerRelTableEntry] = getIndexTableCatalogEntries(
-        catalog::Catalog::Get(*context), &DUMMY_CHECKPOINT_TRANSACTION, indexInfo);
+        catalog::Catalog::Get(*context), extension::getExtensionCheckpointTransaction(), indexInfo);
     upperRelTable->checkpoint(context, upperRelTableEntry, pageAllocator);
     lowerRelTable->checkpoint(context, lowerRelTableEntry, pageAllocator);
 }