mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-11-30 02:48:45 +08:00
Remove index dir (#15467)
Signed-off-by: yun.zhang <yun.zhang@zilliz.com>
This commit is contained in:
jaime
GitHub
parent
b978d555ba
commit
cb9c63d039
1
internal/core/src/index/.gitignore
vendored
1
internal/core/src/index/.gitignore
vendored
@ -1 +0,0 @@
|
||||
cmake_build
|
||||
@ -1,109 +0,0 @@
|
||||
#-------------------------------------------------------------------------------
|
||||
# Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
# or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
#-------------------------------------------------------------------------------
|
||||
|
||||
cmake_minimum_required(VERSION 3.18)
|
||||
message(STATUS "------------------------------KNOWHERE-----------------------------------")
|
||||
message(STATUS "Building using CMake version: ${CMAKE_VERSION}")
|
||||
|
||||
project(knowhere LANGUAGES C CXX)
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
|
||||
# if no build build type is specified, default to release builds
|
||||
if (NOT CMAKE_BUILD_TYPE)
|
||||
set(CMAKE_BUILD_TYPE Release)
|
||||
endif (NOT CMAKE_BUILD_TYPE)
|
||||
|
||||
if (CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")
|
||||
message(STATUS "building milvus_engine on x86 architecture")
|
||||
set(KNOWHERE_BUILD_ARCH x86_64)
|
||||
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "(ppc)")
|
||||
message(STATUS "building milvus_engine on ppc architecture")
|
||||
set(KNOWHERE_BUILD_ARCH ppc64le)
|
||||
else ()
|
||||
message(WARNING "unknown processor type")
|
||||
message(WARNING "CMAKE_SYSTEM_PROCESSOR=${CMAKE_SYSTEM_PROCESSOR}")
|
||||
set(KNOWHERE_BUILD_ARCH unknown)
|
||||
endif ()
|
||||
|
||||
if (CMAKE_BUILD_TYPE STREQUAL "Release")
|
||||
set(BUILD_TYPE "release")
|
||||
else ()
|
||||
set(BUILD_TYPE "debug")
|
||||
endif ()
|
||||
message(STATUS "Build type = ${BUILD_TYPE}")
|
||||
|
||||
include(GNUInstallDirs)
|
||||
set(CMAKE_INSTALL_LIBDIR lib)
|
||||
set(INDEX_SOURCE_DIR ${PROJECT_SOURCE_DIR})
|
||||
set(INDEX_BINARY_DIR ${PROJECT_BINARY_DIR})
|
||||
set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${INDEX_SOURCE_DIR}/cmake")
|
||||
|
||||
# This will set RPATH to all excutable TARGET
|
||||
# self-installed dynamic libraries will be correctly linked by excutable
|
||||
set( CMAKE_INSTALL_RPATH "/usr/lib" "${CMAKE_INSTALL_PREFIX}/${CMAKE_INSTALL_LIBDIR}" )
|
||||
set( CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE )
|
||||
|
||||
include(ExternalProject)
|
||||
include(DefineOptionsCore)
|
||||
include(BuildUtilsCore)
|
||||
|
||||
using_ccache_if_defined( KNOWHERE_USE_CCACHE )
|
||||
set_directory_properties(PROPERTIES RULE_LAUNCH_COMPILE "")
|
||||
|
||||
if (MILVUS_GPU_VERSION)
|
||||
message(STATUS "Building Knowhere GPU version")
|
||||
add_compile_definitions("MILVUS_GPU_VERSION")
|
||||
enable_language(CUDA)
|
||||
find_package(CUDA 10 REQUIRED)
|
||||
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -Xcompiler -fPIC -std=c++11 -D_FORCE_INLINES --expt-extended-lambda")
|
||||
if ( CCACHE_FOUND )
|
||||
set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_FOUND}")
|
||||
endif()
|
||||
else ()
|
||||
message(STATUS "Building Knowhere CPU version")
|
||||
endif ()
|
||||
|
||||
if (MILVUS_SUPPORT_SPTAG)
|
||||
message(STATUS "Building Knowhere with SPTAG supported")
|
||||
add_compile_definitions("MILVUS_SUPPORT_SPTAG")
|
||||
endif ()
|
||||
|
||||
include(ThirdPartyPackagesCore)
|
||||
add_subdirectory(thirdparty)
|
||||
|
||||
if (CMAKE_BUILD_TYPE STREQUAL "Release")
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp")
|
||||
if (MILVUS_GPU_VERSION)
|
||||
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3")
|
||||
endif ()
|
||||
else ()
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp")
|
||||
if (MILVUS_GPU_VERSION)
|
||||
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g")
|
||||
endif ()
|
||||
endif ()
|
||||
|
||||
add_subdirectory(knowhere)
|
||||
|
||||
if (BUILD_COVERAGE STREQUAL "ON")
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fprofile-arcs -ftest-coverage")
|
||||
endif ()
|
||||
|
||||
set(INDEX_INCLUDE_DIRS ${INDEX_INCLUDE_DIRS} PARENT_SCOPE)
|
||||
|
||||
if (KNOWHERE_BUILD_TESTS)
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DELPP_DISABLE_LOGS")
|
||||
add_subdirectory(unittest)
|
||||
endif ()
|
||||
|
||||
config_summary()
|
||||
@ -1,75 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
BUILD_TYPE="Debug"
|
||||
BUILD_UNITTEST="OFF"
|
||||
INSTALL_PREFIX=$(pwd)/cmake_build
|
||||
MAKE_CLEAN="OFF"
|
||||
PROFILING="OFF"
|
||||
|
||||
while getopts "p:d:t:uhrcgm" arg
|
||||
do
|
||||
case $arg in
|
||||
t)
|
||||
BUILD_TYPE=$OPTARG # BUILD_TYPE
|
||||
;;
|
||||
u)
|
||||
echo "Build and run unittest cases" ;
|
||||
BUILD_UNITTEST="ON";
|
||||
;;
|
||||
p)
|
||||
INSTALL_PREFIX=$OPTARG
|
||||
;;
|
||||
r)
|
||||
if [[ -d cmake_build ]]; then
|
||||
cd cmake_build
|
||||
make clean
|
||||
fi
|
||||
;;
|
||||
g)
|
||||
PROFILING="ON"
|
||||
;;
|
||||
h) # help
|
||||
echo "
|
||||
|
||||
parameter:
|
||||
-t: build type(default: Debug)
|
||||
-u: building unit test options(default: OFF)
|
||||
-p: install prefix(default: $(pwd)/knowhere)
|
||||
-r: remove previous build directory(default: OFF)
|
||||
-g: profiling(default: OFF)
|
||||
|
||||
usage:
|
||||
./build.sh -t \${BUILD_TYPE} [-u] [-h] [-g] [-r] [-c]
|
||||
"
|
||||
exit 0
|
||||
;;
|
||||
?)
|
||||
echo "unknown argument"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [[ ! -d cmake_build ]]; then
|
||||
mkdir cmake_build
|
||||
fi
|
||||
|
||||
cd cmake_build
|
||||
|
||||
CUDA_COMPILER=/usr/local/cuda/bin/nvcc
|
||||
|
||||
CMAKE_CMD="cmake -DKNOWHERE_BUILD_TESTS=${BUILD_UNITTEST} \
|
||||
-DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}
|
||||
-DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
|
||||
-DCMAKE_CUDA_COMPILER=${CUDA_COMPILER} \
|
||||
-DMILVUS_ENABLE_PROFILING=${PROFILING} \
|
||||
../"
|
||||
echo ${CMAKE_CMD}
|
||||
|
||||
${CMAKE_CMD}
|
||||
|
||||
make -j 8 ||exit 1
|
||||
|
||||
make install || exit 1
|
||||
|
||||
|
||||
@ -1,218 +0,0 @@
|
||||
# Define a function that check last file modification
|
||||
function(Check_Last_Modify cache_check_lists_file_path working_dir last_modified_commit_id)
|
||||
if (EXISTS "${working_dir}")
|
||||
if (EXISTS "${cache_check_lists_file_path}")
|
||||
set(GIT_LOG_SKIP_NUM 0)
|
||||
set(_MATCH_ALL ON CACHE BOOL "Match all")
|
||||
set(_LOOP_STATUS ON CACHE BOOL "Whether out of loop")
|
||||
file(STRINGS ${cache_check_lists_file_path} CACHE_IGNORE_TXT)
|
||||
while (_LOOP_STATUS)
|
||||
foreach (_IGNORE_ENTRY ${CACHE_IGNORE_TXT})
|
||||
if (NOT _IGNORE_ENTRY MATCHES "^[^#]+")
|
||||
continue()
|
||||
endif ()
|
||||
|
||||
set(_MATCH_ALL OFF)
|
||||
execute_process(COMMAND git log --no-merges -1 --skip=${GIT_LOG_SKIP_NUM} --name-status --pretty= WORKING_DIRECTORY ${working_dir} OUTPUT_VARIABLE CHANGE_FILES)
|
||||
if (NOT CHANGE_FILES STREQUAL "")
|
||||
string(REPLACE "\n" ";" _CHANGE_FILES ${CHANGE_FILES})
|
||||
foreach (_FILE_ENTRY ${_CHANGE_FILES})
|
||||
string(REGEX MATCH "[^ \t]+$" _FILE_NAME ${_FILE_ENTRY})
|
||||
execute_process(COMMAND sh -c "echo ${_FILE_NAME} | grep ${_IGNORE_ENTRY}" RESULT_VARIABLE return_code)
|
||||
if (return_code EQUAL 0)
|
||||
execute_process(COMMAND git log --no-merges -1 --skip=${GIT_LOG_SKIP_NUM} --pretty=%H WORKING_DIRECTORY ${working_dir} OUTPUT_VARIABLE LAST_MODIFIED_COMMIT_ID)
|
||||
set(${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
|
||||
set(_LOOP_STATUS OFF)
|
||||
endif ()
|
||||
endforeach ()
|
||||
else ()
|
||||
set(_LOOP_STATUS OFF)
|
||||
endif ()
|
||||
endforeach ()
|
||||
|
||||
if (_MATCH_ALL)
|
||||
execute_process(COMMAND git log --no-merges -1 --skip=${GIT_LOG_SKIP_NUM} --pretty=%H WORKING_DIRECTORY ${working_dir} OUTPUT_VARIABLE LAST_MODIFIED_COMMIT_ID)
|
||||
set(${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
|
||||
set(_LOOP_STATUS OFF)
|
||||
endif ()
|
||||
|
||||
math(EXPR GIT_LOG_SKIP_NUM "${GIT_LOG_SKIP_NUM} + 1")
|
||||
endwhile (_LOOP_STATUS)
|
||||
else ()
|
||||
execute_process(COMMAND git log --no-merges -1 --skip=${GIT_LOG_SKIP_NUM} --pretty=%H WORKING_DIRECTORY ${working_dir} OUTPUT_VARIABLE LAST_MODIFIED_COMMIT_ID)
|
||||
set(${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
|
||||
endif ()
|
||||
else ()
|
||||
message(FATAL_ERROR "The directory ${working_dir} does not exist")
|
||||
endif ()
|
||||
endfunction()
|
||||
|
||||
# Define a function that extracts a cached package
|
||||
function(ExternalProject_Use_Cache project_name package_file install_path)
|
||||
message(STATUS "Will use cached package file: ${package_file}")
|
||||
|
||||
ExternalProject_Add(${project_name}
|
||||
DOWNLOAD_COMMAND ${CMAKE_COMMAND} -E echo
|
||||
"No download step needed (using cached package)"
|
||||
CONFIGURE_COMMAND ${CMAKE_COMMAND} -E echo
|
||||
"No configure step needed (using cached package)"
|
||||
BUILD_COMMAND ${CMAKE_COMMAND} -E echo
|
||||
"No build step needed (using cached package)"
|
||||
INSTALL_COMMAND ${CMAKE_COMMAND} -E echo
|
||||
"No install step needed (using cached package)"
|
||||
)
|
||||
|
||||
# We want our tar files to contain the Install/<package> prefix (not for any
|
||||
# very special reason, only for consistency and so that we can identify them
|
||||
# in the extraction logs) which means that we must extract them in the
|
||||
# binary (top-level build) directory to have them installed in the right
|
||||
# place for subsequent ExternalProjects to pick them up. It seems that the
|
||||
# only way to control the working directory is with Add_Step!
|
||||
ExternalProject_Add_Step(${project_name} extract
|
||||
ALWAYS 1
|
||||
COMMAND
|
||||
${CMAKE_COMMAND} -E echo
|
||||
"Extracting ${package_file} to ${install_path}"
|
||||
COMMAND
|
||||
${CMAKE_COMMAND} -E tar xzf ${package_file} ${install_path}
|
||||
WORKING_DIRECTORY ${INDEX_BINARY_DIR}
|
||||
)
|
||||
|
||||
ExternalProject_Add_StepTargets(${project_name} extract)
|
||||
endfunction()
|
||||
|
||||
# Define a function that to create a new cached package
|
||||
function(ExternalProject_Create_Cache project_name package_file install_path cache_username cache_password cache_path)
|
||||
if (EXISTS ${package_file})
|
||||
message(STATUS "Removing existing package file: ${package_file}")
|
||||
file(REMOVE ${package_file})
|
||||
endif ()
|
||||
|
||||
string(REGEX REPLACE "(.+)/.+$" "\\1" package_dir ${package_file})
|
||||
if (NOT EXISTS ${package_dir})
|
||||
file(MAKE_DIRECTORY ${package_dir})
|
||||
endif ()
|
||||
|
||||
message(STATUS "Will create cached package file: ${package_file}")
|
||||
|
||||
ExternalProject_Add_Step(${project_name} package
|
||||
DEPENDEES install
|
||||
BYPRODUCTS ${package_file}
|
||||
COMMAND ${CMAKE_COMMAND} -E echo "Updating cached package file: ${package_file}"
|
||||
COMMAND ${CMAKE_COMMAND} -E tar czvf ${package_file} ${install_path}
|
||||
COMMAND ${CMAKE_COMMAND} -E echo "Uploading package file ${package_file} to ${cache_path}"
|
||||
COMMAND curl -u${cache_username}:${cache_password} -T ${package_file} ${cache_path}
|
||||
)
|
||||
|
||||
ExternalProject_Add_StepTargets(${project_name} package)
|
||||
endfunction()
|
||||
|
||||
function(ADD_THIRDPARTY_LIB LIB_NAME)
|
||||
set(options)
|
||||
set(one_value_args SHARED_LIB STATIC_LIB)
|
||||
set(multi_value_args DEPS INCLUDE_DIRECTORIES)
|
||||
cmake_parse_arguments(ARG
|
||||
"${options}"
|
||||
"${one_value_args}"
|
||||
"${multi_value_args}"
|
||||
${ARGN})
|
||||
if (ARG_UNPARSED_ARGUMENTS)
|
||||
message(SEND_ERROR "Error: unrecognized arguments: ${ARG_UNPARSED_ARGUMENTS}")
|
||||
endif ()
|
||||
|
||||
if (ARG_STATIC_LIB AND ARG_SHARED_LIB)
|
||||
if (NOT ARG_STATIC_LIB)
|
||||
message(FATAL_ERROR "No static or shared library provided for ${LIB_NAME}")
|
||||
endif ()
|
||||
|
||||
set(AUG_LIB_NAME "${LIB_NAME}_static")
|
||||
add_library(${AUG_LIB_NAME} STATIC IMPORTED)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES IMPORTED_LOCATION "${ARG_STATIC_LIB}")
|
||||
if (ARG_DEPS)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES INTERFACE_LINK_LIBRARIES "${ARG_DEPS}")
|
||||
endif ()
|
||||
message(STATUS "Added static library dependency ${AUG_LIB_NAME}: ${ARG_STATIC_LIB}")
|
||||
if (ARG_INCLUDE_DIRECTORIES)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES INTERFACE_INCLUDE_DIRECTORIES
|
||||
"${ARG_INCLUDE_DIRECTORIES}")
|
||||
endif ()
|
||||
|
||||
set(AUG_LIB_NAME "${LIB_NAME}_shared")
|
||||
add_library(${AUG_LIB_NAME} SHARED IMPORTED)
|
||||
|
||||
if (WIN32)
|
||||
# Mark the ".lib" location as part of a Windows DLL
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES IMPORTED_IMPLIB "${ARG_SHARED_LIB}")
|
||||
else ()
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES IMPORTED_LOCATION "${ARG_SHARED_LIB}")
|
||||
endif ()
|
||||
if (ARG_DEPS)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES INTERFACE_LINK_LIBRARIES "${ARG_DEPS}")
|
||||
endif ()
|
||||
message(STATUS "Added shared library dependency ${AUG_LIB_NAME}: ${ARG_SHARED_LIB}")
|
||||
if (ARG_INCLUDE_DIRECTORIES)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES INTERFACE_INCLUDE_DIRECTORIES
|
||||
"${ARG_INCLUDE_DIRECTORIES}")
|
||||
endif ()
|
||||
elseif (ARG_STATIC_LIB)
|
||||
set(AUG_LIB_NAME "${LIB_NAME}_static")
|
||||
add_library(${AUG_LIB_NAME} STATIC IMPORTED)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES IMPORTED_LOCATION "${ARG_STATIC_LIB}")
|
||||
if (ARG_DEPS)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES INTERFACE_LINK_LIBRARIES "${ARG_DEPS}")
|
||||
endif ()
|
||||
message(STATUS "Added static library dependency ${AUG_LIB_NAME}: ${ARG_STATIC_LIB}")
|
||||
if (ARG_INCLUDE_DIRECTORIES)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES INTERFACE_INCLUDE_DIRECTORIES
|
||||
"${ARG_INCLUDE_DIRECTORIES}")
|
||||
endif ()
|
||||
elseif (ARG_SHARED_LIB)
|
||||
set(AUG_LIB_NAME "${LIB_NAME}_shared")
|
||||
add_library(${AUG_LIB_NAME} SHARED IMPORTED)
|
||||
|
||||
if (WIN32)
|
||||
# Mark the ".lib" location as part of a Windows DLL
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES IMPORTED_IMPLIB "${ARG_SHARED_LIB}")
|
||||
else ()
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES IMPORTED_LOCATION "${ARG_SHARED_LIB}")
|
||||
endif ()
|
||||
message(STATUS "Added shared library dependency ${AUG_LIB_NAME}: ${ARG_SHARED_LIB}")
|
||||
if (ARG_DEPS)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES INTERFACE_LINK_LIBRARIES "${ARG_DEPS}")
|
||||
endif ()
|
||||
if (ARG_INCLUDE_DIRECTORIES)
|
||||
set_target_properties(${AUG_LIB_NAME}
|
||||
PROPERTIES INTERFACE_INCLUDE_DIRECTORIES
|
||||
"${ARG_INCLUDE_DIRECTORIES}")
|
||||
endif ()
|
||||
else ()
|
||||
message(FATAL_ERROR "No static or shared library provided for ${LIB_NAME}")
|
||||
endif ()
|
||||
endfunction()
|
||||
|
||||
MACRO(using_ccache_if_defined KNOWHERE_USE_CCACHE)
|
||||
if (MILVUS_USE_CCACHE)
|
||||
find_program(CCACHE_FOUND ccache)
|
||||
if (CCACHE_FOUND)
|
||||
message(STATUS "Using ccache: ${CCACHE_FOUND}")
|
||||
set_property(GLOBAL PROPERTY RULE_LAUNCH_COMPILE ${CCACHE_FOUND})
|
||||
set_property(GLOBAL PROPERTY RULE_LAUNCH_LINK ${CCACHE_FOUND})
|
||||
# let ccache preserve C++ comments, because some of them may be
|
||||
# meaningful to the compiler
|
||||
set(ENV{CCACHE_COMMENTS} "1")
|
||||
endif (CCACHE_FOUND)
|
||||
endif ()
|
||||
ENDMACRO(using_ccache_if_defined)
|
||||
@ -1,163 +0,0 @@
|
||||
#-------------------------------------------------------------------------------
|
||||
# Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
# or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
#-------------------------------------------------------------------------------
|
||||
|
||||
macro(set_option_category name)
|
||||
set(KNOWHERE_OPTION_CATEGORY ${name})
|
||||
list(APPEND "KNOWHERE_OPTION_CATEGORIES" ${name})
|
||||
endmacro()
|
||||
|
||||
macro(define_option name description default)
|
||||
option(${name} ${description} ${default})
|
||||
list(APPEND "KNOWHERE_${KNOWHERE_OPTION_CATEGORY}_OPTION_NAMES" ${name})
|
||||
set("${name}_OPTION_DESCRIPTION" ${description})
|
||||
set("${name}_OPTION_DEFAULT" ${default})
|
||||
set("${name}_OPTION_TYPE" "bool")
|
||||
endmacro()
|
||||
|
||||
function(list_join lst glue out)
|
||||
if ("${${lst}}" STREQUAL "")
|
||||
set(${out} "" PARENT_SCOPE)
|
||||
return()
|
||||
endif ()
|
||||
|
||||
list(GET ${lst} 0 joined)
|
||||
list(REMOVE_AT ${lst} 0)
|
||||
foreach (item ${${lst}})
|
||||
set(joined "${joined}${glue}${item}")
|
||||
endforeach ()
|
||||
set(${out} ${joined} PARENT_SCOPE)
|
||||
endfunction()
|
||||
|
||||
macro(define_option_string name description default)
|
||||
set(${name} ${default} CACHE STRING ${description})
|
||||
list(APPEND "KNOWHERE_${KNOWHERE_OPTION_CATEGORY}_OPTION_NAMES" ${name})
|
||||
set("${name}_OPTION_DESCRIPTION" ${description})
|
||||
set("${name}_OPTION_DEFAULT" "\"${default}\"")
|
||||
set("${name}_OPTION_TYPE" "string")
|
||||
|
||||
set("${name}_OPTION_ENUM" ${ARGN})
|
||||
list_join("${name}_OPTION_ENUM" "|" "${name}_OPTION_ENUM")
|
||||
if (NOT ("${${name}_OPTION_ENUM}" STREQUAL ""))
|
||||
set_property(CACHE ${name} PROPERTY STRINGS ${ARGN})
|
||||
endif ()
|
||||
endmacro()
|
||||
|
||||
#----------------------------------------------------------------------
|
||||
set_option_category("Thirdparty")
|
||||
|
||||
set(KNOWHERE_DEPENDENCY_SOURCE "AUTO")
|
||||
|
||||
define_option_string(KNOWHERE_DEPENDENCY_SOURCE
|
||||
"Method to use for acquiring KNOWHERE's build dependencies"
|
||||
"AUTO"
|
||||
"BUNDLED"
|
||||
"SYSTEM")
|
||||
|
||||
|
||||
define_option(KNOWHERE_USE_CCACHE "Use ccache when compiling (if available)" ON)
|
||||
|
||||
define_option(KNOWHERE_VERBOSE_THIRDPARTY_BUILD
|
||||
"Show output from ExternalProjects rather than just logging to files" ON)
|
||||
|
||||
define_option(KNOWHERE_BOOST_USE_SHARED "Rely on boost shared libraries where relevant" OFF)
|
||||
|
||||
define_option(KNOWHERE_BOOST_VENDORED "Use vendored Boost instead of existing Boost. \
|
||||
Note that this requires linking Boost statically" OFF)
|
||||
|
||||
define_option(KNOWHERE_BOOST_HEADER_ONLY "Use only BOOST headers" OFF)
|
||||
|
||||
define_option(KNOWHERE_WITH_ARROW "Build with ARROW" OFF)
|
||||
|
||||
define_option(KNOWHERE_WITH_OPENBLAS "Build with OpenBLAS library" ON)
|
||||
|
||||
define_option(KNOWHERE_WITH_FAISS "Build with FAISS library" ON)
|
||||
|
||||
define_option(KNOWHERE_WITH_FAISS_GPU_VERSION "Build with FAISS GPU version" OFF)
|
||||
|
||||
define_option(FAISS_WITH_MKL "Build FAISS with MKL" OFF)
|
||||
|
||||
define_option(MILVUS_CUDA_ARCH "Build with CUDA arch" "DEFAULT")
|
||||
|
||||
#----------------------------------------------------------------------
|
||||
set_option_category("Test and benchmark")
|
||||
|
||||
#----------------------------------------------------------------------
|
||||
macro(config_summary)
|
||||
message(STATUS "---------------------------------------------------------------------")
|
||||
message(STATUS "KNOWHERE version: ${KNOWHERE_VERSION}")
|
||||
message(STATUS)
|
||||
message(STATUS "Build configuration summary:")
|
||||
|
||||
message(STATUS " Generator: ${CMAKE_GENERATOR}")
|
||||
message(STATUS " Build type: ${CMAKE_BUILD_TYPE}")
|
||||
message(STATUS " Source directory: ${CMAKE_CURRENT_SOURCE_DIR}")
|
||||
if (${CMAKE_EXPORT_COMPILE_COMMANDS})
|
||||
message(
|
||||
STATUS " Compile commands: ${CMAKE_CURRENT_BINARY_DIR}/compile_commands.json")
|
||||
endif ()
|
||||
|
||||
foreach (category ${KNOWHERE_OPTION_CATEGORIES})
|
||||
|
||||
message(STATUS)
|
||||
message(STATUS "${category} options:")
|
||||
|
||||
set(option_names ${KNOWHERE_${category}_OPTION_NAMES})
|
||||
|
||||
set(max_value_length 0)
|
||||
foreach (name ${option_names})
|
||||
string(LENGTH "\"${${name}}\"" value_length)
|
||||
if (${max_value_length} LESS ${value_length})
|
||||
set(max_value_length ${value_length})
|
||||
endif ()
|
||||
endforeach ()
|
||||
|
||||
foreach (name ${option_names})
|
||||
if ("${${name}_OPTION_TYPE}" STREQUAL "string")
|
||||
set(value "\"${${name}}\"")
|
||||
else ()
|
||||
set(value "${${name}}")
|
||||
endif ()
|
||||
|
||||
set(default ${${name}_OPTION_DEFAULT})
|
||||
set(description ${${name}_OPTION_DESCRIPTION})
|
||||
string(LENGTH ${description} description_length)
|
||||
if (${description_length} LESS 70)
|
||||
string(
|
||||
SUBSTRING
|
||||
" "
|
||||
${description_length} -1 description_padding)
|
||||
else ()
|
||||
set(description_padding "
|
||||
")
|
||||
endif ()
|
||||
|
||||
set(comment "[${name}]")
|
||||
|
||||
if ("${value}" STREQUAL "${default}")
|
||||
set(comment "[default] ${comment}")
|
||||
endif ()
|
||||
|
||||
if (NOT ("${${name}_OPTION_ENUM}" STREQUAL ""))
|
||||
set(comment "${comment} [${${name}_OPTION_ENUM}]")
|
||||
endif ()
|
||||
|
||||
string(
|
||||
SUBSTRING "${value} "
|
||||
0 ${max_value_length} value)
|
||||
|
||||
message(STATUS " ${description} ${description_padding} ${value} ${comment}")
|
||||
endforeach ()
|
||||
|
||||
endforeach ()
|
||||
|
||||
endmacro()
|
||||
@ -1,244 +0,0 @@
|
||||
#-------------------------------------------------------------------------------
|
||||
# Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
# or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
#-------------------------------------------------------------------------------
|
||||
set(KNOWHERE_THIRDPARTY_DEPENDENCIES
|
||||
OpenBLAS
|
||||
FAISS
|
||||
)
|
||||
|
||||
message(STATUS "Using ${KNOWHERE_DEPENDENCY_SOURCE} approach to find dependencies")
|
||||
|
||||
# For each dependency, set dependency source to global default, if unset
|
||||
foreach (DEPENDENCY ${KNOWHERE_THIRDPARTY_DEPENDENCIES})
|
||||
if ("${${DEPENDENCY}_SOURCE}" STREQUAL "")
|
||||
set(${DEPENDENCY}_SOURCE ${KNOWHERE_DEPENDENCY_SOURCE})
|
||||
endif ()
|
||||
endforeach ()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Identify OS
|
||||
if (UNIX)
|
||||
## Check for Debian GNU/Linux ________________
|
||||
find_file(DEBIAN_FOUND debian_version debconf.conf
|
||||
PATHS /etc
|
||||
)
|
||||
if (DEBIAN_FOUND)
|
||||
set(CMAKE_OS_NAME "debian" CACHE STRING "Operating system name" FORCE)
|
||||
endif (DEBIAN_FOUND)
|
||||
## Check for Fedora _________________________
|
||||
find_file(FEDORA_FOUND fedora-release
|
||||
PATHS /etc
|
||||
)
|
||||
if (FEDORA_FOUND)
|
||||
set(CMAKE_OS_NAME "fedora" CACHE STRING "Operating system name" FORCE)
|
||||
endif (FEDORA_FOUND)
|
||||
## Check for RedHat _________________________
|
||||
find_file(REDHAT_FOUND redhat-release inittab.RH
|
||||
PATHS /etc
|
||||
)
|
||||
if (REDHAT_FOUND)
|
||||
set(CMAKE_OS_NAME "redhat" CACHE STRING "Operating system name" FORCE)
|
||||
endif (REDHAT_FOUND)
|
||||
## Extra check for Ubuntu ____________________
|
||||
if (DEBIAN_FOUND)
|
||||
## At its core Ubuntu is a Debian system, with
|
||||
## a slightly altered configuration; hence from
|
||||
## a first superficial inspection a system will
|
||||
## be considered as Debian, which signifies an
|
||||
## extra check is required.
|
||||
find_file(UBUNTU_EXTRA legal issue
|
||||
PATHS /etc
|
||||
)
|
||||
if (UBUNTU_EXTRA)
|
||||
## Scan contents of file
|
||||
file(STRINGS ${UBUNTU_EXTRA} UBUNTU_FOUND
|
||||
REGEX Ubuntu
|
||||
)
|
||||
## Check result of string search
|
||||
if (UBUNTU_FOUND)
|
||||
set(CMAKE_OS_NAME "ubuntu" CACHE STRING "Operating system name" FORCE)
|
||||
set(DEBIAN_FOUND FALSE)
|
||||
endif (UBUNTU_FOUND)
|
||||
endif (UBUNTU_EXTRA)
|
||||
endif (DEBIAN_FOUND)
|
||||
endif (UNIX)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# thirdparty directory
|
||||
set(THIRDPARTY_DIR "${INDEX_SOURCE_DIR}/thirdparty")
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# ExternalProject options
|
||||
|
||||
string(TOUPPER ${CMAKE_BUILD_TYPE} UPPERCASE_BUILD_TYPE)
|
||||
|
||||
set(FAISS_FLAGS "-DELPP_THREAD_SAFE -fopenmp -Werror=return-type")
|
||||
set(EP_CXX_FLAGS "${FAISS_FLAGS} ${CMAKE_CXX_FLAGS_${UPPERCASE_BUILD_TYPE}}")
|
||||
set(EP_C_FLAGS "${FAISS_FLAGS} ${CMAKE_C_FLAGS_${UPPERCASE_BUILD_TYPE}}")
|
||||
|
||||
if (NOT MSVC)
|
||||
# Set -fPIC on all external projects
|
||||
set(EP_CXX_FLAGS "${EP_CXX_FLAGS} -fPIC")
|
||||
set(EP_C_FLAGS "${EP_C_FLAGS} -fPIC")
|
||||
endif ()
|
||||
|
||||
# CC/CXX environment variables are captured on the first invocation of the
|
||||
# builder (e.g make or ninja) instead of when CMake is invoked into to build
|
||||
# directory. This leads to issues if the variables are exported in a subshell
|
||||
# and the invocation of make/ninja is in distinct subshell without the same
|
||||
# environment (CC/CXX).
|
||||
set(EP_COMMON_TOOLCHAIN -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER}
|
||||
-DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER})
|
||||
|
||||
if (CMAKE_AR)
|
||||
set(EP_COMMON_TOOLCHAIN ${EP_COMMON_TOOLCHAIN} -DCMAKE_AR=${CMAKE_AR})
|
||||
endif ()
|
||||
|
||||
if (CMAKE_RANLIB)
|
||||
set(EP_COMMON_TOOLCHAIN ${EP_COMMON_TOOLCHAIN} -DCMAKE_RANLIB=${CMAKE_RANLIB})
|
||||
endif ()
|
||||
|
||||
# External projects are still able to override the following declarations.
|
||||
# cmake command line will favor the last defined variable when a duplicate is
|
||||
# encountered. This requires that `EP_COMMON_CMAKE_ARGS` is always the first
|
||||
# argument.
|
||||
set(EP_COMMON_CMAKE_ARGS
|
||||
${EP_COMMON_TOOLCHAIN}
|
||||
-DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
|
||||
-DCMAKE_C_FLAGS=${EP_C_FLAGS}
|
||||
-DCMAKE_C_FLAGS_${UPPERCASE_BUILD_TYPE}=${EP_C_FLAGS}
|
||||
-DCMAKE_CXX_FLAGS=${EP_CXX_FLAGS}
|
||||
-DCMAKE_CXX_FLAGS_${UPPERCASE_BUILD_TYPE}=${EP_CXX_FLAGS})
|
||||
|
||||
if (NOT KNOWHERE_VERBOSE_THIRDPARTY_BUILD)
|
||||
set(EP_LOG_OPTIONS LOG_CONFIGURE 1 LOG_BUILD 1 LOG_INSTALL 1 LOG_DOWNLOAD 1)
|
||||
else ()
|
||||
set(EP_LOG_OPTIONS)
|
||||
endif ()
|
||||
|
||||
# Ensure that a default make is set
|
||||
if ("${MAKE}" STREQUAL "")
|
||||
if (NOT MSVC)
|
||||
find_program(MAKE make)
|
||||
endif ()
|
||||
endif ()
|
||||
|
||||
set(MAKE_BUILD_ARGS "-j6")
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Versions and URLs for toolchain builds, which also can be used to configure
|
||||
# offline builds
|
||||
|
||||
# Read toolchain versions from cpp/thirdparty/versions.txt
|
||||
file(STRINGS "${THIRDPARTY_DIR}/versions.txt" TOOLCHAIN_VERSIONS_TXT)
|
||||
foreach (_VERSION_ENTRY ${TOOLCHAIN_VERSIONS_TXT})
|
||||
# Exclude comments
|
||||
if (NOT _VERSION_ENTRY MATCHES "^[^#][A-Za-z0-9-_]+_VERSION=")
|
||||
continue()
|
||||
endif ()
|
||||
|
||||
string(REGEX MATCH "^[^=]*" _LIB_NAME ${_VERSION_ENTRY})
|
||||
string(REPLACE "${_LIB_NAME}=" "" _LIB_VERSION ${_VERSION_ENTRY})
|
||||
|
||||
# Skip blank or malformed lines
|
||||
if (${_LIB_VERSION} STREQUAL "")
|
||||
continue()
|
||||
endif ()
|
||||
|
||||
# For debugging
|
||||
#message(STATUS "${_LIB_NAME}: ${_LIB_VERSION}")
|
||||
|
||||
set(${_LIB_NAME} "${_LIB_VERSION}")
|
||||
endforeach ()
|
||||
|
||||
if (DEFINED ENV{KNOWHERE_OPENBLAS_URL})
|
||||
set(OPENBLAS_SOURCE_URL "$ENV{KNOWHERE_OPENBLAS_URL}")
|
||||
else ()
|
||||
set(OPENBLAS_SOURCE_URL
|
||||
"https://github.com.cnpmjs.org/xianyi/OpenBLAS/archive/v${OPENBLAS_VERSION}.tar.gz")
|
||||
endif ()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Openblas
|
||||
macro(build_openblas)
|
||||
message(STATUS "Building OpenBLAS-${OPENBLAS_VERSION} from source")
|
||||
set (KNOWHERE_INSTALL_PREFIX ${CMAKE_INSTALL_PREFIX})
|
||||
|
||||
set(OPENBLAS_CMAKE_ARGS
|
||||
${EP_COMMON_CMAKE_ARGS}
|
||||
-DCMAKE_BUILD_TYPE=Release
|
||||
-DBUILD_SHARED_LIBS=ON
|
||||
-DBUILD_STATIC_LIBS=ON
|
||||
-DTARGET=CORE2
|
||||
-DDYNAMIC_ARCH=1
|
||||
-DDYNAMIC_OLDER=1
|
||||
-DUSE_THREAD=0
|
||||
-DUSE_OPENMP=0
|
||||
-DFC=gfortran
|
||||
-DCC=gcc
|
||||
-DINTERFACE64=0
|
||||
-DNUM_THREADS=128
|
||||
-DNO_LAPACKE=0
|
||||
"-DVERSION=${OPENBLAS_VERSION}"
|
||||
"-DCMAKE_INSTALL_PREFIX=${KNOWHERE_INSTALL_PREFIX}"
|
||||
)
|
||||
|
||||
externalproject_add(openblas_ep
|
||||
URL ${OPENBLAS_SOURCE_URL}
|
||||
URL_MD5 "28cc19a6acbf636f5aab5f10b9a0dfe1"
|
||||
CMAKE_ARGS ${OPENBLAS_CMAKE_ARGS}
|
||||
BUILD_COMMAND ${MAKE} ${MAKE_BUILD_ARGS}
|
||||
PREFIX ${CMAKE_BINARY_DIR}/3rdparty_download/openblas-subbuild
|
||||
BINARY_DIR openblas-bin
|
||||
INSTALL_DIR ${KNOWHERE_INSTALL_PREFIX}
|
||||
)
|
||||
|
||||
ExternalProject_Get_Property(openblas_ep INSTALL_DIR)
|
||||
|
||||
if( NOT IS_DIRECTORY ${INSTALL_DIR}/include )
|
||||
file( MAKE_DIRECTORY "${INSTALL_DIR}/include" )
|
||||
endif()
|
||||
|
||||
|
||||
add_library(openblas SHARED IMPORTED)
|
||||
set_target_properties( openblas
|
||||
PROPERTIES
|
||||
IMPORTED_GLOBAL TRUE
|
||||
IMPORTED_LOCATION ${INSTALL_DIR}/${CMAKE_INSTALL_LIBDIR}/libopenblas.so
|
||||
INTERFACE_INCLUDE_DIRECTORIES ${INSTALL_DIR}/${CMAKE_INSTALL_INCLUDEDIR})
|
||||
|
||||
add_dependencies(openblas openblas_ep)
|
||||
endmacro()
|
||||
|
||||
if (KNOWHERE_WITH_OPENBLAS)
|
||||
if (OpenBLAS_SOURCE STREQUAL "AUTO")
|
||||
set (BLA_VENDOR OpenBLAS)
|
||||
find_package(BLAS)
|
||||
|
||||
message(STATUS "Knowhere openblas libraries: ${BLAS_LIBRARIES}")
|
||||
message(STATUS "Knowhere openblas found: ${BLAS_FOUND}")
|
||||
|
||||
if (BLAS_FOUND)
|
||||
add_library(openblas ALIAS BLAS::BLAS)
|
||||
else()
|
||||
build_openblas()
|
||||
endif()
|
||||
|
||||
elseif (OpenBLAS_SOURCE STREQUAL "BUNDLED")
|
||||
build_openblas()
|
||||
elseif (OpenBLAS_SOURCE STREQUAL "SYSTEM")
|
||||
set (BLA_VENDOR OpenBLAS)
|
||||
find_package(BLAS REQUIRED)
|
||||
add_library(openblas ALIAS BLAS::BLAS)
|
||||
endif ()
|
||||
endif()
|
||||
@ -1,183 +0,0 @@
|
||||
#-------------------------------------------------------------------------------
|
||||
# Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
# or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
#-------------------------------------------------------------------------------
|
||||
|
||||
include_directories(${INDEX_SOURCE_DIR}/knowhere)
|
||||
include_directories(${INDEX_SOURCE_DIR}/thirdparty)
|
||||
include_directories(${INDEX_SOURCE_DIR}/../../thirdparty)
|
||||
include_directories(${INDEX_SOURCE_DIR}/thirdparty/NGT/lib)
|
||||
|
||||
set(KNOWHERE_THIRDPARTY_SRC ${INDEX_SOURCE_DIR}/thirdparty)
|
||||
if (MILVUS_SUPPORT_SPTAG)
|
||||
include_directories(${INDEX_SOURCE_DIR}/thirdparty/SPTAG/AnnService)
|
||||
|
||||
set(SPTAG_SOURCE_DIR ${INDEX_SOURCE_DIR}/thirdparty/SPTAG)
|
||||
file(GLOB HDR_FILES
|
||||
${SPTAG_SOURCE_DIR}/AnnService/inc/Core/*.h
|
||||
${SPTAG_SOURCE_DIR}/AnnService/inc/Core/Common/*.h
|
||||
${SPTAG_SOURCE_DIR}/AnnService/inc/Core/BKT/*.h
|
||||
${SPTAG_SOURCE_DIR}/AnnService/inc/Core/KDT/*.h
|
||||
${SPTAG_SOURCE_DIR}/AnnService/inc/Helper/*.h)
|
||||
file(GLOB SRC_FILES
|
||||
${SPTAG_SOURCE_DIR}/AnnService/src/Core/*.cpp
|
||||
${SPTAG_SOURCE_DIR}/AnnService/src/Core/Common/*.cpp
|
||||
${SPTAG_SOURCE_DIR}/AnnService/src/Core/BKT/*.cpp
|
||||
${SPTAG_SOURCE_DIR}/AnnService/src/Core/KDT/*.cpp
|
||||
${SPTAG_SOURCE_DIR}/AnnService/src/Helper/*.cpp)
|
||||
|
||||
if (NOT TARGET SPTAGLibStatic)
|
||||
add_library(SPTAGLibStatic STATIC ${SRC_FILES} ${HDR_FILES})
|
||||
endif ()
|
||||
endif ()
|
||||
|
||||
set(external_srcs
|
||||
knowhere/common/Exception.cpp
|
||||
knowhere/common/Timer.cpp
|
||||
knowhere/common/Utils.cpp
|
||||
)
|
||||
|
||||
set (LOG_SRC
|
||||
knowhere/common/Log.cpp
|
||||
${INDEX_SOURCE_DIR}/../../thirdparty/easyloggingpp/easylogging++.cc
|
||||
)
|
||||
add_library(index_log STATIC ${LOG_SRC})
|
||||
set_target_properties(index_log PROPERTIES RULE_LAUNCH_COMPILE "")
|
||||
set_target_properties(index_log PROPERTIES RULE_LAUNCH_LINK "")
|
||||
|
||||
set(config_srcs
|
||||
knowhere/archive/KnowhereConfig.cpp
|
||||
)
|
||||
|
||||
set(vector_index_srcs
|
||||
knowhere/index/IndexType.cpp
|
||||
knowhere/index/vector_index/adapter/VectorAdapter.cpp
|
||||
knowhere/index/vector_index/helpers/FaissIO.cpp
|
||||
knowhere/index/vector_index/helpers/IndexParameter.cpp
|
||||
knowhere/index/vector_index/helpers/DynamicResultSet.cpp
|
||||
knowhere/index/vector_index/impl/nsg/Distance.cpp
|
||||
knowhere/index/vector_index/impl/nsg/NSG.cpp
|
||||
knowhere/index/vector_index/impl/nsg/NSGHelper.cpp
|
||||
knowhere/index/vector_index/impl/nsg/NSGIO.cpp
|
||||
knowhere/index/vector_index/ConfAdapter.cpp
|
||||
knowhere/index/vector_index/ConfAdapterMgr.cpp
|
||||
knowhere/index/vector_index/FaissBaseBinaryIndex.cpp
|
||||
knowhere/index/vector_index/FaissBaseIndex.cpp
|
||||
knowhere/index/vector_index/IndexBinaryIDMAP.cpp
|
||||
knowhere/index/vector_index/IndexBinaryIVF.cpp
|
||||
knowhere/index/vector_index/IndexIDMAP.cpp
|
||||
knowhere/index/vector_index/IndexIVF.cpp
|
||||
knowhere/index/vector_index/IndexIVFPQ.cpp
|
||||
knowhere/index/vector_index/IndexIVFSQ.cpp
|
||||
knowhere/index/vector_index/IndexIVFHNSW.cpp
|
||||
knowhere/index/vector_index/IndexAnnoy.cpp
|
||||
knowhere/index/vector_index/IndexRHNSW.cpp
|
||||
knowhere/index/vector_index/IndexHNSW.cpp
|
||||
knowhere/index/vector_index/IndexRHNSWFlat.cpp
|
||||
knowhere/index/vector_index/IndexRHNSWSQ.cpp
|
||||
knowhere/index/vector_index/IndexRHNSWPQ.cpp
|
||||
knowhere/index/vector_index/IndexNGT.cpp
|
||||
knowhere/index/vector_index/IndexNGTPANNG.cpp
|
||||
knowhere/index/vector_index/IndexNGTONNG.cpp
|
||||
knowhere/index/vector_index/Statistics.cpp
|
||||
knowhere/index/vector_index/VecIndexFactory.cpp
|
||||
)
|
||||
|
||||
set(vector_offset_index_srcs
|
||||
knowhere/index/vector_offset_index/OffsetBaseIndex.cpp
|
||||
knowhere/index/vector_offset_index/IndexIVF_NM.cpp
|
||||
knowhere/index/vector_offset_index/IndexNSG_NM.cpp
|
||||
)
|
||||
|
||||
if (MILVUS_SUPPORT_SPTAG)
|
||||
set(vector_index_srcs
|
||||
knowhere/index/vector_index/adapter/SptagAdapter.cpp
|
||||
knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp
|
||||
knowhere/index/vector_index/IndexSPTAG.cpp
|
||||
${vector_index_srcs}
|
||||
)
|
||||
endif ()
|
||||
|
||||
set(depend_libs
|
||||
faiss
|
||||
gomp
|
||||
gfortran
|
||||
pthread
|
||||
fiu
|
||||
ngt
|
||||
index_log
|
||||
)
|
||||
|
||||
if (MILVUS_SUPPORT_SPTAG)
|
||||
set(depend_libs
|
||||
SPTAGLibStatic
|
||||
${depend_libs}
|
||||
)
|
||||
endif ()
|
||||
|
||||
|
||||
if (MILVUS_GPU_VERSION)
|
||||
include_directories(${CUDA_INCLUDE_DIRS})
|
||||
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
|
||||
set(cuda_lib
|
||||
cudart
|
||||
cublas
|
||||
)
|
||||
set(depend_libs ${depend_libs}
|
||||
${cuda_lib}
|
||||
)
|
||||
|
||||
set(vector_index_srcs ${vector_index_srcs}
|
||||
knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp
|
||||
knowhere/index/vector_index/gpu/IndexGPUIVF.cpp
|
||||
knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp
|
||||
knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp
|
||||
knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp
|
||||
knowhere/index/vector_index/helpers/Cloner.cpp
|
||||
knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
|
||||
)
|
||||
|
||||
set(vector_offset_index_srcs ${vector_offset_index_srcs}
|
||||
knowhere/index/vector_offset_index/gpu/IndexGPUIVF_NM.cpp
|
||||
)
|
||||
endif ()
|
||||
|
||||
if (NOT TARGET knowhere)
|
||||
add_library(
|
||||
knowhere STATIC
|
||||
${external_srcs}
|
||||
${config_srcs}
|
||||
${vector_index_srcs}
|
||||
${vector_offset_index_srcs}
|
||||
)
|
||||
endif ()
|
||||
|
||||
target_link_libraries(
|
||||
knowhere
|
||||
${depend_libs}
|
||||
)
|
||||
|
||||
set(INDEX_INCLUDE_DIRS
|
||||
${INDEX_SOURCE_DIR}/knowhere
|
||||
${INDEX_SOURCE_DIR}/thirdparty
|
||||
${FAISS_INCLUDE_DIR}
|
||||
${OpenBLAS_INCLUDE_DIR}
|
||||
${LAPACK_INCLUDE_DIR}
|
||||
)
|
||||
|
||||
if (MILVUS_SUPPORT_SPTAG)
|
||||
set(INDEX_INCLUDE_DIRS
|
||||
${INDEX_SOURCE_DIR}/thirdparty/SPTAG/AnnService
|
||||
${INDEX_INCLUDE_DIRS}
|
||||
)
|
||||
endif ()
|
||||
|
||||
set(INDEX_INCLUDE_DIRS ${INDEX_INCLUDE_DIRS} PARENT_SCOPE)
|
||||
29
internal/core/src/index/knowhere/cache/DataObj.h
vendored
29
internal/core/src/index/knowhere/cache/DataObj.h
vendored
@ -1,29 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
|
||||
namespace milvus {
|
||||
namespace cache {
|
||||
|
||||
class DataObj {
|
||||
public:
|
||||
virtual int64_t
|
||||
Size() = 0;
|
||||
virtual ~DataObj() = default;
|
||||
};
|
||||
|
||||
using DataObjPtr = std::shared_ptr<DataObj>;
|
||||
|
||||
} // namespace cache
|
||||
} // namespace milvus
|
||||
@ -1,118 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "KnowhereConfig.h"
|
||||
#include "NGT/lib/NGT/defines.h"
|
||||
#include "faiss/Clustering.h"
|
||||
#include "faiss/FaissHook.h"
|
||||
#include "faiss/common.h"
|
||||
#include "faiss/utils/distances.h"
|
||||
#include "faiss/utils/utils.h"
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/IndexHNSW.h"
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
|
||||
#endif
|
||||
|
||||
namespace milvus {
|
||||
namespace engine {
|
||||
|
||||
constexpr int64_t M_BYTE = 1024 * 1024;
|
||||
|
||||
std::string
|
||||
KnowhereConfig::SetSimdType(const SimdType simd_type) {
|
||||
if (simd_type == SimdType::AUTO) {
|
||||
faiss::faiss_use_avx512 = true;
|
||||
faiss::faiss_use_avx2 = true;
|
||||
faiss::faiss_use_sse4_2 = true;
|
||||
LOG_KNOWHERE_DEBUG_ << "FAISS expect simdType::AUTO";
|
||||
} else if (simd_type == SimdType::AVX512) {
|
||||
faiss::faiss_use_avx512 = true;
|
||||
faiss::faiss_use_avx2 = true;
|
||||
faiss::faiss_use_sse4_2 = true;
|
||||
LOG_KNOWHERE_DEBUG_ << "FAISS expect simdType::AVX512";
|
||||
} else if (simd_type == SimdType::AVX2) {
|
||||
faiss::faiss_use_avx512 = false;
|
||||
faiss::faiss_use_avx2 = true;
|
||||
faiss::faiss_use_sse4_2 = true;
|
||||
LOG_KNOWHERE_DEBUG_ << "FAISS expect simdType::AVX2";
|
||||
} else if (simd_type == SimdType::SSE4_2) {
|
||||
faiss::faiss_use_avx512 = false;
|
||||
faiss::faiss_use_avx2 = false;
|
||||
faiss::faiss_use_sse4_2 = true;
|
||||
LOG_KNOWHERE_DEBUG_ << "FAISS expect simdType::SSE4_2";
|
||||
}
|
||||
|
||||
std::string cpu_flag;
|
||||
faiss::hook_init(cpu_flag);
|
||||
LOG_KNOWHERE_DEBUG_ << "FAISS hook " << cpu_flag;
|
||||
return cpu_flag;
|
||||
}
|
||||
|
||||
void
|
||||
KnowhereConfig::SetBlasThreshold(const int64_t use_blas_threshold) {
|
||||
faiss::distance_compute_blas_threshold = static_cast<int>(use_blas_threshold);
|
||||
}
|
||||
|
||||
void
|
||||
KnowhereConfig::SetEarlyStopThreshold(const double early_stop_threshold) {
|
||||
faiss::early_stop_threshold = early_stop_threshold;
|
||||
}
|
||||
|
||||
void
|
||||
KnowhereConfig::SetClusteringType(const ClusteringType clustering_type) {
|
||||
switch (clustering_type) {
|
||||
case ClusteringType::K_MEANS:
|
||||
default:
|
||||
faiss::clustering_type = faiss::ClusteringType::K_MEANS;
|
||||
break;
|
||||
case ClusteringType::K_MEANS_PLUS_PLUS:
|
||||
faiss::clustering_type = faiss::ClusteringType::K_MEANS_PLUS_PLUS;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
KnowhereConfig::SetStatisticsLevel(const int64_t stat_level) {
|
||||
milvus::knowhere::STATISTICS_LEVEL = stat_level;
|
||||
faiss::STATISTICS_LEVEL = stat_level;
|
||||
}
|
||||
|
||||
void
|
||||
KnowhereConfig::SetLogHandler() {
|
||||
faiss::LOG_ERROR_ = &knowhere::log_error_;
|
||||
faiss::LOG_WARNING_ = &knowhere::log_warning_;
|
||||
// faiss::LOG_DEBUG_ = &knowhere::log_debug_;
|
||||
NGT_LOG_ERROR_ = &knowhere::log_error_;
|
||||
NGT_LOG_WARNING_ = &knowhere::log_warning_;
|
||||
// NGT_LOG_DEBUG_ = &knowhere::log_debug_;
|
||||
}
|
||||
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
void
|
||||
KnowhereConfig::InitGPUResource(const std::vector<int64_t>& gpu_ids) {
|
||||
for (auto id : gpu_ids) {
|
||||
// device_id, pinned_memory, temp_memory, resource_num
|
||||
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(id, 256 * M_BYTE, 256 * M_BYTE, 2);
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
KnowhereConfig::FreeGPUResource() {
|
||||
knowhere::FaissGpuResourceMgr::GetInstance().Free(); // Release gpu resources.
|
||||
}
|
||||
#endif
|
||||
|
||||
} // namespace engine
|
||||
} // namespace milvus
|
||||
@ -1,95 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <vector>
|
||||
#include <string>
|
||||
|
||||
#include "utils/Status.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace engine {
|
||||
|
||||
class KnowhereConfig {
|
||||
public:
|
||||
/**
|
||||
* set SIMD type
|
||||
*/
|
||||
enum SimdType {
|
||||
AUTO = 0, // enable all and depend on the system
|
||||
SSE4_2, // only enable SSE4_2
|
||||
AVX2, // only enable AVX2
|
||||
AVX512, // only enable AVX512
|
||||
};
|
||||
|
||||
static std::string
|
||||
SetSimdType(const SimdType simd_type);
|
||||
|
||||
/**
|
||||
* Set openblas threshold
|
||||
* if nq < use_blas_threshold, calculated by omp
|
||||
* else, calculated by openblas
|
||||
*/
|
||||
static void
|
||||
SetBlasThreshold(const int64_t use_blas_threshold);
|
||||
|
||||
/**
|
||||
* set Clustering early stop [0, 100]
|
||||
* It is to reduce the number of iterations of K-means.
|
||||
* Between each two iterations, if the optimization rate < early_stop_threshold, stop
|
||||
* And if early_stop_threshold = 0, won't early stop
|
||||
*/
|
||||
static void
|
||||
SetEarlyStopThreshold(const double early_stop_threshold);
|
||||
|
||||
/**
|
||||
* set Clustering type
|
||||
*/
|
||||
enum ClusteringType {
|
||||
K_MEANS, // k-means (default)
|
||||
K_MEANS_PLUS_PLUS, // k-means++
|
||||
};
|
||||
|
||||
static void
|
||||
SetClusteringType(const ClusteringType clustering_type);
|
||||
|
||||
/**
|
||||
* set Statistics Level [0, 3]
|
||||
*/
|
||||
static void
|
||||
SetStatisticsLevel(const int64_t stat_level);
|
||||
|
||||
// todo: add log level?
|
||||
/**
|
||||
* set Log handler
|
||||
*/
|
||||
static void
|
||||
SetLogHandler();
|
||||
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
// todo: move to ohter file?
|
||||
/**
|
||||
* init GPU Resource
|
||||
*/
|
||||
static void
|
||||
InitGPUResource(const std::vector<int64_t>& gpu_ids);
|
||||
|
||||
/**
|
||||
* free GPU Resource
|
||||
*/
|
||||
static void
|
||||
FreeGPUResource();
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace engine
|
||||
} // namespace milvus
|
||||
@ -1,92 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string.h>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
struct Binary {
|
||||
std::shared_ptr<uint8_t[]> data;
|
||||
int64_t size = 0;
|
||||
};
|
||||
using BinaryPtr = std::shared_ptr<Binary>;
|
||||
|
||||
inline uint8_t*
|
||||
CopyBinary(const BinaryPtr& bin) {
|
||||
uint8_t* newdata = new uint8_t[bin->size];
|
||||
memcpy(newdata, bin->data.get(), bin->size);
|
||||
return newdata;
|
||||
}
|
||||
|
||||
class BinarySet {
|
||||
public:
|
||||
BinaryPtr
|
||||
GetByName(const std::string& name) const {
|
||||
return binary_map_.at(name);
|
||||
}
|
||||
|
||||
void
|
||||
Append(const std::string& name, BinaryPtr binary) {
|
||||
binary_map_[name] = std::move(binary);
|
||||
}
|
||||
|
||||
void
|
||||
Append(const std::string& name, std::shared_ptr<uint8_t[]> data, int64_t size) {
|
||||
auto binary = std::make_shared<Binary>();
|
||||
binary->data = data;
|
||||
binary->size = size;
|
||||
binary_map_[name] = std::move(binary);
|
||||
}
|
||||
|
||||
// void
|
||||
// Append(const std::string &name, void *data, int64_t size, ID id) {
|
||||
// Binary binary;
|
||||
// binary.data = data;
|
||||
// binary.size = size;
|
||||
// binary.id = id;
|
||||
// binary_map_[name] = binary;
|
||||
//}
|
||||
|
||||
BinaryPtr
|
||||
Erase(const std::string& name) {
|
||||
BinaryPtr result = nullptr;
|
||||
auto it = binary_map_.find(name);
|
||||
if (it != binary_map_.end()) {
|
||||
result = it->second;
|
||||
binary_map_.erase(it);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
void
|
||||
clear() {
|
||||
binary_map_.clear();
|
||||
}
|
||||
|
||||
bool
|
||||
Contains(const std::string& key) const {
|
||||
return binary_map_.find(key) != binary_map_.end();
|
||||
}
|
||||
|
||||
public:
|
||||
std::map<std::string, BinaryPtr> binary_map_;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,22 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "utils/Json.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
using Config = milvus::json;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,75 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <any>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
using Value = std::any;
|
||||
using ValuePtr = std::shared_ptr<Value>;
|
||||
|
||||
class Dataset {
|
||||
public:
|
||||
Dataset() = default;
|
||||
~Dataset() {
|
||||
for (auto const& d : data_) {
|
||||
if (d.first == meta::IDS) {
|
||||
auto row_data = Get<int64_t*>(milvus::knowhere::meta::IDS);
|
||||
// the space of ids must be allocated through malloc
|
||||
free(row_data);
|
||||
}
|
||||
if (d.first == meta::DISTANCE) {
|
||||
auto row_data = Get<float*>(milvus::knowhere::meta::DISTANCE);
|
||||
// the space of distance must be allocated through malloc
|
||||
free(row_data);
|
||||
}
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
void
|
||||
Set(const std::string& k, T&& v) {
|
||||
std::lock_guard<std::mutex> lk(mutex_);
|
||||
data_[k] = std::make_shared<Value>(std::forward<T>(v));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T
|
||||
Get(const std::string& k) {
|
||||
std::lock_guard<std::mutex> lk(mutex_);
|
||||
try {
|
||||
return std::any_cast<T>(*(data_.at(k)));
|
||||
} catch (...) {
|
||||
throw std::logic_error("Can't find this key");
|
||||
}
|
||||
}
|
||||
|
||||
const std::map<std::string, ValuePtr>&
|
||||
data() const {
|
||||
return data_;
|
||||
}
|
||||
|
||||
private:
|
||||
std::mutex mutex_;
|
||||
std::map<std::string, ValuePtr> data_;
|
||||
};
|
||||
using DatasetPtr = std::shared_ptr<Dataset>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,46 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <cstdio>
|
||||
#include <utility>
|
||||
|
||||
#include "Log.h"
|
||||
#include "knowhere/common/Exception.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
KnowhereException::KnowhereException(std::string msg) : msg_(std::move(msg)) {
|
||||
}
|
||||
|
||||
KnowhereException::KnowhereException(const std::string& m, const char* funcName, const char* file, int line) {
|
||||
std::string filename;
|
||||
try {
|
||||
size_t pos;
|
||||
std::string file_path(file);
|
||||
pos = file_path.find_last_of('/');
|
||||
filename = file_path.substr(pos + 1);
|
||||
} catch (std::exception& e) {
|
||||
LOG_KNOWHERE_ERROR_ << e.what();
|
||||
}
|
||||
|
||||
int size = snprintf(nullptr, 0, "Error in %s at %s:%d: %s", funcName, filename.c_str(), line, m.c_str());
|
||||
msg_.resize(size + 1);
|
||||
snprintf(&msg_[0], msg_.size(), "Error in %s at %s:%d: %s", funcName, filename.c_str(), line, m.c_str());
|
||||
}
|
||||
|
||||
const char*
|
||||
KnowhereException::what() const noexcept {
|
||||
return msg_.c_str();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,49 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <exception>
|
||||
#include <string>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class KnowhereException : public std::exception {
|
||||
public:
|
||||
explicit KnowhereException(std::string msg);
|
||||
|
||||
KnowhereException(const std::string& msg, const char* funName, const char* file, int line);
|
||||
|
||||
const char*
|
||||
what() const noexcept override;
|
||||
|
||||
std::string msg_;
|
||||
};
|
||||
|
||||
#define KNOHWERE_ERROR_MSG(MSG) printf("%s", KnowhereException(MSG, __PRETTY_FUNCTION__, __FILE__, __LINE__).what())
|
||||
|
||||
#define KNOWHERE_THROW_MSG(MSG) \
|
||||
do { \
|
||||
throw milvus::knowhere::KnowhereException(MSG, __PRETTY_FUNCTION__, __FILE__, __LINE__); \
|
||||
} while (false)
|
||||
|
||||
#define KNOHERE_THROW_FORMAT(FMT, ...) \
|
||||
do { \
|
||||
std::string __s; \
|
||||
int __size = snprintf(nullptr, 0, FMT, __VA_ARGS__); \
|
||||
__s.resize(__size + 1); \
|
||||
snprintf(&__s[0], __s.size(), FMT, __VA_ARGS__); \
|
||||
throw faiss::FaissException(__s, __PRETTY_FUNCTION__, __FILE__, __LINE__); \
|
||||
} while (false)
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,85 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/common/Log.h"
|
||||
|
||||
#include <cstdarg>
|
||||
#include <cstdio>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
std::string
|
||||
LogOut(const char* pattern, ...) {
|
||||
size_t len = strnlen(pattern, 1024) + 256;
|
||||
auto str_p = std::make_unique<char[]>(len);
|
||||
memset(str_p.get(), 0, len);
|
||||
|
||||
va_list vl;
|
||||
va_start(vl, pattern);
|
||||
vsnprintf(str_p.get(), len - 1, pattern, vl); // NOLINT
|
||||
va_end(vl);
|
||||
|
||||
return std::string(str_p.get());
|
||||
}
|
||||
|
||||
void
|
||||
SetThreadName(const std::string& name) {
|
||||
pthread_setname_np(pthread_self(), name.c_str());
|
||||
}
|
||||
|
||||
std::string
|
||||
GetThreadName() {
|
||||
std::string thread_name = "unamed";
|
||||
char name[16];
|
||||
size_t len = 16;
|
||||
auto err = pthread_getname_np(pthread_self(), name, len);
|
||||
if (not err) {
|
||||
thread_name = name;
|
||||
}
|
||||
|
||||
return thread_name;
|
||||
}
|
||||
|
||||
void
|
||||
log_trace_(const std::string& s) {
|
||||
LOG_KNOWHERE_TRACE_ << s;
|
||||
}
|
||||
|
||||
void
|
||||
log_debug_(const std::string& s) {
|
||||
LOG_KNOWHERE_DEBUG_ << s;
|
||||
}
|
||||
|
||||
void
|
||||
log_info_(const std::string& s) {
|
||||
LOG_KNOWHERE_INFO_ << s;
|
||||
}
|
||||
|
||||
void
|
||||
log_warning_(const std::string& s) {
|
||||
LOG_KNOWHERE_WARNING_ << s;
|
||||
}
|
||||
|
||||
void
|
||||
log_error_(const std::string& s) {
|
||||
LOG_KNOWHERE_ERROR_ << s;
|
||||
}
|
||||
|
||||
void
|
||||
log_fatal_(const std::string& s) {
|
||||
LOG_KNOWHERE_FATAL_ << s;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,77 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "easyloggingpp/easylogging++.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
std::string
|
||||
LogOut(const char* pattern, ...);
|
||||
|
||||
void
|
||||
SetThreadName(const std::string& name);
|
||||
|
||||
std::string
|
||||
GetThreadName();
|
||||
|
||||
void
|
||||
log_trace_(const std::string&);
|
||||
|
||||
void
|
||||
log_debug_(const std::string&);
|
||||
|
||||
void
|
||||
log_info_(const std::string&);
|
||||
|
||||
void
|
||||
log_warning_(const std::string&);
|
||||
|
||||
void
|
||||
log_error_(const std::string&);
|
||||
|
||||
void
|
||||
log_fatal_(const std::string&);
|
||||
|
||||
/*
|
||||
* Please use LOG_MODULE_LEVEL_C macro in member function of class
|
||||
* and LOG_MODULE_LEVEL_ macro in other functions.
|
||||
*/
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
#define KNOWHERE_MODULE_NAME "KNOWHERE"
|
||||
#define KNOWHERE_MODULE_CLASS_FUNCTION \
|
||||
milvus::knowhere::LogOut("[%s][%s::%s][%s] ", KNOWHERE_MODULE_NAME, (typeid(*this).name()), __FUNCTION__, \
|
||||
milvus::knowhere::GetThreadName().c_str())
|
||||
#define KNOWHERE_MODULE_FUNCTION \
|
||||
milvus::knowhere::LogOut("[%s][%s][%s] ", KNOWHERE_MODULE_NAME, __FUNCTION__, \
|
||||
milvus::knowhere::GetThreadName().c_str())
|
||||
|
||||
#define LOG_KNOWHERE_TRACE_C LOG(TRACE) << KNOWHERE_MODULE_CLASS_FUNCTION
|
||||
#define LOG_KNOWHERE_DEBUG_C LOG(DEBUG) << KNOWHERE_MODULE_CLASS_FUNCTION
|
||||
#define LOG_KNOWHERE_INFO_C LOG(INFO) << KNOWHERE_MODULE_CLASS_FUNCTION
|
||||
#define LOG_KNOWHERE_WARNING_C LOG(WARNING) << KNOWHERE_MODULE_CLASS_FUNCTION
|
||||
#define LOG_KNOWHERE_ERROR_C LOG(ERROR) << KNOWHERE_MODULE_CLASS_FUNCTION
|
||||
#define LOG_KNOWHERE_FATAL_C LOG(FATAL) << KNOWHERE_MODULE_CLASS_FUNCTION
|
||||
|
||||
#define LOG_KNOWHERE_TRACE_ LOG(TRACE) << KNOWHERE_MODULE_FUNCTION
|
||||
#define LOG_KNOWHERE_DEBUG_ LOG(DEBUG) << KNOWHERE_MODULE_FUNCTION
|
||||
#define LOG_KNOWHERE_INFO_ LOG(INFO) << KNOWHERE_MODULE_FUNCTION
|
||||
#define LOG_KNOWHERE_WARNING_ LOG(WARNING) << KNOWHERE_MODULE_FUNCTION
|
||||
#define LOG_KNOWHERE_ERROR_ LOG(ERROR) << KNOWHERE_MODULE_FUNCTION
|
||||
#define LOG_KNOWHERE_FATAL_ LOG(FATAL) << KNOWHERE_MODULE_FUNCTION
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,74 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <iostream>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/common/Timer.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
TimeRecorder::TimeRecorder(std::string hdr, int64_t log_level) : header_(std::move(hdr)), log_level_(log_level) {
|
||||
start_ = last_ = stdclock::now();
|
||||
}
|
||||
|
||||
std::string
|
||||
TimeRecorder::GetTimeSpanStr(double span) {
|
||||
std::string str_sec = std::to_string(span * 0.000001) + ((span > 1000000) ? " seconds" : " second");
|
||||
std::string str_ms = std::to_string(span * 0.001) + " ms";
|
||||
|
||||
return str_sec + " [" + str_ms + "]";
|
||||
}
|
||||
|
||||
void
|
||||
TimeRecorder::PrintTimeRecord(const std::string& msg, double span) {
|
||||
std::string str_log;
|
||||
if (!header_.empty()) {
|
||||
str_log += header_ + ": ";
|
||||
}
|
||||
str_log += msg;
|
||||
str_log += " (";
|
||||
str_log += TimeRecorder::GetTimeSpanStr(span);
|
||||
str_log += ")";
|
||||
|
||||
switch (log_level_) {
|
||||
case 0:
|
||||
std::cout << str_log << std::endl;
|
||||
break;
|
||||
default:
|
||||
LOG_KNOWHERE_DEBUG_ << str_log;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
double
|
||||
TimeRecorder::RecordSection(const std::string& msg) {
|
||||
stdclock::time_point curr = stdclock::now();
|
||||
double span = (std::chrono::duration<double, std::micro>(curr - last_)).count();
|
||||
last_ = curr;
|
||||
|
||||
PrintTimeRecord(msg, span);
|
||||
return span;
|
||||
}
|
||||
|
||||
double
|
||||
TimeRecorder::ElapseFromBegin(const std::string& msg) {
|
||||
stdclock::time_point curr = stdclock::now();
|
||||
double span = (std::chrono::duration<double, std::micro>(curr - start_)).count();
|
||||
|
||||
PrintTimeRecord(msg, span);
|
||||
return span;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,49 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <chrono>
|
||||
#include <string>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class TimeRecorder {
|
||||
using stdclock = std::chrono::high_resolution_clock;
|
||||
|
||||
public:
|
||||
// trace = 0, debug = 1, info = 2, warn = 3, error = 4, critical = 5
|
||||
explicit TimeRecorder(std::string hdr, int64_t log_level = 0);
|
||||
virtual ~TimeRecorder() = default;
|
||||
|
||||
double
|
||||
RecordSection(const std::string& msg);
|
||||
|
||||
double
|
||||
ElapseFromBegin(const std::string& msg);
|
||||
|
||||
static std::string
|
||||
GetTimeSpanStr(double span);
|
||||
|
||||
private:
|
||||
void
|
||||
PrintTimeRecord(const std::string& msg, double span);
|
||||
|
||||
private:
|
||||
std::string header_;
|
||||
stdclock::time_point start_;
|
||||
stdclock::time_point last_;
|
||||
int64_t log_level_;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,28 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
using MetricType = std::string;
|
||||
// using IndexType = std::string;
|
||||
using IDType = int64_t;
|
||||
using FloatType = float;
|
||||
using BinaryType = uint8_t;
|
||||
using GraphType = std::vector<std::vector<IDType>>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,111 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/common/Utils.h"
|
||||
#include <algorithm>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
const char* INDEX_FILE_SLICE_SIZE_IN_MEGABYTE = "SLICE_SIZE";
|
||||
const char* INDEX_FILE_SLICE_META = "SLICE_META";
|
||||
|
||||
static const char* META = "meta";
|
||||
static const char* NAME = "name";
|
||||
static const char* SLICE_NUM = "slice_num";
|
||||
static const char* TOTAL_LEN = "total_len";
|
||||
|
||||
void
|
||||
Slice(const std::string& prefix,
|
||||
const BinaryPtr& data_src,
|
||||
const int64_t& slice_len,
|
||||
BinarySet& binarySet,
|
||||
milvus::json& ret) {
|
||||
if (!data_src) {
|
||||
return;
|
||||
}
|
||||
|
||||
int slice_num = 0;
|
||||
for (int64_t i = 0; i < data_src->size; ++slice_num) {
|
||||
int64_t ri = std::min(i + slice_len, data_src->size);
|
||||
auto size = static_cast<size_t>(ri - i);
|
||||
auto slice_i = std::shared_ptr<uint8_t[]>(new uint8_t[size]);
|
||||
memcpy(slice_i.get(), data_src->data.get() + i, size);
|
||||
binarySet.Append(prefix + "_" + std::to_string(slice_num), slice_i, ri - i);
|
||||
i = ri;
|
||||
}
|
||||
ret[NAME] = prefix;
|
||||
ret[SLICE_NUM] = slice_num;
|
||||
ret[TOTAL_LEN] = data_src->size;
|
||||
}
|
||||
|
||||
void
|
||||
Assemble(BinarySet& binarySet) {
|
||||
auto slice_meta = binarySet.Erase(INDEX_FILE_SLICE_META);
|
||||
if (slice_meta == nullptr) {
|
||||
return;
|
||||
}
|
||||
|
||||
milvus::json meta_data =
|
||||
milvus::json::parse(std::string(reinterpret_cast<char*>(slice_meta->data.get()), slice_meta->size));
|
||||
|
||||
for (auto& item : meta_data[META]) {
|
||||
std::string prefix = item[NAME];
|
||||
int slice_num = item[SLICE_NUM];
|
||||
auto total_len = static_cast<size_t>(item[TOTAL_LEN]);
|
||||
auto p_data = std::shared_ptr<uint8_t[]>(new uint8_t[total_len]);
|
||||
int64_t pos = 0;
|
||||
for (auto i = 0; i < slice_num; ++i) {
|
||||
auto slice_i_sp = binarySet.Erase(prefix + "_" + std::to_string(i));
|
||||
memcpy(p_data.get() + pos, slice_i_sp->data.get(), static_cast<size_t>(slice_i_sp->size));
|
||||
pos += slice_i_sp->size;
|
||||
}
|
||||
binarySet.Append(prefix, p_data, total_len);
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
Disassemble(const int64_t& slice_size_in_byte, BinarySet& binarySet) {
|
||||
milvus::json meta_info;
|
||||
auto slice_meta = binarySet.Erase(INDEX_FILE_SLICE_META);
|
||||
if (slice_meta != nullptr) {
|
||||
milvus::json last_meta_data =
|
||||
milvus::json::parse(std::string(reinterpret_cast<char*>(slice_meta->data.get()), slice_meta->size));
|
||||
for (auto& item : last_meta_data[META]) {
|
||||
meta_info[META].emplace_back(item);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::string> slice_key_list;
|
||||
for (auto& kv : binarySet.binary_map_) {
|
||||
if (kv.second->size > slice_size_in_byte) {
|
||||
slice_key_list.push_back(kv.first);
|
||||
}
|
||||
}
|
||||
for (auto& key : slice_key_list) {
|
||||
milvus::json slice_i;
|
||||
Slice(key, binarySet.Erase(key), slice_size_in_byte, binarySet, slice_i);
|
||||
meta_info[META].emplace_back(slice_i);
|
||||
}
|
||||
if (!slice_key_list.empty()) {
|
||||
auto meta_str = meta_info.dump();
|
||||
std::shared_ptr<uint8_t[]> meta_data(new uint8_t[meta_str.length() + 1], std::default_delete<uint8_t[]>());
|
||||
memcpy(meta_data.get(), meta_str.data(), meta_str.length());
|
||||
meta_data.get()[meta_str.length()] = 0;
|
||||
binarySet.Append(INDEX_FILE_SLICE_META, meta_data, meta_str.length() + 1);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,33 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include "BinarySet.h"
|
||||
#include "Config.h"
|
||||
#include "Exception.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
extern const char* INDEX_FILE_SLICE_SIZE_IN_MEGABYTE;
|
||||
extern const char* INDEX_FILE_SLICE_META;
|
||||
|
||||
void
|
||||
Assemble(BinarySet& binarySet);
|
||||
|
||||
void
|
||||
Disassemble(const int64_t& slice_size_in_byte, BinarySet& binarySet);
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,35 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "cache/DataObj.h"
|
||||
#include "knowhere/common/BinarySet.h"
|
||||
#include "knowhere/common/Config.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class Index : public milvus::cache::DataObj {
|
||||
public:
|
||||
virtual BinarySet
|
||||
Serialize(const Config& config) = 0;
|
||||
|
||||
virtual void
|
||||
Load(const BinarySet&) = 0;
|
||||
};
|
||||
|
||||
using IndexPtr = std::shared_ptr<Index>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,46 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <unordered_map>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
/* used in 0.8.0 */
|
||||
namespace IndexEnum {
|
||||
const char* INVALID = "";
|
||||
const char* INDEX_FAISS_IDMAP = "FLAT";
|
||||
const char* INDEX_FAISS_IVFFLAT = "IVF_FLAT";
|
||||
const char* INDEX_FAISS_IVFPQ = "IVF_PQ";
|
||||
const char* INDEX_FAISS_IVFSQ8 = "IVF_SQ8";
|
||||
const char* INDEX_FAISS_IVFSQ8H = "IVF_SQ8_HYBRID";
|
||||
const char* INDEX_FAISS_IVFHNSW = "IVF_HNSW";
|
||||
const char* INDEX_FAISS_BIN_IDMAP = "BIN_FLAT";
|
||||
const char* INDEX_FAISS_BIN_IVFFLAT = "BIN_IVF_FLAT";
|
||||
const char* INDEX_NSG = "NSG";
|
||||
#ifdef MILVUS_SUPPORT_SPTAG
|
||||
const char* INDEX_SPTAG_KDT_RNT = "SPTAG_KDT_RNT";
|
||||
const char* INDEX_SPTAG_BKT_RNT = "SPTAG_BKT_RNT";
|
||||
#endif
|
||||
const char* INDEX_HNSW = "HNSW";
|
||||
const char* INDEX_RHNSWFlat = "RHNSW_FLAT";
|
||||
const char* INDEX_RHNSWPQ = "RHNSW_PQ";
|
||||
const char* INDEX_RHNSWSQ = "RHNSW_SQ";
|
||||
const char* INDEX_ANNOY = "ANNOY";
|
||||
const char* INDEX_NGTPANNG = "NGT_PANNG";
|
||||
const char* INDEX_NGTONNG = "NGT_ONNG";
|
||||
} // namespace IndexEnum
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,77 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
/* used in 0.7.0 */
|
||||
enum class OldIndexType {
|
||||
INVALID = 0,
|
||||
FAISS_IDMAP = 1,
|
||||
FAISS_IVFFLAT_CPU,
|
||||
FAISS_IVFFLAT_GPU,
|
||||
FAISS_IVFFLAT_MIX, // build on gpu and search on cpu
|
||||
FAISS_IVFPQ_CPU,
|
||||
FAISS_IVFPQ_GPU,
|
||||
SPTAG_KDT_RNT_CPU,
|
||||
FAISS_IVFSQ8_MIX,
|
||||
FAISS_IVFSQ8_CPU,
|
||||
FAISS_IVFSQ8_GPU,
|
||||
FAISS_IVFSQ8_HYBRID, // only support build on gpu.
|
||||
NSG_MIX,
|
||||
FAISS_IVFPQ_MIX,
|
||||
SPTAG_BKT_RNT_CPU,
|
||||
HNSW,
|
||||
ANNOY,
|
||||
RHNSW_FLAT,
|
||||
RHNSW_PQ,
|
||||
RHNSW_SQ,
|
||||
NGTPANNG,
|
||||
NGTONNG,
|
||||
FAISS_BIN_IDMAP = 100,
|
||||
FAISS_BIN_IVFLAT_CPU = 101,
|
||||
};
|
||||
|
||||
using IndexType = std::string;
|
||||
|
||||
/* used in 0.8.0 */
|
||||
namespace IndexEnum {
|
||||
extern const char* INVALID;
|
||||
extern const char* INDEX_FAISS_IDMAP;
|
||||
extern const char* INDEX_FAISS_IVFFLAT;
|
||||
extern const char* INDEX_FAISS_IVFPQ;
|
||||
extern const char* INDEX_FAISS_IVFSQ8;
|
||||
extern const char* INDEX_FAISS_IVFSQ8H;
|
||||
extern const char* INDEX_FAISS_IVFHNSW;
|
||||
extern const char* INDEX_FAISS_BIN_IDMAP;
|
||||
extern const char* INDEX_FAISS_BIN_IVFFLAT;
|
||||
extern const char* INDEX_NSG;
|
||||
#ifdef MILVUS_SUPPORT_SPTAG
|
||||
extern const char* INDEX_SPTAG_KDT_RNT;
|
||||
extern const char* INDEX_SPTAG_BKT_RNT;
|
||||
#endif
|
||||
extern const char* INDEX_HNSW;
|
||||
extern const char* INDEX_RHNSWFlat;
|
||||
extern const char* INDEX_RHNSWPQ;
|
||||
extern const char* INDEX_RHNSWSQ;
|
||||
extern const char* INDEX_ANNOY;
|
||||
extern const char* INDEX_NGTPANNG;
|
||||
extern const char* INDEX_NGTONNG;
|
||||
} // namespace IndexEnum
|
||||
|
||||
enum class IndexMode { MODE_CPU = 0, MODE_GPU = 1 };
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,30 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "knowhere/common/Dataset.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class Preprocessor {
|
||||
public:
|
||||
virtual DatasetPtr
|
||||
Preprocess(const DatasetPtr& input) = 0;
|
||||
};
|
||||
|
||||
using PreprocessorPtr = std::shared_ptr<Preprocessor>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,84 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include "faiss/utils/BitsetView.h"
|
||||
#include "knowhere/index/Index.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
enum OperatorType { LT = 0, LE = 1, GT = 3, GE = 4 };
|
||||
|
||||
static std::map<std::string, OperatorType> s_map_operator_type = {
|
||||
{"LT", OperatorType::LT}, {"LTE", OperatorType::LE}, {"GT", OperatorType::GT}, {"GTE", OperatorType::GE},
|
||||
{"lt", OperatorType::LT}, {"lte", OperatorType::LE}, {"gt", OperatorType::GT}, {"gte", OperatorType::GE},
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct IndexStructure {
|
||||
IndexStructure() : a_(0), idx_(0) {
|
||||
}
|
||||
explicit IndexStructure(const T a) : a_(a), idx_(0) {
|
||||
}
|
||||
IndexStructure(const T a, const size_t idx) : a_(a), idx_(idx) {
|
||||
}
|
||||
bool
|
||||
operator<(const IndexStructure& b) const {
|
||||
return a_ < b.a_;
|
||||
}
|
||||
bool
|
||||
operator<=(const IndexStructure& b) const {
|
||||
return a_ <= b.a_;
|
||||
}
|
||||
bool
|
||||
operator>(const IndexStructure& b) const {
|
||||
return a_ > b.a_;
|
||||
}
|
||||
bool
|
||||
operator>=(const IndexStructure& b) const {
|
||||
return a_ >= b.a_;
|
||||
}
|
||||
bool
|
||||
operator==(const IndexStructure& b) const {
|
||||
return a_ == b.a_;
|
||||
}
|
||||
T a_;
|
||||
size_t idx_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class StructuredIndex : public Index {
|
||||
public:
|
||||
virtual void
|
||||
Build(const size_t n, const T* values) = 0;
|
||||
|
||||
virtual const faiss::ConcurrentBitsetPtr
|
||||
In(const size_t n, const T* values) = 0;
|
||||
|
||||
virtual const faiss::ConcurrentBitsetPtr
|
||||
NotIn(const size_t n, const T* values) = 0;
|
||||
|
||||
virtual const faiss::ConcurrentBitsetPtr
|
||||
Range(const T value, const OperatorType op) = 0;
|
||||
|
||||
virtual const faiss::ConcurrentBitsetPtr
|
||||
Range(const T lower_bound_value, bool lb_inclusive, const T upper_bound_value, bool ub_inclusive) = 0;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using StructuredIndexPtr = std::shared_ptr<StructuredIndex<T>>;
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,153 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/structured_index/StructuredIndexFlat.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexFlat<T>::StructuredIndexFlat() : is_built_(false), data_() {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexFlat<T>::StructuredIndexFlat(const size_t n, const T* values) : is_built_(false) {
|
||||
Build(n, values);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexFlat<T>::~StructuredIndexFlat() {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
BinarySet
|
||||
StructuredIndexFlat<T>::Serialize(const milvus::knowhere::Config& config) {
|
||||
// TODO
|
||||
return BinarySet();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexFlat<T>::Load(const milvus::knowhere::BinarySet& index_binary) {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexFlat<T>::Build(const size_t n, const T* values) {
|
||||
data_.reserve(n);
|
||||
T* p = const_cast<T*>(values);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
data_.emplace_back(IndexStructure(*p++, i));
|
||||
}
|
||||
is_built_ = true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
StructuredIndexFlat<T>::In(const size_t n, const T* values) {
|
||||
faiss::ConcurrentBitsetPtr bitset = std::make_shared<faiss::ConcurrentBitset>(data_.size());
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
for (const auto& index : data_) {
|
||||
if (index.a_ == *(values + i)) {
|
||||
bitset->set(index.idx_);
|
||||
}
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
StructuredIndexFlat<T>::NotIn(const size_t n, const T* values) {
|
||||
faiss::ConcurrentBitsetPtr bitset = std::make_shared<faiss::ConcurrentBitset>(data_.size(), 0xff);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
for (const auto& index : data_) {
|
||||
if (index.a_ == *(values + i)) {
|
||||
bitset->clear(index.idx_);
|
||||
}
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
StructuredIndexFlat<T>::Range(const T value, const OperatorType op) {
|
||||
faiss::ConcurrentBitsetPtr bitset = std::make_shared<faiss::ConcurrentBitset>(data_.size());
|
||||
auto lb = data_.begin();
|
||||
auto ub = data_.end();
|
||||
for (; lb < ub; lb++) {
|
||||
switch (op) {
|
||||
case OperatorType::LT:
|
||||
if (*lb < IndexStructure<T>(value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
break;
|
||||
case OperatorType::LE:
|
||||
if (*lb <= IndexStructure<T>(value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
break;
|
||||
case OperatorType::GT:
|
||||
if (*lb > IndexStructure<T>(value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
break;
|
||||
case OperatorType::GE:
|
||||
if (*lb >= IndexStructure<T>(value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
break;
|
||||
default:
|
||||
KNOWHERE_THROW_MSG("Invalid OperatorType:" + std::to_string((int)op) + "!");
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
StructuredIndexFlat<T>::Range(T lower_bound_value, bool lb_inclusive, T upper_bound_value, bool ub_inclusive) {
|
||||
faiss::ConcurrentBitsetPtr bitset = std::make_shared<faiss::ConcurrentBitset>(data_.size());
|
||||
if (lower_bound_value > upper_bound_value) {
|
||||
std::swap(lower_bound_value, upper_bound_value);
|
||||
std::swap(lb_inclusive, ub_inclusive);
|
||||
}
|
||||
auto lb = data_.begin();
|
||||
auto ub = data_.end();
|
||||
for (; lb < ub; ++lb) {
|
||||
if (lb_inclusive && ub_inclusive) {
|
||||
if (*lb >= IndexStructure<T>(lower_bound_value) && *lb <= IndexStructure<T>(upper_bound_value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
} else if (lb_inclusive && !ub_inclusive) {
|
||||
if (*lb >= IndexStructure<T>(lower_bound_value) && *lb < IndexStructure<T>(upper_bound_value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
} else if (!lb_inclusive && ub_inclusive) {
|
||||
if (*lb > IndexStructure<T>(lower_bound_value) && *lb <= IndexStructure<T>(upper_bound_value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
} else {
|
||||
if (*lb > IndexStructure<T>(lower_bound_value) && *lb < IndexStructure<T>(upper_bound_value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,77 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/structured_index/StructuredIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
template <typename T>
|
||||
class StructuredIndexFlat : public StructuredIndex<T> {
|
||||
public:
|
||||
StructuredIndexFlat();
|
||||
StructuredIndexFlat(const size_t n, const T* values);
|
||||
~StructuredIndexFlat();
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config = Config()) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Build(const size_t n, const T* values) override;
|
||||
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
In(const size_t n, const T* values) override;
|
||||
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
NotIn(const size_t n, const T* values) override;
|
||||
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
Range(const T value, const OperatorType op) override;
|
||||
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
Range(T lower_bound_value, bool lb_inclusive, T upper_bound_value, bool ub_inclusive) override;
|
||||
|
||||
const std::vector<IndexStructure<T>>&
|
||||
GetData() {
|
||||
return data_;
|
||||
}
|
||||
|
||||
int64_t
|
||||
Size() override {
|
||||
return (int64_t)data_.size() * sizeof(IndexStructure<T>);
|
||||
}
|
||||
|
||||
bool
|
||||
IsBuilt() const {
|
||||
return is_built_;
|
||||
}
|
||||
|
||||
private:
|
||||
bool is_built_;
|
||||
std::vector<IndexStructure<T>> data_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using StructuredIndexFlatPtr = std::shared_ptr<StructuredIndexFlat<T>>;
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
|
||||
#include "knowhere/index/structured_index/StructuredIndexFlat-inl.h"
|
||||
@ -1,199 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/structured_index/StructuredIndexSort.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexSort<T>::StructuredIndexSort() : is_built_(false), data_() {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexSort<T>::StructuredIndexSort(const size_t n, const T* values) : is_built_(false) {
|
||||
StructuredIndexSort<T>::Build(n, values);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexSort<T>::~StructuredIndexSort() {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexSort<T>::Build(const size_t n, const T* values) {
|
||||
data_.reserve(n);
|
||||
T* p = const_cast<T*>(values);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
data_.emplace_back(IndexStructure(*p++, i));
|
||||
}
|
||||
build();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexSort<T>::build() {
|
||||
if (is_built_)
|
||||
return;
|
||||
if (data_.size() == 0) {
|
||||
// todo: throw an exception
|
||||
KNOWHERE_THROW_MSG("StructuredIndexSort cannot build null values!");
|
||||
}
|
||||
std::sort(data_.begin(), data_.end());
|
||||
is_built_ = true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
BinarySet
|
||||
StructuredIndexSort<T>::Serialize(const milvus::knowhere::Config& config) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
|
||||
auto index_data_size = data_.size() * sizeof(IndexStructure<T>);
|
||||
std::shared_ptr<uint8_t[]> index_data(new uint8_t[index_data_size]);
|
||||
memcpy(index_data.get(), data_.data(), index_data_size);
|
||||
|
||||
std::shared_ptr<uint8_t[]> index_length(new uint8_t[sizeof(size_t)]);
|
||||
auto index_size = data_.size();
|
||||
memcpy(index_length.get(), &index_size, sizeof(size_t));
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("index_data", index_data, index_data_size);
|
||||
res_set.Append("index_length", index_length, sizeof(size_t));
|
||||
return res_set;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexSort<T>::Load(const milvus::knowhere::BinarySet& index_binary) {
|
||||
try {
|
||||
size_t index_size;
|
||||
auto index_length = index_binary.GetByName("index_length");
|
||||
memcpy(&index_size, index_length->data.get(), (size_t)index_length->size);
|
||||
|
||||
auto index_data = index_binary.GetByName("index_data");
|
||||
data_.resize(index_size);
|
||||
memcpy(data_.data(), index_data->data.get(), (size_t)index_data->size);
|
||||
is_built_ = true;
|
||||
} catch (...) {
|
||||
KNOHWERE_ERROR_MSG("StructuredIndexSort Load failed!");
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
StructuredIndexSort<T>::In(const size_t n, const T* values) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
faiss::ConcurrentBitsetPtr bitset = std::make_shared<faiss::ConcurrentBitset>(data_.size());
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
auto lb = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(*(values + i)));
|
||||
auto ub = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(*(values + i)));
|
||||
for (; lb < ub; ++lb) {
|
||||
if (lb->a_ != *(values + i)) {
|
||||
LOG_KNOWHERE_ERROR_ << "error happens in StructuredIndexSort<T>::In, experted value is: "
|
||||
<< *(values + i) << ", but real value is: " << lb->a_;
|
||||
}
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
StructuredIndexSort<T>::NotIn(const size_t n, const T* values) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
faiss::ConcurrentBitsetPtr bitset = std::make_shared<faiss::ConcurrentBitset>(data_.size(), 0xff);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
auto lb = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(*(values + i)));
|
||||
auto ub = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(*(values + i)));
|
||||
for (; lb < ub; ++lb) {
|
||||
if (lb->a_ != *(values + i)) {
|
||||
LOG_KNOWHERE_ERROR_ << "error happens in StructuredIndexSort<T>::NotIn, experted value is: "
|
||||
<< *(values + i) << ", but real value is: " << lb->a_;
|
||||
}
|
||||
bitset->clear(lb->idx_);
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
StructuredIndexSort<T>::Range(const T value, const OperatorType op) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
faiss::ConcurrentBitsetPtr bitset = std::make_shared<faiss::ConcurrentBitset>(data_.size());
|
||||
auto lb = data_.begin();
|
||||
auto ub = data_.end();
|
||||
switch (op) {
|
||||
case OperatorType::LT:
|
||||
ub = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(value));
|
||||
break;
|
||||
case OperatorType::LE:
|
||||
ub = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(value));
|
||||
break;
|
||||
case OperatorType::GT:
|
||||
lb = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(value));
|
||||
break;
|
||||
case OperatorType::GE:
|
||||
lb = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(value));
|
||||
break;
|
||||
default:
|
||||
KNOWHERE_THROW_MSG("Invalid OperatorType:" + std::to_string((int)op) + "!");
|
||||
}
|
||||
for (; lb < ub; ++lb) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
StructuredIndexSort<T>::Range(T lower_bound_value, bool lb_inclusive, T upper_bound_value, bool ub_inclusive) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
faiss::ConcurrentBitsetPtr bitset = std::make_shared<faiss::ConcurrentBitset>(data_.size());
|
||||
if (lower_bound_value > upper_bound_value) {
|
||||
std::swap(lower_bound_value, upper_bound_value);
|
||||
std::swap(lb_inclusive, ub_inclusive);
|
||||
}
|
||||
auto lb = data_.begin();
|
||||
auto ub = data_.end();
|
||||
if (lb_inclusive) {
|
||||
lb = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(lower_bound_value));
|
||||
} else {
|
||||
lb = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(lower_bound_value));
|
||||
}
|
||||
if (ub_inclusive) {
|
||||
ub = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(upper_bound_value));
|
||||
} else {
|
||||
ub = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(upper_bound_value));
|
||||
}
|
||||
for (; lb < ub; ++lb) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,80 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/structured_index/StructuredIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
template <typename T>
|
||||
class StructuredIndexSort : public StructuredIndex<T> {
|
||||
public:
|
||||
StructuredIndexSort();
|
||||
StructuredIndexSort(const size_t n, const T* values);
|
||||
~StructuredIndexSort();
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config = Config()) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Build(const size_t n, const T* values) override;
|
||||
|
||||
void
|
||||
build();
|
||||
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
In(const size_t n, const T* values) override;
|
||||
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
NotIn(const size_t n, const T* values) override;
|
||||
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
Range(const T value, const OperatorType op) override;
|
||||
|
||||
const faiss::ConcurrentBitsetPtr
|
||||
Range(T lower_bound_value, bool lb_inclusive, T upper_bound_value, bool ub_inclusive) override;
|
||||
|
||||
const std::vector<IndexStructure<T>>&
|
||||
GetData() {
|
||||
return data_;
|
||||
}
|
||||
|
||||
int64_t
|
||||
Size() override {
|
||||
return (int64_t)data_.size() * sizeof(IndexStructure<T>);
|
||||
}
|
||||
|
||||
bool
|
||||
IsBuilt() const {
|
||||
return is_built_;
|
||||
}
|
||||
|
||||
private:
|
||||
bool is_built_;
|
||||
std::vector<IndexStructure<T>> data_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using StructuredIndexSortPtr = std::shared_ptr<StructuredIndexSort<T>>;
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
|
||||
#include "knowhere/index/structured_index/StructuredIndexSort-inl.h"
|
||||
@ -1,87 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include "faiss/utils/BitsetView.h"
|
||||
#include "knowhere/index/Index.h"
|
||||
#include <boost/dynamic_bitset.hpp>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere::scalar {
|
||||
|
||||
enum OperatorType { LT = 0, LE = 1, GT = 3, GE = 4 };
|
||||
|
||||
static std::map<std::string, OperatorType> s_map_operator_type = {
|
||||
{"LT", OperatorType::LT}, {"LTE", OperatorType::LE}, {"GT", OperatorType::GT}, {"GTE", OperatorType::GE},
|
||||
{"lt", OperatorType::LT}, {"lte", OperatorType::LE}, {"gt", OperatorType::GT}, {"gte", OperatorType::GE},
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct IndexStructure {
|
||||
IndexStructure() : a_(0), idx_(0) {
|
||||
}
|
||||
explicit IndexStructure(const T a) : a_(a), idx_(0) {
|
||||
}
|
||||
IndexStructure(const T a, const size_t idx) : a_(a), idx_(idx) {
|
||||
}
|
||||
bool
|
||||
operator<(const IndexStructure& b) const {
|
||||
return a_ < b.a_;
|
||||
}
|
||||
bool
|
||||
operator<=(const IndexStructure& b) const {
|
||||
return a_ <= b.a_;
|
||||
}
|
||||
bool
|
||||
operator>(const IndexStructure& b) const {
|
||||
return a_ > b.a_;
|
||||
}
|
||||
bool
|
||||
operator>=(const IndexStructure& b) const {
|
||||
return a_ >= b.a_;
|
||||
}
|
||||
bool
|
||||
operator==(const IndexStructure& b) const {
|
||||
return a_ == b.a_;
|
||||
}
|
||||
T a_;
|
||||
size_t idx_;
|
||||
};
|
||||
using TargetBitmap = boost::dynamic_bitset<>;
|
||||
using TargetBitmapPtr = std::unique_ptr<TargetBitmap>;
|
||||
|
||||
template <typename T>
|
||||
class StructuredIndex : public Index {
|
||||
public:
|
||||
virtual void
|
||||
Build(const size_t n, const T* values) = 0;
|
||||
|
||||
virtual const TargetBitmapPtr
|
||||
In(const size_t n, const T* values) = 0;
|
||||
|
||||
virtual const TargetBitmapPtr
|
||||
NotIn(const size_t n, const T* values) = 0;
|
||||
|
||||
virtual const TargetBitmapPtr
|
||||
Range(const T value, const OperatorType op) = 0;
|
||||
|
||||
virtual const TargetBitmapPtr
|
||||
Range(const T lower_bound_value, bool lb_inclusive, const T upper_bound_value, bool ub_inclusive) = 0;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using StructuredIndexPtr = std::shared_ptr<StructuredIndex<T>>;
|
||||
} // namespace knowhere::scalar
|
||||
} // namespace milvus
|
||||
@ -1,154 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/structured_index_simple/StructuredIndexFlat.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere::scalar {
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexFlat<T>::StructuredIndexFlat() : is_built_(false), data_() {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexFlat<T>::StructuredIndexFlat(const size_t n, const T* values) : is_built_(false) {
|
||||
Build(n, values);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexFlat<T>::~StructuredIndexFlat() {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexFlat<T>::Build(const size_t n, const T* values) {
|
||||
data_.reserve(n);
|
||||
T* p = const_cast<T*>(values);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
data_.emplace_back(IndexStructure(*p++, i));
|
||||
}
|
||||
is_built_ = true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const TargetBitmapPtr
|
||||
StructuredIndexFlat<T>::In(const size_t n, const T* values) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
TargetBitmapPtr bitset = std::make_unique<TargetBitmap>(data_.size());
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
for (const auto& index : data_) {
|
||||
if (index->a_ == *(values + i)) {
|
||||
bitset->set(index->idx_);
|
||||
}
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const TargetBitmapPtr
|
||||
StructuredIndexFlat<T>::NotIn(const size_t n, const T* values) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
TargetBitmapPtr bitset = std::make_unique<TargetBitmap>(data_.size());
|
||||
bitset->set();
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
for (const auto& index : data_) {
|
||||
if (index->a_ == *(values + i)) {
|
||||
bitset->reset(index->idx_);
|
||||
}
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const TargetBitmapPtr
|
||||
StructuredIndexFlat<T>::Range(const T value, const OperatorType op) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
TargetBitmapPtr bitset = std::make_unique<TargetBitmap>(data_.size());
|
||||
auto lb = data_.begin();
|
||||
auto ub = data_.end();
|
||||
for (; lb < ub; lb++) {
|
||||
switch (op) {
|
||||
case OperatorType::LT:
|
||||
if (lb < IndexStructure<T>(value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
break;
|
||||
case OperatorType::LE:
|
||||
if (lb <= IndexStructure<T>(value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
break;
|
||||
case OperatorType::GT:
|
||||
if (lb > IndexStructure<T>(value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
break;
|
||||
case OperatorType::GE:
|
||||
if (lb >= IndexStructure<T>(value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
break;
|
||||
default:
|
||||
KNOWHERE_THROW_MSG("Invalid OperatorType:" + std::to_string((int)op) + "!");
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const TargetBitmapPtr
|
||||
StructuredIndexFlat<T>::Range(T lower_bound_value, bool lb_inclusive, T upper_bound_value, bool ub_inclusive) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
TargetBitmapPtr bitset = std::make_unique<TargetBitmap>(data_.size());
|
||||
if (lower_bound_value > upper_bound_value) {
|
||||
std::swap(lower_bound_value, upper_bound_value);
|
||||
std::swap(lb_inclusive, ub_inclusive);
|
||||
}
|
||||
auto lb = data_.begin();
|
||||
auto ub = data_.end();
|
||||
for (; lb < ub; ++lb) {
|
||||
if (lb_inclusive && ub_inclusive) {
|
||||
if (lb >= IndexStructure<T>(lower_bound_value) && lb <= IndexStructure<T>(upper_bound_value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
} else if (lb_inclusive && !ub_inclusive) {
|
||||
if (lb >= IndexStructure<T>(lower_bound_value) && lb < IndexStructure<T>(upper_bound_value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
} else if (!lb_inclusive && ub_inclusive) {
|
||||
if (lb > IndexStructure<T>(lower_bound_value) && lb <= IndexStructure<T>(upper_bound_value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
} else {
|
||||
if (lb > IndexStructure<T>(lower_bound_value) && lb < IndexStructure<T>(upper_bound_value)) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
} // namespace knowhere::scalar
|
||||
} // namespace milvus
|
||||
@ -1,80 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/structured_index_simple/StructuredIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere::scalar {
|
||||
|
||||
template <typename T>
|
||||
class StructuredIndexFlat : public StructuredIndex<T> {
|
||||
public:
|
||||
StructuredIndexFlat();
|
||||
StructuredIndexFlat(const size_t n, const T* values);
|
||||
~StructuredIndexFlat();
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config = Config()) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Build(const size_t n, const T* values) override;
|
||||
|
||||
void
|
||||
build();
|
||||
|
||||
const TargetBitmapPtr
|
||||
In(const size_t n, const T* values) override;
|
||||
|
||||
const TargetBitmapPtr
|
||||
NotIn(const size_t n, const T* values) override;
|
||||
|
||||
const TargetBitmapPtr
|
||||
Range(const T value, const OperatorType op) override;
|
||||
|
||||
const TargetBitmapPtr
|
||||
Range(T lower_bound_value, bool lb_inclusive, T upper_bound_value, bool ub_inclusive) override;
|
||||
|
||||
const std::vector<IndexStructure<T>>&
|
||||
GetData() {
|
||||
return data_;
|
||||
}
|
||||
|
||||
int64_t
|
||||
Size() override {
|
||||
return (int64_t)data_.size();
|
||||
}
|
||||
|
||||
bool
|
||||
IsBuilt() const {
|
||||
return is_built_;
|
||||
}
|
||||
|
||||
private:
|
||||
bool is_built_;
|
||||
std::vector<IndexStructure<T>> data_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using StructuredIndexFlatPtr = std::shared_ptr<StructuredIndexFlat<T>>;
|
||||
} // namespace knowhere::scalar
|
||||
} // namespace milvus
|
||||
|
||||
#include "knowhere/index/structured_index_simple/StructuredIndexFlat-inl.h"
|
||||
@ -1,200 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/structured_index_simple/StructuredIndexSort.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere::scalar {
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexSort<T>::StructuredIndexSort() : is_built_(false), data_() {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexSort<T>::StructuredIndexSort(const size_t n, const T* values) : is_built_(false) {
|
||||
StructuredIndexSort<T>::Build(n, values);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
StructuredIndexSort<T>::~StructuredIndexSort() {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexSort<T>::Build(const size_t n, const T* values) {
|
||||
data_.reserve(n);
|
||||
T* p = const_cast<T*>(values);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
data_.emplace_back(IndexStructure(*p++, i));
|
||||
}
|
||||
build();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexSort<T>::build() {
|
||||
if (is_built_)
|
||||
return;
|
||||
if (data_.size() == 0) {
|
||||
// todo: throw an exception
|
||||
KNOWHERE_THROW_MSG("StructuredIndexSort cannot build null values!");
|
||||
}
|
||||
std::sort(data_.begin(), data_.end());
|
||||
is_built_ = true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
BinarySet
|
||||
StructuredIndexSort<T>::Serialize(const milvus::knowhere::Config& config) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
|
||||
auto index_data_size = data_.size() * sizeof(IndexStructure<T>);
|
||||
std::shared_ptr<uint8_t[]> index_data(new uint8_t[index_data_size]);
|
||||
memcpy(index_data.get(), data_.data(), index_data_size);
|
||||
|
||||
std::shared_ptr<uint8_t[]> index_length(new uint8_t[sizeof(size_t)]);
|
||||
auto index_size = data_.size();
|
||||
memcpy(index_length.get(), &index_size, sizeof(size_t));
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("index_data", index_data, index_data_size);
|
||||
res_set.Append("index_length", index_length, sizeof(size_t));
|
||||
return res_set;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void
|
||||
StructuredIndexSort<T>::Load(const milvus::knowhere::BinarySet& index_binary) {
|
||||
try {
|
||||
size_t index_size;
|
||||
auto index_length = index_binary.GetByName("index_length");
|
||||
memcpy(&index_size, index_length->data.get(), (size_t)index_length->size);
|
||||
|
||||
auto index_data = index_binary.GetByName("index_data");
|
||||
data_.resize(index_size);
|
||||
memcpy(data_.data(), index_data->data.get(), (size_t)index_data->size);
|
||||
is_built_ = true;
|
||||
} catch (...) {
|
||||
KNOHWERE_ERROR_MSG("StructuredIndexSort Load failed!");
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const TargetBitmapPtr
|
||||
StructuredIndexSort<T>::In(const size_t n, const T* values) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
TargetBitmapPtr bitset = std::make_unique<TargetBitmap>(data_.size());
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
auto lb = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(*(values + i)));
|
||||
auto ub = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(*(values + i)));
|
||||
for (; lb < ub; ++lb) {
|
||||
if (lb->a_ != *(values + i)) {
|
||||
LOG_KNOWHERE_ERROR_ << "error happens in StructuredIndexSort<T>::In, experted value is: "
|
||||
<< *(values + i) << ", but real value is: " << lb->a_;
|
||||
}
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const TargetBitmapPtr
|
||||
StructuredIndexSort<T>::NotIn(const size_t n, const T* values) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
TargetBitmapPtr bitset = std::make_unique<TargetBitmap>(data_.size());
|
||||
bitset->set();
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
auto lb = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(*(values + i)));
|
||||
auto ub = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(*(values + i)));
|
||||
for (; lb < ub; ++lb) {
|
||||
if (lb->a_ != *(values + i)) {
|
||||
LOG_KNOWHERE_ERROR_ << "error happens in StructuredIndexSort<T>::NotIn, experted value is: "
|
||||
<< *(values + i) << ", but real value is: " << lb->a_;
|
||||
}
|
||||
bitset->reset(lb->idx_);
|
||||
}
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const TargetBitmapPtr
|
||||
StructuredIndexSort<T>::Range(const T value, const OperatorType op) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
TargetBitmapPtr bitset = std::make_unique<TargetBitmap>(data_.size());
|
||||
auto lb = data_.begin();
|
||||
auto ub = data_.end();
|
||||
switch (op) {
|
||||
case OperatorType::LT:
|
||||
ub = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(value));
|
||||
break;
|
||||
case OperatorType::LE:
|
||||
ub = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(value));
|
||||
break;
|
||||
case OperatorType::GT:
|
||||
lb = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(value));
|
||||
break;
|
||||
case OperatorType::GE:
|
||||
lb = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(value));
|
||||
break;
|
||||
default:
|
||||
KNOWHERE_THROW_MSG("Invalid OperatorType:" + std::to_string((int)op) + "!");
|
||||
}
|
||||
for (; lb < ub; ++lb) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const TargetBitmapPtr
|
||||
StructuredIndexSort<T>::Range(T lower_bound_value, bool lb_inclusive, T upper_bound_value, bool ub_inclusive) {
|
||||
if (!is_built_) {
|
||||
build();
|
||||
}
|
||||
TargetBitmapPtr bitset = std::make_unique<TargetBitmap>(data_.size());
|
||||
if (lower_bound_value > upper_bound_value) {
|
||||
std::swap(lower_bound_value, upper_bound_value);
|
||||
std::swap(lb_inclusive, ub_inclusive);
|
||||
}
|
||||
auto lb = data_.begin();
|
||||
auto ub = data_.end();
|
||||
if (lb_inclusive) {
|
||||
lb = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(lower_bound_value));
|
||||
} else {
|
||||
lb = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(lower_bound_value));
|
||||
}
|
||||
if (ub_inclusive) {
|
||||
ub = std::upper_bound(data_.begin(), data_.end(), IndexStructure<T>(upper_bound_value));
|
||||
} else {
|
||||
ub = std::lower_bound(data_.begin(), data_.end(), IndexStructure<T>(upper_bound_value));
|
||||
}
|
||||
for (; lb < ub; ++lb) {
|
||||
bitset->set(lb->idx_);
|
||||
}
|
||||
return bitset;
|
||||
}
|
||||
|
||||
} // namespace knowhere::scalar
|
||||
} // namespace milvus
|
||||
@ -1,80 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/structured_index_simple/StructuredIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere::scalar {
|
||||
|
||||
template <typename T>
|
||||
class StructuredIndexSort : public StructuredIndex<T> {
|
||||
public:
|
||||
StructuredIndexSort();
|
||||
StructuredIndexSort(const size_t n, const T* values);
|
||||
~StructuredIndexSort();
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config = Config()) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Build(const size_t n, const T* values) override;
|
||||
|
||||
void
|
||||
build();
|
||||
|
||||
const TargetBitmapPtr
|
||||
In(size_t n, const T* values) override;
|
||||
|
||||
const TargetBitmapPtr
|
||||
NotIn(size_t n, const T* values) override;
|
||||
|
||||
const TargetBitmapPtr
|
||||
Range(T value, OperatorType op) override;
|
||||
|
||||
const TargetBitmapPtr
|
||||
Range(T lower_bound_value, bool lb_inclusive, T upper_bound_value, bool ub_inclusive) override;
|
||||
|
||||
const std::vector<IndexStructure<T>>&
|
||||
GetData() {
|
||||
return data_;
|
||||
}
|
||||
|
||||
int64_t
|
||||
Size() override {
|
||||
return (int64_t)data_.size();
|
||||
}
|
||||
|
||||
bool
|
||||
IsBuilt() const {
|
||||
return is_built_;
|
||||
}
|
||||
|
||||
private:
|
||||
bool is_built_;
|
||||
std::vector<IndexStructure<T>> data_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using StructuredIndexSortPtr = std::shared_ptr<StructuredIndexSort<T>>;
|
||||
} // namespace knowhere::scalar
|
||||
} // namespace milvus
|
||||
|
||||
#include "knowhere/index/structured_index_simple/StructuredIndexSort-inl.h"
|
||||
@ -1,427 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/ConfAdapter.h"
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include "faiss/gpu/utils/DeviceUtils.h"
|
||||
#endif
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
static const int64_t MIN_NBITS = 1;
|
||||
static const int64_t MAX_NBITS = 16;
|
||||
static const int64_t DEFAULT_NBITS = 8;
|
||||
static const int64_t MIN_NLIST = 1;
|
||||
static const int64_t MAX_NLIST = 65536;
|
||||
static const int64_t MIN_NPROBE = 1;
|
||||
static const int64_t MAX_NPROBE = MAX_NLIST;
|
||||
static const int64_t DEFAULT_MIN_DIM = 1;
|
||||
static const int64_t DEFAULT_MAX_DIM = 32768;
|
||||
static const int64_t NGT_MIN_EDGE_SIZE = 1;
|
||||
static const int64_t NGT_MAX_EDGE_SIZE = 200;
|
||||
static const int64_t HNSW_MIN_EFCONSTRUCTION = 8;
|
||||
static const int64_t HNSW_MAX_EFCONSTRUCTION = 512;
|
||||
static const int64_t HNSW_MIN_M = 4;
|
||||
static const int64_t HNSW_MAX_M = 64;
|
||||
static const int64_t HNSW_MAX_EF = 32768;
|
||||
static const std::vector<std::string> METRICS{knowhere::Metric::L2, knowhere::Metric::IP};
|
||||
|
||||
#define CheckIntByRange(key, min, max) \
|
||||
if (!oricfg.contains(key) || !oricfg[key].is_number_integer() || oricfg[key].get<int64_t>() > max || \
|
||||
oricfg[key].get<int64_t>() < min) { \
|
||||
return false; \
|
||||
}
|
||||
|
||||
#define CheckFloatByRange(key, min, max) \
|
||||
if (!oricfg.contains(key) || !oricfg[key].is_number_float() || oricfg[key].get<float>() > max || \
|
||||
oricfg[key].get<float>() < min) { \
|
||||
return false; \
|
||||
}
|
||||
|
||||
#define CheckIntByValues(key, container) \
|
||||
if (!oricfg.contains(key) || !oricfg[key].is_number_integer()) { \
|
||||
return false; \
|
||||
} else { \
|
||||
auto finder = std::find(std::begin(container), std::end(container), oricfg[key].get<int64_t>()); \
|
||||
if (finder == std::end(container)) { \
|
||||
return false; \
|
||||
} \
|
||||
}
|
||||
|
||||
#define CheckStrByValues(key, container) \
|
||||
if (!oricfg.contains(key) || !oricfg[key].is_string()) { \
|
||||
return false; \
|
||||
} else { \
|
||||
auto finder = std::find(std::begin(container), std::end(container), oricfg[key].get<std::string>()); \
|
||||
if (finder == std::end(container)) { \
|
||||
return false; \
|
||||
} \
|
||||
}
|
||||
|
||||
bool
|
||||
ConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::meta::DIM, DEFAULT_MIN_DIM, DEFAULT_MAX_DIM);
|
||||
CheckStrByValues(knowhere::Metric::TYPE, METRICS);
|
||||
return true;
|
||||
}
|
||||
|
||||
bool
|
||||
ConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
const int64_t DEFAULT_MIN_K = 1;
|
||||
const int64_t DEFAULT_MAX_K = 16384;
|
||||
CheckIntByRange(knowhere::meta::TOPK, DEFAULT_MIN_K - 1, DEFAULT_MAX_K);
|
||||
return true;
|
||||
}
|
||||
|
||||
int64_t
|
||||
MatchNlist(int64_t size, int64_t nlist) {
|
||||
const int64_t MIN_POINTS_PER_CENTROID = 39;
|
||||
|
||||
if (nlist * MIN_POINTS_PER_CENTROID > size) {
|
||||
// nlist is too large, adjust to a proper value
|
||||
nlist = std::max(static_cast<int64_t>(1), size / MIN_POINTS_PER_CENTROID);
|
||||
LOG_KNOWHERE_WARNING_ << "Row num " << size << " match nlist " << nlist;
|
||||
}
|
||||
return nlist;
|
||||
}
|
||||
|
||||
int64_t
|
||||
MatchNbits(int64_t size, int64_t nbits) {
|
||||
if (size < (1 << nbits)) {
|
||||
// nbits is too large, adjust to a proper value
|
||||
if (size >= (1 << 8)) {
|
||||
nbits = 8;
|
||||
} else if (size >= (1 << 4)) {
|
||||
nbits = 4;
|
||||
} else if (size >= (1 << 2)) {
|
||||
nbits = 2;
|
||||
} else {
|
||||
nbits = 1;
|
||||
}
|
||||
LOG_KNOWHERE_WARNING_ << "Row num " << size << " match nbits " << nbits;
|
||||
}
|
||||
return nbits;
|
||||
}
|
||||
|
||||
bool
|
||||
IVFConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::nlist, MIN_NLIST, MAX_NLIST);
|
||||
|
||||
// auto tune params
|
||||
auto rows = oricfg[knowhere::meta::ROWS].get<int64_t>();
|
||||
auto nlist = oricfg[knowhere::IndexParams::nlist].get<int64_t>();
|
||||
oricfg[knowhere::IndexParams::nlist] = MatchNlist(rows, nlist);
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
IVFConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
int64_t max_nprobe = MAX_NPROBE;
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
if (mode == IndexMode::MODE_GPU) {
|
||||
max_nprobe = faiss::gpu::getMaxKSelection();
|
||||
}
|
||||
#endif
|
||||
CheckIntByRange(knowhere::IndexParams::nprobe, MIN_NPROBE, max_nprobe);
|
||||
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
IVFSQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
oricfg[knowhere::IndexParams::nbits] = DEFAULT_NBITS;
|
||||
return IVFConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
IVFPQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
if (!IVFConfAdapter::CheckTrain(oricfg, mode)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
CheckIntByRange(knowhere::IndexParams::nbits, MIN_NBITS, MAX_NBITS);
|
||||
|
||||
auto rows = oricfg[knowhere::meta::ROWS].get<int64_t>();
|
||||
auto nbits = oricfg.count(IndexParams::nbits) ? oricfg[IndexParams::nbits].get<int64_t>() : DEFAULT_NBITS;
|
||||
oricfg[knowhere::IndexParams::nbits] = MatchNbits(rows, nbits);
|
||||
|
||||
auto m = oricfg[knowhere::IndexParams::m].get<int64_t>();
|
||||
auto dimension = oricfg[knowhere::meta::DIM].get<int64_t>();
|
||||
|
||||
IndexMode ivfpq_mode = mode;
|
||||
return CheckPQParams(dimension, m, nbits, ivfpq_mode);
|
||||
}
|
||||
|
||||
bool
|
||||
IVFPQConfAdapter::CheckPQParams(int64_t dimension, int64_t m, int64_t nbits, IndexMode& mode) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
if (mode == knowhere::IndexMode::MODE_GPU && !IVFPQConfAdapter::CheckGPUPQParams(dimension, m, nbits)) {
|
||||
mode = knowhere::IndexMode::MODE_CPU;
|
||||
}
|
||||
#endif
|
||||
if (mode == knowhere::IndexMode::MODE_CPU && !IVFPQConfAdapter::CheckCPUPQParams(dimension, m)) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool
|
||||
IVFPQConfAdapter::CheckGPUPQParams(int64_t dimension, int64_t m, int64_t nbits) {
|
||||
/*
|
||||
* Faiss 1.6
|
||||
* Only 1, 2, 3, 4, 6, 8, 10, 12, 16, 20, 24, 28, 32 dims per sub-quantizer are currently supported with
|
||||
* no precomputed codes. Precomputed codes supports any number of dimensions, but will involve memory overheads.
|
||||
*/
|
||||
static const std::vector<int64_t> support_dim_per_subquantizer{32, 28, 24, 20, 16, 12, 10, 8, 6, 4, 3, 2, 1};
|
||||
static const std::vector<int64_t> support_subquantizer{96, 64, 56, 48, 40, 32, 28, 24, 20, 16, 12, 8, 4, 3, 2, 1};
|
||||
|
||||
if (!CheckCPUPQParams(dimension, m)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
int64_t sub_dim = dimension / m;
|
||||
return (std::find(std::begin(support_subquantizer), std::end(support_subquantizer), m) !=
|
||||
support_subquantizer.end()) &&
|
||||
(std::find(std::begin(support_dim_per_subquantizer), std::end(support_dim_per_subquantizer), sub_dim) !=
|
||||
support_dim_per_subquantizer.end()) &&
|
||||
(nbits == 8);
|
||||
}
|
||||
|
||||
bool
|
||||
IVFPQConfAdapter::CheckCPUPQParams(int64_t dimension, int64_t m) {
|
||||
return (dimension % m == 0);
|
||||
}
|
||||
|
||||
bool
|
||||
IVFHNSWConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
// HNSW param check
|
||||
CheckIntByRange(knowhere::IndexParams::efConstruction, HNSW_MIN_EFCONSTRUCTION, HNSW_MAX_EFCONSTRUCTION);
|
||||
CheckIntByRange(knowhere::IndexParams::M, HNSW_MIN_M, HNSW_MAX_M);
|
||||
|
||||
// IVF param check
|
||||
CheckIntByRange(knowhere::IndexParams::nlist, MIN_NLIST, MAX_NLIST);
|
||||
|
||||
// auto tune params
|
||||
auto rows = oricfg[knowhere::meta::ROWS].get<int64_t>();
|
||||
auto nlist = oricfg[knowhere::IndexParams::nlist].get<int64_t>();
|
||||
oricfg[knowhere::IndexParams::nlist] = MatchNlist(rows, nlist);
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
IVFHNSWConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
// HNSW param check
|
||||
CheckIntByRange(knowhere::IndexParams::ef, oricfg[knowhere::meta::TOPK], HNSW_MAX_EF);
|
||||
|
||||
// IVF param check
|
||||
CheckIntByRange(knowhere::IndexParams::nprobe, MIN_NPROBE, MAX_NPROBE);
|
||||
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
NSGConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
const int64_t MIN_KNNG = 5;
|
||||
const int64_t MAX_KNNG = 300;
|
||||
const int64_t MIN_SEARCH_LENGTH = 10;
|
||||
const int64_t MAX_SEARCH_LENGTH = 300;
|
||||
const int64_t MIN_OUT_DEGREE = 5;
|
||||
const int64_t MAX_OUT_DEGREE = 300;
|
||||
const int64_t MIN_CANDIDATE_POOL_SIZE = 50;
|
||||
const int64_t MAX_CANDIDATE_POOL_SIZE = 1000;
|
||||
|
||||
CheckStrByValues(knowhere::Metric::TYPE, METRICS);
|
||||
CheckIntByRange(knowhere::IndexParams::knng, MIN_KNNG, MAX_KNNG);
|
||||
CheckIntByRange(knowhere::IndexParams::search_length, MIN_SEARCH_LENGTH, MAX_SEARCH_LENGTH);
|
||||
CheckIntByRange(knowhere::IndexParams::out_degree, MIN_OUT_DEGREE, MAX_OUT_DEGREE);
|
||||
CheckIntByRange(knowhere::IndexParams::candidate, MIN_CANDIDATE_POOL_SIZE, MAX_CANDIDATE_POOL_SIZE);
|
||||
|
||||
// auto tune params
|
||||
oricfg[knowhere::IndexParams::nlist] = MatchNlist(oricfg[knowhere::meta::ROWS].get<int64_t>(), 8192);
|
||||
|
||||
int64_t nprobe = int(oricfg[knowhere::IndexParams::nlist].get<int64_t>() * 0.1);
|
||||
oricfg[knowhere::IndexParams::nprobe] = nprobe < 1 ? 1 : nprobe;
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool
|
||||
NSGConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
static int64_t MIN_SEARCH_LENGTH = 1;
|
||||
static int64_t MAX_SEARCH_LENGTH = 300;
|
||||
|
||||
CheckIntByRange(knowhere::IndexParams::search_length, MIN_SEARCH_LENGTH, MAX_SEARCH_LENGTH);
|
||||
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
HNSWConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::efConstruction, HNSW_MIN_EFCONSTRUCTION, HNSW_MAX_EFCONSTRUCTION);
|
||||
CheckIntByRange(knowhere::IndexParams::M, HNSW_MIN_M, HNSW_MAX_M);
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
HNSWConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::ef, oricfg[knowhere::meta::TOPK], HNSW_MAX_EF);
|
||||
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
RHNSWFlatConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::efConstruction, HNSW_MIN_EFCONSTRUCTION, HNSW_MAX_EFCONSTRUCTION);
|
||||
CheckIntByRange(knowhere::IndexParams::M, HNSW_MIN_M, HNSW_MAX_M);
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
RHNSWFlatConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::ef, oricfg[knowhere::meta::TOPK], HNSW_MAX_EF);
|
||||
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
RHNSWPQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::efConstruction, HNSW_MIN_EFCONSTRUCTION, HNSW_MAX_EFCONSTRUCTION);
|
||||
CheckIntByRange(knowhere::IndexParams::M, HNSW_MIN_M, HNSW_MAX_M);
|
||||
|
||||
auto dimension = oricfg[knowhere::meta::DIM].get<int64_t>();
|
||||
|
||||
if (!IVFPQConfAdapter::CheckCPUPQParams(dimension, oricfg[knowhere::IndexParams::PQM].get<int64_t>())) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
RHNSWPQConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::ef, oricfg[knowhere::meta::TOPK], HNSW_MAX_EF);
|
||||
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
RHNSWSQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::efConstruction, HNSW_MIN_EFCONSTRUCTION, HNSW_MAX_EFCONSTRUCTION);
|
||||
CheckIntByRange(knowhere::IndexParams::M, HNSW_MIN_M, HNSW_MAX_M);
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
RHNSWSQConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::ef, oricfg[knowhere::meta::TOPK], HNSW_MAX_EF);
|
||||
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
BinIDMAPConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
static const std::vector<std::string> METRICS{knowhere::Metric::HAMMING, knowhere::Metric::JACCARD,
|
||||
knowhere::Metric::TANIMOTO, knowhere::Metric::SUBSTRUCTURE,
|
||||
knowhere::Metric::SUPERSTRUCTURE};
|
||||
|
||||
CheckIntByRange(knowhere::meta::DIM, DEFAULT_MIN_DIM, DEFAULT_MAX_DIM);
|
||||
CheckStrByValues(knowhere::Metric::TYPE, METRICS);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool
|
||||
BinIVFConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
static const std::vector<std::string> METRICS{knowhere::Metric::HAMMING, knowhere::Metric::JACCARD,
|
||||
knowhere::Metric::TANIMOTO};
|
||||
|
||||
CheckIntByRange(knowhere::meta::DIM, DEFAULT_MIN_DIM, DEFAULT_MAX_DIM);
|
||||
CheckIntByRange(knowhere::IndexParams::nlist, MIN_NLIST, MAX_NLIST);
|
||||
CheckStrByValues(knowhere::Metric::TYPE, METRICS);
|
||||
|
||||
// auto tune params
|
||||
auto rows = oricfg[knowhere::meta::ROWS].get<int64_t>();
|
||||
auto nlist = oricfg[knowhere::IndexParams::nlist].get<int64_t>();
|
||||
oricfg[knowhere::IndexParams::nlist] = MatchNlist(rows, nlist);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool
|
||||
ANNOYConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
static int64_t MIN_NTREES = 1;
|
||||
// too large of n_trees takes much time, if there is real requirement, change this threshold.
|
||||
static int64_t MAX_NTREES = 1024;
|
||||
|
||||
CheckIntByRange(knowhere::IndexParams::n_trees, MIN_NTREES, MAX_NTREES);
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
ANNOYConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::search_k, std::numeric_limits<int64_t>::min(),
|
||||
std::numeric_limits<int64_t>::max());
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
NGTPANNGConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::edge_size, NGT_MIN_EDGE_SIZE, NGT_MAX_EDGE_SIZE);
|
||||
CheckIntByRange(knowhere::IndexParams::forcedly_pruned_edge_size, NGT_MIN_EDGE_SIZE, NGT_MAX_EDGE_SIZE);
|
||||
CheckIntByRange(knowhere::IndexParams::selectively_pruned_edge_size, NGT_MIN_EDGE_SIZE, NGT_MAX_EDGE_SIZE);
|
||||
if (oricfg[knowhere::IndexParams::selectively_pruned_edge_size].get<int64_t>() >=
|
||||
oricfg[knowhere::IndexParams::forcedly_pruned_edge_size].get<int64_t>()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
NGTPANNGConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::max_search_edges, -1, NGT_MAX_EDGE_SIZE);
|
||||
CheckFloatByRange(knowhere::IndexParams::epsilon, -1.0, 1.0);
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
NGTONNGConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::edge_size, NGT_MIN_EDGE_SIZE, NGT_MAX_EDGE_SIZE);
|
||||
CheckIntByRange(knowhere::IndexParams::outgoing_edge_size, NGT_MIN_EDGE_SIZE, NGT_MAX_EDGE_SIZE);
|
||||
CheckIntByRange(knowhere::IndexParams::incoming_edge_size, NGT_MIN_EDGE_SIZE, NGT_MAX_EDGE_SIZE);
|
||||
|
||||
return ConfAdapter::CheckTrain(oricfg, mode);
|
||||
}
|
||||
|
||||
bool
|
||||
NGTONNGConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) {
|
||||
CheckIntByRange(knowhere::IndexParams::max_search_edges, -1, NGT_MAX_EDGE_SIZE);
|
||||
CheckFloatByRange(knowhere::IndexParams::epsilon, -1.0, 1.0);
|
||||
return ConfAdapter::CheckSearch(oricfg, type, mode);
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,158 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Config.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class ConfAdapter {
|
||||
public:
|
||||
virtual bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode);
|
||||
|
||||
virtual bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode);
|
||||
};
|
||||
using ConfAdapterPtr = std::shared_ptr<ConfAdapter>;
|
||||
|
||||
class IVFConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class IVFSQConfAdapter : public IVFConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class IVFPQConfAdapter : public IVFConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
static bool
|
||||
CheckPQParams(int64_t dimension, int64_t m, int64_t nbits, IndexMode& mode);
|
||||
|
||||
static bool
|
||||
CheckGPUPQParams(int64_t dimension, int64_t m, int64_t nbits);
|
||||
|
||||
static bool
|
||||
CheckCPUPQParams(int64_t dimension, int64_t m);
|
||||
};
|
||||
|
||||
class IVFHNSWConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class NSGConfAdapter : public IVFConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class BinIDMAPConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class BinIVFConfAdapter : public IVFConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class HNSWConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class ANNOYConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class RHNSWFlatConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class RHNSWPQConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class RHNSWSQConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class NGTPANNGConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
class NGTONNGConfAdapter : public ConfAdapter {
|
||||
public:
|
||||
bool
|
||||
CheckTrain(Config& oricfg, const IndexMode mode) override;
|
||||
|
||||
bool
|
||||
CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,62 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/ConfAdapterMgr.h"
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
ConfAdapterPtr
|
||||
AdapterMgr::GetAdapter(const IndexType type) {
|
||||
auto register_wrapper = [&, this]() { RegisterAdapter(); };
|
||||
std::call_once(once_, register_wrapper);
|
||||
|
||||
try {
|
||||
return collection_.at(type)();
|
||||
} catch (...) {
|
||||
KNOWHERE_THROW_MSG("Can not find confadapter: " + type);
|
||||
}
|
||||
}
|
||||
|
||||
#define REGISTER_CONF_ADAPTER(T, TYPE, NAME) static AdapterMgr::register_t<T> reg_##NAME##_(TYPE)
|
||||
|
||||
void
|
||||
AdapterMgr::RegisterAdapter() {
|
||||
REGISTER_CONF_ADAPTER(ConfAdapter, IndexEnum::INDEX_FAISS_IDMAP, idmap_adapter);
|
||||
REGISTER_CONF_ADAPTER(IVFConfAdapter, IndexEnum::INDEX_FAISS_IVFFLAT, ivf_adapter);
|
||||
REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexEnum::INDEX_FAISS_IVFPQ, ivfpq_adapter);
|
||||
REGISTER_CONF_ADAPTER(IVFSQConfAdapter, IndexEnum::INDEX_FAISS_IVFSQ8, ivfsq8_adapter);
|
||||
REGISTER_CONF_ADAPTER(IVFSQConfAdapter, IndexEnum::INDEX_FAISS_IVFSQ8H, ivfsq8h_adapter);
|
||||
REGISTER_CONF_ADAPTER(IVFHNSWConfAdapter, IndexEnum::INDEX_FAISS_IVFHNSW, ivfhnsw_adapter);
|
||||
REGISTER_CONF_ADAPTER(BinIDMAPConfAdapter, IndexEnum::INDEX_FAISS_BIN_IDMAP, idmap_bin_adapter);
|
||||
REGISTER_CONF_ADAPTER(BinIVFConfAdapter, IndexEnum::INDEX_FAISS_BIN_IVFFLAT, ivf_bin_adapter);
|
||||
REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexEnum::INDEX_NSG, nsg_adapter);
|
||||
#ifdef MILVUS_SUPPORT_SPTAG
|
||||
REGISTER_CONF_ADAPTER(ConfAdapter, IndexEnum::INDEX_SPTAG_KDT_RNT, sptag_kdt_adapter);
|
||||
REGISTER_CONF_ADAPTER(ConfAdapter, IndexEnum::INDEX_SPTAG_BKT_RNT, sptag_bkt_adapter);
|
||||
#endif
|
||||
REGISTER_CONF_ADAPTER(HNSWConfAdapter, IndexEnum::INDEX_HNSW, hnsw_adapter);
|
||||
REGISTER_CONF_ADAPTER(ANNOYConfAdapter, IndexEnum::INDEX_ANNOY, annoy_adapter);
|
||||
REGISTER_CONF_ADAPTER(RHNSWFlatConfAdapter, IndexEnum::INDEX_RHNSWFlat, rhnswflat_adapter);
|
||||
REGISTER_CONF_ADAPTER(RHNSWPQConfAdapter, IndexEnum::INDEX_RHNSWPQ, rhnswpq_adapter);
|
||||
REGISTER_CONF_ADAPTER(RHNSWSQConfAdapter, IndexEnum::INDEX_RHNSWSQ, rhnswsq_adapter);
|
||||
REGISTER_CONF_ADAPTER(NGTPANNGConfAdapter, IndexEnum::INDEX_NGTPANNG, ngtpanng_adapter);
|
||||
REGISTER_CONF_ADAPTER(NGTONNGConfAdapter, IndexEnum::INDEX_NGTONNG, ngtonng_adapter);
|
||||
|
||||
// though init_ won't be used later, it's better to set `init_` to true after registration was done.
|
||||
init_ = true;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,53 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <unordered_map>
|
||||
#include <mutex>
|
||||
|
||||
#include "knowhere/index/IndexType.h"
|
||||
#include "knowhere/index/vector_index/ConfAdapter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class AdapterMgr {
|
||||
public:
|
||||
template <typename T>
|
||||
struct register_t {
|
||||
explicit register_t(const IndexType type) {
|
||||
AdapterMgr::GetInstance().collection_[type] = ([] { return std::make_shared<T>(); });
|
||||
}
|
||||
};
|
||||
|
||||
static AdapterMgr&
|
||||
GetInstance() {
|
||||
static AdapterMgr instance;
|
||||
return instance;
|
||||
}
|
||||
|
||||
ConfAdapterPtr
|
||||
GetAdapter(const IndexType indexType);
|
||||
|
||||
void
|
||||
RegisterAdapter();
|
||||
|
||||
protected:
|
||||
bool init_ = false;
|
||||
std::unordered_map<IndexType, std::function<ConfAdapterPtr()>> collection_;
|
||||
std::once_flag once_;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,51 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include <faiss/index_io.h>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/FaissBaseBinaryIndex.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
BinarySet
|
||||
FaissBaseBinaryIndex::SerializeImpl(const IndexType& type) {
|
||||
try {
|
||||
faiss::IndexBinary* index = index_.get();
|
||||
|
||||
MemoryIOWriter writer;
|
||||
faiss::write_index_binary(index, &writer);
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("BinaryIVF", data, writer.rp);
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
FaissBaseBinaryIndex::LoadImpl(const BinarySet& index_binary, const IndexType& type) {
|
||||
auto binary = index_binary.GetByName("BinaryIVF");
|
||||
|
||||
MemoryIOReader reader;
|
||||
reader.total = binary->size;
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
faiss::IndexBinary* index = faiss::read_index_binary(&reader);
|
||||
index_.reset(index);
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,42 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include <faiss/IndexBinary.h>
|
||||
|
||||
#include "knowhere/common/BinarySet.h"
|
||||
#include "knowhere/common/Dataset.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class FaissBaseBinaryIndex {
|
||||
protected:
|
||||
explicit FaissBaseBinaryIndex(std::shared_ptr<faiss::IndexBinary> index) : index_(std::move(index)) {
|
||||
}
|
||||
|
||||
virtual BinarySet
|
||||
SerializeImpl(const IndexType& type);
|
||||
|
||||
virtual void
|
||||
LoadImpl(const BinarySet& index_binary, const IndexType& type);
|
||||
|
||||
public:
|
||||
std::shared_ptr<faiss::IndexBinary> index_ = nullptr;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,63 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <faiss/index_io.h>
|
||||
#include <fiu/fiu-local.h>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
#include "knowhere/index/vector_index/FaissBaseIndex.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
BinarySet
|
||||
FaissBaseIndex::SerializeImpl(const IndexType& type) {
|
||||
try {
|
||||
fiu_do_on("FaissBaseIndex.SerializeImpl.throw_exception", throw std::exception());
|
||||
faiss::Index* index = index_.get();
|
||||
|
||||
MemoryIOWriter writer;
|
||||
faiss::write_index(index, &writer);
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
BinarySet res_set;
|
||||
// TODO(linxj): use virtual func Name() instead of raw string.
|
||||
res_set.Append("IVF", data, writer.rp);
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
FaissBaseIndex::LoadImpl(const BinarySet& binary_set, const IndexType& type) {
|
||||
auto binary = binary_set.GetByName("IVF");
|
||||
|
||||
MemoryIOReader reader;
|
||||
reader.total = binary->size;
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
faiss::Index* index = faiss::read_index(&reader);
|
||||
index_.reset(index);
|
||||
|
||||
SealImpl();
|
||||
}
|
||||
|
||||
void
|
||||
FaissBaseIndex::SealImpl() {
|
||||
}
|
||||
|
||||
// FaissBaseIndex::~FaissBaseIndex() {}
|
||||
//
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,44 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include <faiss/Index.h>
|
||||
|
||||
#include "knowhere/common/BinarySet.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class FaissBaseIndex {
|
||||
protected:
|
||||
explicit FaissBaseIndex(std::shared_ptr<faiss::Index> index) : index_(std::move(index)) {
|
||||
}
|
||||
|
||||
virtual BinarySet
|
||||
SerializeImpl(const IndexType& type);
|
||||
|
||||
virtual void
|
||||
LoadImpl(const BinarySet&, const IndexType& type);
|
||||
|
||||
virtual void
|
||||
SealImpl();
|
||||
|
||||
public:
|
||||
std::shared_ptr<faiss::Index> index_ = nullptr;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,178 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexAnnoy.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cassert>
|
||||
#include <iterator>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
BinarySet
|
||||
IndexAnnoy::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
auto metric_type_length = metric_type_.length();
|
||||
std::shared_ptr<uint8_t[]> metric_type(new uint8_t[metric_type_length]);
|
||||
memcpy(metric_type.get(), metric_type_.data(), metric_type_.length());
|
||||
|
||||
auto dim = Dim();
|
||||
std::shared_ptr<uint8_t[]> dim_data(new uint8_t[sizeof(uint64_t)]);
|
||||
memcpy(dim_data.get(), &dim, sizeof(uint64_t));
|
||||
|
||||
size_t index_length = index_->get_index_length();
|
||||
std::shared_ptr<uint8_t[]> index_data(new uint8_t[index_length]);
|
||||
memcpy(index_data.get(), index_->get_index(), index_length);
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("annoy_metric_type", metric_type, metric_type_length);
|
||||
res_set.Append("annoy_dim", dim_data, sizeof(uint64_t));
|
||||
res_set.Append("annoy_index_data", index_data, index_length);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, res_set);
|
||||
}
|
||||
return res_set;
|
||||
}
|
||||
|
||||
void
|
||||
IndexAnnoy::Load(const BinarySet& index_binary) {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
auto metric_type = index_binary.GetByName("annoy_metric_type");
|
||||
metric_type_.resize(static_cast<size_t>(metric_type->size));
|
||||
memcpy(metric_type_.data(), metric_type->data.get(), static_cast<size_t>(metric_type->size));
|
||||
|
||||
auto dim_data = index_binary.GetByName("annoy_dim");
|
||||
uint64_t dim;
|
||||
memcpy(&dim, dim_data->data.get(), static_cast<size_t>(dim_data->size));
|
||||
|
||||
if (metric_type_ == Metric::L2) {
|
||||
index_ = std::make_shared<AnnoyIndex<int64_t, float, ::Euclidean, ::Kiss64Random>>(dim);
|
||||
} else if (metric_type_ == Metric::IP) {
|
||||
index_ = std::make_shared<AnnoyIndex<int64_t, float, ::DotProduct, ::Kiss64Random>>(dim);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("metric not supported " + metric_type_);
|
||||
}
|
||||
|
||||
auto index_data = index_binary.GetByName("annoy_index_data");
|
||||
char* p = nullptr;
|
||||
if (!index_->load_index(reinterpret_cast<void*>(index_data->data.get()), index_data->size, &p)) {
|
||||
std::string error_msg(p);
|
||||
free(p);
|
||||
KNOWHERE_THROW_MSG(error_msg);
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexAnnoy::BuildAll(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (index_) {
|
||||
// it is builded all
|
||||
LOG_KNOWHERE_DEBUG_ << "IndexAnnoy::BuildAll: index_ has been built!";
|
||||
return;
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
|
||||
metric_type_ = config[Metric::TYPE];
|
||||
if (metric_type_ == Metric::L2) {
|
||||
index_ = std::make_shared<AnnoyIndex<int64_t, float, ::Euclidean, ::Kiss64Random>>(dim);
|
||||
} else if (metric_type_ == Metric::IP) {
|
||||
index_ = std::make_shared<AnnoyIndex<int64_t, float, ::DotProduct, ::Kiss64Random>>(dim);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("metric not supported " + metric_type_);
|
||||
}
|
||||
|
||||
for (int i = 0; i < rows; ++i) {
|
||||
index_->add_item(i, static_cast<const float*>(p_data) + dim * i);
|
||||
}
|
||||
|
||||
index_->build(config[IndexParams::n_trees].get<int64_t>());
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
IndexAnnoy::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
auto k = config[meta::TOPK].get<int64_t>();
|
||||
auto search_k = config[IndexParams::search_k].get<int64_t>();
|
||||
auto all_num = rows * k;
|
||||
auto p_id = static_cast<int64_t*>(malloc(all_num * sizeof(int64_t)));
|
||||
auto p_dist = static_cast<float*>(malloc(all_num * sizeof(float)));
|
||||
|
||||
#pragma omp parallel for
|
||||
for (unsigned int i = 0; i < rows; ++i) {
|
||||
std::vector<int64_t> result;
|
||||
result.reserve(k);
|
||||
std::vector<float> distances;
|
||||
distances.reserve(k);
|
||||
index_->get_nns_by_vector(static_cast<const float*>(p_data) + i * dim, k, search_k, &result, &distances,
|
||||
bitset);
|
||||
|
||||
size_t result_num = result.size();
|
||||
auto local_p_id = p_id + k * i;
|
||||
auto local_p_dist = p_dist + k * i;
|
||||
memcpy(local_p_id, result.data(), result_num * sizeof(int64_t));
|
||||
memcpy(local_p_dist, distances.data(), result_num * sizeof(float));
|
||||
|
||||
MapOffsetToUid(local_p_id, result_num);
|
||||
|
||||
for (; result_num < k; result_num++) {
|
||||
local_p_id[result_num] = -1;
|
||||
local_p_dist[result_num] = 1.0 / 0.0;
|
||||
}
|
||||
}
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
return ret_ds;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IndexAnnoy::Count() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->get_n_items();
|
||||
}
|
||||
|
||||
int64_t
|
||||
IndexAnnoy::Dim() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->get_dim();
|
||||
}
|
||||
|
||||
void
|
||||
IndexAnnoy::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
index_size_ = index_->cal_size();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,68 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "annoy/src/annoylib.h"
|
||||
#include "annoy/src/kissrandom.h"
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexAnnoy : public VecIndex {
|
||||
public:
|
||||
IndexAnnoy() {
|
||||
index_type_ = IndexEnum::INDEX_ANNOY;
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
BuildAll(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override {
|
||||
KNOWHERE_THROW_MSG("Annoy not support build item dynamically, please invoke BuildAll interface.");
|
||||
}
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override {
|
||||
KNOWHERE_THROW_MSG("Incremental index is not supported");
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) override;
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
private:
|
||||
MetricType metric_type_;
|
||||
std::shared_ptr<AnnoyIndexInterface<int64_t, float>> index_ = nullptr;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,177 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexBinaryIDMAP.h"
|
||||
|
||||
#include <faiss/IndexBinaryFlat.h>
|
||||
#include <faiss/MetaIndexes.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
BinarySet
|
||||
BinaryIDMAP::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
|
||||
auto ret = SerializeImpl(index_type_);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, ret);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIDMAP::Load(const BinarySet& index_binary) {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
LoadImpl(index_binary, index_type_);
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
BinaryIDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
|
||||
auto k = config[meta::TOPK].get<int64_t>();
|
||||
auto elems = rows * k;
|
||||
size_t p_id_size = sizeof(int64_t) * elems;
|
||||
size_t p_dist_size = sizeof(float) * elems;
|
||||
auto p_id = static_cast<int64_t*>(malloc(p_id_size));
|
||||
auto p_dist = static_cast<float*>(malloc(p_dist_size));
|
||||
|
||||
QueryImpl(rows, reinterpret_cast<const uint8_t*>(p_data), k, p_dist, p_id, config, bitset);
|
||||
MapOffsetToUid(p_id, static_cast<size_t>(elems));
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
|
||||
return ret_ds;
|
||||
}
|
||||
|
||||
DynamicResultSegment
|
||||
BinaryIDMAP::QueryByDistance(const milvus::knowhere::DatasetPtr& dataset,
|
||||
const milvus::knowhere::Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
GET_TENSOR_DATA(dataset)
|
||||
if (rows != 1) {
|
||||
KNOWHERE_THROW_MSG("QueryByDistance only accept nq = 1!");
|
||||
}
|
||||
|
||||
auto default_type = index_->metric_type;
|
||||
if (config.contains(Metric::TYPE)) {
|
||||
index_->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
}
|
||||
std::vector<faiss::RangeSearchPartialResult*> res;
|
||||
DynamicResultSegment result;
|
||||
float radius = 0.0;
|
||||
if (index_->metric_type != faiss::MetricType::METRIC_Substructure &&
|
||||
index_->metric_type != faiss::METRIC_Superstructure) {
|
||||
radius = config[IndexParams::range_search_radius].get<float>();
|
||||
}
|
||||
auto buffer_size = config.contains(IndexParams::range_search_buffer_size)
|
||||
? config[IndexParams::range_search_buffer_size].get<size_t>()
|
||||
: 16384;
|
||||
auto real_idx = dynamic_cast<faiss::IndexBinaryFlat*>(index_.get());
|
||||
if (real_idx == nullptr) {
|
||||
KNOWHERE_THROW_MSG("Cannot dynamic_cast the index to faiss::IndexBinaryFlat type!");
|
||||
}
|
||||
real_idx->range_search(rows, reinterpret_cast<const uint8_t*>(p_data), radius, res, buffer_size, bitset);
|
||||
ExchangeDataset(result, res);
|
||||
MapUids(result);
|
||||
index_->metric_type = default_type;
|
||||
return result;
|
||||
}
|
||||
|
||||
int64_t
|
||||
BinaryIDMAP::Count() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->ntotal;
|
||||
}
|
||||
|
||||
int64_t
|
||||
BinaryIDMAP::Dim() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->d;
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIDMAP::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
|
||||
index_->add(rows, reinterpret_cast<const uint8_t*>(p_data));
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIDMAP::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
// users will assign the metric type when querying
|
||||
// so we let Tanimoto be the default type
|
||||
constexpr faiss::MetricType metric_type = faiss::METRIC_Tanimoto;
|
||||
|
||||
auto dim = config[meta::DIM].get<int64_t>();
|
||||
auto index = std::make_shared<faiss::IndexBinaryFlat>(dim, metric_type);
|
||||
index_ = index;
|
||||
}
|
||||
|
||||
const uint8_t*
|
||||
BinaryIDMAP::GetRawVectors() {
|
||||
try {
|
||||
auto flat_index = dynamic_cast<faiss::IndexBinaryFlat*>(index_.get());
|
||||
return flat_index->xb.data();
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIDMAP::QueryImpl(int64_t n,
|
||||
const uint8_t* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
// assign the metric type
|
||||
index_->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
|
||||
auto i_distances = reinterpret_cast<int32_t*>(distances);
|
||||
index_->search(n, data, k, i_distances, labels, bitset);
|
||||
|
||||
// if hamming, it need transform int32 to float
|
||||
if (index_->metric_type == faiss::METRIC_Hamming) {
|
||||
int64_t num = n * k;
|
||||
for (int64_t i = 0; i < num; i++) {
|
||||
distances[i] = static_cast<float>(i_distances[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,81 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/index/vector_index/FaissBaseBinaryIndex.h"
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
#include "knowhere/index/vector_index/helpers/DynamicResultSet.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class BinaryIDMAP : public VecIndex, public FaissBaseBinaryIndex {
|
||||
public:
|
||||
BinaryIDMAP() : FaissBaseBinaryIndex(nullptr) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_BIN_IDMAP;
|
||||
}
|
||||
|
||||
explicit BinaryIDMAP(std::shared_ptr<faiss::IndexBinary> index) : FaissBaseBinaryIndex(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_BIN_IDMAP;
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config&) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet&) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override;
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr&, const Config&, const faiss::BitsetView bitset) override;
|
||||
|
||||
DynamicResultSegment
|
||||
QueryByDistance(const DatasetPtr& dataset, const Config& config, const faiss::BitsetView bitset);
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
int64_t
|
||||
IndexSize() override {
|
||||
return Count() * Dim() / 8;
|
||||
}
|
||||
|
||||
virtual const uint8_t*
|
||||
GetRawVectors();
|
||||
|
||||
protected:
|
||||
virtual void
|
||||
QueryImpl(int64_t n,
|
||||
const uint8_t* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset);
|
||||
};
|
||||
|
||||
using BinaryIDMAPPtr = std::shared_ptr<BinaryIDMAP>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,234 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexBinaryIVF.h"
|
||||
|
||||
#include <faiss/IndexBinaryFlat.h>
|
||||
#include <faiss/IndexBinaryIVF.h>
|
||||
|
||||
#include <chrono>
|
||||
#include <string>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
using stdclock = std::chrono::high_resolution_clock;
|
||||
|
||||
BinarySet
|
||||
BinaryIVF::Serialize(const Config& config) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
auto ret = SerializeImpl(index_type_);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, ret);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIVF::Load(const BinarySet& index_binary) {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
LoadImpl(index_binary, index_type_);
|
||||
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
auto ivf_index = static_cast<faiss::IndexBinaryIVF*>(index_.get());
|
||||
ivf_index->nprobe_statistics.resize(ivf_index->nlist, 0);
|
||||
}
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
BinaryIVF::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
|
||||
int64_t* p_id = nullptr;
|
||||
float* p_dist = nullptr;
|
||||
auto release_when_exception = [&]() {
|
||||
if (p_id != nullptr) {
|
||||
free(p_id);
|
||||
}
|
||||
if (p_dist != nullptr) {
|
||||
free(p_dist);
|
||||
}
|
||||
};
|
||||
|
||||
try {
|
||||
auto k = config[meta::TOPK].get<int64_t>();
|
||||
auto elems = rows * k;
|
||||
|
||||
size_t p_id_size = sizeof(int64_t) * elems;
|
||||
size_t p_dist_size = sizeof(float) * elems;
|
||||
p_id = static_cast<int64_t*>(malloc(p_id_size));
|
||||
p_dist = static_cast<float*>(malloc(p_dist_size));
|
||||
|
||||
QueryImpl(rows, reinterpret_cast<const uint8_t*>(p_data), k, p_dist, p_id, config, bitset);
|
||||
MapOffsetToUid(p_id, static_cast<size_t>(elems));
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
|
||||
return ret_ds;
|
||||
} catch (faiss::FaissException& e) {
|
||||
release_when_exception();
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
} catch (std::exception& e) {
|
||||
release_when_exception();
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
int64_t
|
||||
BinaryIVF::Count() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->ntotal;
|
||||
}
|
||||
|
||||
int64_t
|
||||
BinaryIVF::Dim() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->d;
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIVF::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
auto bin_ivf_index = dynamic_cast<faiss::IndexBinaryIVF*>(index_.get());
|
||||
auto nb = bin_ivf_index->invlists->compute_ntotal();
|
||||
auto nlist = bin_ivf_index->nlist;
|
||||
auto code_size = bin_ivf_index->code_size;
|
||||
|
||||
// binary ivf codes, ids and quantizer
|
||||
index_size_ = nb * code_size + nb * sizeof(int64_t) + nlist * code_size;
|
||||
}
|
||||
|
||||
StatisticsPtr
|
||||
BinaryIVF::GetStatistics() {
|
||||
if (!STATISTICS_LEVEL) {
|
||||
return stats;
|
||||
}
|
||||
auto ivf_stats = std::dynamic_pointer_cast<IVFStatistics>(stats);
|
||||
auto ivf_index = dynamic_cast<faiss::IndexBinaryIVF*>(index_.get());
|
||||
auto lock = ivf_stats->Lock();
|
||||
ivf_stats->update_ivf_access_stats(ivf_index->nprobe_statistics);
|
||||
return ivf_stats;
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIVF::ClearStatistics() {
|
||||
if (!STATISTICS_LEVEL) {
|
||||
return;
|
||||
}
|
||||
auto ivf_stats = std::dynamic_pointer_cast<IVFStatistics>(stats);
|
||||
auto ivf_index = dynamic_cast<faiss::IndexBinaryIVF*>(index_.get());
|
||||
ivf_index->clear_nprobe_statistics();
|
||||
ivf_index->index_ivf_stats.reset();
|
||||
auto lock = ivf_stats->Lock();
|
||||
ivf_stats->clear();
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIVF::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
|
||||
int64_t nlist = config[IndexParams::nlist];
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
faiss::IndexBinary* coarse_quantizer = new faiss::IndexBinaryFlat(dim, metric_type);
|
||||
auto index = std::make_shared<faiss::IndexBinaryIVF>(coarse_quantizer, dim, nlist, metric_type);
|
||||
index->own_fields = true;
|
||||
index->train(rows, static_cast<const uint8_t*>(p_data));
|
||||
index_ = index;
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIVF::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
index_->add(rows, reinterpret_cast<const uint8_t*>(p_data));
|
||||
}
|
||||
|
||||
std::shared_ptr<faiss::IVFSearchParameters>
|
||||
BinaryIVF::GenParams(const Config& config) {
|
||||
auto params = std::make_shared<faiss::IVFSearchParameters>();
|
||||
params->nprobe = config[IndexParams::nprobe];
|
||||
// params->max_codes = config["max_code"];
|
||||
return params;
|
||||
}
|
||||
|
||||
void
|
||||
BinaryIVF::QueryImpl(int64_t n,
|
||||
const uint8_t* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
auto params = GenParams(config);
|
||||
auto ivf_index = dynamic_cast<faiss::IndexBinaryIVF*>(index_.get());
|
||||
ivf_index->nprobe = params->nprobe;
|
||||
|
||||
stdclock::time_point before = stdclock::now();
|
||||
auto i_distances = reinterpret_cast<int32_t*>(distances);
|
||||
|
||||
index_->search(n, data, k, i_distances, labels, bitset);
|
||||
|
||||
stdclock::time_point after = stdclock::now();
|
||||
double search_cost = (std::chrono::duration<double, std::micro>(after - before)).count();
|
||||
LOG_KNOWHERE_DEBUG_ << "IVF_NM search cost: " << search_cost
|
||||
<< ", quantization cost: " << ivf_index->index_ivf_stats.quantization_time
|
||||
<< ", data search cost: " << ivf_index->index_ivf_stats.search_time;
|
||||
|
||||
if (STATISTICS_LEVEL) {
|
||||
auto ivf_stats = std::dynamic_pointer_cast<IVFStatistics>(stats);
|
||||
auto lock = ivf_stats->Lock();
|
||||
if (STATISTICS_LEVEL >= 1) {
|
||||
ivf_stats->update_nq(n);
|
||||
ivf_stats->count_nprobe(ivf_index->nprobe);
|
||||
ivf_stats->update_total_query_time(ivf_index->index_ivf_stats.quantization_time +
|
||||
ivf_index->index_ivf_stats.search_time);
|
||||
ivf_index->index_ivf_stats.quantization_time = 0;
|
||||
ivf_index->index_ivf_stats.search_time = 0;
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 2) {
|
||||
ivf_stats->update_filter_percentage(bitset);
|
||||
}
|
||||
}
|
||||
|
||||
// if hamming, it need transform int32 to float
|
||||
if (ivf_index->metric_type == faiss::METRIC_Hamming) {
|
||||
int64_t num = n * k;
|
||||
for (int64_t i = 0; i < num; i++) {
|
||||
distances[i] = static_cast<float>(i_distances[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,86 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include <faiss/IndexIVF.h>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/FaissBaseBinaryIndex.h"
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class BinaryIVF : public VecIndex, public FaissBaseBinaryIndex {
|
||||
public:
|
||||
BinaryIVF() : FaissBaseBinaryIndex(nullptr) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_BIN_IVFFLAT;
|
||||
stats = std::make_shared<milvus::knowhere::IVFStatistics>(index_type_);
|
||||
}
|
||||
|
||||
explicit BinaryIVF(std::shared_ptr<faiss::IndexBinary> index) : FaissBaseBinaryIndex(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_BIN_IVFFLAT;
|
||||
stats = std::make_shared<milvus::knowhere::IVFStatistics>(index_type_);
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override;
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) override;
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
StatisticsPtr
|
||||
GetStatistics() override;
|
||||
|
||||
void
|
||||
ClearStatistics() override;
|
||||
|
||||
protected:
|
||||
virtual std::shared_ptr<faiss::IVFSearchParameters>
|
||||
GenParams(const Config& config);
|
||||
|
||||
virtual void
|
||||
QueryImpl(int64_t n,
|
||||
const uint8_t* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset);
|
||||
};
|
||||
|
||||
using BinaryIVFIndexPtr = std::shared_ptr<BinaryIVF>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,259 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexHNSW.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cassert>
|
||||
#include <chrono>
|
||||
#include <iterator>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "hnswlib/hnswalg.h"
|
||||
#include "hnswlib/hnswlib.h"
|
||||
#include "hnswlib/space_ip.h"
|
||||
#include "hnswlib/space_l2.h"
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
// void
|
||||
// normalize_vector(float* data, float* norm_array, size_t dim) {
|
||||
// float norm = 0.0f;
|
||||
// for (int i = 0; i < dim; i++) norm += data[i] * data[i];
|
||||
// norm = 1.0f / (sqrtf(norm) + 1e-30f);
|
||||
// for (int i = 0; i < dim; i++) norm_array[i] = data[i] * norm;
|
||||
// }
|
||||
|
||||
BinarySet
|
||||
IndexHNSW::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
try {
|
||||
MemoryIOWriter writer;
|
||||
index_->saveIndex(writer);
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("HNSW", data, writer.rp);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, res_set);
|
||||
}
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexHNSW::Load(const BinarySet& index_binary) {
|
||||
try {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
auto binary = index_binary.GetByName("HNSW");
|
||||
|
||||
MemoryIOReader reader;
|
||||
reader.total = binary->size;
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
hnswlib::SpaceInterface<float>* space = nullptr;
|
||||
index_ = std::make_shared<hnswlib::HierarchicalNSW<float>>(space);
|
||||
index_->stats_enable = (STATISTICS_LEVEL >= 3);
|
||||
index_->loadIndex(reader);
|
||||
auto hnsw_stats = std::static_pointer_cast<LibHNSWStatistics>(stats);
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
auto lock = hnsw_stats->Lock();
|
||||
hnsw_stats->update_level_distribution(index_->maxlevel_, index_->level_stats_);
|
||||
}
|
||||
// LOG_KNOWHERE_DEBUG_ << "IndexHNSW::Load finished, show statistics:";
|
||||
// LOG_KNOWHERE_DEBUG_ << hnsw_stats->ToString();
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexHNSW::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
try {
|
||||
auto dim = dataset_ptr->Get<int64_t>(meta::DIM);
|
||||
auto rows = dataset_ptr->Get<int64_t>(meta::ROWS);
|
||||
|
||||
hnswlib::SpaceInterface<float>* space;
|
||||
std::string metric_type = config[Metric::TYPE];
|
||||
if (metric_type == Metric::L2) {
|
||||
space = new hnswlib::L2Space(dim);
|
||||
} else if (metric_type == Metric::IP) {
|
||||
space = new hnswlib::InnerProductSpace(dim);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Metric type not supported: " + metric_type);
|
||||
}
|
||||
index_ = std::make_shared<hnswlib::HierarchicalNSW<float>>(space, rows, config[IndexParams::M].get<int64_t>(),
|
||||
config[IndexParams::efConstruction].get<int64_t>());
|
||||
index_->stats_enable = (STATISTICS_LEVEL >= 3);
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexHNSW::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
|
||||
index_->addPoint(p_data, 0);
|
||||
#pragma omp parallel for
|
||||
for (int i = 1; i < rows; ++i) {
|
||||
index_->addPoint((reinterpret_cast<const float*>(p_data) + Dim() * i), i);
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
auto hnsw_stats = std::static_pointer_cast<LibHNSWStatistics>(stats);
|
||||
auto lock = hnsw_stats->Lock();
|
||||
hnsw_stats->update_level_distribution(index_->maxlevel_, index_->level_stats_);
|
||||
}
|
||||
// LOG_KNOWHERE_DEBUG_ << "IndexHNSW::Train finished, show statistics:";
|
||||
// LOG_KNOWHERE_DEBUG_ << GetStatistics()->ToString();
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
IndexHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
|
||||
size_t k = config[meta::TOPK].get<int64_t>();
|
||||
size_t id_size = sizeof(int64_t) * k;
|
||||
size_t dist_size = sizeof(float) * k;
|
||||
auto p_id = static_cast<int64_t*>(malloc(id_size * rows));
|
||||
auto p_dist = static_cast<float*>(malloc(dist_size * rows));
|
||||
std::vector<hnswlib::StatisticsInfo> query_stats;
|
||||
auto hnsw_stats = std::dynamic_pointer_cast<LibHNSWStatistics>(stats);
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
query_stats.resize(rows);
|
||||
for (auto i = 0; i < rows; ++i) {
|
||||
query_stats[i].target_level = hnsw_stats->target_level;
|
||||
}
|
||||
}
|
||||
|
||||
index_->setEf(config[IndexParams::ef].get<int64_t>());
|
||||
bool transform = (index_->metric_type_ == 1); // InnerProduct: 1
|
||||
|
||||
std::chrono::high_resolution_clock::time_point query_start, query_end;
|
||||
query_start = std::chrono::high_resolution_clock::now();
|
||||
|
||||
#pragma omp parallel for
|
||||
for (unsigned int i = 0; i < rows; ++i) {
|
||||
auto single_query = (float*)p_data + i * dim;
|
||||
std::priority_queue<std::pair<float, hnswlib::labeltype>> rst;
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
rst = index_->searchKnn(single_query, k, bitset, query_stats[i]);
|
||||
} else {
|
||||
auto dummy_stat = hnswlib::StatisticsInfo();
|
||||
rst = index_->searchKnn(single_query, k, bitset, dummy_stat);
|
||||
}
|
||||
size_t rst_size = rst.size();
|
||||
|
||||
auto p_single_dis = p_dist + i * k;
|
||||
auto p_single_id = p_id + i * k;
|
||||
size_t idx = rst_size - 1;
|
||||
while (!rst.empty()) {
|
||||
auto& it = rst.top();
|
||||
p_single_dis[idx] = transform ? (1 - it.first) : it.first;
|
||||
p_single_id[idx] = it.second;
|
||||
rst.pop();
|
||||
idx--;
|
||||
}
|
||||
MapOffsetToUid(p_single_id, rst_size);
|
||||
|
||||
for (idx = rst_size; idx < k; idx++) {
|
||||
p_single_dis[idx] = float(1.0 / 0.0);
|
||||
p_single_id[idx] = -1;
|
||||
}
|
||||
}
|
||||
query_end = std::chrono::high_resolution_clock::now();
|
||||
|
||||
if (STATISTICS_LEVEL) {
|
||||
auto lock = hnsw_stats->Lock();
|
||||
if (STATISTICS_LEVEL >= 1) {
|
||||
hnsw_stats->update_nq(rows);
|
||||
hnsw_stats->update_ef_sum(index_->ef_ * rows);
|
||||
hnsw_stats->update_total_query_time(
|
||||
std::chrono::duration_cast<std::chrono::milliseconds>(query_end - query_start).count());
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 2) {
|
||||
hnsw_stats->update_filter_percentage(bitset);
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
for (auto i = 0; i < rows; ++i) {
|
||||
for (auto j = 0; j < query_stats[i].accessed_points.size(); ++j) {
|
||||
auto tgt = hnsw_stats->access_cnt_map.find(query_stats[i].accessed_points[j]);
|
||||
if (tgt == hnsw_stats->access_cnt_map.end())
|
||||
hnsw_stats->access_cnt_map[query_stats[i].accessed_points[j]] = 1;
|
||||
else
|
||||
tgt->second += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// LOG_KNOWHERE_DEBUG_ << "IndexHNSW::Query finished, show statistics:";
|
||||
// LOG_KNOWHERE_DEBUG_ << GetStatistics()->ToString();
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
return ret_ds;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IndexHNSW::Count() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->cur_element_count;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IndexHNSW::Dim() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return (*static_cast<size_t*>(index_->dist_func_param_));
|
||||
}
|
||||
|
||||
void
|
||||
IndexHNSW::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
index_size_ = index_->cal_size();
|
||||
}
|
||||
|
||||
void
|
||||
IndexHNSW::ClearStatistics() {
|
||||
if (!STATISTICS_LEVEL)
|
||||
return;
|
||||
auto hnsw_stats = std::static_pointer_cast<LibHNSWStatistics>(stats);
|
||||
auto lock = hnsw_stats->Lock();
|
||||
hnsw_stats->clear();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,63 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "hnswlib/hnswlib.h"
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexHNSW : public VecIndex {
|
||||
public:
|
||||
IndexHNSW() {
|
||||
index_type_ = IndexEnum::INDEX_HNSW;
|
||||
stats = std::make_shared<milvus::knowhere::LibHNSWStatistics>(index_type_);
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override;
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) override;
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
void
|
||||
ClearStatistics() override;
|
||||
|
||||
private:
|
||||
std::shared_ptr<hnswlib::HierarchicalNSW<float>> index_;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,195 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include "knowhere/index/vector_index/IndexIDMAP.h"
|
||||
|
||||
#include <faiss/AutoTune.h>
|
||||
#include <faiss/IndexFlat.h>
|
||||
#include <faiss/MetaIndexes.h>
|
||||
#include <faiss/clone_index.h>
|
||||
#include <faiss/index_io.h>
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#endif
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
|
||||
#endif
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
BinarySet
|
||||
IDMAP::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
|
||||
auto ret = SerializeImpl(index_type_);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, ret);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
void
|
||||
IDMAP::Load(const BinarySet& binary_set) {
|
||||
Assemble(const_cast<BinarySet&>(binary_set));
|
||||
LoadImpl(binary_set, index_type_);
|
||||
}
|
||||
|
||||
void
|
||||
IDMAP::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
// users will assign the metric type when querying
|
||||
// so we let L2 be the default type
|
||||
constexpr faiss::MetricType metric_type = faiss::METRIC_L2;
|
||||
|
||||
auto dim = config[meta::DIM].get<int64_t>();
|
||||
auto index = std::make_shared<faiss::IndexFlat>(dim, metric_type);
|
||||
index_ = index;
|
||||
}
|
||||
|
||||
void
|
||||
IDMAP::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
index_->add(rows, reinterpret_cast<const float*>(p_data));
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
IDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
|
||||
auto k = config[meta::TOPK].get<int64_t>();
|
||||
auto elems = rows * k;
|
||||
size_t p_id_size = sizeof(int64_t) * elems;
|
||||
size_t p_dist_size = sizeof(float) * elems;
|
||||
auto p_id = static_cast<int64_t*>(malloc(p_id_size));
|
||||
auto p_dist = static_cast<float*>(malloc(p_dist_size));
|
||||
|
||||
QueryImpl(rows, reinterpret_cast<const float*>(p_data), k, p_dist, p_id, config, bitset);
|
||||
MapOffsetToUid(p_id, static_cast<size_t>(elems));
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
return ret_ds;
|
||||
}
|
||||
|
||||
DynamicResultSegment
|
||||
IDMAP::QueryByDistance(const milvus::knowhere::DatasetPtr& dataset,
|
||||
const milvus::knowhere::Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
GET_TENSOR_DATA(dataset)
|
||||
if (rows != 1) {
|
||||
KNOWHERE_THROW_MSG("QueryByDistance only accept nq = 1!");
|
||||
}
|
||||
|
||||
auto default_type = index_->metric_type;
|
||||
if (config.contains(Metric::TYPE)) {
|
||||
index_->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
}
|
||||
std::vector<faiss::RangeSearchPartialResult*> res;
|
||||
DynamicResultSegment result;
|
||||
auto radius = config[IndexParams::range_search_radius].get<float>();
|
||||
auto buffer_size = config.contains(IndexParams::range_search_buffer_size)
|
||||
? config[IndexParams::range_search_buffer_size].get<size_t>()
|
||||
: 16384;
|
||||
auto real_idx = dynamic_cast<faiss::IndexFlat*>(index_.get());
|
||||
if (real_idx == nullptr) {
|
||||
KNOWHERE_THROW_MSG("Cannot dynamic_cast the index to faiss::IndexFlat type!");
|
||||
}
|
||||
if (index_->metric_type == faiss::MetricType::METRIC_L2) {
|
||||
radius *= radius;
|
||||
}
|
||||
real_idx->range_search(rows, reinterpret_cast<const float*>(p_data), radius, res, buffer_size, bitset);
|
||||
ExchangeDataset(result, res);
|
||||
MapUids(result);
|
||||
index_->metric_type = default_type;
|
||||
return result;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IDMAP::Count() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->ntotal;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IDMAP::Dim() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->d;
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
IDMAP::CopyCpuToGpu(const int64_t device_id, const Config& config) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
|
||||
ResScope rs(res, device_id, false);
|
||||
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get());
|
||||
|
||||
std::shared_ptr<faiss::Index> device_index;
|
||||
device_index.reset(gpu_index);
|
||||
return std::make_shared<GPUIDMAP>(device_index, device_id, res);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
|
||||
}
|
||||
#else
|
||||
KNOWHERE_THROW_MSG("Calling IDMAP::CopyCpuToGpu when we are using CPU version");
|
||||
#endif
|
||||
}
|
||||
|
||||
const float*
|
||||
IDMAP::GetRawVectors() {
|
||||
try {
|
||||
auto flat_index = dynamic_cast<faiss::IndexFlat*>(index_.get());
|
||||
return flat_index->xb.data();
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IDMAP::QueryImpl(int64_t n,
|
||||
const float* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
// assign the metric type
|
||||
index_->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
index_->search(n, data, k, distances, labels, bitset);
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,77 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/index/vector_index/FaissBaseIndex.h"
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
#include "knowhere/index/vector_index/helpers/DynamicResultSet.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IDMAP : public VecIndex, public FaissBaseIndex {
|
||||
public:
|
||||
IDMAP() : FaissBaseIndex(nullptr) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IDMAP;
|
||||
}
|
||||
|
||||
explicit IDMAP(std::shared_ptr<faiss::Index> index) : FaissBaseIndex(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IDMAP;
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config&) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet&) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override;
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr&, const Config&, const faiss::BitsetView) override;
|
||||
|
||||
DynamicResultSegment
|
||||
QueryByDistance(const DatasetPtr& dataset, const Config& config, const faiss::BitsetView bitset);
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
int64_t
|
||||
IndexSize() override {
|
||||
return Count() * Dim() * sizeof(FloatType);
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
CopyCpuToGpu(const int64_t, const Config&);
|
||||
|
||||
virtual const float*
|
||||
GetRawVectors();
|
||||
|
||||
protected:
|
||||
virtual void
|
||||
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::BitsetView);
|
||||
};
|
||||
|
||||
using IDMAPPtr = std::shared_ptr<IDMAP>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,394 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <faiss/AutoTune.h>
|
||||
#include <faiss/IVFlib.h>
|
||||
#include <faiss/IndexFlat.h>
|
||||
#include <faiss/IndexIVF.h>
|
||||
#include <faiss/IndexIVFFlat.h>
|
||||
#include <faiss/IndexIVFPQ.h>
|
||||
#include <faiss/clone_index.h>
|
||||
#include <faiss/index_io.h>
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include <faiss/gpu/GpuAutoTune.h>
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#endif
|
||||
|
||||
#include <fiu/fiu-local.h>
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/IndexIVF.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
|
||||
#endif
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
using stdclock = std::chrono::high_resolution_clock;
|
||||
|
||||
BinarySet
|
||||
IVF::Serialize(const Config& config) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
auto ret = SerializeImpl(index_type_);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, ret);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
void
|
||||
IVF::Load(const BinarySet& binary_set) {
|
||||
Assemble(const_cast<BinarySet&>(binary_set));
|
||||
LoadImpl(binary_set, index_type_);
|
||||
|
||||
if (IndexMode() == IndexMode::MODE_CPU && STATISTICS_LEVEL >= 3) {
|
||||
auto ivf_index = static_cast<faiss::IndexIVFFlat*>(index_.get());
|
||||
ivf_index->nprobe_statistics.resize(ivf_index->nlist, 0);
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IVF::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
|
||||
auto nlist = config[IndexParams::nlist].get<int64_t>();
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
faiss::Index* coarse_quantizer = new faiss::IndexFlat(dim, metric_type);
|
||||
auto index = std::make_shared<faiss::IndexIVFFlat>(coarse_quantizer, dim, nlist, metric_type);
|
||||
index->own_fields = true;
|
||||
index->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
index_ = index;
|
||||
}
|
||||
|
||||
void
|
||||
IVF::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
index_->add(rows, reinterpret_cast<const float*>(p_data));
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
IVF::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
|
||||
int64_t* p_id = nullptr;
|
||||
float* p_dist = nullptr;
|
||||
auto release_when_exception = [&]() {
|
||||
if (p_id != nullptr) {
|
||||
free(p_id);
|
||||
}
|
||||
if (p_dist != nullptr) {
|
||||
free(p_dist);
|
||||
}
|
||||
};
|
||||
|
||||
try {
|
||||
fiu_do_on("IVF.Search.throw_std_exception", throw std::exception());
|
||||
fiu_do_on("IVF.Search.throw_faiss_exception", throw faiss::FaissException(""));
|
||||
auto k = config[meta::TOPK].get<int64_t>();
|
||||
auto elems = rows * k;
|
||||
|
||||
size_t p_id_size = sizeof(int64_t) * elems;
|
||||
size_t p_dist_size = sizeof(float) * elems;
|
||||
p_id = static_cast<int64_t*>(malloc(p_id_size));
|
||||
p_dist = static_cast<float*>(malloc(p_dist_size));
|
||||
|
||||
QueryImpl(rows, reinterpret_cast<const float*>(p_data), k, p_dist, p_id, config, bitset);
|
||||
MapOffsetToUid(p_id, static_cast<size_t>(elems));
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
return ret_ds;
|
||||
} catch (faiss::FaissException& e) {
|
||||
release_when_exception();
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
} catch (std::exception& e) {
|
||||
release_when_exception();
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
#if 0
|
||||
DatasetPtr
|
||||
IVF::QueryById(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
auto rows = dataset_ptr->Get<int64_t>(meta::ROWS);
|
||||
auto p_data = dataset_ptr->Get<const int64_t*>(meta::IDS);
|
||||
|
||||
try {
|
||||
int64_t k = config[meta::TOPK].get<int64_t>();
|
||||
auto elems = rows * k;
|
||||
|
||||
size_t p_id_size = sizeof(int64_t) * elems;
|
||||
size_t p_dist_size = sizeof(float) * elems;
|
||||
auto p_id = (int64_t*)malloc(p_id_size);
|
||||
auto p_dist = (float*)malloc(p_dist_size);
|
||||
|
||||
// todo: enable search by id (zhiru)
|
||||
// auto blacklist = dataset_ptr->Get<faiss::ConcurrentBitsetPtr>("bitset");
|
||||
auto index_ivf = std::static_pointer_cast<faiss::IndexIVF>(index_);
|
||||
index_ivf->search_by_id(rows, p_data, k, p_dist, p_id, bitset_);
|
||||
|
||||
// std::stringstream ss_res_id, ss_res_dist;
|
||||
// for (int i = 0; i < 10; ++i) {
|
||||
// printf("%llu", res_ids[i]);
|
||||
// printf("\n");
|
||||
// printf("%.6f", res_dis[i]);
|
||||
// printf("\n");
|
||||
// ss_res_id << res_ids[i] << " ";
|
||||
// ss_res_dist << res_dis[i] << " ";
|
||||
// }
|
||||
// std::cout << std::endl << "after search: " << std::endl;
|
||||
// std::cout << ss_res_id.str() << std::endl;
|
||||
// std::cout << ss_res_dist.str() << std::endl << std::endl;
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
return ret_ds;
|
||||
} catch (faiss::FaissException& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
IVF::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
auto p_data = dataset_ptr->Get<const int64_t*>(meta::IDS);
|
||||
auto elems = dataset_ptr->Get<int64_t>(meta::DIM);
|
||||
|
||||
try {
|
||||
size_t p_x_size = sizeof(float) * elems;
|
||||
auto p_x = (float*)malloc(p_x_size);
|
||||
|
||||
auto index_ivf = std::static_pointer_cast<faiss::IndexIVF>(index_);
|
||||
index_ivf->get_vector_by_id(1, p_data, p_x, bitset_);
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::TENSOR, p_x);
|
||||
return ret_ds;
|
||||
} catch (faiss::FaissException& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
int64_t
|
||||
IVF::Count() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->ntotal;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IVF::Dim() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->d;
|
||||
}
|
||||
|
||||
void
|
||||
IVF::Seal() {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
SealImpl();
|
||||
}
|
||||
|
||||
void
|
||||
IVF::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
auto ivf_index = static_cast<faiss::IndexIVFFlat*>(index_.get());
|
||||
auto nb = ivf_index->invlists->compute_ntotal();
|
||||
auto nlist = ivf_index->nlist;
|
||||
auto code_size = ivf_index->code_size;
|
||||
// ivf codes, ivf ids and quantizer
|
||||
index_size_ = nb * code_size + nb * sizeof(int64_t) + nlist * code_size;
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
IVF::CopyCpuToGpu(const int64_t device_id, const Config& config) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
|
||||
ResScope rs(res, device_id, false);
|
||||
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get());
|
||||
|
||||
std::shared_ptr<faiss::Index> device_index;
|
||||
device_index.reset(gpu_index);
|
||||
return std::make_shared<GPUIVF>(device_index, device_id, res);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
|
||||
}
|
||||
|
||||
#else
|
||||
KNOWHERE_THROW_MSG("Calling IVF::CopyCpuToGpu when we are using CPU version");
|
||||
#endif
|
||||
}
|
||||
|
||||
void
|
||||
IVF::GenGraph(const float* data, const int64_t k, GraphType& graph, const Config& config) {
|
||||
int64_t K = k + 1;
|
||||
auto ntotal = Count();
|
||||
|
||||
size_t dim = config[meta::DIM];
|
||||
auto batch_size = 1000;
|
||||
auto tail_batch_size = ntotal % batch_size;
|
||||
auto batch_search_count = ntotal / batch_size;
|
||||
auto total_search_count = tail_batch_size == 0 ? batch_search_count : batch_search_count + 1;
|
||||
|
||||
std::vector<float> res_dis(K * batch_size);
|
||||
graph.resize(ntotal);
|
||||
GraphType res_vec(total_search_count);
|
||||
for (int i = 0; i < total_search_count; ++i) {
|
||||
auto b_size = (i == (total_search_count - 1)) && tail_batch_size != 0 ? tail_batch_size : batch_size;
|
||||
|
||||
auto& res = res_vec[i];
|
||||
res.resize(K * b_size);
|
||||
|
||||
const float* xq = data + batch_size * dim * i;
|
||||
QueryImpl(b_size, xq, K, res_dis.data(), res.data(), config, nullptr);
|
||||
|
||||
for (int j = 0; j < b_size; ++j) {
|
||||
auto& node = graph[batch_size * i + j];
|
||||
node.resize(k);
|
||||
auto start_pos = j * K + 1;
|
||||
for (int m = 0, cursor = start_pos; m < k && cursor < start_pos + k; ++m, ++cursor) {
|
||||
node[m] = res[cursor];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::shared_ptr<faiss::IVFSearchParameters>
|
||||
IVF::GenParams(const Config& config) {
|
||||
auto params = std::make_shared<faiss::IVFSearchParameters>();
|
||||
params->nprobe = config[IndexParams::nprobe];
|
||||
// params->max_codes = config["max_codes"];
|
||||
return params;
|
||||
}
|
||||
|
||||
void
|
||||
IVF::QueryImpl(int64_t n,
|
||||
const float* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
auto params = GenParams(config);
|
||||
auto ivf_index = dynamic_cast<faiss::IndexIVF*>(index_.get());
|
||||
ivf_index->nprobe = std::min(params->nprobe, ivf_index->invlists->nlist);
|
||||
stdclock::time_point before = stdclock::now();
|
||||
if (params->nprobe > 1 && n <= 4) {
|
||||
ivf_index->parallel_mode = 1;
|
||||
} else {
|
||||
ivf_index->parallel_mode = 0;
|
||||
}
|
||||
auto ivf_stats = std::dynamic_pointer_cast<IVFStatistics>(stats);
|
||||
ivf_index->search(n, data, k, distances, labels, bitset);
|
||||
stdclock::time_point after = stdclock::now();
|
||||
double search_cost = (std::chrono::duration<double, std::micro>(after - before)).count();
|
||||
if (STATISTICS_LEVEL) {
|
||||
auto lock = ivf_stats->Lock();
|
||||
if (STATISTICS_LEVEL >= 1) {
|
||||
ivf_stats->update_nq(n);
|
||||
ivf_stats->count_nprobe(ivf_index->nprobe);
|
||||
|
||||
LOG_KNOWHERE_DEBUG_ << "IVF search cost: " << search_cost
|
||||
<< ", quantization cost: " << ivf_index->index_ivf_stats.quantization_time
|
||||
<< ", data search cost: " << ivf_index->index_ivf_stats.search_time;
|
||||
ivf_stats->update_total_query_time(ivf_index->index_ivf_stats.quantization_time +
|
||||
ivf_index->index_ivf_stats.search_time);
|
||||
ivf_index->index_ivf_stats.quantization_time = 0;
|
||||
ivf_index->index_ivf_stats.search_time = 0;
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 2) {
|
||||
ivf_stats->update_filter_percentage(bitset);
|
||||
}
|
||||
}
|
||||
// LOG_KNOWHERE_DEBUG_ << "IndexIVF::QueryImpl finished, show statistics:";
|
||||
// LOG_KNOWHERE_DEBUG_ << GetStatistics()->ToString();
|
||||
}
|
||||
|
||||
void
|
||||
IVF::SealImpl() {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
faiss::Index* index = index_.get();
|
||||
auto idx = dynamic_cast<faiss::IndexIVF*>(index);
|
||||
if (idx != nullptr) {
|
||||
idx->to_readonly();
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
StatisticsPtr
|
||||
IVF::GetStatistics() {
|
||||
if (IndexMode() != IndexMode::MODE_CPU || !STATISTICS_LEVEL) {
|
||||
return stats;
|
||||
}
|
||||
auto ivf_stats = std::static_pointer_cast<IVFStatistics>(stats);
|
||||
auto ivf_index = static_cast<faiss::IndexIVF*>(index_.get());
|
||||
auto lock = ivf_stats->Lock();
|
||||
ivf_stats->update_ivf_access_stats(ivf_index->nprobe_statistics);
|
||||
return ivf_stats;
|
||||
}
|
||||
|
||||
void
|
||||
IVF::ClearStatistics() {
|
||||
if (IndexMode() != IndexMode::MODE_CPU || !STATISTICS_LEVEL) {
|
||||
return;
|
||||
}
|
||||
auto ivf_stats = std::static_pointer_cast<IVFStatistics>(stats);
|
||||
auto ivf_index = static_cast<faiss::IndexIVF*>(index_.get());
|
||||
ivf_index->clear_nprobe_statistics();
|
||||
ivf_index->index_ivf_stats.reset();
|
||||
auto lock = ivf_stats->Lock();
|
||||
ivf_stats->clear();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,102 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include <faiss/IndexIVF.h>
|
||||
|
||||
#include "knowhere/common/Typedef.h"
|
||||
#include "knowhere/index/vector_index/FaissBaseIndex.h"
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IVF : public VecIndex, public FaissBaseIndex {
|
||||
public:
|
||||
IVF() : FaissBaseIndex(nullptr) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFFLAT;
|
||||
stats = std::make_shared<milvus::knowhere::IVFStatistics>(index_type_);
|
||||
}
|
||||
|
||||
explicit IVF(std::shared_ptr<faiss::Index> index) : FaissBaseIndex(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFFLAT;
|
||||
stats = std::make_shared<milvus::knowhere::IVFStatistics>(index_type_);
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config&) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet&) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override;
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr&, const Config&, const faiss::BitsetView) override;
|
||||
|
||||
#if 0
|
||||
DatasetPtr
|
||||
QueryById(const DatasetPtr& dataset, const Config& config) override;
|
||||
#endif
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
StatisticsPtr
|
||||
GetStatistics() override;
|
||||
|
||||
void
|
||||
ClearStatistics() override;
|
||||
|
||||
#if 0
|
||||
DatasetPtr
|
||||
GetVectorById(const DatasetPtr& dataset, const Config& config) override;
|
||||
#endif
|
||||
|
||||
virtual void
|
||||
Seal();
|
||||
|
||||
virtual VecIndexPtr
|
||||
CopyCpuToGpu(const int64_t, const Config&);
|
||||
|
||||
virtual void
|
||||
GenGraph(const float* data, const int64_t k, GraphType& graph, const Config& config);
|
||||
|
||||
protected:
|
||||
virtual std::shared_ptr<faiss::IVFSearchParameters>
|
||||
GenParams(const Config&);
|
||||
|
||||
virtual void
|
||||
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::BitsetView);
|
||||
|
||||
void
|
||||
SealImpl() override;
|
||||
};
|
||||
|
||||
using IVFPtr = std::shared_ptr<IVF>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,136 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include <faiss/IndexFlat.h>
|
||||
#include <faiss/IndexIVFFlat.h>
|
||||
#include <faiss/clone_index.h>
|
||||
#include <faiss/index_io.h>
|
||||
|
||||
#include "faiss/IndexRHNSW.h"
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFHNSW.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
BinarySet
|
||||
IVFHNSW::Serialize(const Config& config) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
try {
|
||||
// Serialize IVF index and HNSW data
|
||||
auto res_set = SerializeImpl(index_type_);
|
||||
auto index = dynamic_cast<faiss::IndexIVFFlat*>(index_.get());
|
||||
auto real_idx = dynamic_cast<faiss::IndexRHNSWFlat*>(index->quantizer);
|
||||
if (real_idx == nullptr) {
|
||||
KNOWHERE_THROW_MSG("Quantizer index is not a faiss::IndexRHNSWFlat");
|
||||
}
|
||||
|
||||
MemoryIOWriter writer;
|
||||
faiss::write_index(real_idx->storage, &writer);
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
res_set.Append("HNSW_STORAGE", data, writer.rp);
|
||||
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, res_set);
|
||||
}
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IVFHNSW::Load(const BinarySet& binary_set) {
|
||||
try {
|
||||
// Load IVF index and HNSW data
|
||||
Assemble(const_cast<BinarySet&>(binary_set));
|
||||
LoadImpl(binary_set, index_type_);
|
||||
|
||||
auto index = dynamic_cast<faiss::IndexIVFFlat*>(index_.get());
|
||||
MemoryIOReader reader;
|
||||
auto binary = binary_set.GetByName("HNSW_STORAGE");
|
||||
reader.total = static_cast<size_t>(binary->size);
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
auto real_idx = dynamic_cast<faiss::IndexRHNSWFlat*>(index->quantizer);
|
||||
real_idx->storage = faiss::read_index(&reader);
|
||||
real_idx->init_hnsw();
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IVFHNSW::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
auto coarse_quantizer = new faiss::IndexRHNSWFlat(dim, config[IndexParams::M], metric_type);
|
||||
coarse_quantizer->hnsw.efConstruction = config[IndexParams::efConstruction];
|
||||
auto index = std::make_shared<faiss::IndexIVFFlat>(coarse_quantizer, dim, config[IndexParams::nlist].get<int64_t>(),
|
||||
metric_type);
|
||||
index->own_fields = true;
|
||||
index->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
index_ = index;
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
IVFHNSW::CopyCpuToGpu(const int64_t device_id, const Config& config) {
|
||||
KNOWHERE_THROW_MSG("IVFHNSW::CopyCpuToGpu not supported.");
|
||||
}
|
||||
|
||||
void
|
||||
IVFHNSW::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
auto ivf_index = static_cast<faiss::IndexIVFFlat*>(index_.get());
|
||||
auto nb = ivf_index->invlists->compute_ntotal();
|
||||
auto code_size = ivf_index->code_size;
|
||||
auto hnsw_quantizer = dynamic_cast<faiss::IndexRHNSWFlat*>(ivf_index->quantizer);
|
||||
// ivf codes, ivf ids and hnsw_flat quantizer
|
||||
index_size_ = nb * code_size + nb * sizeof(int64_t) + hnsw_quantizer->cal_size();
|
||||
}
|
||||
|
||||
void
|
||||
IVFHNSW::QueryImpl(int64_t n,
|
||||
const float* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
auto params = GenParams(config);
|
||||
auto ivf_index = dynamic_cast<faiss::IndexIVF*>(index_.get());
|
||||
ivf_index->nprobe = std::min(params->nprobe, ivf_index->invlists->nlist);
|
||||
if (params->nprobe > 1 && n <= 4) {
|
||||
ivf_index->parallel_mode = 1;
|
||||
} else {
|
||||
ivf_index->parallel_mode = 0;
|
||||
}
|
||||
// Update HNSW quantizer search param
|
||||
auto hnsw_quantizer = dynamic_cast<faiss::IndexRHNSWFlat*>(ivf_index->quantizer);
|
||||
hnsw_quantizer->hnsw.efSearch = config[IndexParams::ef].get<int64_t>();
|
||||
ivf_index->search(n, data, k, distances, labels, bitset);
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,61 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/index/vector_index/IndexIVF.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IVFHNSW : public IVF {
|
||||
public:
|
||||
IVFHNSW() : IVF() {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFHNSW;
|
||||
}
|
||||
|
||||
explicit IVFHNSW(std::shared_ptr<faiss::Index> index) : IVF(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFHNSW;
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config&) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet&) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyCpuToGpu(const int64_t, const Config&) override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
protected:
|
||||
void
|
||||
QueryImpl(int64_t n,
|
||||
const float* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) override;
|
||||
};
|
||||
|
||||
using IVFHNSWPtr = std::shared_ptr<IVFHNSW>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,109 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <string>
|
||||
|
||||
#include <faiss/IndexFlat.h>
|
||||
#include <faiss/IndexIVFPQ.h>
|
||||
#include <faiss/clone_index.h>
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#endif
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFPQ.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include "knowhere/index/vector_index/ConfAdapter.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h"
|
||||
#endif
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
void
|
||||
IVFPQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
faiss::Index* coarse_quantizer = new faiss::IndexFlat(dim, metric_type);
|
||||
auto index = std::make_shared<faiss::IndexIVFPQ>(coarse_quantizer, dim, config[IndexParams::nlist].get<int64_t>(),
|
||||
config[IndexParams::m].get<int64_t>(),
|
||||
config[IndexParams::nbits].get<int64_t>(), metric_type);
|
||||
index->own_fields = true;
|
||||
index->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
index_ = index;
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
IVFPQ::CopyCpuToGpu(const int64_t device_id, const Config& config) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
auto ivfpq_index = dynamic_cast<faiss::IndexIVFPQ*>(index_.get());
|
||||
int64_t dim = ivfpq_index->d;
|
||||
int64_t m = ivfpq_index->pq.M;
|
||||
int64_t nbits = ivfpq_index->pq.nbits;
|
||||
if (!IVFPQConfAdapter::CheckGPUPQParams(dim, m, nbits)) {
|
||||
return nullptr;
|
||||
}
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
|
||||
ResScope rs(res, device_id, false);
|
||||
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get());
|
||||
|
||||
std::shared_ptr<faiss::Index> device_index;
|
||||
device_index.reset(gpu_index);
|
||||
return std::make_shared<GPUIVFPQ>(device_index, device_id, res);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
|
||||
}
|
||||
#else
|
||||
KNOWHERE_THROW_MSG("Calling IVFPQ::CopyCpuToGpu when we are using CPU version");
|
||||
#endif
|
||||
}
|
||||
|
||||
std::shared_ptr<faiss::IVFSearchParameters>
|
||||
IVFPQ::GenParams(const Config& config) {
|
||||
auto params = std::make_shared<faiss::IVFPQSearchParameters>();
|
||||
params->nprobe = config[IndexParams::nprobe];
|
||||
// params->scan_table_threshold = config["scan_table_threhold"]
|
||||
// params->polysemous_ht = config["polysemous_ht"]
|
||||
// params->max_codes = config["max_codes"]
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
void
|
||||
IVFPQ::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
auto ivfpq_index = dynamic_cast<faiss::IndexIVFPQ*>(index_.get());
|
||||
auto nb = ivfpq_index->invlists->compute_ntotal();
|
||||
auto code_size = ivfpq_index->code_size;
|
||||
auto pq = ivfpq_index->pq;
|
||||
auto nlist = ivfpq_index->nlist;
|
||||
auto d = ivfpq_index->d;
|
||||
|
||||
// ivf codes, ivf ids and quantizer
|
||||
auto capacity = nb * code_size + nb * sizeof(int64_t) + nlist * d * sizeof(float);
|
||||
auto centroid_table = pq.M * pq.ksub * pq.dsub * sizeof(float);
|
||||
auto precomputed_table = nlist * pq.M * pq.ksub * sizeof(float);
|
||||
if (precomputed_table > ivfpq_index->precomputed_table_max_bytes) {
|
||||
// will not precompute table
|
||||
precomputed_table = 0;
|
||||
}
|
||||
index_size_ = capacity + centroid_table + precomputed_table;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,51 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/index/vector_index/IndexIVF.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IVFPQ : public IVF {
|
||||
public:
|
||||
IVFPQ() : IVF() {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFPQ;
|
||||
stats = std::make_shared<milvus::knowhere::IVFStatistics>(index_type_);
|
||||
}
|
||||
|
||||
explicit IVFPQ(std::shared_ptr<faiss::Index> index) : IVF(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFPQ;
|
||||
stats = std::make_shared<milvus::knowhere::IVFStatistics>(index_type_);
|
||||
}
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyCpuToGpu(const int64_t, const Config&) override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
protected:
|
||||
std::shared_ptr<faiss::IVFSearchParameters>
|
||||
GenParams(const Config& config) override;
|
||||
};
|
||||
|
||||
using IVFPQPtr = std::shared_ptr<IVFPQ>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,82 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include <faiss/gpu/GpuAutoTune.h>
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#endif
|
||||
#include <faiss/IndexFlat.h>
|
||||
#include <faiss/IndexScalarQuantizer.h>
|
||||
#include <faiss/clone_index.h>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFSQ.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
|
||||
#endif
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
void
|
||||
IVFSQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
faiss::Index* coarse_quantizer = new faiss::IndexFlat(dim, metric_type);
|
||||
auto index = std::make_shared<faiss::IndexIVFScalarQuantizer>(
|
||||
coarse_quantizer, dim, config[IndexParams::nlist].get<int64_t>(), faiss::QuantizerType::QT_8bit, metric_type);
|
||||
index->own_fields = true;
|
||||
index->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
index_ = index;
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
IVFSQ::CopyCpuToGpu(const int64_t device_id, const Config& config) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
|
||||
ResScope rs(res, device_id, false);
|
||||
|
||||
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get());
|
||||
|
||||
std::shared_ptr<faiss::Index> device_index;
|
||||
device_index.reset(gpu_index);
|
||||
return std::make_shared<GPUIVFSQ>(device_index, device_id, res);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
|
||||
}
|
||||
#else
|
||||
KNOWHERE_THROW_MSG("Calling IVFSQ::CopyCpuToGpu when we are using CPU version");
|
||||
#endif
|
||||
}
|
||||
|
||||
void
|
||||
IVFSQ::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
auto ivfsq_index = dynamic_cast<faiss::IndexIVFScalarQuantizer*>(index_.get());
|
||||
auto nb = ivfsq_index->invlists->compute_ntotal();
|
||||
auto code_size = ivfsq_index->code_size;
|
||||
auto nlist = ivfsq_index->nlist;
|
||||
auto d = ivfsq_index->d;
|
||||
// ivf codes, ivf ids, sq trained vectors and quantizer
|
||||
index_size_ = nb * code_size + nb * sizeof(int64_t) + 2 * d * sizeof(float) + nlist * d * sizeof(float);
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,47 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/index/vector_index/IndexIVF.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IVFSQ : public IVF {
|
||||
public:
|
||||
IVFSQ() : IVF() {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8;
|
||||
stats = std::make_shared<milvus::knowhere::IVFStatistics>(index_type_);
|
||||
}
|
||||
|
||||
explicit IVFSQ(std::shared_ptr<faiss::Index> index) : IVF(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8;
|
||||
stats = std::make_shared<milvus::knowhere::IVFStatistics>(index_type_);
|
||||
}
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyCpuToGpu(const int64_t, const Config&) override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
};
|
||||
|
||||
using IVFSQPtr = std::shared_ptr<IVFSQ>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,221 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexNGT.h"
|
||||
|
||||
#include <omp.h>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
BinarySet
|
||||
IndexNGT::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
std::stringstream obj, grp, prf, tre;
|
||||
index_->saveIndex(obj, grp, prf, tre);
|
||||
|
||||
auto obj_str = obj.str();
|
||||
auto grp_str = grp.str();
|
||||
auto prf_str = prf.str();
|
||||
auto tre_str = tre.str();
|
||||
uint64_t obj_size = obj_str.size();
|
||||
uint64_t grp_size = grp_str.size();
|
||||
uint64_t prf_size = prf_str.size();
|
||||
uint64_t tre_size = tre_str.size();
|
||||
|
||||
std::shared_ptr<uint8_t[]> obj_data(new uint8_t[obj_size]);
|
||||
memcpy(obj_data.get(), obj_str.data(), obj_size);
|
||||
std::shared_ptr<uint8_t[]> grp_data(new uint8_t[grp_size]);
|
||||
memcpy(grp_data.get(), grp_str.data(), grp_size);
|
||||
std::shared_ptr<uint8_t[]> prf_data(new uint8_t[prf_size]);
|
||||
memcpy(prf_data.get(), prf_str.data(), prf_size);
|
||||
std::shared_ptr<uint8_t[]> tre_data(new uint8_t[tre_size]);
|
||||
memcpy(tre_data.get(), tre_str.data(), tre_size);
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("ngt_obj_data", obj_data, obj_size);
|
||||
res_set.Append("ngt_grp_data", grp_data, grp_size);
|
||||
res_set.Append("ngt_prf_data", prf_data, prf_size);
|
||||
res_set.Append("ngt_tre_data", tre_data, tre_size);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, res_set);
|
||||
}
|
||||
return res_set;
|
||||
}
|
||||
|
||||
void
|
||||
IndexNGT::Load(const BinarySet& index_binary) {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
auto obj_data = index_binary.GetByName("ngt_obj_data");
|
||||
std::string obj_str(reinterpret_cast<char*>(obj_data->data.get()), obj_data->size);
|
||||
|
||||
auto grp_data = index_binary.GetByName("ngt_grp_data");
|
||||
std::string grp_str(reinterpret_cast<char*>(grp_data->data.get()), grp_data->size);
|
||||
|
||||
auto prf_data = index_binary.GetByName("ngt_prf_data");
|
||||
std::string prf_str(reinterpret_cast<char*>(prf_data->data.get()), prf_data->size);
|
||||
|
||||
auto tre_data = index_binary.GetByName("ngt_tre_data");
|
||||
std::string tre_str(reinterpret_cast<char*>(tre_data->data.get()), tre_data->size);
|
||||
|
||||
std::stringstream obj(obj_str);
|
||||
std::stringstream grp(grp_str);
|
||||
std::stringstream prf(prf_str);
|
||||
std::stringstream tre(tre_str);
|
||||
|
||||
index_ = std::shared_ptr<NGT::Index>(NGT::Index::loadIndex(obj, grp, prf, tre));
|
||||
}
|
||||
|
||||
void
|
||||
IndexNGT::BuildAll(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
KNOWHERE_THROW_MSG("IndexNGT has no implementation of BuildAll, please use IndexNGT(PANNG/ONNG) instead!");
|
||||
}
|
||||
|
||||
#if 0
|
||||
void
|
||||
IndexNGT::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
KNOWHERE_THROW_MSG("IndexNGT has no implementation of Train, please use IndexNGT(PANNG/ONNG) instead!");
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr);
|
||||
|
||||
NGT::Property prop;
|
||||
prop.setDefaultForCreateIndex();
|
||||
prop.dimension = dim;
|
||||
|
||||
MetricType metric_type = config[Metric::TYPE];
|
||||
|
||||
if (metric_type == Metric::L2)
|
||||
prop.distanceType = NGT::Index::Property::DistanceType::DistanceTypeL2;
|
||||
else if (metric_type == Metric::HAMMING)
|
||||
prop.distanceType = NGT::Index::Property::DistanceType::DistanceTypeHamming;
|
||||
else if (metric_type == Metric::JACCARD)
|
||||
prop.distanceType = NGT::Index::Property::DistanceType::DistanceTypeJaccard;
|
||||
else
|
||||
KNOWHERE_THROW_MSG("Metric type not supported: " + metric_type);
|
||||
index_ =
|
||||
std::shared_ptr<NGT::Index>(NGT::Index::createGraphAndTree(reinterpret_cast<const float*>(p_data), prop, rows));
|
||||
}
|
||||
|
||||
void
|
||||
IndexNGT::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
GET_TENSOR_DATA(dataset_ptr);
|
||||
|
||||
index_->append(reinterpret_cast<const float*>(p_data), rows);
|
||||
}
|
||||
#endif
|
||||
|
||||
DatasetPtr
|
||||
IndexNGT::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
GET_TENSOR_DATA(dataset_ptr);
|
||||
|
||||
int k = config[meta::TOPK].get<int>();
|
||||
auto epsilon = config[IndexParams::epsilon].get<float>();
|
||||
auto edge_size = config[IndexParams::max_search_edges].get<int>();
|
||||
if (edge_size == -1) { // pass -1
|
||||
edge_size--;
|
||||
}
|
||||
size_t id_size = sizeof(int64_t) * k;
|
||||
size_t dist_size = sizeof(float) * k;
|
||||
auto p_id = static_cast<int64_t*>(malloc(id_size * rows));
|
||||
auto p_dist = static_cast<float*>(malloc(dist_size * rows));
|
||||
|
||||
NGT::Command::SearchParameter sp;
|
||||
sp.size = k;
|
||||
|
||||
#pragma omp parallel for
|
||||
for (unsigned int i = 0; i < rows; ++i) {
|
||||
const float* single_query = reinterpret_cast<float*>(const_cast<void*>(p_data)) + i * Dim();
|
||||
|
||||
NGT::Object* object = index_->allocateObject(single_query, Dim());
|
||||
NGT::SearchContainer sc(*object);
|
||||
|
||||
// double epsilon = sp.beginOfEpsilon;
|
||||
|
||||
NGT::ObjectDistances res;
|
||||
sc.setResults(&res);
|
||||
sc.setSize(static_cast<size_t>(sp.size));
|
||||
sc.setRadius(sp.radius);
|
||||
|
||||
if (sp.accuracy > 0.0) {
|
||||
sc.setExpectedAccuracy(sp.accuracy);
|
||||
} else {
|
||||
sc.setEpsilon(epsilon);
|
||||
}
|
||||
// sc.setEdgeSize(sp.edgeSize);
|
||||
sc.setEdgeSize(edge_size);
|
||||
|
||||
try {
|
||||
index_->search(sc, bitset);
|
||||
} catch (NGT::Exception& err) {
|
||||
KNOWHERE_THROW_MSG("Query failed");
|
||||
}
|
||||
|
||||
auto local_id = p_id + i * k;
|
||||
auto local_dist = p_dist + i * k;
|
||||
|
||||
int64_t res_num = res.size();
|
||||
float dis_coefficient = 1.0;
|
||||
if (index_->getObjectSpace().getDistanceType() == NGT::ObjectSpace::DistanceType::DistanceTypeIP) {
|
||||
dis_coefficient = -1.0;
|
||||
}
|
||||
for (int64_t idx = 0; idx < res_num; ++idx) {
|
||||
*(local_id + idx) = res[idx].id - 1;
|
||||
*(local_dist + idx) = res[idx].distance * dis_coefficient;
|
||||
}
|
||||
MapOffsetToUid(local_id, res_num);
|
||||
while (res_num < static_cast<int64_t>(k)) {
|
||||
*(local_id + res_num) = -1;
|
||||
*(local_dist + res_num) = 1.0 / 0.0;
|
||||
}
|
||||
index_->deleteObject(object);
|
||||
}
|
||||
|
||||
auto res_ds = std::make_shared<Dataset>();
|
||||
res_ds->Set(meta::IDS, p_id);
|
||||
res_ds->Set(meta::DISTANCE, p_dist);
|
||||
return res_ds;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IndexNGT::Count() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->getNumberOfVectors();
|
||||
}
|
||||
|
||||
int64_t
|
||||
IndexNGT::Dim() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->getDimension();
|
||||
}
|
||||
|
||||
void
|
||||
IndexNGT::UpdateIndexSize() {
|
||||
KNOWHERE_THROW_MSG("IndexNGT has no implementation of UpdateIndexSize, please use IndexNGT(PANNG/ONNG) instead!");
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,68 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <NGT/lib/NGT/Command.h>
|
||||
#include <NGT/lib/NGT/Common.h>
|
||||
#include <NGT/lib/NGT/Index.h>
|
||||
|
||||
#include <knowhere/common/Exception.h>
|
||||
#include <knowhere/index/IndexType.h>
|
||||
#include <knowhere/index/vector_index/VecIndex.h>
|
||||
#include <memory>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexNGT : public VecIndex {
|
||||
public:
|
||||
IndexNGT() {
|
||||
index_type_ = IndexEnum::INVALID;
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
BuildAll(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override {
|
||||
KNOWHERE_THROW_MSG("NGT not support add item dynamically, please invoke BuildAll interface.");
|
||||
}
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) override {
|
||||
KNOWHERE_THROW_MSG("Incremental index is not supported");
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) override;
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
protected:
|
||||
std::shared_ptr<NGT::Index> index_ = nullptr;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,77 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexNGTONNG.h"
|
||||
|
||||
#include "NGT/lib/NGT/GraphOptimizer.h"
|
||||
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
#include <cstddef>
|
||||
#include <memory>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
void
|
||||
IndexNGTONNG::BuildAll(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr);
|
||||
|
||||
NGT::Property prop;
|
||||
prop.setDefaultForCreateIndex();
|
||||
prop.dimension = dim;
|
||||
|
||||
auto edge_size = config[IndexParams::edge_size].get<int64_t>();
|
||||
prop.edgeSizeForCreation = edge_size;
|
||||
prop.insertionRadiusCoefficient = 1.0;
|
||||
|
||||
MetricType metric_type = config[Metric::TYPE];
|
||||
|
||||
if (metric_type == Metric::L2) {
|
||||
prop.distanceType = NGT::Index::Property::DistanceType::DistanceTypeL2;
|
||||
} else if (metric_type == Metric::IP) {
|
||||
prop.distanceType = NGT::Index::Property::DistanceType::DistanceTypeIP;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Metric type not supported: " + metric_type);
|
||||
}
|
||||
|
||||
index_ =
|
||||
std::shared_ptr<NGT::Index>(NGT::Index::createGraphAndTree(reinterpret_cast<const float*>(p_data), prop, rows));
|
||||
|
||||
// reconstruct graph
|
||||
NGT::GraphOptimizer graphOptimizer(false);
|
||||
|
||||
auto number_of_outgoing_edges = config[IndexParams::outgoing_edge_size].get<size_t>();
|
||||
auto number_of_incoming_edges = config[IndexParams::incoming_edge_size].get<size_t>();
|
||||
|
||||
graphOptimizer.shortcutReduction = true;
|
||||
graphOptimizer.searchParameterOptimization = false;
|
||||
graphOptimizer.prefetchParameterOptimization = false;
|
||||
graphOptimizer.accuracyTableGeneration = false;
|
||||
graphOptimizer.margin = 0.2;
|
||||
graphOptimizer.gtEpsilon = 0.1;
|
||||
|
||||
graphOptimizer.set(number_of_outgoing_edges, number_of_incoming_edges, 1000, 20);
|
||||
|
||||
graphOptimizer.execute(*index_);
|
||||
}
|
||||
|
||||
void
|
||||
IndexNGTONNG::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
index_size_ = index_->memSize();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,33 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "knowhere/index/vector_index/IndexNGT.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexNGTONNG : public IndexNGT {
|
||||
public:
|
||||
IndexNGTONNG() {
|
||||
index_type_ = IndexEnum::INDEX_NGTONNG;
|
||||
}
|
||||
|
||||
void
|
||||
BuildAll(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,138 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexNGTPANNG.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
#include <memory>
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
void
|
||||
IndexNGTPANNG::BuildAll(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr);
|
||||
|
||||
NGT::Property prop;
|
||||
prop.setDefaultForCreateIndex();
|
||||
prop.dimension = dim;
|
||||
|
||||
auto edge_size = config[IndexParams::edge_size].get<int64_t>();
|
||||
prop.edgeSizeForCreation = edge_size;
|
||||
|
||||
MetricType metric_type = config[Metric::TYPE];
|
||||
|
||||
if (metric_type == Metric::L2) {
|
||||
prop.distanceType = NGT::Index::Property::DistanceType::DistanceTypeL2;
|
||||
} else if (metric_type == Metric::IP) {
|
||||
prop.distanceType = NGT::Index::Property::DistanceType::DistanceTypeIP;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Metric type not supported: " + metric_type);
|
||||
}
|
||||
|
||||
index_ =
|
||||
std::shared_ptr<NGT::Index>(NGT::Index::createGraphAndTree(reinterpret_cast<const float*>(p_data), prop, rows));
|
||||
|
||||
auto forcedly_pruned_edge_size = config[IndexParams::forcedly_pruned_edge_size].get<int64_t>();
|
||||
auto selectively_pruned_edge_size = config[IndexParams::selectively_pruned_edge_size].get<int64_t>();
|
||||
|
||||
if (!forcedly_pruned_edge_size && !selectively_pruned_edge_size) {
|
||||
KNOWHERE_THROW_MSG(
|
||||
"a lack of parameters forcedly_pruned_edge_size and selectively_pruned_edge_size 4 index NGTPANNG");
|
||||
return;
|
||||
}
|
||||
|
||||
if (forcedly_pruned_edge_size && selectively_pruned_edge_size &&
|
||||
selectively_pruned_edge_size >= forcedly_pruned_edge_size) {
|
||||
KNOWHERE_THROW_MSG("Selectively pruned edge size should less than remaining edge size");
|
||||
}
|
||||
|
||||
// std::map<size_t, size_t> stats;
|
||||
// size_t max_len = 0;
|
||||
|
||||
// prune
|
||||
auto& graph = dynamic_cast<NGT::GraphIndex&>(index_->getIndex());
|
||||
for (size_t id = 1; id < graph.repository.size(); id++) {
|
||||
try {
|
||||
NGT::GraphNode& node = *graph.getNode(id);
|
||||
// auto sz = node.size();
|
||||
// if (max_len < sz)
|
||||
// max_len = sz;
|
||||
// auto fd = stats.find(sz);
|
||||
// if (fd != stats.end()) {
|
||||
// fd->second ++;
|
||||
// } else {
|
||||
// stats[sz] = 1;
|
||||
// }
|
||||
if (node.size() >= forcedly_pruned_edge_size) {
|
||||
node.resize(forcedly_pruned_edge_size);
|
||||
}
|
||||
if (node.size() >= selectively_pruned_edge_size) {
|
||||
size_t rank = 0;
|
||||
for (auto i = node.begin(); i != node.end(); ++rank) {
|
||||
if (rank >= selectively_pruned_edge_size) {
|
||||
bool found = false;
|
||||
for (size_t t1 = 0; t1 < node.size() && found == false; ++t1) {
|
||||
if (t1 >= selectively_pruned_edge_size) {
|
||||
break;
|
||||
}
|
||||
if (rank == t1) { // can't reach here
|
||||
continue;
|
||||
}
|
||||
NGT::GraphNode& node2 = *graph.getNode(node[t1].id);
|
||||
for (size_t t2 = 0; t2 < node2.size(); ++t2) {
|
||||
if (t2 >= selectively_pruned_edge_size) {
|
||||
break;
|
||||
}
|
||||
if (node2[t2].id == (*i).id) {
|
||||
found = true;
|
||||
break;
|
||||
}
|
||||
} // for
|
||||
} // for
|
||||
if (found) {
|
||||
// remove
|
||||
i = node.erase(i);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
i++;
|
||||
} // for
|
||||
}
|
||||
} catch (NGT::Exception& err) {
|
||||
std::cerr << "Graph::search: Warning. Cannot get the node. ID=" << id << ":" << err.what() << std::endl;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
/*
|
||||
std::vector<size_t> cnt(max_len, 0);
|
||||
for (auto &pr : stats) {
|
||||
cnt[pr.first] = pr.second;
|
||||
}
|
||||
for (auto i = 0; i < cnt.size(); ++ i) {
|
||||
if (cnt[i]) {
|
||||
std::cout << "len = " << i << ", cnt = " << cnt[i] << std::endl;
|
||||
}
|
||||
}
|
||||
*/
|
||||
}
|
||||
|
||||
void
|
||||
IndexNGTPANNG::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
index_size_ = index_->memSize();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,33 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "knowhere/index/vector_index/IndexNGT.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexNGTPANNG : public IndexNGT {
|
||||
public:
|
||||
IndexNGTPANNG() {
|
||||
index_type_ = IndexEnum::INDEX_NGTPANNG;
|
||||
}
|
||||
|
||||
void
|
||||
BuildAll(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,206 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexRHNSW.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cassert>
|
||||
#include <chrono>
|
||||
#include <iterator>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
BinarySet
|
||||
IndexRHNSW::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
try {
|
||||
MemoryIOWriter writer;
|
||||
writer.name = this->index_type() + "_Index";
|
||||
faiss::write_index(index_.get(), &writer);
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append(writer.name, data, writer.rp);
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSW::Load(const BinarySet& index_binary) {
|
||||
try {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
MemoryIOReader reader;
|
||||
reader.name = this->index_type() + "_Index";
|
||||
auto binary = index_binary.GetByName(reader.name);
|
||||
|
||||
reader.total = static_cast<size_t>(binary->size);
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
auto idx = faiss::read_index(&reader);
|
||||
auto hnsw_stats = std::static_pointer_cast<RHNSWStatistics>(stats);
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
auto real_idx = static_cast<faiss::IndexRHNSW*>(idx);
|
||||
auto lock = hnsw_stats->Lock();
|
||||
hnsw_stats->update_level_distribution(real_idx->hnsw.max_level, real_idx->hnsw.level_stats);
|
||||
real_idx->set_target_level(hnsw_stats->target_level);
|
||||
// LOG_KNOWHERE_DEBUG_ << "IndexRHNSW::Load finished, show statistics:";
|
||||
// LOG_KNOWHERE_DEBUG_ << hnsw_stats->ToString();
|
||||
}
|
||||
index_.reset(idx);
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSW::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
KNOWHERE_THROW_MSG("IndexRHNSW has no implementation of Train, please use IndexRHNSW(Flat/SQ/PQ) instead!");
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSW::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
|
||||
index_->add(rows, reinterpret_cast<const float*>(p_data));
|
||||
auto hnsw_stats = std::static_pointer_cast<RHNSWStatistics>(stats);
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
auto real_idx = static_cast<faiss::IndexRHNSW*>(index_.get());
|
||||
auto lock = hnsw_stats->Lock();
|
||||
hnsw_stats->update_level_distribution(real_idx->hnsw.max_level, real_idx->hnsw.level_stats);
|
||||
real_idx->set_target_level(hnsw_stats->target_level);
|
||||
}
|
||||
// LOG_KNOWHERE_DEBUG_ << "IndexRHNSW::Load finished, show statistics:";
|
||||
// LOG_KNOWHERE_DEBUG_ << GetStatistics()->ToString();
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
IndexRHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
GET_TENSOR_DATA(dataset_ptr)
|
||||
auto k = config[meta::TOPK].get<int64_t>();
|
||||
auto result_count = rows * k;
|
||||
|
||||
auto p_id = static_cast<int64_t*>(malloc(result_count * sizeof(int64_t)));
|
||||
auto p_dist = static_cast<float*>(malloc(result_count * sizeof(float)));
|
||||
for (int64_t i = 0; i < result_count; ++i) {
|
||||
p_id[i] = -1;
|
||||
p_dist[i] = -1;
|
||||
}
|
||||
|
||||
auto real_index = dynamic_cast<faiss::IndexRHNSW*>(index_.get());
|
||||
|
||||
real_index->hnsw.efSearch = (config[IndexParams::ef].get<int64_t>());
|
||||
|
||||
std::chrono::high_resolution_clock::time_point query_start, query_end;
|
||||
query_start = std::chrono::high_resolution_clock::now();
|
||||
real_index->search(rows, reinterpret_cast<const float*>(p_data), k, p_dist, p_id, bitset);
|
||||
query_end = std::chrono::high_resolution_clock::now();
|
||||
if (STATISTICS_LEVEL) {
|
||||
auto hnsw_stats = std::dynamic_pointer_cast<RHNSWStatistics>(stats);
|
||||
auto lock = hnsw_stats->Lock();
|
||||
if (STATISTICS_LEVEL >= 1) {
|
||||
hnsw_stats->update_nq(rows);
|
||||
hnsw_stats->update_ef_sum(real_index->hnsw.efSearch * rows);
|
||||
hnsw_stats->update_total_query_time(
|
||||
std::chrono::duration_cast<std::chrono::milliseconds>(query_end - query_start).count());
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 2) {
|
||||
hnsw_stats->update_filter_percentage(bitset);
|
||||
}
|
||||
}
|
||||
// LOG_KNOWHERE_DEBUG_ << "IndexRHNSW::Load finished, show statistics:";
|
||||
// LOG_KNOWHERE_DEBUG_ << GetStatistics()->ToString();
|
||||
|
||||
MapOffsetToUid(p_id, result_count);
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
return ret_ds;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IndexRHNSW::Count() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->ntotal;
|
||||
}
|
||||
|
||||
int64_t
|
||||
IndexRHNSW::Dim() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_->d;
|
||||
}
|
||||
|
||||
StatisticsPtr
|
||||
IndexRHNSW::GetStatistics() {
|
||||
if (!STATISTICS_LEVEL) {
|
||||
return stats;
|
||||
}
|
||||
auto hnsw_stats = std::static_pointer_cast<RHNSWStatistics>(stats);
|
||||
auto real_index = static_cast<faiss::IndexRHNSW*>(index_.get());
|
||||
auto lock = hnsw_stats->Lock();
|
||||
real_index->get_sorted_access_counts(hnsw_stats->access_cnt, hnsw_stats->access_total);
|
||||
return hnsw_stats;
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSW::ClearStatistics() {
|
||||
if (!STATISTICS_LEVEL) {
|
||||
return;
|
||||
}
|
||||
auto hnsw_stats = std::static_pointer_cast<RHNSWStatistics>(stats);
|
||||
auto real_index = static_cast<faiss::IndexRHNSW*>(index_.get());
|
||||
real_index->clear_stats();
|
||||
auto lock = hnsw_stats->Lock();
|
||||
hnsw_stats->clear();
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSW::UpdateIndexSize() {
|
||||
KNOWHERE_THROW_MSG(
|
||||
"IndexRHNSW has no implementation of UpdateIndexSize, please use IndexRHNSW(Flat/SQ/PQ) instead!");
|
||||
}
|
||||
|
||||
/*
|
||||
BinarySet
|
||||
IndexRHNSW::SerializeImpl(const milvus::knowhere::IndexType &type) { return BinarySet(); }
|
||||
|
||||
void
|
||||
IndexRHNSW::SealImpl() {}
|
||||
|
||||
void
|
||||
IndexRHNSW::LoadImpl(const milvus::knowhere::BinarySet &, const milvus::knowhere::IndexType &type) {}
|
||||
*/
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,72 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/FaissBaseIndex.h"
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
|
||||
#include <faiss/index_io.h>
|
||||
#include "faiss/IndexRHNSW.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexRHNSW : public VecIndex, public FaissBaseIndex {
|
||||
public:
|
||||
IndexRHNSW() : FaissBaseIndex(nullptr) {
|
||||
index_type_ = IndexEnum::INVALID;
|
||||
stats = std::make_shared<milvus::knowhere::RHNSWStatistics>(index_type_);
|
||||
}
|
||||
|
||||
explicit IndexRHNSW(std::shared_ptr<faiss::Index> index) : FaissBaseIndex(std::move(index)) {
|
||||
index_type_ = IndexEnum::INVALID;
|
||||
stats = std::make_shared<milvus::knowhere::RHNSWStatistics>(index_type_);
|
||||
}
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override;
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) override;
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
StatisticsPtr
|
||||
GetStatistics() override;
|
||||
|
||||
void
|
||||
ClearStatistics() override;
|
||||
};
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,109 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexRHNSWFlat.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cassert>
|
||||
#include <iterator>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
IndexRHNSWFlat::IndexRHNSWFlat(int d, int M, milvus::knowhere::MetricType metric) {
|
||||
faiss::MetricType mt =
|
||||
metric == Metric::L2 ? faiss::MetricType::METRIC_L2 : faiss::MetricType::METRIC_INNER_PRODUCT;
|
||||
index_ = std::shared_ptr<faiss::Index>(new faiss::IndexRHNSWFlat(d, M, mt));
|
||||
}
|
||||
|
||||
BinarySet
|
||||
IndexRHNSWFlat::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
try {
|
||||
auto res_set = IndexRHNSW::Serialize(config);
|
||||
auto real_idx = dynamic_cast<faiss::IndexRHNSWFlat*>(index_.get());
|
||||
if (real_idx == nullptr) {
|
||||
KNOWHERE_THROW_MSG("index is not a faiss::IndexRHNSWFlat");
|
||||
}
|
||||
|
||||
int64_t meta_info[3] = {real_idx->storage->metric_type, real_idx->storage->d, real_idx->storage->ntotal};
|
||||
auto meta_space = new uint8_t[sizeof(meta_info)];
|
||||
memcpy(meta_space, meta_info, sizeof(meta_info));
|
||||
std::shared_ptr<uint8_t[]> space_sp(meta_space, std::default_delete<uint8_t[]>());
|
||||
res_set.Append("META", space_sp, sizeof(meta_info));
|
||||
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, res_set);
|
||||
}
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWFlat::Load(const BinarySet& index_binary) {
|
||||
try {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
IndexRHNSW::Load(index_binary);
|
||||
|
||||
int64_t meta_info[3]; // = {metric_type, dim, ntotal}
|
||||
auto meta_data = index_binary.GetByName("META");
|
||||
memcpy(meta_info, meta_data->data.get(), meta_data->size);
|
||||
|
||||
auto real_idx = dynamic_cast<faiss::IndexRHNSWFlat*>(index_.get());
|
||||
real_idx->storage =
|
||||
new faiss::IndexFlat(static_cast<faiss::idx_t>(meta_info[1]), static_cast<faiss::MetricType>(meta_info[0]));
|
||||
auto binary_data = index_binary.GetByName(RAW_DATA);
|
||||
real_idx->storage->add(meta_info[2], reinterpret_cast<const float*>(binary_data->data.get()));
|
||||
real_idx->init_hnsw();
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWFlat::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
try {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
int32_t efConstruction = config[IndexParams::efConstruction];
|
||||
|
||||
auto idx = new faiss::IndexRHNSWFlat(int(dim), config[IndexParams::M], metric_type);
|
||||
idx->hnsw.efConstruction = efConstruction;
|
||||
index_ = std::shared_ptr<faiss::Index>(idx);
|
||||
index_->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWFlat::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
index_size_ = dynamic_cast<faiss::IndexRHNSWFlat*>(index_.get())->cal_size();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,51 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <utility>
|
||||
|
||||
#include "IndexRHNSW.h"
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexRHNSWFlat : public IndexRHNSW {
|
||||
public:
|
||||
IndexRHNSWFlat() : IndexRHNSW() {
|
||||
index_type_ = IndexEnum::INDEX_RHNSWFlat;
|
||||
}
|
||||
|
||||
explicit IndexRHNSWFlat(std::shared_ptr<faiss::Index> index) : IndexRHNSW(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_RHNSWFlat;
|
||||
}
|
||||
|
||||
IndexRHNSWFlat(int d, int M, MetricType metric = Metric::L2);
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,106 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexRHNSWPQ.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cassert>
|
||||
#include <iterator>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
IndexRHNSWPQ::IndexRHNSWPQ(int d, int pq_m, int M) {
|
||||
index_ = std::shared_ptr<faiss::Index>(new faiss::IndexRHNSWPQ(d, pq_m, M));
|
||||
}
|
||||
|
||||
BinarySet
|
||||
IndexRHNSWPQ::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
try {
|
||||
auto res_set = IndexRHNSW::Serialize(config);
|
||||
MemoryIOWriter writer;
|
||||
writer.name = QUANTIZATION_DATA;
|
||||
auto real_idx = dynamic_cast<faiss::IndexRHNSWPQ*>(index_.get());
|
||||
if (real_idx == nullptr) {
|
||||
KNOWHERE_THROW_MSG("dynamic_cast<faiss::IndexRHNSWPQ*>(index_) failed during Serialize!");
|
||||
}
|
||||
faiss::write_index(real_idx->storage, &writer);
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
res_set.Append(writer.name, data, writer.rp);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, res_set);
|
||||
}
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWPQ::Load(const BinarySet& index_binary) {
|
||||
try {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
IndexRHNSW::Load(index_binary);
|
||||
MemoryIOReader reader;
|
||||
reader.name = QUANTIZATION_DATA;
|
||||
auto binary = index_binary.GetByName(reader.name);
|
||||
|
||||
reader.total = static_cast<size_t>(binary->size);
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
auto real_idx = dynamic_cast<faiss::IndexRHNSWPQ*>(index_.get());
|
||||
if (real_idx == nullptr) {
|
||||
KNOWHERE_THROW_MSG("dynamic_cast<faiss::IndexRHNSWPQ*>(index_) failed during Load!");
|
||||
}
|
||||
real_idx->storage = faiss::read_index(&reader);
|
||||
real_idx->init_hnsw();
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWPQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
try {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
int32_t efConstruction = config[IndexParams::efConstruction];
|
||||
|
||||
auto idx = new faiss::IndexRHNSWPQ(int(dim), config[IndexParams::PQM], config[IndexParams::M]);
|
||||
idx->hnsw.efConstruction = efConstruction;
|
||||
index_ = std::shared_ptr<faiss::Index>(idx);
|
||||
index_->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWPQ::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
index_size_ = dynamic_cast<faiss::IndexRHNSWPQ*>(index_.get())->cal_size();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,52 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <utility>
|
||||
|
||||
#include "IndexRHNSW.h"
|
||||
#include "knowhere/common/Exception.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexRHNSWPQ : public IndexRHNSW {
|
||||
public:
|
||||
IndexRHNSWPQ() : IndexRHNSW() {
|
||||
index_type_ = IndexEnum::INDEX_RHNSWPQ;
|
||||
}
|
||||
|
||||
explicit IndexRHNSWPQ(std::shared_ptr<faiss::Index> index) : IndexRHNSW(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_RHNSWPQ;
|
||||
}
|
||||
|
||||
IndexRHNSWPQ(int d, int pq_m, int M);
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
private:
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,111 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/IndexRHNSWSQ.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cassert>
|
||||
#include <iterator>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
IndexRHNSWSQ::IndexRHNSWSQ(int d, faiss::QuantizerType qtype, int M, milvus::knowhere::MetricType metric) {
|
||||
faiss::MetricType mt =
|
||||
metric == Metric::L2 ? faiss::MetricType::METRIC_L2 : faiss::MetricType::METRIC_INNER_PRODUCT;
|
||||
index_ = std::shared_ptr<faiss::Index>(new faiss::IndexRHNSWSQ(d, qtype, M, mt));
|
||||
}
|
||||
|
||||
BinarySet
|
||||
IndexRHNSWSQ::Serialize(const Config& config) {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
try {
|
||||
auto res_set = IndexRHNSW::Serialize(config);
|
||||
MemoryIOWriter writer;
|
||||
writer.name = QUANTIZATION_DATA;
|
||||
auto real_idx = dynamic_cast<faiss::IndexRHNSWSQ*>(index_.get());
|
||||
if (real_idx == nullptr) {
|
||||
KNOWHERE_THROW_MSG("dynamic_cast<faiss::IndexRHNSWSQ*>(index_) failed during Serialize!");
|
||||
}
|
||||
faiss::write_index(real_idx->storage, &writer);
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
res_set.Append(writer.name, data, writer.rp);
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, res_set);
|
||||
}
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWSQ::Load(const BinarySet& index_binary) {
|
||||
try {
|
||||
Assemble(const_cast<BinarySet&>(index_binary));
|
||||
IndexRHNSW::Load(index_binary);
|
||||
MemoryIOReader reader;
|
||||
reader.name = QUANTIZATION_DATA;
|
||||
auto binary = index_binary.GetByName(reader.name);
|
||||
|
||||
reader.total = static_cast<size_t>(binary->size);
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
auto real_idx = dynamic_cast<faiss::IndexRHNSWSQ*>(index_.get());
|
||||
if (real_idx == nullptr) {
|
||||
KNOWHERE_THROW_MSG("dynamic_cast<faiss::IndexRHNSWSQ*>(index_) failed during Load!");
|
||||
}
|
||||
real_idx->storage = faiss::read_index(&reader);
|
||||
real_idx->init_hnsw();
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWSQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
try {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
int32_t efConstruction = config[IndexParams::efConstruction];
|
||||
|
||||
auto idx =
|
||||
new faiss::IndexRHNSWSQ(int(dim), faiss::QuantizerType::QT_8bit, config[IndexParams::M], metric_type);
|
||||
idx->hnsw.efConstruction = efConstruction;
|
||||
index_ = std::shared_ptr<faiss::Index>(idx);
|
||||
index_->train(rows, static_cast<const float*>(p_data));
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IndexRHNSWSQ::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
index_size_ = dynamic_cast<faiss::IndexRHNSWSQ*>(index_.get())->cal_size();
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,52 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <utility>
|
||||
|
||||
#include "IndexRHNSW.h"
|
||||
#include "knowhere/common/Exception.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class IndexRHNSWSQ : public IndexRHNSW {
|
||||
public:
|
||||
IndexRHNSWSQ() : IndexRHNSW() {
|
||||
index_type_ = IndexEnum::INDEX_RHNSWSQ;
|
||||
}
|
||||
|
||||
explicit IndexRHNSWSQ(std::shared_ptr<faiss::Index> index) : IndexRHNSW(std::move(index)) {
|
||||
index_type_ = IndexEnum::INDEX_RHNSWSQ;
|
||||
}
|
||||
|
||||
IndexRHNSWSQ(int d, faiss::QuantizerType qtype, int M, MetricType metric = Metric::L2);
|
||||
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_binary) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
private:
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,245 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include <SPTAG/AnnService/inc/Core/Common.h>
|
||||
#include <SPTAG/AnnService/inc/Core/VectorSet.h>
|
||||
#include <SPTAG/AnnService/inc/Server/QueryParser.h>
|
||||
|
||||
#include <array>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
|
||||
#undef mkdir
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/IndexSPTAG.h"
|
||||
#include "knowhere/index/vector_index/adapter/SptagAdapter.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/SPTAGParameterMgr.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
CPUSPTAGRNG::CPUSPTAGRNG(const std::string& IndexType) {
|
||||
if (IndexType == "KDT") {
|
||||
index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::KDT, SPTAG::VectorValueType::Float);
|
||||
index_ptr_->SetParameter("DistCalcMethod", "L2");
|
||||
index_type_ = IndexEnum::INDEX_SPTAG_KDT_RNT;
|
||||
} else {
|
||||
index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::BKT, SPTAG::VectorValueType::Float);
|
||||
index_ptr_->SetParameter("DistCalcMethod", "L2");
|
||||
index_type_ = IndexEnum::INDEX_SPTAG_BKT_RNT;
|
||||
}
|
||||
}
|
||||
|
||||
BinarySet
|
||||
CPUSPTAGRNG::Serialize(const Config& config) {
|
||||
std::string index_config;
|
||||
std::vector<SPTAG::ByteArray> index_blobs;
|
||||
|
||||
std::shared_ptr<std::vector<std::uint64_t>> buffersize = index_ptr_->CalculateBufferSize();
|
||||
std::vector<char*> res(buffersize->size() + 1);
|
||||
for (uint64_t i = 1; i < res.size(); i++) {
|
||||
res[i] = new char[buffersize->at(i - 1)];
|
||||
auto ptr = &res[i][0];
|
||||
index_blobs.emplace_back(SPTAG::ByteArray((std::uint8_t*)ptr, buffersize->at(i - 1), false));
|
||||
}
|
||||
|
||||
index_ptr_->SaveIndex(index_config, index_blobs);
|
||||
|
||||
size_t length = index_config.length();
|
||||
char* cstr = new char[length];
|
||||
snprintf(cstr, length, "%s", index_config.c_str());
|
||||
|
||||
BinarySet binary_set;
|
||||
std::shared_ptr<uint8_t[]> sample;
|
||||
sample.reset(static_cast<uint8_t*>(index_blobs[0].Data()));
|
||||
std::shared_ptr<uint8_t[]> tree;
|
||||
tree.reset(static_cast<uint8_t*>(index_blobs[1].Data()));
|
||||
std::shared_ptr<uint8_t[]> graph;
|
||||
graph.reset(static_cast<uint8_t*>(index_blobs[2].Data()));
|
||||
std::shared_ptr<uint8_t[]> deleteid;
|
||||
deleteid.reset(static_cast<uint8_t*>(index_blobs[3].Data()));
|
||||
std::shared_ptr<uint8_t[]> metadata1;
|
||||
metadata1.reset(static_cast<uint8_t*>(index_blobs[4].Data()));
|
||||
std::shared_ptr<uint8_t[]> metadata2;
|
||||
metadata2.reset(static_cast<uint8_t*>(index_blobs[5].Data()));
|
||||
std::shared_ptr<uint8_t[]> x_cfg;
|
||||
x_cfg.reset(static_cast<uint8_t*>((void*)cstr));
|
||||
|
||||
binary_set.Append("samples", sample, index_blobs[0].Length());
|
||||
binary_set.Append("tree", tree, index_blobs[1].Length());
|
||||
binary_set.Append("deleteid", deleteid, index_blobs[3].Length());
|
||||
binary_set.Append("metadata1", metadata1, index_blobs[4].Length());
|
||||
binary_set.Append("metadata2", metadata2, index_blobs[5].Length());
|
||||
binary_set.Append("config", x_cfg, length);
|
||||
binary_set.Append("graph", graph, index_blobs[2].Length());
|
||||
|
||||
if (config.contains(INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
|
||||
Disassemble(config[INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<int64_t>() * 1024 * 1024, binary_set);
|
||||
}
|
||||
return binary_set;
|
||||
}
|
||||
|
||||
void
|
||||
CPUSPTAGRNG::Load(const BinarySet& binary_set) {
|
||||
Assemble(const_cast<BinarySet&>(binary_set));
|
||||
std::string index_config;
|
||||
std::vector<SPTAG::ByteArray> index_blobs;
|
||||
|
||||
auto samples = binary_set.GetByName("samples");
|
||||
index_blobs.push_back(SPTAG::ByteArray(samples->data.get(), samples->size, false));
|
||||
|
||||
auto tree = binary_set.GetByName("tree");
|
||||
index_blobs.push_back(SPTAG::ByteArray(tree->data.get(), tree->size, false));
|
||||
|
||||
auto graph = binary_set.GetByName("graph");
|
||||
index_blobs.push_back(SPTAG::ByteArray(graph->data.get(), graph->size, false));
|
||||
|
||||
auto deleteid = binary_set.GetByName("deleteid");
|
||||
index_blobs.push_back(SPTAG::ByteArray(deleteid->data.get(), deleteid->size, false));
|
||||
|
||||
auto metadata1 = binary_set.GetByName("metadata1");
|
||||
index_blobs.push_back(SPTAG::ByteArray(CopyBinary(metadata1), metadata1->size, true));
|
||||
|
||||
auto metadata2 = binary_set.GetByName("metadata2");
|
||||
index_blobs.push_back(SPTAG::ByteArray(metadata2->data.get(), metadata2->size, false));
|
||||
|
||||
auto config = binary_set.GetByName("config");
|
||||
index_config = reinterpret_cast<char*>(config->data.get());
|
||||
|
||||
index_ptr_->LoadIndex(index_config, index_blobs);
|
||||
}
|
||||
|
||||
void
|
||||
CPUSPTAGRNG::BuildAll(const DatasetPtr& origin, const Config& train_config) {
|
||||
SetParameters(train_config);
|
||||
|
||||
DatasetPtr dataset = origin;
|
||||
|
||||
auto vectorset = ConvertToVectorSet(dataset);
|
||||
auto metaset = ConvertToMetadataSet(dataset);
|
||||
index_ptr_->BuildIndex(vectorset, metaset);
|
||||
}
|
||||
|
||||
void
|
||||
CPUSPTAGRNG::SetParameters(const Config& config) {
|
||||
#define Assign(param_name, str_name) \
|
||||
index_ptr_->SetParameter(str_name, std::to_string(build_cfg[param_name].get<int64_t>()))
|
||||
|
||||
if (index_type_ == IndexEnum::INDEX_SPTAG_KDT_RNT) {
|
||||
auto build_cfg = SPTAGParameterMgr::GetInstance().GetKDTParameters();
|
||||
|
||||
Assign("kdtnumber", "KDTNumber");
|
||||
Assign("numtopdimensionkdtsplit", "NumTopDimensionKDTSplit");
|
||||
Assign("samples", "Samples");
|
||||
Assign("tptnumber", "TPTNumber");
|
||||
Assign("tptleafsize", "TPTLeafSize");
|
||||
Assign("numtopdimensiontptsplit", "NumTopDimensionTPTSplit");
|
||||
Assign("neighborhoodsize", "NeighborhoodSize");
|
||||
Assign("graphneighborhoodscale", "GraphNeighborhoodScale");
|
||||
Assign("graphcefscale", "GraphCEFScale");
|
||||
Assign("refineiterations", "RefineIterations");
|
||||
Assign("cef", "CEF");
|
||||
Assign("maxcheckforrefinegraph", "MaxCheckForRefineGraph");
|
||||
Assign("numofthreads", "NumberOfThreads");
|
||||
Assign("maxcheck", "MaxCheck");
|
||||
Assign("thresholdofnumberofcontinuousnobetterpropagation", "ThresholdOfNumberOfContinuousNoBetterPropagation");
|
||||
Assign("numberofinitialdynamicpivots", "NumberOfInitialDynamicPivots");
|
||||
Assign("numberofotherdynamicpivots", "NumberOfOtherDynamicPivots");
|
||||
} else {
|
||||
auto build_cfg = SPTAGParameterMgr::GetInstance().GetBKTParameters();
|
||||
|
||||
Assign("bktnumber", "BKTNumber");
|
||||
Assign("bktkmeansk", "BKTKMeansK");
|
||||
Assign("bktleafsize", "BKTLeafSize");
|
||||
Assign("samples", "Samples");
|
||||
Assign("tptnumber", "TPTNumber");
|
||||
Assign("tptleafsize", "TPTLeafSize");
|
||||
Assign("numtopdimensiontptsplit", "NumTopDimensionTPTSplit");
|
||||
Assign("neighborhoodsize", "NeighborhoodSize");
|
||||
Assign("graphneighborhoodscale", "GraphNeighborhoodScale");
|
||||
Assign("graphcefscale", "GraphCEFScale");
|
||||
Assign("refineiterations", "RefineIterations");
|
||||
Assign("cef", "CEF");
|
||||
Assign("maxcheckforrefinegraph", "MaxCheckForRefineGraph");
|
||||
Assign("numofthreads", "NumberOfThreads");
|
||||
Assign("maxcheck", "MaxCheck");
|
||||
Assign("thresholdofnumberofcontinuousnobetterpropagation", "ThresholdOfNumberOfContinuousNoBetterPropagation");
|
||||
Assign("numberofinitialdynamicpivots", "NumberOfInitialDynamicPivots");
|
||||
Assign("numberofotherdynamicpivots", "NumberOfOtherDynamicPivots");
|
||||
}
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
CPUSPTAGRNG::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
|
||||
SetParameters(config);
|
||||
|
||||
float* p_data = (float*)dataset_ptr->Get<const void*>(meta::TENSOR);
|
||||
for (auto i = 0; i < 10; ++i) {
|
||||
for (auto j = 0; j < 10; ++j) {
|
||||
std::cout << p_data[i * 10 + j] << " ";
|
||||
}
|
||||
std::cout << std::endl;
|
||||
}
|
||||
std::vector<SPTAG::QueryResult> query_results = ConvertToQueryResult(dataset_ptr, config);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (auto i = 0; i < query_results.size(); ++i) {
|
||||
auto target = (float*)query_results[i].GetTarget();
|
||||
std::cout << target[0] << ", " << target[1] << ", " << target[2] << std::endl;
|
||||
index_ptr_->SearchIndex(query_results[i]);
|
||||
}
|
||||
|
||||
return ConvertToDataset(query_results, uids_);
|
||||
}
|
||||
|
||||
int64_t
|
||||
CPUSPTAGRNG::Count() {
|
||||
if (!index_ptr_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_ptr_->GetNumSamples();
|
||||
}
|
||||
|
||||
int64_t
|
||||
CPUSPTAGRNG::Dim() {
|
||||
if (!index_ptr_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
return index_ptr_->GetFeatureDim();
|
||||
}
|
||||
|
||||
void
|
||||
CPUSPTAGRNG::UpdateIndexSize() {
|
||||
if (!index_ptr_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
index_size_ = index_ptr_->GetIndexSize();
|
||||
}
|
||||
|
||||
// void
|
||||
// CPUSPTAGRNG::Add(const DatasetPtr& origin, const Config& add_config) {
|
||||
// SetParameters(add_config);
|
||||
// DatasetPtr dataset = origin->Clone();
|
||||
|
||||
// // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
|
||||
// // && preprocessor_) {
|
||||
// // preprocessor_->Preprocess(dataset);
|
||||
// //}
|
||||
|
||||
// auto vectorset = ConvertToVectorSet(dataset);
|
||||
// auto metaset = ConvertToMetadataSet(dataset);
|
||||
// index_ptr_->AddIndex(vectorset, metaset);
|
||||
// }
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,72 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <SPTAG/AnnService/inc/Core/VectorIndex.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class CPUSPTAGRNG : public VecIndex {
|
||||
public:
|
||||
explicit CPUSPTAGRNG(const std::string& IndexType);
|
||||
|
||||
public:
|
||||
BinarySet
|
||||
Serialize(const Config& config) override;
|
||||
|
||||
void
|
||||
Load(const BinarySet& index_array) override;
|
||||
|
||||
void
|
||||
BuildAll(const DatasetPtr&, const Config&) override;
|
||||
|
||||
void
|
||||
Train(const DatasetPtr& dataset_ptr, const Config& config) override {
|
||||
KNOWHERE_THROW_MSG("SPTAGRNG not support build item dynamically, please invoke BuildAll interface.");
|
||||
}
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override {
|
||||
KNOWHERE_THROW_MSG("Incremental index SPTAGRNG is not supported");
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) override;
|
||||
|
||||
int64_t
|
||||
Count() override;
|
||||
|
||||
int64_t
|
||||
Dim() override;
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
private:
|
||||
void
|
||||
SetParameters(const Config& config);
|
||||
|
||||
private:
|
||||
std::shared_ptr<SPTAG::VectorIndex> index_ptr_;
|
||||
};
|
||||
|
||||
using CPUSPTAGRNGPtr = std::shared_ptr<CPUSPTAGRNG>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,203 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "IndexIVF.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/Statistics.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
int32_t STATISTICS_LEVEL = 0;
|
||||
|
||||
std::string
|
||||
Statistics::ToString() {
|
||||
std::ostringstream ret;
|
||||
|
||||
if (STATISTICS_LEVEL == 0) {
|
||||
ret << "There is nothing because configuration STATISTICS_LEVEL = 0" << std::endl;
|
||||
return ret.str();
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 1) {
|
||||
ret << "Total batches: " << batch_cnt << std::endl;
|
||||
ret << "Total queries: " << nq_cnt << std::endl;
|
||||
ret << "Qps: " << Qps() << std::endl;
|
||||
|
||||
ret << "The frequency distribution of the num of queries:" << std::endl;
|
||||
size_t left = 1, right = 1;
|
||||
for (size_t i = 0; i < NQ_Histogram_Slices - 1; i++) {
|
||||
ret << "[" << left << ", " << right << "].count = " << nq_stat[i] << std::endl;
|
||||
left = right + 1;
|
||||
right <<= 1;
|
||||
}
|
||||
ret << "[" << left << ", +00).count = " << nq_stat.back() << std::endl;
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 2) {
|
||||
ret << "The frequency distribution of filter: " << std::endl;
|
||||
for (auto i = 0; i < 20; ++i) {
|
||||
ret << "[" << i * 5 << "%, " << i * 5 + 5 << "%).count = " << filter_stat[i] << std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
return ret.str();
|
||||
}
|
||||
|
||||
std::string
|
||||
HNSWStatistics::ToString() {
|
||||
std::ostringstream ret;
|
||||
|
||||
if (STATISTICS_LEVEL >= 1) {
|
||||
ret << "Avg Ef: " << AvgSearchEf() << std::endl;
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
std::vector<size_t> axis_x = {5, 10, 20, 40};
|
||||
std::vector<double> access_cdf = AccessCDF(axis_x);
|
||||
ret << "There are " << access_total << " times point-access at level " << target_level << std::endl;
|
||||
ret << "The CDF at level " << target_level << ":" << std::endl;
|
||||
for (auto i = 0; i < axis_x.size(); ++i) {
|
||||
ret << "(" << axis_x[i] << "," << access_cdf[i] << ") ";
|
||||
}
|
||||
ret << std::endl;
|
||||
ret << "Level distribution: " << std::endl;
|
||||
size_t point_cnt = 0;
|
||||
for (int i = distribution.size() - 1; i >= 0; i--) {
|
||||
point_cnt += distribution[i];
|
||||
ret << "Level " << i << " has " << point_cnt << " points" << std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
return Statistics::ToString() + ret.str();
|
||||
}
|
||||
|
||||
std::vector<size_t>
|
||||
GenSplitIndex(size_t size, const std::vector<size_t>& axis_x) {
|
||||
// Gen split index
|
||||
std::vector<size_t> split_idx(axis_x.size());
|
||||
for (size_t i = 0; i < axis_x.size(); i++) {
|
||||
if (axis_x[i] >= 100) {
|
||||
// for safe, not to let idx be larger than size
|
||||
split_idx[i] = size;
|
||||
} else {
|
||||
split_idx[i] = (axis_x[i] * size + 50) / 100;
|
||||
}
|
||||
}
|
||||
return split_idx;
|
||||
}
|
||||
|
||||
std::vector<double>
|
||||
CaculateCDF(size_t access_total, const std::vector<size_t>& access_cnt, const std::vector<size_t>& axis_x) {
|
||||
auto split_idx = GenSplitIndex(access_cnt.size(), axis_x);
|
||||
|
||||
// count cdf
|
||||
std::vector<double> access_cdf;
|
||||
access_cdf.resize(split_idx.size(), 0.0);
|
||||
|
||||
size_t idx = 0;
|
||||
size_t tmp_cnt = 0;
|
||||
for (size_t i = 0; i < split_idx.size(); ++i) {
|
||||
if (i != 0 && split_idx[i] < split_idx[i - 1]) {
|
||||
// wrong split_idx
|
||||
// Todo: log output
|
||||
access_cdf[i] = 0;
|
||||
} else {
|
||||
while (idx < split_idx[i]) {
|
||||
tmp_cnt += access_cnt[idx];
|
||||
idx++;
|
||||
}
|
||||
access_cdf[i] = static_cast<double>(tmp_cnt) / static_cast<double>(access_total);
|
||||
}
|
||||
}
|
||||
|
||||
return access_cdf;
|
||||
}
|
||||
|
||||
std::vector<double>
|
||||
LibHNSWStatistics::AccessCDF(const std::vector<size_t>& axis_x) {
|
||||
// copy from std::map to std::vector
|
||||
std::vector<size_t> access_cnt;
|
||||
access_cnt.reserve(access_cnt_map.size());
|
||||
access_total = 0;
|
||||
for (auto& elem : access_cnt_map) {
|
||||
access_cnt.push_back(elem.second);
|
||||
access_total += elem.second;
|
||||
}
|
||||
std::sort(access_cnt.begin(), access_cnt.end(), std::greater<>());
|
||||
|
||||
return CaculateCDF(access_total, access_cnt, axis_x);
|
||||
}
|
||||
|
||||
std::vector<double>
|
||||
RHNSWStatistics::AccessCDF(const std::vector<size_t>& axis_x) {
|
||||
return CaculateCDF(access_total, access_cnt, axis_x);
|
||||
}
|
||||
|
||||
std::string
|
||||
IVFStatistics::ToString() {
|
||||
std::ostringstream ret;
|
||||
|
||||
if (STATISTICS_LEVEL >= 1) {
|
||||
ret << "nlist " << Nlist() << std::endl;
|
||||
ret << "(nprobe, count): " << std::endl;
|
||||
auto nprobe = SearchNprobe();
|
||||
for (auto& it : nprobe) {
|
||||
ret << "(" << it.first << ", " << it.second << ") ";
|
||||
}
|
||||
ret << std::endl;
|
||||
}
|
||||
if (STATISTICS_LEVEL >= 3) {
|
||||
std::vector<size_t> axis_x = {5, 10, 20, 40};
|
||||
ret << "Bucket CDF " << std::endl;
|
||||
auto output = AccessCDF(axis_x);
|
||||
for (int i = 0; i < output.size(); i++) {
|
||||
ret << "Top " << axis_x[i] << "% access count " << output[i] << std::endl;
|
||||
}
|
||||
ret << std::endl;
|
||||
}
|
||||
return Statistics::ToString() + ret.str();
|
||||
}
|
||||
|
||||
void
|
||||
IVFStatistics::count_nprobe(const int64_t nprobe) {
|
||||
// nprobe count
|
||||
auto it = nprobe_count.find(nprobe);
|
||||
if (it == nprobe_count.end()) {
|
||||
nprobe_count[nprobe] = 1;
|
||||
} else {
|
||||
it->second++;
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IVFStatistics::update_ivf_access_stats(const std::vector<size_t>& nprobe_statistics) {
|
||||
nlist = nprobe_statistics.size();
|
||||
access_total = 0;
|
||||
access_cnt = nprobe_statistics;
|
||||
|
||||
std::sort(access_cnt.begin(), access_cnt.end(), std::greater<>());
|
||||
// access total
|
||||
for (auto& cnt : access_cnt) {
|
||||
access_total += cnt;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<double>
|
||||
IVFStatistics::AccessCDF(const std::vector<size_t>& axis_x) {
|
||||
return CaculateCDF(access_total, access_cnt, axis_x);
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,376 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
extern int32_t STATISTICS_LEVEL;
|
||||
|
||||
inline uint64_t
|
||||
upper_bound_of_pow2(uint64_t x) {
|
||||
--x;
|
||||
x |= (x >> 1);
|
||||
x |= (x >> 2);
|
||||
x |= (x >> 4);
|
||||
x |= (x >> 8);
|
||||
x |= (x >> 16);
|
||||
x |= (x >> 32);
|
||||
return x + 1;
|
||||
}
|
||||
|
||||
inline int
|
||||
len_of_pow2(uint64_t x) {
|
||||
return __builtin_popcountl(x - 1);
|
||||
}
|
||||
|
||||
/*
|
||||
* class: Statistics
|
||||
*/
|
||||
class Statistics {
|
||||
public:
|
||||
static const size_t NQ_Histogram_Slices = 13;
|
||||
static const size_t Filter_Histogram_Slices = 21;
|
||||
|
||||
explicit Statistics(std::string& idx_t)
|
||||
: index_type(idx_t),
|
||||
nq_cnt(0),
|
||||
batch_cnt(0),
|
||||
total_query_time(0.0),
|
||||
nq_stat(NQ_Histogram_Slices, 0),
|
||||
filter_stat(Filter_Histogram_Slices, 0),
|
||||
update_lock() {
|
||||
}
|
||||
|
||||
/*
|
||||
* Get index type
|
||||
* @retval: index type in string
|
||||
*/
|
||||
const std::string&
|
||||
IndexType() {
|
||||
return index_type;
|
||||
}
|
||||
|
||||
/*
|
||||
* To string (may be for log output)
|
||||
* @retval: string output
|
||||
*/
|
||||
virtual std::string
|
||||
ToString();
|
||||
|
||||
virtual ~Statistics() = default;
|
||||
|
||||
/*
|
||||
* Get batch count of the queries (Level 1)
|
||||
* @retval: query batch count
|
||||
*/
|
||||
size_t
|
||||
BatchCount() {
|
||||
return batch_cnt;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the statistics of the nq (Level 1)
|
||||
* @retval: count nq 1, 2, 3~4, 5~8, 9~16,…, 1024~2048, larger than 2048 (13 slices)
|
||||
*/
|
||||
const std::vector<size_t>&
|
||||
NQHistogram() {
|
||||
return nq_stat;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get query response per-second (Level 1)
|
||||
* @retval: Qps
|
||||
*/
|
||||
double
|
||||
Qps() {
|
||||
// ms -> s
|
||||
return total_query_time ? (nq_cnt * 1000.0 / total_query_time) : 0.0;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the statistics of the filter for each batch (Level 2)
|
||||
* @retval: count 0~5%, 5~10%, 10~15%, ...95~100%, 100% (21 slices)
|
||||
*/
|
||||
const std::vector<size_t>&
|
||||
FilterHistograms() {
|
||||
return filter_stat;
|
||||
}
|
||||
|
||||
std::unique_lock<std::mutex>
|
||||
Lock() {
|
||||
return std::unique_lock<std::mutex>(update_lock);
|
||||
}
|
||||
|
||||
public:
|
||||
void
|
||||
update_nq(const int64_t nq) {
|
||||
// batch
|
||||
batch_cnt++;
|
||||
|
||||
// nq_cnt
|
||||
nq_cnt += static_cast<size_t>(nq);
|
||||
|
||||
// nq_stat
|
||||
if (nq > 2048) {
|
||||
nq_stat[12]++;
|
||||
} else {
|
||||
nq_stat[len_of_pow2(upper_bound_of_pow2(static_cast<size_t>(nq)))]++;
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
update_total_query_time(const double query_time) {
|
||||
total_query_time += query_time;
|
||||
}
|
||||
|
||||
void
|
||||
update_filter_percentage(const faiss::BitsetView bitset) {
|
||||
double fps = !bitset.empty() ? static_cast<double>(bitset.count_1()) / bitset.size() : 0.0;
|
||||
filter_stat[static_cast<int>(fps * 100) / 5] += 1;
|
||||
}
|
||||
|
||||
virtual void
|
||||
clear() {
|
||||
total_query_time = 0.0;
|
||||
nq_cnt = 0;
|
||||
batch_cnt = 0;
|
||||
nq_stat.resize(NQ_Histogram_Slices, 0);
|
||||
filter_stat.resize(Filter_Histogram_Slices, 0);
|
||||
}
|
||||
|
||||
public:
|
||||
std::string& index_type;
|
||||
size_t batch_cnt; // updated in query
|
||||
size_t nq_cnt; // updated in query
|
||||
double total_query_time; // updated in query (unit: ms)
|
||||
std::vector<size_t> nq_stat; // updated in query
|
||||
std::vector<size_t> filter_stat; // updated in query
|
||||
std::mutex update_lock;
|
||||
};
|
||||
using StatisticsPtr = std::shared_ptr<Statistics>;
|
||||
|
||||
/*
|
||||
* class: HNSWStatistics
|
||||
*/
|
||||
class HNSWStatistics : public Statistics {
|
||||
public:
|
||||
explicit HNSWStatistics(std::string& idx_t)
|
||||
: Statistics(idx_t), distribution(), target_level(1), access_total(0), ef_sum(0) {
|
||||
}
|
||||
|
||||
~HNSWStatistics() override = default;
|
||||
|
||||
/*
|
||||
* To string (may be for log output)
|
||||
* @retval: string output
|
||||
*/
|
||||
std::string
|
||||
ToString() override;
|
||||
|
||||
/*
|
||||
* Get nodes count in each level
|
||||
* @retval: none
|
||||
*/
|
||||
const std::vector<size_t>&
|
||||
LevelNodesNum() {
|
||||
return distribution;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get average search parameter ‘ef’ (average for batches) (Level 1)
|
||||
* @retval: avg Ef
|
||||
*/
|
||||
double
|
||||
AvgSearchEf() {
|
||||
return nq_cnt ? ef_sum / nq_cnt : 0;
|
||||
}
|
||||
|
||||
/*
|
||||
* Cumulative distribution function of nodes access (Level 3)
|
||||
* @param: none (axis_x = {5,10,15,20,...100} by default)
|
||||
* @retval: Access CDF
|
||||
*/
|
||||
virtual std::vector<double>
|
||||
AccessCDF() {
|
||||
std::vector<size_t> axis_x(20);
|
||||
for (size_t i = 0; i < 20; ++i) {
|
||||
axis_x[i] = (i + 1) * 5;
|
||||
}
|
||||
|
||||
return AccessCDF(axis_x);
|
||||
}
|
||||
|
||||
/*
|
||||
* Cumulative distribution function of nodes access
|
||||
* @param: axis_x[in] specified by users and should be in ascending order
|
||||
* @retval: Access CDF
|
||||
*/
|
||||
virtual std::vector<double>
|
||||
AccessCDF(const std::vector<size_t>& axis_x) = 0;
|
||||
|
||||
public:
|
||||
void
|
||||
update_ef_sum(const int64_t ef) {
|
||||
ef_sum += ef;
|
||||
}
|
||||
|
||||
void
|
||||
update_level_distribution(const int max_level, const std::vector<int>& levels) {
|
||||
distribution.resize(max_level + 1);
|
||||
for (auto i = 0; i <= max_level; ++i) {
|
||||
distribution[i] = levels[i];
|
||||
if (distribution[i] >= 1000 && distribution[i] < 10000) {
|
||||
target_level = i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
clear() override {
|
||||
Statistics::clear();
|
||||
access_total = 0;
|
||||
ef_sum = 0;
|
||||
}
|
||||
|
||||
public:
|
||||
std::vector<size_t> distribution;
|
||||
size_t target_level;
|
||||
size_t access_total; // depend on subclass type
|
||||
size_t ef_sum; // updated in query
|
||||
};
|
||||
|
||||
/*
|
||||
* class: LibHNSWStatistics
|
||||
* for index: HNSW
|
||||
*/
|
||||
class LibHNSWStatistics : public HNSWStatistics {
|
||||
public:
|
||||
explicit LibHNSWStatistics(std::string& idx_t) : HNSWStatistics(idx_t), access_cnt_map() {
|
||||
}
|
||||
|
||||
~LibHNSWStatistics() override = default;
|
||||
|
||||
std::vector<double>
|
||||
AccessCDF(const std::vector<size_t>& axis_x) override;
|
||||
|
||||
public:
|
||||
void
|
||||
clear() override {
|
||||
HNSWStatistics::clear();
|
||||
access_cnt_map.clear();
|
||||
}
|
||||
|
||||
public:
|
||||
std::unordered_map<int64_t, size_t> access_cnt_map; // updated in query
|
||||
};
|
||||
|
||||
/*
|
||||
* class: RHNSWStatistics
|
||||
* for index: RHNSW_FLAT, RHNSE_SQ, RHNSW_PQ
|
||||
*/
|
||||
class RHNSWStatistics : public HNSWStatistics {
|
||||
public:
|
||||
explicit RHNSWStatistics(std::string& idx_t) : HNSWStatistics(idx_t), access_cnt() {
|
||||
}
|
||||
|
||||
~RHNSWStatistics() override = default;
|
||||
|
||||
std::vector<double>
|
||||
AccessCDF(const std::vector<size_t>& axis_x) override;
|
||||
|
||||
public:
|
||||
std::vector<size_t> access_cnt; // prepared in GetStatistics
|
||||
};
|
||||
|
||||
/*
|
||||
* class: IVFStatistics
|
||||
* for index: IVF_FLAT, IVF_PQ, IVF_SQ8
|
||||
*/
|
||||
class IVFStatistics : public Statistics {
|
||||
public:
|
||||
explicit IVFStatistics(std::string& idx_t) : Statistics(idx_t), nprobe_count(), access_cnt(), nlist(0) {
|
||||
}
|
||||
|
||||
~IVFStatistics() override = default;
|
||||
|
||||
/*
|
||||
* To string (may be for log output)
|
||||
* @retval: string output
|
||||
*/
|
||||
std::string
|
||||
ToString() override;
|
||||
|
||||
/*
|
||||
* Get the statistics of the search parameter nprboe (count of batches) (Level 1)
|
||||
* @retval: nprobe
|
||||
*/
|
||||
int64_t
|
||||
Nlist() {
|
||||
return nlist;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the statistics of the search parameter nprboe (count of batches) (Level 1)
|
||||
* @retval: <nprobe, count>
|
||||
*/
|
||||
std::unordered_map<int64_t, size_t>
|
||||
SearchNprobe() {
|
||||
auto lock = Lock();
|
||||
auto rst = nprobe_count;
|
||||
lock.unlock();
|
||||
return rst;
|
||||
}
|
||||
|
||||
/*
|
||||
* Cumulative distribution function of bucket access (Level 3)
|
||||
* @param: axis_x[in] specified by users and should be in ascending order
|
||||
* @retval: Access CDF
|
||||
*/
|
||||
std::vector<double>
|
||||
AccessCDF(const std::vector<size_t>& axis_x);
|
||||
|
||||
public:
|
||||
void
|
||||
count_nprobe(const int64_t nprobe);
|
||||
|
||||
void
|
||||
update_ivf_access_stats(const std::vector<size_t>& nprobe_statistics);
|
||||
|
||||
void
|
||||
clear() override {
|
||||
Statistics::clear();
|
||||
nprobe_count.clear();
|
||||
access_total = 0;
|
||||
}
|
||||
|
||||
public:
|
||||
std::unordered_map<int64_t, size_t> nprobe_count; // updated in query
|
||||
std::vector<size_t> access_cnt; // prepared in GetStatistics
|
||||
size_t access_total; // updated in query
|
||||
size_t nlist;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,152 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <faiss/utils/BitsetView.h>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Dataset.h"
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Typedef.h"
|
||||
#include "knowhere/common/Utils.h"
|
||||
#include "knowhere/index/Index.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
#include "knowhere/index/vector_index/Statistics.h"
|
||||
#include "knowhere/index/vector_index/helpers/DynamicResultSet.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
#define RAW_DATA "RAW_DATA"
|
||||
#define QUANTIZATION_DATA "QUANTIZATION_DATA"
|
||||
|
||||
class VecIndex : public Index {
|
||||
public:
|
||||
virtual void
|
||||
BuildAll(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
Train(dataset_ptr, config);
|
||||
AddWithoutIds(dataset_ptr, config);
|
||||
}
|
||||
|
||||
virtual void
|
||||
Train(const DatasetPtr& dataset, const Config& config) = 0;
|
||||
|
||||
virtual void
|
||||
AddWithoutIds(const DatasetPtr& dataset, const Config& config) = 0;
|
||||
|
||||
virtual DatasetPtr
|
||||
Query(const DatasetPtr& dataset, const Config& config, const faiss::BitsetView bitset) = 0;
|
||||
|
||||
virtual int64_t
|
||||
Dim() = 0;
|
||||
|
||||
virtual int64_t
|
||||
Count() = 0;
|
||||
|
||||
virtual StatisticsPtr
|
||||
GetStatistics() {
|
||||
return stats;
|
||||
}
|
||||
|
||||
virtual void
|
||||
ClearStatistics() {
|
||||
}
|
||||
|
||||
virtual IndexType
|
||||
index_type() const {
|
||||
return index_type_;
|
||||
}
|
||||
|
||||
virtual IndexMode
|
||||
index_mode() const {
|
||||
return index_mode_;
|
||||
}
|
||||
|
||||
std::shared_ptr<std::vector<IDType>>
|
||||
GetUids() const {
|
||||
return uids_;
|
||||
}
|
||||
|
||||
void
|
||||
SetUids(std::shared_ptr<std::vector<IDType>> uids) {
|
||||
uids_ = uids;
|
||||
}
|
||||
|
||||
void
|
||||
MapOffsetToUid(IDType* id, size_t n) {
|
||||
if (uids_) {
|
||||
for (size_t i = 0; i < n; i++) {
|
||||
if (id[i] >= 0) {
|
||||
id[i] = uids_->at(id[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
MapUids(DynamicResultSegment& milvus_dataset) {
|
||||
if (uids_) {
|
||||
for (auto& mrspr : milvus_dataset) {
|
||||
for (auto j = 0; j < mrspr->buffers.size(); ++j) {
|
||||
auto buf = mrspr->buffers[j];
|
||||
auto len = j + 1 == mrspr->buffers.size() ? mrspr->wp : mrspr->buffer_size;
|
||||
for (auto i = 0; i < len; ++i) {
|
||||
if (buf.ids[i] >= 0) {
|
||||
buf.ids[i] = uids_->at(buf.ids[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
size_t
|
||||
UidsSize() {
|
||||
return uids_ ? uids_->size() * sizeof(IDType) : 0;
|
||||
}
|
||||
|
||||
virtual int64_t
|
||||
IndexSize() {
|
||||
if (index_size_ == -1) {
|
||||
KNOWHERE_THROW_MSG("Index size not set");
|
||||
}
|
||||
return index_size_;
|
||||
}
|
||||
|
||||
void
|
||||
SetIndexSize(int64_t size) {
|
||||
index_size_ = size;
|
||||
}
|
||||
|
||||
virtual void
|
||||
UpdateIndexSize() {
|
||||
}
|
||||
|
||||
int64_t
|
||||
Size() override {
|
||||
return UidsSize() + IndexSize();
|
||||
}
|
||||
|
||||
protected:
|
||||
IndexType index_type_ = "";
|
||||
IndexMode index_mode_ = IndexMode::MODE_CPU;
|
||||
std::shared_ptr<std::vector<IDType>> uids_ = nullptr;
|
||||
int64_t index_size_ = -1;
|
||||
StatisticsPtr stats = nullptr;
|
||||
};
|
||||
|
||||
using VecIndexPtr = std::shared_ptr<VecIndex>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,114 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include "knowhere/index/vector_index/VecIndexFactory.h"
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/common/Log.h"
|
||||
#include "knowhere/index/vector_index/IndexAnnoy.h"
|
||||
#include "knowhere/index/vector_index/IndexBinaryIDMAP.h"
|
||||
#include "knowhere/index/vector_index/IndexBinaryIVF.h"
|
||||
#include "knowhere/index/vector_index/IndexHNSW.h"
|
||||
#include "knowhere/index/vector_index/IndexIDMAP.h"
|
||||
#include "knowhere/index/vector_index/IndexIVF.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFPQ.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFSQ.h"
|
||||
#include "knowhere/index/vector_index/IndexNGTONNG.h"
|
||||
#include "knowhere/index/vector_index/IndexNGTPANNG.h"
|
||||
#include "knowhere/index/vector_index/IndexRHNSWFlat.h"
|
||||
#include "knowhere/index/vector_index/IndexRHNSWPQ.h"
|
||||
#include "knowhere/index/vector_index/IndexRHNSWSQ.h"
|
||||
#include "knowhere/index/vector_offset_index/IndexIVF_NM.h"
|
||||
#include "knowhere/index/vector_offset_index/IndexNSG_NM.h"
|
||||
|
||||
#ifdef MILVUS_SUPPORT_SPTAG
|
||||
#include "knowhere/index/vector_index/IndexSPTAG.h"
|
||||
#endif
|
||||
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include <cuda.h>
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h"
|
||||
#include "knowhere/index/vector_index/helpers/Cloner.h"
|
||||
#include "knowhere/index/vector_offset_index/gpu/IndexGPUIVF_NM.h"
|
||||
#endif
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
VecIndexPtr
|
||||
VecIndexFactory::CreateVecIndex(const IndexType& type, const IndexMode mode) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
auto gpu_device = -1; // TODO: remove hardcode here, get from invoker
|
||||
#endif
|
||||
if (type == IndexEnum::INDEX_FAISS_IDMAP) {
|
||||
return std::make_shared<knowhere::IDMAP>();
|
||||
} else if (type == IndexEnum::INDEX_FAISS_IVFFLAT) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
if (mode == IndexMode::MODE_GPU) {
|
||||
return std::make_shared<knowhere::GPUIVF_NM>(gpu_device);
|
||||
}
|
||||
#endif
|
||||
return std::make_shared<knowhere::IVF_NM>();
|
||||
} else if (type == IndexEnum::INDEX_FAISS_IVFPQ) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
if (mode == IndexMode::MODE_GPU) {
|
||||
return std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
|
||||
}
|
||||
#endif
|
||||
return std::make_shared<knowhere::IVFPQ>();
|
||||
} else if (type == IndexEnum::INDEX_FAISS_IVFSQ8) {
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
if (mode == IndexMode::MODE_GPU) {
|
||||
return std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
|
||||
}
|
||||
#endif
|
||||
return std::make_shared<knowhere::IVFSQ>();
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
} else if (type == IndexEnum::INDEX_FAISS_IVFSQ8H) {
|
||||
return std::make_shared<knowhere::IVFSQHybrid>(gpu_device);
|
||||
#endif
|
||||
} else if (type == IndexEnum::INDEX_FAISS_BIN_IDMAP) {
|
||||
return std::make_shared<knowhere::BinaryIDMAP>();
|
||||
} else if (type == IndexEnum::INDEX_FAISS_BIN_IVFFLAT) {
|
||||
return std::make_shared<knowhere::BinaryIVF>();
|
||||
} else if (type == IndexEnum::INDEX_NSG) {
|
||||
return std::make_shared<knowhere::NSG_NM>(-1);
|
||||
#ifdef MILVUS_SUPPORT_SPTAG
|
||||
} else if (type == IndexEnum::INDEX_SPTAG_KDT_RNT) {
|
||||
return std::make_shared<knowhere::CPUSPTAGRNG>("KDT");
|
||||
} else if (type == IndexEnum::INDEX_SPTAG_BKT_RNT) {
|
||||
return std::make_shared<knowhere::CPUSPTAGRNG>("BKT");
|
||||
#endif
|
||||
} else if (type == IndexEnum::INDEX_HNSW) {
|
||||
return std::make_shared<knowhere::IndexHNSW>();
|
||||
} else if (type == IndexEnum::INDEX_ANNOY) {
|
||||
return std::make_shared<knowhere::IndexAnnoy>();
|
||||
} else if (type == IndexEnum::INDEX_RHNSWFlat) {
|
||||
return std::make_shared<knowhere::IndexRHNSWFlat>();
|
||||
} else if (type == IndexEnum::INDEX_RHNSWPQ) {
|
||||
return std::make_shared<knowhere::IndexRHNSWPQ>();
|
||||
} else if (type == IndexEnum::INDEX_RHNSWSQ) {
|
||||
return std::make_shared<knowhere::IndexRHNSWSQ>();
|
||||
} else if (type == IndexEnum::INDEX_NGTPANNG) {
|
||||
return std::make_shared<knowhere::IndexNGTPANNG>();
|
||||
} else if (type == IndexEnum::INDEX_NGTONNG) {
|
||||
return std::make_shared<knowhere::IndexNGTONNG>();
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,41 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "knowhere/index/IndexType.h"
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class VecIndexFactory {
|
||||
private:
|
||||
VecIndexFactory() = default;
|
||||
VecIndexFactory(const VecIndexFactory&) = delete;
|
||||
VecIndexFactory
|
||||
operator=(const VecIndexFactory&) = delete;
|
||||
|
||||
public:
|
||||
static VecIndexFactory&
|
||||
GetInstance() {
|
||||
static VecIndexFactory inst;
|
||||
return inst;
|
||||
}
|
||||
|
||||
knowhere::VecIndexPtr
|
||||
CreateVecIndex(const IndexType& type, const IndexMode mode = IndexMode::MODE_CPU);
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,92 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include "knowhere/index/vector_index/adapter/SptagAdapter.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
std::shared_ptr<SPTAG::MetadataSet>
|
||||
ConvertToMetadataSet(const DatasetPtr& dataset_ptr) {
|
||||
auto elems = dataset_ptr->Get<int64_t>(meta::ROWS);
|
||||
|
||||
auto p_id = new int64_t[elems];
|
||||
for (int64_t i = 0; i < elems; ++i) p_id[i] = i;
|
||||
|
||||
auto p_offset = new int64_t[elems + 1];
|
||||
for (int64_t i = 0; i <= elems; ++i) p_offset[i] = i * 8;
|
||||
|
||||
std::shared_ptr<SPTAG::MetadataSet> metaset(
|
||||
new SPTAG::MemMetadataSet(SPTAG::ByteArray((std::uint8_t*)p_id, elems * sizeof(int64_t), true),
|
||||
SPTAG::ByteArray((std::uint8_t*)p_offset, elems * sizeof(int64_t), true), elems));
|
||||
|
||||
return metaset;
|
||||
}
|
||||
|
||||
std::shared_ptr<SPTAG::VectorSet>
|
||||
ConvertToVectorSet(const DatasetPtr& dataset_ptr) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
size_t num_bytes = rows * dim * sizeof(float);
|
||||
SPTAG::ByteArray byte_array((uint8_t*)p_data, num_bytes, false);
|
||||
|
||||
auto vectorset = std::make_shared<SPTAG::BasicVectorSet>(byte_array, SPTAG::VectorValueType::Float, dim, rows);
|
||||
return vectorset;
|
||||
}
|
||||
|
||||
std::vector<SPTAG::QueryResult>
|
||||
ConvertToQueryResult(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr);
|
||||
|
||||
int64_t k = config[meta::TOPK].get<int64_t>();
|
||||
std::vector<SPTAG::QueryResult> query_results(rows, SPTAG::QueryResult(nullptr, k, true));
|
||||
for (auto i = 0; i < rows; ++i) {
|
||||
query_results[i].SetTarget((float*)p_data + i * dim);
|
||||
}
|
||||
|
||||
return query_results;
|
||||
}
|
||||
|
||||
DatasetPtr
|
||||
ConvertToDataset(std::vector<SPTAG::QueryResult> query_results, std::shared_ptr<std::vector<int64_t>> uid) {
|
||||
auto k = query_results[0].GetResultNum();
|
||||
auto elems = query_results.size() * k;
|
||||
|
||||
size_t p_id_size = sizeof(int64_t) * elems;
|
||||
size_t p_dist_size = sizeof(float) * elems;
|
||||
auto p_id = (int64_t*)malloc(p_id_size);
|
||||
auto p_dist = (float*)malloc(p_dist_size);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (size_t i = 0; i < query_results.size(); ++i) {
|
||||
auto results = query_results[i].GetResults();
|
||||
auto num_result = query_results[i].GetResultNum();
|
||||
for (auto j = 0; j < num_result; ++j) {
|
||||
// p_id[i * k + j] = results[j].VID;
|
||||
auto id = *(int64_t*)query_results[i].GetMetadata(j).Data();
|
||||
if (uid != nullptr) {
|
||||
if (id >= 0) {
|
||||
id = uid->at(id);
|
||||
}
|
||||
}
|
||||
p_id[i * k + j] = id;
|
||||
p_dist[i * k + j] = results[j].Dist;
|
||||
}
|
||||
}
|
||||
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::IDS, p_id);
|
||||
ret_ds->Set(meta::DISTANCE, p_dist);
|
||||
return ret_ds;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,37 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <SPTAG/AnnService/inc/Core/VectorIndex.h>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/common/Config.h"
|
||||
#include "knowhere/common/Dataset.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
std::shared_ptr<SPTAG::VectorSet>
|
||||
ConvertToVectorSet(const DatasetPtr& dataset_ptr);
|
||||
|
||||
std::shared_ptr<SPTAG::MetadataSet>
|
||||
ConvertToMetadataSet(const DatasetPtr& dataset_ptr);
|
||||
|
||||
std::vector<SPTAG::QueryResult>
|
||||
ConvertToQueryResult(const DatasetPtr& dataset_ptr, const Config& config);
|
||||
|
||||
DatasetPtr
|
||||
ConvertToDataset(std::vector<SPTAG::QueryResult> query_results, std::shared_ptr<std::vector<int64_t>> uid);
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,31 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "knowhere/common/Dataset.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
DatasetPtr
|
||||
GenDataset(const int64_t nb, const int64_t dim, const void* xb) {
|
||||
auto ret_ds = std::make_shared<Dataset>();
|
||||
ret_ds->Set(meta::ROWS, nb);
|
||||
ret_ds->Set(meta::DIM, dim);
|
||||
ret_ds->Set(meta::TENSOR, xb);
|
||||
return ret_ds;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,33 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include "knowhere/common/Dataset.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
#define GET_TENSOR_DATA(dataset_ptr) \
|
||||
int64_t rows = dataset_ptr->Get<int64_t>(meta::ROWS); \
|
||||
const void* p_data = dataset_ptr->Get<const void*>(meta::TENSOR);
|
||||
|
||||
#define GET_TENSOR_DATA_DIM(dataset_ptr) \
|
||||
GET_TENSOR_DATA(dataset_ptr) \
|
||||
int64_t dim = dataset_ptr->Get<int64_t>(meta::DIM);
|
||||
|
||||
extern DatasetPtr
|
||||
GenDataset(const int64_t nb, const int64_t dim, const void* xb);
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,50 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class GPUIndex {
|
||||
public:
|
||||
explicit GPUIndex(const int& device_id) : gpu_id_(device_id) {
|
||||
}
|
||||
|
||||
GPUIndex(const int& device_id, const ResPtr& resource) : gpu_id_(device_id), res_(resource) {
|
||||
}
|
||||
|
||||
virtual VecIndexPtr
|
||||
CopyGpuToCpu(const Config&) = 0;
|
||||
|
||||
virtual VecIndexPtr
|
||||
CopyGpuToGpu(const int64_t, const Config&) = 0;
|
||||
|
||||
void
|
||||
SetGpuDevice(const int& gpu_id) {
|
||||
gpu_id_ = gpu_id;
|
||||
}
|
||||
|
||||
const int64_t
|
||||
GetGpuDevice() {
|
||||
return gpu_id_;
|
||||
}
|
||||
|
||||
protected:
|
||||
int64_t gpu_id_;
|
||||
ResWPtr res_;
|
||||
};
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,149 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <faiss/AutoTune.h>
|
||||
#include <faiss/IndexFlat.h>
|
||||
#include <faiss/MetaIndexes.h>
|
||||
#include <faiss/index_io.h>
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#endif
|
||||
#include <fiu/fiu-local.h>
|
||||
#include <string>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/IndexType.h"
|
||||
#include "knowhere/index/vector_index/IndexIDMAP.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
VecIndexPtr
|
||||
GPUIDMAP::CopyGpuToCpu(const Config& config) {
|
||||
faiss::Index* device_index = index_.get();
|
||||
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
|
||||
|
||||
std::shared_ptr<faiss::Index> new_index;
|
||||
new_index.reset(host_index);
|
||||
return std::make_shared<IDMAP>(new_index);
|
||||
}
|
||||
|
||||
BinarySet
|
||||
GPUIDMAP::SerializeImpl(const IndexType& type) {
|
||||
try {
|
||||
fiu_do_on("GPUIDMP.SerializeImpl.throw_exception", throw std::exception());
|
||||
MemoryIOWriter writer;
|
||||
{
|
||||
faiss::Index* index = index_.get();
|
||||
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(index);
|
||||
|
||||
faiss::write_index(host_index, &writer);
|
||||
delete host_index;
|
||||
}
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("IVF", data, writer.rp);
|
||||
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
GPUIDMAP::LoadImpl(const BinarySet& index_binary, const IndexType& type) {
|
||||
auto binary = index_binary.GetByName("IVF");
|
||||
MemoryIOReader reader;
|
||||
{
|
||||
reader.total = binary->size;
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
faiss::Index* index = faiss::read_index(&reader);
|
||||
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) {
|
||||
ResScope rs(res, gpu_id_, false);
|
||||
auto device_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index);
|
||||
index_.reset(device_index);
|
||||
res_ = res;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Load error, can't get gpu resource");
|
||||
}
|
||||
|
||||
delete index;
|
||||
}
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
GPUIDMAP::CopyGpuToGpu(const int64_t device_id, const Config& config) {
|
||||
auto cpu_index = CopyGpuToCpu(config);
|
||||
return std::static_pointer_cast<IDMAP>(cpu_index)->CopyCpuToGpu(device_id, config);
|
||||
}
|
||||
|
||||
const float*
|
||||
GPUIDMAP::GetRawVectors() {
|
||||
KNOWHERE_THROW_MSG("Not support");
|
||||
}
|
||||
|
||||
void
|
||||
GPUIDMAP::QueryImpl(int64_t n,
|
||||
const float* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
ResScope rs(res_, gpu_id_);
|
||||
|
||||
// assign the metric type
|
||||
index_->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
index_->search(n, data, k, distances, labels, bitset);
|
||||
}
|
||||
|
||||
void
|
||||
GPUIDMAP::GenGraph(const float* data, const int64_t k, GraphType& graph, const Config& config) {
|
||||
int64_t K = k + 1;
|
||||
auto ntotal = Count();
|
||||
|
||||
size_t dim = config[meta::DIM];
|
||||
auto batch_size = 1000;
|
||||
auto tail_batch_size = ntotal % batch_size;
|
||||
auto batch_search_count = ntotal / batch_size;
|
||||
auto total_search_count = tail_batch_size == 0 ? batch_search_count : batch_search_count + 1;
|
||||
|
||||
std::vector<float> res_dis(K * batch_size);
|
||||
graph.resize(ntotal);
|
||||
Graph res_vec(total_search_count);
|
||||
for (int i = 0; i < total_search_count; ++i) {
|
||||
auto b_size = (i == (total_search_count - 1)) && tail_batch_size != 0 ? tail_batch_size : batch_size;
|
||||
|
||||
auto& res = res_vec[i];
|
||||
res.resize(K * b_size);
|
||||
|
||||
const float* xq = data + batch_size * dim * i;
|
||||
QueryImpl(b_size, xq, K, res_dis.data(), res.data(), config, nullptr);
|
||||
|
||||
for (int j = 0; j < b_size; ++j) {
|
||||
auto& node = graph[batch_size * i + j];
|
||||
node.resize(k);
|
||||
auto start_pos = j * K + 1;
|
||||
for (int m = 0, cursor = start_pos; m < k && cursor < start_pos + k; ++m, ++cursor) {
|
||||
node[m] = res[cursor];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,61 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "knowhere/index/vector_index/IndexIDMAP.h"
|
||||
#include "knowhere/index/vector_index/gpu/GPUIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
using Graph = std::vector<std::vector<int64_t>>;
|
||||
|
||||
class GPUIDMAP : public IDMAP, public GPUIndex {
|
||||
public:
|
||||
explicit GPUIDMAP(std::shared_ptr<faiss::Index> index, const int64_t device_id, ResPtr& res)
|
||||
: IDMAP(std::move(index)), GPUIndex(device_id, res) {
|
||||
index_mode_ = IndexMode::MODE_GPU;
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
CopyGpuToCpu(const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyGpuToGpu(const int64_t, const Config&) override;
|
||||
|
||||
const float*
|
||||
GetRawVectors() override;
|
||||
|
||||
void
|
||||
GenGraph(const float*, const int64_t, GraphType&, const Config&);
|
||||
|
||||
virtual ~GPUIDMAP() = default;
|
||||
|
||||
protected:
|
||||
BinarySet
|
||||
SerializeImpl(const IndexType&) override;
|
||||
|
||||
void
|
||||
LoadImpl(const BinarySet&, const IndexType&) override;
|
||||
|
||||
void
|
||||
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::BitsetView) override;
|
||||
};
|
||||
|
||||
using GPUIDMAPPtr = std::shared_ptr<GPUIDMAP>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,162 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#include <faiss/gpu/GpuIndexIVF.h>
|
||||
#include <faiss/gpu/GpuIndexIVFFlat.h>
|
||||
#include <faiss/index_io.h>
|
||||
#include <fiu/fiu-local.h>
|
||||
#include <string>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h"
|
||||
#include "knowhere/index/vector_index/helpers/Cloner.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
void
|
||||
GPUIVF::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
gpu_id_ = config[knowhere::meta::DEVICEID];
|
||||
|
||||
auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
|
||||
if (gpu_res != nullptr) {
|
||||
ResScope rs(gpu_res, gpu_id_, true);
|
||||
faiss::gpu::GpuIndexIVFFlatConfig idx_config;
|
||||
idx_config.device = static_cast<int32_t>(gpu_id_);
|
||||
int32_t nlist = config[IndexParams::nlist];
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
index_ = std::make_shared<faiss::gpu::GpuIndexIVFFlat>(gpu_res->faiss_res.get(), dim, nlist, metric_type,
|
||||
idx_config);
|
||||
index_->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
res_ = gpu_res;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Build IVF can't get gpu resource");
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
GPUIVF::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
if (auto spt = res_.lock()) {
|
||||
ResScope rs(res_, gpu_id_);
|
||||
IVF::AddWithoutIds(dataset_ptr, config);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Add IVF can't get gpu resource");
|
||||
}
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
GPUIVF::CopyGpuToCpu(const Config& config) {
|
||||
auto device_idx = std::dynamic_pointer_cast<faiss::gpu::GpuIndexIVF>(index_);
|
||||
if (device_idx != nullptr) {
|
||||
faiss::Index* device_index = index_.get();
|
||||
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
|
||||
|
||||
std::shared_ptr<faiss::Index> new_index;
|
||||
new_index.reset(host_index);
|
||||
return std::make_shared<IVF>(new_index);
|
||||
} else {
|
||||
return std::make_shared<IVF>(index_);
|
||||
}
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
GPUIVF::CopyGpuToGpu(const int64_t device_id, const Config& config) {
|
||||
auto host_index = CopyGpuToCpu(config);
|
||||
return std::static_pointer_cast<IVF>(host_index)->CopyCpuToGpu(device_id, config);
|
||||
}
|
||||
|
||||
BinarySet
|
||||
GPUIVF::SerializeImpl(const IndexType& type) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
try {
|
||||
fiu_do_on("GPUIVF.SerializeImpl.throw_exception", throw std::exception());
|
||||
MemoryIOWriter writer;
|
||||
{
|
||||
faiss::Index* index = index_.get();
|
||||
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(index);
|
||||
|
||||
faiss::write_index(host_index, &writer);
|
||||
delete host_index;
|
||||
}
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("IVF", data, writer.rp);
|
||||
|
||||
return res_set;
|
||||
} catch (std::exception& e) {
|
||||
KNOWHERE_THROW_MSG(e.what());
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
GPUIVF::LoadImpl(const BinarySet& binary_set, const IndexType& type) {
|
||||
auto binary = binary_set.GetByName("IVF");
|
||||
MemoryIOReader reader;
|
||||
{
|
||||
reader.total = binary->size;
|
||||
reader.data_ = binary->data.get();
|
||||
|
||||
faiss::Index* index = faiss::read_index(&reader);
|
||||
|
||||
if (auto temp_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) {
|
||||
ResScope rs(temp_res, gpu_id_, false);
|
||||
auto device_index = faiss::gpu::index_cpu_to_gpu(temp_res->faiss_res.get(), gpu_id_, index);
|
||||
index_.reset(device_index);
|
||||
res_ = temp_res;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Load error, can't get gpu resource");
|
||||
}
|
||||
|
||||
delete index;
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
GPUIVF::QueryImpl(int64_t n,
|
||||
const float* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
auto device_index = std::dynamic_pointer_cast<faiss::gpu::GpuIndexIVF>(index_);
|
||||
fiu_do_on("GPUIVF.search_impl.invald_index", device_index = nullptr);
|
||||
if (device_index) {
|
||||
device_index->nprobe = std::min(static_cast<int>(config[IndexParams::nprobe]), device_index->nlist);
|
||||
ResScope rs(res_, gpu_id_);
|
||||
|
||||
// if query size > 2048 we search by blocks to avoid malloc issue
|
||||
const int64_t block_size = 2048;
|
||||
int64_t dim = device_index->d;
|
||||
for (int64_t i = 0; i < n; i += block_size) {
|
||||
int64_t search_size = (n - i > block_size) ? block_size : (n - i);
|
||||
device_index->search(search_size, reinterpret_cast<const float*>(data) + i * dim, k, distances + i * k,
|
||||
labels + i * k, bitset);
|
||||
}
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Not a GpuIndexIVF type.");
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,60 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/index/vector_index/IndexIVF.h"
|
||||
#include "knowhere/index/vector_index/gpu/GPUIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class GPUIVF : public IVF, public GPUIndex {
|
||||
public:
|
||||
explicit GPUIVF(const int& device_id) : IVF(), GPUIndex(device_id) {
|
||||
index_mode_ = IndexMode::MODE_GPU;
|
||||
}
|
||||
|
||||
explicit GPUIVF(std::shared_ptr<faiss::Index> index, const int64_t device_id, ResPtr& res)
|
||||
: IVF(std::move(index)), GPUIndex(device_id, res) {
|
||||
index_mode_ = IndexMode::MODE_GPU;
|
||||
}
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
void
|
||||
AddWithoutIds(const DatasetPtr&, const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyGpuToCpu(const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyGpuToGpu(const int64_t, const Config&) override;
|
||||
|
||||
protected:
|
||||
BinarySet
|
||||
SerializeImpl(const IndexType&) override;
|
||||
|
||||
void
|
||||
LoadImpl(const BinarySet&, const IndexType&) override;
|
||||
|
||||
void
|
||||
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::BitsetView) override;
|
||||
};
|
||||
|
||||
using GPUIVFPtr = std::shared_ptr<GPUIVF>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,71 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <string>
|
||||
|
||||
#include <faiss/IndexIVFPQ.h>
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#include <faiss/gpu/GpuIndexIVFPQ.h>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFPQ.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
void
|
||||
GPUIVFPQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
gpu_id_ = config[knowhere::meta::DEVICEID];
|
||||
auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
|
||||
if (gpu_res != nullptr) {
|
||||
ResScope rs(gpu_res, gpu_id_, true);
|
||||
faiss::gpu::GpuIndexIVFPQConfig idx_config;
|
||||
idx_config.device = static_cast<int32_t>(gpu_id_);
|
||||
int32_t nlist = config[IndexParams::nlist];
|
||||
int32_t m = config[IndexParams::m];
|
||||
int32_t nbits = config[IndexParams::nbits];
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
index_ = std::make_shared<faiss::gpu::GpuIndexIVFPQ>(gpu_res->faiss_res.get(), dim, nlist, m, nbits,
|
||||
metric_type, idx_config);
|
||||
device_index->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
res_ = gpu_res;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Build IVFPQ can't get gpu resource");
|
||||
}
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
GPUIVFPQ::CopyGpuToCpu(const Config& config) {
|
||||
faiss::Index* device_index = index_.get();
|
||||
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
|
||||
|
||||
std::shared_ptr<faiss::Index> new_index;
|
||||
new_index.reset(host_index);
|
||||
return std::make_shared<IVFPQ>(new_index);
|
||||
}
|
||||
|
||||
std::shared_ptr<faiss::IVFSearchParameters>
|
||||
GPUIVFPQ::GenParams(const Config& config) {
|
||||
auto params = std::make_shared<faiss::IVFPQSearchParameters>();
|
||||
params->nprobe = config[IndexParams::nprobe];
|
||||
// params->scan_table_threshold = config["scan_table_threhold"]
|
||||
// params->polysemous_ht = config["polysemous_ht"]
|
||||
// params->max_codes = config["max_codes"]
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,49 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class GPUIVFPQ : public GPUIVF {
|
||||
public:
|
||||
explicit GPUIVFPQ(const int& device_id) : GPUIVF(device_id) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFPQ;
|
||||
}
|
||||
|
||||
GPUIVFPQ(std::shared_ptr<faiss::Index> index, const int64_t device_id, ResPtr& res)
|
||||
: GPUIVF(std::move(index), device_id, res) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFPQ;
|
||||
}
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyGpuToCpu(const Config&) override;
|
||||
|
||||
virtual ~GPUIVFPQ() = default;
|
||||
|
||||
protected:
|
||||
std::shared_ptr<faiss::IVFSearchParameters>
|
||||
GenParams(const Config& config) override;
|
||||
};
|
||||
|
||||
using GPUIVFPQPtr = std::shared_ptr<GPUIVFPQ>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,60 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <faiss/IndexFlat.h>
|
||||
#include <faiss/IndexScalarQuantizer.h>
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFSQ.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
void
|
||||
GPUIVFSQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
gpu_id_ = config[knowhere::meta::DEVICEID];
|
||||
auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
|
||||
if (gpu_res != nullptr) {
|
||||
ResScope rs(gpu_res, gpu_id_, true);
|
||||
faiss::gpu::GpuIndexIVFScalarQuantizerConfig idx_config;
|
||||
idx_config.device = static_cast<int32_t>(gpu_id_);
|
||||
int32_t nlist = config[IndexParams::nlist];
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
index_ = std::make_shared<faiss::gpu::GpuIndexIVFScalarQuantizer>(
|
||||
gpu_res->faiss_res.get(), dim, nlist, faiss::QuantizerType::QT_8bit, metric_type, true, idx_config);
|
||||
index_->train(rows, (float*)p_data);
|
||||
res_ = gpu_res;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource");
|
||||
}
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
GPUIVFSQ::CopyGpuToCpu(const Config& config) {
|
||||
faiss::Index* device_index = index_.get();
|
||||
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
|
||||
|
||||
std::shared_ptr<faiss::Index> new_index;
|
||||
new_index.reset(host_index);
|
||||
return std::make_shared<IVFSQ>(new_index);
|
||||
}
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,45 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
class GPUIVFSQ : public GPUIVF {
|
||||
public:
|
||||
explicit GPUIVFSQ(const int& device_id) : GPUIVF(device_id) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8;
|
||||
}
|
||||
|
||||
explicit GPUIVFSQ(std::shared_ptr<faiss::Index> index, const int64_t device_id, ResPtr& res)
|
||||
: GPUIVF(std::move(index), device_id, res) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8;
|
||||
}
|
||||
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyGpuToCpu(const Config&) override;
|
||||
|
||||
virtual ~GPUIVFSQ() = default;
|
||||
};
|
||||
|
||||
using GPUIVFSQPtr = std::shared_ptr<GPUIVFSQ>;
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,284 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <faiss/IndexSQHybrid.h>
|
||||
#include <faiss/gpu/GpuCloner.h>
|
||||
#include <faiss/gpu/GpuIndexIVF.h>
|
||||
#include <faiss/index_factory.h>
|
||||
#include <fiu/fiu-local.h>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h"
|
||||
#include "knowhere/index/vector_index/helpers/FaissIO.h"
|
||||
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
|
||||
void
|
||||
IVFSQHybrid::Train(const DatasetPtr& dataset_ptr, const Config& config) {
|
||||
GET_TENSOR_DATA_DIM(dataset_ptr)
|
||||
gpu_id_ = config[knowhere::meta::DEVICEID];
|
||||
|
||||
auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
|
||||
if (gpu_res != nullptr) {
|
||||
ResScope rs(gpu_res, gpu_id_, true);
|
||||
faiss::gpu::GpuIndexIVFSQHybridConfig idx_config;
|
||||
idx_config.device = static_cast<int32_t>(gpu_id_);
|
||||
int32_t nlist = config[IndexParams::nlist];
|
||||
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
|
||||
index_ = std::make_shared<faiss::gpu::GpuIndexIVFSQHybrid>(
|
||||
gpu_res->faiss_res.get(), dim, nlist, faiss::QuantizerType::QT_8bit, metric_type, true, idx_config);
|
||||
index_->train(rows, reinterpret_cast<const float*>(p_data));
|
||||
res_ = gpu_res;
|
||||
gpu_mode_ = 2;
|
||||
index_mode_ = IndexMode::MODE_GPU;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Build IVFSQHybrid can't get gpu resource");
|
||||
}
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
IVFSQHybrid::CopyGpuToCpu(const Config& config) {
|
||||
if (gpu_mode_ == 0) {
|
||||
return std::make_shared<IVFSQHybrid>(index_);
|
||||
}
|
||||
|
||||
faiss::Index* device_index = index_.get();
|
||||
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
|
||||
|
||||
if (auto* ivf_index = dynamic_cast<faiss::IndexIVF*>(host_index)) {
|
||||
ivf_index->to_readonly();
|
||||
ivf_index->backup_quantizer();
|
||||
}
|
||||
|
||||
std::shared_ptr<faiss::Index> new_index;
|
||||
new_index.reset(host_index);
|
||||
return std::make_shared<IVFSQHybrid>(new_index);
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
IVFSQHybrid::CopyCpuToGpu(const int64_t device_id, const Config& config) {
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
|
||||
ResScope rs(res, device_id, false);
|
||||
faiss::gpu::GpuClonerOptions option;
|
||||
option.allInGpu = true;
|
||||
|
||||
auto idx = dynamic_cast<faiss::IndexIVF*>(index_.get());
|
||||
idx->restore_quantizer();
|
||||
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get(), &option);
|
||||
std::shared_ptr<faiss::Index> device_index = std::shared_ptr<faiss::Index>(gpu_index);
|
||||
auto new_idx = std::make_shared<IVFSQHybrid>(device_index, device_id, res);
|
||||
return new_idx;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu: " + std::to_string(gpu_id_) + "resource");
|
||||
}
|
||||
}
|
||||
|
||||
std::pair<VecIndexPtr, FaissIVFQuantizerPtr>
|
||||
IVFSQHybrid::CopyCpuToGpuWithQuantizer(const int64_t device_id, const Config& config) {
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
|
||||
ResScope rs(res, device_id, false);
|
||||
faiss::gpu::GpuClonerOptions option;
|
||||
option.allInGpu = true;
|
||||
|
||||
faiss::IndexComposition index_composition;
|
||||
index_composition.index = index_.get();
|
||||
index_composition.quantizer = nullptr;
|
||||
index_composition.mode = 0; // copy all
|
||||
|
||||
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, &index_composition, &option);
|
||||
|
||||
std::shared_ptr<faiss::Index> device_index;
|
||||
device_index.reset(gpu_index);
|
||||
auto new_idx = std::make_shared<IVFSQHybrid>(device_index, device_id, res);
|
||||
|
||||
auto q = std::make_shared<FaissIVFQuantizer>();
|
||||
q->quantizer = index_composition.quantizer;
|
||||
q->size = index_composition.quantizer->d * index_composition.quantizer->getNumVecs() * sizeof(float);
|
||||
q->gpu_id = device_id;
|
||||
return std::make_pair(new_idx, q);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu: " + std::to_string(gpu_id_) + "resource");
|
||||
}
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
IVFSQHybrid::LoadData(const FaissIVFQuantizerPtr& quantizer_ptr, const Config& config) {
|
||||
int64_t gpu_id = config[knowhere::meta::DEVICEID];
|
||||
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id)) {
|
||||
ResScope rs(res, gpu_id, false);
|
||||
faiss::gpu::GpuClonerOptions option;
|
||||
option.allInGpu = true;
|
||||
|
||||
auto ivf_quantizer = std::dynamic_pointer_cast<FaissIVFQuantizer>(quantizer_ptr);
|
||||
if (ivf_quantizer == nullptr)
|
||||
KNOWHERE_THROW_MSG("quantizer type not faissivfquantizer");
|
||||
|
||||
faiss::IndexComposition index_composition;
|
||||
index_composition.index = index_.get();
|
||||
index_composition.quantizer = ivf_quantizer->quantizer;
|
||||
index_composition.mode = 2; // only copy data
|
||||
|
||||
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id, &index_composition, &option);
|
||||
std::shared_ptr<faiss::Index> new_idx;
|
||||
new_idx.reset(gpu_index);
|
||||
auto sq_idx = std::make_shared<IVFSQHybrid>(new_idx, gpu_id, res);
|
||||
return sq_idx;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu: " + std::to_string(gpu_id) + "resource");
|
||||
}
|
||||
}
|
||||
|
||||
FaissIVFQuantizerPtr
|
||||
IVFSQHybrid::LoadQuantizer(const Config& config) {
|
||||
auto gpu_id = config[knowhere::meta::DEVICEID].get<int64_t>();
|
||||
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id)) {
|
||||
ResScope rs(res, gpu_id, false);
|
||||
faiss::gpu::GpuClonerOptions option;
|
||||
option.allInGpu = true;
|
||||
|
||||
faiss::IndexComposition index_composition;
|
||||
index_composition.index = index_.get();
|
||||
index_composition.quantizer = nullptr;
|
||||
index_composition.mode = 1; // only copy quantizer
|
||||
|
||||
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id, &index_composition, &option);
|
||||
delete gpu_index;
|
||||
|
||||
auto q = std::make_shared<FaissIVFQuantizer>();
|
||||
|
||||
auto q_ptr = index_composition.quantizer;
|
||||
q->size = q_ptr->d * q_ptr->getNumVecs() * sizeof(float);
|
||||
q->quantizer = q_ptr;
|
||||
q->gpu_id = gpu_id;
|
||||
return q;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu: " + std::to_string(gpu_id) + "resource");
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IVFSQHybrid::SetQuantizer(const FaissIVFQuantizerPtr& quantizer_ptr) {
|
||||
faiss::IndexIVF* ivf_index = dynamic_cast<faiss::IndexIVF*>(index_.get());
|
||||
if (ivf_index == nullptr) {
|
||||
KNOWHERE_THROW_MSG("Index type error");
|
||||
}
|
||||
|
||||
// Once SetQuantizer() is called, make sure UnsetQuantizer() is also called before destructuring.
|
||||
// Otherwise, ivf_index->quantizer will be double free.
|
||||
|
||||
quantizer_ = quantizer_ptr;
|
||||
ivf_index->quantizer = quantizer_->quantizer;
|
||||
gpu_mode_ = 1;
|
||||
}
|
||||
|
||||
void
|
||||
IVFSQHybrid::UnsetQuantizer() {
|
||||
auto* ivf_index = dynamic_cast<faiss::IndexIVF*>(index_.get());
|
||||
if (ivf_index == nullptr) {
|
||||
KNOWHERE_THROW_MSG("Index type error");
|
||||
}
|
||||
|
||||
// set back to cpu mode
|
||||
ivf_index->restore_quantizer();
|
||||
quantizer_ = nullptr;
|
||||
gpu_mode_ = 0;
|
||||
}
|
||||
|
||||
BinarySet
|
||||
IVFSQHybrid::SerializeImpl(const IndexType& type) {
|
||||
if (!index_ || !index_->is_trained) {
|
||||
KNOWHERE_THROW_MSG("index not initialize or trained");
|
||||
}
|
||||
|
||||
fiu_do_on("IVFSQHybrid.SerializeImpl.zero_gpu_mode", gpu_mode_ = 0);
|
||||
if (gpu_mode_ == 0) {
|
||||
MemoryIOWriter writer;
|
||||
faiss::write_index(index_.get(), &writer);
|
||||
|
||||
std::shared_ptr<uint8_t[]> data(writer.data_);
|
||||
|
||||
BinarySet res_set;
|
||||
res_set.Append("IVF", data, writer.rp);
|
||||
|
||||
return res_set;
|
||||
} else if (gpu_mode_ == 2) {
|
||||
return GPUIVF::SerializeImpl(type);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Can't serialize IVFSQ8Hybrid");
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IVFSQHybrid::LoadImpl(const BinarySet& binary_set, const IndexType& type) {
|
||||
FaissBaseIndex::LoadImpl(binary_set, index_type_); // load on cpu
|
||||
auto* ivf_index = dynamic_cast<faiss::IndexIVF*>(index_.get());
|
||||
ivf_index->backup_quantizer();
|
||||
gpu_mode_ = 0;
|
||||
}
|
||||
|
||||
void
|
||||
IVFSQHybrid::QueryImpl(int64_t n,
|
||||
const float* data,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* labels,
|
||||
const Config& config,
|
||||
const faiss::BitsetView bitset) {
|
||||
if (gpu_mode_ == 2) {
|
||||
GPUIVF::QueryImpl(n, data, k, distances, labels, config, bitset);
|
||||
// index_->search(n, (float*)data, k, distances, labels);
|
||||
} else if (gpu_mode_ == 1) { // hybrid
|
||||
auto gpu_id = quantizer_->gpu_id;
|
||||
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id)) {
|
||||
ResScope rs(res, gpu_id, true);
|
||||
IVF::QueryImpl(n, data, k, distances, labels, config, bitset);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("Hybrid Search Error, can't get gpu: " + std::to_string(gpu_id) + "resource");
|
||||
}
|
||||
} else if (gpu_mode_ == 0) {
|
||||
IVF::QueryImpl(n, data, k, distances, labels, config, bitset);
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
IVFSQHybrid::UpdateIndexSize() {
|
||||
if (!index_) {
|
||||
KNOWHERE_THROW_MSG("index not initialize");
|
||||
}
|
||||
auto ivfsqh_index = dynamic_cast<faiss::IndexIVFSQHybrid*>(index_.get());
|
||||
auto nb = ivfsqh_index->invlists->compute_ntotal();
|
||||
auto code_size = ivfsqh_index->code_size;
|
||||
auto nlist = ivfsqh_index->nlist;
|
||||
auto d = ivfsqh_index->d;
|
||||
// ivf codes, ivf ids, sq trained vectors and quantizer
|
||||
index_size_ = nb * code_size + nb * sizeof(int64_t) + 2 * d * sizeof(float) + nlist * d * sizeof(float);
|
||||
}
|
||||
|
||||
FaissIVFQuantizer::~FaissIVFQuantizer() {
|
||||
if (quantizer != nullptr) {
|
||||
delete quantizer;
|
||||
quantizer = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,105 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <faiss/gpu/GpuIndexFlat.h>
|
||||
#include <faiss/index_io.h>
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
|
||||
struct FaissIVFQuantizer {
|
||||
faiss::gpu::GpuIndexFlat* quantizer = nullptr;
|
||||
int64_t gpu_id;
|
||||
int64_t size = -1;
|
||||
|
||||
~FaissIVFQuantizer();
|
||||
};
|
||||
using FaissIVFQuantizerPtr = std::shared_ptr<FaissIVFQuantizer>;
|
||||
|
||||
class IVFSQHybrid : public GPUIVFSQ {
|
||||
public:
|
||||
explicit IVFSQHybrid(const int& device_id) : GPUIVFSQ(device_id) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8H;
|
||||
gpu_mode_ = 0;
|
||||
index_mode_ = IndexMode::MODE_CPU;
|
||||
}
|
||||
|
||||
explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index) : GPUIVFSQ(-1) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8H;
|
||||
index_ = index;
|
||||
gpu_mode_ = 0;
|
||||
index_mode_ = IndexMode::MODE_CPU;
|
||||
}
|
||||
|
||||
explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index, const int64_t device_id, ResPtr& resource)
|
||||
: GPUIVFSQ(index, device_id, resource) {
|
||||
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8H;
|
||||
gpu_mode_ = 2;
|
||||
}
|
||||
|
||||
public:
|
||||
void
|
||||
Train(const DatasetPtr&, const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyGpuToCpu(const Config&) override;
|
||||
|
||||
VecIndexPtr
|
||||
CopyCpuToGpu(const int64_t, const Config&) override;
|
||||
|
||||
std::pair<VecIndexPtr, FaissIVFQuantizerPtr>
|
||||
CopyCpuToGpuWithQuantizer(const int64_t, const Config&);
|
||||
|
||||
VecIndexPtr
|
||||
LoadData(const FaissIVFQuantizerPtr&, const Config&);
|
||||
|
||||
FaissIVFQuantizerPtr
|
||||
LoadQuantizer(const Config& conf);
|
||||
|
||||
void
|
||||
SetQuantizer(const FaissIVFQuantizerPtr& q);
|
||||
|
||||
void
|
||||
UnsetQuantizer();
|
||||
|
||||
void
|
||||
UpdateIndexSize() override;
|
||||
|
||||
protected:
|
||||
BinarySet
|
||||
SerializeImpl(const IndexType&) override;
|
||||
|
||||
void
|
||||
LoadImpl(const BinarySet&, const IndexType&) override;
|
||||
|
||||
void
|
||||
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::BitsetView) override;
|
||||
|
||||
protected:
|
||||
int64_t gpu_mode_ = 0; // 0: CPU, 1: Hybrid, 2: GPU
|
||||
FaissIVFQuantizerPtr quantizer_ = nullptr;
|
||||
};
|
||||
|
||||
using IVFSQHybridPtr = std::shared_ptr<IVFSQHybrid>;
|
||||
|
||||
#endif
|
||||
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
@ -1,74 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#ifdef MILVUS_GPU_VERSION
|
||||
#include "knowhere/index/vector_index/helpers/Cloner.h"
|
||||
#include "knowhere/common/Exception.h"
|
||||
#include "knowhere/index/vector_index/IndexIDMAP.h"
|
||||
#include "knowhere/index/vector_index/IndexIVF.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFPQ.h"
|
||||
#include "knowhere/index/vector_index/IndexIVFSQ.h"
|
||||
#include "knowhere/index/vector_index/gpu/GPUIndex.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h"
|
||||
#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h"
|
||||
#include "knowhere/index/vector_offset_index/IndexIVF_NM.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
namespace cloner {
|
||||
|
||||
void
|
||||
CopyIndexData(const VecIndexPtr& dst_index, const VecIndexPtr& src_index) {
|
||||
dst_index->SetUids(src_index->GetUids());
|
||||
dst_index->SetIndexSize(src_index->IndexSize());
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
CopyGpuToCpu(const VecIndexPtr& index, const Config& config) {
|
||||
if (auto device_index = std::dynamic_pointer_cast<GPUIndex>(index)) {
|
||||
VecIndexPtr result = device_index->CopyGpuToCpu(config);
|
||||
CopyIndexData(result, index);
|
||||
return result;
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("index type is not gpuindex");
|
||||
}
|
||||
}
|
||||
|
||||
VecIndexPtr
|
||||
CopyCpuToGpu(const VecIndexPtr& index, const int64_t device_id, const Config& config) {
|
||||
VecIndexPtr result;
|
||||
if (auto device_index = std::dynamic_pointer_cast<IVFSQHybrid>(index)) {
|
||||
result = device_index->CopyCpuToGpu(device_id, config);
|
||||
} else if (auto cpu_index = std::dynamic_pointer_cast<IVF_NM>(index)) {
|
||||
result = cpu_index->CopyCpuToGpu(device_id, config);
|
||||
} else if (auto device_index = std::dynamic_pointer_cast<GPUIndex>(index)) {
|
||||
result = device_index->CopyGpuToGpu(device_id, config);
|
||||
} else if (auto cpu_index = std::dynamic_pointer_cast<IVFSQ>(index)) {
|
||||
result = cpu_index->CopyCpuToGpu(device_id, config);
|
||||
} else if (auto cpu_index = std::dynamic_pointer_cast<IVFPQ>(index)) {
|
||||
result = cpu_index->CopyCpuToGpu(device_id, config);
|
||||
} else if (auto cpu_index = std::dynamic_pointer_cast<IVF>(index)) {
|
||||
result = cpu_index->CopyCpuToGpu(device_id, config);
|
||||
} else if (auto cpu_index = std::dynamic_pointer_cast<IDMAP>(index)) {
|
||||
result = cpu_index->CopyCpuToGpu(device_id, config);
|
||||
} else {
|
||||
KNOWHERE_THROW_MSG("this index type not support transfer to gpu");
|
||||
}
|
||||
if (result != nullptr) {
|
||||
CopyIndexData(result, index);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace cloner
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
#endif
|
||||
@ -1,28 +0,0 @@
|
||||
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
||||
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "knowhere/index/vector_index/VecIndex.h"
|
||||
|
||||
namespace milvus {
|
||||
namespace knowhere {
|
||||
namespace cloner {
|
||||
|
||||
extern VecIndexPtr
|
||||
CopyCpuToGpu(const VecIndexPtr& index, const int64_t device_id, const Config& config);
|
||||
|
||||
extern VecIndexPtr
|
||||
CopyGpuToCpu(const VecIndexPtr& index, const Config& config);
|
||||
|
||||
} // namespace cloner
|
||||
} // namespace knowhere
|
||||
} // namespace milvus
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue
Block a user