storage: Optimize Vector Index (pingcap#232)

Signed-off-by: Wish <breezewish@outlook.com>
Co-authored-by: Lloyd-Pottiger <60744015+Lloyd-Pottiger@users.noreply.github.com>

.gitmodules

@@ -152,3 +152,6 @@
 [submodule "contrib/simsimd"]
 	path = contrib/simsimd
 	url = https://github.com/ashvardanian/SimSIMD
+[submodule "contrib/highfive"]
+	path = contrib/highfive
+	url = https://github.com/BlueBrain/HighFive

contrib/CMakeLists.txt

@@ -202,3 +202,7 @@ add_subdirectory(fastpforlib)
 add_subdirectory(usearch-cmake)
 
 add_subdirectory(simsimd-cmake)
+
+add_subdirectory(hdf5-cmake)
+
+add_subdirectory(highfive-cmake)

contrib/hdf5-cmake/.gitignore

@@ -0,0 +1 @@
+/download/*

contrib/hdf5-cmake/CMakeLists.txt

@@ -0,0 +1,41 @@
+include(ExternalProject)
+
+# hdf5 is too large. Instead of adding as a submodule, let's simply download from GitHub.
+ExternalProject_Add(hdf5-external
+    PREFIX          ${CMAKE_CURRENT_BINARY_DIR}
+    DOWNLOAD_DIR    ${TiFlash_SOURCE_DIR}/contrib/hdf5-cmake/download
+    URL             https://github.com/HDFGroup/hdf5/archive/refs/tags/hdf5_1.14.4.3.zip
+    URL_HASH        MD5=bc987d22e787290127aacd7b99b4f31e
+    CMAKE_ARGS
+        -DCMAKE_BUILD_TYPE=Release
+        -DCMAKE_INSTALL_PREFIX=<INSTALL_DIR>
+        -DBUILD_STATIC_LIBS=ON
+        -DBUILD_SHARED_LIBS=OFF
+        -DBUILD_TESTING=OFF
+        -DHDF5_BUILD_HL_LIB=OFF
+        -DHDF5_BUILD_TOOLS=OFF
+        -DHDF5_BUILD_CPP_LIB=ON
+        -DHDF5_BUILD_EXAMPLES=OFF
+        -DHDF5_ENABLE_Z_LIB_SUPPORT=OFF
+        -DHDF5_ENABLE_SZIP_SUPPORT=OFF
+    BUILD_BYPRODUCTS           <INSTALL_DIR>/lib/${CMAKE_FIND_LIBRARY_PREFIXES}hdf5.a  # Workaround for Ninja
+    USES_TERMINAL_DOWNLOAD     TRUE
+    USES_TERMINAL_CONFIGURE    TRUE
+    USES_TERMINAL_BUILD        TRUE
+    USES_TERMINAL_INSTALL      TRUE
+    EXCLUDE_FROM_ALL           TRUE
+    DOWNLOAD_EXTRACT_TIMESTAMP TRUE
+)
+
+ExternalProject_Get_Property(hdf5-external INSTALL_DIR)
+
+add_library(tiflash_contrib::hdf5 STATIC IMPORTED GLOBAL)
+set_target_properties(tiflash_contrib::hdf5 PROPERTIES
+    IMPORTED_LOCATION       ${INSTALL_DIR}/lib/${CMAKE_FIND_LIBRARY_PREFIXES}hdf5.a
+)
+add_dependencies(tiflash_contrib::hdf5 hdf5-external)
+
+file(MAKE_DIRECTORY ${INSTALL_DIR}/include)
+target_include_directories(tiflash_contrib::hdf5 SYSTEM INTERFACE
+    ${INSTALL_DIR}/include
+)

contrib/highfive-cmake/CMakeLists.txt

@@ -0,0 +1,18 @@
+set(HIGHFIVE_PROJECT_DIR "${TiFlash_SOURCE_DIR}/contrib/highfive")
+set(HIGHFIVE_SOURCE_DIR "${HIGHFIVE_PROJECT_DIR}/include")
+
+if (NOT EXISTS "${HIGHFIVE_SOURCE_DIR}/highfive/highfive.hpp")
+    message (FATAL_ERROR "submodule contrib/highfive not found")
+endif()
+
+add_library(_highfive INTERFACE)
+
+target_include_directories(_highfive SYSTEM INTERFACE
+    ${HIGHFIVE_SOURCE_DIR}
+)
+
+target_link_libraries(_highfive INTERFACE
+    tiflash_contrib::hdf5
+)
+
+add_library(tiflash_contrib::highfive ALIAS _highfive)

dbms/CMakeLists.txt

@@ -388,7 +388,18 @@ if (ENABLE_TESTS)
         )
     target_include_directories(bench_dbms BEFORE PRIVATE ${SPARCEHASH_INCLUDE_DIR} ${benchmark_SOURCE_DIR}/include)
     target_compile_definitions(bench_dbms PUBLIC DBMS_PUBLIC_GTEST)
-    target_link_libraries(bench_dbms gtest dbms test_util_bench_main benchmark tiflash_functions server_for_test delta_merge kvstore tiflash_aggregate_functions)
+    target_link_libraries(bench_dbms
+        gtest
+        benchmark
+        tiflash_contrib::highfive
+
+        dbms
+        test_util_bench_main
+        tiflash_functions
+        server_for_test
+        delta_merge
+        tiflash_aggregate_functions
+        kvstore)
 
     add_check(bench_dbms)
 endif ()

dbms/src/Storages/DeltaMerge/BitmapFilter/BitmapFilterView.h

@@ -37,11 +37,19 @@ class BitmapFilterView
         RUNTIME_CHECK(filter_offset + filter_size <= filter->size(), filter_offset, filter_size, filter->size());
     }
 
+    /**
+     * @brief Create a BitmapFilter and construct a BitmapFilterView with it.
+     * Should be only used in tests.
+     */
+    static BitmapFilterView createWithFilter(UInt32 size, bool default_value)
+    {
+        return BitmapFilterView(std::make_shared<BitmapFilter>(size, default_value), 0, size);
+    }
+
     // Caller should ensure n in [0, size).
     inline bool get(UInt32 n) const { return filter->get(filter_offset + n); }
 
     inline bool operator[](UInt32 n) const { return get(n); }
-
     inline UInt32 size() const { return filter_size; }
 
     inline UInt32 offset() const { return filter_offset; }

dbms/src/Storages/DeltaMerge/DeltaMergeStore_InternalSegment.cpp

@@ -39,6 +39,11 @@ extern const Metric DT_SnapshotOfSegmentIngest;
 extern const Metric DT_SnapshotOfSegmentIngestIndex;
 } // namespace CurrentMetrics
 
+namespace DB::ErrorCodes
+{
+extern const int ABORTED;
+}
+
 namespace DB::DM
 {
 
@@ -673,18 +678,20 @@ void DeltaMergeStore::segmentEnsureStableIndex(
     RUNTIME_CHECK(dm_files.size() == 1); // size > 1 is currently not supported.
     const auto & dm_file = dm_files[0];
 
-    // 2. Check whether the DMFile has been referenced by any valid segment.
-    {
+    auto is_file_valid = [this, dm_file] {
         std::shared_lock lock(read_write_mutex);
         auto segment_ids = dmfile_id_to_segment_ids.get(dm_file->fileId());
-        if (segment_ids.empty())
-        {
-            LOG_DEBUG(
-                log,
-                "EnsureStableIndex - Give up because no segment to update, source_segment={}",
-                source_segment_info);
-            return;
-        }
+        return !segment_ids.empty();
+    };
+
+    // 2. Check whether the DMFile has been referenced by any valid segment.
+    if (!is_file_valid())
+    {
+        LOG_DEBUG(
+            log,
+            "EnsureStableIndex - Give up because no segment to update, source_segment={}",
+            source_segment_info);
+        return;
     }
 
     LOG_INFO(
@@ -700,7 +707,29 @@ void DeltaMergeStore::segmentEnsureStableIndex(
         .dm_files = dm_files,
         .dm_context = dm_context,
     });
-    auto new_dmfiles = iw.build();
+
+    DMFiles new_dmfiles{};
+
+    try
+    {
+        // When file is not valid we need to abort the index build.
+        new_dmfiles = iw.build(is_file_valid);
+    }
+    catch (const Exception & e)
+    {
+        if (e.code() == ErrorCodes::ABORTED)
+        {
+            LOG_INFO(
+                log,
+                "EnsureStableIndex - Build index aborted because DMFile is no longer valid, dm_files={} "
+                "source_segment={}",
+                DMFile::info(dm_files),
+                source_segment_info);
+            return;
+        }
+        throw;
+    }
+
     RUNTIME_CHECK(!new_dmfiles.empty());
 
     LOG_INFO(

dbms/src/Storages/DeltaMerge/File/DMFileIndexWriter.cpp

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Common/Exception.h>
 #include <Interpreters/SharedContexts/Disagg.h>
 #include <Storages/DeltaMerge/DeltaMergeDefines.h>
 #include <Storages/DeltaMerge/File/DMFile.h>
@@ -22,6 +23,11 @@
 #include <Storages/DeltaMerge/ScanContext.h>
 #include <Storages/PathPool.h>
 
+namespace DB::ErrorCodes
+{
+extern const int ABORTED;
+}
+
 namespace DB::DM
 {
 
@@ -63,7 +69,7 @@ DMFileIndexWriter::LocalIndexBuildInfo DMFileIndexWriter::getLocalIndexBuildInfo
     return build;
 }
 
-size_t DMFileIndexWriter::buildIndexForFile(const DMFilePtr & dm_file_mutable) const
+size_t DMFileIndexWriter::buildIndexForFile(const DMFilePtr & dm_file_mutable, ProceedCheckFn should_proceed) const
 {
     const auto column_defines = dm_file_mutable->getColumnDefines();
     const auto del_cd_iter = std::find_if(column_defines.cbegin(), column_defines.cend(), [](const ColumnDefine & cd) {
@@ -128,6 +134,9 @@ size_t DMFileIndexWriter::buildIndexForFile(const DMFilePtr & dm_file_mutable) c
     // Read all blocks and build index
     while (true)
     {
+        if (!should_proceed())
+            throw Exception(ErrorCodes::ABORTED, "Index build is interrupted");
+
         auto block = read_stream->read();
         if (!block)
             break;
@@ -146,7 +155,7 @@ size_t DMFileIndexWriter::buildIndexForFile(const DMFilePtr & dm_file_mutable) c
             const auto & col_with_type_and_name = block.safeGetByPosition(col_idx + 1);
             RUNTIME_CHECK(col_with_type_and_name.column_id == read_columns[col_idx + 1].id);
             const auto & col = col_with_type_and_name.column;
-            index_builder->addBlock(*col, del_mark);
+            index_builder->addBlock(*col, del_mark, should_proceed);
         }
     }
 
@@ -187,7 +196,7 @@ size_t DMFileIndexWriter::buildIndexForFile(const DMFilePtr & dm_file_mutable) c
     return total_built_index_bytes;
 }
 
-DMFiles DMFileIndexWriter::build() const
+DMFiles DMFileIndexWriter::build(ProceedCheckFn should_proceed) const
 {
     RUNTIME_CHECK(!built);
     // Create a clone of existing DMFile instances by using DMFile::restore,
@@ -214,7 +223,7 @@ DMFiles DMFileIndexWriter::build() const
 
     for (const auto & cloned_dmfile : cloned_dm_files)
     {
-        auto index_bytes = buildIndexForFile(cloned_dmfile);
+        auto index_bytes = buildIndexForFile(cloned_dmfile, should_proceed);
         if (auto data_store = options.dm_context.global_context.getSharedContextDisagg()->remote_data_store;
             !data_store)
         {

dbms/src/Storages/DeltaMerge/File/DMFileIndexWriter.h

@@ -57,16 +57,23 @@ class DMFileIndexWriter
         const DMContext & dm_context;
     };
 
+    using ProceedCheckFn = std::function<bool()>;
+
     explicit DMFileIndexWriter(const Options & options)
         : logger(Logger::get())
         , options(options)
     {}
 
-    // Note: This method can only be called once.
-    DMFiles build() const;
+    // Note: You cannot call build() multiple times, as duplicate meta version will result in exceptions.
+    DMFiles build(ProceedCheckFn should_proceed) const;
+
+    DMFiles build() const
+    {
+        return build([]() { return true; });
+    }
 
 private:
-    size_t buildIndexForFile(const DMFilePtr & dm_file_mutable) const;
+    size_t buildIndexForFile(const DMFilePtr & dm_file_mutable, ProceedCheckFn should_proceed) const;
 
 private:
     const LoggerPtr logger;

dbms/src/Storages/DeltaMerge/Index/VectorIndex.h

@@ -35,6 +35,8 @@ class VectorIndexBuilder
     /// The key is the row's offset in the DMFile.
     using Key = UInt32;
 
+    using ProceedCheckFn = std::function<bool()>;
+
 public:
     static VectorIndexBuilderPtr create(const TiDB::VectorIndexDefinitionPtr & definition);
 
@@ -47,7 +49,11 @@ class VectorIndexBuilder
 
     virtual ~VectorIndexBuilder() = default;
 
-    virtual void addBlock(const IColumn & column, const ColumnVector<UInt8> * del_mark) = 0;
+    virtual void addBlock( //
+        const IColumn & column,
+        const ColumnVector<UInt8> * del_mark,
+        ProceedCheckFn should_proceed)
+        = 0;
 
     virtual void save(std::string_view path) const = 0;
 
@@ -80,6 +86,8 @@ class VectorIndexViewer
     // Invalid rows in `valid_rows` will be discared when applying the search
     virtual std::vector<Key> search(const ANNQueryInfoPtr & queryInfo, const RowFilter & valid_rows) const = 0;
 
+    virtual size_t size() const = 0;
+
     // Get the value (i.e. vector content) of a Key.
     virtual void get(Key key, std::vector<Float32> & out) const = 0;