You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@doris.apache.org by yi...@apache.org on 2025/04/01 02:51:00 UTC
(doris) branch vector_index updated: [feature](ANN) implemented the import process of vector index (#49703)
This is an automated email from the ASF dual-hosted git repository.
yiguolei pushed a commit to branch vector_index
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/vector_index by this push:
new f83f50380b7 [feature](ANN) implemented the import process of vector index (#49703)
f83f50380b7 is described below
commit f83f50380b7aac7e3790b11b1a4d23ca99af69e6
Author: chenlinzhong <49...@qq.com>
AuthorDate: Tue Apr 1 10:50:52 2025 +0800
[feature](ANN) implemented the import process of vector index (#49703)
### What problem does this PR solve?
Issue Number: close #xxx
Related PR: #xxx
Problem Summary:
```
CREATE TABLE `vector_table` (
`siteid` int(11) NULL DEFAULT "10" COMMENT "",
`embedding` array<float> NOT NULL COMMENT "",
`comment` text NULL,
INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES(
"index_type"="diskann",
"metric_type"="l2",
"dim"="8",
"search_list"="32",
"max_degree"="100"
) COMMENT 'test diskann index',
INDEX idx_comment (`comment`) USING INVERTED PROPERTIES("support_phrase" = "true", "parser" = "english", "lower_case" = "true") COMMENT 'inverted index for comment'
) ENGINE=OLAP
duplicate KEY(`siteid`) COMMENT "OLAP"
DISTRIBUTED BY HASH(`siteid`) BUCKETS 1
PROPERTIES (
"replication_num" = "1",
"storage_vault_name" = "s3_vault"
);
INSERT INTO `vector_table` (`siteid`, `embedding`,`comment`) VALUES
(10, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,20],"emb1"),
(20, [7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0,30],"emb2");
MySQL [clz]> select * from vector_table;
+--------+---------------------------+---------+
| siteid | embedding | comment |
+--------+---------------------------+---------+
| 10 | [1, 2, 3, 4, 5, 6, 7, 20] | emb1 |
| 20 | [7, 6, 5, 4, 3, 2, 1, 30] | emb2 |
+--------+---------------------------+---------+
编译镜像：https://clz-repo.cdn.bcebos.com/doris_ann.image.tar
编译命令：
export JAVA_HOME=/usr/lib/jvm/jdk-17.0.2/;
DISABLE_BE_JAVA_EXTENSIONS=ON BUILD_TYPE=Debug DORIS_TOOLCHAIN=gcc sh build.sh --be --fe
```
### Release note
None
### Check List (For Author)
- Test <!-- At least one of them must be included. -->
- [ ] Regression test
- [ ] Unit Test
- [ ] Manual test (add detailed scripts or steps below)
- [ ] No need to test or manual test. Explain why:
- [ ] This is a refactor/code format and no logic has been changed.
- [ ] Previous test can cover this change.
- [ ] No code files have been changed.
- [ ] Other reason <!-- Add your reason? -->
- Behavior changed:
- [ ] No.
- [ ] Yes. <!-- Explain the behavior change -->
- Does this need documentation?
- [ ] No.
- [ ] Yes. <!-- Add document PR link here. eg:
https://github.com/apache/doris-website/pull/1214 -->
### Check List (For Reviewer who merge this PR)
- [ ] Confirm the release note
- [ ] Confirm test cases
- [ ] Confirm document
- [ ] Add branch pick label <!-- Add branch pick label that this PR
should merge into -->
---
be/CMakeLists.txt | 40 +-
be/src/cloud/CMakeLists.txt | 3 +
be/src/common/status.h | 2 +-
be/src/extern/diskann/CMakeLists.txt | 173 +
be/src/extern/diskann/include/ThreadPool.h | 98 +
.../extern/diskann/include/abstract_data_store.h | 127 +
.../extern/diskann/include/abstract_graph_store.h | 75 +
be/src/extern/diskann/include/abstract_index.h | 129 +
be/src/extern/diskann/include/abstract_scratch.h | 35 +
.../extern/diskann/include/aligned_file_reader.h | 70 +
be/src/extern/diskann/include/ann_exception.h | 36 +
be/src/extern/diskann/include/any_wrappers.h | 53 +
.../diskann/include/boost_dynamic_bitset_fwd.h | 11 +
be/src/extern/diskann/include/cached_io.h | 217 ++
be/src/extern/diskann/include/combined_file.h | 105 +
be/src/extern/diskann/include/common_includes.h | 27 +
be/src/extern/diskann/include/concurrent_queue.h | 132 +
be/src/extern/diskann/include/cosine_similarity.h | 283 ++
be/src/extern/diskann/include/defaults.h | 34 +
be/src/extern/diskann/include/disk_utils.h | 135 +
be/src/extern/diskann/include/distance.h | 235 ++
be/src/extern/diskann/include/exceptions.h | 17 +
be/src/extern/diskann/include/filter_utils.h | 221 ++
be/src/extern/diskann/include/in_mem_data_store.h | 89 +
be/src/extern/diskann/include/in_mem_graph_store.h | 56 +
be/src/extern/diskann/include/index.h | 442 +++
be/src/extern/diskann/include/index_build_params.h | 73 +
be/src/extern/diskann/include/index_config.h | 256 ++
be/src/extern/diskann/include/index_factory.h | 51 +
be/src/extern/diskann/include/io_reader.h | 0
.../diskann/include/linux_aligned_file_reader.h | 39 +
be/src/extern/diskann/include/locking.h | 20 +
be/src/extern/diskann/include/logger.h | 35 +
be/src/extern/diskann/include/logger_impl.h | 61 +
be/src/extern/diskann/include/math_utils.h | 87 +
be/src/extern/diskann/include/memory_mapper.h | 43 +
be/src/extern/diskann/include/natural_number_map.h | 86 +
be/src/extern/diskann/include/natural_number_set.h | 50 +
be/src/extern/diskann/include/neighbor.h | 160 +
be/src/extern/diskann/include/parameters.h | 119 +
be/src/extern/diskann/include/partition.h | 53 +
be/src/extern/diskann/include/percentile_stats.h | 82 +
be/src/extern/diskann/include/pq.h | 117 +
be/src/extern/diskann/include/pq_common.h | 30 +
be/src/extern/diskann/include/pq_data_store.h | 97 +
be/src/extern/diskann/include/pq_flash_index.h | 286 ++
be/src/extern/diskann/include/pq_l2_distance.h | 87 +
be/src/extern/diskann/include/pq_scratch.h | 23 +
.../diskann/include/program_options_utils.hpp | 81 +
be/src/extern/diskann/include/quantized_distance.h | 56 +
be/src/extern/diskann/include/restapi/common.h | 18 +
.../diskann/include/restapi/search_wrapper.h | 140 +
be/src/extern/diskann/include/restapi/server.h | 45 +
be/src/extern/diskann/include/scratch.h | 216 ++
be/src/extern/diskann/include/simd_utils.h | 106 +
be/src/extern/diskann/include/tag_uint128.h | 68 +
be/src/extern/diskann/include/timer.h | 40 +
be/src/extern/diskann/include/tsl/.clang-format | 2 +
.../diskann/include/tsl/robin_growth_policy.h | 330 ++
be/src/extern/diskann/include/tsl/robin_hash.h | 1285 +++++++
be/src/extern/diskann/include/tsl/robin_map.h | 668 ++++
be/src/extern/diskann/include/tsl/robin_set.h | 535 +++
.../diskann/include/tsl/sparse_growth_policy.h | 301 ++
be/src/extern/diskann/include/tsl/sparse_hash.h | 2215 ++++++++++++
be/src/extern/diskann/include/tsl/sparse_map.h | 800 +++++
be/src/extern/diskann/include/tsl/sparse_set.h | 655 ++++
be/src/extern/diskann/include/types.h | 22 +
be/src/extern/diskann/include/utils.h | 1370 ++++++++
.../diskann/include/windows_aligned_file_reader.h | 57 +
.../diskann/include/windows_customizations.h | 16 +
be/src/extern/diskann/include/windows_slim_lock.h | 73 +
be/src/extern/diskann/src/abstract_data_store.cpp | 45 +
be/src/extern/diskann/src/abstract_index.cpp | 334 ++
be/src/extern/diskann/src/ann_exception.cpp | 36 +
be/src/extern/diskann/src/combined_file.cpp | 218 ++
be/src/extern/diskann/src/disk_utils.cpp | 1841 ++++++++++
be/src/extern/diskann/src/distance.cpp | 733 ++++
be/src/extern/diskann/src/filter_utils.cpp | 355 ++
be/src/extern/diskann/src/in_mem_data_store.cpp | 401 +++
be/src/extern/diskann/src/in_mem_graph_store.cpp | 284 ++
be/src/extern/diskann/src/index.cpp | 3627 ++++++++++++++++++++
be/src/extern/diskann/src/index_factory.cpp | 213 ++
.../diskann/src/linux_aligned_file_reader.cpp | 228 ++
be/src/extern/diskann/src/logger.cpp | 97 +
be/src/extern/diskann/src/math_utils.cpp | 458 +++
be/src/extern/diskann/src/memory_mapper.cpp | 107 +
be/src/extern/diskann/src/natural_number_map.cpp | 116 +
be/src/extern/diskann/src/natural_number_set.cpp | 70 +
be/src/extern/diskann/src/partition.cpp | 706 ++++
be/src/extern/diskann/src/pq.cpp | 1821 ++++++++++
be/src/extern/diskann/src/pq_data_store.cpp | 260 ++
be/src/extern/diskann/src/pq_flash_index.cpp | 1851 ++++++++++
be/src/extern/diskann/src/pq_l2_distance.cpp | 284 ++
be/src/extern/diskann/src/scratch.cpp | 181 +
be/src/extern/diskann/src/utils.cpp | 477 +++
.../diskann/src/windows_aligned_file_reader.cpp | 189 +
be/src/index-tools/ann_tool.cpp | 361 ++
be/src/index-tools/index_tool.cpp | 10 +-
be/src/io/CMakeLists.txt | 5 +
be/src/olap/CMakeLists.txt | 3 +
be/src/olap/compaction.cpp | 28 +-
be/src/olap/rowset/beta_rowset_writer.cpp | 28 +-
be/src/olap/rowset/beta_rowset_writer.h | 20 +-
be/src/olap/rowset/rowset_writer.h | 8 +-
be/src/olap/rowset/segcompaction.cpp | 4 +-
be/src/olap/rowset/segcompaction.h | 8 +-
be/src/olap/rowset/segment_creator.cpp | 24 +-
be/src/olap/rowset/segment_creator.h | 8 +-
be/src/olap/rowset/segment_v2/ann_index_writer.cpp | 115 +
be/src/olap/rowset/segment_v2/ann_index_writer.h | 111 +
be/src/olap/rowset/segment_v2/column_writer.cpp | 41 +-
be/src/olap/rowset/segment_v2/column_writer.h | 17 +-
be/src/olap/rowset/segment_v2/index_writer.cpp | 169 +
.../{inverted_index_writer.h => index_writer.h} | 20 +-
.../segment_v2/inverted_index_compaction.cpp | 2 +-
.../rowset/segment_v2/inverted_index_compaction.h | 2 +-
.../rowset/segment_v2/inverted_index_file_reader.h | 2 +-
.../rowset/segment_v2/inverted_index_writer.cpp | 297 +-
.../olap/rowset/segment_v2/inverted_index_writer.h | 181 +-
be/src/olap/rowset/segment_v2/segment_writer.cpp | 31 +-
be/src/olap/rowset/segment_v2/segment_writer.h | 13 +-
.../rowset/segment_v2/vertical_segment_writer.cpp | 29 +-
.../rowset/segment_v2/vertical_segment_writer.h | 15 +-
...dex_file_writer.cpp => x_index_file_writer.cpp} | 66 +-
...d_index_file_writer.h => x_index_file_writer.h} | 10 +-
be/src/olap/rowset/vertical_beta_rowset_writer.cpp | 12 +-
be/src/olap/schema_change.cpp | 2 +-
be/src/olap/schema_change.h | 2 +-
be/src/olap/tablet_meta.cpp | 3 +
be/src/olap/tablet_schema.cpp | 36 +-
be/src/olap/tablet_schema.h | 24 +
be/src/olap/task/index_builder.cpp | 46 +-
be/src/olap/task/index_builder.h | 14 +-
be/src/pipeline/CMakeLists.txt | 3 +
be/src/runtime/CMakeLists.txt | 3 +
be/src/runtime/exec_env.h | 2 +-
be/src/runtime/exec_env_init.cpp | 3 +
be/src/service/CMakeLists.txt | 5 +
be/src/{runtime => vector}/CMakeLists.txt | 26 +-
be/src/vector/ann_parse.h | 143 +
be/src/vector/diskann_vector_index.cpp | 206 ++
be/src/vector/diskann_vector_index.h | 274 ++
be/src/vector/stream_wrapper.h | 142 +
be/src/vector/vector_index.h | 94 +
.../antlr4/org/apache/doris/nereids/DorisLexer.g4 | 1 +
.../antlr4/org/apache/doris/nereids/DorisParser.g4 | 5 +-
fe/fe-core/src/main/cup/sql_parser.cup | 7 +
.../java/org/apache/doris/analysis/IndexDef.java | 9 +-
.../main/java/org/apache/doris/catalog/Index.java | 4 +
.../trees/plans/commands/info/IndexDefinition.java | 10 +-
fe/fe-core/src/main/jflex/sql_scanner.flex | 1 +
gensrc/proto/olap_file.proto | 1 +
gensrc/thrift/Descriptors.thrift | 3 +-
153 files changed, 30949 insertions(+), 491 deletions(-)
diff --git a/be/CMakeLists.txt b/be/CMakeLists.txt
index d476af8e211..1cd1ac08ee4 100644
--- a/be/CMakeLists.txt
+++ b/be/CMakeLists.txt
@@ -421,7 +421,7 @@ else()
endif()
# Add flags that are common across build types
-set(CMAKE_CXX_FLAGS "${CXX_COMMON_FLAGS} ${CMAKE_CXX_FLAGS} ${EXTRA_CXX_FLAGS}")
+set(CMAKE_CXX_FLAGS "${CXX_COMMON_FLAGS} ${CMAKE_CXX_FLAGS} ${EXTRA_CXX_FLAGS} -fopenmp -fopenmp-simd ")
set(CMAKE_C_FLAGS ${CMAKE_CXX_FLAGS})
@@ -511,6 +511,8 @@ set(DORIS_LINK_LIBS
Cloud
${WL_END_GROUP}
CommonCPP
+ diskann_s
+ vector
)
set(absl_DIR ${THIRDPARTY_DIR}/lib/cmake/absl)
@@ -750,20 +752,28 @@ function(pch_reuse target)
endif()
endfunction(pch_reuse target)
-add_subdirectory(${SRC_DIR}/agent)
-add_subdirectory(${SRC_DIR}/common)
-add_subdirectory(${SRC_DIR}/exec)
-add_subdirectory(${SRC_DIR}/exprs)
-add_subdirectory(${SRC_DIR}/gen_cpp)
-add_subdirectory(${SRC_DIR}/geo)
-add_subdirectory(${SRC_DIR}/gutil)
-add_subdirectory(${SRC_DIR}/http)
-add_subdirectory(${SRC_DIR}/io)
-add_subdirectory(${SRC_DIR}/olap)
-add_subdirectory(${SRC_DIR}/runtime)
-add_subdirectory(${SRC_DIR}/service) # this include doris_be
-add_subdirectory(${SRC_DIR}/udf)
-add_subdirectory(${SRC_DIR}/cloud)
+macro(add_subdirectory_with_log subdir)
+ message(STATUS "Start compiling ${subdir} at ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}")
+ add_subdirectory(${subdir})
+endmacro()
+
+add_subdirectory_with_log(${SRC_DIR}/agent)
+add_subdirectory_with_log(${SRC_DIR}/common)
+add_subdirectory_with_log(${SRC_DIR}/exec)
+add_subdirectory_with_log(${SRC_DIR}/exprs)
+add_subdirectory_with_log(${SRC_DIR}/gen_cpp)
+add_subdirectory_with_log(${SRC_DIR}/geo)
+add_subdirectory_with_log(${SRC_DIR}/gutil)
+add_subdirectory_with_log(${SRC_DIR}/http)
+add_subdirectory_with_log(${SRC_DIR}/io)
+add_subdirectory_with_log(${SRC_DIR}/olap)
+add_subdirectory_with_log(${SRC_DIR}/runtime)
+add_subdirectory_with_log(${SRC_DIR}/service) # this include doris_be
+add_subdirectory_with_log(${SRC_DIR}/udf)
+add_subdirectory_with_log(${SRC_DIR}/cloud)
+add_subdirectory_with_log(${SRC_DIR}/extern/diskann)
+add_subdirectory_with_log(${SRC_DIR}/vector)
+
option(BUILD_META_TOOL "Build meta tool" OFF)
if (BUILD_META_TOOL)
diff --git a/be/src/cloud/CMakeLists.txt b/be/src/cloud/CMakeLists.txt
index dbe8160c2f8..fd259cbb3e4 100644
--- a/be/src/cloud/CMakeLists.txt
+++ b/be/src/cloud/CMakeLists.txt
@@ -15,6 +15,9 @@
# specific language governing permissions and limitations
# under the License.
+set(CMAKE_COMPILE_WARNING_AS_ERROR OFF)
+add_compile_options(-Wno-error=attributes)
+
# where to put generated libraries
set(LIBRARY_OUTPUT_PATH "${BUILD_DIR}/src/cloud")
diff --git a/be/src/common/status.h b/be/src/common/status.h
index 0252ec8564f..c7fd15b7b86 100644
--- a/be/src/common/status.h
+++ b/be/src/common/status.h
@@ -640,7 +640,7 @@ inline std::string Status::to_string_no_stack() const {
// some generally useful macros
#define RETURN_IF_ERROR(stmt) \
do { \
- Status _status_ = (stmt); \
+ doris::Status _status_ = (stmt); \
if (UNLIKELY(!_status_.ok())) { \
return _status_; \
} \
diff --git a/be/src/extern/diskann/CMakeLists.txt b/be/src/extern/diskann/CMakeLists.txt
new file mode 100644
index 00000000000..f2e2415841d
--- /dev/null
+++ b/be/src/extern/diskann/CMakeLists.txt
@@ -0,0 +1,173 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+# Parameters:
+#
+# BOOST_ROOT:
+# Specify root of the Boost library if Boost cannot be auto-detected.
+#
+# DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS:
+# This is a work-in-progress feature, not completed yet. The core DiskANN library will be split into
+# build-related and search-related functionality. In build-related functionality, when using tcmalloc,
+# it's possible to release memory that's free but reserved by tcmalloc. Setting this to true enables
+# such behavior.
+# Contact for this feature: gopalrs.
+cmake_minimum_required(VERSION 3.15)
+set(LIBRARY_OUTPUT_PATH "${BUILD_DIR}/src/extern")
+project(diskann)
+set(CMAKE_STANDARD 17)
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+set(CMAKE_COMPILE_WARNING_AS_ERROR OFF)
+add_compile_options(-Wno-error=attributes)
+add_compile_options(-Wno-deprecated-copy)
+add_compile_options(-Wno-reorder)
+add_compile_options(-Wno-unused-but-set-variable)
+add_compile_options(-Wno-error=unused-variable)
+
+
+message(STATUS "CMAKE_C_FLAGS in be: ${CMAKE_C_FLAGS}")
+message(STATUS "CMAKE_CXX_FLAGS in be: ${CMAKE_CXX_FLAGS}")
+message(STATUS "CMAKE_PREFIX_PATH in be: ${CMAKE_PREFIX_PATH}")
+message(STATUS "LD_LIBRARY_PATH in be: $ENV{LD_LIBRARY_PATH}")
+message(STATUS "CMAKE_MODULE_PATH in be: $ENV")
+
+
+set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake;${CMAKE_MODULE_PATH}")
+
+
+
+include_directories(${PROJECT_SOURCE_DIR}/include)
+include_directories(/home/users/clz/baidu/third-party/doris-diskann/build)
+
+if(NOT PYBIND)
+ set(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS ON)
+endif()
+# It's necessary to include tcmalloc headers only if calling into MallocExtension interface.
+# For using tcmalloc in DiskANN tools, it's enough to just link with tcmalloc.
+if (DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS)
+ include_directories(${PROJECT_SOURCE_DIR}/gperftools/src)
+endif()
+
+#OpenMP
+find_package(OpenMP)
+message(" OpenMP1 ${OPENMP_FOUND} ${CMAKE_CXX_COMPILER}")
+if (OPENMP_FOUND)
+ set (CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
+ set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
+else()
+ message(FATAL_ERROR "No OpenMP support")
+endif()
+
+# DiskANN core uses header-only libraries. Only DiskANN tools need program_options which has a linker library,
+# but its size is small. Reduce number of dependent DLLs by linking statically.
+
+# find_package(Boost COMPONENTS program_options)
+# message("${Boost_PROGRAM_OPTIONS_INCLUDE_DIRS},${Boost_PROGRAM_OPTIONS_LIBRARIES}")
+# if (NOT Boost_FOUND)
+# message(FATAL_ERROR "Couldn't find Boost dependency")
+# endif()
+include_directories(/home/users/clz/baidu/third-party/palo/thirdparty/installed/include/)
+link_directories(/home/users/clz/baidu/third-party/palo/thirdparty/installed/lib64)
+
+#MKL Config
+# expected path for manual intel mkl installs
+set(POSSIBLE_OMP_PATHS "/opt/intel/oneapi/compiler/latest/linux/compiler/lib/intel64_lin/libiomp5.so;/usr/lib/x86_64-linux-gnu/libiomp5.so;/opt/intel/lib/intel64_lin/libiomp5.so")
+foreach(POSSIBLE_OMP_PATH ${POSSIBLE_OMP_PATHS})
+ if (EXISTS ${POSSIBLE_OMP_PATH})
+ get_filename_component(OMP_PATH ${POSSIBLE_OMP_PATH} DIRECTORY)
+ endif()
+endforeach()
+
+if(NOT OMP_PATH)
+ message(FATAL_ERROR "Could not find Intel OMP in standard locations; use -DOMP_PATH to specify the install location for your environment")
+endif()
+link_directories(${OMP_PATH})
+
+
+set(POSSIBLE_MKL_LIB_PATHS "/opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_core.so;/usr/lib/x86_64-linux-gnu/libmkl_core.so;/opt/intel/mkl/lib/intel64/libmkl_core.so")
+foreach(POSSIBLE_MKL_LIB_PATH ${POSSIBLE_MKL_LIB_PATHS})
+ if (EXISTS ${POSSIBLE_MKL_LIB_PATH})
+ get_filename_component(MKL_PATH ${POSSIBLE_MKL_LIB_PATH} DIRECTORY)
+ endif()
+endforeach()
+
+set(POSSIBLE_MKL_INCLUDE_PATHS "/opt/intel/oneapi/mkl/latest/include;/usr/include/mkl;/opt/intel/mkl/include/;")
+foreach(POSSIBLE_MKL_INCLUDE_PATH ${POSSIBLE_MKL_INCLUDE_PATHS})
+ if (EXISTS ${POSSIBLE_MKL_INCLUDE_PATH})
+ set(MKL_INCLUDE_PATH ${POSSIBLE_MKL_INCLUDE_PATH})
+ endif()
+endforeach()
+if(NOT MKL_PATH)
+ message(FATAL_ERROR "Could not find Intel MKL in standard locations; use -DMKL_PATH to specify the install location for your environment")
+elseif(NOT MKL_INCLUDE_PATH)
+ message(FATAL_ERROR "Could not find Intel MKL in standard locations; use -DMKL_INCLUDE_PATH to specify the install location for headers for your environment")
+endif()
+if (EXISTS ${MKL_PATH}/libmkl_def.so.2)
+ set(MKL_DEF_SO ${MKL_PATH}/libmkl_def.so.2)
+elseif(EXISTS ${MKL_PATH}/libmkl_def.so)
+ set(MKL_DEF_SO ${MKL_PATH}/libmkl_def.so)
+else()
+ message(FATAL_ERROR "Despite finding MKL, libmkl_def.so was not found in expected locations.")
+endif()
+link_directories(${MKL_PATH})
+message("mkl ${MKL_PATH}")
+include_directories(${MKL_INCLUDE_PATH})
+
+# compile flags and link libraries
+add_compile_options(-m64 -Wl,--no-as-needed)
+# if (NOT PYBIND)
+# link_libraries(mkl_intel_ilp64 mkl_intel_thread mkl_core iomp5 pthread m dl)
+# else()
+ # static linking for python so as to minimize customer dependency issues
+ link_libraries(
+ ${MKL_PATH}/libmkl_intel_ilp64.a
+ ${MKL_PATH}/libmkl_intel_thread.a
+ ${MKL_PATH}/libmkl_core.a
+ ${MKL_DEF_SO}
+ /opt/intel/lib/intel64_lin/libiomp5.a
+ pthread
+ m
+ dl
+ )
+# endif()
+add_definitions(-DMKL_ILP64)
+
+# Section for tcmalloc. The DiskANN tools are always linked to tcmalloc.
+if(NOT PYBIND)
+ set(DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS "-ltcmalloc")
+endif()
+
+if (DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS)
+ add_definitions(-DRELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS)
+endif()
+
+set(DISKANN_ASYNC_LIB aio)
+
+
+#Main compiler/linker settings
+set(ENV{TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD} 500000000000)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mavx2 -mfma -msse2 -ftree-vectorize -fno-builtin-malloc -fno-builtin-calloc -fno-builtin-realloc -fno-builtin-free -fopenmp -fopenmp-simd -funroll-loops -Wfatal-errors -DUSE_AVX2")
+set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DNDEBUG")
+if (NOT PYBIND)
+ set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -DNDEBUG -Ofast")
+ if (NOT PORTABLE)
+ set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -march=native -mtune=native")
+ endif()
+else()
+ # -Ofast is not supported in a python extension module
+ set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -DNDEBUG -fPIC")
+endif()
+
+
+
+
+set(CMAKE_COMPILE_WARNING_AS_ERROR OFF)
+
+file(GLOB CPP_SOURCES src/*.cpp)
+add_library(${PROJECT_NAME} ${CPP_SOURCES})
+add_library(${PROJECT_NAME}_s STATIC ${CPP_SOURCES})
+
+
+
diff --git a/be/src/extern/diskann/include/ThreadPool.h b/be/src/extern/diskann/include/ThreadPool.h
new file mode 100644
index 00000000000..3c051dc3fd3
--- /dev/null
+++ b/be/src/extern/diskann/include/ThreadPool.h
@@ -0,0 +1,98 @@
+#ifndef THREAD_POOL_H
+#define THREAD_POOL_H
+
+#include <vector>
+#include <queue>
+#include <memory>
+#include <thread>
+#include <mutex>
+#include <condition_variable>
+#include <future>
+#include <functional>
+#include <stdexcept>
+
+class ThreadPool {
+public:
+ ThreadPool(size_t);
+ template<class F, class... Args>
+ auto enqueue(F&& f, Args&&... args)
+ -> std::future<typename std::result_of<F(Args...)>::type>;
+ ~ThreadPool();
+private:
+ // need to keep track of threads so we can join them
+ std::vector< std::thread > workers;
+ // the task queue
+ std::queue< std::function<void()> > tasks;
+
+ // synchronization
+ std::mutex queue_mutex;
+ std::condition_variable condition;
+ bool stop;
+};
+
+// the constructor just launches some amount of workers
+inline ThreadPool::ThreadPool(size_t threads)
+ : stop(false)
+{
+ for(size_t i = 0;i<threads;++i)
+ workers.emplace_back(
+ [this]
+ {
+ for(;;)
+ {
+ std::function<void()> task;
+
+ {
+ std::unique_lock<std::mutex> lock(this->queue_mutex);
+ this->condition.wait(lock,
+ [this]{ return this->stop || !this->tasks.empty(); });
+ if(this->stop && this->tasks.empty())
+ return;
+ task = std::move(this->tasks.front());
+ this->tasks.pop();
+ }
+
+ task();
+ }
+ }
+ );
+}
+
+// add new work item to the pool
+template<class F, class... Args>
+auto ThreadPool::enqueue(F&& f, Args&&... args)
+ -> std::future<typename std::result_of<F(Args...)>::type>
+{
+ using return_type = typename std::result_of<F(Args...)>::type;
+
+ auto task = std::make_shared< std::packaged_task<return_type()> >(
+ std::bind(std::forward<F>(f), std::forward<Args>(args)...)
+ );
+
+ std::future<return_type> res = task->get_future();
+ {
+ std::unique_lock<std::mutex> lock(queue_mutex);
+
+ // don't allow enqueueing after stopping the pool
+ if(stop)
+ throw std::runtime_error("enqueue on stopped ThreadPool");
+
+ tasks.emplace([task](){ (*task)(); });
+ }
+ condition.notify_one();
+ return res;
+}
+
+// the destructor joins all threads
+inline ThreadPool::~ThreadPool()
+{
+ {
+ std::unique_lock<std::mutex> lock(queue_mutex);
+ stop = true;
+ }
+ condition.notify_all();
+ for(std::thread &worker: workers)
+ worker.join();
+}
+
+#endif
\ No newline at end of file
diff --git a/be/src/extern/diskann/include/abstract_data_store.h b/be/src/extern/diskann/include/abstract_data_store.h
new file mode 100644
index 00000000000..89856f1fa20
--- /dev/null
+++ b/be/src/extern/diskann/include/abstract_data_store.h
@@ -0,0 +1,127 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <vector>
+#include <string>
+
+#include "types.h"
+#include "windows_customizations.h"
+#include "distance.h"
+
+namespace diskann
+{
+
+template <typename data_t> class AbstractScratch;
+
+template <typename data_t> class AbstractDataStore
+{
+ public:
+ AbstractDataStore(const location_t capacity, const size_t dim);
+
+ virtual ~AbstractDataStore() = default;
+
+ // Return number of points returned
+ virtual location_t load(const std::string &filename) = 0;
+
+ // Why does store take num_pts? Since store only has capacity, but we allow
+ // resizing we can end up in a situation where the store has spare capacity.
+ // To optimize disk utilization, we pass the number of points that are "true"
+ // points, so that the store can discard the empty locations before saving.
+ virtual size_t save(const std::string &filename, const location_t num_pts) = 0;
+
+ DISKANN_DLLEXPORT virtual location_t capacity() const;
+
+ DISKANN_DLLEXPORT virtual size_t get_dims() const;
+
+ // Implementers can choose to return _dim if they are not
+ // concerned about memory alignment.
+ // Some distance metrics (like l2) need data vectors to be aligned, so we
+ // align the dimension by padding zeros.
+ virtual size_t get_aligned_dim() const = 0;
+
+ // populate the store with vectors (either from a pointer or bin file),
+ // potentially after pre-processing the vectors if the metric deems so
+ // e.g., normalizing vectors for cosine distance over floating-point vectors
+ // useful for bulk or static index building.
+ virtual void populate_data(const data_t *vectors, const location_t num_pts) = 0;
+ virtual void populate_data(const std::string &filename, const size_t offset) = 0;
+
+ // save the first num_pts many vectors back to bin file
+ // note: cannot undo the pre-processing done in populate data
+ virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) = 0;
+
+ // Returns the updated capacity of the datastore. Clients should check
+ // if resize actually changed the capacity to new_num_points before
+ // proceeding with operations. See the code below:
+ // auto new_capcity = data_store->resize(new_num_points);
+ // if ( new_capacity >= new_num_points) {
+ // //PROCEED
+ // else
+ // //ERROR.
+ virtual location_t resize(const location_t new_num_points);
+
+ // operations on vectors
+ // like populate_data function, but over one vector at a time useful for
+ // streaming setting
+ virtual void get_vector(const location_t i, data_t *dest) const = 0;
+ virtual void set_vector(const location_t i, const data_t *const vector) = 0;
+ virtual void prefetch_vector(const location_t loc) = 0;
+
+ // internal shuffle operations to move around vectors
+ // will bulk-move all the vectors in [old_start_loc, old_start_loc +
+ // num_points) to [new_start_loc, new_start_loc + num_points) and set the old
+ // positions to zero vectors.
+ virtual void move_vectors(const location_t old_start_loc, const location_t new_start_loc,
+ const location_t num_points) = 0;
+
+ // same as above, without resetting the vectors in [from_loc, from_loc +
+ // num_points) to zero
+ virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0;
+
+ // With the PQ Data Store PR, we have also changed iterate_to_fixed_point to NOT take the query
+ // from the scratch object. Therefore every data store has to implement preprocess_query which
+ // at the least will be to copy the query into the scratch object. So making this pure virtual.
+ virtual void preprocess_query(const data_t *aligned_query,
+ AbstractScratch<data_t> *query_scratch = nullptr) const = 0;
+ // distance functions.
+ virtual float get_distance(const data_t *query, const location_t loc) const = 0;
+ virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count,
+ float *distances, AbstractScratch<data_t> *scratch_space = nullptr) const = 0;
+ // Specific overload for index.cpp.
+ virtual void get_distance(const data_t *preprocessed_query, const std::vector<location_t> &ids,
+ std::vector<float> &distances, AbstractScratch<data_t> *scratch_space) const = 0;
+ virtual float get_distance(const location_t loc1, const location_t loc2) const = 0;
+
+ // stats of the data stored in store
+ // Returns the point in the dataset that is closest to the mean of all points
+ // in the dataset
+ virtual location_t calculate_medoid() const = 0;
+
+ // REFACTOR PQ TODO: Each data store knows about its distance function, so this is
+ // redundant. However, we don't have an OptmizedDataStore yet, and to preserve code
+ // compability, we are exposing this function.
+ virtual Distance<data_t> *get_dist_fn() const = 0;
+
+ // search helpers
+ // if the base data is aligned per the request of the metric, this will tell
+ // how to align the query vector in a consistent manner
+ virtual size_t get_alignment_factor() const = 0;
+
+ protected:
+ // Expand the datastore to new_num_points. Returns the new capacity created,
+ // which should be == new_num_points in the normal case. Implementers can also
+ // return _capacity to indicate that there are not implementing this method.
+ virtual location_t expand(const location_t new_num_points) = 0;
+
+ // Shrink the datastore to new_num_points. It is NOT an error if shrink
+ // doesn't reduce the capacity so callers need to check this correctly. See
+ // also for "default" implementation
+ virtual location_t shrink(const location_t new_num_points) = 0;
+
+ location_t _capacity;
+ size_t _dim;
+};
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/abstract_graph_store.h b/be/src/extern/diskann/include/abstract_graph_store.h
new file mode 100644
index 00000000000..4ef09c724d0
--- /dev/null
+++ b/be/src/extern/diskann/include/abstract_graph_store.h
@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "types.h"
+
+namespace diskann
+{
+
+class AbstractGraphStore
+{
+ public:
+ AbstractGraphStore(const size_t total_pts, const size_t reserve_graph_degree)
+ : _capacity(total_pts), _reserve_graph_degree(reserve_graph_degree)
+ {
+ }
+
+ virtual ~AbstractGraphStore() = default;
+
+ // returns tuple of <nodes_read, start, num_frozen_points>
+ virtual std::tuple<uint32_t, uint32_t, size_t> load(const std::string &index_path_prefix,
+ const size_t num_points) = 0;
+ virtual int store(const std::string &index_path_prefix, const size_t num_points, const size_t num_fz_points,
+ const uint32_t start) = 0;
+
+ virtual int store(std::stringstream &index_stream, const size_t num_points,
+ const size_t num_frozen_points, const uint32_t start){
+ return 0;
+ }
+ virtual int save_graph(std::stringstream &out, const size_t num_points,
+ const size_t num_frozen_points, const uint32_t start){ return 0;}
+
+ // not synchronised, user should use lock when necvessary.
+ virtual const std::vector<location_t> &get_neighbours(const location_t i) const = 0;
+ virtual void add_neighbour(const location_t i, location_t neighbour_id) = 0;
+ virtual void clear_neighbours(const location_t i) = 0;
+ virtual void swap_neighbours(const location_t a, location_t b) = 0;
+
+ virtual void set_neighbours(const location_t i, std::vector<location_t> &neighbours) = 0;
+
+ virtual size_t resize_graph(const size_t new_size) = 0;
+ virtual void clear_graph() = 0;
+
+ virtual uint32_t get_max_observed_degree() = 0;
+
+ // set during load
+ virtual size_t get_max_range_of_graph() = 0;
+
+ // Total internal points _max_points + _num_frozen_points
+ size_t get_total_points()
+ {
+ return _capacity;
+ }
+
+ protected:
+ // Internal function, changes total points when resize_graph is called.
+ void set_total_points(size_t new_capacity)
+ {
+ _capacity = new_capacity;
+ }
+
+ size_t get_reserve_graph_degree()
+ {
+ return _reserve_graph_degree;
+ }
+
+ private:
+ size_t _capacity;
+ size_t _reserve_graph_degree;
+};
+
+} // namespace diskann
\ No newline at end of file
diff --git a/be/src/extern/diskann/include/abstract_index.h b/be/src/extern/diskann/include/abstract_index.h
new file mode 100644
index 00000000000..059866f7c1a
--- /dev/null
+++ b/be/src/extern/diskann/include/abstract_index.h
@@ -0,0 +1,129 @@
+#pragma once
+#include "distance.h"
+#include "parameters.h"
+#include "utils.h"
+#include "types.h"
+#include "index_config.h"
+#include "index_build_params.h"
+#include <any>
+
+namespace diskann
+{
+struct consolidation_report
+{
+ enum status_code
+ {
+ SUCCESS = 0,
+ FAIL = 1,
+ LOCK_FAIL = 2,
+ INCONSISTENT_COUNT_ERROR = 3
+ };
+ status_code _status;
+ size_t _active_points, _max_points, _empty_slots, _slots_released, _delete_set_size, _num_calls_to_process_delete;
+ double _time;
+
+ consolidation_report(status_code status, size_t active_points, size_t max_points, size_t empty_slots,
+ size_t slots_released, size_t delete_set_size, size_t num_calls_to_process_delete,
+ double time_secs)
+ : _status(status), _active_points(active_points), _max_points(max_points), _empty_slots(empty_slots),
+ _slots_released(slots_released), _delete_set_size(delete_set_size),
+ _num_calls_to_process_delete(num_calls_to_process_delete), _time(time_secs)
+ {
+ }
+};
+
+/* A templated independent class for intercation with Index. Uses Type Erasure to add virtual implemetation of methods
+that can take any type(using std::any) and Provides a clean API that can be inherited by different type of Index.
+*/
+class AbstractIndex
+{
+ public:
+ AbstractIndex() = default;
+ virtual ~AbstractIndex() = default;
+
+ virtual void build(const std::string &data_file, const size_t num_points_to_load,
+ IndexFilterParams &build_params) = 0;
+
+ template <typename data_type, typename tag_type>
+ void build(const data_type *data, const size_t num_points_to_load, const std::vector<tag_type> &tags);
+
+ virtual void save(const char *filename, bool compact_before_save = false) = 0;
+
+#ifdef EXEC_ENV_OLS
+ virtual void load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l) = 0;
+#else
+ virtual void load(const char *index_file, uint32_t num_threads, uint32_t search_l) = 0;
+#endif
+
+ // For FastL2 search on optimized layout
+ template <typename data_type>
+ void search_with_optimized_layout(const data_type *query, size_t K, size_t L, uint32_t *indices);
+
+ // Initialize space for res_vectors before calling.
+ template <typename data_type, typename tag_type>
+ size_t search_with_tags(const data_type *query, const uint64_t K, const uint32_t L, tag_type *tags,
+ float *distances, std::vector<data_type *> &res_vectors, bool use_filters = false,
+ const std::string filter_label = "");
+
+ // Added search overload that takes L as parameter, so that we
+ // can customize L on a per-query basis without tampering with "Parameters"
+ // IDtype is either uint32_t or uint64_t
+ template <typename data_type, typename IDType>
+ std::pair<uint32_t, uint32_t> search(const data_type *query, const size_t K, const uint32_t L, IDType *indices,
+ float *distances = nullptr);
+
+ // Filter support search
+ // IndexType is either uint32_t or uint64_t
+ template <typename IndexType>
+ std::pair<uint32_t, uint32_t> search_with_filters(const DataType &query, const std::string &raw_label,
+ const size_t K, const uint32_t L, IndexType *indices,
+ float *distances);
+
+ // insert points with labels, labels should be present for filtered index
+ template <typename data_type, typename tag_type, typename label_type>
+ int insert_point(const data_type *point, const tag_type tag, const std::vector<label_type> &labels);
+
+ // insert point for unfiltered index build. do not use with filtered index
+ template <typename data_type, typename tag_type> int insert_point(const data_type *point, const tag_type tag);
+
+ // delete point with tag, or return -1 if point can not be deleted
+ template <typename tag_type> int lazy_delete(const tag_type &tag);
+
+ // batch delete tags and populates failed tags if unabke to delete given tags.
+ template <typename tag_type>
+ void lazy_delete(const std::vector<tag_type> &tags, std::vector<tag_type> &failed_tags);
+
+ template <typename tag_type> void get_active_tags(tsl::robin_set<tag_type> &active_tags);
+
+ template <typename data_type> void set_start_points_at_random(data_type radius, uint32_t random_seed = 0);
+
+ virtual consolidation_report consolidate_deletes(const IndexWriteParameters ¶meters) = 0;
+
+ virtual void optimize_index_layout() = 0;
+
+ // memory should be allocated for vec before calling this function
+ template <typename tag_type, typename data_type> int get_vector_by_tag(tag_type &tag, data_type *vec);
+
+ template <typename label_type> void set_universal_label(const label_type universal_label);
+
+ private:
+ virtual void _build(const DataType &data, const size_t num_points_to_load, TagVector &tags) = 0;
+ virtual std::pair<uint32_t, uint32_t> _search(const DataType &query, const size_t K, const uint32_t L,
+ std::any &indices, float *distances = nullptr) = 0;
+ virtual std::pair<uint32_t, uint32_t> _search_with_filters(const DataType &query, const std::string &filter_label,
+ const size_t K, const uint32_t L, std::any &indices,
+ float *distances) = 0;
+ virtual int _insert_point(const DataType &data_point, const TagType tag, Labelvector &labels) = 0;
+ virtual int _insert_point(const DataType &data_point, const TagType tag) = 0;
+ virtual int _lazy_delete(const TagType &tag) = 0;
+ virtual void _lazy_delete(TagVector &tags, TagVector &failed_tags) = 0;
+ virtual void _get_active_tags(TagRobinSet &active_tags) = 0;
+ virtual void _set_start_points_at_random(DataType radius, uint32_t random_seed = 0) = 0;
+ virtual int _get_vector_by_tag(TagType &tag, DataType &vec) = 0;
+ virtual size_t _search_with_tags(const DataType &query, const uint64_t K, const uint32_t L, const TagType &tags,
+ float *distances, DataVector &res_vectors, bool use_filters = false,
+ const std::string filter_label = "") = 0;
+ virtual void _search_with_optimized_layout(const DataType &query, size_t K, size_t L, uint32_t *indices) = 0;
+ virtual void _set_universal_label(const LabelType universal_label) = 0;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/abstract_scratch.h b/be/src/extern/diskann/include/abstract_scratch.h
new file mode 100644
index 00000000000..b42a836f683
--- /dev/null
+++ b/be/src/extern/diskann/include/abstract_scratch.h
@@ -0,0 +1,35 @@
+#pragma once
+namespace diskann
+{
+
+template <typename data_t> class PQScratch;
+
+// By somewhat more than a coincidence, it seems that both InMemQueryScratch
+// and SSDQueryScratch have the aligned query and PQScratch objects. So we
+// can put them in a neat hierarchy and keep PQScratch as a standalone class.
+template <typename data_t> class AbstractScratch
+{
+ public:
+ AbstractScratch() = default;
+ // This class does not take any responsibilty for memory management of
+ // its members. It is the responsibility of the derived classes to do so.
+ virtual ~AbstractScratch() = default;
+
+ // Scratch objects should not be copied
+ AbstractScratch(const AbstractScratch &) = delete;
+ AbstractScratch &operator=(const AbstractScratch &) = delete;
+
+ data_t *aligned_query_T()
+ {
+ return _aligned_query_T;
+ }
+ PQScratch<data_t> *pq_scratch()
+ {
+ return _pq_scratch;
+ }
+
+ protected:
+ data_t *_aligned_query_T = nullptr;
+ PQScratch<data_t> *_pq_scratch = nullptr;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/aligned_file_reader.h b/be/src/extern/diskann/include/aligned_file_reader.h
new file mode 100644
index 00000000000..f8e9c302385
--- /dev/null
+++ b/be/src/extern/diskann/include/aligned_file_reader.h
@@ -0,0 +1,70 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#define MAX_IO_DEPTH 128
+
+#include <vector>
+#include <atomic>
+
+#include <fcntl.h>
+#include <libaio.h>
+#include <unistd.h>
+typedef io_context_t IOContext;
+
+#include <malloc.h>
+#include <cstdio>
+#include <mutex>
+#include <thread>
+#include "tsl/robin_map.h"
+#include "utils.h"
+
+// NOTE :: all 3 fields must be 512-aligned
+struct AlignedRead
+{
+ uint64_t offset; // where to read from
+ uint64_t len; // how much to read
+ void *buf; // where to read into
+
+ AlignedRead() : offset(0), len(0), buf(nullptr)
+ {
+ }
+
+ AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf)
+ {
+ assert(IS_512_ALIGNED(offset));
+ assert(IS_512_ALIGNED(len));
+ assert(IS_512_ALIGNED(buf));
+ // assert(malloc_usable_size(buf) >= len);
+ }
+};
+
+class AlignedFileReader
+{
+ protected:
+ tsl::robin_map<std::thread::id, IOContext> ctx_map;
+ std::mutex ctx_mut;
+
+ public:
+ // returns the thread-specific context
+ // returns (io_context_t)(-1) if thread is not registered
+ virtual IOContext &get_ctx() = 0;
+
+ virtual ~AlignedFileReader(){};
+
+ // register thread-id for a context
+ virtual void register_thread() = 0;
+ // de-register thread-id for a context
+ virtual void deregister_thread() = 0;
+ virtual void deregister_all_threads() = 0;
+
+ // Open & close ops
+ // Blocking calls
+ virtual void open(const std::string &fname) = 0;
+ virtual void close() = 0;
+
+ // process batch of aligned requests in parallel
+ // NOTE :: blocking call
+ virtual void read(std::vector<AlignedRead> &read_reqs, IOContext &ctx, bool async = false) = 0;
+};
\ No newline at end of file
diff --git a/be/src/extern/diskann/include/ann_exception.h b/be/src/extern/diskann/include/ann_exception.h
new file mode 100644
index 00000000000..9f55c8c1855
--- /dev/null
+++ b/be/src/extern/diskann/include/ann_exception.h
@@ -0,0 +1,36 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <string>
+#include <stdexcept>
+#include <system_error>
+#include <cstdint>
+
+#include "windows_customizations.h"
+
+#ifndef _WINDOWS
+#define __FUNCSIG__ __PRETTY_FUNCTION__
+#endif
+
+namespace diskann
+{
+
+class ANNException : public std::runtime_error
+{
+ public:
+ DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode);
+ DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode, const std::string &funcSig,
+ const std::string &fileName, uint32_t lineNum);
+
+ private:
+ int _errorCode;
+};
+
+class FileException : public ANNException
+{
+ public:
+ DISKANN_DLLEXPORT FileException(const std::string &filename, std::system_error &e, const std::string &funcSig,
+ const std::string &fileName, uint32_t lineNum);
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/any_wrappers.h b/be/src/extern/diskann/include/any_wrappers.h
new file mode 100644
index 00000000000..da9005cfbd7
--- /dev/null
+++ b/be/src/extern/diskann/include/any_wrappers.h
@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <cstdint>
+#include <cstddef>
+#include <vector>
+#include <any>
+#include "tsl/robin_set.h"
+
+namespace AnyWrapper
+{
+
+/*
+ * Base Struct to hold refrence to the data.
+ * Note: No memory mamagement, caller need to keep object alive.
+ */
+struct AnyReference
+{
+ template <typename Ty> AnyReference(Ty &reference) : _data(&reference)
+ {
+ }
+
+ template <typename Ty> Ty &get()
+ {
+ auto ptr = std::any_cast<Ty *>(_data);
+ return *ptr;
+ }
+
+ private:
+ std::any _data;
+};
+struct AnyRobinSet : public AnyReference
+{
+ template <typename T> AnyRobinSet(const tsl::robin_set<T> &robin_set) : AnyReference(robin_set)
+ {
+ }
+ template <typename T> AnyRobinSet(tsl::robin_set<T> &robin_set) : AnyReference(robin_set)
+ {
+ }
+};
+
+struct AnyVector : public AnyReference
+{
+ template <typename T> AnyVector(const std::vector<T> &vector) : AnyReference(vector)
+ {
+ }
+ template <typename T> AnyVector(std::vector<T> &vector) : AnyReference(vector)
+ {
+ }
+};
+} // namespace AnyWrapper
diff --git a/be/src/extern/diskann/include/boost_dynamic_bitset_fwd.h b/be/src/extern/diskann/include/boost_dynamic_bitset_fwd.h
new file mode 100644
index 00000000000..5aebb2bc220
--- /dev/null
+++ b/be/src/extern/diskann/include/boost_dynamic_bitset_fwd.h
@@ -0,0 +1,11 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+namespace boost
+{
+#ifndef BOOST_DYNAMIC_BITSET_FWD_HPP
+template <typename Block = unsigned long, typename Allocator = std::allocator<Block>> class dynamic_bitset;
+#endif
+} // namespace boost
diff --git a/be/src/extern/diskann/include/cached_io.h b/be/src/extern/diskann/include/cached_io.h
new file mode 100644
index 00000000000..daef2f2f715
--- /dev/null
+++ b/be/src/extern/diskann/include/cached_io.h
@@ -0,0 +1,217 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <sstream>
+
+#include "logger.h"
+#include "ann_exception.h"
+
+// sequential cached reads
+class cached_ifstream
+{
+ public:
+ cached_ifstream()
+ {
+ }
+ cached_ifstream(const std::string &filename, uint64_t cacheSize) : cache_size(cacheSize), cur_off(0)
+ {
+ reader.exceptions(std::ifstream::failbit | std::ifstream::badbit);
+ this->open(filename, cache_size);
+ }
+ ~cached_ifstream()
+ {
+ delete[] cache_buf;
+ reader.close();
+ }
+
+ void open(const std::string &filename, uint64_t cacheSize)
+ {
+ this->cur_off = 0;
+
+ try
+ {
+ reader.open(filename, std::ios::binary | std::ios::ate);
+ fsize = reader.tellg();
+ reader.seekg(0, std::ios::beg);
+ assert(reader.is_open());
+ assert(cacheSize > 0);
+ cacheSize = (std::min)(cacheSize, fsize);
+ this->cache_size = cacheSize;
+ cache_buf = new char[cacheSize];
+ reader.read(cache_buf, cacheSize);
+ diskann::cout << "Opened: " << filename.c_str() << ", size: " << fsize << ", cache_size: " << cacheSize
+ << std::endl;
+ }
+ catch (std::system_error &e)
+ {
+ throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, __LINE__);
+ }
+ }
+
+ size_t get_file_size()
+ {
+ return fsize;
+ }
+
+ void read(char *read_buf, uint64_t n_bytes)
+ {
+ assert(cache_buf != nullptr);
+ assert(read_buf != nullptr);
+
+ if (n_bytes <= (cache_size - cur_off))
+ {
+ // case 1: cache contains all data
+ memcpy(read_buf, cache_buf + cur_off, n_bytes);
+ cur_off += n_bytes;
+ }
+ else
+ {
+ // case 2: cache contains some data
+ uint64_t cached_bytes = cache_size - cur_off;
+ if (n_bytes - cached_bytes > fsize - reader.tellg())
+ {
+ std::stringstream stream;
+ stream << "Reading beyond end of file" << std::endl;
+ stream << "n_bytes: " << n_bytes << " cached_bytes: " << cached_bytes << " fsize: " << fsize
+ << " current pos:" << reader.tellg() << std::endl;
+ diskann::cout << stream.str() << std::endl;
+ throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__);
+ }
+ memcpy(read_buf, cache_buf + cur_off, cached_bytes);
+
+ // go to disk and fetch more data
+ reader.read(read_buf + cached_bytes, n_bytes - cached_bytes);
+ // reset cur off
+ cur_off = cache_size;
+
+ uint64_t size_left = fsize - reader.tellg();
+
+ if (size_left >= cache_size)
+ {
+ reader.read(cache_buf, cache_size);
+ cur_off = 0;
+ }
+ // note that if size_left < cache_size, then cur_off = cache_size,
+ // so subsequent reads will all be directly from file
+ }
+ }
+
+ private:
+ // underlying ifstream
+ std::ifstream reader;
+ // # bytes to cache in one shot read
+ uint64_t cache_size = 0;
+ // underlying buf for cache
+ char *cache_buf = nullptr;
+ // offset into cache_buf for cur_pos
+ uint64_t cur_off = 0;
+ // file size
+ uint64_t fsize = 0;
+};
+
+// sequential cached writes
+class cached_ofstream
+{
+ public:
+ cached_ofstream(const std::string &filename, uint64_t cache_size) : cache_size(cache_size), cur_off(0)
+ {
+ writer.exceptions(std::ifstream::failbit | std::ifstream::badbit);
+ try
+ {
+ writer.open(filename, std::ios::binary);
+ assert(writer.is_open());
+ assert(cache_size > 0);
+ cache_buf = new char[cache_size];
+ diskann::cout << "Opened: " << filename.c_str() << ", cache_size: " << cache_size << std::endl;
+ }
+ catch (std::system_error &e)
+ {
+ throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, __LINE__);
+ }
+ }
+
+ ~cached_ofstream()
+ {
+ this->close();
+ }
+
+ void close()
+ {
+ // dump any remaining data in memory
+ if (cur_off > 0)
+ {
+ this->flush_cache();
+ }
+
+ if (cache_buf != nullptr)
+ {
+ delete[] cache_buf;
+ cache_buf = nullptr;
+ }
+
+ if (writer.is_open())
+ writer.close();
+ diskann::cout << "Finished writing " << fsize << "B" << std::endl;
+ }
+
+ size_t get_file_size()
+ {
+ return fsize;
+ }
+ // writes n_bytes from write_buf to the underlying ofstream/cache
+ void write(char *write_buf, uint64_t n_bytes)
+ {
+ assert(cache_buf != nullptr);
+ if (n_bytes <= (cache_size - cur_off))
+ {
+ // case 1: cache can take all data
+ memcpy(cache_buf + cur_off, write_buf, n_bytes);
+ cur_off += n_bytes;
+ }
+ else
+ {
+ // case 2: cache cant take all data
+ // go to disk and write existing cache data
+ writer.write(cache_buf, cur_off);
+ fsize += cur_off;
+ // write the new data to disk
+ writer.write(write_buf, n_bytes);
+ fsize += n_bytes;
+ // memset all cache data and reset cur_off
+ memset(cache_buf, 0, cache_size);
+ cur_off = 0;
+ }
+ }
+
+ void flush_cache()
+ {
+ assert(cache_buf != nullptr);
+ writer.write(cache_buf, cur_off);
+ fsize += cur_off;
+ memset(cache_buf, 0, cache_size);
+ cur_off = 0;
+ }
+
+ void reset()
+ {
+ flush_cache();
+ writer.seekp(0);
+ }
+
+ private:
+ // underlying ofstream
+ std::ofstream writer;
+ // # bytes to cache for one shot write
+ uint64_t cache_size = 0;
+ // underlying buf for cache
+ char *cache_buf = nullptr;
+ // offset into cache_buf for cur_pos
+ uint64_t cur_off = 0;
+
+ // file size
+ uint64_t fsize = 0;
+};
diff --git a/be/src/extern/diskann/include/combined_file.h b/be/src/extern/diskann/include/combined_file.h
new file mode 100644
index 00000000000..d5285f3fdf2
--- /dev/null
+++ b/be/src/extern/diskann/include/combined_file.h
@@ -0,0 +1,105 @@
+#pragma once
+
+#include <iostream>
+#include <fstream>
+#include <string>
+#include <vector>
+#include <cstring>
+#include <cstdio>
+#include <map>
+#include <memory>
+#include <mutex>
+
+namespace diskann {
+
+// 抽象的文件读取器基类
+class Reader {
+ public:
+ virtual ~Reader() {}
+ virtual size_t read(char* buffer, uint64_t offset, size_t size) = 0;
+ virtual void seek(uint64_t offset) = 0;
+ virtual bool open(const std::string& file_path) = 0;
+ virtual void close() = 0;
+ virtual uint64_t get_file_size() = 0;
+ virtual uint64_t get_base_offset() = 0;
+};
+
+class IOWriter {
+ public:
+ virtual ~IOWriter() {}
+ virtual void write(const char* data, size_t size) = 0;
+ virtual void close() = 0;
+};
+
+
+
+class LocalFileReader : public Reader {
+ private:
+ std::ifstream in;
+ std::string _file_path;
+ //文件起始
+ uint64_t _base_offset = 0;
+ uint64_t _size = 0 ;
+ std::mutex _mutex;
+ public:
+ LocalFileReader(const std::string& file_path, uint64_t start = 0, uint64_t size = 0);
+
+ LocalFileReader();
+
+ bool open(const std::string& file_path);
+
+ uint64_t get_file_size();
+
+ size_t read(char* buffer, uint64_t offset, size_t size);
+
+ void seek(uint64_t offset);
+
+ void close();
+
+ uint64_t get_base_offset() {
+ return _base_offset;
+ }
+};
+
+class LocalFileIOWriter : public IOWriter {
+ public:
+ LocalFileIOWriter(const std::string& file_path);
+ void write(const char* data, size_t size);
+ void close();
+ private:
+ std::ofstream out;
+};
+
+class FileMerger {
+ private:
+ struct FileInfo {
+ std::string alias;
+ std::string path;
+ uint64_t offset;
+ uint64_t size;
+ };
+ std::vector<FileInfo> files;
+ uint64_t current_offset;
+
+ // 使用map来缓存已经实例化过的reader，键为文件别名，值为对应的reader智能指针
+ std::map<std::string, std::shared_ptr<Reader>> reader_cache;
+
+ public:
+ FileMerger() : current_offset(0) {}
+
+ // 添加文件方法
+ bool add(const std::string& alias, const std::string& path);
+ void clear();
+
+ // 把最终合并的文件写到磁盘，文件的meta信息要放到最后，且方便反序列化，同时记录total_meta_size并写入磁盘末尾
+ void save(IOWriter* writer);
+
+ // 根据文件名原始文件别名， 获取base_offset和文件大小，并根据base_offset和文件创建一个reader，返回这个reader
+ // 这里从磁盘读取文件相关offset和size信息（假设磁盘文件格式符合之前定义的规则），先获取total_meta_size来定位meta起始位置
+ // 同时添加了缓存功能，记录已经找到的reader，如果发现之前实例化过则直接返回
+ template <typename ReaderType>
+ std::shared_ptr<ReaderType> get_reader(const std::string& alias, const std::string& merged_file_path);
+};
+
+
+}
diff --git a/be/src/extern/diskann/include/common_includes.h b/be/src/extern/diskann/include/common_includes.h
new file mode 100644
index 00000000000..e1a51bdec34
--- /dev/null
+++ b/be/src/extern/diskann/include/common_includes.h
@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <algorithm>
+#include <atomic>
+#include <cassert>
+#include <chrono>
+#include <climits>
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <ctime>
+#include <fcntl.h>
+#include <fstream>
+#include <iostream>
+#include <iomanip>
+#include <omp.h>
+#include <queue>
+#include <random>
+#include <set>
+#include <shared_mutex>
+#include <sys/stat.h>
+#include <sstream>
+#include <unordered_map>
+#include <vector>
diff --git a/be/src/extern/diskann/include/concurrent_queue.h b/be/src/extern/diskann/include/concurrent_queue.h
new file mode 100644
index 00000000000..1e57bbf0f03
--- /dev/null
+++ b/be/src/extern/diskann/include/concurrent_queue.h
@@ -0,0 +1,132 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
+#include <mutex>
+#include <queue>
+#include <thread>
+#include <type_traits>
+#include <unordered_set>
+
+namespace diskann
+{
+
+template <typename T> class ConcurrentQueue
+{
+ typedef std::chrono::microseconds chrono_us_t;
+ typedef std::unique_lock<std::mutex> mutex_locker;
+
+ std::queue<T> q;
+ std::mutex mut;
+ std::mutex push_mut;
+ std::mutex pop_mut;
+ std::condition_variable push_cv;
+ std::condition_variable pop_cv;
+ T null_T;
+
+ public:
+ ConcurrentQueue()
+ {
+ }
+
+ ConcurrentQueue(T nullT)
+ {
+ this->null_T = nullT;
+ }
+
+ ~ConcurrentQueue()
+ {
+ this->push_cv.notify_all();
+ this->pop_cv.notify_all();
+ }
+
+ // queue stats
+ uint64_t size()
+ {
+ mutex_locker lk(this->mut);
+ uint64_t ret = q.size();
+ lk.unlock();
+ return ret;
+ }
+
+ bool empty()
+ {
+ return (this->size() == 0);
+ }
+
+ // PUSH BACK
+ void push(T &new_val)
+ {
+ mutex_locker lk(this->mut);
+ this->q.push(new_val);
+ lk.unlock();
+ }
+
+ template <class Iterator> void insert(Iterator iter_begin, Iterator iter_end)
+ {
+ mutex_locker lk(this->mut);
+ for (Iterator it = iter_begin; it != iter_end; it++)
+ {
+ this->q.push(*it);
+ }
+ lk.unlock();
+ }
+
+ // POP FRONT
+ T pop()
+ {
+ mutex_locker lk(this->mut);
+ if (this->q.empty())
+ {
+ lk.unlock();
+ return this->null_T;
+ }
+ else
+ {
+ T ret = this->q.front();
+ this->q.pop();
+ // diskann::cout << "thread_id: " << std::this_thread::get_id() <<
+ // ", ctx: "
+ // << ret.ctx << "\n";
+ lk.unlock();
+ return ret;
+ }
+ }
+
+ // register for notifications
+ void wait_for_push_notify(chrono_us_t wait_time = chrono_us_t{10})
+ {
+ mutex_locker lk(this->push_mut);
+ this->push_cv.wait_for(lk, wait_time);
+ lk.unlock();
+ }
+
+ void wait_for_pop_notify(chrono_us_t wait_time = chrono_us_t{10})
+ {
+ mutex_locker lk(this->pop_mut);
+ this->pop_cv.wait_for(lk, wait_time);
+ lk.unlock();
+ }
+
+ // just notify functions
+ void push_notify_one()
+ {
+ this->push_cv.notify_one();
+ }
+ void push_notify_all()
+ {
+ this->push_cv.notify_all();
+ }
+ void pop_notify_one()
+ {
+ this->pop_cv.notify_one();
+ }
+ void pop_notify_all()
+ {
+ this->pop_cv.notify_all();
+ }
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/cosine_similarity.h b/be/src/extern/diskann/include/cosine_similarity.h
new file mode 100644
index 00000000000..dc51f6c0ab5
--- /dev/null
+++ b/be/src/extern/diskann/include/cosine_similarity.h
@@ -0,0 +1,283 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <immintrin.h>
+#include <smmintrin.h>
+#include <tmmintrin.h>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <vector>
+#include <limits>
+#include <algorithm>
+#include <stdexcept>
+
+#include "simd_utils.h"
+
+extern bool Avx2SupportedCPU;
+
+#ifdef _WINDOWS
+// SIMD implementation of Cosine similarity. Taken from hnsw library.
+
+/**
+ * Non-metric Space Library
+ *
+ * Authors: Bilegsaikhan Naidan (https://github.com/bileg), Leonid Boytsov
+ * (http://boytsov.info). With contributions from Lawrence Cayton
+ * (http://lcayton.com/) and others.
+ *
+ * For the complete list of contributors and further details see:
+ * https://github.com/searchivarius/NonMetricSpaceLib
+ *
+ * Copyright (c) 2014
+ *
+ * This code is released under the
+ * Apache License Version 2.0 http://www.apache.org/licenses/.
+ *
+ */
+
+namespace diskann
+{
+
+using namespace std;
+
+#define PORTABLE_ALIGN16 __declspec(align(16))
+
+static float NormScalarProductSIMD2(const int8_t *pVect1, const int8_t *pVect2, uint32_t qty)
+{
+ if (Avx2SupportedCPU)
+ {
+ __m256 cos, p1Len, p2Len;
+ cos = p1Len = p2Len = _mm256_setzero_ps();
+ while (qty >= 32)
+ {
+ __m256i rx = _mm256_load_si256((__m256i *)pVect1), ry = _mm256_load_si256((__m256i *)pVect2);
+ cos = _mm256_add_ps(cos, _mm256_mul_epi8(rx, ry));
+ p1Len = _mm256_add_ps(p1Len, _mm256_mul_epi8(rx, rx));
+ p2Len = _mm256_add_ps(p2Len, _mm256_mul_epi8(ry, ry));
+ pVect1 += 32;
+ pVect2 += 32;
+ qty -= 32;
+ }
+ while (qty > 0)
+ {
+ __m128i rx = _mm_load_si128((__m128i *)pVect1), ry = _mm_load_si128((__m128i *)pVect2);
+ cos = _mm256_add_ps(cos, _mm256_mul32_pi8(rx, ry));
+ p1Len = _mm256_add_ps(p1Len, _mm256_mul32_pi8(rx, rx));
+ p2Len = _mm256_add_ps(p2Len, _mm256_mul32_pi8(ry, ry));
+ pVect1 += 4;
+ pVect2 += 4;
+ qty -= 4;
+ }
+ cos = _mm256_hadd_ps(_mm256_hadd_ps(cos, cos), cos);
+ p1Len = _mm256_hadd_ps(_mm256_hadd_ps(p1Len, p1Len), p1Len);
+ p2Len = _mm256_hadd_ps(_mm256_hadd_ps(p2Len, p2Len), p2Len);
+ float denominator = max(numeric_limits<float>::min() * 2, sqrt(p1Len.m256_f32[0] + p1Len.m256_f32[4]) *
+ sqrt(p2Len.m256_f32[0] + p2Len.m256_f32[4]));
+ float cosine = (cos.m256_f32[0] + cos.m256_f32[4]) / denominator;
+
+ return max(float(-1), min(float(1), cosine));
+ }
+
+ __m128 cos, p1Len, p2Len;
+ cos = p1Len = p2Len = _mm_setzero_ps();
+ __m128i rx, ry;
+ while (qty >= 16)
+ {
+ rx = _mm_load_si128((__m128i *)pVect1);
+ ry = _mm_load_si128((__m128i *)pVect2);
+ cos = _mm_add_ps(cos, _mm_mul_epi8(rx, ry));
+ p1Len = _mm_add_ps(p1Len, _mm_mul_epi8(rx, rx));
+ p2Len = _mm_add_ps(p2Len, _mm_mul_epi8(ry, ry));
+ pVect1 += 16;
+ pVect2 += 16;
+ qty -= 16;
+ }
+ while (qty > 0)
+ {
+ rx = _mm_load_si128((__m128i *)pVect1);
+ ry = _mm_load_si128((__m128i *)pVect2);
+ cos = _mm_add_ps(cos, _mm_mul32_pi8(rx, ry));
+ p1Len = _mm_add_ps(p1Len, _mm_mul32_pi8(rx, rx));
+ p2Len = _mm_add_ps(p2Len, _mm_mul32_pi8(ry, ry));
+ pVect1 += 4;
+ pVect2 += 4;
+ qty -= 4;
+ }
+ cos = _mm_hadd_ps(_mm_hadd_ps(cos, cos), cos);
+ p1Len = _mm_hadd_ps(_mm_hadd_ps(p1Len, p1Len), p1Len);
+ p2Len = _mm_hadd_ps(_mm_hadd_ps(p2Len, p2Len), p2Len);
+ float norm1 = p1Len.m128_f32[0];
+ float norm2 = p2Len.m128_f32[0];
+
+ static const float eps = numeric_limits<float>::min() * 2;
+
+ if (norm1 < eps)
+ { /*
+ * This shouldn't normally happen for this space, but
+ * if it does, we don't want to get NANs
+ */
+ if (norm2 < eps)
+ {
+ return 1;
+ }
+ return 0;
+ }
+ /*
+ * Sometimes due to rounding errors, we get values > 1 or < -1.
+ * This throws off other functions that use scalar product, e.g., acos
+ */
+ return max(float(-1), min(float(1), cos.m128_f32[0] / sqrt(norm1) / sqrt(norm2)));
+}
+
+static float NormScalarProductSIMD(const float *pVect1, const float *pVect2, uint32_t qty)
+{
+ // Didn't get significant performance gain compared with 128bit version.
+ static const float eps = numeric_limits<float>::min() * 2;
+
+ if (Avx2SupportedCPU)
+ {
+ uint32_t qty8 = qty / 8;
+
+ const float *pEnd1 = pVect1 + 8 * qty8;
+ const float *pEnd2 = pVect1 + qty;
+
+ __m256 v1, v2;
+ __m256 sum_prod = _mm256_set_ps(0, 0, 0, 0, 0, 0, 0, 0);
+ __m256 sum_square1 = sum_prod;
+ __m256 sum_square2 = sum_prod;
+
+ while (pVect1 < pEnd1)
+ {
+ v1 = _mm256_loadu_ps(pVect1);
+ pVect1 += 8;
+ v2 = _mm256_loadu_ps(pVect2);
+ pVect2 += 8;
+ sum_prod = _mm256_add_ps(sum_prod, _mm256_mul_ps(v1, v2));
+ sum_square1 = _mm256_add_ps(sum_square1, _mm256_mul_ps(v1, v1));
+ sum_square2 = _mm256_add_ps(sum_square2, _mm256_mul_ps(v2, v2));
+ }
+
+ float PORTABLE_ALIGN16 TmpResProd[8];
+ float PORTABLE_ALIGN16 TmpResSquare1[8];
+ float PORTABLE_ALIGN16 TmpResSquare2[8];
+
+ _mm256_store_ps(TmpResProd, sum_prod);
+ _mm256_store_ps(TmpResSquare1, sum_square1);
+ _mm256_store_ps(TmpResSquare2, sum_square2);
+
+ float sum = 0.0f;
+ float norm1 = 0.0f;
+ float norm2 = 0.0f;
+ for (uint32_t i = 0; i < 8; ++i)
+ {
+ sum += TmpResProd[i];
+ norm1 += TmpResSquare1[i];
+ norm2 += TmpResSquare2[i];
+ }
+
+ while (pVect1 < pEnd2)
+ {
+ sum += (*pVect1) * (*pVect2);
+ norm1 += (*pVect1) * (*pVect1);
+ norm2 += (*pVect2) * (*pVect2);
+
+ ++pVect1;
+ ++pVect2;
+ }
+
+ if (norm1 < eps)
+ {
+ return norm2 < eps ? 1.0f : 0.0f;
+ }
+
+ return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2)));
+ }
+
+ __m128 v1, v2;
+ __m128 sum_prod = _mm_set1_ps(0);
+ __m128 sum_square1 = sum_prod;
+ __m128 sum_square2 = sum_prod;
+
+ while (qty >= 4)
+ {
+ v1 = _mm_loadu_ps(pVect1);
+ pVect1 += 4;
+ v2 = _mm_loadu_ps(pVect2);
+ pVect2 += 4;
+ sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+ sum_square1 = _mm_add_ps(sum_square1, _mm_mul_ps(v1, v1));
+ sum_square2 = _mm_add_ps(sum_square2, _mm_mul_ps(v2, v2));
+
+ qty -= 4;
+ }
+
+ float sum = sum_prod.m128_f32[0] + sum_prod.m128_f32[1] + sum_prod.m128_f32[2] + sum_prod.m128_f32[3];
+ float norm1 = sum_square1.m128_f32[0] + sum_square1.m128_f32[1] + sum_square1.m128_f32[2] + sum_square1.m128_f32[3];
+ float norm2 = sum_square2.m128_f32[0] + sum_square2.m128_f32[1] + sum_square2.m128_f32[2] + sum_square2.m128_f32[3];
+
+ if (norm1 < eps)
+ {
+ return norm2 < eps ? 1.0f : 0.0f;
+ }
+
+ return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2)));
+}
+
+static float NormScalarProductSIMD2(const float *pVect1, const float *pVect2, uint32_t qty)
+{
+ return NormScalarProductSIMD(pVect1, pVect2, qty);
+}
+
+template <class T> static float CosineSimilarity2(const T *p1, const T *p2, uint32_t qty)
+{
+ return std::max(0.0f, 1.0f - NormScalarProductSIMD2(p1, p2, qty));
+}
+
+// static template float CosineSimilarity2<__int8>(const __int8* pVect1,
+// const __int8* pVect2, size_t qty);
+
+// static template float CosineSimilarity2<float>(const float* pVect1,
+// const float* pVect2, size_t qty);
+
+template <class T> static void CosineSimilarityNormalize(T *pVector, uint32_t qty)
+{
+ T sum = 0;
+ for (uint32_t i = 0; i < qty; ++i)
+ {
+ sum += pVector[i] * pVector[i];
+ }
+ sum = 1 / sqrt(sum);
+ if (sum == 0)
+ {
+ sum = numeric_limits<T>::min();
+ }
+ for (uint32_t i = 0; i < qty; ++i)
+ {
+ pVector[i] *= sum;
+ }
+}
+
+// template static void CosineSimilarityNormalize<float>(float* pVector,
+// size_t qty);
+// template static void CosineSimilarityNormalize<double>(double* pVector,
+// size_t qty);
+
+template <> void CosineSimilarityNormalize(__int8 * /*pVector*/, uint32_t /*qty*/)
+{
+ throw std::runtime_error("For int8 type vector, you can not use cosine distance!");
+}
+
+template <> void CosineSimilarityNormalize(__int16 * /*pVector*/, uint32_t /*qty*/)
+{
+ throw std::runtime_error("For int16 type vector, you can not use cosine distance!");
+}
+
+template <> void CosineSimilarityNormalize(int * /*pVector*/, uint32_t /*qty*/)
+{
+ throw std::runtime_error("For int type vector, you can not use cosine distance!");
+}
+} // namespace diskann
+#endif
\ No newline at end of file
diff --git a/be/src/extern/diskann/include/defaults.h b/be/src/extern/diskann/include/defaults.h
new file mode 100644
index 00000000000..ef1750fcf49
--- /dev/null
+++ b/be/src/extern/diskann/include/defaults.h
@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <stdint.h>
+
+namespace diskann
+{
+namespace defaults
+{
+const float ALPHA = 1.2f;
+const uint32_t NUM_THREADS = 0;
+const uint32_t MAX_OCCLUSION_SIZE = 750;
+const bool HAS_LABELS = false;
+const uint32_t FILTER_LIST_SIZE = 0;
+const uint32_t NUM_FROZEN_POINTS_STATIC = 0;
+const uint32_t NUM_FROZEN_POINTS_DYNAMIC = 1;
+
+// In-mem index related limits
+const float GRAPH_SLACK_FACTOR = 1.3f;
+
+// SSD Index related limits
+const uint64_t MAX_GRAPH_DEGREE = 512;
+const uint64_t SECTOR_LEN = 4096;
+const uint64_t MAX_N_SECTOR_READS = 128;
+
+// following constants should always be specified, but are useful as a
+// sensible default at cli / python boundaries
+const uint32_t MAX_DEGREE = 64;
+const uint32_t BUILD_LIST_SIZE = 100;
+const uint32_t SATURATE_GRAPH = false;
+const uint32_t SEARCH_LIST_SIZE = 100;
+} // namespace defaults
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/disk_utils.h b/be/src/extern/diskann/include/disk_utils.h
new file mode 100644
index 00000000000..ab633597eae
--- /dev/null
+++ b/be/src/extern/diskann/include/disk_utils.h
@@ -0,0 +1,135 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <algorithm>
+#include <fcntl.h>
+#include <cassert>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <random>
+#include <set>
+#ifdef __APPLE__
+#else
+#include <malloc.h>
+#endif
+
+#ifdef _WINDOWS
+#include <Windows.h>
+typedef HANDLE FileHandle;
+#else
+#include <unistd.h>
+typedef int FileHandle;
+#endif
+
+#include "cached_io.h"
+#include "common_includes.h"
+
+#include "utils.h"
+#include "windows_customizations.h"
+
+namespace diskann
+{
+const size_t MAX_SAMPLE_POINTS_FOR_WARMUP = 100000;
+const double PQ_TRAINING_SET_FRACTION = 0.1;
+const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25;
+const double THRESHOLD_FOR_CACHING_IN_GB = 1.0;
+const uint32_t NUM_NODES_TO_CACHE = 250000;
+const uint32_t WARMUP_L = 20;
+const uint32_t NUM_KMEANS_REPS = 12;
+
+const std::string PQ_PIVOTS = "pq_pivots";
+const std::string PQ_COMPRESSED = "pq_compressed";
+const std::string MEM_INDEX = "mem_index";
+const std::string DISK_INDEX_PATH = "disk_index_path";
+const std::string MEM_INDEX_DATA = "mem_index_data";
+
+template <typename T, typename LabelT> class PQFlashIndex;
+
+DISKANN_DLLEXPORT double get_memory_budget(const std::string &mem_budget_str);
+DISKANN_DLLEXPORT double get_memory_budget(double search_ram_budget_in_gb);
+DISKANN_DLLEXPORT void add_new_file_to_single_index(std::string index_file, std::string new_file);
+
+DISKANN_DLLEXPORT size_t calculate_num_pq_chunks(double final_index_ram_limit, size_t points_num, uint32_t dim);
+
+DISKANN_DLLEXPORT void read_idmap(const std::string &fname, std::vector<uint32_t> &ivecs);
+
+#ifdef EXEC_ENV_OLS
+template <typename T>
+DISKANN_DLLEXPORT T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, uint64_t &warmup_num,
+ uint64_t warmup_dim, uint64_t warmup_aligned_dim);
+#else
+template <typename T>
+DISKANN_DLLEXPORT T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, uint64_t warmup_dim,
+ uint64_t warmup_aligned_dim);
+#endif
+
+DISKANN_DLLEXPORT int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix,
+ const std::string &idmaps_prefix, const std::string &idmaps_suffix,
+ const uint64_t nshards, uint32_t max_degree, const std::string &output_vamana,
+ const std::string &medoids_file, bool use_filters = false,
+ const std::string &labels_to_medoids_file = std::string(""));
+
+DISKANN_DLLEXPORT void extract_shard_labels(const std::string &in_label_file, const std::string &shard_ids_bin,
+ const std::string &shard_label_file);
+
+template <typename T>
+DISKANN_DLLEXPORT std::string preprocess_base_file(const std::string &infile, const std::string &indexPrefix,
+ diskann::Metric &distMetric);
+
+template <typename T, typename LabelT = uint32_t>
+DISKANN_DLLEXPORT int build_merged_vamana_index(std::string base_file, diskann::Metric _compareMetric, uint32_t L,
+ uint32_t R, double sampling_rate, double ram_budget,
+ std::string mem_index_path, std::string medoids_file,
+ std::string centroids_file, size_t build_pq_bytes, bool use_opq,
+ uint32_t num_threads, bool use_filters = false,
+ const std::string &label_file = std::string(""),
+ const std::string &labels_to_medoids_file = std::string(""),
+ const std::string &universal_label = "", const uint32_t Lf = 0);
+
+template <typename T, typename LabelT = uint32_t>
+DISKANN_DLLEXPORT int build_merged_vamana_index(std::stringstream & data_stream, diskann::Metric _compareMetric, uint32_t L,
+ uint32_t R, double sampling_rate, double ram_budget,
+ std::stringstream &index_stream, std::string medoids_file,
+ std::string centroids_file, size_t build_pq_bytes, bool use_opq,
+ uint32_t num_threads, bool use_filters = false,
+ const std::string &label_file = std::string(""),
+ const std::string &labels_to_medoids_file = std::string(""),
+ const std::string &universal_label = "", const uint32_t Lf = 0);
+
+
+template <typename T, typename LabelT>
+DISKANN_DLLEXPORT uint32_t optimize_beamwidth(std::unique_ptr<diskann::PQFlashIndex<T, LabelT>> &_pFlashIndex,
+ T *tuning_sample, uint64_t tuning_sample_num,
+ uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads,
+ uint32_t start_bw = 2);
+
+template <typename T, typename LabelT = uint32_t>
+DISKANN_DLLEXPORT int build_disk_index(
+ const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters,
+ diskann::Metric _compareMetric, bool use_opq = false,
+ const std::string &codebook_prefix = "", // default is empty for no codebook pass in
+ bool use_filters = false,
+ const std::string &label_file = std::string(""), // default is empty string for no label_file
+ const std::string &universal_label = "", const uint32_t filter_threshold = 0,
+ const uint32_t Lf = 0); // default is empty string for no universal label
+
+template <typename T>
+DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, const std::string mem_index_file,
+ const std::string output_file,
+ const std::string reorder_data_file = std::string(""));
+
+template <typename T>
+DISKANN_DLLEXPORT void create_disk_layout(const T &data, const std::string mem_index_file,
+ const std::string output_file,
+ const std::string reorder_data_file = std::string(""));
+
+template <typename T>
+DISKANN_DLLEXPORT void create_disk_layout(std::stringstream & data_stream, std::stringstream &index_stream,
+ std::stringstream &disklayout_stream,
+ const std::string reorder_data_file = std::string(""));
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/distance.h b/be/src/extern/diskann/include/distance.h
new file mode 100644
index 00000000000..f3b1de25aaa
--- /dev/null
+++ b/be/src/extern/diskann/include/distance.h
@@ -0,0 +1,235 @@
+#pragma once
+#include "windows_customizations.h"
+#include <cstring>
+
+namespace diskann
+{
+enum Metric
+{
+ L2 = 0,
+ INNER_PRODUCT = 1,
+ COSINE = 2,
+ FAST_L2 = 3
+};
+
+template <typename T> class Distance
+{
+ public:
+ DISKANN_DLLEXPORT Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric)
+ {
+ }
+
+ // distance comparison function
+ DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const = 0;
+
+ // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE
+ DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, const float normA, const float normB,
+ uint32_t length) const;
+
+ // For MIPS, normalization adds an extra dimension to the vectors.
+ // This function lets callers know if the normalization process
+ // changes the dimension.
+ DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const;
+
+ DISKANN_DLLEXPORT virtual diskann::Metric get_metric() const;
+
+ // This is for efficiency. If no normalization is required, the callers
+ // can simply ignore the normalize_data_for_build() function.
+ DISKANN_DLLEXPORT virtual bool preprocessing_required() const;
+
+ // Check the preprocessing_required() function before calling this.
+ // Clients can call the function like this:
+ //
+ // if (metric->preprocessing_required()){
+ // T* normalized_data_batch;
+ // Split data into batches of batch_size and for each, call:
+ // metric->preprocess_base_points(data_batch, batch_size);
+ //
+ // TODO: This does not take into account the case for SSD inner product
+ // where the dimensions change after normalization.
+ DISKANN_DLLEXPORT virtual void preprocess_base_points(T *original_data, const size_t orig_dim,
+ const size_t num_points);
+
+ // Invokes normalization for a single vector during search. The scratch space
+ // has to be created by the caller keeping track of the fact that
+ // normalization might change the dimension of the query vector.
+ DISKANN_DLLEXPORT virtual void preprocess_query(const T *query_vec, const size_t query_dim, T *scratch_query);
+
+ // If an algorithm has a requirement that some data be aligned to a certain
+ // boundary it can use this function to indicate that requirement. Currently,
+ // we are setting it to 8 because that works well for AVX2. If we have AVX512
+ // implementations of distance algos, they might have to set this to 16
+ // (depending on how they are implemented)
+ DISKANN_DLLEXPORT virtual size_t get_required_alignment() const;
+
+ // Providing a default implementation for the virtual destructor because we
+ // don't expect most metric implementations to need it.
+ DISKANN_DLLEXPORT virtual ~Distance() = default;
+
+ protected:
+ diskann::Metric _distance_metric;
+ size_t _alignment_factor = 8;
+};
+
+class DistanceCosineInt8 : public Distance<int8_t>
+{
+ public:
+ DistanceCosineInt8() : Distance<int8_t>(diskann::Metric::COSINE)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const;
+};
+
+class DistanceL2Int8 : public Distance<int8_t>
+{
+ public:
+ DistanceL2Int8() : Distance<int8_t>(diskann::Metric::L2)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t size) const;
+};
+
+// AVX implementations. Borrowed from HNSW code.
+class AVXDistanceL2Int8 : public Distance<int8_t>
+{
+ public:
+ AVXDistanceL2Int8() : Distance<int8_t>(diskann::Metric::L2)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const;
+};
+
+class DistanceCosineFloat : public Distance<float>
+{
+ public:
+ DistanceCosineFloat() : Distance<float>(diskann::Metric::COSINE)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const;
+};
+
+class DistanceL2Float : public Distance<float>
+{
+ public:
+ DistanceL2Float() : Distance<float>(diskann::Metric::L2)
+ {
+ }
+
+#ifdef _WINDOWS
+ DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const;
+#else
+ DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const __attribute__((hot));
+#endif
+};
+
+class AVXDistanceL2Float : public Distance<float>
+{
+ public:
+ AVXDistanceL2Float() : Distance<float>(diskann::Metric::L2)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const;
+};
+
+template <typename T> class SlowDistanceL2 : public Distance<T>
+{
+ public:
+ SlowDistanceL2() : Distance<T>(diskann::Metric::L2)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const;
+};
+
+class SlowDistanceCosineUInt8 : public Distance<uint8_t>
+{
+ public:
+ SlowDistanceCosineUInt8() : Distance<uint8_t>(diskann::Metric::COSINE)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t length) const;
+};
+
+class DistanceL2UInt8 : public Distance<uint8_t>
+{
+ public:
+ DistanceL2UInt8() : Distance<uint8_t>(diskann::Metric::L2)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t size) const;
+};
+
+template <typename T> class DistanceInnerProduct : public Distance<T>
+{
+ public:
+ DistanceInnerProduct() : Distance<T>(diskann::Metric::INNER_PRODUCT)
+ {
+ }
+
+ DistanceInnerProduct(diskann::Metric metric) : Distance<T>(metric)
+ {
+ }
+ inline float inner_product(const T *a, const T *b, unsigned size) const;
+
+ inline float compare(const T *a, const T *b, unsigned size) const
+ {
+ float result = inner_product(a, b, size);
+ // if (result < 0)
+ // return std::numeric_limits<float>::max();
+ // else
+ return -result;
+ }
+};
+
+template <typename T> class DistanceFastL2 : public DistanceInnerProduct<T>
+{
+ // currently defined only for float.
+ // templated for future use.
+ public:
+ DistanceFastL2() : DistanceInnerProduct<T>(diskann::Metric::FAST_L2)
+ {
+ }
+ float norm(const T *a, unsigned size) const;
+ float compare(const T *a, const T *b, float norm, unsigned size) const;
+};
+
+class AVXDistanceInnerProductFloat : public Distance<float>
+{
+ public:
+ AVXDistanceInnerProductFloat() : Distance<float>(diskann::Metric::INNER_PRODUCT)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const;
+};
+
+class AVXNormalizedCosineDistanceFloat : public Distance<float>
+{
+ private:
+ AVXDistanceInnerProductFloat _innerProduct;
+
+ protected:
+ void normalize_and_copy(const float *a, uint32_t length, float *a_norm) const;
+
+ public:
+ AVXNormalizedCosineDistanceFloat() : Distance<float>(diskann::Metric::COSINE)
+ {
+ }
+ DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const
+ {
+ // Inner product returns negative values to indicate distance.
+ // This will ensure that cosine is between -1 and 1.
+ return 1.0f + _innerProduct.compare(a, b, length);
+ }
+ DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const override;
+
+ DISKANN_DLLEXPORT virtual bool preprocessing_required() const;
+
+ DISKANN_DLLEXPORT virtual void preprocess_base_points(float *original_data, const size_t orig_dim,
+ const size_t num_points) override;
+
+ DISKANN_DLLEXPORT virtual void preprocess_query(const float *query_vec, const size_t query_dim,
+ float *scratch_query_vector) override;
+};
+
+template <typename T> Distance<T> *get_distance_function(Metric m);
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/exceptions.h b/be/src/extern/diskann/include/exceptions.h
new file mode 100644
index 00000000000..99e4e736117
--- /dev/null
+++ b/be/src/extern/diskann/include/exceptions.h
@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <stdexcept>
+
+namespace diskann
+{
+
+class NotImplementedException : public std::logic_error
+{
+ public:
+ NotImplementedException() : std::logic_error("Function not yet implemented.")
+ {
+ }
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/filter_utils.h b/be/src/extern/diskann/include/filter_utils.h
new file mode 100644
index 00000000000..55f7aed28f1
--- /dev/null
+++ b/be/src/extern/diskann/include/filter_utils.h
@@ -0,0 +1,221 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <algorithm>
+#include <fcntl.h>
+#include <cassert>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <random>
+#include <set>
+#include <tuple>
+#include <string>
+#include <tsl/robin_map.h>
+#include <tsl/robin_set.h>
+#ifdef __APPLE__
+#else
+#include <malloc.h>
+#endif
+
+#ifdef _WINDOWS
+#include <Windows.h>
+typedef HANDLE FileHandle;
+#else
+#include <unistd.h>
+typedef int FileHandle;
+#endif
+
+#ifndef _WINDOWS
+#include <sys/uio.h>
+#endif
+
+#include "cached_io.h"
+#include "common_includes.h"
+#include "memory_mapper.h"
+#include "utils.h"
+#include "windows_customizations.h"
+
+// custom types (for readability)
+typedef tsl::robin_set<std::string> label_set;
+typedef std::string path;
+
+// structs for returning multiple items from a function
+typedef std::tuple<std::vector<label_set>, tsl::robin_map<std::string, uint32_t>, tsl::robin_set<std::string>>
+ parse_label_file_return_values;
+typedef std::tuple<std::vector<std::vector<uint32_t>>, uint64_t> load_label_index_return_values;
+
+namespace diskann
+{
+template <typename T>
+DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, label_set all_labels,
+ unsigned R, unsigned L, float alpha, unsigned num_threads);
+
+DISKANN_DLLEXPORT load_label_index_return_values load_label_index(path label_index_path,
+ uint32_t label_number_of_points);
+
+template <typename LabelT>
+DISKANN_DLLEXPORT std::tuple<std::vector<std::vector<LabelT>>, tsl::robin_set<LabelT>> parse_formatted_label_file(
+ path label_file);
+
+DISKANN_DLLEXPORT parse_label_file_return_values parse_label_file(path label_data_path, std::string universal_label);
+
+template <typename T>
+DISKANN_DLLEXPORT tsl::robin_map<std::string, std::vector<uint32_t>> generate_label_specific_vector_files_compat(
+ path input_data_path, tsl::robin_map<std::string, uint32_t> labels_to_number_of_points,
+ std::vector<label_set> point_ids_to_labels, label_set all_labels);
+
+/*
+ * For each label, generates a file containing all vectors that have said label.
+ * Also copies data from original bin file to new dimension-aligned file.
+ *
+ * Utilizes POSIX functions mmap and writev in order to minimize memory
+ * overhead, so we include an STL version as well.
+ *
+ * Each data file is saved under the following format:
+ * input_data_path + "_" + label
+ */
+#ifndef _WINDOWS
+template <typename T>
+inline tsl::robin_map<std::string, std::vector<uint32_t>> generate_label_specific_vector_files(
+ path input_data_path, tsl::robin_map<std::string, uint32_t> labels_to_number_of_points,
+ std::vector<label_set> point_ids_to_labels, label_set all_labels)
+{
+#ifndef _WINDOWS
+ auto file_writing_timer = std::chrono::high_resolution_clock::now();
+ diskann::MemoryMapper input_data(input_data_path);
+ char *input_start = input_data.getBuf();
+
+ uint32_t number_of_points, dimension;
+ std::memcpy(&number_of_points, input_start, sizeof(uint32_t));
+ std::memcpy(&dimension, input_start + sizeof(uint32_t), sizeof(uint32_t));
+ const uint32_t VECTOR_SIZE = dimension * sizeof(T);
+ const size_t METADATA = 2 * sizeof(uint32_t);
+ if (number_of_points != point_ids_to_labels.size())
+ {
+ std::cerr << "Error: number of points in labels file and data file differ." << std::endl;
+ throw;
+ }
+
+ tsl::robin_map<std::string, iovec *> label_to_iovec_map;
+ tsl::robin_map<std::string, uint32_t> label_to_curr_iovec;
+ tsl::robin_map<std::string, std::vector<uint32_t>> label_id_to_orig_id;
+
+ // setup iovec list for each label
+ for (const auto &lbl : all_labels)
+ {
+ iovec *label_iovecs = (iovec *)malloc(labels_to_number_of_points[lbl] * sizeof(iovec));
+ if (label_iovecs == nullptr)
+ {
+ throw;
+ }
+ label_to_iovec_map[lbl] = label_iovecs;
+ label_to_curr_iovec[lbl] = 0;
+ label_id_to_orig_id[lbl].reserve(labels_to_number_of_points[lbl]);
+ }
+
+ // each point added to corresponding per-label iovec list
+ for (uint32_t point_id = 0; point_id < number_of_points; point_id++)
+ {
+ char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id);
+ iovec curr_iovec;
+
+ curr_iovec.iov_base = curr_point;
+ curr_iovec.iov_len = VECTOR_SIZE;
+ for (const auto &lbl : point_ids_to_labels[point_id])
+ {
+ *(label_to_iovec_map[lbl] + label_to_curr_iovec[lbl]) = curr_iovec;
+ label_to_curr_iovec[lbl]++;
+ label_id_to_orig_id[lbl].push_back(point_id);
+ }
+ }
+
+ // write each label iovec to resp. file
+ for (const auto &lbl : all_labels)
+ {
+ int label_input_data_fd;
+ path curr_label_input_data_path(input_data_path + "_" + lbl);
+ uint32_t curr_num_pts = labels_to_number_of_points[lbl];
+
+ label_input_data_fd =
+ open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t)0644);
+ if (label_input_data_fd == -1)
+ throw;
+
+ // write metadata
+ uint32_t metadata[2] = {curr_num_pts, dimension};
+ int return_value = write(label_input_data_fd, metadata, sizeof(uint32_t) * 2);
+ if (return_value == -1)
+ {
+ throw;
+ }
+
+ // limits on number of iovec structs per writev means we need to perform
+ // multiple writevs
+ size_t i = 0;
+ while (curr_num_pts > IOV_MAX)
+ {
+ return_value = writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), IOV_MAX);
+ if (return_value == -1)
+ {
+ close(label_input_data_fd);
+ throw;
+ }
+ curr_num_pts -= IOV_MAX;
+ i += 1;
+ }
+ return_value = writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), curr_num_pts);
+ if (return_value == -1)
+ {
+ close(label_input_data_fd);
+ throw;
+ }
+
+ free(label_to_iovec_map[lbl]);
+ close(label_input_data_fd);
+ }
+
+ std::chrono::duration<double> file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer;
+ std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time "
+ << file_writing_time.count() << "\n"
+ << std::endl;
+
+ return label_id_to_orig_id;
+#endif
+}
+#endif
+
+inline std::vector<uint32_t> loadTags(const std::string &tags_file, const std::string &base_file)
+{
+ const bool tags_enabled = tags_file.empty() ? false : true;
+ std::vector<uint32_t> location_to_tag;
+ if (tags_enabled)
+ {
+ size_t tag_file_ndims, tag_file_npts;
+ std::uint32_t *tag_data;
+ diskann::load_bin<std::uint32_t>(tags_file, tag_data, tag_file_npts, tag_file_ndims);
+ if (tag_file_ndims != 1)
+ {
+ diskann::cerr << "tags file error" << std::endl;
+ throw diskann::ANNException("tag file error", -1, __FUNCSIG__, __FILE__, __LINE__);
+ }
+
+ // check if the point count match
+ size_t base_file_npts, base_file_ndims;
+ diskann::get_bin_metadata(base_file, base_file_npts, base_file_ndims);
+ if (base_file_npts != tag_file_npts)
+ {
+ diskann::cerr << "point num in tags file mismatch" << std::endl;
+ throw diskann::ANNException("point num in tags file mismatch", -1, __FUNCSIG__, __FILE__, __LINE__);
+ }
+
+ location_to_tag.assign(tag_data, tag_data + tag_file_npts);
+ delete[] tag_data;
+ }
+ return location_to_tag;
+}
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/in_mem_data_store.h b/be/src/extern/diskann/include/in_mem_data_store.h
new file mode 100644
index 00000000000..0a0a617daef
--- /dev/null
+++ b/be/src/extern/diskann/include/in_mem_data_store.h
@@ -0,0 +1,89 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+#pragma once
+
+#include <shared_mutex>
+#include <memory>
+
+#include "tsl/robin_map.h"
+#include "tsl/robin_set.h"
+#include "tsl/sparse_map.h"
+// #include "boost/dynamic_bitset.hpp"
+
+#include "abstract_data_store.h"
+
+#include "distance.h"
+#include "natural_number_map.h"
+#include "natural_number_set.h"
+#include "aligned_file_reader.h"
+
+namespace diskann
+{
+template <typename data_t> class InMemDataStore : public AbstractDataStore<data_t>
+{
+ public:
+ InMemDataStore(const location_t capacity, const size_t dim, std::unique_ptr<Distance<data_t>> distance_fn);
+ virtual ~InMemDataStore();
+
+ virtual location_t load(const std::string &filename) override;
+ virtual size_t save(const std::string &filename, const location_t num_points) override;
+
+ virtual size_t get_aligned_dim() const override;
+
+ // Populate internal data from unaligned data while doing alignment and any
+ // normalization that is required.
+ virtual void populate_data(const data_t *vectors, const location_t num_pts) override;
+ virtual void populate_data(const std::string &filename, const size_t offset) override;
+
+ virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) override;
+
+ virtual void get_vector(const location_t i, data_t *target) const override;
+ virtual void set_vector(const location_t i, const data_t *const vector) override;
+ virtual void prefetch_vector(const location_t loc) override;
+
+ virtual void move_vectors(const location_t old_location_start, const location_t new_location_start,
+ const location_t num_points) override;
+ virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) override;
+
+ virtual void preprocess_query(const data_t *query, AbstractScratch<data_t> *query_scratch) const override;
+
+ virtual float get_distance(const data_t *preprocessed_query, const location_t loc) const override;
+ virtual float get_distance(const location_t loc1, const location_t loc2) const override;
+
+ virtual void get_distance(const data_t *preprocessed_query, const location_t *locations,
+ const uint32_t location_count, float *distances,
+ AbstractScratch<data_t> *scratch) const override;
+ virtual void get_distance(const data_t *preprocessed_query, const std::vector<location_t> &ids,
+ std::vector<float> &distances, AbstractScratch<data_t> *scratch_space) const override;
+
+ virtual location_t calculate_medoid() const override;
+
+ virtual Distance<data_t> *get_dist_fn() const override;
+
+ virtual size_t get_alignment_factor() const override;
+
+ protected:
+ virtual location_t expand(const location_t new_size) override;
+ virtual location_t shrink(const location_t new_size) override;
+
+ virtual location_t load_impl(const std::string &filename);
+#ifdef EXEC_ENV_OLS
+ virtual location_t load_impl(AlignedFileReader &reader);
+#endif
+
+ private:
+ data_t *_data = nullptr;
+
+ size_t _aligned_dim;
+
+ // It may seem weird to put distance metric along with the data store class,
+ // but this gives us perf benefits as the datastore can do distance
+ // computations during search and compute norms of vectors internally without
+ // have to copy data back and forth.
+ std::unique_ptr<Distance<data_t>> _distance_fn;
+
+ // in case we need to save vector norms for optimization
+ std::shared_ptr<float[]> _pre_computed_norms;
+};
+
+} // namespace diskann
\ No newline at end of file
diff --git a/be/src/extern/diskann/include/in_mem_graph_store.h b/be/src/extern/diskann/include/in_mem_graph_store.h
new file mode 100644
index 00000000000..3952b17ebd8
--- /dev/null
+++ b/be/src/extern/diskann/include/in_mem_graph_store.h
@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include "abstract_graph_store.h"
+
+namespace diskann
+{
+
+class InMemGraphStore : public AbstractGraphStore
+{
+ public:
+ InMemGraphStore(const size_t total_pts, const size_t reserve_graph_degree);
+
+ // returns tuple of <nodes_read, start, num_frozen_points>
+ virtual std::tuple<uint32_t, uint32_t, size_t> load(const std::string &index_path_prefix,
+ const size_t num_points) override;
+ virtual int store(const std::string &index_path_prefix, const size_t num_points, const size_t num_frozen_points,
+ const uint32_t start) override;
+ virtual int store(std::stringstream &index_stream, const size_t num_points,
+ const size_t num_frozen_points, const uint32_t start) override;
+
+ virtual const std::vector<location_t> &get_neighbours(const location_t i) const override;
+ virtual void add_neighbour(const location_t i, location_t neighbour_id) override;
+ virtual void clear_neighbours(const location_t i) override;
+ virtual void swap_neighbours(const location_t a, location_t b) override;
+
+ virtual void set_neighbours(const location_t i, std::vector<location_t> &neighbors) override;
+
+ virtual size_t resize_graph(const size_t new_size) override;
+ virtual void clear_graph() override;
+
+ virtual size_t get_max_range_of_graph() override;
+ virtual uint32_t get_max_observed_degree() override;
+
+ protected:
+ virtual std::tuple<uint32_t, uint32_t, size_t> load_impl(const std::string &filename, size_t expected_num_points);
+#ifdef EXEC_ENV_OLS
+ virtual std::tuple<uint32_t, uint32_t, size_t> load_impl(AlignedFileReader &reader, size_t expected_num_points);
+#endif
+
+ int save_graph(const std::string &index_path_prefix, const size_t active_points, const size_t num_frozen_points,
+ const uint32_t start);
+
+ virtual int save_graph(std::stringstream &out, const size_t num_points,
+ const size_t num_frozen_points, const uint32_t start) override;
+
+ private:
+ size_t _max_range_of_graph = 0;
+ uint32_t _max_observed_degree = 0;
+
+ std::vector<std::vector<uint32_t>> _graph;
+};
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/index.h b/be/src/extern/diskann/include/index.h
new file mode 100644
index 00000000000..40f17d40a91
--- /dev/null
+++ b/be/src/extern/diskann/include/index.h
@@ -0,0 +1,442 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include "common_includes.h"
+
+#ifdef EXEC_ENV_OLS
+#include "aligned_file_reader.h"
+#endif
+
+#include "distance.h"
+#include "locking.h"
+#include "natural_number_map.h"
+#include "natural_number_set.h"
+#include "neighbor.h"
+#include "parameters.h"
+#include "utils.h"
+#include "windows_customizations.h"
+#include "scratch.h"
+#include "in_mem_data_store.h"
+#include "in_mem_graph_store.h"
+#include "abstract_index.h"
+
+#include "quantized_distance.h"
+#include "pq_data_store.h"
+
+#define OVERHEAD_FACTOR 1.1
+#define EXPAND_IF_FULL 0
+#define DEFAULT_MAXC 750
+
+namespace diskann
+{
+
+inline double estimate_ram_usage(size_t size, uint32_t dim, uint32_t datasize, uint32_t degree)
+{
+ double size_of_data = ((double)size) * ROUND_UP(dim, 8) * datasize;
+ double size_of_graph = ((double)size) * degree * sizeof(uint32_t) * defaults::GRAPH_SLACK_FACTOR;
+ double size_of_locks = ((double)size) * sizeof(non_recursive_mutex);
+ double size_of_outer_vector = ((double)size) * sizeof(ptrdiff_t);
+
+ return OVERHEAD_FACTOR * (size_of_data + size_of_graph + size_of_locks + size_of_outer_vector);
+}
+
+template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> class Index : public AbstractIndex
+{
+ /**************************************************************************
+ *
+ * Public functions acquire one or more of _update_lock, _consolidate_lock,
+ * _tag_lock, _delete_lock before calling protected functions which DO NOT
+ * acquire these locks. They might acquire locks on _locks[i]
+ *
+ **************************************************************************/
+
+ public:
+ // Constructor for Bulk operations and for creating the index object solely
+ // for loading a prexisting index.
+ DISKANN_DLLEXPORT Index(const IndexConfig &index_config, std::shared_ptr<AbstractDataStore<T>> data_store,
+ std::unique_ptr<AbstractGraphStore> graph_store,
+ std::shared_ptr<AbstractDataStore<T>> pq_data_store = nullptr);
+
+ // Constructor for incremental index
+ DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points,
+ const std::shared_ptr<IndexWriteParameters> index_parameters,
+ const std::shared_ptr<IndexSearchParams> index_search_params,
+ const size_t num_frozen_pts = 0, const bool dynamic_index = false,
+ const bool enable_tags = false, const bool concurrent_consolidate = false,
+ const bool pq_dist_build = false, const size_t num_pq_chunks = 0,
+ const bool use_opq = false, const bool filtered_index = false);
+
+ DISKANN_DLLEXPORT ~Index();
+
+ // Saves graph, data, metadata and associated tags.
+ DISKANN_DLLEXPORT void save(const char *filename, bool compact_before_save = false);
+ DISKANN_DLLEXPORT void save(std::stringstream &mem_index_stream, bool compact_before_save = false);
+
+ // Load functions
+#ifdef EXEC_ENV_OLS
+ DISKANN_DLLEXPORT void load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l);
+#else
+ // Reads the number of frozen points from graph's metadata file section.
+ DISKANN_DLLEXPORT static size_t get_graph_num_frozen_points(const std::string &graph_file);
+
+ DISKANN_DLLEXPORT void load(const char *index_file, uint32_t num_threads, uint32_t search_l);
+#endif
+
+ // get some private variables
+ DISKANN_DLLEXPORT size_t get_num_points();
+ DISKANN_DLLEXPORT size_t get_max_points();
+
+ DISKANN_DLLEXPORT bool detect_common_filters(uint32_t point_id, bool search_invocation,
+ const std::vector<LabelT> &incoming_labels);
+
+ // Batch build from a file. Optionally pass tags vector.
+ DISKANN_DLLEXPORT void build(const char *filename, const size_t num_points_to_load,
+ const std::vector<TagT> &tags = std::vector<TagT>());
+
+ // Batch build from a file. Optionally pass tags file.
+ DISKANN_DLLEXPORT void build(const char *filename, const size_t num_points_to_load, const char *tag_filename);
+
+ // Batch build from a data array, which must pad vectors to aligned_dim
+ DISKANN_DLLEXPORT void build(const T *data, const size_t num_points_to_load, const std::vector<TagT> &tags);
+
+ // Based on filter params builds a filtered or unfiltered index
+ DISKANN_DLLEXPORT void build(const std::string &data_file, const size_t num_points_to_load,
+ IndexFilterParams &filter_params);
+
+ // Filtered Support
+ DISKANN_DLLEXPORT void build_filtered_index(const char *filename, const std::string &label_file,
+ const size_t num_points_to_load,
+ const std::vector<TagT> &tags = std::vector<TagT>());
+
+ DISKANN_DLLEXPORT void set_universal_label(const LabelT &label);
+
+ // Get converted integer label from string to int map (_label_map)
+ DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &raw_label);
+
+ // Set starting point of an index before inserting any points incrementally.
+ // The data count should be equal to _num_frozen_pts * _aligned_dim.
+ DISKANN_DLLEXPORT void set_start_points(const T *data, size_t data_count);
+ // Set starting points to random points on a sphere of certain radius.
+ // A fixed random seed can be specified for scenarios where it's important
+ // to have higher consistency between index builds.
+ DISKANN_DLLEXPORT void set_start_points_at_random(T radius, uint32_t random_seed = 0);
+
+ // For FastL2 search on a static index, we interleave the data with graph
+ DISKANN_DLLEXPORT void optimize_index_layout();
+
+ // For FastL2 search on optimized layout
+ DISKANN_DLLEXPORT void search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices);
+
+ // Added search overload that takes L as parameter, so that we
+ // can customize L on a per-query basis without tampering with "Parameters"
+ template <typename IDType>
+ DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> search(const T *query, const size_t K, const uint32_t L,
+ IDType *indices, float *distances = nullptr);
+
+ // Initialize space for res_vectors before calling.
+ DISKANN_DLLEXPORT size_t search_with_tags(const T *query, const uint64_t K, const uint32_t L, TagT *tags,
+ float *distances, std::vector<T *> &res_vectors, bool use_filters = false,
+ const std::string filter_label = "");
+
+ // Filter support search
+ template <typename IndexType>
+ DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> search_with_filters(const T *query, const LabelT &filter_label,
+ const size_t K, const uint32_t L,
+ IndexType *indices, float *distances);
+
+ // Will fail if tag already in the index or if tag=0.
+ DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag);
+
+ // Will fail if tag already in the index or if tag=0.
+ DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag, const std::vector<LabelT> &label);
+
+ // call this before issuing deletions to sets relevant flags
+ DISKANN_DLLEXPORT int enable_delete();
+
+ // Record deleted point now and restructure graph later. Return -1 if tag
+ // not found, 0 if OK.
+ DISKANN_DLLEXPORT int lazy_delete(const TagT &tag);
+
+ // Record deleted points now and restructure graph later. Add to failed_tags
+ // if tag not found.
+ DISKANN_DLLEXPORT void lazy_delete(const std::vector<TagT> &tags, std::vector<TagT> &failed_tags);
+
+ // Call after a series of lazy deletions
+ // Returns number of live points left after consolidation
+ // If _conc_consolidates is set in the ctor, then this call can be invoked
+ // alongside inserts and lazy deletes, else it acquires _update_lock
+ DISKANN_DLLEXPORT consolidation_report consolidate_deletes(const IndexWriteParameters ¶meters);
+
+ DISKANN_DLLEXPORT void prune_all_neighbors(const uint32_t max_degree, const uint32_t max_occlusion,
+ const float alpha);
+
+ DISKANN_DLLEXPORT bool is_index_saved();
+
+ // repositions frozen points to the end of _data - if they have been moved
+ // during deletion
+ DISKANN_DLLEXPORT void reposition_frozen_point_to_end();
+ DISKANN_DLLEXPORT void reposition_points(uint32_t old_location_start, uint32_t new_location_start,
+ uint32_t num_locations);
+
+ // DISKANN_DLLEXPORT void save_index_as_one_file(bool flag);
+
+ DISKANN_DLLEXPORT void get_active_tags(tsl::robin_set<TagT> &active_tags);
+
+ // memory should be allocated for vec before calling this function
+ DISKANN_DLLEXPORT int get_vector_by_tag(TagT &tag, T *vec);
+
+ DISKANN_DLLEXPORT void print_status();
+
+ DISKANN_DLLEXPORT void count_nodes_at_bfs_levels();
+
+ // This variable MUST be updated if the number of entries in the metadata
+ // change.
+ DISKANN_DLLEXPORT static const int METADATA_ROWS = 5;
+
+ // ********************************
+ //
+ // Internals of the library
+ //
+ // ********************************
+
+ protected:
+ // overload of abstract index virtual methods
+ virtual void _build(const DataType &data, const size_t num_points_to_load, TagVector &tags) override;
+
+ virtual std::pair<uint32_t, uint32_t> _search(const DataType &query, const size_t K, const uint32_t L,
+ std::any &indices, float *distances = nullptr) override;
+ virtual std::pair<uint32_t, uint32_t> _search_with_filters(const DataType &query,
+ const std::string &filter_label_raw, const size_t K,
+ const uint32_t L, std::any &indices,
+ float *distances) override;
+
+ virtual int _insert_point(const DataType &data_point, const TagType tag) override;
+ virtual int _insert_point(const DataType &data_point, const TagType tag, Labelvector &labels) override;
+
+ virtual int _lazy_delete(const TagType &tag) override;
+
+ virtual void _lazy_delete(TagVector &tags, TagVector &failed_tags) override;
+
+ virtual void _get_active_tags(TagRobinSet &active_tags) override;
+
+ virtual void _set_start_points_at_random(DataType radius, uint32_t random_seed = 0) override;
+
+ virtual int _get_vector_by_tag(TagType &tag, DataType &vec) override;
+
+ virtual void _search_with_optimized_layout(const DataType &query, size_t K, size_t L, uint32_t *indices) override;
+
+ virtual size_t _search_with_tags(const DataType &query, const uint64_t K, const uint32_t L, const TagType &tags,
+ float *distances, DataVector &res_vectors, bool use_filters = false,
+ const std::string filter_label = "") override;
+
+ virtual void _set_universal_label(const LabelType universal_label) override;
+
+ // No copy/assign.
+ Index(const Index<T, TagT, LabelT> &) = delete;
+ Index<T, TagT, LabelT> &operator=(const Index<T, TagT, LabelT> &) = delete;
+
+ // Use after _data and _nd have been populated
+ // Acquire exclusive _update_lock before calling
+ void build_with_data_populated(const std::vector<TagT> &tags);
+
+ // generates 1 frozen point that will never be deleted from the graph
+ // This is not visible to the user
+ void generate_frozen_point();
+
+ // determines navigating node of the graph by calculating medoid of datafopt
+ uint32_t calculate_entry_point();
+
+ void parse_label_file(const std::string &label_file, size_t &num_pts_labels);
+
+ std::unordered_map<std::string, LabelT> load_label_map(const std::string &map_file);
+
+ // Returns the locations of start point and frozen points suitable for use
+ // with iterate_to_fixed_point.
+ std::vector<uint32_t> get_init_ids();
+
+ // The query to use is placed in scratch->aligned_query
+ std::pair<uint32_t, uint32_t> iterate_to_fixed_point(InMemQueryScratch<T> *scratch, const uint32_t Lindex,
+ const std::vector<uint32_t> &init_ids, bool use_filter,
+ const std::vector<LabelT> &filters, bool search_invocation);
+
+ void search_for_point_and_prune(int location, uint32_t Lindex, std::vector<uint32_t> &pruned_list,
+ InMemQueryScratch<T> *scratch, bool use_filter = false,
+ uint32_t filteredLindex = 0);
+
+ void prune_neighbors(const uint32_t location, std::vector<Neighbor> &pool, std::vector<uint32_t> &pruned_list,
+ InMemQueryScratch<T> *scratch);
+
+ void prune_neighbors(const uint32_t location, std::vector<Neighbor> &pool, const uint32_t range,
+ const uint32_t max_candidate_size, const float alpha, std::vector<uint32_t> &pruned_list,
+ InMemQueryScratch<T> *scratch);
+
+ // Prunes candidates in @pool to a shorter list @result
+ // @pool must be sorted before calling
+ void occlude_list(const uint32_t location, std::vector<Neighbor> &pool, const float alpha, const uint32_t degree,
+ const uint32_t maxc, std::vector<uint32_t> &result, InMemQueryScratch<T> *scratch,
+ const tsl::robin_set<uint32_t> *const delete_set_ptr = nullptr);
+
+ // add reverse links from all the visited nodes to node n.
+ void inter_insert(uint32_t n, std::vector<uint32_t> &pruned_list, const uint32_t range,
+ InMemQueryScratch<T> *scratch);
+
+ void inter_insert(uint32_t n, std::vector<uint32_t> &pruned_list, InMemQueryScratch<T> *scratch);
+
+ // Acquire exclusive _update_lock before calling
+ void link();
+
+ // Acquire exclusive _tag_lock and _delete_lock before calling
+ int reserve_location();
+
+ // Acquire exclusive _tag_lock before calling
+ size_t release_location(int location);
+ size_t release_locations(const tsl::robin_set<uint32_t> &locations);
+
+ // Resize the index when no slots are left for insertion.
+ // Acquire exclusive _update_lock and _tag_lock before calling.
+ void resize(size_t new_max_points);
+
+ // Acquire unique lock on _update_lock, _consolidate_lock, _tag_lock
+ // and _delete_lock before calling these functions.
+ // Renumber nodes, update tag and location maps and compact the
+ // graph, mode = _consolidated_order in case of lazy deletion and
+ // _compacted_order in case of eager deletion
+ DISKANN_DLLEXPORT void compact_data();
+ DISKANN_DLLEXPORT void compact_frozen_point();
+
+ // Remove deleted nodes from adjacency list of node loc
+ // Replace removed neighbors with second order neighbors.
+ // Also acquires _locks[i] for i = loc and out-neighbors of loc.
+ void process_delete(const tsl::robin_set<uint32_t> &old_delete_set, size_t loc, const uint32_t range,
+ const uint32_t maxc, const float alpha, InMemQueryScratch<T> *scratch);
+
+ void initialize_query_scratch(uint32_t num_threads, uint32_t search_l, uint32_t indexing_l, uint32_t r,
+ uint32_t maxc, size_t dim);
+
+ // Do not call without acquiring appropriate locks
+ // call public member functions save and load to invoke these.
+ DISKANN_DLLEXPORT size_t save_graph(std::string filename);
+ DISKANN_DLLEXPORT size_t save_graph(std::stringstream &index_stream);
+ DISKANN_DLLEXPORT size_t save_data(std::string filename);
+ DISKANN_DLLEXPORT size_t save_tags(std::string filename);
+ DISKANN_DLLEXPORT size_t save_delete_list(const std::string &filename);
+ DISKANN_DLLEXPORT size_t load_graph(const std::string filename, size_t expected_num_points);
+ DISKANN_DLLEXPORT size_t load_data(std::string filename0);
+ DISKANN_DLLEXPORT size_t load_tags(const std::string tag_file_name);
+ DISKANN_DLLEXPORT size_t load_delete_set(const std::string &filename);
+ private:
+ // Distance functions
+ Metric _dist_metric = diskann::L2;
+
+ // Data
+ std::shared_ptr<AbstractDataStore<T>> _data_store;
+
+ // Graph related data structures
+ std::unique_ptr<AbstractGraphStore> _graph_store;
+
+ char *_opt_graph = nullptr;
+
+ // Dimensions
+ size_t _dim = 0;
+ size_t _nd = 0; // number of active points i.e. existing in the graph
+ size_t _max_points = 0; // total number of points in given data set
+
+ // _num_frozen_pts is the number of points which are used as initial
+ // candidates when iterating to closest point(s). These are not visible
+ // externally and won't be returned by search. At least 1 frozen point is
+ // needed for a dynamic index. The frozen points have consecutive locations.
+ // See also _start below.
+ size_t _num_frozen_pts = 0;
+ size_t _frozen_pts_used = 0;
+ size_t _node_size;
+ size_t _data_len;
+ size_t _neighbor_len;
+
+ // Start point of the search. When _num_frozen_pts is greater than zero,
+ // this is the location of the first frozen point. Otherwise, this is a
+ // location of one of the points in index.
+ uint32_t _start = 0;
+
+ bool _has_built = false;
+ bool _saturate_graph = false;
+ bool _save_as_one_file = false; // plan to support in next version
+ bool _dynamic_index = false;
+ bool _enable_tags = false;
+ bool _normalize_vecs = false; // Using normalied L2 for cosine.
+ bool _deletes_enabled = false;
+
+ // Filter Support
+
+ bool _filtered_index = false;
+ // Location to label is only updated during insert_point(), all other reads are protected by
+ // default as a location can only be released at end of consolidate deletes
+ std::vector<std::vector<LabelT>> _location_to_labels;
+ tsl::robin_set<LabelT> _labels;
+ std::string _labels_file;
+ std::unordered_map<LabelT, uint32_t> _label_to_start_id;
+ std::unordered_map<uint32_t, uint32_t> _medoid_counts;
+
+ bool _use_universal_label = false;
+ LabelT _universal_label = 0;
+ uint32_t _filterIndexingQueueSize;
+ std::unordered_map<std::string, LabelT> _label_map;
+
+ // Indexing parameters
+ uint32_t _indexingQueueSize;
+ uint32_t _indexingRange;
+ uint32_t _indexingMaxC;
+ float _indexingAlpha;
+ uint32_t _indexingThreads;
+
+ // Query scratch data structures
+ ConcurrentQueue<InMemQueryScratch<T> *> _query_scratch;
+
+ // Flags for PQ based distance calculation
+ bool _pq_dist = false;
+ bool _use_opq = false;
+ size_t _num_pq_chunks = 0;
+ // REFACTOR
+ // uint8_t *_pq_data = nullptr;
+ std::shared_ptr<QuantizedDistance<T>> _pq_distance_fn = nullptr;
+ std::shared_ptr<AbstractDataStore<T>> _pq_data_store = nullptr;
+ bool _pq_generated = false;
+ FixedChunkPQTable _pq_table;
+
+ //
+ // Data structures, locks and flags for dynamic indexing and tags
+ //
+
+ // lazy_delete removes entry from _location_to_tag and _tag_to_location. If
+ // _location_to_tag does not resolve a location, infer that it was deleted.
+ tsl::sparse_map<TagT, uint32_t> _tag_to_location;
+ natural_number_map<uint32_t, TagT> _location_to_tag;
+
+ // _empty_slots has unallocated slots and those freed by consolidate_delete.
+ // _delete_set has locations marked deleted by lazy_delete. Will not be
+ // immediately available for insert. consolidate_delete will release these
+ // slots to _empty_slots.
+ natural_number_set<uint32_t> _empty_slots;
+ std::unique_ptr<tsl::robin_set<uint32_t>> _delete_set;
+
+ bool _data_compacted = true; // true if data has been compacted
+ bool _is_saved = false; // Checking if the index is already saved.
+ bool _conc_consolidate = false; // use _lock while searching
+
+ // Acquire locks in the order below when acquiring multiple locks
+ std::shared_timed_mutex // RW mutex between save/load (exclusive lock) and
+ _update_lock; // search/inserts/deletes/consolidate (shared lock)
+ std::shared_timed_mutex // Ensure only one consolidate or compact_data is
+ _consolidate_lock; // ever active
+ std::shared_timed_mutex // RW lock for _tag_to_location,
+ _tag_lock; // _location_to_tag, _empty_slots, _nd, _max_points, _label_to_start_id
+ std::shared_timed_mutex // RW Lock on _delete_set and _data_compacted
+ _delete_lock; // variable
+
+ // Per node lock, cardinality=_max_points + _num_frozen_points
+ std::vector<non_recursive_mutex> _locks;
+
+ static const float INDEX_GROWTH_FACTOR;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/index_build_params.h b/be/src/extern/diskann/include/index_build_params.h
new file mode 100644
index 00000000000..d4f454830ac
--- /dev/null
+++ b/be/src/extern/diskann/include/index_build_params.h
@@ -0,0 +1,73 @@
+#pragma once
+
+#include "common_includes.h"
+#include "parameters.h"
+
+namespace diskann
+{
+struct IndexFilterParams
+{
+ public:
+ std::string save_path_prefix;
+ std::string label_file;
+ std::string tags_file;
+ std::string universal_label;
+ uint32_t filter_threshold = 0;
+
+ private:
+ IndexFilterParams(const std::string &save_path_prefix, const std::string &label_file,
+ const std::string &universal_label, uint32_t filter_threshold)
+ : save_path_prefix(save_path_prefix), label_file(label_file), universal_label(universal_label),
+ filter_threshold(filter_threshold)
+ {
+ }
+
+ friend class IndexFilterParamsBuilder;
+};
+class IndexFilterParamsBuilder
+{
+ public:
+ IndexFilterParamsBuilder() = default;
+
+ IndexFilterParamsBuilder &with_save_path_prefix(const std::string &save_path_prefix)
+ {
+ if (save_path_prefix.empty() || save_path_prefix == "")
+ throw ANNException("Error: save_path_prefix can't be empty", -1);
+ this->_save_path_prefix = save_path_prefix;
+ return *this;
+ }
+
+ IndexFilterParamsBuilder &with_label_file(const std::string &label_file)
+ {
+ this->_label_file = label_file;
+ return *this;
+ }
+
+ IndexFilterParamsBuilder &with_universal_label(const std::string &univeral_label)
+ {
+ this->_universal_label = univeral_label;
+ return *this;
+ }
+
+ IndexFilterParamsBuilder &with_filter_threshold(const std::uint32_t &filter_threshold)
+ {
+ this->_filter_threshold = filter_threshold;
+ return *this;
+ }
+
+ IndexFilterParams build()
+ {
+ return IndexFilterParams(_save_path_prefix, _label_file, _universal_label, _filter_threshold);
+ }
+
+ IndexFilterParamsBuilder(const IndexFilterParamsBuilder &) = delete;
+ IndexFilterParamsBuilder &operator=(const IndexFilterParamsBuilder &) = delete;
+
+ private:
+ std::string _save_path_prefix;
+ std::string _label_file;
+ std::string _tags_file;
+ std::string _universal_label;
+ uint32_t _filter_threshold = 0;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/index_config.h b/be/src/extern/diskann/include/index_config.h
new file mode 100644
index 00000000000..a8e64d08cec
--- /dev/null
+++ b/be/src/extern/diskann/include/index_config.h
@@ -0,0 +1,256 @@
+#pragma once
+
+#include "common_includes.h"
+#include "parameters.h"
+
+namespace diskann
+{
+enum class DataStoreStrategy
+{
+ MEMORY
+};
+
+enum class GraphStoreStrategy
+{
+ MEMORY
+};
+
+struct IndexConfig
+{
+ DataStoreStrategy data_strategy;
+ GraphStoreStrategy graph_strategy;
+
+ Metric metric;
+ size_t dimension;
+ size_t max_points;
+
+ bool dynamic_index;
+ bool enable_tags;
+ bool pq_dist_build;
+ bool concurrent_consolidate;
+ bool use_opq;
+ bool filtered_index;
+
+ size_t num_pq_chunks;
+ size_t num_frozen_pts;
+
+ std::string label_type;
+ std::string tag_type;
+ std::string data_type;
+
+ // Params for building index
+ std::shared_ptr<IndexWriteParameters> index_write_params;
+ // Params for searching index
+ std::shared_ptr<IndexSearchParams> index_search_params;
+
+ private:
+ IndexConfig(DataStoreStrategy data_strategy, GraphStoreStrategy graph_strategy, Metric metric, size_t dimension,
+ size_t max_points, size_t num_pq_chunks, size_t num_frozen_points, bool dynamic_index, bool enable_tags,
+ bool pq_dist_build, bool concurrent_consolidate, bool use_opq, bool filtered_index,
+ std::string &data_type, const std::string &tag_type, const std::string &label_type,
+ std::shared_ptr<IndexWriteParameters> index_write_params,
+ std::shared_ptr<IndexSearchParams> index_search_params)
+ : data_strategy(data_strategy), graph_strategy(graph_strategy), metric(metric), dimension(dimension),
+ max_points(max_points), dynamic_index(dynamic_index), enable_tags(enable_tags), pq_dist_build(pq_dist_build),
+ concurrent_consolidate(concurrent_consolidate), use_opq(use_opq), filtered_index(filtered_index),
+ num_pq_chunks(num_pq_chunks), num_frozen_pts(num_frozen_points), label_type(label_type), tag_type(tag_type),
+ data_type(data_type), index_write_params(index_write_params), index_search_params(index_search_params)
+ {
+ }
+
+ friend class IndexConfigBuilder;
+};
+
+class IndexConfigBuilder
+{
+ public:
+ IndexConfigBuilder() = default;
+
+ IndexConfigBuilder &with_metric(Metric m)
+ {
+ this->_metric = m;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_graph_load_store_strategy(GraphStoreStrategy graph_strategy)
+ {
+ this->_graph_strategy = graph_strategy;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_data_load_store_strategy(DataStoreStrategy data_strategy)
+ {
+ this->_data_strategy = data_strategy;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_dimension(size_t dimension)
+ {
+ this->_dimension = dimension;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_max_points(size_t max_points)
+ {
+ this->_max_points = max_points;
+ return *this;
+ }
+
+ IndexConfigBuilder &is_dynamic_index(bool dynamic_index)
+ {
+ this->_dynamic_index = dynamic_index;
+ return *this;
+ }
+
+ IndexConfigBuilder &is_enable_tags(bool enable_tags)
+ {
+ this->_enable_tags = enable_tags;
+ return *this;
+ }
+
+ IndexConfigBuilder &is_pq_dist_build(bool pq_dist_build)
+ {
+ this->_pq_dist_build = pq_dist_build;
+ return *this;
+ }
+
+ IndexConfigBuilder &is_concurrent_consolidate(bool concurrent_consolidate)
+ {
+ this->_concurrent_consolidate = concurrent_consolidate;
+ return *this;
+ }
+
+ IndexConfigBuilder &is_use_opq(bool use_opq)
+ {
+ this->_use_opq = use_opq;
+ return *this;
+ }
+
+ IndexConfigBuilder &is_filtered(bool is_filtered)
+ {
+ this->_filtered_index = is_filtered;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_num_pq_chunks(size_t num_pq_chunks)
+ {
+ this->_num_pq_chunks = num_pq_chunks;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_num_frozen_pts(size_t num_frozen_pts)
+ {
+ this->_num_frozen_pts = num_frozen_pts;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_label_type(const std::string &label_type)
+ {
+ this->_label_type = label_type;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_tag_type(const std::string &tag_type)
+ {
+ this->_tag_type = tag_type;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_data_type(const std::string &data_type)
+ {
+ this->_data_type = data_type;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_index_write_params(IndexWriteParameters &index_write_params)
+ {
+ this->_index_write_params = std::make_shared<IndexWriteParameters>(index_write_params);
+ return *this;
+ }
+
+ IndexConfigBuilder &with_index_write_params(std::shared_ptr<IndexWriteParameters> index_write_params_ptr)
+ {
+ if (index_write_params_ptr == nullptr)
+ {
+ diskann::cout << "Passed, empty build_params while creating index config" << std::endl;
+ return *this;
+ }
+ this->_index_write_params = index_write_params_ptr;
+ return *this;
+ }
+
+ IndexConfigBuilder &with_index_search_params(IndexSearchParams &search_params)
+ {
+ this->_index_search_params = std::make_shared<IndexSearchParams>(search_params);
+ return *this;
+ }
+
+ IndexConfigBuilder &with_index_search_params(std::shared_ptr<IndexSearchParams> search_params_ptr)
+ {
+ if (search_params_ptr == nullptr)
+ {
+ diskann::cout << "Passed, empty search_params while creating index config" << std::endl;
+ return *this;
+ }
+ this->_index_search_params = search_params_ptr;
+ return *this;
+ }
+
+ IndexConfig build()
+ {
+ if (_data_type == "" || _data_type.empty())
+ throw ANNException("Error: data_type can not be empty", -1);
+
+ if (_dynamic_index && _num_frozen_pts == 0)
+ {
+ _num_frozen_pts = 1;
+ }
+
+ if (_dynamic_index)
+ {
+ if (_index_search_params != nullptr && _index_search_params->initial_search_list_size == 0)
+ throw ANNException("Error: please pass initial_search_list_size for building dynamic index.", -1);
+ }
+
+ // sanity check
+ if (_dynamic_index && _num_frozen_pts == 0)
+ {
+ diskann::cout << "_num_frozen_pts passed as 0 for dynamic_index. Setting it to 1 for safety." << std::endl;
+ _num_frozen_pts = 1;
+ }
+
+ return IndexConfig(_data_strategy, _graph_strategy, _metric, _dimension, _max_points, _num_pq_chunks,
+ _num_frozen_pts, _dynamic_index, _enable_tags, _pq_dist_build, _concurrent_consolidate,
+ _use_opq, _filtered_index, _data_type, _tag_type, _label_type, _index_write_params,
+ _index_search_params);
+ }
+
+ IndexConfigBuilder(const IndexConfigBuilder &) = delete;
+ IndexConfigBuilder &operator=(const IndexConfigBuilder &) = delete;
+
+ private:
+ DataStoreStrategy _data_strategy;
+ GraphStoreStrategy _graph_strategy;
+
+ Metric _metric;
+ size_t _dimension;
+ size_t _max_points;
+
+ bool _dynamic_index = false;
+ bool _enable_tags = false;
+ bool _pq_dist_build = false;
+ bool _concurrent_consolidate = false;
+ bool _use_opq = false;
+ bool _filtered_index{defaults::HAS_LABELS};
+
+ size_t _num_pq_chunks = 0;
+ size_t _num_frozen_pts{defaults::NUM_FROZEN_POINTS_STATIC};
+
+ std::string _label_type{"uint32"};
+ std::string _tag_type{"uint32"};
+ std::string _data_type;
+
+ std::shared_ptr<IndexWriteParameters> _index_write_params;
+ std::shared_ptr<IndexSearchParams> _index_search_params;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/index_factory.h b/be/src/extern/diskann/include/index_factory.h
new file mode 100644
index 00000000000..76fb0b978eb
--- /dev/null
+++ b/be/src/extern/diskann/include/index_factory.h
@@ -0,0 +1,51 @@
+#pragma once
+
+#include "index.h"
+#include "abstract_graph_store.h"
+#include "in_mem_graph_store.h"
+#include "pq_data_store.h"
+
+namespace diskann
+{
+class IndexFactory
+{
+ public:
+ DISKANN_DLLEXPORT explicit IndexFactory(const IndexConfig &config);
+ DISKANN_DLLEXPORT std::unique_ptr<AbstractIndex> create_instance();
+
+ DISKANN_DLLEXPORT static std::unique_ptr<AbstractGraphStore> construct_graphstore(
+ const GraphStoreStrategy stratagy, const size_t size, const size_t reserve_graph_degree);
+
+ template <typename T>
+ DISKANN_DLLEXPORT static std::shared_ptr<AbstractDataStore<T>> construct_datastore(DataStoreStrategy stratagy,
+ size_t num_points,
+ size_t dimension, Metric m);
+ // For now PQDataStore incorporates within itself all variants of quantization that we support. In the
+ // future it may be necessary to introduce an AbstractPQDataStore class to spearate various quantization
+ // flavours.
+ template <typename T>
+ DISKANN_DLLEXPORT static std::shared_ptr<PQDataStore<T>> construct_pq_datastore(DataStoreStrategy strategy,
+ size_t num_points, size_t dimension,
+ Metric m, size_t num_pq_chunks,
+ bool use_opq);
+ template <typename T> static Distance<T> *construct_inmem_distance_fn(Metric m);
+
+ private:
+ void check_config();
+
+ template <typename data_type, typename tag_type, typename label_type>
+ std::unique_ptr<AbstractIndex> create_instance();
+
+ std::unique_ptr<AbstractIndex> create_instance(const std::string &data_type, const std::string &tag_type,
+ const std::string &label_type);
+
+ template <typename data_type>
+ std::unique_ptr<AbstractIndex> create_instance(const std::string &tag_type, const std::string &label_type);
+
+ template <typename data_type, typename tag_type>
+ std::unique_ptr<AbstractIndex> create_instance(const std::string &label_type);
+
+ std::unique_ptr<IndexConfig> _config;
+};
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/io_reader.h b/be/src/extern/diskann/include/io_reader.h
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/be/src/extern/diskann/include/linux_aligned_file_reader.h b/be/src/extern/diskann/include/linux_aligned_file_reader.h
new file mode 100644
index 00000000000..7620e31942c
--- /dev/null
+++ b/be/src/extern/diskann/include/linux_aligned_file_reader.h
@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#ifndef _WINDOWS
+
+#include "aligned_file_reader.h"
+
+class LinuxAlignedFileReader : public AlignedFileReader
+{
+ private:
+ uint64_t file_sz;
+ FileHandle file_desc;
+ io_context_t bad_ctx = (io_context_t)-1;
+
+ public:
+ LinuxAlignedFileReader();
+ ~LinuxAlignedFileReader();
+
+ IOContext &get_ctx();
+
+ // register thread-id for a context
+ void register_thread();
+
+ // de-register thread-id for a context
+ void deregister_thread();
+ void deregister_all_threads();
+
+ // Open & close ops
+ // Blocking calls
+ void open(const std::string &fname);
+ void close();
+
+ // process batch of aligned requests in parallel
+ // NOTE :: blocking call
+ void read(std::vector<AlignedRead> &read_reqs, IOContext &ctx, bool async = false);
+};
+
+#endif
diff --git a/be/src/extern/diskann/include/locking.h b/be/src/extern/diskann/include/locking.h
new file mode 100644
index 00000000000..890c24a2b70
--- /dev/null
+++ b/be/src/extern/diskann/include/locking.h
@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+#pragma once
+
+#include <mutex>
+
+#ifdef _WINDOWS
+#include "windows_slim_lock.h"
+#endif
+
+namespace diskann
+{
+#ifdef _WINDOWS
+using non_recursive_mutex = windows_exclusive_slim_lock;
+using LockGuard = windows_exclusive_slim_lock_guard;
+#else
+using non_recursive_mutex = std::mutex;
+using LockGuard = std::lock_guard<non_recursive_mutex>;
+#endif
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/logger.h b/be/src/extern/diskann/include/logger.h
new file mode 100644
index 00000000000..0b17807dbf4
--- /dev/null
+++ b/be/src/extern/diskann/include/logger.h
@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+#pragma once
+
+#include <functional>
+#include <iostream>
+#include "windows_customizations.h"
+
+#ifdef EXEC_ENV_OLS
+#ifndef ENABLE_CUSTOM_LOGGER
+#define ENABLE_CUSTOM_LOGGER
+#endif // !ENABLE_CUSTOM_LOGGER
+#endif // EXEC_ENV_OLS
+
+namespace diskann
+{
+#ifdef ENABLE_CUSTOM_LOGGER
+DISKANN_DLLEXPORT extern std::basic_ostream<char> cout;
+DISKANN_DLLEXPORT extern std::basic_ostream<char> cerr;
+#else
+using std::cerr;
+using std::cout;
+#endif
+
+enum class DISKANN_DLLEXPORT LogLevel
+{
+ LL_Info = 0,
+ LL_Error,
+ LL_Count
+};
+
+#ifdef ENABLE_CUSTOM_LOGGER
+DISKANN_DLLEXPORT void SetCustomLogger(std::function<void(LogLevel, const char *)> logger);
+#endif
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/logger_impl.h b/be/src/extern/diskann/include/logger_impl.h
new file mode 100644
index 00000000000..03c65e0ced0
--- /dev/null
+++ b/be/src/extern/diskann/include/logger_impl.h
@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <sstream>
+#include <mutex>
+
+#include "ann_exception.h"
+#include "logger.h"
+
+namespace diskann
+{
+#ifdef ENABLE_CUSTOM_LOGGER
+class ANNStreamBuf : public std::basic_streambuf<char>
+{
+ public:
+ DISKANN_DLLEXPORT explicit ANNStreamBuf(FILE *fp);
+ DISKANN_DLLEXPORT ~ANNStreamBuf();
+
+ DISKANN_DLLEXPORT bool is_open() const
+ {
+ return true; // because stdout and stderr are always open.
+ }
+ DISKANN_DLLEXPORT void close();
+ DISKANN_DLLEXPORT virtual int underflow();
+ DISKANN_DLLEXPORT virtual int overflow(int c);
+ DISKANN_DLLEXPORT virtual int sync();
+
+ private:
+ FILE *_fp;
+ char *_buf;
+ int _bufIndex;
+ std::mutex _mutex;
+ LogLevel _logLevel;
+
+ int flush();
+ void logImpl(char *str, int numchars);
+
+ // Why the two buffer-sizes? If we are running normally, we are basically
+ // interacting with a character output system, so we short-circuit the
+ // output process by keeping an empty buffer and writing each character
+ // to stdout/stderr. But if we are running in OLS, we have to take all
+ // the text that is written to diskann::cout/diskann:cerr, consolidate it
+ // and push it out in one-shot, because the OLS infra does not give us
+ // character based output. Therefore, we use a larger buffer that is large
+ // enough to store the longest message, and continuously add characters
+ // to it. When the calling code outputs a std::endl or std::flush, sync()
+ // will be called and will output a log level, component name, and the text
+ // that has been collected. (sync() is also called if the buffer is full, so
+ // overflows/missing text are not a concern).
+ // This implies calling code _must_ either print std::endl or std::flush
+ // to ensure that the message is written immediately.
+
+ static const int BUFFER_SIZE = 1024;
+
+ ANNStreamBuf(const ANNStreamBuf &);
+ ANNStreamBuf &operator=(const ANNStreamBuf &);
+};
+#endif
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/math_utils.h b/be/src/extern/diskann/include/math_utils.h
new file mode 100644
index 00000000000..83d189f7077
--- /dev/null
+++ b/be/src/extern/diskann/include/math_utils.h
@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include "common_includes.h"
+#include "utils.h"
+
+namespace math_utils
+{
+
+float calc_distance(float *vec_1, float *vec_2, size_t dim);
+
+// compute l2-squared norms of data stored in row major num_points * dim,
+// needs
+// to be pre-allocated
+void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, const size_t dim);
+
+void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot_mat, float *&new_mat,
+ bool transpose_rot = false);
+
+// calculate closest center to data of num_points * dim (row major)
+// centers is num_centers * dim (row major)
+// data_l2sq has pre-computed squared norms of data
+// centers_l2sq has pre-computed squared norms of centers
+// pre-allocated center_index will contain id of k nearest centers
+// pre-allocated dist_matrix shound be num_points * num_centers and contain
+// squared distances
+
+// Ideally used only by compute_closest_centers
+void compute_closest_centers_in_block(const float *const data, const size_t num_points, const size_t dim,
+ const float *const centers, const size_t num_centers,
+ const float *const docs_l2sq, const float *const centers_l2sq,
+ uint32_t *center_index, float *const dist_matrix, size_t k = 1);
+
+// Given data in num_points * new_dim row major
+// Pivots stored in full_pivot_data as k * new_dim row major
+// Calculate the closest pivot for each point and store it in vector
+// closest_centers_ivf (which needs to be allocated outside)
+// Additionally, if inverted index is not null (and pre-allocated), it will
+// return inverted index for each center Additionally, if pts_norms_squared is
+// not null, then it will assume that point norms are pre-computed and use
+// those
+// values
+
+void compute_closest_centers(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers,
+ size_t k, uint32_t *closest_centers_ivf, std::vector<size_t> *inverted_index = NULL,
+ float *pts_norms_squared = NULL);
+
+// if to_subtract is 1, will subtract nearest center from each row. Else will
+// add. Output will be in data_load iself.
+// Nearest centers need to be provided in closst_centers.
+
+void process_residuals(float *data_load, size_t num_points, size_t dim, float *cur_pivot_data, size_t num_centers,
+ uint32_t *closest_centers, bool to_subtract);
+
+} // namespace math_utils
+
+namespace kmeans
+{
+
+// run Lloyds one iteration
+// Given data in row major num_points * dim, and centers in row major
+// num_centers * dim
+// And squared lengths of data points, output the closest center to each data
+// point, update centers, and also return inverted index.
+// If closest_centers == NULL, will allocate memory and return.
+// Similarly, if closest_docs == NULL, will allocate memory and return.
+
+float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, size_t num_centers, float *docs_l2sq,
+ std::vector<size_t> *closest_docs, uint32_t *&closest_center);
+
+// Run Lloyds until max_reps or stopping criterion
+// If you pass NULL for closest_docs and closest_center, it will NOT return
+// the results, else it will assume appriate allocation as closest_docs = new
+// vector<size_t> [num_centers], and closest_center = new size_t[num_points]
+// Final centers are output in centers as row major num_centers * dim
+//
+float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, const size_t num_centers,
+ const size_t max_reps, std::vector<size_t> *closest_docs, uint32_t *closest_center);
+
+// assumes already memory allocated for pivot_data as new
+// float[num_centers*dim] and select randomly num_centers points as pivots
+void selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers);
+
+void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers);
+} // namespace kmeans
diff --git a/be/src/extern/diskann/include/memory_mapper.h b/be/src/extern/diskann/include/memory_mapper.h
new file mode 100644
index 00000000000..75faca1bb73
--- /dev/null
+++ b/be/src/extern/diskann/include/memory_mapper.h
@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#ifndef _WINDOWS
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#else
+#include <Windows.h>
+#endif
+#include <string>
+
+namespace diskann
+{
+class MemoryMapper
+{
+ private:
+#ifndef _WINDOWS
+ int _fd;
+#else
+ HANDLE _bareFile;
+ HANDLE _fd;
+
+#endif
+ char *_buf;
+ size_t _fileSize;
+ const char *_fileName;
+
+ public:
+ MemoryMapper(const char *filename);
+ MemoryMapper(const std::string &filename);
+
+ char *getBuf();
+ size_t getFileSize();
+
+ ~MemoryMapper();
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/natural_number_map.h b/be/src/extern/diskann/include/natural_number_map.h
new file mode 100644
index 00000000000..e846882a847
--- /dev/null
+++ b/be/src/extern/diskann/include/natural_number_map.h
@@ -0,0 +1,86 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <memory>
+#include <type_traits>
+#include <vector>
+
+#include <boost/dynamic_bitset.hpp>
+
+namespace diskann
+{
+// A map whose key is a natural number (from 0 onwards) and maps to a value.
+// Made as both memory and performance efficient map for scenario such as
+// DiskANN location-to-tag map. There, the pool of numbers is consecutive from
+// zero to some max value, and it's expected that most if not all keys from 0
+// up to some current maximum will be present in the map. The memory usage of
+// the map is determined by the largest inserted key since it uses vector as a
+// backing store and bitset for presence indication.
+//
+// Thread-safety: this class is not thread-safe in general.
+// Exception: multiple read-only operations are safe on the object only if
+// there are no writers to it in parallel.
+template <typename Key, typename Value> class natural_number_map
+{
+ public:
+ static_assert(std::is_trivial<Key>::value, "Key must be a trivial type");
+
+ // Represents a reference to a element in the map. Used while iterating
+ // over map entries.
+ struct position
+ {
+ size_t _key;
+ // The number of keys that were enumerated when iterating through the
+ // map so far. Used to early-terminate enumeration when ithere are no
+ // more entries in the map.
+ size_t _keys_already_enumerated;
+
+ // Returns whether it's valid to access the element at this position in
+ // the map.
+ bool is_valid() const;
+ };
+
+ natural_number_map();
+
+ void reserve(size_t count);
+ size_t size() const;
+
+ void set(Key key, Value value);
+ void erase(Key key);
+
+ bool contains(Key key) const;
+ bool try_get(Key key, Value &value) const;
+
+ // Returns the value at the specified position. Prerequisite: position is
+ // valid.
+ Value get(const position &pos) const;
+
+ // Finds the first element in the map, if any. Invalidated by changes in the
+ // map.
+ position find_first() const;
+
+ // Finds the next element in the map after the specified position.
+ // Invalidated by changes in the map.
+ position find_next(const position &after_position) const;
+
+ void clear();
+
+ private:
+ // Number of entries in the map. Not the same as size() of the
+ // _values_vector below.
+ size_t _size;
+
+ // Array of values. The key is the index of the value.
+ std::vector<Value> _values_vector;
+
+ // Values that are in the set have the corresponding bit index set
+ // to 1.
+ //
+ // Use a pointer here to allow for forward declaration of dynamic_bitset
+ // in public headers to avoid making boost a dependency for clients
+ // of DiskANN.
+ std::unique_ptr<boost::dynamic_bitset<>> _values_bitset;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/natural_number_set.h b/be/src/extern/diskann/include/natural_number_set.h
new file mode 100644
index 00000000000..ec5b827e615
--- /dev/null
+++ b/be/src/extern/diskann/include/natural_number_set.h
@@ -0,0 +1,50 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <memory>
+#include <type_traits>
+
+#include "boost_dynamic_bitset_fwd.h"
+
+namespace diskann
+{
+// A set of natural numbers (from 0 onwards). Made for scenario where the
+// pool of numbers is consecutive from zero to some max value and very
+// efficient methods for "add to set", "get any value from set", "is in set"
+// are needed. The memory usage of the set is determined by the largest
+// number of inserted entries (uses a vector as a backing store) as well as
+// the largest value to be placed in it (uses bitset as well).
+//
+// Thread-safety: this class is not thread-safe in general.
+// Exception: multiple read-only operations (e.g. is_in_set, empty, size) are
+// safe on the object only if there are no writers to it in parallel.
+template <typename T> class natural_number_set
+{
+ public:
+ static_assert(std::is_trivial<T>::value, "Identifier must be a trivial type");
+
+ natural_number_set();
+
+ bool is_empty() const;
+ void reserve(size_t count);
+ void insert(T id);
+ T pop_any();
+ void clear();
+ size_t size() const;
+ bool is_in_set(T id) const;
+
+ private:
+ // Values that are currently in set.
+ std::vector<T> _values_vector;
+
+ // Values that are in the set have the corresponding bit index set
+ // to 1.
+ //
+ // Use a pointer here to allow for forward declaration of dynamic_bitset
+ // in public headers to avoid making boost a dependency for clients
+ // of DiskANN.
+ std::unique_ptr<boost::dynamic_bitset<>> _values_bitset;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/neighbor.h b/be/src/extern/diskann/include/neighbor.h
new file mode 100644
index 00000000000..c58e65070c3
--- /dev/null
+++ b/be/src/extern/diskann/include/neighbor.h
@@ -0,0 +1,160 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <cstddef>
+#include <mutex>
+#include <vector>
+#include "utils.h"
+
+namespace diskann
+{
+
+struct Neighbor
+{
+ unsigned id;
+ float distance;
+ bool expanded;
+
+ Neighbor() = default;
+
+ Neighbor(unsigned id, float distance) : id{id}, distance{distance}, expanded(false)
+ {
+ }
+
+ inline bool operator<(const Neighbor &other) const
+ {
+ return distance < other.distance || (distance == other.distance && id < other.id);
+ }
+
+ inline bool operator==(const Neighbor &other) const
+ {
+ return (id == other.id);
+ }
+};
+
+// Invariant: after every `insert` and `closest_unexpanded()`, `_cur` points to
+// the first Neighbor which is unexpanded.
+class NeighborPriorityQueue
+{
+ public:
+ NeighborPriorityQueue() : _size(0), _capacity(0), _cur(0), _pre(0)
+ {
+ }
+
+ explicit NeighborPriorityQueue(size_t capacity) : _size(0), _capacity(capacity), _cur(0), _data(capacity + 1)
+ {
+ }
+
+ // Inserts the item ordered into the set up to the sets capacity.
+ // The item will be dropped if it is the same id as an exiting
+ // set item or it has a greated distance than the final
+ // item in the set. The set cursor that is used to pop() the
+ // next item will be set to the lowest index of an uncheck item
+ void insert(const Neighbor &nbr)
+ {
+ if (_size == _capacity && _data[_size - 1] < nbr)
+ {
+ return;
+ }
+
+ size_t lo = 0, hi = _size;
+ while (lo < hi)
+ {
+ size_t mid = (lo + hi) >> 1;
+ if (nbr < _data[mid])
+ {
+ hi = mid;
+ // Make sure the same id isn't inserted into the set
+ }
+ else if (_data[mid].id == nbr.id)
+ {
+ return;
+ }
+ else
+ {
+ lo = mid + 1;
+ }
+ }
+
+ if (lo < _capacity)
+ {
+ std::memmove(&_data[lo + 1], &_data[lo], (_size - lo) * sizeof(Neighbor));
+ }
+ _data[lo] = {nbr.id, nbr.distance};
+ if (_size < _capacity)
+ {
+ _size++;
+ }
+ if (lo < _cur)
+ {
+ _cur = lo;
+ }
+ }
+
+ Neighbor closest_unexpanded()
+ {
+ _data[_cur].expanded = true;
+ _pre = _cur;
+ while (_cur < _size && _data[_cur].expanded)
+ {
+ _cur++;
+ }
+ return _data[_pre];
+ }
+
+ void remove_pre_expanded_node(){
+ int len = (_size - _pre - 1) * sizeof(Neighbor);
+ if (len > 0)
+ std::memmove(&_data[_pre], &_data[_pre + 1], len);
+ _cur--;
+ _size--;
+ }
+
+ bool has_unexpanded_node() const
+ {
+ return _cur < _size;
+ }
+
+ size_t size() const
+ {
+ return _size;
+ }
+
+ size_t capacity() const
+ {
+ return _capacity;
+ }
+
+ void reserve(size_t capacity)
+ {
+ if (capacity + 1 > _data.size())
+ {
+ _data.resize(capacity + 1);
+ }
+ _capacity = capacity;
+ }
+
+ Neighbor &operator[](size_t i)
+ {
+ return _data[i];
+ }
+
+ Neighbor operator[](size_t i) const
+ {
+ return _data[i];
+ }
+
+ void clear()
+ {
+ _size = 0;
+ _cur = 0;
+ }
+
+ private:
+ size_t _size, _capacity, _cur, _pre;
+ std::vector<Neighbor> _data;
+};
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/parameters.h b/be/src/extern/diskann/include/parameters.h
new file mode 100644
index 00000000000..0206814bd27
--- /dev/null
+++ b/be/src/extern/diskann/include/parameters.h
@@ -0,0 +1,119 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <sstream>
+#include <typeinfo>
+#include <unordered_map>
+
+#include "omp.h"
+#include "defaults.h"
+
+namespace diskann
+{
+
+class IndexWriteParameters
+
+{
+ public:
+ const uint32_t search_list_size; // L
+ const uint32_t max_degree; // R
+ const bool saturate_graph;
+ const uint32_t max_occlusion_size; // C
+ const float alpha;
+ const uint32_t num_threads;
+ const uint32_t filter_list_size; // Lf
+
+ IndexWriteParameters(const uint32_t search_list_size, const uint32_t max_degree, const bool saturate_graph,
+ const uint32_t max_occlusion_size, const float alpha, const uint32_t num_threads,
+ const uint32_t filter_list_size)
+ : search_list_size(search_list_size), max_degree(max_degree), saturate_graph(saturate_graph),
+ max_occlusion_size(max_occlusion_size), alpha(alpha), num_threads(num_threads),
+ filter_list_size(filter_list_size)
+ {
+ }
+
+ friend class IndexWriteParametersBuilder;
+};
+
+class IndexSearchParams
+{
+ public:
+ IndexSearchParams(const uint32_t initial_search_list_size, const uint32_t num_search_threads)
+ : initial_search_list_size(initial_search_list_size), num_search_threads(num_search_threads)
+ {
+ }
+ const uint32_t initial_search_list_size; // search L
+ const uint32_t num_search_threads; // search threads
+};
+
+class IndexWriteParametersBuilder
+{
+ /**
+ * Fluent builder pattern to keep track of the 7 non-default properties
+ * and their order. The basic ctor was getting unwieldy.
+ */
+ public:
+ IndexWriteParametersBuilder(const uint32_t search_list_size, // L
+ const uint32_t max_degree // R
+ )
+ : _search_list_size(search_list_size), _max_degree(max_degree)
+ {
+ }
+
+ IndexWriteParametersBuilder &with_max_occlusion_size(const uint32_t max_occlusion_size)
+ {
+ _max_occlusion_size = max_occlusion_size;
+ return *this;
+ }
+
+ IndexWriteParametersBuilder &with_saturate_graph(const bool saturate_graph)
+ {
+ _saturate_graph = saturate_graph;
+ return *this;
+ }
+
+ IndexWriteParametersBuilder &with_alpha(const float alpha)
+ {
+ _alpha = alpha;
+ return *this;
+ }
+
+ IndexWriteParametersBuilder &with_num_threads(const uint32_t num_threads)
+ {
+ _num_threads = num_threads == 0 ? omp_get_num_procs() : num_threads;
+ return *this;
+ }
+
+ IndexWriteParametersBuilder &with_filter_list_size(const uint32_t filter_list_size)
+ {
+ _filter_list_size = filter_list_size == 0 ? _search_list_size : filter_list_size;
+ return *this;
+ }
+
+ IndexWriteParameters build() const
+ {
+ return IndexWriteParameters(_search_list_size, _max_degree, _saturate_graph, _max_occlusion_size, _alpha,
+ _num_threads, _filter_list_size);
+ }
+
+ IndexWriteParametersBuilder(const IndexWriteParameters &wp)
+ : _search_list_size(wp.search_list_size), _max_degree(wp.max_degree),
+ _max_occlusion_size(wp.max_occlusion_size), _saturate_graph(wp.saturate_graph), _alpha(wp.alpha),
+ _filter_list_size(wp.filter_list_size)
+ {
+ }
+ IndexWriteParametersBuilder(const IndexWriteParametersBuilder &) = delete;
+ IndexWriteParametersBuilder &operator=(const IndexWriteParametersBuilder &) = delete;
+
+ private:
+ uint32_t _search_list_size{};
+ uint32_t _max_degree{};
+ uint32_t _max_occlusion_size{defaults::MAX_OCCLUSION_SIZE};
+ bool _saturate_graph{defaults::SATURATE_GRAPH};
+ float _alpha{defaults::ALPHA};
+ uint32_t _num_threads{defaults::NUM_THREADS};
+ uint32_t _filter_list_size{defaults::FILTER_LIST_SIZE};
+};
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/partition.h b/be/src/extern/diskann/include/partition.h
new file mode 100644
index 00000000000..ccba399474d
--- /dev/null
+++ b/be/src/extern/diskann/include/partition.h
@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <cassert>
+#include <sstream>
+#include <stack>
+#include <string>
+#include <unordered_map>
+
+#include "neighbor.h"
+#include "parameters.h"
+#include "tsl/robin_set.h"
+#include "utils.h"
+
+#include "windows_customizations.h"
+
+template <typename T>
+void gen_random_slice(const std::string base_file, const std::string output_prefix, double sampling_rate);
+
+template <typename T>
+void gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, size_t &slice_size,
+ size_t &ndims);
+
+template <typename T>
+void gen_random_slice(std::stringstream &_data_stream, double p_val, float *&sampled_data, size_t &slice_size,
+ size_t &ndims);
+
+template <typename T>
+void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data,
+ size_t &slice_size);
+
+int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivots, const size_t num_centers,
+ const size_t dim, const size_t k_base, std::vector<size_t> &cluster_sizes);
+
+template <typename T>
+int shard_data_into_clusters(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim,
+ const size_t k_base, std::string prefix_path);
+
+template <typename T>
+int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots, const size_t num_centers,
+ const size_t dim, const size_t k_base, std::string prefix_path);
+
+template <typename T>
+int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_filename, std::string data_filename);
+
+template <typename T>
+int partition(const std::string data_file, const float sampling_rate, size_t num_centers, size_t max_k_means_reps,
+ const std::string prefix_path, size_t k_base);
+
+template <typename T>
+int partition_with_ram_budget(const std::string data_file, const double sampling_rate, double ram_budget,
+ size_t graph_degree, const std::string prefix_path, size_t k_base);
diff --git a/be/src/extern/diskann/include/percentile_stats.h b/be/src/extern/diskann/include/percentile_stats.h
new file mode 100644
index 00000000000..344b091ce92
--- /dev/null
+++ b/be/src/extern/diskann/include/percentile_stats.h
@@ -0,0 +1,82 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <fstream>
+#include <functional>
+#ifdef _WINDOWS
+#include <numeric>
+#endif
+#include <string>
+#include <vector>
+
+#include "distance.h"
+#include "parameters.h"
+
+namespace diskann
+{
+struct QueryStats
+{
+ float total_us = 0; // total time to process query in micros
+ float io_us = 0; // total time spent in IO
+ float cpu_us = 0; // total time spent in CPU
+
+ unsigned n_4k = 0; // # of 4kB reads
+ unsigned n_8k = 0; // # of 8kB reads
+ unsigned n_12k = 0; // # of 12kB reads
+ unsigned n_ios = 0; // total # of IOs issued
+ unsigned read_size = 0; // total # of bytes read
+ unsigned n_cmps_saved = 0; // # cmps saved
+ unsigned n_cmps = 0; // # cmps
+ unsigned n_cache_hits = 0; // # cache_hits
+ unsigned n_hops = 0; // # search hops
+
+ std::string to_string(){
+ std::stringstream ss;
+ ss << "total_us:" << total_us
+ << " io_us:" << io_us
+ << " cpu_us:" << cpu_us
+ << " n_4k:" << n_4k
+ << " n_8k:" << n_8k
+ << " n_12k:" << n_12k
+ << " n_ios:" << n_ios
+ << " read_size:" << read_size
+ << " n_cmps_saved:" << n_cmps_saved
+ << " n_cmps:" << n_cmps
+ << " n_cache_hits:" << n_cache_hits
+ << " n_hops:" << n_hops;
+ return ss.str();
+ }
+};
+
+template <typename T>
+inline T get_percentile_stats(QueryStats *stats, uint64_t len, float percentile,
+ const std::function<T(const QueryStats &)> &member_fn)
+{
+ std::vector<T> vals(len);
+ for (uint64_t i = 0; i < len; i++)
+ {
+ vals[i] = member_fn(stats[i]);
+ }
+
+ std::sort(vals.begin(), vals.end(), [](const T &left, const T &right) { return left < right; });
+
+ auto retval = vals[(uint64_t)(percentile * len)];
+ vals.clear();
+ return retval;
+}
+
+template <typename T>
+inline double get_mean_stats(QueryStats *stats, uint64_t len, const std::function<T(const QueryStats &)> &member_fn)
+{
+ double avg = 0;
+ for (uint64_t i = 0; i < len; i++)
+ {
+ avg += (double)member_fn(stats[i]);
+ }
+ return avg / len;
+}
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/pq.h b/be/src/extern/diskann/include/pq.h
new file mode 100644
index 00000000000..a65b9049a6a
--- /dev/null
+++ b/be/src/extern/diskann/include/pq.h
@@ -0,0 +1,117 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include "utils.h"
+#include "pq_common.h"
+#include "vector/stream_wrapper.h"
+
+namespace diskann
+{
+class FixedChunkPQTable
+{
+ float *tables = nullptr; // pq_tables = float array of size [256 * ndims]
+ uint64_t ndims = 0; // ndims = true dimension of vectors
+ uint64_t n_chunks = 0;
+ bool use_rotation = false;
+ uint32_t *chunk_offsets = nullptr;
+ float *centroid = nullptr;
+ float *tables_tr = nullptr; // same as pq_tables, but col-major
+ float *rotmat_tr = nullptr;
+
+ public:
+ FixedChunkPQTable();
+
+ virtual ~FixedChunkPQTable();
+
+#ifdef EXEC_ENV_OLS
+ void load_pq_centroid_bin(MemoryMappedFiles &files, const char *pq_table_file, size_t num_chunks);
+#else
+ void load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks);
+ void load_pq_centroid_bin(std::shared_ptr<Reader> reader, size_t num_chunks);
+ void load_pq_centroid_bin(IReaderWrapperSPtr reader, size_t num_chunks);
+#endif
+
+ uint32_t get_num_chunks();
+
+ void preprocess_query(float *query_vec);
+
+ // assumes pre-processed query
+ void populate_chunk_distances(const float *query_vec, float *dist_vec);
+
+ float l2_distance(const float *query_vec, uint8_t *base_vec);
+
+ float inner_product(const float *query_vec, uint8_t *base_vec);
+
+ // assumes no rotation is involved
+ void inflate_vector(uint8_t *base_vec, float *out_vec);
+
+ void populate_chunk_inner_products(const float *query_vec, float *dist_vec);
+};
+
+void aggregate_coords(const std::vector<unsigned> &ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out);
+
+void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists,
+ std::vector<float> &dists_out);
+
+// Need to replace calls to these with calls to vector& based functions above
+void aggregate_coords(const unsigned *ids, const uint64_t n_ids, const uint8_t *all_coords, const uint64_t ndims,
+ uint8_t *out);
+
+void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists,
+ float *dists_out);
+
+DISKANN_DLLEXPORT int generate_pq_pivots(const float *const train_data, size_t num_train, unsigned dim,
+ unsigned num_centers, unsigned num_pq_chunks, unsigned max_k_means_reps,
+ std::string pq_pivots_path, bool make_zero_mean = false);
+
+DISKANN_DLLEXPORT int generate_pq_pivots(const float *const train_data, size_t num_train, unsigned dim,
+ unsigned num_centers, unsigned num_pq_chunks, unsigned max_k_means_reps,
+ std::stringstream &pq_pivots_stream, bool make_zero_mean = false);
+
+DISKANN_DLLEXPORT int generate_opq_pivots(const float *train_data, size_t num_train, unsigned dim, unsigned num_centers,
+ unsigned num_pq_chunks, std::string opq_pivots_path,
+ bool make_zero_mean = false);
+
+DISKANN_DLLEXPORT int generate_opq_pivots(const float *train_data, size_t num_train, unsigned dim, unsigned num_centers,
+ unsigned num_pq_chunks, std::stringstream &opq_pivots_stream,
+ bool make_zero_mean = false);
+
+DISKANN_DLLEXPORT int generate_pq_pivots_simplified(const float *train_data, size_t num_train, size_t dim,
+ size_t num_pq_chunks, std::vector<float> &pivot_data_vector);
+
+template <typename T>
+int generate_pq_data_from_pivots(const std::string &data_file, unsigned num_centers, unsigned num_pq_chunks,
+ const std::string &pq_pivots_path, const std::string &pq_compressed_vectors_path,
+ bool use_opq = false);
+
+template <typename T>
+int generate_pq_data_from_pivots(std::stringstream &data_stream, unsigned num_centers, unsigned num_pq_chunks,
+ std::stringstream &pq_pivots_stream, std::stringstream &pq_compressed_stream,
+ bool use_opq = false);
+
+DISKANN_DLLEXPORT int generate_pq_data_from_pivots_simplified(const float *data, const size_t num,
+ const float *pivot_data, const size_t pivots_num,
+ const size_t dim, const size_t num_pq_chunks,
+ std::vector<uint8_t> &pq);
+
+template <typename T>
+void generate_disk_quantized_data(const std::string &data_file_to_use, const std::string &disk_pq_pivots_path,
+ const std::string &disk_pq_compressed_vectors_path,
+ const diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims);
+
+template <typename T>
+void generate_quantized_data(const std::string &data_file_to_use, const std::string &pq_pivots_path,
+ const std::string &pq_compressed_vectors_path, const diskann::Metric compareMetric,
+ const double p_val, const uint64_t num_pq_chunks, const bool use_opq,
+ const std::string &codebook_prefix = "");
+
+template <typename T>
+void generate_quantized_data(std::stringstream & _data_stream, std::stringstream &pq_pivots_stream,
+ std::stringstream &pq_compressed_stream, const diskann::Metric compareMetric,
+ const double p_val, const uint64_t num_pq_chunks, const bool use_opq,
+ const std::string &codebook_prefix = "");
+
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/pq_common.h b/be/src/extern/diskann/include/pq_common.h
new file mode 100644
index 00000000000..c6a3a573983
--- /dev/null
+++ b/be/src/extern/diskann/include/pq_common.h
@@ -0,0 +1,30 @@
+#pragma once
+
+#include <string>
+#include <sstream>
+
+#define NUM_PQ_BITS 8
+#define NUM_PQ_CENTROIDS (1 << NUM_PQ_BITS)
+#define MAX_OPQ_ITERS 20
+#define NUM_KMEANS_REPS_PQ 12
+#define MAX_PQ_TRAINING_SET_SIZE 256000
+#define MAX_PQ_CHUNKS 512
+
+namespace diskann
+{
+inline std::string get_quantized_vectors_filename(const std::string &prefix, bool use_opq, uint32_t num_chunks)
+{
+ return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_compressed.bin";
+}
+
+inline std::string get_pivot_data_filename(const std::string &prefix, bool use_opq, uint32_t num_chunks)
+{
+ return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_pivots.bin";
+}
+
+inline std::string get_rotation_matrix_suffix(const std::string &pivot_data_filename)
+{
+ return pivot_data_filename + "_rotation_matrix.bin";
+}
+
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/pq_data_store.h b/be/src/extern/diskann/include/pq_data_store.h
new file mode 100644
index 00000000000..7c0cb5fe0a6
--- /dev/null
+++ b/be/src/extern/diskann/include/pq_data_store.h
@@ -0,0 +1,97 @@
+#pragma once
+#include <memory>
+#include "distance.h"
+#include "quantized_distance.h"
+#include "pq.h"
+#include "abstract_data_store.h"
+
+namespace diskann
+{
+// REFACTOR TODO: By default, the PQDataStore is an in-memory datastore because both Vamana and
+// DiskANN treat it the same way. But with DiskPQ, that may need to change.
+template <typename data_t> class PQDataStore : public AbstractDataStore<data_t>
+{
+
+ public:
+ PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::unique_ptr<Distance<data_t>> distance_fn,
+ std::unique_ptr<QuantizedDistance<data_t>> pq_distance_fn);
+ PQDataStore(const PQDataStore &) = delete;
+ PQDataStore &operator=(const PQDataStore &) = delete;
+ ~PQDataStore();
+
+ // Load quantized vectors from a set of files. Here filename is treated
+ // as a prefix and the files are assumed to be named with DiskANN
+ // conventions.
+ virtual location_t load(const std::string &file_prefix) override;
+
+ // Save quantized vectors to a set of files whose names start with
+ // file_prefix.
+ // Currently, the plan is to save the quantized vectors to the quantized
+ // vectors file.
+ virtual size_t save(const std::string &file_prefix, const location_t num_points) override;
+
+ // Since base class function is pure virtual, we need to declare it here, even though alignent concept is not needed
+ // for Quantized data stores.
+ virtual size_t get_aligned_dim() const override;
+
+ // Populate quantized data from unaligned data using PQ functionality
+ virtual void populate_data(const data_t *vectors, const location_t num_pts) override;
+ virtual void populate_data(const std::string &filename, const size_t offset) override;
+
+ virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) override;
+
+ virtual void get_vector(const location_t i, data_t *target) const override;
+ virtual void set_vector(const location_t i, const data_t *const vector) override;
+ virtual void prefetch_vector(const location_t loc) override;
+
+ virtual void move_vectors(const location_t old_location_start, const location_t new_location_start,
+ const location_t num_points) override;
+ virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) override;
+
+ virtual void preprocess_query(const data_t *query, AbstractScratch<data_t> *scratch) const override;
+
+ virtual float get_distance(const data_t *query, const location_t loc) const override;
+ virtual float get_distance(const location_t loc1, const location_t loc2) const override;
+
+ // NOTE: Caller must invoke "PQDistance->preprocess_query" ONCE before calling
+ // this function.
+ virtual void get_distance(const data_t *preprocessed_query, const location_t *locations,
+ const uint32_t location_count, float *distances,
+ AbstractScratch<data_t> *scratch_space) const override;
+
+ // NOTE: Caller must invoke "PQDistance->preprocess_query" ONCE before calling
+ // this function.
+ virtual void get_distance(const data_t *preprocessed_query, const std::vector<location_t> &ids,
+ std::vector<float> &distances, AbstractScratch<data_t> *scratch_space) const override;
+
+ // We are returning the distance function that is used for full precision
+ // vectors here, not the PQ distance function. This is because the callers
+ // all are expecting a Distance<T> not QuantizedDistance<T>.
+ virtual Distance<data_t> *get_dist_fn() const override;
+
+ virtual location_t calculate_medoid() const override;
+
+ virtual size_t get_alignment_factor() const override;
+
+ protected:
+ virtual location_t expand(const location_t new_size) override;
+ virtual location_t shrink(const location_t new_size) override;
+
+ virtual location_t load_impl(const std::string &filename);
+#ifdef EXEC_ENV_OLS
+ virtual location_t load_impl(AlignedFileReader &reader);
+#endif
+
+ private:
+ uint8_t *_quantized_data = nullptr;
+ size_t _num_chunks = 0;
+
+ // REFACTOR TODO: Doing this temporarily before refactoring OPQ into
+ // its own class. Remove later.
+ bool _use_opq = false;
+
+ Metric _distance_metric;
+ std::unique_ptr<Distance<data_t>> _distance_fn = nullptr;
+ std::unique_ptr<QuantizedDistance<data_t>> _pq_distance_fn = nullptr;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/pq_flash_index.h b/be/src/extern/diskann/include/pq_flash_index.h
new file mode 100644
index 00000000000..e5550a33c75
--- /dev/null
+++ b/be/src/extern/diskann/include/pq_flash_index.h
@@ -0,0 +1,286 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include "common_includes.h"
+
+#include "aligned_file_reader.h"
+#include "concurrent_queue.h"
+#include "neighbor.h"
+#include "parameters.h"
+#include "percentile_stats.h"
+#include "pq.h"
+#include "utils.h"
+#include "windows_customizations.h"
+#include "scratch.h"
+#include "tsl/robin_map.h"
+#include "tsl/robin_set.h"
+#include "ThreadPool.h"
+
+#include "vector/stream_wrapper.h"
+
+
+#define FULL_PRECISION_REORDER_MULTIPLIER 3
+
+namespace diskann
+{
+
+struct Filter {
+ virtual bool is_member(uint32_t idx) {
+ return false;
+ }
+ virtual ~Filter() = default;
+};
+
+class BatchReader{
+ public:
+ BatchReader(uint32_t threads){
+ _pool = std::make_shared<ThreadPool>(threads);
+ }
+ void read(std::vector<AlignedRead> requests){
+ if (_reader) {
+ std::vector<std::future<void>> futures;
+ for (const auto& req : requests) {
+ auto future = _pool->enqueue([&, req]() {
+ _reader->read((char*)req.buf, req.len, req.offset);
+ });
+ futures.push_back(std::move(future));
+ }
+ for (auto& fut : futures) {
+ fut.get();
+ }
+ } else {
+ throw std::runtime_error("Reader not initialized properly.");
+ }
+ }
+
+ void set_reader(IReaderWrapperSPtr reader){
+ _reader = reader;
+ }
+
+ private:
+ std::shared_ptr<ThreadPool> _pool;
+ IReaderWrapperSPtr _reader;
+};
+
+template <typename T, typename LabelT = uint32_t> class PQFlashIndex
+{
+ public:
+ DISKANN_DLLEXPORT PQFlashIndex(IReaderWrapperSPtr reader,
+ diskann::Metric metric = diskann::Metric::L2);
+ DISKANN_DLLEXPORT PQFlashIndex(std::shared_ptr<AlignedFileReader> &fileReader,
+ diskann::Metric metric = diskann::Metric::L2);
+ DISKANN_DLLEXPORT ~PQFlashIndex();
+
+ DISKANN_DLLEXPORT int load(uint32_t num_threads,
+ IReaderWrapperSPtr pq_pivots_reader,
+ IReaderWrapperSPtr pq_compressed_reader,
+ IReaderWrapperSPtr vamana_index_reader,
+ IReaderWrapperSPtr disk_layout_reader,
+ IReaderWrapperSPtr tag_reader);
+
+#ifdef EXEC_ENV_OLS
+ DISKANN_DLLEXPORT int load(diskann::MemoryMappedFiles &files, uint32_t num_threads, const char *index_prefix);
+#else
+ // load compressed data, and obtains the handle to the disk-resident index
+ DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *index_prefix);
+#endif
+
+#ifdef EXEC_ENV_OLS
+ DISKANN_DLLEXPORT int load_from_separate_paths(diskann::MemoryMappedFiles &files, uint32_t num_threads,
+ const char *index_filepath, const char *pivots_filepath,
+ const char *compressed_filepath);
+#else
+ DISKANN_DLLEXPORT int load_from_separate_paths(uint32_t num_threads, const char *index_filepath,
+ const char *pivots_filepath, const char *compressed_filepath);
+ DISKANN_DLLEXPORT int load_from_compound_file(uint32_t num_threads, const char *compound_filepath);
+#endif
+
+ DISKANN_DLLEXPORT void load_cache_list(std::vector<uint32_t> &node_list);
+
+#ifdef EXEC_ENV_OLS
+ DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(MemoryMappedFiles &files, std::string sample_bin,
+ uint64_t l_search, uint64_t beamwidth,
+ uint64_t num_nodes_to_cache, uint32_t nthreads,
+ std::vector<uint32_t> &node_list);
+#else
+ DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search,
+ uint64_t beamwidth, uint64_t num_nodes_to_cache,
+ uint32_t num_threads,
+ std::vector<uint32_t> &node_list);
+#endif
+
+ DISKANN_DLLEXPORT void cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector<uint32_t> &node_list,
+ const bool shuffle = false);
+
+ DISKANN_DLLEXPORT uint32_t cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search,
+ uint64_t *res_ids, float *res_dists, const uint64_t beam_width,
+ Filter *filter = nullptr,
+ QueryStats *stats = nullptr);
+
+ DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &filter_label);
+
+ DISKANN_DLLEXPORT uint32_t range_search(const T *query1, const double range, const uint64_t min_l_search,
+ const uint64_t max_l_search, std::vector<uint64_t> &indices,
+ std::vector<float> &distances, const uint64_t min_beam_width,
+ QueryStats *stats = nullptr);
+
+ DISKANN_DLLEXPORT uint64_t get_data_dim();
+ std::shared_ptr<AlignedFileReader> reader;
+ IReaderWrapperSPtr customReader;
+
+ DISKANN_DLLEXPORT diskann::Metric get_metric();
+
+ //
+ // node_ids: input list of node_ids to be read
+ // coord_buffers: pointers to pre-allocated buffers that coords need to copied to. If null, dont copy.
+ // nbr_buffers: pre-allocated buffers to copy neighbors into
+ //
+ // returns a vector of bool one for each node_id: true if read is success, else false
+ //
+ DISKANN_DLLEXPORT std::vector<bool> read_nodes(const std::vector<uint32_t> &node_ids,
+ std::vector<T *> &coord_buffers,
+ std::vector<std::pair<uint32_t, uint32_t *>> &nbr_buffers);
+
+ DISKANN_DLLEXPORT std::vector<std::uint8_t> get_pq_vector(std::uint64_t vid);
+ DISKANN_DLLEXPORT uint64_t get_num_points();
+
+ protected:
+ DISKANN_DLLEXPORT void use_medoids_data_as_centroids();
+ DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, uint64_t visited_reserve = 4096);
+ DISKANN_DLLEXPORT void setup_thread_data_without_ctx(uint64_t nthreads, uint64_t visited_reserve = 4096);
+
+
+ DISKANN_DLLEXPORT void set_universal_label(const LabelT &label);
+
+ private:
+ DISKANN_DLLEXPORT inline bool point_has_label(uint32_t point_id, LabelT label_id);
+ std::unordered_map<std::string, LabelT> load_label_map(std::basic_istream<char> &infile);
+ DISKANN_DLLEXPORT void parse_label_file(std::basic_istream<char> &infile, size_t &num_pts_labels);
+ DISKANN_DLLEXPORT void get_label_file_metadata(const std::string &fileContent, uint32_t &num_pts,
+ uint32_t &num_total_labels);
+ DISKANN_DLLEXPORT void generate_random_labels(std::vector<LabelT> &labels, const uint32_t num_labels,
+ const uint32_t nthreads);
+ void reset_stream_for_reading(std::basic_istream<char> &infile);
+
+ // sector # on disk where node_id is present with in the graph part
+ DISKANN_DLLEXPORT uint64_t get_node_sector(uint64_t node_id);
+
+ // ptr to start of the node
+ DISKANN_DLLEXPORT char *offset_to_node(char *sector_buf, uint64_t node_id);
+
+ // returns region of `node_buf` containing [NNBRS][NBR_ID(uint32_t)]
+ DISKANN_DLLEXPORT uint32_t *offset_to_node_nhood(char *node_buf);
+
+ // returns region of `node_buf` containing [COORD(T)]
+ DISKANN_DLLEXPORT T *offset_to_node_coords(char *node_buf);
+
+ // index info for multi-node sectors
+ // nhood of node `i` is in sector: [i / nnodes_per_sector]
+ // offset in sector: [(i % nnodes_per_sector) * max_node_len]
+ //
+ // index info for multi-sector nodes
+ // nhood of node `i` is in sector: [i * DIV_ROUND_UP(_max_node_len, SECTOR_LEN)]
+ // offset in sector: [0]
+ //
+ // Common info
+ // coords start at ofsset
+ // #nbrs of node `i`: *(unsigned*) (offset + disk_bytes_per_point)
+ // nbrs of node `i` : (unsigned*) (offset + disk_bytes_per_point + 1)
+
+ uint64_t _max_node_len = 0;
+ uint64_t _nnodes_per_sector = 0; // 0 for multi-sector nodes, >0 for multi-node sectors
+ uint64_t _max_degree = 0;
+
+ // Data used for searching with re-order vectors
+ uint64_t _ndims_reorder_vecs = 0;
+ uint64_t _reorder_data_start_sector = 0;
+ uint64_t _nvecs_per_sector = 0;
+
+ diskann::Metric metric = diskann::Metric::L2;
+
+ // used only for inner product search to re-scale the result value
+ // (due to the pre-processing of base during index build)
+ float _max_base_norm = 0.0f;
+
+ // data info
+ uint64_t _num_points = 0;
+ uint64_t _num_frozen_points = 0;
+ uint64_t _frozen_location = 0;
+ uint64_t _data_dim = 0;
+ uint64_t _aligned_dim = 0;
+ uint64_t _disk_bytes_per_point = 0; // Number of bytes
+
+ std::string _disk_index_file;
+ std::vector<std::pair<uint32_t, uint32_t>> _node_visit_counter;
+
+ // PQ data
+ // _n_chunks = # of chunks ndims is split into
+ // data: char * _n_chunks
+ // chunk_size = chunk size of each dimension chunk
+ // pq_tables = float* [[2^8 * [chunk_size]] * _n_chunks]
+ uint8_t *data = nullptr;
+ uint64_t _n_chunks;
+ FixedChunkPQTable _pq_table;
+
+ // distance comparator
+ std::shared_ptr<Distance<T>> _dist_cmp;
+ std::shared_ptr<Distance<float>> _dist_cmp_float;
+
+ // for very large datasets: we use PQ even for the disk resident index
+ bool _use_disk_index_pq = false;
+ uint64_t _disk_pq_n_chunks = 0;
+ FixedChunkPQTable _disk_pq_table;
+
+ // medoid/start info
+
+ // graph has one entry point by default,
+ // we can optionally have multiple starting points
+ uint32_t *_medoids = nullptr;
+ // defaults to 1
+ size_t _num_medoids;
+ // by default, it is empty. If there are multiple
+ // centroids, we pick the medoid corresponding to the
+ // closest centroid as the starting point of search
+ float *_centroid_data = nullptr;
+
+ // nhood_cache; the uint32_t in nhood_Cache are offsets into nhood_cache_buf
+ unsigned *_nhood_cache_buf = nullptr;
+ tsl::robin_map<uint32_t, std::pair<uint32_t, uint32_t *>> _nhood_cache;
+
+ // coord_cache; The T* in coord_cache are offsets into coord_cache_buf
+ T *_coord_cache_buf = nullptr;
+ tsl::robin_map<uint32_t, T *> _coord_cache;
+
+ // thread-specific scratch
+ ConcurrentQueue<SSDThreadData<T> *> _thread_data;
+ uint64_t _max_nthreads;
+ bool _load_flag = false;
+ bool _count_visited_nodes = false;
+ bool _reorder_data_exists = false;
+ uint64_t _reoreder_data_offset = 0;
+
+ // filter support
+ uint32_t *_pts_to_label_offsets = nullptr;
+ uint32_t *_pts_to_label_counts = nullptr;
+ LabelT *_pts_to_labels = nullptr;
+ std::unordered_map<LabelT, std::vector<uint32_t>> _filter_to_medoid_ids;
+ bool _use_universal_label = false;
+ LabelT _universal_filter_label;
+ tsl::robin_set<uint32_t> _dummy_pts;
+ tsl::robin_set<uint32_t> _has_dummy_pts;
+ tsl::robin_map<uint32_t, uint32_t> _dummy_to_real_map;
+ tsl::robin_map<uint32_t, std::vector<uint32_t>> _real_to_dummy_map;
+ std::unordered_map<std::string, LabelT> _label_map;
+
+ std::shared_ptr<BatchReader> _batch_reader;
+
+#ifdef EXEC_ENV_OLS
+ // Set to a larger value than the actual header to accommodate
+ // any additions we make to the header. This is an outer limit
+ // on how big the header can be.
+ static const int HEADER_SIZE = defaults::SECTOR_LEN;
+ char *getHeaderBytes();
+#endif
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/pq_l2_distance.h b/be/src/extern/diskann/include/pq_l2_distance.h
new file mode 100644
index 00000000000..e6fc6e41bec
--- /dev/null
+++ b/be/src/extern/diskann/include/pq_l2_distance.h
@@ -0,0 +1,87 @@
+#pragma once
+#include "quantized_distance.h"
+
+namespace diskann
+{
+template <typename data_t> class PQL2Distance : public QuantizedDistance<data_t>
+{
+ public:
+ // REFACTOR TODO: We could take a file prefix here and load the
+ // PQ pivots file, so that the distance object is initialized
+ // immediately after construction. But this would not work well
+ // with our data store concept where the store is created first
+ // and data populated after.
+ // REFACTOR TODO: Ideally, we should only read the num_chunks from
+ // the pivots file. However, we read the pivots file only later, but
+ // clients can call functions like get_<xxx>_filename without calling
+ // load_pivot_data. Hence this. The TODO is whether we should check
+ // that the num_chunks from the file is the same as this one.
+
+ PQL2Distance(uint32_t num_chunks, bool use_opq = false);
+
+ virtual ~PQL2Distance() override;
+
+ virtual bool is_opq() const override;
+
+ virtual std::string get_quantized_vectors_filename(const std::string &prefix) const override;
+ virtual std::string get_pivot_data_filename(const std::string &prefix) const override;
+ virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const override;
+
+#ifdef EXEC_ENV_OLS
+ virtual void load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file,
+ size_t num_chunks) override;
+#else
+ virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) override;
+#endif
+
+ // Number of chunks in the PQ table. Depends on the compression level used.
+ // Has to be < ndim
+ virtual uint32_t get_num_chunks() const override;
+
+ // Preprocess the query by computing chunk distances from the query vector to
+ // various centroids. Since we don't want this class to do scratch management,
+ // we will take a PQScratch object which can come either from Index class or
+ // PQFlashIndex class.
+ virtual void preprocess_query(const data_t *aligned_query, uint32_t original_dim,
+ PQScratch<data_t> &pq_scratch) override;
+
+ // Distance function used for graph traversal. This function must be called
+ // after
+ // preprocess_query. The reason we do not call preprocess ourselves is because
+ // that function has to be called once per query, while this function is
+ // called at each iteration of the graph walk. NOTE: This function expects
+ // 1. the query to be preprocessed using preprocess_query()
+ // 2. the scratch object to contain the quantized vectors corresponding to ids
+ // in aligned_pq_coord_scratch. Done by calling aggregate_coords()
+ //
+ virtual void preprocessed_distance(PQScratch<data_t> &pq_scratch, const uint32_t id_count,
+ float *dists_out) override;
+
+ // Same as above, but returns the distances in a vector instead of an array.
+ // Convenience function for index.cpp.
+ virtual void preprocessed_distance(PQScratch<data_t> &pq_scratch, const uint32_t n_ids,
+ std::vector<float> &dists_out) override;
+
+ // Currently this function is required for DiskPQ. However, it too can be
+ // subsumed under preprocessed_distance if we add the appropriate scratch
+ // variables to PQScratch and initialize them in
+ // pq_flash_index.cpp::disk_iterate_to_fixed_point()
+ virtual float brute_force_distance(const float *query_vec, uint8_t *base_vec) override;
+
+ protected:
+ // assumes pre-processed query
+ virtual void prepopulate_chunkwise_distances(const float *query_vec, float *dist_vec);
+
+ // assumes no rotation is involved
+ // virtual void inflate_vector(uint8_t *base_vec, float *out_vec);
+
+ float *_tables = nullptr; // pq_tables = float array of size [256 * ndims]
+ uint64_t _ndims = 0; // ndims = true dimension of vectors
+ uint64_t _num_chunks = 0;
+ bool _is_opq = false;
+ uint32_t *_chunk_offsets = nullptr;
+ float *_centroid = nullptr;
+ float *_tables_tr = nullptr; // same as pq_tables, but col-major
+ float *_rotmat_tr = nullptr;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/pq_scratch.h b/be/src/extern/diskann/include/pq_scratch.h
new file mode 100644
index 00000000000..95f1b13954c
--- /dev/null
+++ b/be/src/extern/diskann/include/pq_scratch.h
@@ -0,0 +1,23 @@
+#pragma once
+#include <cstdint>
+#include "pq_common.h"
+#include "utils.h"
+
+namespace diskann
+{
+
+template <typename T> class PQScratch
+{
+ public:
+ float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS]
+ float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE
+ uint8_t *aligned_pq_coord_scratch = nullptr; // AT LEAST [N_CHUNKS * MAX_DEGREE]
+ float *rotated_query = nullptr;
+ float *aligned_query_float = nullptr;
+
+ PQScratch(size_t graph_degree, size_t aligned_dim);
+ void initialize(size_t dim, const T *query, const float norm = 1.0f);
+ virtual ~PQScratch();
+};
+
+} // namespace diskann
\ No newline at end of file
diff --git a/be/src/extern/diskann/include/program_options_utils.hpp b/be/src/extern/diskann/include/program_options_utils.hpp
new file mode 100644
index 00000000000..2be60595b22
--- /dev/null
+++ b/be/src/extern/diskann/include/program_options_utils.hpp
@@ -0,0 +1,81 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <string.h>
+
+namespace program_options_utils
+{
+const std::string make_program_description(const char *executable_name, const char *description)
+{
+ return std::string("\n")
+ .append(description)
+ .append("\n\n")
+ .append("Usage: ")
+ .append(executable_name)
+ .append(" [OPTIONS]");
+}
+
+// Required parameters
+const char *DATA_TYPE_DESCRIPTION = "data type, one of {int8, uint8, float} - float is single precision (32 bit)";
+const char *DISTANCE_FUNCTION_DESCRIPTION =
+ "distance function {l2, mips, fast_l2, cosine}. 'fast l2' and 'mips' only support data_type float";
+const char *INDEX_PATH_PREFIX_DESCRIPTION = "Path prefix to the index, e.g. '/mnt/data/my_ann_index'";
+const char *RESULT_PATH_DESCRIPTION =
+ "Path prefix for saving results of the queries, e.g. '/mnt/data/query_file_X.bin'";
+const char *QUERY_FILE_DESCRIPTION = "Query file in binary format, e.g. '/mnt/data/query_file_X.bin'";
+const char *NUMBER_OF_RESULTS_DESCRIPTION = "Number of neighbors to be returned (K in the DiskANN white paper)";
+const char *SEARCH_LIST_DESCRIPTION =
+ "Size of search list to use. This value is the number of neighbor/distance pairs to keep in memory at the same "
+ "time while performing a query. This can also be described as the size of the working set at query time. This "
+ "must be greater than or equal to the number of results/neighbors to return (K in the white paper). Corresponds "
+ "to L in the DiskANN white paper.";
+const char *INPUT_DATA_PATH = "Input data file in bin format. This is the file you want to build the index over. "
+ "File format: Shape of the vector followed by the vector of embeddings as binary data.";
+
+// Optional parameters
+const char *FILTER_LABEL_DESCRIPTION =
+ "Filter to use when running a query. 'filter_label' and 'query_filters_file' are mutually exclusive.";
+const char *FILTERS_FILE_DESCRIPTION =
+ "Filter file for Queries for Filtered Search. File format is text with one filter per line. File must "
+ "have exactly one filter OR the same number of filters as there are queries in the 'query_file'.";
+const char *LABEL_TYPE_DESCRIPTION =
+ "Storage type of Labels {uint/uint32, ushort/uint16}, default value is uint which will consume memory 4 bytes per "
+ "filter. 'uint' is an alias for 'uint32' and 'ushort' is an alias for 'uint16'.";
+const char *GROUND_TRUTH_FILE_DESCRIPTION =
+ "ground truth file for the queryset"; // what's the format, what's the requirements? does it need to include an
+ // entry for every item or just a small subset? I have so many questions about
+ // this file
+const char *NUMBER_THREADS_DESCRIPTION = "Number of threads used for building index. Defaults to number of logical "
+ "processor cores on your this machine returned by omp_get_num_procs()";
+const char *FAIL_IF_RECALL_BELOW =
+ "Value between 0 (inclusive) and 100 (exclusive) indicating the recall tolerance percentage threshold before "
+ "program fails with a non-zero exit code. The default value of 0 means that the program will complete "
+ "successfully with any recall value. A non-zero value indicates the floor for acceptable recall values. If the "
+ "calculated recall value is below this threshold then the program will write out the results but return a non-zero "
+ "exit code as a signal that the recall was not acceptable."; // does it continue running or die immediately? Will I
+ // still get my results even if the return code is -1?
+
+const char *NUMBER_OF_NODES_TO_CACHE = "Number of BFS nodes around medoid(s) to cache. Default value: 0";
+const char *BEAMWIDTH = "Beamwidth for search. Set 0 to optimize internally. Default value: 2";
+const char *MAX_BUILD_DEGREE = "Maximum graph degree";
+const char *GRAPH_BUILD_COMPLEXITY =
+ "Size of the search working set during build time. This is the numer of neighbor/distance pairs to keep in memory "
+ "while building the index. Higher value results in a higher quality graph but it will take more time to build the "
+ "graph.";
+const char *GRAPH_BUILD_ALPHA = "Alpha controls density and diameter of graph, set 1 for sparse graph, 1.2 or 1.4 for "
+ "denser graphs with lower diameter";
+const char *BUIlD_GRAPH_PQ_BYTES = "Number of PQ bytes to build the index; 0 for full precision build";
+const char *USE_OPQ = "Use Optimized Product Quantization (OPQ).";
+const char *LABEL_FILE = "Input label file in txt format for Filtered Index build. The file should contain comma "
+ "separated filters for each node with each line corresponding to a graph node";
+const char *UNIVERSAL_LABEL =
+ "Universal label, Use only in conjunction with label file for filtered index build. If a "
+ "graph node has all the labels against it, we can assign a special universal filter to the "
+ "point instead of comma separated filters for that point. The universal label should be assigned to nodes "
+ "in the labels file instead of listing all labels for a node. DiskANN will not automatically assign a "
+ "universal label to a node.";
+const char *FILTERED_LBUILD = "Build complexity for filtered points, higher value results in better graphs";
+
+} // namespace program_options_utils
diff --git a/be/src/extern/diskann/include/quantized_distance.h b/be/src/extern/diskann/include/quantized_distance.h
new file mode 100644
index 00000000000..cc4aea92969
--- /dev/null
+++ b/be/src/extern/diskann/include/quantized_distance.h
@@ -0,0 +1,56 @@
+#pragma once
+#include <memory>
+#include <string>
+#include <vector>
+#include "abstract_scratch.h"
+
+namespace diskann
+{
+template <typename data_t> class PQScratch;
+
+template <typename data_t> class QuantizedDistance
+{
+ public:
+ QuantizedDistance() = default;
+ QuantizedDistance(const QuantizedDistance &) = delete;
+ QuantizedDistance &operator=(const QuantizedDistance &) = delete;
+ virtual ~QuantizedDistance() = default;
+
+ virtual bool is_opq() const = 0;
+ virtual std::string get_quantized_vectors_filename(const std::string &prefix) const = 0;
+ virtual std::string get_pivot_data_filename(const std::string &prefix) const = 0;
+ virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const = 0;
+
+ // Loading the PQ centroid table need not be part of the abstract class.
+ // However, we want to indicate that this function will change once we have a
+ // file reader hierarchy, so leave it here as-is.
+#ifdef EXEC_ENV_OLS
+ virtual void load_pivot_data(MemoryMappedFiles &files, const std::String &pq_table_file, size_t num_chunks) = 0;
+#else
+ virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) = 0;
+#endif
+
+ // Number of chunks in the PQ table. Depends on the compression level used.
+ // Has to be < ndim
+ virtual uint32_t get_num_chunks() const = 0;
+
+ // Preprocess the query by computing chunk distances from the query vector to
+ // various centroids. Since we don't want this class to do scratch management,
+ // we will take a PQScratch object which can come either from Index class or
+ // PQFlashIndex class.
+ virtual void preprocess_query(const data_t *query_vec, uint32_t query_dim, PQScratch<data_t> &pq_scratch) = 0;
+
+ // Workhorse
+ // This function must be called after preprocess_query
+ virtual void preprocessed_distance(PQScratch<data_t> &pq_scratch, const uint32_t id_count, float *dists_out) = 0;
+
+ // Same as above, but convenience function for index.cpp.
+ virtual void preprocessed_distance(PQScratch<data_t> &pq_scratch, const uint32_t n_ids,
+ std::vector<float> &dists_out) = 0;
+
+ // Currently this function is required for DiskPQ. However, it too can be subsumed
+ // under preprocessed_distance if we add the appropriate scratch variables to
+ // PQScratch and initialize them in pq_flash_index.cpp::disk_iterate_to_fixed_point()
+ virtual float brute_force_distance(const float *query_vec, uint8_t *base_vec) = 0;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/restapi/common.h b/be/src/extern/diskann/include/restapi/common.h
new file mode 100644
index 00000000000..b8339635a86
--- /dev/null
+++ b/be/src/extern/diskann/include/restapi/common.h
@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <cpprest/base_uri.h>
+#include <restapi/search_wrapper.h>
+
+namespace diskann
+{
+// Constants
+static const std::string VECTOR_KEY = "query", K_KEY = "k", INDICES_KEY = "indices", DISTANCES_KEY = "distances",
+ TAGS_KEY = "tags", QUERY_ID_KEY = "query_id", ERROR_MESSAGE_KEY = "error", L_KEY = "Ls",
+ TIME_TAKEN_KEY = "time_taken_in_us", PARTITION_KEY = "partition",
+ UNKNOWN_ERROR = "unknown_error";
+const unsigned int DEFAULT_L = 100;
+
+} // namespace diskann
\ No newline at end of file
diff --git a/be/src/extern/diskann/include/restapi/search_wrapper.h b/be/src/extern/diskann/include/restapi/search_wrapper.h
new file mode 100644
index 00000000000..ebd067d8aa2
--- /dev/null
+++ b/be/src/extern/diskann/include/restapi/search_wrapper.h
@@ -0,0 +1,140 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include <stdexcept>
+
+#include <index.h>
+#include <pq_flash_index.h>
+
+namespace diskann
+{
+class SearchResult
+{
+ public:
+ SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, const unsigned *const indices,
+ const float *const distances, const std::string *const tags = nullptr,
+ const unsigned *const partitions = nullptr);
+
+ const std::vector<unsigned int> &get_indices() const
+ {
+ return _indices;
+ }
+ const std::vector<float> &get_distances() const
+ {
+ return _distances;
+ }
+ bool tags_enabled() const
+ {
+ return _tags_enabled;
+ }
+ const std::vector<std::string> &get_tags() const
+ {
+ return _tags;
+ }
+ bool partitions_enabled() const
+ {
+ return _partitions_enabled;
+ }
+ const std::vector<unsigned> &get_partitions() const
+ {
+ return _partitions;
+ }
+ unsigned get_time() const
+ {
+ return _search_time_in_ms;
+ }
+
+ private:
+ unsigned int _K;
+ unsigned int _search_time_in_ms;
+ std::vector<unsigned int> _indices;
+ std::vector<float> _distances;
+
+ bool _tags_enabled;
+ std::vector<std::string> _tags;
+
+ bool _partitions_enabled;
+ std::vector<unsigned> _partitions;
+};
+
+class SearchNotImplementedException : public std::logic_error
+{
+ private:
+ std::string _errormsg;
+
+ public:
+ SearchNotImplementedException(const char *type) : std::logic_error("Not Implemented")
+ {
+ _errormsg = "Search with data type ";
+ _errormsg += std::string(type);
+ _errormsg += " not implemented : ";
+ _errormsg += __FUNCTION__;
+ }
+
+ virtual const char *what() const throw()
+ {
+ return _errormsg.c_str();
+ }
+};
+
+class BaseSearch
+{
+ public:
+ BaseSearch(const std::string &tagsFile = nullptr);
+ virtual SearchResult search(const float *query, const unsigned int dimensions, const unsigned int K,
+ const unsigned int Ls)
+ {
+ throw SearchNotImplementedException("float");
+ }
+ virtual SearchResult search(const int8_t *query, const unsigned int dimensions, const unsigned int K,
+ const unsigned int Ls)
+ {
+ throw SearchNotImplementedException("int8_t");
+ }
+
+ virtual SearchResult search(const uint8_t *query, const unsigned int dimensions, const unsigned int K,
+ const unsigned int Ls)
+ {
+ throw SearchNotImplementedException("uint8_t");
+ }
+
+ void lookup_tags(const unsigned K, const unsigned *indices, std::string *ret_tags);
+
+ protected:
+ bool _tags_enabled;
+ std::vector<std::string> _tags_str;
+};
+
+template <typename T> class InMemorySearch : public BaseSearch
+{
+ public:
+ InMemorySearch(const std::string &baseFile, const std::string &indexFile, const std::string &tagsFile, Metric m,
+ uint32_t num_threads, uint32_t search_l);
+ virtual ~InMemorySearch();
+
+ SearchResult search(const T *query, const unsigned int dimensions, const unsigned int K, const unsigned int Ls);
+
+ private:
+ unsigned int _dimensions, _numPoints;
+ std::unique_ptr<diskann::Index<T>> _index;
+};
+
+template <typename T> class PQFlashSearch : public BaseSearch
+{
+ public:
+ PQFlashSearch(const std::string &indexPrefix, const unsigned num_nodes_to_cache, const unsigned num_threads,
+ const std::string &tagsFile, Metric m);
+ virtual ~PQFlashSearch();
+
+ SearchResult search(const T *query, const unsigned int dimensions, const unsigned int K, const unsigned int Ls);
+
+ private:
+ unsigned int _dimensions, _numPoints;
+ std::unique_ptr<diskann::PQFlashIndex<T>> _index;
+ std::shared_ptr<AlignedFileReader> reader;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/restapi/server.h b/be/src/extern/diskann/include/restapi/server.h
new file mode 100644
index 00000000000..1d75847a2c8
--- /dev/null
+++ b/be/src/extern/diskann/include/restapi/server.h
@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <restapi/common.h>
+#include <cpprest/http_listener.h>
+
+namespace diskann
+{
+class Server
+{
+ public:
+ Server(web::uri &url, std::vector<std::unique_ptr<diskann::BaseSearch>> &multi_searcher,
+ const std::string &typestring);
+ virtual ~Server();
+
+ pplx::task<void> open();
+ pplx::task<void> close();
+
+ protected:
+ template <class T> void handle_post(web::http::http_request message);
+
+ template <typename T>
+ web::json::value toJsonArray(const std::vector<T> &v, std::function<web::json::value(const T &)> valConverter);
+ web::json::value prepareResponse(const int64_t &queryId, const int k);
+
+ template <class T>
+ void parseJson(const utility::string_t &body, unsigned int &k, int64_t &queryId, T *&queryVector,
+ unsigned int &dimensions, unsigned &Ls);
+
+ web::json::value idsToJsonArray(const diskann::SearchResult &result);
+ web::json::value distancesToJsonArray(const diskann::SearchResult &result);
+ web::json::value tagsToJsonArray(const diskann::SearchResult &result);
+ web::json::value partitionsToJsonArray(const diskann::SearchResult &result);
+
+ SearchResult aggregate_results(const unsigned K, const std::vector<diskann::SearchResult> &results);
+
+ private:
+ bool _isDebug;
+ std::unique_ptr<web::http::experimental::listener::http_listener> _listener;
+ const bool _multi_search;
+ std::vector<std::unique_ptr<diskann::BaseSearch>> _multi_searcher;
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/scratch.h b/be/src/extern/diskann/include/scratch.h
new file mode 100644
index 00000000000..2f43e3365a6
--- /dev/null
+++ b/be/src/extern/diskann/include/scratch.h
@@ -0,0 +1,216 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <vector>
+
+#include "boost_dynamic_bitset_fwd.h"
+// #include "boost/dynamic_bitset.hpp"
+#include "tsl/robin_set.h"
+#include "tsl/robin_map.h"
+#include "tsl/sparse_map.h"
+
+#include "aligned_file_reader.h"
+#include "abstract_scratch.h"
+#include "neighbor.h"
+#include "defaults.h"
+#include "concurrent_queue.h"
+
+namespace diskann
+{
+template <typename T> class PQScratch;
+
+//
+// AbstractScratch space for in-memory index based search
+//
+template <typename T> class InMemQueryScratch : public AbstractScratch<T>
+{
+ public:
+ ~InMemQueryScratch();
+ InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, size_t aligned_dim,
+ size_t alignment_factor, bool init_pq_scratch = false);
+ void resize_for_new_L(uint32_t new_search_l);
+ void clear();
+
+ inline uint32_t get_L()
+ {
+ return _L;
+ }
+ inline uint32_t get_R()
+ {
+ return _R;
+ }
+ inline uint32_t get_maxc()
+ {
+ return _maxc;
+ }
+ inline T *aligned_query()
+ {
+ return this->_aligned_query_T;
+ }
+ inline PQScratch<T> *pq_scratch()
+ {
+ return this->_pq_scratch;
+ }
+ inline std::vector<Neighbor> &pool()
+ {
+ return _pool;
+ }
+ inline NeighborPriorityQueue &best_l_nodes()
+ {
+ return _best_l_nodes;
+ }
+ inline std::vector<float> &occlude_factor()
+ {
+ return _occlude_factor;
+ }
+ inline tsl::robin_set<uint32_t> &inserted_into_pool_rs()
+ {
+ return _inserted_into_pool_rs;
+ }
+ inline boost::dynamic_bitset<> &inserted_into_pool_bs()
+ {
+ return *_inserted_into_pool_bs;
+ }
+ inline std::vector<uint32_t> &id_scratch()
+ {
+ return _id_scratch;
+ }
+ inline std::vector<float> &dist_scratch()
+ {
+ return _dist_scratch;
+ }
+ inline tsl::robin_set<uint32_t> &expanded_nodes_set()
+ {
+ return _expanded_nodes_set;
+ }
+ inline std::vector<Neighbor> &expanded_nodes_vec()
+ {
+ return _expanded_nghrs_vec;
+ }
+ inline std::vector<uint32_t> &occlude_list_output()
+ {
+ return _occlude_list_output;
+ }
+
+ private:
+ uint32_t _L;
+ uint32_t _R;
+ uint32_t _maxc;
+
+ // _pool stores all neighbors explored from best_L_nodes.
+ // Usually around L+R, but could be higher.
+ // Initialized to 3L+R for some slack, expands as needed.
+ std::vector<Neighbor> _pool;
+
+ // _best_l_nodes is reserved for storing best L entries
+ // Underlying storage is L+1 to support inserts
+ NeighborPriorityQueue _best_l_nodes;
+
+ // _occlude_factor.size() >= pool.size() in occlude_list function
+ // _pool is clipped to maxc in occlude_list before affecting _occlude_factor
+ // _occlude_factor is initialized to maxc size
+ std::vector<float> _occlude_factor;
+
+ // Capacity initialized to 20L
+ tsl::robin_set<uint32_t> _inserted_into_pool_rs;
+
+ // Use a pointer here to allow for forward declaration of dynamic_bitset
+ // in public headers to avoid making boost a dependency for clients
+ // of DiskANN.
+ boost::dynamic_bitset<> *_inserted_into_pool_bs;
+
+ // _id_scratch.size() must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp
+ std::vector<uint32_t> _id_scratch;
+
+ // _dist_scratch must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp
+ // _dist_scratch should be at least the size of id_scratch
+ std::vector<float> _dist_scratch;
+
+ // Buffers used in process delete, capacity increases as needed
+ tsl::robin_set<uint32_t> _expanded_nodes_set;
+ std::vector<Neighbor> _expanded_nghrs_vec;
+ std::vector<uint32_t> _occlude_list_output;
+};
+
+//
+// AbstractScratch space for SSD index based search
+//
+
+template <typename T> class SSDQueryScratch : public AbstractScratch<T>
+{
+ public:
+ T *coord_scratch = nullptr; // MUST BE AT LEAST [sizeof(T) * data_dim]
+
+ char *sector_scratch = nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN]
+ size_t sector_idx = 0; // index of next [SECTOR_LEN] scratch to use
+
+ tsl::robin_set<size_t> visited;
+ NeighborPriorityQueue retset;
+ std::vector<Neighbor> full_retset;
+
+ SSDQueryScratch(size_t aligned_dim, size_t visited_reserve);
+ ~SSDQueryScratch();
+
+ void reset();
+};
+
+template <typename T> class SSDThreadData
+{
+ public:
+ SSDQueryScratch<T> scratch;
+ IOContext ctx;
+
+ SSDThreadData(size_t aligned_dim, size_t visited_reserve);
+ void clear();
+};
+
+//
+// Class to avoid the hassle of pushing and popping the query scratch.
+//
+template <typename T> class ScratchStoreManager
+{
+ public:
+ ScratchStoreManager(ConcurrentQueue<T *> &query_scratch) : _scratch_pool(query_scratch)
+ {
+ _scratch = query_scratch.pop();
+ while (_scratch == nullptr)
+ {
+ query_scratch.wait_for_push_notify();
+ _scratch = query_scratch.pop();
+ }
+ }
+ T *scratch_space()
+ {
+ return _scratch;
+ }
+
+ ~ScratchStoreManager()
+ {
+ _scratch->clear();
+ _scratch_pool.push(_scratch);
+ _scratch_pool.push_notify_all();
+ }
+
+ void destroy()
+ {
+ while (!_scratch_pool.empty())
+ {
+ auto scratch = _scratch_pool.pop();
+ while (scratch == nullptr)
+ {
+ _scratch_pool.wait_for_push_notify();
+ scratch = _scratch_pool.pop();
+ }
+ delete scratch;
+ }
+ }
+
+ private:
+ T *_scratch;
+ ConcurrentQueue<T *> &_scratch_pool;
+ ScratchStoreManager(const ScratchStoreManager<T> &);
+ ScratchStoreManager &operator=(const ScratchStoreManager<T> &);
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/simd_utils.h b/be/src/extern/diskann/include/simd_utils.h
new file mode 100644
index 00000000000..4b07369983c
--- /dev/null
+++ b/be/src/extern/diskann/include/simd_utils.h
@@ -0,0 +1,106 @@
+#pragma once
+
+#ifdef _WINDOWS
+#include <immintrin.h>
+#include <smmintrin.h>
+#include <tmmintrin.h>
+#include <intrin.h>
+#else
+#include <immintrin.h>
+#endif
+
+namespace diskann
+{
+static inline __m256 _mm256_mul_epi8(__m256i X)
+{
+ __m256i zero = _mm256_setzero_si256();
+
+ __m256i sign_x = _mm256_cmpgt_epi8(zero, X);
+
+ __m256i xlo = _mm256_unpacklo_epi8(X, sign_x);
+ __m256i xhi = _mm256_unpackhi_epi8(X, sign_x);
+
+ return _mm256_cvtepi32_ps(_mm256_add_epi32(_mm256_madd_epi16(xlo, xlo), _mm256_madd_epi16(xhi, xhi)));
+}
+
+static inline __m128 _mm_mulhi_epi8(__m128i X)
+{
+ __m128i zero = _mm_setzero_si128();
+ __m128i sign_x = _mm_cmplt_epi8(X, zero);
+ __m128i xhi = _mm_unpackhi_epi8(X, sign_x);
+
+ return _mm_cvtepi32_ps(_mm_add_epi32(_mm_setzero_si128(), _mm_madd_epi16(xhi, xhi)));
+}
+
+static inline __m128 _mm_mulhi_epi8_shift32(__m128i X)
+{
+ __m128i zero = _mm_setzero_si128();
+ X = _mm_srli_epi64(X, 32);
+ __m128i sign_x = _mm_cmplt_epi8(X, zero);
+ __m128i xhi = _mm_unpackhi_epi8(X, sign_x);
+
+ return _mm_cvtepi32_ps(_mm_add_epi32(_mm_setzero_si128(), _mm_madd_epi16(xhi, xhi)));
+}
+static inline __m128 _mm_mul_epi8(__m128i X, __m128i Y)
+{
+ __m128i zero = _mm_setzero_si128();
+
+ __m128i sign_x = _mm_cmplt_epi8(X, zero);
+ __m128i sign_y = _mm_cmplt_epi8(Y, zero);
+
+ __m128i xlo = _mm_unpacklo_epi8(X, sign_x);
+ __m128i xhi = _mm_unpackhi_epi8(X, sign_x);
+ __m128i ylo = _mm_unpacklo_epi8(Y, sign_y);
+ __m128i yhi = _mm_unpackhi_epi8(Y, sign_y);
+
+ return _mm_cvtepi32_ps(_mm_add_epi32(_mm_madd_epi16(xlo, ylo), _mm_madd_epi16(xhi, yhi)));
+}
+static inline __m128 _mm_mul_epi8(__m128i X)
+{
+ __m128i zero = _mm_setzero_si128();
+ __m128i sign_x = _mm_cmplt_epi8(X, zero);
+ __m128i xlo = _mm_unpacklo_epi8(X, sign_x);
+ __m128i xhi = _mm_unpackhi_epi8(X, sign_x);
+
+ return _mm_cvtepi32_ps(_mm_add_epi32(_mm_madd_epi16(xlo, xlo), _mm_madd_epi16(xhi, xhi)));
+}
+
+static inline __m128 _mm_mul32_pi8(__m128i X, __m128i Y)
+{
+ __m128i xlo = _mm_cvtepi8_epi16(X), ylo = _mm_cvtepi8_epi16(Y);
+ return _mm_cvtepi32_ps(_mm_unpacklo_epi32(_mm_madd_epi16(xlo, ylo), _mm_setzero_si128()));
+}
+
+static inline __m256 _mm256_mul_epi8(__m256i X, __m256i Y)
+{
+ __m256i zero = _mm256_setzero_si256();
+
+ __m256i sign_x = _mm256_cmpgt_epi8(zero, X);
+ __m256i sign_y = _mm256_cmpgt_epi8(zero, Y);
+
+ __m256i xlo = _mm256_unpacklo_epi8(X, sign_x);
+ __m256i xhi = _mm256_unpackhi_epi8(X, sign_x);
+ __m256i ylo = _mm256_unpacklo_epi8(Y, sign_y);
+ __m256i yhi = _mm256_unpackhi_epi8(Y, sign_y);
+
+ return _mm256_cvtepi32_ps(_mm256_add_epi32(_mm256_madd_epi16(xlo, ylo), _mm256_madd_epi16(xhi, yhi)));
+}
+
+static inline __m256 _mm256_mul32_pi8(__m128i X, __m128i Y)
+{
+ __m256i xlo = _mm256_cvtepi8_epi16(X), ylo = _mm256_cvtepi8_epi16(Y);
+ return _mm256_blend_ps(_mm256_cvtepi32_ps(_mm256_madd_epi16(xlo, ylo)), _mm256_setzero_ps(), 252);
+}
+
+static inline float _mm256_reduce_add_ps(__m256 x)
+{
+ /* ( x3+x7, x2+x6, x1+x5, x0+x4 ) */
+ const __m128 x128 = _mm_add_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x));
+ /* ( -, -, x1+x3+x5+x7, x0+x2+x4+x6 ) */
+ const __m128 x64 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128));
+ /* ( -, -, -, x0+x1+x2+x3+x4+x5+x6+x7 ) */
+ const __m128 x32 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55));
+ /* Conversion to float is a no-op on x86-64 */
+ return _mm_cvtss_f32(x32);
+}
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/tag_uint128.h b/be/src/extern/diskann/include/tag_uint128.h
new file mode 100644
index 00000000000..642de315941
--- /dev/null
+++ b/be/src/extern/diskann/include/tag_uint128.h
@@ -0,0 +1,68 @@
+#pragma once
+#include <cstdint>
+#include <type_traits>
+
+namespace diskann
+{
+#pragma pack(push, 1)
+
+struct tag_uint128
+{
+ std::uint64_t _data1 = 0;
+ std::uint64_t _data2 = 0;
+
+ bool operator==(const tag_uint128 &other) const
+ {
+ return _data1 == other._data1 && _data2 == other._data2;
+ }
+
+ bool operator==(std::uint64_t other) const
+ {
+ return _data1 == other && _data2 == 0;
+ }
+
+ tag_uint128 &operator=(const tag_uint128 &other)
+ {
+ _data1 = other._data1;
+ _data2 = other._data2;
+
+ return *this;
+ }
+
+ tag_uint128 &operator=(std::uint64_t other)
+ {
+ _data1 = other;
+ _data2 = 0;
+
+ return *this;
+ }
+};
+
+#pragma pack(pop)
+} // namespace diskann
+
+namespace std
+{
+// Hash 128 input bits down to 64 bits of output.
+// This is intended to be a reasonably good hash function.
+inline std::uint64_t Hash128to64(const std::uint64_t &low, const std::uint64_t &high)
+{
+ // Murmur-inspired hashing.
+ const std::uint64_t kMul = 0x9ddfea08eb382d69ULL;
+ std::uint64_t a = (low ^ high) * kMul;
+ a ^= (a >> 47);
+ std::uint64_t b = (high ^ a) * kMul;
+ b ^= (b >> 47);
+ b *= kMul;
+ return b;
+}
+
+template <> struct hash<diskann::tag_uint128>
+{
+ size_t operator()(const diskann::tag_uint128 &key) const noexcept
+ {
+ return Hash128to64(key._data1, key._data2); // map -0 to 0
+ }
+};
+
+} // namespace std
\ No newline at end of file
diff --git a/be/src/extern/diskann/include/timer.h b/be/src/extern/diskann/include/timer.h
new file mode 100644
index 00000000000..325edf30297
--- /dev/null
+++ b/be/src/extern/diskann/include/timer.h
@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+#pragma once
+
+#include <chrono>
+
+namespace diskann
+{
+class Timer
+{
+ typedef std::chrono::high_resolution_clock _clock;
+ std::chrono::time_point<_clock> check_point;
+
+ public:
+ Timer() : check_point(_clock::now())
+ {
+ }
+
+ void reset()
+ {
+ check_point = _clock::now();
+ }
+
+ long long elapsed() const
+ {
+ return std::chrono::duration_cast<std::chrono::microseconds>(_clock::now() - check_point).count();
+ }
+
+ float elapsed_seconds() const
+ {
+ return (float)elapsed() / 1000000.0f;
+ }
+
+ std::string elapsed_seconds_for_step(const std::string &step) const
+ {
+ return std::string("Time for ") + step + std::string(": ") + std::to_string(elapsed_seconds()) +
+ std::string(" seconds");
+ }
+};
+} // namespace diskann
diff --git a/be/src/extern/diskann/include/tsl/.clang-format b/be/src/extern/diskann/include/tsl/.clang-format
new file mode 100644
index 00000000000..9d159247d51
--- /dev/null
+++ b/be/src/extern/diskann/include/tsl/.clang-format
@@ -0,0 +1,2 @@
+DisableFormat: true
+SortIncludes: false
diff --git a/be/src/extern/diskann/include/tsl/robin_growth_policy.h b/be/src/extern/diskann/include/tsl/robin_growth_policy.h
new file mode 100644
index 00000000000..6bfa9e5f927
--- /dev/null
+++ b/be/src/extern/diskann/include/tsl/robin_growth_policy.h
@@ -0,0 +1,330 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2017 Tessil
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#ifndef TSL_ROBIN_GROWTH_POLICY_H
+#define TSL_ROBIN_GROWTH_POLICY_H
+
+
+#include <algorithm>
+#include <array>
+#include <climits>
+#include <cmath>
+#include <cstddef>
+#include <iterator>
+#include <limits>
+#include <ratio>
+#include <stdexcept>
+
+
+#ifndef tsl_assert
+# ifdef TSL_DEBUG
+# define tsl_assert(expr) assert(expr)
+# else
+# define tsl_assert(expr) (static_cast<void>(0))
+# endif
+#endif
+
+
+/**
+ * If exceptions are enabled, throw the exception passed in parameter, otherwise call std::terminate.
+ */
+#ifndef TSL_THROW_OR_TERMINATE
+# if (defined(__cpp_exceptions) || defined(__EXCEPTIONS) || (defined (_MSC_VER) && defined (_CPPUNWIND))) && !defined(TSL_NO_EXCEPTIONS)
+# define TSL_THROW_OR_TERMINATE(ex, msg) throw ex(msg)
+# else
+# ifdef NDEBUG
+# define TSL_THROW_OR_TERMINATE(ex, msg) std::terminate()
+# else
+# include <cstdio>
+# define TSL_THROW_OR_TERMINATE(ex, msg) do { std::fprintf(stderr, msg); std::terminate(); } while(0)
+# endif
+# endif
+#endif
+
+
+#ifndef TSL_LIKELY
+# if defined(__GNUC__) || defined(__clang__)
+# define TSL_LIKELY(exp) (__builtin_expect(!!(exp), true))
+# else
+# define TSL_LIKELY(exp) (exp)
+# endif
+#endif
+
+
+namespace tsl {
+namespace rh {
+
+/**
+ * Grow the hash table by a factor of GrowthFactor keeping the bucket count to a power of two. It allows
+ * the table to use a mask operation instead of a modulo operation to map a hash to a bucket.
+ *
+ * GrowthFactor must be a power of two >= 2.
+ */
+template<std::size_t GrowthFactor>
+class power_of_two_growth_policy {
+public:
+ /**
+ * Called on the hash table creation and on rehash. The number of buckets for the table is passed in parameter.
+ * This number is a minimum, the policy may update this value with a higher value if needed (but not lower).
+ *
+ * If 0 is given, min_bucket_count_in_out must still be 0 after the policy creation and
+ * bucket_for_hash must always return 0 in this case.
+ */
+ explicit power_of_two_growth_policy(std::size_t& min_bucket_count_in_out) {
+ if(min_bucket_count_in_out > max_bucket_count()) {
+ TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size.");
+ }
+
+ if(min_bucket_count_in_out > 0) {
+ min_bucket_count_in_out = round_up_to_power_of_two(min_bucket_count_in_out);
+ m_mask = min_bucket_count_in_out - 1;
+ }
+ else {
+ m_mask = 0;
+ }
+ }
+
+ /**
+ * Return the bucket [0, bucket_count()) to which the hash belongs.
+ * If bucket_count() is 0, it must always return 0.
+ */
+ std::size_t bucket_for_hash(std::size_t hash) const noexcept {
+ return hash & m_mask;
+ }
+
+ /**
+ * Return the number of buckets that should be used on next growth.
+ */
+ std::size_t next_bucket_count() const {
+ if((m_mask + 1) > max_bucket_count() / GrowthFactor) {
+ TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size.");
+ }
+
+ return (m_mask + 1) * GrowthFactor;
+ }
+
+ /**
+ * Return the maximum number of buckets supported by the policy.
+ */
+ std::size_t max_bucket_count() const {
+ // Largest power of two.
+ return ((std::numeric_limits<std::size_t>::max)() / 2) + 1;
+ }
+
+ /**
+ * Reset the growth policy as if it was created with a bucket count of 0.
+ * After a clear, the policy must always return 0 when bucket_for_hash is called.
+ */
+ void clear() noexcept {
+ m_mask = 0;
+ }
+
+private:
+ static std::size_t round_up_to_power_of_two(std::size_t value) {
+ if(is_power_of_two(value)) {
+ return value;
+ }
+
+ if(value == 0) {
+ return 1;
+ }
+
+ --value;
+ for(std::size_t i = 1; i < sizeof(std::size_t) * CHAR_BIT; i *= 2) {
+ value |= value >> i;
+ }
+
+ return value + 1;
+ }
+
+ static constexpr bool is_power_of_two(std::size_t value) {
+ return value != 0 && (value & (value - 1)) == 0;
+ }
+
+protected:
+ static_assert(is_power_of_two(GrowthFactor) && GrowthFactor >= 2, "GrowthFactor must be a power of two >= 2.");
+
+ std::size_t m_mask;
+};
+
+
+/**
+ * Grow the hash table by GrowthFactor::num / GrowthFactor::den and use a modulo to map a hash
+ * to a bucket. Slower but it can be useful if you want a slower growth.
+ */
+template<class GrowthFactor = std::ratio<3, 2>>
+class mod_growth_policy {
+public:
+ explicit mod_growth_policy(std::size_t& min_bucket_count_in_out) {
+ if(min_bucket_count_in_out > max_bucket_count()) {
+ TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size.");
+ }
+
+ if(min_bucket_count_in_out > 0) {
+ m_mod = min_bucket_count_in_out;
+ }
+ else {
+ m_mod = 1;
+ }
+ }
+
+ std::size_t bucket_for_hash(std::size_t hash) const noexcept {
+ return hash % m_mod;
+ }
+
+ std::size_t next_bucket_count() const {
+ if(m_mod == max_bucket_count()) {
+ TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size.");
+ }
+
+ const double next_bucket_count = std::ceil(double(m_mod) * REHASH_SIZE_MULTIPLICATION_FACTOR);
+ if(!std::isnormal(next_bucket_count)) {
+ TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size.");
+ }
+
+ if(next_bucket_count > double(max_bucket_count())) {
+ return max_bucket_count();
+ }
+ else {
+ return std::size_t(next_bucket_count);
+ }
+ }
+
+ std::size_t max_bucket_count() const {
+ return MAX_BUCKET_COUNT;
+ }
+
+ void clear() noexcept {
+ m_mod = 1;
+ }
+
+private:
+ static constexpr double REHASH_SIZE_MULTIPLICATION_FACTOR = 1.0 * GrowthFactor::num / GrowthFactor::den;
+ static const std::size_t MAX_BUCKET_COUNT =
+ std::size_t(double(
+ (std::numeric_limits<std::size_t>::max)() / REHASH_SIZE_MULTIPLICATION_FACTOR
+ ));
+
+ static_assert(REHASH_SIZE_MULTIPLICATION_FACTOR >= 1.1, "Growth factor should be >= 1.1.");
+
+ std::size_t m_mod;
+};
+
+
+
+namespace detail {
+
+static constexpr const std::array<std::size_t, 40> PRIMES = {{
+ 1ul, 5ul, 17ul, 29ul, 37ul, 53ul, 67ul, 79ul, 97ul, 131ul, 193ul, 257ul, 389ul, 521ul, 769ul, 1031ul,
+ 1543ul, 2053ul, 3079ul, 6151ul, 12289ul, 24593ul, 49157ul, 98317ul, 196613ul, 393241ul, 786433ul,
+ 1572869ul, 3145739ul, 6291469ul, 12582917ul, 25165843ul, 50331653ul, 100663319ul, 201326611ul,
+ 402653189ul, 805306457ul, 1610612741ul, 3221225473ul, 4294967291ul
+}};
+
+template<unsigned int IPrime>
+static constexpr std::size_t mod(std::size_t hash) { return hash % PRIMES[IPrime]; }
+
+// MOD_PRIME[iprime](hash) returns hash % PRIMES[iprime]. This table allows for faster modulo as the
+// compiler can optimize the modulo code better with a constant known at the compilation.
+static constexpr const std::array<std::size_t(*)(std::size_t), 40> MOD_PRIME = {{
+ &mod<0>, &mod<1>, &mod<2>, &mod<3>, &mod<4>, &mod<5>, &mod<6>, &mod<7>, &mod<8>, &mod<9>, &mod<10>,
+ &mod<11>, &mod<12>, &mod<13>, &mod<14>, &mod<15>, &mod<16>, &mod<17>, &mod<18>, &mod<19>, &mod<20>,
+ &mod<21>, &mod<22>, &mod<23>, &mod<24>, &mod<25>, &mod<26>, &mod<27>, &mod<28>, &mod<29>, &mod<30>,
+ &mod<31>, &mod<32>, &mod<33>, &mod<34>, &mod<35>, &mod<36>, &mod<37> , &mod<38>, &mod<39>
+}};
+
+}
+
+/**
+ * Grow the hash table by using prime numbers as bucket count. Slower than tsl::rh::power_of_two_growth_policy in
+ * general but will probably distribute the values around better in the buckets with a poor hash function.
+ *
+ * To allow the compiler to optimize the modulo operation, a lookup table is used with constant primes numbers.
+ *
+ * With a switch the code would look like:
+ * \code
+ * switch(iprime) { // iprime is the current prime of the hash table
+ * case 0: hash % 5ul;
+ * break;
+ * case 1: hash % 17ul;
+ * break;
+ * case 2: hash % 29ul;
+ * break;
+ * ...
+ * }
+ * \endcode
+ *
+ * Due to the constant variable in the modulo the compiler is able to optimize the operation
+ * by a series of multiplications, substractions and shifts.
+ *
+ * The 'hash % 5' could become something like 'hash - (hash * 0xCCCCCCCD) >> 34) * 5' in a 64 bits environement.
+ */
+class prime_growth_policy {
+public:
+ explicit prime_growth_policy(std::size_t& min_bucket_count_in_out) {
+ auto it_prime = std::lower_bound(detail::PRIMES.begin(),
+ detail::PRIMES.end(), min_bucket_count_in_out);
+ if(it_prime == detail::PRIMES.end()) {
+ TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size.");
+ }
+
+ m_iprime = static_cast<unsigned int>(std::distance(detail::PRIMES.begin(), it_prime));
+ if(min_bucket_count_in_out > 0) {
+ min_bucket_count_in_out = *it_prime;
+ }
+ else {
+ min_bucket_count_in_out = 0;
+ }
+ }
+
+ std::size_t bucket_for_hash(std::size_t hash) const noexcept {
+ return detail::MOD_PRIME[m_iprime](hash);
+ }
+
+ std::size_t next_bucket_count() const {
+ if(m_iprime + 1 >= detail::PRIMES.size()) {
+ TSL_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maxmimum size.");
+ }
+
+ return detail::PRIMES[m_iprime + 1];
+ }
+
+ std::size_t max_bucket_count() const {
+ return detail::PRIMES.back();
+ }
+
+ void clear() noexcept {
+ m_iprime = 0;
+ }
+
+private:
+ unsigned int m_iprime;
+
+ static_assert((std::numeric_limits<decltype(m_iprime)>::max)() >= detail::PRIMES.size(),
+ "The type of m_iprime is not big enough.");
+};
+
+}
+}
+
+#endif
diff --git a/be/src/extern/diskann/include/tsl/robin_hash.h b/be/src/extern/diskann/include/tsl/robin_hash.h
new file mode 100644
index 00000000000..5ecc9622c87
--- /dev/null
+++ b/be/src/extern/diskann/include/tsl/robin_hash.h
@@ -0,0 +1,1285 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2017 Tessil
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#ifndef TSL_ROBIN_HASH_H
+#define TSL_ROBIN_HASH_H
+
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <exception>
+#include <iterator>
+#include <limits>
+#include <memory>
+#include <stdexcept>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+#include <vector>
+#include "robin_growth_policy.h"
+
+
+namespace tsl {
+
+namespace detail_robin_hash {
+
+template<typename T>
+struct make_void {
+ using type = void;
+};
+
+template<typename T, typename = void>
+struct has_is_transparent: std::false_type {
+};
+
+template<typename T>
+struct has_is_transparent<T, typename make_void<typename T::is_transparent>::type>: std::true_type {
+};
+
+template<typename U>
+struct is_power_of_two_policy: std::false_type {
+};
+
+template<std::size_t GrowthFactor>
+struct is_power_of_two_policy<tsl::rh::power_of_two_growth_policy<GrowthFactor>>: std::true_type {
+};
+
+
+
+using truncated_hash_type = std::uint_least32_t;
+
+/**
+ * Helper class that store a truncated hash if StoreHash is true and nothing otherwise.
+ */
+template<bool StoreHash>
+class bucket_entry_hash {
+public:
+ bool bucket_hash_equal(std::size_t /*hash*/) const noexcept {
+ return true;
+ }
+
+ truncated_hash_type truncated_hash() const noexcept {
+ return 0;
+ }
+
+protected:
+ void set_hash(truncated_hash_type /*hash*/) noexcept {
+ }
+};
+
+template<>
+class bucket_entry_hash<true> {
+public:
+ bool bucket_hash_equal(std::size_t hash) const noexcept {
+ return m_hash == truncated_hash_type(hash);
+ }
+
+ truncated_hash_type truncated_hash() const noexcept {
+ return m_hash;
+ }
+
+protected:
+ void set_hash(truncated_hash_type hash) noexcept {
+ m_hash = truncated_hash_type(hash);
+ }
+
+private:
+ truncated_hash_type m_hash;
+};
+
+
+/**
+ * Each bucket entry has:
+ * - A value of type `ValueType`.
+ * - An integer to store how far the value of the bucket, if any, is from its ideal bucket
+ * (ex: if the current bucket 5 has the value 'foo' and `hash('foo') % nb_buckets` == 3,
+ * `dist_from_ideal_bucket()` will return 2 as the current value of the bucket is two
+ * buckets away from its ideal bucket)
+ * If there is no value in the bucket (i.e. `empty()` is true) `dist_from_ideal_bucket()` will be < 0.
+ * - A marker which tells us if the bucket is the last bucket of the bucket array (useful for the
+ * iterator of the hash table).
+ * - If `StoreHash` is true, 32 bits of the hash of the value, if any, are also stored in the bucket.
+ * If the size of the hash is more than 32 bits, it is truncated. We don't store the full hash
+ * as storing the hash is a potential opportunity to use the unused space due to the alignement
+ * of the bucket_entry structure. We can thus potentially store the hash without any extra space
+ * (which would not be possible with 64 bits of the hash).
+ */
+template<typename ValueType, bool StoreHash>
+class bucket_entry: public bucket_entry_hash<StoreHash> {
+ using bucket_hash = bucket_entry_hash<StoreHash>;
+
+public:
+ using value_type = ValueType;
+ using distance_type = std::int_least16_t;
+
+
+ bucket_entry() noexcept: bucket_hash(), m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET),
+ m_last_bucket(false)
+ {
+ tsl_assert(empty());
+ }
+
+ bucket_entry(bool last_bucket) noexcept: bucket_hash(), m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET),
+ m_last_bucket(last_bucket)
+ {
+ tsl_assert(empty());
+ }
+
+ bucket_entry(const bucket_entry& other) noexcept(std::is_nothrow_copy_constructible<value_type>::value):
+ bucket_hash(other),
+ m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET),
+ m_last_bucket(other.m_last_bucket)
+ {
+ if(!other.empty()) {
+ ::new (static_cast<void*>(std::addressof(m_value))) value_type(other.value());
+ m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket;
+ }
+ }
+
+ /**
+ * Never really used, but still necessary as we must call resize on an empty `std::vector<bucket_entry>`.
+ * and we need to support move-only types. See robin_hash constructor for details.
+ */
+ bucket_entry(bucket_entry&& other) noexcept(std::is_nothrow_move_constructible<value_type>::value):
+ bucket_hash(std::move(other)),
+ m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET),
+ m_last_bucket(other.m_last_bucket)
+ {
+ if(!other.empty()) {
+ ::new (static_cast<void*>(std::addressof(m_value))) value_type(std::move(other.value()));
+ m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket;
+ }
+ }
+
+ bucket_entry& operator=(const bucket_entry& other)
+ noexcept(std::is_nothrow_copy_constructible<value_type>::value)
+ {
+ if(this != &other) {
+ clear();
+
+ bucket_hash::operator=(other);
+ if(!other.empty()) {
+ ::new (static_cast<void*>(std::addressof(m_value))) value_type(other.value());
+ }
+
+ m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket;
+ m_last_bucket = other.m_last_bucket;
+ }
+
+ return *this;
+ }
+
+ bucket_entry& operator=(bucket_entry&& ) = delete;
+
+ ~bucket_entry() noexcept {
+ clear();
+ }
+
+ void clear() noexcept {
+ if(!empty()) {
+ destroy_value();
+ m_dist_from_ideal_bucket = EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET;
+ }
+ }
+
+ bool empty() const noexcept {
+ return m_dist_from_ideal_bucket == EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET;
+ }
+
+ value_type& value() noexcept {
+ tsl_assert(!empty());
+ return *reinterpret_cast<value_type*>(std::addressof(m_value));
+ }
+
+ const value_type& value() const noexcept {
+ tsl_assert(!empty());
+ return *reinterpret_cast<const value_type*>(std::addressof(m_value));
+ }
+
+ distance_type dist_from_ideal_bucket() const noexcept {
+ return m_dist_from_ideal_bucket;
+ }
+
+ bool last_bucket() const noexcept {
+ return m_last_bucket;
+ }
+
+ void set_as_last_bucket() noexcept {
+ m_last_bucket = true;
+ }
+
+ template<typename... Args>
+ void set_value_of_empty_bucket(distance_type dist_from_ideal_bucket,
+ truncated_hash_type hash, Args&&... value_type_args)
+ {
+ tsl_assert(dist_from_ideal_bucket >= 0);
+ tsl_assert(empty());
+
+ ::new (static_cast<void*>(std::addressof(m_value))) value_type(std::forward<Args>(value_type_args)...);
+ this->set_hash(hash);
+ m_dist_from_ideal_bucket = dist_from_ideal_bucket;
+
+ tsl_assert(!empty());
+ }
+
+ void swap_with_value_in_bucket(distance_type& dist_from_ideal_bucket,
+ truncated_hash_type& hash, value_type& value)
+ {
+ tsl_assert(!empty());
+
+ using std::swap;
+ swap(value, this->value());
+ swap(dist_from_ideal_bucket, m_dist_from_ideal_bucket);
+
+ // Avoid warning of unused variable if StoreHash is false
+ (void) hash;
+ if(StoreHash) {
+ const truncated_hash_type tmp_hash = this->truncated_hash();
+ this->set_hash(hash);
+ hash = tmp_hash;
+ }
+ }
+
+ static truncated_hash_type truncate_hash(std::size_t hash) noexcept {
+ return truncated_hash_type(hash);
+ }
+
+private:
+ void destroy_value() noexcept {
+ tsl_assert(!empty());
+ value().~value_type();
+ }
+
+private:
+ using storage = typename std::aligned_storage<sizeof(value_type), alignof(value_type)>::type;
+
+ static const distance_type EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET = -1;
+
+ distance_type m_dist_from_ideal_bucket;
+ bool m_last_bucket;
+ storage m_value;
+};
+
+
+
+/**
+ * Internal common class used by `robin_map` and `robin_set`.
+ *
+ * ValueType is what will be stored by `robin_hash` (usually `std::pair<Key, T>` for map and `Key` for set).
+ *
+ * `KeySelect` should be a `FunctionObject` which takes a `ValueType` in parameter and returns a
+ * reference to the key.
+ *
+ * `ValueSelect` should be a `FunctionObject` which takes a `ValueType` in parameter and returns a
+ * reference to the value. `ValueSelect` should be void if there is no value (in a set for example).
+ *
+ * The strong exception guarantee only holds if the expression
+ * `std::is_nothrow_swappable<ValueType>::value && std::is_nothrow_move_constructible<ValueType>::value` is true.
+ *
+ * Behaviour is undefined if the destructor of `ValueType` throws.
+ */
+template<class ValueType,
+ class KeySelect,
+ class ValueSelect,
+ class Hash,
+ class KeyEqual,
+ class Allocator,
+ bool StoreHash,
+ class GrowthPolicy>
+class robin_hash: private Hash, private KeyEqual, private GrowthPolicy {
+private:
+ template<typename U>
+ using has_mapped_type = typename std::integral_constant<bool, !std::is_same<U, void>::value>;
+
+ static_assert(noexcept(std::declval<GrowthPolicy>().bucket_for_hash(std::size_t(0))), "GrowthPolicy::bucket_for_hash must be noexcept.");
+ static_assert(noexcept(std::declval<GrowthPolicy>().clear()), "GrowthPolicy::clear must be noexcept.");
+
+public:
+ template<bool IsConst>
+ class robin_iterator;
+
+ using key_type = typename KeySelect::key_type;
+ using value_type = ValueType;
+ using size_type = std::size_t;
+ using difference_type = std::ptrdiff_t;
+ using hasher = Hash;
+ using key_equal = KeyEqual;
+ using allocator_type = Allocator;
+ using reference = value_type&;
+ using const_reference = const value_type&;
+ using pointer = value_type*;
+ using const_pointer = const value_type*;
+ using iterator = robin_iterator<false>;
+ using const_iterator = robin_iterator<true>;
+
+
+private:
+ /**
+ * Either store the hash because we are asked by the `StoreHash` template parameter
+ * or store the hash because it doesn't cost us anything in size and can be used to speed up rehash.
+ */
+ static constexpr bool STORE_HASH = StoreHash ||
+ (
+ (sizeof(tsl::detail_robin_hash::bucket_entry<value_type, true>) ==
+ sizeof(tsl::detail_robin_hash::bucket_entry<value_type, false>))
+ &&
+ (sizeof(std::size_t) == sizeof(truncated_hash_type) ||
+ is_power_of_two_policy<GrowthPolicy>::value)
+ &&
+ // Don't store the hash for primitive types with default hash.
+ (!std::is_arithmetic<key_type>::value ||
+ !std::is_same<Hash, std::hash<key_type>>::value)
+ );
+
+ /**
+ * Only use the stored hash on lookup if we are explictly asked. We are not sure how slow
+ * the KeyEqual operation is. An extra comparison may slow things down with a fast KeyEqual.
+ */
+ static constexpr bool USE_STORED_HASH_ON_LOOKUP = StoreHash;
+
+ /**
+ * We can only use the hash on rehash if the size of the hash type is the same as the stored one or
+ * if we use a power of two modulo. In the case of the power of two modulo, we just mask
+ * the least significant bytes, we just have to check that the truncated_hash_type didn't truncated
+ * more bytes.
+ */
+ static bool USE_STORED_HASH_ON_REHASH(size_type bucket_count) {
+ (void) bucket_count;
+ if(STORE_HASH && sizeof(std::size_t) == sizeof(truncated_hash_type)) {
+ return true;
+ }
+ else if(STORE_HASH && is_power_of_two_policy<GrowthPolicy>::value) {
+ tsl_assert(bucket_count > 0);
+ return (bucket_count - 1) <= (std::numeric_limits<truncated_hash_type>::max)();
+ }
+ else {
+ return false;
+ }
+ }
+
+ using bucket_entry = tsl::detail_robin_hash::bucket_entry<value_type, STORE_HASH>;
+ using distance_type = typename bucket_entry::distance_type;
+
+ using buckets_allocator = typename std::allocator_traits<allocator_type>::template rebind_alloc<bucket_entry>;
+ using buckets_container_type = std::vector<bucket_entry, buckets_allocator>;
+
+
+public:
+ /**
+ * The 'operator*()' and 'operator->()' methods return a const reference and const pointer respectively to the
+ * stored value type.
+ *
+ * In case of a map, to get a mutable reference to the value associated to a key (the '.second' in the
+ * stored pair), you have to call 'value()'.
+ *
+ * The main reason for this is that if we returned a `std::pair<Key, T>&` instead
+ * of a `const std::pair<Key, T>&`, the user may modify the key which will put the map in a undefined state.
+ */
+ template<bool IsConst>
+ class robin_iterator {
+ friend class robin_hash;
+
+ private:
+ using iterator_bucket = typename std::conditional<IsConst,
+ typename buckets_container_type::const_iterator,
+ typename buckets_container_type::iterator>::type;
+
+
+ robin_iterator(iterator_bucket it) noexcept: m_iterator(it) {
+ }
+
+ public:
+ using iterator_category = std::forward_iterator_tag;
+ using value_type = const typename robin_hash::value_type;
+ using difference_type = std::ptrdiff_t;
+ using reference = value_type&;
+ using pointer = value_type*;
+
+
+ robin_iterator() noexcept {
+ }
+
+ robin_iterator(const robin_iterator<false>& other) noexcept: m_iterator(other.m_iterator) {
+ }
+
+ const typename robin_hash::key_type& key() const {
+ return KeySelect()(m_iterator->value());
+ }
+
+ template<class U = ValueSelect, typename std::enable_if<has_mapped_type<U>::value && IsConst>::type* = nullptr>
+ const typename U::value_type& value() const {
+ return U()(m_iterator->value());
+ }
+
+ template<class U = ValueSelect, typename std::enable_if<has_mapped_type<U>::value && !IsConst>::type* = nullptr>
+ typename U::value_type& value() {
+ return U()(m_iterator->value());
+ }
+
+ reference operator*() const {
+ return m_iterator->value();
+ }
+
+ pointer operator->() const {
+ return std::addressof(m_iterator->value());
+ }
+
+ robin_iterator& operator++() {
+ while(true) {
+ if(m_iterator->last_bucket()) {
+ ++m_iterator;
+ return *this;
+ }
+
+ ++m_iterator;
+ if(!m_iterator->empty()) {
+ return *this;
+ }
+ }
+ }
+
+ robin_iterator operator++(int) {
+ robin_iterator tmp(*this);
+ ++*this;
+
+ return tmp;
+ }
+
+ friend bool operator==(const robin_iterator& lhs, const robin_iterator& rhs) {
+ return lhs.m_iterator == rhs.m_iterator;
+ }
+
+ friend bool operator!=(const robin_iterator& lhs, const robin_iterator& rhs) {
+ return !(lhs == rhs);
+ }
+
+ private:
+ iterator_bucket m_iterator;
+ };
+
+
+public:
+ robin_hash(size_type bucket_count,
+ const Hash& hash,
+ const KeyEqual& equal,
+ const Allocator& alloc,
+ float max_load_factor): Hash(hash),
+ KeyEqual(equal),
+ GrowthPolicy(bucket_count),
+ m_buckets(alloc),
+ m_first_or_empty_bucket(static_empty_bucket_ptr()),
+ m_bucket_count(bucket_count),
+ m_nb_elements(0),
+ m_grow_on_next_insert(false)
+ {
+ if(bucket_count > max_bucket_count()) {
+ TSL_THROW_OR_TERMINATE(std::length_error, "The map exceeds its maxmimum size.");
+ }
+
+ if(m_bucket_count > 0) {
+ /*
+ * We can't use the `vector(size_type count, const Allocator& alloc)` constructor
+ * as it's only available in C++14 and we need to support C++11. We thus must resize after using
+ * the `vector(const Allocator& alloc)` constructor.
+ *
+ * We can't use `vector(size_type count, const T& value, const Allocator& alloc)` as it requires the
+ * value T to be copyable.
+ */
+ m_buckets.resize(m_bucket_count);
+ m_first_or_empty_bucket = m_buckets.data();
+
+ tsl_assert(!m_buckets.empty());
+ m_buckets.back().set_as_last_bucket();
+ }
+
+
+ this->max_load_factor(max_load_factor);
+ }
+
+ robin_hash(const robin_hash& other): Hash(other),
+ KeyEqual(other),
+ GrowthPolicy(other),
+ m_buckets(other.m_buckets),
+ m_first_or_empty_bucket(m_buckets.empty()?static_empty_bucket_ptr():m_buckets.data()),
+ m_bucket_count(other.m_bucket_count),
+ m_nb_elements(other.m_nb_elements),
+ m_load_threshold(other.m_load_threshold),
+ m_max_load_factor(other.m_max_load_factor),
+ m_grow_on_next_insert(other.m_grow_on_next_insert)
+ {
+ }
+
+ robin_hash(robin_hash&& other) noexcept(std::is_nothrow_move_constructible<Hash>::value &&
+ std::is_nothrow_move_constructible<KeyEqual>::value &&
+ std::is_nothrow_move_constructible<GrowthPolicy>::value &&
+ std::is_nothrow_move_constructible<buckets_container_type>::value)
+ : Hash(std::move(static_cast<Hash&>(other))),
+ KeyEqual(std::move(static_cast<KeyEqual&>(other))),
+ GrowthPolicy(std::move(static_cast<GrowthPolicy&>(other))),
+ m_buckets(std::move(other.m_buckets)),
+ m_first_or_empty_bucket(m_buckets.empty()?static_empty_bucket_ptr():m_buckets.data()),
+ m_bucket_count(other.m_bucket_count),
+ m_nb_elements(other.m_nb_elements),
+ m_load_threshold(other.m_load_threshold),
+ m_max_load_factor(other.m_max_load_factor),
+ m_grow_on_next_insert(other.m_grow_on_next_insert)
+ {
+ other.GrowthPolicy::clear();
+ other.m_buckets.clear();
+ other.m_first_or_empty_bucket = static_empty_bucket_ptr();
+ other.m_bucket_count = 0;
+ other.m_nb_elements = 0;
+ other.m_load_threshold = 0;
+ other.m_grow_on_next_insert = false;
+ }
+
+ robin_hash& operator=(const robin_hash& other) {
+ if(&other != this) {
+ Hash::operator=(other);
+ KeyEqual::operator=(other);
+ GrowthPolicy::operator=(other);
+
+ m_buckets = other.m_buckets;
+ m_first_or_empty_bucket = m_buckets.empty()?static_empty_bucket_ptr():
+ m_buckets.data();
+ m_bucket_count = other.m_bucket_count;
+ m_nb_elements = other.m_nb_elements;
+ m_load_threshold = other.m_load_threshold;
+ m_max_load_factor = other.m_max_load_factor;
+ m_grow_on_next_insert = other.m_grow_on_next_insert;
+ }
+
+ return *this;
+ }
+
+ robin_hash& operator=(robin_hash&& other) {
+ other.swap(*this);
+ other.clear();
+
+ return *this;
+ }
+
+ allocator_type get_allocator() const {
+ return m_buckets.get_allocator();
+ }
+
+
+ /*
+ * Iterators
+ */
+ iterator begin() noexcept {
+ auto begin = m_buckets.begin();
+ while(begin != m_buckets.end() && begin->empty()) {
+ ++begin;
+ }
+
+ return iterator(begin);
+ }
+
+ const_iterator begin() const noexcept {
+ return cbegin();
+ }
+
+ const_iterator cbegin() const noexcept {
+ auto begin = m_buckets.cbegin();
+ while(begin != m_buckets.cend() && begin->empty()) {
+ ++begin;
+ }
+
+ return const_iterator(begin);
+ }
+
+ iterator end() noexcept {
+ return iterator(m_buckets.end());
+ }
+
+ const_iterator end() const noexcept {
+ return cend();
+ }
+
+ const_iterator cend() const noexcept {
+ return const_iterator(m_buckets.cend());
+ }
+
+
+ /*
+ * Capacity
+ */
+ bool empty() const noexcept {
+ return m_nb_elements == 0;
+ }
+
+ size_type size() const noexcept {
+ return m_nb_elements;
+ }
+
+ size_type max_size() const noexcept {
+ return m_buckets.max_size();
+ }
+
+ /*
+ * Modifiers
+ */
+ void clear() noexcept {
+ for(auto& bucket: m_buckets) {
+ bucket.clear();
+ }
+
+ m_nb_elements = 0;
+ m_grow_on_next_insert = false;
+ }
+
+
+
+ template<typename P>
+ std::pair<iterator, bool> insert(P&& value) {
+ return insert_impl(KeySelect()(value), std::forward<P>(value));
+ }
+
+ template<typename P>
+ iterator insert(const_iterator hint, P&& value) {
+ if(hint != cend() && compare_keys(KeySelect()(*hint), KeySelect()(value))) {
+ return mutable_iterator(hint);
+ }
+
+ return insert(std::forward<P>(value)).first;
+ }
+
+ template<class InputIt>
+ void insert(InputIt first, InputIt last) {
+ if(std::is_base_of<std::forward_iterator_tag,
+ typename std::iterator_traits<InputIt>::iterator_category>::value)
+ {
+ const auto nb_elements_insert = std::distance(first, last);
+ const size_type nb_free_buckets = m_load_threshold - size();
+ tsl_assert(m_load_threshold >= size());
+
+ if(nb_elements_insert > 0 && nb_free_buckets < size_type(nb_elements_insert)) {
+ reserve(size() + size_type(nb_elements_insert));
+ }
+ }
+
+ for(; first != last; ++first) {
+ insert(*first);
+ }
+ }
+
+
+
+ template<class K, class M>
+ std::pair<iterator, bool> insert_or_assign(K&& key, M&& obj) {
+ auto it = try_emplace(std::forward<K>(key), std::forward<M>(obj));
+ if(!it.second) {
+ it.first.value() = std::forward<M>(obj);
+ }
+
+ return it;
+ }
+
+ template<class K, class M>
+ iterator insert_or_assign(const_iterator hint, K&& key, M&& obj) {
+ if(hint != cend() && compare_keys(KeySelect()(*hint), key)) {
+ auto it = mutable_iterator(hint);
+ it.value() = std::forward<M>(obj);
+
+ return it;
+ }
+
+ return insert_or_assign(std::forward<K>(key), std::forward<M>(obj)).first;
+ }
+
+
+ template<class... Args>
+ std::pair<iterator, bool> emplace(Args&&... args) {
+ return insert(value_type(std::forward<Args>(args)...));
+ }
+
+ template<class... Args>
+ iterator emplace_hint(const_iterator hint, Args&&... args) {
+ return insert(hint, value_type(std::forward<Args>(args)...));
+ }
+
+
+
+ template<class K, class... Args>
+ std::pair<iterator, bool> try_emplace(K&& key, Args&&... args) {
+ return insert_impl(key, std::piecewise_construct,
+ std::forward_as_tuple(std::forward<K>(key)),
+ std::forward_as_tuple(std::forward<Args>(args)...));
+ }
+
+ template<class K, class... Args>
+ iterator try_emplace(const_iterator hint, K&& key, Args&&... args) {
+ if(hint != cend() && compare_keys(KeySelect()(*hint), key)) {
+ return mutable_iterator(hint);
+ }
+
+ return try_emplace(std::forward<K>(key), std::forward<Args>(args)...).first;
+ }
+
+ /**
+ * Here to avoid `template<class K> size_type erase(const K& key)` being used when
+ * we use an `iterator` instead of a `const_iterator`.
+ */
+ iterator erase(iterator pos) {
+ erase_from_bucket(pos);
+
+ /**
+ * Erase bucket used a backward shift after clearing the bucket.
+ * Check if there is a new value in the bucket, if not get the next non-empty.
+ */
+ if(pos.m_iterator->empty()) {
+ ++pos;
+ }
+
+ return pos;
+ }
+
+ iterator erase(const_iterator pos) {
+ return erase(mutable_iterator(pos));
+ }
+
+ iterator erase(const_iterator first, const_iterator last) {
+ if(first == last) {
+ return mutable_iterator(first);
+ }
+
+ auto first_mutable = mutable_iterator(first);
+ auto last_mutable = mutable_iterator(last);
+ for(auto it = first_mutable.m_iterator; it != last_mutable.m_iterator; ++it) {
+ if(!it->empty()) {
+ it->clear();
+ m_nb_elements--;
+ }
+ }
+
+ if(last_mutable == end()) {
+ return end();
+ }
+
+
+ /*
+ * Backward shift on the values which come after the deleted values.
+ * We try to move the values closer to their ideal bucket.
+ */
+ std::size_t icloser_bucket = std::size_t(std::distance(m_buckets.begin(), first_mutable.m_iterator));
+ std::size_t ito_move_closer_value = std::size_t(std::distance(m_buckets.begin(), last_mutable.m_iterator));
+ tsl_assert(ito_move_closer_value > icloser_bucket);
+
+ const std::size_t ireturn_bucket = ito_move_closer_value -
+ (std::min)(ito_move_closer_value - icloser_bucket,
+ std::size_t(m_buckets[ito_move_closer_value].dist_from_ideal_bucket()));
+
+ while(ito_move_closer_value < m_buckets.size() && m_buckets[ito_move_closer_value].dist_from_ideal_bucket() > 0) {
+ icloser_bucket = ito_move_closer_value -
+ (std::min)(ito_move_closer_value - icloser_bucket,
+ std::size_t(m_buckets[ito_move_closer_value].dist_from_ideal_bucket()));
+
+
+ tsl_assert(m_buckets[icloser_bucket].empty());
+ const distance_type new_distance = distance_type(m_buckets[ito_move_closer_value].dist_from_ideal_bucket() -
+ (ito_move_closer_value - icloser_bucket));
+ m_buckets[icloser_bucket].set_value_of_empty_bucket(new_distance,
+ m_buckets[ito_move_closer_value].truncated_hash(),
+ std::move(m_buckets[ito_move_closer_value].value()));
+ m_buckets[ito_move_closer_value].clear();
+
+
+ ++icloser_bucket;
+ ++ito_move_closer_value;
+ }
+
+
+ return iterator(m_buckets.begin() + ireturn_bucket);
+ }
+
+
+ template<class K>
+ size_type erase(const K& key) {
+ return erase(key, hash_key(key));
+ }
+
+ template<class K>
+ size_type erase(const K& key, std::size_t hash) {
+ auto it = find(key, hash);
+ if(it != end()) {
+ erase_from_bucket(it);
+
+ return 1;
+ }
+ else {
+ return 0;
+ }
+ }
+
+
+
+
+
+ void swap(robin_hash& other) {
+ using std::swap;
+
+ swap(static_cast<Hash&>(*this), static_cast<Hash&>(other));
+ swap(static_cast<KeyEqual&>(*this), static_cast<KeyEqual&>(other));
+ swap(static_cast<GrowthPolicy&>(*this), static_cast<GrowthPolicy&>(other));
+ swap(m_buckets, other.m_buckets);
+ swap(m_first_or_empty_bucket, other.m_first_or_empty_bucket);
+ swap(m_bucket_count, other.m_bucket_count);
+ swap(m_nb_elements, other.m_nb_elements);
+ swap(m_load_threshold, other.m_load_threshold);
+ swap(m_max_load_factor, other.m_max_load_factor);
+ swap(m_grow_on_next_insert, other.m_grow_on_next_insert);
+ }
+
+
+ /*
+ * Lookup
+ */
+ template<class K, class U = ValueSelect, typename std::enable_if<has_mapped_type<U>::value>::type* = nullptr>
+ typename U::value_type& at(const K& key) {
+ return at(key, hash_key(key));
+ }
+
+ template<class K, class U = ValueSelect, typename std::enable_if<has_mapped_type<U>::value>::type* = nullptr>
+ typename U::value_type& at(const K& key, std::size_t hash) {
+ return const_cast<typename U::value_type&>(static_cast<const robin_hash*>(this)->at(key, hash));
+ }
+
+
+ template<class K, class U = ValueSelect, typename std::enable_if<has_mapped_type<U>::value>::type* = nullptr>
+ const typename U::value_type& at(const K& key) const {
+ return at(key, hash_key(key));
+ }
+
+ template<class K, class U = ValueSelect, typename std::enable_if<has_mapped_type<U>::value>::type* = nullptr>
+ const typename U::value_type& at(const K& key, std::size_t hash) const {
+ auto it = find(key, hash);
+ if(it != cend()) {
+ return it.value();
+ }
+ else {
+ TSL_THROW_OR_TERMINATE(std::out_of_range, "Couldn't find key.");
+ }
+ }
+
+ template<class K, class U = ValueSelect, typename std::enable_if<has_mapped_type<U>::value>::type* = nullptr>
+ typename U::value_type& operator[](K&& key) {
+ return try_emplace(std::forward<K>(key)).first.value();
+ }
+
+
+ template<class K>
+ size_type count(const K& key) const {
+ return count(key, hash_key(key));
+ }
+
+ template<class K>
+ size_type count(const K& key, std::size_t hash) const {
+ if(find(key, hash) != cend()) {
+ return 1;
+ }
+ else {
+ return 0;
+ }
+ }
+
+
+ template<class K>
+ iterator find(const K& key) {
+ return find_impl(key, hash_key(key));
+ }
+
+ template<class K>
+ iterator find(const K& key, std::size_t hash) {
+ return find_impl(key, hash);
+ }
+
+
+ template<class K>
+ const_iterator find(const K& key) const {
+ return find_impl(key, hash_key(key));
+ }
+
+ template<class K>
+ const_iterator find(const K& key, std::size_t hash) const {
+ return find_impl(key, hash);
+ }
+
+
+ template<class K>
+ std::pair<iterator, iterator> equal_range(const K& key) {
+ return equal_range(key, hash_key(key));
+ }
+
+ template<class K>
+ std::pair<iterator, iterator> equal_range(const K& key, std::size_t hash) {
+ iterator it = find(key, hash);
+ return std::make_pair(it, (it == end())?it:std::next(it));
+ }
+
+
+ template<class K>
+ std::pair<const_iterator, const_iterator> equal_range(const K& key) const {
+ return equal_range(key, hash_key(key));
+ }
+
+ template<class K>
+ std::pair<const_iterator, const_iterator> equal_range(const K& key, std::size_t hash) const {
+ const_iterator it = find(key, hash);
+ return std::make_pair(it, (it == cend())?it:std::next(it));
+ }
+
+ /*
+ * Bucket interface
+ */
+ size_type bucket_count() const {
+ return m_bucket_count;
+ }
+
+ size_type max_bucket_count() const {
+ return (std::min)(GrowthPolicy::max_bucket_count(), m_buckets.max_size());
+ }
+
+ /*
+ * Hash policy
+ */
+ float load_factor() const {
+ if(bucket_count() == 0) {
+ return 0;
+ }
+
+ return float(m_nb_elements)/float(bucket_count());
+ }
+
+ float max_load_factor() const {
+ return m_max_load_factor;
+ }
+
+ void max_load_factor(float ml) {
+ m_max_load_factor = (std::max)(0.1f, (std::min)(ml, 0.95f));
+ m_load_threshold = size_type(float(bucket_count())*m_max_load_factor);
+ }
+
+ void rehash(size_type count) {
+ count = (std::max)(count, size_type(std::ceil(float(size())/max_load_factor())));
+ rehash_impl(count);
+ }
+
+ void reserve(size_type count) {
+ rehash(size_type(std::ceil(float(count)/max_load_factor())));
+ }
+
+ /*
+ * Observers
+ */
+ hasher hash_function() const {
+ return static_cast<const Hash&>(*this);
+ }
+
+ key_equal key_eq() const {
+ return static_cast<const KeyEqual&>(*this);
+ }
+
+
+ /*
+ * Other
+ */
+ iterator mutable_iterator(const_iterator pos) {
+ return iterator(m_buckets.begin() + std::distance(m_buckets.cbegin(), pos.m_iterator));
+ }
+
+private:
+ template<class K>
+ std::size_t hash_key(const K& key) const {
+ return Hash::operator()(key);
+ }
+
+ template<class K1, class K2>
+ bool compare_keys(const K1& key1, const K2& key2) const {
+ return KeyEqual::operator()(key1, key2);
+ }
+
+ std::size_t bucket_for_hash(std::size_t hash) const {
+ const std::size_t bucket = GrowthPolicy::bucket_for_hash(hash);
+ tsl_assert(bucket < m_buckets.size() || (bucket == 0 && m_buckets.empty()));
+
+ return bucket;
+ }
+
+ template<class U = GrowthPolicy, typename std::enable_if<is_power_of_two_policy<U>::value>::type* = nullptr>
+ std::size_t next_bucket(std::size_t index) const noexcept {
+ tsl_assert(index < bucket_count());
+
+ return (index + 1) & this->m_mask;
+ }
+
+ template<class U = GrowthPolicy, typename std::enable_if<!is_power_of_two_policy<U>::value>::type* = nullptr>
+ std::size_t next_bucket(std::size_t index) const noexcept {
+ tsl_assert(index < bucket_count());
+
+ index++;
+ return (index != bucket_count())?index:0;
+ }
+
+
+
+ template<class K>
+ iterator find_impl(const K& key, std::size_t hash) {
+ return mutable_iterator(static_cast<const robin_hash*>(this)->find(key, hash));
+ }
+
+ template<class K>
+ const_iterator find_impl(const K& key, std::size_t hash) const {
+ std::size_t ibucket = bucket_for_hash(hash);
+ distance_type dist_from_ideal_bucket = 0;
+
+ while(dist_from_ideal_bucket <= (m_first_or_empty_bucket + ibucket)->dist_from_ideal_bucket()) {
+ if(TSL_LIKELY((!USE_STORED_HASH_ON_LOOKUP || (m_first_or_empty_bucket + ibucket)->bucket_hash_equal(hash)) &&
+ compare_keys(KeySelect()((m_first_or_empty_bucket + ibucket)->value()), key)))
+ {
+ return const_iterator(m_buckets.begin() + ibucket);
+ }
+
+ ibucket = next_bucket(ibucket);
+ dist_from_ideal_bucket++;
+ }
+
+ return cend();
+ }
+
+ void erase_from_bucket(iterator pos) {
+ pos.m_iterator->clear();
+ m_nb_elements--;
+
+ /**
+ * Backward shift, swap the empty bucket, previous_ibucket, with the values on its right, ibucket,
+ * until we cross another empty bucket or if the other bucket has a distance_from_ideal_bucket == 0.
+ *
+ * We try to move the values closer to their ideal bucket.
+ */
+ std::size_t previous_ibucket = std::size_t(std::distance(m_buckets.begin(), pos.m_iterator));
+ std::size_t ibucket = next_bucket(previous_ibucket);
+
+ while(m_buckets[ibucket].dist_from_ideal_bucket() > 0) {
+ tsl_assert(m_buckets[previous_ibucket].empty());
+
+ const distance_type new_distance = distance_type(m_buckets[ibucket].dist_from_ideal_bucket() - 1);
+ m_buckets[previous_ibucket].set_value_of_empty_bucket(new_distance, m_buckets[ibucket].truncated_hash(),
+ std::move(m_buckets[ibucket].value()));
+ m_buckets[ibucket].clear();
+
+ previous_ibucket = ibucket;
+ ibucket = next_bucket(ibucket);
+ }
+ }
+
+ template<class K, class... Args>
+ std::pair<iterator, bool> insert_impl(const K& key, Args&&... value_type_args) {
+ const std::size_t hash = hash_key(key);
+
+ std::size_t ibucket = bucket_for_hash(hash);
+ distance_type dist_from_ideal_bucket = 0;
+
+ while(dist_from_ideal_bucket <= (m_first_or_empty_bucket + ibucket)->dist_from_ideal_bucket()) {
+ if((!USE_STORED_HASH_ON_LOOKUP || (m_first_or_empty_bucket + ibucket)->bucket_hash_equal(hash)) &&
+ compare_keys(KeySelect()((m_first_or_empty_bucket + ibucket)->value()), key))
+ {
+ return std::make_pair(iterator(m_buckets.begin() + ibucket), false);
+ }
+
+ ibucket = next_bucket(ibucket);
+ dist_from_ideal_bucket++;
+ }
+
+ if(grow_on_high_load()) {
+ ibucket = bucket_for_hash(hash);
+ dist_from_ideal_bucket = 0;
+
+ while(dist_from_ideal_bucket <= (m_first_or_empty_bucket + ibucket)->dist_from_ideal_bucket()) {
+ ibucket = next_bucket(ibucket);
+ dist_from_ideal_bucket++;
+ }
+ }
+
+
+ if((m_first_or_empty_bucket + ibucket)->empty()) {
+ (m_first_or_empty_bucket + ibucket)->set_value_of_empty_bucket(dist_from_ideal_bucket, bucket_entry::truncate_hash(hash),
+ std::forward<Args>(value_type_args)...);
+ }
+ else {
+ insert_value(ibucket, dist_from_ideal_bucket, bucket_entry::truncate_hash(hash),
+ std::forward<Args>(value_type_args)...);
+ }
+
+
+ m_nb_elements++;
+ /*
+ * The value will be inserted in ibucket in any case, either because it was
+ * empty or by stealing the bucket (robin hood).
+ */
+ return std::make_pair(iterator(m_buckets.begin() + ibucket), true);
+ }
+
+
+ template<class... Args>
+ void insert_value(std::size_t ibucket, distance_type dist_from_ideal_bucket,
+ truncated_hash_type hash, Args&&... value_type_args)
+ {
+ insert_value(ibucket, dist_from_ideal_bucket, hash, value_type(std::forward<Args>(value_type_args)...));
+ }
+
+ void insert_value(std::size_t ibucket, distance_type dist_from_ideal_bucket,
+ truncated_hash_type hash, value_type&& value)
+ {
+ m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value);
+ ibucket = next_bucket(ibucket);
+ dist_from_ideal_bucket++;
+
+ while(!m_buckets[ibucket].empty()) {
+ if(dist_from_ideal_bucket > m_buckets[ibucket].dist_from_ideal_bucket()) {
+ if(dist_from_ideal_bucket >= REHASH_ON_HIGH_NB_PROBES__NPROBES &&
+ load_factor() >= REHASH_ON_HIGH_NB_PROBES__MIN_LOAD_FACTOR)
+ {
+ /**
+ * The number of probes is really high, rehash the map on the next insert.
+ * Difficult to do now as rehash may throw an exception.
+ */
+ m_grow_on_next_insert = true;
+ }
+
+ m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value);
+ }
+
+ ibucket = next_bucket(ibucket);
+ dist_from_ideal_bucket++;
+ }
+
+ m_buckets[ibucket].set_value_of_empty_bucket(dist_from_ideal_bucket, hash, std::move(value));
+ }
+
+
+ void rehash_impl(size_type count) {
+ robin_hash new_table(count, static_cast<Hash&>(*this), static_cast<KeyEqual&>(*this),
+ get_allocator(), m_max_load_factor);
+
+ const bool use_stored_hash = USE_STORED_HASH_ON_REHASH(new_table.bucket_count());
+ for(auto& bucket: m_buckets) {
+ if(bucket.empty()) {
+ continue;
+ }
+
+ const std::size_t hash = use_stored_hash?bucket.truncated_hash():
+ new_table.hash_key(KeySelect()(bucket.value()));
+
+ new_table.insert_value_on_rehash(new_table.bucket_for_hash(hash), 0,
+ bucket_entry::truncate_hash(hash), std::move(bucket.value()));
+ }
+
+ new_table.m_nb_elements = m_nb_elements;
+ new_table.swap(*this);
+ }
+
+ void insert_value_on_rehash(std::size_t ibucket, distance_type dist_from_ideal_bucket,
+ truncated_hash_type hash, value_type&& value)
+ {
+ while(true) {
+ if(dist_from_ideal_bucket > m_buckets[ibucket].dist_from_ideal_bucket()) {
+ if(m_buckets[ibucket].empty()) {
+ m_buckets[ibucket].set_value_of_empty_bucket(dist_from_ideal_bucket, hash, std::move(value));
+ return;
+ }
+ else {
+ m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value);
+ }
+ }
+
+ dist_from_ideal_bucket++;
+ ibucket = next_bucket(ibucket);
+ }
+ }
+
+
+
+ /**
+ * Return true if the map has been rehashed.
+ */
+ bool grow_on_high_load() {
+ if(m_grow_on_next_insert || size() >= m_load_threshold) {
+ rehash_impl(GrowthPolicy::next_bucket_count());
+ m_grow_on_next_insert = false;
+
+ return true;
+ }
+
+ return false;
+ }
+
+
+public:
+ static const size_type DEFAULT_INIT_BUCKETS_SIZE = 16;
+ static constexpr float DEFAULT_MAX_LOAD_FACTOR = 0.5f;
+
+private:
+ static const distance_type REHASH_ON_HIGH_NB_PROBES__NPROBES = 128;
+ static constexpr float REHASH_ON_HIGH_NB_PROBES__MIN_LOAD_FACTOR = 0.15f;
+
+
+ /**
+ * Return an always valid pointer to an static empty bucket_entry with last_bucket() == true.
+ */
+ bucket_entry* static_empty_bucket_ptr() {
+ static bucket_entry empty_bucket(true);
+ return &empty_bucket;
+ }
+
+private:
+ buckets_container_type m_buckets;
+
+ /**
+ * Points to m_buckets.data() if !m_buckets.empty() otherwise points to static_empty_bucket_ptr.
+ * This variable is useful to avoid the cost of checking if m_buckets is empty when trying
+ * to find an element.
+ */
+ bucket_entry* m_first_or_empty_bucket;
+
+ /**
+ * Used a lot in find, avoid the call to m_buckets.size() which is a bit slower.
+ */
+ size_type m_bucket_count;
+
+ size_type m_nb_elements;
+
+ size_type m_load_threshold;
+ float m_max_load_factor;
+
+ bool m_grow_on_next_insert;
+};
+
+}
+
+}
+
+#endif
diff --git a/be/src/extern/diskann/include/tsl/robin_map.h b/be/src/extern/diskann/include/tsl/robin_map.h
new file mode 100644
index 00000000000..5958e70f028
--- /dev/null
+++ b/be/src/extern/diskann/include/tsl/robin_map.h
@@ -0,0 +1,668 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2017 Tessil
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#ifndef TSL_ROBIN_MAP_H
+#define TSL_ROBIN_MAP_H
+
+
+#include <cstddef>
+#include <functional>
+#include <initializer_list>
+#include <memory>
+#include <type_traits>
+#include <utility>
+#include "robin_hash.h"
+
+
+namespace tsl {
+
+
+/**
+ * Implementation of a hash map using open-adressing and the robin hood hashing algorithm with backward shift deletion.
+ *
+ * For operations modifying the hash map (insert, erase, rehash, ...), the strong exception guarantee
+ * is only guaranteed when the expression `std::is_nothrow_swappable<std::pair<Key, T>>::value &&
+ * std::is_nothrow_move_constructible<std::pair<Key, T>>::value` is true, otherwise if an exception
+ * is thrown during the swap or the move, the hash map may end up in a undefined state. Per the standard
+ * a `Key` or `T` with a noexcept copy constructor and no move constructor also satisfies the
+ * `std::is_nothrow_move_constructible<std::pair<Key, T>>::value` criterion (and will thus guarantee the
+ * strong exception for the map).
+ *
+ * When `StoreHash` is true, 32 bits of the hash are stored alongside the values. It can improve
+ * the performance during lookups if the `KeyEqual` function takes time (if it engenders a cache-miss for example)
+ * as we then compare the stored hashes before comparing the keys. When `tsl::rh::power_of_two_growth_policy` is used
+ * as `GrowthPolicy`, it may also speed-up the rehash process as we can avoid to recalculate the hash.
+ * When it is detected that storing the hash will not incur any memory penality due to alignement (i.e.
+ * `sizeof(tsl::detail_robin_hash::bucket_entry<ValueType, true>) ==
+ * sizeof(tsl::detail_robin_hash::bucket_entry<ValueType, false>)`) and `tsl::rh::power_of_two_growth_policy` is
+ * used, the hash will be stored even if `StoreHash` is false so that we can speed-up the rehash (but it will
+ * not be used on lookups unless `StoreHash` is true).
+ *
+ * `GrowthPolicy` defines how the map grows and consequently how a hash value is mapped to a bucket.
+ * By default the map uses `tsl::rh::power_of_two_growth_policy`. This policy keeps the number of buckets
+ * to a power of two and uses a mask to map the hash to a bucket instead of the slow modulo.
+ * Other growth policies are available and you may define your own growth policy,
+ * check `tsl::rh::power_of_two_growth_policy` for the interface.
+ *
+ * If the destructor of `Key` or `T` throws an exception, the behaviour of the class is undefined.
+ *
+ * Iterators invalidation:
+ * - clear, operator=, reserve, rehash: always invalidate the iterators.
+ * - insert, emplace, emplace_hint, operator[]: if there is an effective insert, invalidate the iterators.
+ * - erase: always invalidate the iterators.
+ */
+template<class Key,
+ class T,
+ class Hash = std::hash<Key>,
+ class KeyEqual = std::equal_to<Key>,
+ class Allocator = std::allocator<std::pair<Key, T>>,
+ bool StoreHash = false,
+ class GrowthPolicy = tsl::rh::power_of_two_growth_policy<2>>
+class robin_map {
+private:
+ template<typename U>
+ using has_is_transparent = tsl::detail_robin_hash::has_is_transparent<U>;
+
+ class KeySelect {
+ public:
+ using key_type = Key;
+
+ const key_type& operator()(const std::pair<Key, T>& key_value) const noexcept {
+ return key_value.first;
+ }
+
+ key_type& operator()(std::pair<Key, T>& key_value) noexcept {
+ return key_value.first;
+ }
+ };
+
+ class ValueSelect {
+ public:
+ using value_type = T;
+
+ const value_type& operator()(const std::pair<Key, T>& key_value) const noexcept {
+ return key_value.second;
+ }
+
+ value_type& operator()(std::pair<Key, T>& key_value) noexcept {
+ return key_value.second;
+ }
+ };
+
+ using ht = detail_robin_hash::robin_hash<std::pair<Key, T>, KeySelect, ValueSelect,
+ Hash, KeyEqual, Allocator, StoreHash, GrowthPolicy>;
+
+public:
+ using key_type = typename ht::key_type;
+ using mapped_type = T;
+ using value_type = typename ht::value_type;
+ using size_type = typename ht::size_type;
+ using difference_type = typename ht::difference_type;
+ using hasher = typename ht::hasher;
+ using key_equal = typename ht::key_equal;
+ using allocator_type = typename ht::allocator_type;
+ using reference = typename ht::reference;
+ using const_reference = typename ht::const_reference;
+ using pointer = typename ht::pointer;
+ using const_pointer = typename ht::const_pointer;
+ using iterator = typename ht::iterator;
+ using const_iterator = typename ht::const_iterator;
+
+
+public:
+ /*
+ * Constructors
+ */
+ robin_map(): robin_map(ht::DEFAULT_INIT_BUCKETS_SIZE) {
+ }
+
+ explicit robin_map(size_type bucket_count,
+ const Hash& hash = Hash(),
+ const KeyEqual& equal = KeyEqual(),
+ const Allocator& alloc = Allocator()):
+ m_ht(bucket_count, hash, equal, alloc, ht::DEFAULT_MAX_LOAD_FACTOR)
+ {
+ }
+
+ robin_map(size_type bucket_count,
+ const Allocator& alloc): robin_map(bucket_count, Hash(), KeyEqual(), alloc)
+ {
+ }
+
+ robin_map(size_type bucket_count,
+ const Hash& hash,
+ const Allocator& alloc): robin_map(bucket_count, hash, KeyEqual(), alloc)
+ {
+ }
+
+ explicit robin_map(const Allocator& alloc): robin_map(ht::DEFAULT_INIT_BUCKETS_SIZE, alloc) {
+ }
+
+ template<class InputIt>
+ robin_map(InputIt first, InputIt last,
+ size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE,
+ const Hash& hash = Hash(),
+ const KeyEqual& equal = KeyEqual(),
+ const Allocator& alloc = Allocator()): robin_map(bucket_count, hash, equal, alloc)
+ {
+ insert(first, last);
+ }
+
+ template<class InputIt>
+ robin_map(InputIt first, InputIt last,
+ size_type bucket_count,
+ const Allocator& alloc): robin_map(first, last, bucket_count, Hash(), KeyEqual(), alloc)
+ {
+ }
+
+ template<class InputIt>
+ robin_map(InputIt first, InputIt last,
+ size_type bucket_count,
+ const Hash& hash,
+ const Allocator& alloc): robin_map(first, last, bucket_count, hash, KeyEqual(), alloc)
+ {
+ }
+
+ robin_map(std::initializer_list<value_type> init,
+ size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE,
+ const Hash& hash = Hash(),
+ const KeyEqual& equal = KeyEqual(),
+ const Allocator& alloc = Allocator()):
+ robin_map(init.begin(), init.end(), bucket_count, hash, equal, alloc)
+ {
+ }
+
+ robin_map(std::initializer_list<value_type> init,
+ size_type bucket_count,
+ const Allocator& alloc):
+ robin_map(init.begin(), init.end(), bucket_count, Hash(), KeyEqual(), alloc)
+ {
+ }
+
+ robin_map(std::initializer_list<value_type> init,
+ size_type bucket_count,
+ const Hash& hash,
+ const Allocator& alloc):
+ robin_map(init.begin(), init.end(), bucket_count, hash, KeyEqual(), alloc)
+ {
+ }
+
+ robin_map& operator=(std::initializer_list<value_type> ilist) {
+ m_ht.clear();
+
+ m_ht.reserve(ilist.size());
+ m_ht.insert(ilist.begin(), ilist.end());
+
+ return *this;
+ }
+
+ allocator_type get_allocator() const { return m_ht.get_allocator(); }
+
+
+ /*
+ * Iterators
+ */
+ iterator begin() noexcept { return m_ht.begin(); }
+ const_iterator begin() const noexcept { return m_ht.begin(); }
+ const_iterator cbegin() const noexcept { return m_ht.cbegin(); }
+
+ iterator end() noexcept { return m_ht.end(); }
+ const_iterator end() const noexcept { return m_ht.end(); }
+ const_iterator cend() const noexcept { return m_ht.cend(); }
+
+
+ /*
+ * Capacity
+ */
+ bool empty() const noexcept { return m_ht.empty(); }
+ size_type size() const noexcept { return m_ht.size(); }
+ size_type max_size() const noexcept { return m_ht.max_size(); }
+
+ /*
+ * Modifiers
+ */
+ void clear() noexcept { m_ht.clear(); }
+
+
+
+ std::pair<iterator, bool> insert(const value_type& value) {
+ return m_ht.insert(value);
+ }
+
+ template<class P, typename std::enable_if<std::is_constructible<value_type, P&&>::value>::type* = nullptr>
+ std::pair<iterator, bool> insert(P&& value) {
+ return m_ht.emplace(std::forward<P>(value));
+ }
+
+ std::pair<iterator, bool> insert(value_type&& value) {
+ return m_ht.insert(std::move(value));
+ }
+
+
+ iterator insert(const_iterator hint, const value_type& value) {
+ return m_ht.insert(hint, value);
+ }
+
+ template<class P, typename std::enable_if<std::is_constructible<value_type, P&&>::value>::type* = nullptr>
+ iterator insert(const_iterator hint, P&& value) {
+ return m_ht.emplace_hint(hint, std::forward<P>(value));
+ }
+
+ iterator insert(const_iterator hint, value_type&& value) {
+ return m_ht.insert(hint, std::move(value));
+ }
+
+
+ template<class InputIt>
+ void insert(InputIt first, InputIt last) {
+ m_ht.insert(first, last);
+ }
+
+ void insert(std::initializer_list<value_type> ilist) {
+ m_ht.insert(ilist.begin(), ilist.end());
+ }
+
+
+
+
+ template<class M>
+ std::pair<iterator, bool> insert_or_assign(const key_type& k, M&& obj) {
+ return m_ht.insert_or_assign(k, std::forward<M>(obj));
+ }
+
+ template<class M>
+ std::pair<iterator, bool> insert_or_assign(key_type&& k, M&& obj) {
+ return m_ht.insert_or_assign(std::move(k), std::forward<M>(obj));
+ }
+
+ template<class M>
+ iterator insert_or_assign(const_iterator hint, const key_type& k, M&& obj) {
+ return m_ht.insert_or_assign(hint, k, std::forward<M>(obj));
+ }
+
+ template<class M>
+ iterator insert_or_assign(const_iterator hint, key_type&& k, M&& obj) {
+ return m_ht.insert_or_assign(hint, std::move(k), std::forward<M>(obj));
+ }
+
+
+
+ /**
+ * Due to the way elements are stored, emplace will need to move or copy the key-value once.
+ * The method is equivalent to insert(value_type(std::forward<Args>(args)...));
+ *
+ * Mainly here for compatibility with the std::unordered_map interface.
+ */
+ template<class... Args>
+ std::pair<iterator, bool> emplace(Args&&... args) {
+ return m_ht.emplace(std::forward<Args>(args)...);
+ }
+
+
+
+ /**
+ * Due to the way elements are stored, emplace_hint will need to move or copy the key-value once.
+ * The method is equivalent to insert(hint, value_type(std::forward<Args>(args)...));
+ *
+ * Mainly here for compatibility with the std::unordered_map interface.
+ */
+ template<class... Args>
+ iterator emplace_hint(const_iterator hint, Args&&... args) {
+ return m_ht.emplace_hint(hint, std::forward<Args>(args)...);
+ }
+
+
+
+
+ template<class... Args>
+ std::pair<iterator, bool> try_emplace(const key_type& k, Args&&... args) {
+ return m_ht.try_emplace(k, std::forward<Args>(args)...);
+ }
+
+ template<class... Args>
+ std::pair<iterator, bool> try_emplace(key_type&& k, Args&&... args) {
+ return m_ht.try_emplace(std::move(k), std::forward<Args>(args)...);
+ }
+
+ template<class... Args>
+ iterator try_emplace(const_iterator hint, const key_type& k, Args&&... args) {
+ return m_ht.try_emplace(hint, k, std::forward<Args>(args)...);
+ }
+
+ template<class... Args>
+ iterator try_emplace(const_iterator hint, key_type&& k, Args&&... args) {
+ return m_ht.try_emplace(hint, std::move(k), std::forward<Args>(args)...);
+ }
+
+
+
+
+ iterator erase(iterator pos) { return m_ht.erase(pos); }
+ iterator erase(const_iterator pos) { return m_ht.erase(pos); }
+ iterator erase(const_iterator first, const_iterator last) { return m_ht.erase(first, last); }
+ size_type erase(const key_type& key) { return m_ht.erase(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup to the value if you already have the hash.
+ */
+ size_type erase(const key_type& key, std::size_t precalculated_hash) {
+ return m_ht.erase(key, precalculated_hash);
+ }
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ size_type erase(const K& key) { return m_ht.erase(key); }
+
+ /**
+ * @copydoc erase(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup to the value if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ size_type erase(const K& key, std::size_t precalculated_hash) {
+ return m_ht.erase(key, precalculated_hash);
+ }
+
+
+
+ void swap(robin_map& other) { other.m_ht.swap(m_ht); }
+
+
+
+ /*
+ * Lookup
+ */
+ T& at(const Key& key) { return m_ht.at(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ T& at(const Key& key, std::size_t precalculated_hash) { return m_ht.at(key, precalculated_hash); }
+
+
+ const T& at(const Key& key) const { return m_ht.at(key); }
+
+ /**
+ * @copydoc at(const Key& key, std::size_t precalculated_hash)
+ */
+ const T& at(const Key& key, std::size_t precalculated_hash) const { return m_ht.at(key, precalculated_hash); }
+
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ T& at(const K& key) { return m_ht.at(key); }
+
+ /**
+ * @copydoc at(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ T& at(const K& key, std::size_t precalculated_hash) { return m_ht.at(key, precalculated_hash); }
+
+
+ /**
+ * @copydoc at(const K& key)
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ const T& at(const K& key) const { return m_ht.at(key); }
+
+ /**
+ * @copydoc at(const K& key, std::size_t precalculated_hash)
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ const T& at(const K& key, std::size_t precalculated_hash) const { return m_ht.at(key, precalculated_hash); }
+
+
+
+
+ T& operator[](const Key& key) { return m_ht[key]; }
+ T& operator[](Key&& key) { return m_ht[std::move(key)]; }
+
+
+
+
+ size_type count(const Key& key) const { return m_ht.count(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ size_type count(const Key& key, std::size_t precalculated_hash) const {
+ return m_ht.count(key, precalculated_hash);
+ }
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ size_type count(const K& key) const { return m_ht.count(key); }
+
+ /**
+ * @copydoc count(const K& key) const
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ size_type count(const K& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); }
+
+
+
+
+ iterator find(const Key& key) { return m_ht.find(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ iterator find(const Key& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); }
+
+ const_iterator find(const Key& key) const { return m_ht.find(key); }
+
+ /**
+ * @copydoc find(const Key& key, std::size_t precalculated_hash)
+ */
+ const_iterator find(const Key& key, std::size_t precalculated_hash) const {
+ return m_ht.find(key, precalculated_hash);
+ }
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ iterator find(const K& key) { return m_ht.find(key); }
+
+ /**
+ * @copydoc find(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ iterator find(const K& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); }
+
+ /**
+ * @copydoc find(const K& key)
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ const_iterator find(const K& key) const { return m_ht.find(key); }
+
+ /**
+ * @copydoc find(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ const_iterator find(const K& key, std::size_t precalculated_hash) const {
+ return m_ht.find(key, precalculated_hash);
+ }
+
+
+
+
+ std::pair<iterator, iterator> equal_range(const Key& key) { return m_ht.equal_range(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ std::pair<iterator, iterator> equal_range(const Key& key, std::size_t precalculated_hash) {
+ return m_ht.equal_range(key, precalculated_hash);
+ }
+
+ std::pair<const_iterator, const_iterator> equal_range(const Key& key) const { return m_ht.equal_range(key); }
+
+ /**
+ * @copydoc equal_range(const Key& key, std::size_t precalculated_hash)
+ */
+ std::pair<const_iterator, const_iterator> equal_range(const Key& key, std::size_t precalculated_hash) const {
+ return m_ht.equal_range(key, precalculated_hash);
+ }
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ std::pair<iterator, iterator> equal_range(const K& key) { return m_ht.equal_range(key); }
+
+
+ /**
+ * @copydoc equal_range(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ std::pair<iterator, iterator> equal_range(const K& key, std::size_t precalculated_hash) {
+ return m_ht.equal_range(key, precalculated_hash);
+ }
+
+ /**
+ * @copydoc equal_range(const K& key)
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ std::pair<const_iterator, const_iterator> equal_range(const K& key) const { return m_ht.equal_range(key); }
+
+ /**
+ * @copydoc equal_range(const K& key, std::size_t precalculated_hash)
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ std::pair<const_iterator, const_iterator> equal_range(const K& key, std::size_t precalculated_hash) const {
+ return m_ht.equal_range(key, precalculated_hash);
+ }
+
+
+
+
+ /*
+ * Bucket interface
+ */
+ size_type bucket_count() const { return m_ht.bucket_count(); }
+ size_type max_bucket_count() const { return m_ht.max_bucket_count(); }
+
+
+ /*
+ * Hash policy
+ */
+ float load_factor() const { return m_ht.load_factor(); }
+ float max_load_factor() const { return m_ht.max_load_factor(); }
+ void max_load_factor(float ml) { m_ht.max_load_factor(ml); }
+
+ void rehash(size_type count) { m_ht.rehash(count); }
+ void reserve(size_type count) { m_ht.reserve(count); }
+
+
+ /*
+ * Observers
+ */
+ hasher hash_function() const { return m_ht.hash_function(); }
+ key_equal key_eq() const { return m_ht.key_eq(); }
+
+ /*
+ * Other
+ */
+
+ /**
+ * Convert a const_iterator to an iterator.
+ */
+ iterator mutable_iterator(const_iterator pos) {
+ return m_ht.mutable_iterator(pos);
+ }
+
+ friend bool operator==(const robin_map& lhs, const robin_map& rhs) {
+ if(lhs.size() != rhs.size()) {
+ return false;
+ }
+
+ for(const auto& element_lhs: lhs) {
+ const auto it_element_rhs = rhs.find(element_lhs.first);
+ if(it_element_rhs == rhs.cend() || element_lhs.second != it_element_rhs->second) {
+ return false;
+ }
+ }
+
+ return true;
+ }
+
+ friend bool operator!=(const robin_map& lhs, const robin_map& rhs) {
+ return !operator==(lhs, rhs);
+ }
+
+ friend void swap(robin_map& lhs, robin_map& rhs) {
+ lhs.swap(rhs);
+ }
+
+private:
+ ht m_ht;
+};
+
+
+/**
+ * Same as `tsl::robin_map<Key, T, Hash, KeyEqual, Allocator, StoreHash, tsl::rh::prime_growth_policy>`.
+ */
+template<class Key,
+ class T,
+ class Hash = std::hash<Key>,
+ class KeyEqual = std::equal_to<Key>,
+ class Allocator = std::allocator<std::pair<Key, T>>,
+ bool StoreHash = false>
+using robin_pg_map = robin_map<Key, T, Hash, KeyEqual, Allocator, StoreHash, tsl::rh::prime_growth_policy>;
+
+} // end namespace tsl
+
+#endif
diff --git a/be/src/extern/diskann/include/tsl/robin_set.h b/be/src/extern/diskann/include/tsl/robin_set.h
new file mode 100644
index 00000000000..4e4667e2606
--- /dev/null
+++ b/be/src/extern/diskann/include/tsl/robin_set.h
@@ -0,0 +1,535 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2017 Tessil
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#ifndef TSL_ROBIN_SET_H
+#define TSL_ROBIN_SET_H
+
+
+#include <cstddef>
+#include <functional>
+#include <initializer_list>
+#include <memory>
+#include <type_traits>
+#include <utility>
+#include "robin_hash.h"
+
+
+namespace tsl {
+
+
+/**
+ * Implementation of a hash set using open-adressing and the robin hood hashing algorithm with backward shift deletion.
+ *
+ * For operations modifying the hash set (insert, erase, rehash, ...), the strong exception guarantee
+ * is only guaranteed when the expression `std::is_nothrow_swappable<Key>::value &&
+ * std::is_nothrow_move_constructible<Key>::value` is true, otherwise if an exception
+ * is thrown during the swap or the move, the hash set may end up in a undefined state. Per the standard
+ * a `Key` with a noexcept copy constructor and no move constructor also satisfies the
+ * `std::is_nothrow_move_constructible<Key>::value` criterion (and will thus guarantee the
+ * strong exception for the set).
+ *
+ * When `StoreHash` is true, 32 bits of the hash are stored alongside the values. It can improve
+ * the performance during lookups if the `KeyEqual` function takes time (or engenders a cache-miss for example)
+ * as we then compare the stored hashes before comparing the keys. When `tsl::rh::power_of_two_growth_policy` is used
+ * as `GrowthPolicy`, it may also speed-up the rehash process as we can avoid to recalculate the hash.
+ * When it is detected that storing the hash will not incur any memory penality due to alignement (i.e.
+ * `sizeof(tsl::detail_robin_hash::bucket_entry<ValueType, true>) ==
+ * sizeof(tsl::detail_robin_hash::bucket_entry<ValueType, false>)`) and `tsl::rh::power_of_two_growth_policy` is
+ * used, the hash will be stored even if `StoreHash` is false so that we can speed-up the rehash (but it will
+ * not be used on lookups unless `StoreHash` is true).
+ *
+ * `GrowthPolicy` defines how the set grows and consequently how a hash value is mapped to a bucket.
+ * By default the set uses `tsl::rh::power_of_two_growth_policy`. This policy keeps the number of buckets
+ * to a power of two and uses a mask to set the hash to a bucket instead of the slow modulo.
+ * Other growth policies are available and you may define your own growth policy,
+ * check `tsl::rh::power_of_two_growth_policy` for the interface.
+ *
+ * If the destructor of `Key` throws an exception, the behaviour of the class is undefined.
+ *
+ * Iterators invalidation:
+ * - clear, operator=, reserve, rehash: always invalidate the iterators.
+ * - insert, emplace, emplace_hint, operator[]: if there is an effective insert, invalidate the iterators.
+ * - erase: always invalidate the iterators.
+ */
+template<class Key,
+ class Hash = std::hash<Key>,
+ class KeyEqual = std::equal_to<Key>,
+ class Allocator = std::allocator<Key>,
+ bool StoreHash = false,
+ class GrowthPolicy = tsl::rh::power_of_two_growth_policy<2>>
+class robin_set {
+private:
+ template<typename U>
+ using has_is_transparent = tsl::detail_robin_hash::has_is_transparent<U>;
+
+ class KeySelect {
+ public:
+ using key_type = Key;
+
+ const key_type& operator()(const Key& key) const noexcept {
+ return key;
+ }
+
+ key_type& operator()(Key& key) noexcept {
+ return key;
+ }
+ };
+
+ using ht = detail_robin_hash::robin_hash<Key, KeySelect, void,
+ Hash, KeyEqual, Allocator, StoreHash, GrowthPolicy>;
+
+public:
+ using key_type = typename ht::key_type;
+ using value_type = typename ht::value_type;
+ using size_type = typename ht::size_type;
+ using difference_type = typename ht::difference_type;
+ using hasher = typename ht::hasher;
+ using key_equal = typename ht::key_equal;
+ using allocator_type = typename ht::allocator_type;
+ using reference = typename ht::reference;
+ using const_reference = typename ht::const_reference;
+ using pointer = typename ht::pointer;
+ using const_pointer = typename ht::const_pointer;
+ using iterator = typename ht::iterator;
+ using const_iterator = typename ht::const_iterator;
+
+
+ /*
+ * Constructors
+ */
+ robin_set(): robin_set(ht::DEFAULT_INIT_BUCKETS_SIZE) {
+ }
+
+ explicit robin_set(size_type bucket_count,
+ const Hash& hash = Hash(),
+ const KeyEqual& equal = KeyEqual(),
+ const Allocator& alloc = Allocator()):
+ m_ht(bucket_count, hash, equal, alloc, ht::DEFAULT_MAX_LOAD_FACTOR)
+ {
+ }
+
+ robin_set(size_type bucket_count,
+ const Allocator& alloc): robin_set(bucket_count, Hash(), KeyEqual(), alloc)
+ {
+ }
+
+ robin_set(size_type bucket_count,
+ const Hash& hash,
+ const Allocator& alloc): robin_set(bucket_count, hash, KeyEqual(), alloc)
+ {
+ }
+
+ explicit robin_set(const Allocator& alloc): robin_set(ht::DEFAULT_INIT_BUCKETS_SIZE, alloc) {
+ }
+
+ template<class InputIt>
+ robin_set(InputIt first, InputIt last,
+ size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE,
+ const Hash& hash = Hash(),
+ const KeyEqual& equal = KeyEqual(),
+ const Allocator& alloc = Allocator()): robin_set(bucket_count, hash, equal, alloc)
+ {
+ insert(first, last);
+ }
+
+ template<class InputIt>
+ robin_set(InputIt first, InputIt last,
+ size_type bucket_count,
+ const Allocator& alloc): robin_set(first, last, bucket_count, Hash(), KeyEqual(), alloc)
+ {
+ }
+
+ template<class InputIt>
+ robin_set(InputIt first, InputIt last,
+ size_type bucket_count,
+ const Hash& hash,
+ const Allocator& alloc): robin_set(first, last, bucket_count, hash, KeyEqual(), alloc)
+ {
+ }
+
+ robin_set(std::initializer_list<value_type> init,
+ size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE,
+ const Hash& hash = Hash(),
+ const KeyEqual& equal = KeyEqual(),
+ const Allocator& alloc = Allocator()):
+ robin_set(init.begin(), init.end(), bucket_count, hash, equal, alloc)
+ {
+ }
+
+ robin_set(std::initializer_list<value_type> init,
+ size_type bucket_count,
+ const Allocator& alloc):
+ robin_set(init.begin(), init.end(), bucket_count, Hash(), KeyEqual(), alloc)
+ {
+ }
+
+ robin_set(std::initializer_list<value_type> init,
+ size_type bucket_count,
+ const Hash& hash,
+ const Allocator& alloc):
+ robin_set(init.begin(), init.end(), bucket_count, hash, KeyEqual(), alloc)
+ {
+ }
+
+
+ robin_set& operator=(std::initializer_list<value_type> ilist) {
+ m_ht.clear();
+
+ m_ht.reserve(ilist.size());
+ m_ht.insert(ilist.begin(), ilist.end());
+
+ return *this;
+ }
+
+ allocator_type get_allocator() const { return m_ht.get_allocator(); }
+
+
+ /*
+ * Iterators
+ */
+ iterator begin() noexcept { return m_ht.begin(); }
+ const_iterator begin() const noexcept { return m_ht.begin(); }
+ const_iterator cbegin() const noexcept { return m_ht.cbegin(); }
+
+ iterator end() noexcept { return m_ht.end(); }
+ const_iterator end() const noexcept { return m_ht.end(); }
+ const_iterator cend() const noexcept { return m_ht.cend(); }
+
+
+ /*
+ * Capacity
+ */
+ bool empty() const noexcept { return m_ht.empty(); }
+ size_type size() const noexcept { return m_ht.size(); }
+ size_type max_size() const noexcept { return m_ht.max_size(); }
+
+ /*
+ * Modifiers
+ */
+ void clear() noexcept { m_ht.clear(); }
+
+
+
+
+ std::pair<iterator, bool> insert(const value_type& value) {
+ return m_ht.insert(value);
+ }
+
+ std::pair<iterator, bool> insert(value_type&& value) {
+ return m_ht.insert(std::move(value));
+ }
+
+ iterator insert(const_iterator hint, const value_type& value) {
+ return m_ht.insert(hint, value);
+ }
+
+ iterator insert(const_iterator hint, value_type&& value) {
+ return m_ht.insert(hint, std::move(value));
+ }
+
+ template<class InputIt>
+ void insert(InputIt first, InputIt last) {
+ m_ht.insert(first, last);
+ }
+
+ void insert(std::initializer_list<value_type> ilist) {
+ m_ht.insert(ilist.begin(), ilist.end());
+ }
+
+
+
+
+ /**
+ * Due to the way elements are stored, emplace will need to move or copy the key-value once.
+ * The method is equivalent to insert(value_type(std::forward<Args>(args)...));
+ *
+ * Mainly here for compatibility with the std::unordered_map interface.
+ */
+ template<class... Args>
+ std::pair<iterator, bool> emplace(Args&&... args) {
+ return m_ht.emplace(std::forward<Args>(args)...);
+ }
+
+
+
+ /**
+ * Due to the way elements are stored, emplace_hint will need to move or copy the key-value once.
+ * The method is equivalent to insert(hint, value_type(std::forward<Args>(args)...));
+ *
+ * Mainly here for compatibility with the std::unordered_map interface.
+ */
+ template<class... Args>
+ iterator emplace_hint(const_iterator hint, Args&&... args) {
+ return m_ht.emplace_hint(hint, std::forward<Args>(args)...);
+ }
+
+
+
+ iterator erase(iterator pos) { return m_ht.erase(pos); }
+ iterator erase(const_iterator pos) { return m_ht.erase(pos); }
+ iterator erase(const_iterator first, const_iterator last) { return m_ht.erase(first, last); }
+ size_type erase(const key_type& key) { return m_ht.erase(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup to the value if you already have the hash.
+ */
+ size_type erase(const key_type& key, std::size_t precalculated_hash) {
+ return m_ht.erase(key, precalculated_hash);
+ }
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ size_type erase(const K& key) { return m_ht.erase(key); }
+
+ /**
+ * @copydoc erase(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup to the value if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ size_type erase(const K& key, std::size_t precalculated_hash) {
+ return m_ht.erase(key, precalculated_hash);
+ }
+
+
+
+ void swap(robin_set& other) { other.m_ht.swap(m_ht); }
+
+
+
+ /*
+ * Lookup
+ */
+ size_type count(const Key& key) const { return m_ht.count(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ size_type count(const Key& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); }
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ size_type count(const K& key) const { return m_ht.count(key); }
+
+ /**
+ * @copydoc count(const K& key) const
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ size_type count(const K& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); }
+
+
+
+
+ iterator find(const Key& key) { return m_ht.find(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ iterator find(const Key& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); }
+
+ const_iterator find(const Key& key) const { return m_ht.find(key); }
+
+ /**
+ * @copydoc find(const Key& key, std::size_t precalculated_hash)
+ */
+ const_iterator find(const Key& key, std::size_t precalculated_hash) const { return m_ht.find(key, precalculated_hash); }
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ iterator find(const K& key) { return m_ht.find(key); }
+
+ /**
+ * @copydoc find(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ iterator find(const K& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); }
+
+ /**
+ * @copydoc find(const K& key)
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ const_iterator find(const K& key) const { return m_ht.find(key); }
+
+ /**
+ * @copydoc find(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ const_iterator find(const K& key, std::size_t precalculated_hash) const { return m_ht.find(key, precalculated_hash); }
+
+
+
+
+ std::pair<iterator, iterator> equal_range(const Key& key) { return m_ht.equal_range(key); }
+
+ /**
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ std::pair<iterator, iterator> equal_range(const Key& key, std::size_t precalculated_hash) {
+ return m_ht.equal_range(key, precalculated_hash);
+ }
+
+ std::pair<const_iterator, const_iterator> equal_range(const Key& key) const { return m_ht.equal_range(key); }
+
+ /**
+ * @copydoc equal_range(const Key& key, std::size_t precalculated_hash)
+ */
+ std::pair<const_iterator, const_iterator> equal_range(const Key& key, std::size_t precalculated_hash) const {
+ return m_ht.equal_range(key, precalculated_hash);
+ }
+
+ /**
+ * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists.
+ * If so, K must be hashable and comparable to Key.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ std::pair<iterator, iterator> equal_range(const K& key) { return m_ht.equal_range(key); }
+
+ /**
+ * @copydoc equal_range(const K& key)
+ *
+ * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same
+ * as hash_function()(key). Usefull to speed-up the lookup if you already have the hash.
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ std::pair<iterator, iterator> equal_range(const K& key, std::size_t precalculated_hash) {
+ return m_ht.equal_range(key, precalculated_hash);
+ }
+
+ /**
+ * @copydoc equal_range(const K& key)
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ std::pair<const_iterator, const_iterator> equal_range(const K& key) const { return m_ht.equal_range(key); }
+
+ /**
+ * @copydoc equal_range(const K& key, std::size_t precalculated_hash)
+ */
+ template<class K, class KE = KeyEqual, typename std::enable_if<has_is_transparent<KE>::value>::type* = nullptr>
+ std::pair<const_iterator, const_iterator> equal_range(const K& key, std::size_t precalculated_hash) const {
+ return m_ht.equal_range(key, precalculated_hash);
+ }
+
+
+
+
+ /*
+ * Bucket interface
+ */
+ size_type bucket_count() const { return m_ht.bucket_count(); }
+ size_type max_bucket_count() const { return m_ht.max_bucket_count(); }
+
+
+ /*
+ * Hash policy
+ */
+ float load_factor() const { return m_ht.load_factor(); }
+ float max_load_factor() const { return m_ht.max_load_factor(); }
+ void max_load_factor(float ml) { m_ht.max_load_factor(ml); }
+
+ void rehash(size_type count) { m_ht.rehash(count); }
+ void reserve(size_type count) { m_ht.reserve(count); }
+
+
+ /*
+ * Observers
+ */
+ hasher hash_function() const { return m_ht.hash_function(); }
+ key_equal key_eq() const { return m_ht.key_eq(); }
+
+
+ /*
+ * Other
+ */
+
+ /**
+ * Convert a const_iterator to an iterator.
+ */
+ iterator mutable_iterator(const_iterator pos) {
+ return m_ht.mutable_iterator(pos);
+ }
+
+ friend bool operator==(const robin_set& lhs, const robin_set& rhs) {
+ if(lhs.size() != rhs.size()) {
+ return false;
+ }
+
+ for(const auto& element_lhs: lhs) {
+ const auto it_element_rhs = rhs.find(element_lhs);
+ if(it_element_rhs == rhs.cend()) {
+ return false;
+ }
+ }
+
+ return true;
+ }
+
+ friend bool operator!=(const robin_set& lhs, const robin_set& rhs) {
+ return !operator==(lhs, rhs);
+ }
+
+ friend void swap(robin_set& lhs, robin_set& rhs) {
+ lhs.swap(rhs);
+ }
+
+private:
+ ht m_ht;
+};
+
+
+/**
+ * Same as `tsl::robin_set<Key, Hash, KeyEqual, Allocator, StoreHash, tsl::rh::prime_growth_policy>`.
+ */
+template<class Key,
+ class Hash = std::hash<Key>,
+ class KeyEqual = std::equal_to<Key>,
+ class Allocator = std::allocator<Key>,
+ bool StoreHash = false>
+using robin_pg_set = robin_set<Key, Hash, KeyEqual, Allocator, StoreHash, tsl::rh::prime_growth_policy>;
+
+} // end namespace tsl
+
+#endif
+
diff --git a/be/src/extern/diskann/include/tsl/sparse_growth_policy.h b/be/src/extern/diskann/include/tsl/sparse_growth_policy.h
new file mode 100644
index 00000000000..d73aaaf420f
--- /dev/null
+++ b/be/src/extern/diskann/include/tsl/sparse_growth_policy.h
@@ -0,0 +1,301 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2017 Thibaut Goetghebuer-Planchon <te...@gmx.com>
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#ifndef TSL_SPARSE_GROWTH_POLICY_H
+#define TSL_SPARSE_GROWTH_POLICY_H
+
+#include <algorithm>
+#include <array>
+#include <climits>
+#include <cmath>
+#include <cstddef>
+#include <iterator>
+#include <limits>
+#include <ratio>
+#include <stdexcept>
+
+namespace tsl {
+namespace sh {
+
+/**
+ * Grow the hash table by a factor of GrowthFactor keeping the bucket count to a
+ * power of two. It allows the table to use a mask operation instead of a modulo
+ * operation to map a hash to a bucket.
+ *
+ * GrowthFactor must be a power of two >= 2.
+ */
+template <std::size_t GrowthFactor>
+class power_of_two_growth_policy {
+ public:
+ /**
+ * Called on the hash table creation and on rehash. The number of buckets for
+ * the table is passed in parameter. This number is a minimum, the policy may
+ * update this value with a higher value if needed (but not lower).
+ *
+ * If 0 is given, min_bucket_count_in_out must still be 0 after the policy
+ * creation and bucket_for_hash must always return 0 in this case.
+ */
+ explicit power_of_two_growth_policy(std::size_t &min_bucket_count_in_out) {
+ if (min_bucket_count_in_out > max_bucket_count()) {
+ throw std::length_error("The hash table exceeds its maximum size.");
+ }
+
+ if (min_bucket_count_in_out > 0) {
+ min_bucket_count_in_out =
+ round_up_to_power_of_two(min_bucket_count_in_out);
+ m_mask = min_bucket_count_in_out - 1;
+ } else {
+ m_mask = 0;
+ }
+ }
+
+ /**
+ * Return the bucket [0, bucket_count()) to which the hash belongs.
+ * If bucket_count() is 0, it must always return 0.
+ */
+ std::size_t bucket_for_hash(std::size_t hash) const noexcept {
+ return hash & m_mask;
+ }
+
+ /**
+ * Return the number of buckets that should be used on next growth.
+ */
+ std::size_t next_bucket_count() const {
+ if ((m_mask + 1) > max_bucket_count() / GrowthFactor) {
+ throw std::length_error("The hash table exceeds its maximum size.");
+ }
+
+ return (m_mask + 1) * GrowthFactor;
+ }
+
+ /**
+ * Return the maximum number of buckets supported by the policy.
+ */
+ std::size_t max_bucket_count() const {
+ // Largest power of two.
+ return (std::numeric_limits<std::size_t>::max() / 2) + 1;
+ }
+
+ /**
+ * Reset the growth policy as if it was created with a bucket count of 0.
+ * After a clear, the policy must always return 0 when bucket_for_hash is
+ * called.
+ */
+ void clear() noexcept { m_mask = 0; }
+
+ private:
+ static std::size_t round_up_to_power_of_two(std::size_t value) {
+ if (is_power_of_two(value)) {
+ return value;
+ }
+
+ if (value == 0) {
+ return 1;
+ }
+
+ --value;
+ for (std::size_t i = 1; i < sizeof(std::size_t) * CHAR_BIT; i *= 2) {
+ value |= value >> i;
+ }
+
+ return value + 1;
+ }
+
+ static constexpr bool is_power_of_two(std::size_t value) {
+ return value != 0 && (value & (value - 1)) == 0;
+ }
+
+ protected:
+ static_assert(is_power_of_two(GrowthFactor) && GrowthFactor >= 2,
+ "GrowthFactor must be a power of two >= 2.");
+
+ std::size_t m_mask;
+};
+
+/**
+ * Grow the hash table by GrowthFactor::num / GrowthFactor::den and use a modulo
+ * to map a hash to a bucket. Slower but it can be useful if you want a slower
+ * growth.
+ */
+template <class GrowthFactor = std::ratio<3, 2>>
+class mod_growth_policy {
+ public:
+ explicit mod_growth_policy(std::size_t &min_bucket_count_in_out) {
+ if (min_bucket_count_in_out > max_bucket_count()) {
+ throw std::length_error("The hash table exceeds its maximum size.");
+ }
+
+ if (min_bucket_count_in_out > 0) {
+ m_mod = min_bucket_count_in_out;
+ } else {
+ m_mod = 1;
+ }
+ }
+
+ std::size_t bucket_for_hash(std::size_t hash) const noexcept {
+ return hash % m_mod;
+ }
+
+ std::size_t next_bucket_count() const {
+ if (m_mod == max_bucket_count()) {
+ throw std::length_error("The hash table exceeds its maximum size.");
+ }
+
+ const double next_bucket_count =
+ std::ceil(double(m_mod) * REHASH_SIZE_MULTIPLICATION_FACTOR);
+ if (!std::isnormal(next_bucket_count)) {
+ throw std::length_error("The hash table exceeds its maximum size.");
+ }
+
+ if (next_bucket_count > double(max_bucket_count())) {
+ return max_bucket_count();
+ } else {
+ return std::size_t(next_bucket_count);
+ }
+ }
+
+ std::size_t max_bucket_count() const { return MAX_BUCKET_COUNT; }
+
+ void clear() noexcept { m_mod = 1; }
+
+ private:
+ static constexpr double REHASH_SIZE_MULTIPLICATION_FACTOR =
+ 1.0 * GrowthFactor::num / GrowthFactor::den;
+ static const std::size_t MAX_BUCKET_COUNT =
+ std::size_t(double(std::numeric_limits<std::size_t>::max() /
+ REHASH_SIZE_MULTIPLICATION_FACTOR));
+
+ static_assert(REHASH_SIZE_MULTIPLICATION_FACTOR >= 1.1,
+ "Growth factor should be >= 1.1.");
+
+ std::size_t m_mod;
+};
+
+/**
+ * Grow the hash table by using prime numbers as bucket count. Slower than
+ * tsl::sh::power_of_two_growth_policy in general but will probably distribute
+ * the values around better in the buckets with a poor hash function.
+ *
+ * To allow the compiler to optimize the modulo operation, a lookup table is
+ * used with constant primes numbers.
+ *
+ * With a switch the code would look like:
+ * \code
+ * switch(iprime) { // iprime is the current prime of the hash table
+ * case 0: hash % 5ul;
+ * break;
+ * case 1: hash % 17ul;
+ * break;
+ * case 2: hash % 29ul;
+ * break;
+ * ...
+ * }
+ * \endcode
+ *
+ * Due to the constant variable in the modulo the compiler is able to optimize
+ * the operation by a series of multiplications, substractions and shifts.
+ *
+ * The 'hash % 5' could become something like 'hash - (hash * 0xCCCCCCCD) >> 34)
+ * * 5' in a 64 bits environment.
+ */
+class prime_growth_policy {
+ public:
+ explicit prime_growth_policy(std::size_t &min_bucket_count_in_out) {
+ auto it_prime = std::lower_bound(primes().begin(), primes().end(),
+ min_bucket_count_in_out);
+ if (it_prime == primes().end()) {
+ throw std::length_error("The hash table exceeds its maximum size.");
+ }
+
+ m_iprime =
+ static_cast<unsigned int>(std::distance(primes().begin(), it_prime));
+ if (min_bucket_count_in_out > 0) {
+ min_bucket_count_in_out = *it_prime;
+ } else {
+ min_bucket_count_in_out = 0;
+ }
+ }
+
+ std::size_t bucket_for_hash(std::size_t hash) const noexcept {
+ return mod_prime()[m_iprime](hash);
+ }
+
+ std::size_t next_bucket_count() const {
+ if (m_iprime + 1 >= primes().size()) {
+ throw std::length_error("The hash table exceeds its maximum size.");
+ }
+
+ return primes()[m_iprime + 1];
+ }
+
+ std::size_t max_bucket_count() const { return primes().back(); }
+
+ void clear() noexcept { m_iprime = 0; }
+
+ private:
+ static const std::array<std::size_t, 40> &primes() {
+ static const std::array<std::size_t, 40> PRIMES = {
+ {1ul, 5ul, 17ul, 29ul, 37ul,
+ 53ul, 67ul, 79ul, 97ul, 131ul,
+ 193ul, 257ul, 389ul, 521ul, 769ul,
+ 1031ul, 1543ul, 2053ul, 3079ul, 6151ul,
+ 12289ul, 24593ul, 49157ul, 98317ul, 196613ul,
+ 393241ul, 786433ul, 1572869ul, 3145739ul, 6291469ul,
+ 12582917ul, 25165843ul, 50331653ul, 100663319ul, 201326611ul,
+ 402653189ul, 805306457ul, 1610612741ul, 3221225473ul, 4294967291ul}};
+
+ static_assert(
+ std::numeric_limits<decltype(m_iprime)>::max() >= PRIMES.size(),
+ "The type of m_iprime is not big enough.");
+
+ return PRIMES;
+ }
+
+ static const std::array<std::size_t (*)(std::size_t), 40> &mod_prime() {
+ // MOD_PRIME[iprime](hash) returns hash % PRIMES[iprime]. This table allows
+ // for faster modulo as the compiler can optimize the modulo code better
+ // with a constant known at the compilation.
+ static const std::array<std::size_t (*)(std::size_t), 40> MOD_PRIME = {
+ {&mod<0>, &mod<1>, &mod<2>, &mod<3>, &mod<4>, &mod<5>, &mod<6>,
+ &mod<7>, &mod<8>, &mod<9>, &mod<10>, &mod<11>, &mod<12>, &mod<13>,
+ &mod<14>, &mod<15>, &mod<16>, &mod<17>, &mod<18>, &mod<19>, &mod<20>,
+ &mod<21>, &mod<22>, &mod<23>, &mod<24>, &mod<25>, &mod<26>, &mod<27>,
+ &mod<28>, &mod<29>, &mod<30>, &mod<31>, &mod<32>, &mod<33>, &mod<34>,
+ &mod<35>, &mod<36>, &mod<37>, &mod<38>, &mod<39>}};
+
+ return MOD_PRIME;
+ }
+
+ template <unsigned int IPrime>
+ static std::size_t mod(std::size_t hash) {
+ return hash % primes()[IPrime];
+ }
+
+ private:
+ unsigned int m_iprime;
+};
+
+} // namespace sh
+} // namespace tsl
+
+#endif
diff --git a/be/src/extern/diskann/include/tsl/sparse_hash.h b/be/src/extern/diskann/include/tsl/sparse_hash.h
new file mode 100644
index 00000000000..e2115b426c6
--- /dev/null
+++ b/be/src/extern/diskann/include/tsl/sparse_hash.h
@@ -0,0 +1,2215 @@
+/**
+ * MIT License
+ *
+ * Copyright (c) 2017 Thibaut Goetghebuer-Planchon <te...@gmx.com>
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#ifndef TSL_SPARSE_HASH_H
+#define TSL_SPARSE_HASH_H
+
+#include <algorithm>
+#include <cassert>
+#include <climits>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <iterator>
+#include <limits>
+#include <memory>
+#include <stdexcept>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#include "sparse_growth_policy.h"
+
+#ifdef __INTEL_COMPILER
+#include <immintrin.h> // For _popcnt32 and _popcnt64
+#endif
+
+#ifdef _MSC_VER
+#include <intrin.h> // For __cpuid, __popcnt and __popcnt64
+#endif
+
+#ifdef TSL_DEBUG
+#define tsl_sh_assert(expr) assert(expr)
+#else
+#define tsl_sh_assert(expr) (static_cast<void>(0))
+#endif
+
+namespace tsl {
+
+namespace sh {
+enum class probing { linear, quadratic };
+
+enum class exception_safety { basic, strong };
+
+enum class sparsity { high, medium, low };
+} // namespace sh
+
+namespace detail_popcount {
+/**
+ * Define the popcount(ll) methods and pick-up the best depending on the
+ * compiler.
+ */
+
+// From Wikipedia: https://en.wikipedia.org/wiki/Hamming_weight
+inline int fallback_popcountll(unsigned long long int x) {
+ static_assert(
+ sizeof(unsigned long long int) == sizeof(std::uint64_t),
+ "sizeof(unsigned long long int) must be equal to sizeof(std::uint64_t). "
+ "Open a feature request if you need support for a platform where it "
+ "isn't the case.");
+
+ const std::uint64_t m1 = 0x5555555555555555ull;
+ const std::uint64_t m2 = 0x3333333333333333ull;
+ const std::uint64_t m4 = 0x0f0f0f0f0f0f0f0full;
+ const std::uint64_t h01 = 0x0101010101010101ull;
+
+ x -= (x >> 1ull) & m1;
+ x = (x & m2) + ((x >> 2ull) & m2);
+ x = (x + (x >> 4ull)) & m4;
+ return static_cast<int>((x * h01) >> (64ull - 8ull));
+}
+
+inline int fallback_popcount(unsigned int x) {
+ static_assert(sizeof(unsigned int) == sizeof(std::uint32_t) ||
+ sizeof(unsigned int) == sizeof(std::uint64_t),
+ "sizeof(unsigned int) must be equal to sizeof(std::uint32_t) "
+ "or sizeof(std::uint64_t). "
+ "Open a feature request if you need support for a platform "
+ "where it isn't the case.");
+
+ if (sizeof(unsigned int) == sizeof(std::uint32_t)) {
+ const std::uint32_t m1 = 0x55555555;
+ const std::uint32_t m2 = 0x33333333;
+ const std::uint32_t m4 = 0x0f0f0f0f;
+ const std::uint32_t h01 = 0x01010101;
+
+ x -= (x >> 1) & m1;
+ x = (x & m2) + ((x >> 2) & m2);
+ x = (x + (x >> 4)) & m4;
+ return static_cast<int>((x * h01) >> (32 - 8));
+ } else {
+ return fallback_popcountll(x);
+ }
+}
+
+#if defined(__clang__) || defined(__GNUC__)
+inline int popcountll(unsigned long long int value) {
+ return __builtin_popcountll(value);
+}
+
+inline int popcount(unsigned int value) { return __builtin_popcount(value); }
+
+#elif defined(_MSC_VER)
+/**
+ * We need to check for popcount support at runtime on Windows with __cpuid
+ * See https://msdn.microsoft.com/en-us/library/bb385231.aspx
+ */
+inline bool has_popcount_support() {
+ int cpu_infos[4];
+ __cpuid(cpu_infos, 1);
+ return (cpu_infos[2] & (1 << 23)) != 0;
+}
+
+inline int popcountll(unsigned long long int value) {
+#ifdef _WIN64
+ static_assert(
+ sizeof(unsigned long long int) == sizeof(std::int64_t),
+ "sizeof(unsigned long long int) must be equal to sizeof(std::int64_t). ");
+
+ static const bool has_popcount = has_popcount_support();
+ return has_popcount
+ ? static_cast<int>(__popcnt64(static_cast<std::int64_t>(value)))
+ : fallback_popcountll(value);
+#else
+ return fallback_popcountll(value);
+#endif
+}
+
+inline int popcount(unsigned int value) {
+ static_assert(sizeof(unsigned int) == sizeof(std::int32_t),
+ "sizeof(unsigned int) must be equal to sizeof(std::int32_t). ");
+
+ static const bool has_popcount = has_popcount_support();
+ return has_popcount
+ ? static_cast<int>(__popcnt(static_cast<std::int32_t>(value)))
+ : fallback_popcount(value);
+}
+
+#elif defined(__INTEL_COMPILER)
+inline int popcountll(unsigned long long int value) {
+ static_assert(sizeof(unsigned long long int) == sizeof(__int64), "");
+ return _popcnt64(static_cast<__int64>(value));
+}
+
+inline int popcount(unsigned int value) {
+ return _popcnt32(static_cast<int>(value));
+}
+
+#else
+inline int popcountll(unsigned long long int x) {
+ return fallback_popcountll(x);
+}
+
+inline int popcount(unsigned int x) { return fallback_popcount(x); }
+
+#endif
+} // namespace detail_popcount
+
+namespace detail_sparse_hash {
+
+template <typename T>
+struct make_void {
+ using type = void;
+};
+
+template <typename T, typename = void>
+struct has_is_transparent : std::false_type {};
+
+template <typename T>
+struct has_is_transparent<T,
+ typename make_void<typename T::is_transparent>::type>
+ : std::true_type {};
+
+template <typename U>
+struct is_power_of_two_policy : std::false_type {};
+
+template <std::size_t GrowthFactor>
+struct is_power_of_two_policy<tsl::sh::power_of_two_growth_policy<GrowthFactor>>
+ : std::true_type {};
+
+inline constexpr bool is_power_of_two(std::size_t value) {
+ return value != 0 && (value & (value - 1)) == 0;
+}
+
+inline std::size_t round_up_to_power_of_two(std::size_t value) {
+ if (is_power_of_two(value)) {
+ return value;
+ }
+
+ if (value == 0) {
+ return 1;
+ }
+
+ --value;
+ for (std::size_t i = 1; i < sizeof(std::size_t) * CHAR_BIT; i *= 2) {
+ value |= value >> i;
+ }
+
+ return value + 1;
+}
+
+template <typename T, typename U>
+static T numeric_cast(U value,
+ const char *error_message = "numeric_cast() failed.") {
+ T ret = static_cast<T>(value);
+ if (static_cast<U>(ret) != value) {
+ throw std::runtime_error(error_message);
+ }
+
+ const bool is_same_signedness =
+ (std::is_unsigned<T>::value && std::is_unsigned<U>::value) ||
+ (std::is_signed<T>::value && std::is_signed<U>::value);
+ if (!is_same_signedness && (ret < T{}) != (value < U{})) {
+ throw std::runtime_error(error_message);
+ }
+
+ return ret;
+}
+
+/**
+ * Fixed size type used to represent size_type values on serialization. Need to
+ * be big enough to represent a std::size_t on 32 and 64 bits platforms, and
+ * must be the same size on both platforms.
+ */
+using slz_size_type = std::uint64_t;
+static_assert(std::numeric_limits<slz_size_type>::max() >=
+ std::numeric_limits<std::size_t>::max(),
+ "slz_size_type must be >= std::size_t");
+
+template <class T, class Deserializer>
+static T deserialize_value(Deserializer &deserializer) {
+ // MSVC < 2017 is not conformant, circumvent the problem by removing the
+ // template keyword
+#if defined(_MSC_VER) && _MSC_VER < 1910
+ return deserializer.Deserializer::operator()<T>();
+#else
+ return deserializer.Deserializer::template operator()<T>();
+#endif
+}
+
+/**
+ * WARNING: the sparse_array class doesn't free the ressources allocated through
+ * the allocator passed in parameter in each method. You have to manually call
+ * `clear(Allocator&)` when you don't need a sparse_array object anymore.
+ *
+ * The reason is that the sparse_array doesn't store the allocator to avoid
+ * wasting space in each sparse_array when the allocator has a size > 0. It only
+ * allocates/deallocates objects with the allocator that is passed in parameter.
+ *
+ *
+ *
+ * Index denotes a value between [0, BITMAP_NB_BITS), it is an index similar to
+ * std::vector. Offset denotes the real position in `m_values` corresponding to
+ * an index.
+ *
+ * We are using raw pointers instead of std::vector to avoid loosing
+ * 2*sizeof(size_t) bytes to store the capacity and size of the vector in each
+ * sparse_array. We know we can only store up to BITMAP_NB_BITS elements in the
+ * array, we don't need such big types.
+ *
+ *
+ * T must be nothrow move constructible and/or copy constructible.
+ * Behaviour is undefined if the destructor of T throws an exception.
+ *
+ * See https://smerity.com/articles/2015/google_sparsehash.html for details on
+ * the idea behinds the implementation.
+ *
+ * TODO Check to use std::realloc and std::memmove when possible
+ */
+template <typename T, typename Allocator, tsl::sh::sparsity Sparsity>
+class sparse_array {
+ public:
+ using value_type = T;
+ using size_type = std::uint_least8_t;
+ using allocator_type = Allocator;
+ using iterator = value_type *;
+ using const_iterator = const value_type *;
+
+ private:
+ static const size_type CAPACITY_GROWTH_STEP =
+ (Sparsity == tsl::sh::sparsity::high) ? 2
+ : (Sparsity == tsl::sh::sparsity::medium)
+ ? 4
+ : 8; // (Sparsity == tsl::sh::sparsity::low)
+
+ /**
+ * Bitmap size configuration.
+ * Use 32 bits for the bitmap on 32-bits or less environnement as popcount on
+ * 64 bits numbers is slow on these environnement. Use 64 bits bitmap
+ * otherwise.
+ */
+#if SIZE_MAX <= UINT32_MAX
+ using bitmap_type = std::uint_least32_t;
+ static const std::size_t BITMAP_NB_BITS = 32;
+ static const std::size_t BUCKET_SHIFT = 5;
+#else
+ using bitmap_type = std::uint_least64_t;
+ static const std::size_t BITMAP_NB_BITS = 64;
+ static const std::size_t BUCKET_SHIFT = 6;
+#endif
+
+ static const std::size_t BUCKET_MASK = BITMAP_NB_BITS - 1;
+
+ static_assert(is_power_of_two(BITMAP_NB_BITS),
+ "BITMAP_NB_BITS must be a power of two.");
+ static_assert(std::numeric_limits<bitmap_type>::digits >= BITMAP_NB_BITS,
+ "bitmap_type must be able to hold at least BITMAP_NB_BITS.");
+ static_assert((std::size_t(1) << BUCKET_SHIFT) == BITMAP_NB_BITS,
+ "(1 << BUCKET_SHIFT) must be equal to BITMAP_NB_BITS.");
+ static_assert(std::numeric_limits<size_type>::max() >= BITMAP_NB_BITS,
+ "size_type must be big enough to hold BITMAP_NB_BITS.");
+ static_assert(std::is_unsigned<bitmap_type>::value,
+ "bitmap_type must be unsigned.");
+ static_assert((std::numeric_limits<bitmap_type>::max() & BUCKET_MASK) ==
+ BITMAP_NB_BITS - 1,
+ "");
+
+ public:
+ /**
+ * Map an ibucket [0, bucket_count) in the hash table to a sparse_ibucket
+ * (a sparse_array holds multiple buckets, so there is less sparse_array than
+ * bucket_count).
+ *
+ * The bucket ibucket is in
+ * m_sparse_buckets[sparse_ibucket(ibucket)][index_in_sparse_bucket(ibucket)]
+ * instead of something like m_buckets[ibucket] in a classical hash table.
+ */
+ static std::size_t sparse_ibucket(std::size_t ibucket) {
+ return ibucket >> BUCKET_SHIFT;
+ }
+
+ /**
+ * Map an ibucket [0, bucket_count) in the hash table to an index in the
+ * sparse_array which corresponds to the bucket.
+ *
+ * The bucket ibucket is in
+ * m_sparse_buckets[sparse_ibucket(ibucket)][index_in_sparse_bucket(ibucket)]
+ * instead of something like m_buckets[ibucket] in a classical hash table.
+ */
+ static typename sparse_array::size_type index_in_sparse_bucket(
+ std::size_t ibucket) {
+ return static_cast<typename sparse_array::size_type>(
+ ibucket & sparse_array::BUCKET_MASK);
+ }
+
+ static std::size_t nb_sparse_buckets(std::size_t bucket_count) noexcept {
+ if (bucket_count == 0) {
+ return 0;
+ }
+
+ return std::max<std::size_t>(
+ 1, sparse_ibucket(tsl::detail_sparse_hash::round_up_to_power_of_two(
+ bucket_count)));
+ }
+
+ public:
+ sparse_array() noexcept
+ : m_values(nullptr),
+ m_bitmap_vals(0),
+ m_bitmap_deleted_vals(0),
+ m_nb_elements(0),
+ m_capacity(0),
+ m_last_array(false) {}
+
+ explicit sparse_array(bool last_bucket) noexcept
+ : m_values(nullptr),
+ m_bitmap_vals(0),
+ m_bitmap_deleted_vals(0),
+ m_nb_elements(0),
+ m_capacity(0),
+ m_last_array(last_bucket) {}
+
+ sparse_array(size_type capacity, Allocator &alloc)
+ : m_values(nullptr),
+ m_bitmap_vals(0),
+ m_bitmap_deleted_vals(0),
+ m_nb_elements(0),
+ m_capacity(capacity),
+ m_last_array(false) {
+ if (m_capacity > 0) {
+ m_values = alloc.allocate(m_capacity);
+ tsl_sh_assert(m_values !=
+ nullptr); // allocate should throw if there is a failure
+ }
+ }
+
+ sparse_array(const sparse_array &other, Allocator &alloc)
+ : m_values(nullptr),
+ m_bitmap_vals(other.m_bitmap_vals),
+ m_bitmap_deleted_vals(other.m_bitmap_deleted_vals),
+ m_nb_elements(0),
+ m_capacity(other.m_capacity),
+ m_last_array(other.m_last_array) {
+ tsl_sh_assert(other.m_capacity >= other.m_nb_elements);
+ if (m_capacity == 0) {
+ return;
+ }
+
+ m_values = alloc.allocate(m_capacity);
+ tsl_sh_assert(m_values !=
+ nullptr); // allocate should throw if there is a failure
+ try {
+ for (size_type i = 0; i < other.m_nb_elements; i++) {
+ construct_value(alloc, m_values + i, other.m_values[i]);
+ m_nb_elements++;
+ }
+ } catch (...) {
+ clear(alloc);
+ throw;
+ }
+ }
+
+ sparse_array(sparse_array &&other) noexcept
+ : m_values(other.m_values),
+ m_bitmap_vals(other.m_bitmap_vals),
+ m_bitmap_deleted_vals(other.m_bitmap_deleted_vals),
+ m_nb_elements(other.m_nb_elements),
+ m_capacity(other.m_capacity),
+ m_last_array(other.m_last_array) {
+ other.m_values = nullptr;
+ other.m_bitmap_vals = 0;
+ other.m_bitmap_deleted_vals = 0;
+ other.m_nb_elements = 0;
+ other.m_capacity = 0;
+ }
+
+ sparse_array(sparse_array &&other, Allocator &alloc)
+ : m_values(nullptr),
+ m_bitmap_vals(other.m_bitmap_vals),
+ m_bitmap_deleted_vals(other.m_bitmap_deleted_vals),
+ m_nb_elements(0),
+ m_capacity(other.m_capacity),
+ m_last_array(other.m_last_array) {
+ tsl_sh_assert(other.m_capacity >= other.m_nb_elements);
+ if (m_capacity == 0) {
+ return;
+ }
+
+ m_values = alloc.allocate(m_capacity);
+ tsl_sh_assert(m_values !=
+ nullptr); // allocate should throw if there is a failure
+ try {
+ for (size_type i = 0; i < other.m_nb_elements; i++) {
+ construct_value(alloc, m_values + i, std::move(other.m_values[i]));
+ m_nb_elements++;
+ }
+ } catch (...) {
+ clear(alloc);
+ throw;
+ }
+ }
+
+ sparse_array &operator=(const sparse_array &) = delete;
+ sparse_array &operator=(sparse_array &&) = delete;
+
+ ~sparse_array() noexcept {
+ // The code that manages the sparse_array must have called clear before
+ // destruction. See documentation of sparse_array for more details.
+ tsl_sh_assert(m_capacity == 0 && m_nb_elements == 0 && m_values == nullptr);
+ }
+
+ iterator begin() noexcept { return m_values; }
+ iterator end() noexcept { return m_values + m_nb_elements; }
+ const_iterator begin() const noexcept { return cbegin(); }
+ const_iterator end() const noexcept { return cend(); }
+ const_iterator cbegin() const noexcept { return m_values; }
+ const_iterator cend() const noexcept { return m_values + m_nb_elements; }
+
+ bool empty() const noexcept { return m_nb_elements == 0; }
+
+ size_type size() const noexcept { return m_nb_elements; }
+
+ void clear(allocator_type &alloc) noexcept {
+ destroy_and_deallocate_values(alloc, m_values, m_nb_elements, m_capacity);
+
+ m_values = nullptr;
+ m_bitmap_vals = 0;
+ m_bitmap_deleted_vals = 0;
+ m_nb_elements = 0;
+ m_capacity = 0;
+ }
+
+ bool last() const noexcept { return m_last_array; }
+
+ void set_as_last() noexcept { m_last_array = true; }
+
+ bool has_value(size_type index) const noexcept {
+ tsl_sh_assert(index < BITMAP_NB_BITS);
+ return (m_bitmap_vals & (bitmap_type(1) << index)) != 0;
+ }
+
+ bool has_deleted_value(size_type index) const noexcept {
+ tsl_sh_assert(index < BITMAP_NB_BITS);
+ return (m_bitmap_deleted_vals & (bitmap_type(1) << index)) != 0;
+ }
+
+ iterator value(size_type index) noexcept {
+ tsl_sh_assert(has_value(index));
+ return m_values + index_to_offset(index);
+ }
+
+ const_iterator value(size_type index) const noexcept {
+ tsl_sh_assert(has_value(index));
+ return m_values + index_to_offset(index);
+ }
+
+ /**
+ * Return iterator to set value.
+ */
+ template <typename... Args>
+ iterator set(allocator_type &alloc, size_type index, Args &&...value_args) {
+ tsl_sh_assert(!has_value(index));
+
+ const size_type offset = index_to_offset(index);
+ insert_at_offset(alloc, offset, std::forward<Args>(value_args)...);
+
+ m_bitmap_vals = (m_bitmap_vals | (bitmap_type(1) << index));
+ m_bitmap_deleted_vals =
+ (m_bitmap_deleted_vals & ~(bitmap_type(1) << index));
+
+ m_nb_elements++;
+
+ tsl_sh_assert(has_value(index));
+ tsl_sh_assert(!has_deleted_value(index));
+
+ return m_values + offset;
+ }
+
+ iterator erase(allocator_type &alloc, iterator position) {
+ const size_type offset =
+ static_cast<size_type>(std::distance(begin(), position));
+ return erase(alloc, position, offset_to_index(offset));
+ }
+
+ // Return the next value or end if no next value
+ iterator erase(allocator_type &alloc, iterator position, size_type index) {
+ tsl_sh_assert(has_value(index));
+ tsl_sh_assert(!has_deleted_value(index));
+
+ const size_type offset =
+ static_cast<size_type>(std::distance(begin(), position));
+ erase_at_offset(alloc, offset);
+
+ m_bitmap_vals = (m_bitmap_vals & ~(bitmap_type(1) << index));
+ m_bitmap_deleted_vals = (m_bitmap_deleted_vals | (bitmap_type(1) << index));
+
+ m_nb_elements--;
+
+ tsl_sh_assert(!has_value(index));
+ tsl_sh_assert(has_deleted_value(index));
+
+ return m_values + offset;
+ }
+
+ void swap(sparse_array &other) {
+ using std::swap;
+
+ swap(m_values, other.m_values);
+ swap(m_bitmap_vals, other.m_bitmap_vals);
+ swap(m_bitmap_deleted_vals, other.m_bitmap_deleted_vals);
+ swap(m_nb_elements, other.m_nb_elements);
+ swap(m_capacity, other.m_capacity);
+ swap(m_last_array, other.m_last_array);
+ }
+
+ static iterator mutable_iterator(const_iterator pos) {
+ return const_cast<iterator>(pos);
+ }
+
+ template <class Serializer>
+ void serialize(Serializer &serializer) const {
+ const slz_size_type sparse_bucket_size = m_nb_elements;
+ serializer(sparse_bucket_size);
+
+ const slz_size_type bitmap_vals = m_bitmap_vals;
+ serializer(bitmap_vals);
+
+ const slz_size_type bitmap_deleted_vals = m_bitmap_deleted_vals;
+ serializer(bitmap_deleted_vals);
+
+ for (const value_type &value : *this) {
+ serializer(value);
+ }
+ }
+
+ template <class Deserializer>
+ static sparse_array deserialize_hash_compatible(Deserializer &deserializer,
+ Allocator &alloc) {
+ const slz_size_type sparse_bucket_size =
+ deserialize_value<slz_size_type>(deserializer);
+ const slz_size_type bitmap_vals =
+ deserialize_value<slz_size_type>(deserializer);
+ const slz_size_type bitmap_deleted_vals =
+ deserialize_value<slz_size_type>(deserializer);
+
+ if (sparse_bucket_size > BITMAP_NB_BITS) {
+ throw std::runtime_error(
+ "Deserialized sparse_bucket_size is too big for the platform. "
+ "Maximum should be BITMAP_NB_BITS.");
+ }
+
+ sparse_array sarray;
+ if (sparse_bucket_size == 0) {
+ return sarray;
+ }
+
+ sarray.m_bitmap_vals = numeric_cast<bitmap_type>(
+ bitmap_vals, "Deserialized bitmap_vals is too big.");
+ sarray.m_bitmap_deleted_vals = numeric_cast<bitmap_type>(
+ bitmap_deleted_vals, "Deserialized bitmap_deleted_vals is too big.");
+
+ sarray.m_capacity = numeric_cast<size_type>(
+ sparse_bucket_size, "Deserialized sparse_bucket_size is too big.");
+ sarray.m_values = alloc.allocate(sarray.m_capacity);
+
+ try {
+ for (size_type ivalue = 0; ivalue < sarray.m_capacity; ivalue++) {
+ construct_value(alloc, sarray.m_values + ivalue,
+ deserialize_value<value_type>(deserializer));
+ sarray.m_nb_elements++;
+ }
+ } catch (...) {
+ sarray.clear(alloc);
+ throw;
+ }
+
+ return sarray;
+ }
+
+ /**
+ * Deserialize the values of the bucket and insert them all in sparse_hash
+ * through sparse_hash.insert(...).
+ */
+ template <class Deserializer, class SparseHash>
+ static void deserialize_values_into_sparse_hash(Deserializer &deserializer,
+ SparseHash &sparse_hash) {
+ const slz_size_type sparse_bucket_size =
+ deserialize_value<slz_size_type>(deserializer);
+
+ const slz_size_type bitmap_vals =
+ deserialize_value<slz_size_type>(deserializer);
+ static_cast<void>(bitmap_vals); // Ignore, not needed
+
+ const slz_size_type bitmap_deleted_vals =
+ deserialize_value<slz_size_type>(deserializer);
+ static_cast<void>(bitmap_deleted_vals); // Ignore, not needed
+
+ for (slz_size_type ivalue = 0; ivalue < sparse_bucket_size; ivalue++) {
+ sparse_hash.insert(deserialize_value<value_type>(deserializer));
+ }
+ }
+
+ private:
+ template <typename... Args>
+ static void construct_value(allocator_type &alloc, value_type *value,
+ Args &&...value_args) {
+ std::allocator_traits<allocator_type>::construct(
+ alloc, value, std::forward<Args>(value_args)...);
+ }
+
+ static void destroy_value(allocator_type &alloc, value_type *value) noexcept {
+ std::allocator_traits<allocator_type>::destroy(alloc, value);
+ }
+
+ static void destroy_and_deallocate_values(
+ allocator_type &alloc, value_type *values, size_type nb_values,
+ size_type capacity_values) noexcept {
+ for (size_type i = 0; i < nb_values; i++) {
+ destroy_value(alloc, values + i);
+ }
+
+ alloc.deallocate(values, capacity_values);
+ }
+
+ static size_type popcount(bitmap_type val) noexcept {
+ if (sizeof(bitmap_type) <= sizeof(unsigned int)) {
+ return static_cast<size_type>(
+ tsl::detail_popcount::popcount(static_cast<unsigned int>(val)));
+ } else {
+ return static_cast<size_type>(tsl::detail_popcount::popcountll(val));
+ }
+ }
+
+ size_type index_to_offset(size_type index) const noexcept {
+ tsl_sh_assert(index < BITMAP_NB_BITS);
+ return popcount(m_bitmap_vals &
+ ((bitmap_type(1) << index) - bitmap_type(1)));
+ }
+
+ // TODO optimize
+ size_type offset_to_index(size_type offset) const noexcept {
+ tsl_sh_assert(offset < m_nb_elements);
+
+ bitmap_type bitmap_vals = m_bitmap_vals;
+ size_type index = 0;
+ size_type nb_ones = 0;
+
+ while (bitmap_vals != 0) {
+ if ((bitmap_vals & 0x1) == 1) {
+ if (nb_ones == offset) {
+ break;
+ }
+
+ nb_ones++;
+ }
+
+ index++;
+ bitmap_vals = bitmap_vals >> 1;
+ }
+
+ return index;
+ }
+
+ size_type next_capacity() const noexcept {
+ return static_cast<size_type>(m_capacity + CAPACITY_GROWTH_STEP);
+ }
+
+ /**
+ * Insertion
+ *
+ * Two situations:
+ * - Either we are in a situation where
+ * std::is_nothrow_move_constructible<value_type>::value is true. In this
+ * case, on insertion we just reallocate m_values when we reach its capacity
+ * (i.e. m_nb_elements == m_capacity), otherwise we just put the new value at
... 24098 lines suppressed ...
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@doris.apache.org
For additional commands, e-mail: commits-help@doris.apache.org