From a4fe9d1a28a3ca2dd0538280456ec1d3470c95ba Mon Sep 17 00:00:00 2001 From: johnny <1667704220@qq.com> Date: Sun, 9 Feb 2025 22:31:09 +0800 Subject: [PATCH 1/2] =?UTF-8?q?=E6=94=AF=E6=8C=81ivfpq=E7=B4=A2=E5=BC=95?= =?UTF-8?q?=E6=9E=84=E5=BB=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/bin/gs_guc/cluster_guc.conf | 2 +- src/common/backend/utils/adt/vector.cpp | 30 +- .../backend/utils/misc/guc/guc_storage.cpp | 6 +- .../process/postmaster/postmaster.cpp | 10 +- .../process/threadpool/knl_instance.cpp | 2 +- .../storage/access/common/reloptions.cpp | 16 +- .../storage/access/datavec/Makefile | 2 +- .../storage/access/datavec/hnswadaptor.cpp | 127 ++++---- .../storage/access/datavec/hnswbuild.cpp | 4 +- .../storage/access/datavec/hnswutils.cpp | 10 +- .../storage/access/datavec/ivfadaptor.cpp | 46 +++ .../storage/access/datavec/ivfbuild.cpp | 308 +++++++++++++++++- .../storage/access/datavec/ivfflat.cpp | 4 + .../storage/access/datavec/ivfutils.cpp | 162 +++++++++ .../storage/access/datavec/utils.cpp | 2 + src/include/access/datavec/hnsw.h | 36 +- src/include/access/datavec/ivfflat.h | 68 ++++ src/include/access/datavec/utils.h | 34 +- src/include/access/datavec/vector.h | 2 + .../knl/knl_guc/knl_instance_attr_storage.h | 2 +- src/include/knl/knl_instance.h | 2 +- 21 files changed, 743 insertions(+), 132 deletions(-) create mode 100644 src/gausskernel/storage/access/datavec/ivfadaptor.cpp diff --git a/src/bin/gs_guc/cluster_guc.conf b/src/bin/gs_guc/cluster_guc.conf index 221de0340e..7908da000d 100755 --- a/src/bin/gs_guc/cluster_guc.conf +++ b/src/bin/gs_guc/cluster_guc.conf @@ -848,7 +848,7 @@ uwal_async_append_switch|bool|0,0|NULL|NULL| enable_gazelle_performance_mode|bool|0,0|NULL|NULL| enable_aggr_coerce_type|bool|0,0|NULL|NULL| enable_default_local_index|bool|0,0|NULL|NULL| -enable_hnswpq|bool|0,0|NULL|NULL| +enable_pq|bool|0,0|NULL|NULL| [cmserver] log_dir|string|0,0|NULL|NULL| log_file_size|int|0,2047|MB|NULL| diff --git a/src/common/backend/utils/adt/vector.cpp b/src/common/backend/utils/adt/vector.cpp index 795dc72d9f..0ec9734b14 100644 --- a/src/common/backend/utils/adt/vector.cpp +++ b/src/common/backend/utils/adt/vector.cpp @@ -821,7 +821,7 @@ Datum vector_l2_squared_distance(PG_FUNCTION_ARGS) } #ifdef __aarch64__ -VECTOR_TARGET_CLONES static float +VECTOR_TARGET_CLONES float VectorInnerProduct(int dim, float *ax, float *bx) { float dis = 0.0f; @@ -850,7 +850,7 @@ VectorInnerProduct(int dim, float *ax, float *bx) } #else -VECTOR_TARGET_CLONES static float VectorInnerProduct(int dim, float *ax, float *bx) +VECTOR_TARGET_CLONES float VectorInnerProduct(int dim, float *ax, float *bx) { float distance = 0.0; @@ -1539,6 +1539,32 @@ Datum sparsevec_to_vector(PG_FUNCTION_ARGS) PG_RETURN_POINTER(result); } +#ifdef __aarch64__ +void VectorMadd(size_t n, const float *ax, float bf, const float *bx, float *cx) +{ + const size_t nSimd = n - (n & 3); + const float32x4_t bfv = vdupq_n_f32(bf); + size_t i; + + for (i = 0; i < nSimd; i += 4) { + const float32x4_t ai = vld1q_f32(ax + i); + const float32x4_t bi = vld1q_f32(bx + i); + const float32x4_t ci = vfmaq_f32(ai, bfv, bi); + vst1q_f32(cx + i, ci); + } + for (; i < n; ++i) { + cx[i] = ax[i] + bf * bx[i]; + } +} +#else +void VectorMadd(size_t n, const float *ax, float bf, const float *bx, float *cx) +{ + for (size_t i = 0; i < n; i++) { + cx[i] = ax[i] + bf * bx[i]; + } +} +#endif + /* * WAL-log a range of blocks in a relation. * diff --git a/src/common/backend/utils/misc/guc/guc_storage.cpp b/src/common/backend/utils/misc/guc/guc_storage.cpp index c50cade0f3..6114ab174d 100755 --- a/src/common/backend/utils/misc/guc/guc_storage.cpp +++ b/src/common/backend/utils/misc/guc/guc_storage.cpp @@ -1409,13 +1409,13 @@ static void InitStorageConfigureNamesBool() NULL, NULL, NULL}, - {{"enable_hnswpq", + {{"enable_pq", PGC_POSTMASTER, NODE_SINGLENODE, QUERY_TUNING_OTHER, - gettext_noop("Whether enable hnswpq"), + gettext_noop("Whether enable pq in datavec"), NULL}, - &g_instance.attr.attr_storage.enable_hnswpq, + &g_instance.attr.attr_storage.enable_pq, false, NULL, NULL, diff --git a/src/gausskernel/process/postmaster/postmaster.cpp b/src/gausskernel/process/postmaster/postmaster.cpp index d7b91c978c..38f42017d1 100644 --- a/src/gausskernel/process/postmaster/postmaster.cpp +++ b/src/gausskernel/process/postmaster/postmaster.cpp @@ -3152,12 +3152,12 @@ int PostmasterMain(int argc, char* argv[]) } /* init datavec hnswpq */ - if (g_instance.attr.attr_storage.enable_hnswpq) { - int ret = HNSWPQInit(); + if (g_instance.attr.attr_storage.enable_pq) { + int ret = PQInit(); if (ret != 0) { - ereport(PANIC, (errmsg("datavec HNSWPQ init failed, ret: %d", ret))); + ereport(PANIC, (errmsg("datavec PQ init failed, ret: %d", ret))); } - ereport(LOG, (errmsg("datavec HNSWPQ init success."))); + ereport(LOG, (errmsg("datavec PQ init success."))); } /* init sharestorge(dorado) */ @@ -9748,7 +9748,7 @@ void ExitPostmaster(int status) * MUST -- vadim 05-10-1999 */ DMSUninit(); - HNSWPQUinit(); + PQUinit(); CloseGaussPidDir(); diff --git a/src/gausskernel/process/threadpool/knl_instance.cpp b/src/gausskernel/process/threadpool/knl_instance.cpp index 1861914cc7..de104572e5 100755 --- a/src/gausskernel/process/threadpool/knl_instance.cpp +++ b/src/gausskernel/process/threadpool/knl_instance.cpp @@ -1167,7 +1167,7 @@ void knl_instance_init() knl_g_datadir_init(&g_instance.datadir_cxt); knl_g_listen_sock_init(&g_instance.listen_cxt); - g_instance.hnswpq_inited = false; + g_instance.pq_inited = false; #ifdef USE_SPQ knl_g_spq_context_init(&g_instance.spq_cxt); diff --git a/src/gausskernel/storage/access/common/reloptions.cpp b/src/gausskernel/storage/access/common/reloptions.cpp index 41a89b3c38..c6e5eb3b69 100644 --- a/src/gausskernel/storage/access/common/reloptions.cpp +++ b/src/gausskernel/storage/access/common/reloptions.cpp @@ -125,7 +125,8 @@ static relopt_bool boolRelOpts[] = { {{"compress_diff_convert", "Whether do diiffer convert in compression", RELOPT_KIND_HEAP | RELOPT_KIND_BTREE}, false}, {{"deduplication", "Enables \"deduplication\" feature for btree index", RELOPT_KIND_BTREE}, false}, - {{ "enable_pq", "Whether to enable PQ", RELOPT_KIND_HNSW }, HNSW_DEFAULT_ENABLE_PQ }, + {{"enable_pq", "Whether to enable PQ", RELOPT_KIND_HNSW | RELOPT_KIND_IVFFLAT }, GENERIC_DEFAULT_ENABLE_PQ }, + {{"by_residual", "Whether to use residual during IVFPQ", RELOPT_KIND_IVFFLAT}, IVFFLAT_DEFAULT_PQ_RESIDUAL}, /* list terminator */ {{NULL}}}; @@ -263,11 +264,14 @@ static relopt_int intRelOpts[] = { HNSW_DEFAULT_EF_CONSTRUCTION, HNSW_MIN_EF_CONSTRUCTION, HNSW_MAX_EF_CONSTRUCTION }, - {{ "pq_m", "Number of PQ subquantizer", RELOPT_KIND_HNSW }, HNSW_DEFAULT_PQ_M, HNSW_MIN_PQ_M, HNSW_MAX_PQ_M }, - {{ "pq_ksub", "Number of centroids for each PQ subquantizer", RELOPT_KIND_HNSW }, - HNSW_DEFAULT_PQ_KSUB, - HNSW_MIN_PQ_KSUB, - HNSW_MAX_PQ_KSUB }, + {{ "pq_m", "Number of PQ subquantizer", RELOPT_KIND_HNSW |RELOPT_KIND_IVFFLAT}, + GENERIC_DEFAULT_PQ_M, + GENERIC_MIN_PQ_M, + GENERIC_MAX_PQ_M }, + {{ "pq_ksub", "Number of centroids for each PQ subquantizer", RELOPT_KIND_HNSW | RELOPT_KIND_IVFFLAT }, + GENERIC_DEFAULT_PQ_KSUB, + GENERIC_MIN_PQ_KSUB, + GENERIC_MAX_PQ_KSUB }, {{ "lists", "Number of inverted lists", RELOPT_KIND_IVFFLAT }, IVFFLAT_DEFAULT_LISTS, IVFFLAT_MIN_LISTS, diff --git a/src/gausskernel/storage/access/datavec/Makefile b/src/gausskernel/storage/access/datavec/Makefile index c99a75053e..c32985f862 100644 --- a/src/gausskernel/storage/access/datavec/Makefile +++ b/src/gausskernel/storage/access/datavec/Makefile @@ -12,6 +12,6 @@ endif OBJS = bitutils.o hnsw.o hnswbuild.o hnswdelete.o hnswinsert.o hnswscan.o hnswutils.o hnswvacuum.o \ ivfbuild.o ivfflat.o ivfinsert.o ivfkmeans.o ivfscan.o ivfutils.o ivfvacuum.o vecindex.o \ - utils.o hnswadaptor.o + utils.o hnswadaptor.o ivfadaptor.o include $(top_srcdir)/src/gausskernel/common.mk diff --git a/src/gausskernel/storage/access/datavec/hnswadaptor.cpp b/src/gausskernel/storage/access/datavec/hnswadaptor.cpp index b53e99da3f..21684b415a 100644 --- a/src/gausskernel/storage/access/datavec/hnswadaptor.cpp +++ b/src/gausskernel/storage/access/datavec/hnswadaptor.cpp @@ -22,166 +22,165 @@ */ #include #include "access/datavec/hnsw.h" +#include "access/datavec/utils.h" -hnswpq_func_t g_hnsw_func = {0}; - -// return HNSWPQ_ERROR if error occurs -#define HNSWPQ_RETURN_IFERR(ret) \ +// return PQ_ERROR if error occurs +#define PQ_RETURN_IFERR(ret) \ do { \ int _status_ = (ret); \ - if (SECUREC_UNLIKELY(_status_ != HNSWPQ_SUCCESS)) { \ + if (SECUREC_UNLIKELY(_status_ != PQ_SUCCESS)) { \ return _status_; \ } \ } while (0) -int hnswpq_resolve_path(char* absolute_path, const char* raw_path, const char* filename) +int pq_resolve_path(char* absolute_path, const char* raw_path, const char* filename) { char path[MAX_PATH_LEN] = { 0 }; if (!realpath(raw_path, path)) { if (errno != ENOENT && errno != EACCES) { - return HNSWPQ_ERROR; + return PQ_ERROR; } } int ret = snprintf_s(absolute_path, MAX_PATH_LEN, MAX_PATH_LEN - 1, "%s/%s", path, filename); if (ret < 0) { - return HNSWPQ_ERROR; + return PQ_ERROR; } - return HNSWPQ_SUCCESS; + return PQ_SUCCESS; } -int hnswpq_load_symbol(char *symbol, void **sym_lib_handle) +int pq_load_symbol(char *symbol, void **sym_lib_handle) { #ifndef WIN32 const char *dlsym_err = NULL; - *sym_lib_handle = dlsym(g_hnsw_func.handle, symbol); + *sym_lib_handle = dlsym(g_pq_func.handle, symbol); dlsym_err = dlerror(); if (dlsym_err != NULL) { ereport(FATAL, (errcode(ERRCODE_INVALID_OPERATION), - errmsg("incompatible library \"%s\", load %s failed, %s", HNSWPQ_SO_NAME, symbol, dlsym_err))); - return HNSWPQ_ERROR; + errmsg("incompatible library \"%s\", load %s failed, %s", PQ_SO_NAME, symbol, dlsym_err))); + return PQ_ERROR; } #endif // !WIN32 - return HNSWPQ_SUCCESS; + return PQ_SUCCESS; } -#define HNSWPQ_LOAD_SYMBOL_FUNC(func) hnswpq_load_symbol(#func, (void **)&g_hnsw_func.func) +#define PQ_LOAD_SYMBOL_FUNC(func) pq_load_symbol(#func, (void **)&g_pq_func.func) -int hnswpq_open_dl(void **lib_handle, char *symbol) +int pq_open_dl(void **lib_handle, char *symbol) { #ifndef WIN32 *lib_handle = dlopen(symbol, RTLD_LAZY); if (*lib_handle == NULL) { - ereport(ERROR, (errcode_for_file_access(), errmsg("could not load library %s, %s", HNSWPQ_SO_NAME, dlerror()))); - return HNSWPQ_ERROR; + ereport(ERROR, (errcode_for_file_access(), errmsg("could not load library %s, %s", PQ_SO_NAME, dlerror()))); + return PQ_ERROR; } - return HNSWPQ_SUCCESS; + return PQ_SUCCESS; #else - return HNSWPQ_ERROR; + return PQ_ERROR; #endif } -void hnswpq_close_dl(void *lib_handle) +void pq_close_dl(void *lib_handle) { #ifndef WIN32 (void)dlclose(lib_handle); #endif } -int hnswpq_load_symbols(char *lib_dl_path) +int pq_load_symbols(char *lib_dl_path) { - HNSWPQ_RETURN_IFERR(hnswpq_open_dl(&g_hnsw_func.handle, lib_dl_path)); + PQ_RETURN_IFERR(pq_open_dl(&g_pq_func.handle, lib_dl_path)); - HNSWPQ_RETURN_IFERR(HNSWPQ_LOAD_SYMBOL_FUNC(ComputePQTable)); - HNSWPQ_RETURN_IFERR(HNSWPQ_LOAD_SYMBOL_FUNC(ComputeVectorPQCode)); - HNSWPQ_RETURN_IFERR(HNSWPQ_LOAD_SYMBOL_FUNC(GetPQDistanceTableSdc)); - HNSWPQ_RETURN_IFERR(HNSWPQ_LOAD_SYMBOL_FUNC(GetPQDistanceTableAdc)); - HNSWPQ_RETURN_IFERR(HNSWPQ_LOAD_SYMBOL_FUNC(GetPQDistance)); + PQ_RETURN_IFERR(PQ_LOAD_SYMBOL_FUNC(ComputePQTable)); + PQ_RETURN_IFERR(PQ_LOAD_SYMBOL_FUNC(ComputeVectorPQCode)); + PQ_RETURN_IFERR(PQ_LOAD_SYMBOL_FUNC(GetPQDistanceTableSdc)); + PQ_RETURN_IFERR(PQ_LOAD_SYMBOL_FUNC(GetPQDistanceTableAdc)); + PQ_RETURN_IFERR(PQ_LOAD_SYMBOL_FUNC(GetPQDistance)); - return HNSWPQ_SUCCESS; + return PQ_SUCCESS; } -int hnswpq_func_init() +int pq_func_init() { - if (g_hnsw_func.inited) { - return HNSWPQ_SUCCESS; + if (g_pq_func.inited) { + return PQ_SUCCESS; } char lib_dl_path[MAX_PATH_LEN] = { 0 }; - char* raw_path = getenv(HNSWPQ_ENV_PATH); + char* raw_path = getenv(PQ_ENV_PATH); if (raw_path == nullptr) { - ereport(ERROR, (errmsg("failed to get DATAVEC_HNSWPQ_LIB_PATH"))); - return HNSWPQ_ERROR; + ereport(ERROR, (errmsg("failed to get DATAVEC_PQ_LIB_PATH"))); + return PQ_ERROR; } - int ret = hnswpq_resolve_path(lib_dl_path, raw_path, HNSWPQ_SO_NAME); - if (ret != HNSWPQ_SUCCESS) { + int ret = pq_resolve_path(lib_dl_path, raw_path, PQ_SO_NAME); + if (ret != PQ_SUCCESS) { ereport(ERROR, (errmsg( "failed to resolve the path of libvecturbo.so, lib_dl_path %s, raw_path %s", lib_dl_path, raw_path))); - return HNSWPQ_ERROR; + return PQ_ERROR; } - ret = hnswpq_load_symbols(lib_dl_path); - if (ret != HNSWPQ_SUCCESS) { - return HNSWPQ_ERROR; + ret = pq_load_symbols(lib_dl_path); + if (ret != PQ_SUCCESS) { + return PQ_ERROR; } - g_hnsw_func.inited = true; - return HNSWPQ_SUCCESS; + g_pq_func.inited = true; + return PQ_SUCCESS; } -int HNSWPQInit() +int PQInit() { #ifdef __x86_64__ - ereport(FATAL, (errmsg("HNSWPQ only support in arm."))); + ereport(FATAL, (errmsg("PQ only support in arm."))); #endif - if (hnswpq_func_init() != HNSWPQ_SUCCESS) { - ereport(FATAL, (errmsg("failed to init HNSWPQ library"))); - return HNSWPQ_ERROR; + if (pq_func_init() != PQ_SUCCESS) { + ereport(FATAL, (errmsg("failed to init PQ library"))); + return PQ_ERROR; } - g_instance.hnswpq_inited = true; - return HNSWPQ_SUCCESS; + g_instance.pq_inited = true; + return PQ_SUCCESS; } -void HNSWPQUinit() +void PQUinit() { - if (!g_instance.attr.attr_storage.enable_hnswpq || ! g_instance.hnswpq_inited) { + if (!g_instance.attr.attr_storage.enable_pq || ! g_instance.pq_inited) { return; } - g_instance.hnswpq_inited = false; - ereport(LOG, (errmsg("datavec HNSWPQ uninit"))); - if (g_hnsw_func.handle != NULL) { - hnswpq_close_dl(g_hnsw_func.handle); - g_hnsw_func.handle = NULL; - g_hnsw_func.inited = false; + g_instance.pq_inited = false; + ereport(LOG, (errmsg("datavec PQ uninit"))); + if (g_pq_func.handle != NULL) { + pq_close_dl(g_pq_func.handle); + g_pq_func.handle = NULL; + g_pq_func.inited = false; } } int ComputePQTable(VectorArray samples, PQParams *params) { - return g_hnsw_func.ComputePQTable(samples, params); + return g_pq_func.ComputePQTable(samples, params); } int ComputeVectorPQCode(float *vector, const PQParams *params, uint8 *pqCode) { - return g_hnsw_func.ComputeVectorPQCode(vector, params, pqCode); + return g_pq_func.ComputeVectorPQCode(vector, params, pqCode); } int GetPQDistanceTableSdc(const PQParams *params, float *pqDistanceTable) { - return g_hnsw_func.GetPQDistanceTableSdc(params, pqDistanceTable); + return g_pq_func.GetPQDistanceTableSdc(params, pqDistanceTable); } int GetPQDistanceTableAdc(float *vector, const PQParams *params, float *pqDistanceTable) { - return g_hnsw_func.GetPQDistanceTableAdc(vector, params, pqDistanceTable); + return g_pq_func.GetPQDistanceTableAdc(vector, params, pqDistanceTable); } int GetPQDistance(const uint8 *basecode, const uint8 *querycode, const PQParams *params, - const float *pqDistanceTable, float *PQDistance) + const float *pqDistanceTable, float *pqDistance) { - return g_hnsw_func.GetPQDistance(basecode, querycode, params, pqDistanceTable, PQDistance); + return g_pq_func.GetPQDistance(basecode, querycode, params, pqDistanceTable, pqDistance); } \ No newline at end of file diff --git a/src/gausskernel/storage/access/datavec/hnswbuild.cpp b/src/gausskernel/storage/access/datavec/hnswbuild.cpp index 41b51ad0f3..6494b16b67 100644 --- a/src/gausskernel/storage/access/datavec/hnswbuild.cpp +++ b/src/gausskernel/storage/access/datavec/hnswbuild.cpp @@ -1233,8 +1233,8 @@ static void InitBuildState(HnswBuildState *buildstate, Relation heap, Relation i if (buildstate->enablePQ && !buildstate->typeInfo->supportPQ) { ereport(ERROR, (errmsg("this data type cannot support hnswpq."))); } - if (buildstate->enablePQ && !g_instance.hnswpq_inited) { - ereport(ERROR, (errmsg("this instance has not currently loaded the hnswpq dynamic library."))); + if (buildstate->enablePQ && !g_instance.pq_inited) { + ereport(ERROR, (errmsg("this instance has not currently loaded the pq dynamic library."))); } buildstate->pqM = HnswGetPqM(index); diff --git a/src/gausskernel/storage/access/datavec/hnswutils.cpp b/src/gausskernel/storage/access/datavec/hnswutils.cpp index 43fe0dab47..75c2498349 100644 --- a/src/gausskernel/storage/access/datavec/hnswutils.cpp +++ b/src/gausskernel/storage/access/datavec/hnswutils.cpp @@ -163,7 +163,7 @@ bool HnswGetEnablePQ(Relation index) return opts->enablePQ; } - return HNSW_DEFAULT_ENABLE_PQ; + return GENERIC_DEFAULT_ENABLE_PQ; } /* @@ -177,7 +177,7 @@ int HnswGetPqM(Relation index) return opts->pqM; } - return HNSW_DEFAULT_PQ_M; + return GENERIC_DEFAULT_PQ_M; } /* @@ -191,7 +191,7 @@ int HnswGetPqKsub(Relation index) return opts->pqKsub; } - return HNSW_DEFAULT_PQ_KSUB; + return GENERIC_DEFAULT_PQ_KSUB; } /* @@ -1593,9 +1593,9 @@ void InitPQParamsOnDisk(PQParams *params, Relation index, FmgrInfo *procinfo, in UnlockReleaseBuffer(buf); int pqMode = HNSW_PQMODE_DEFAULT; - if (*enablePQ && !g_instance.hnswpq_inited) { + if (*enablePQ && !g_instance.pq_inited) { ereport(ERROR, (errmsg("the SQL involves operations related to HNSWPQ, " - "but this instance has not currently loaded the HNSWPQ dynamic library."))); + "but this instance has not currently loaded the PQ dynamic library."))); } if (*enablePQ) { diff --git a/src/gausskernel/storage/access/datavec/ivfadaptor.cpp b/src/gausskernel/storage/access/datavec/ivfadaptor.cpp new file mode 100644 index 0000000000..0ddede4de4 --- /dev/null +++ b/src/gausskernel/storage/access/datavec/ivfadaptor.cpp @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2024 Huawei Technologies Co.,Ltd. + * + * openGauss is licensed under Mulan PSL v2. + * You can use this software according to the terms and conditions of the Mulan PSL v2. + * You may obtain a copy of Mulan PSL v2 at: + * + * http://license.coscl.org.cn/MulanPSL2 + * + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, + * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, + * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. + * See the Mulan PSL v2 for more details. + * ------------------------------------------------------------------------- + * + * ivfadaptor.cpp + * + * IDENTIFICATION + * src/gausskernel/storage/access/datavec/ivfadaptor.cpp + * + * ------------------------------------------------------------------------- + */ +#include +#include "access/datavec/ivfflat.h" +#include "access/datavec/utils.h" + +int IvfComputePQTable(VectorArray samples, PQParams *params) +{ + return g_pq_func.ComputePQTable(samples, params); +} + +int IvfComputeVectorPQCode(float *vector, const PQParams *params, uint8 *pqCode) +{ + return g_pq_func.ComputeVectorPQCode(vector, params, pqCode); +} + +int IvfGetPQDistanceTableAdc(float *vector, const PQParams *params, float *pqDistanceTable) +{ + return g_pq_func.GetPQDistanceTableAdc(vector, params, pqDistanceTable); +} + +int IvfGetPQDistance(const uint8 *basecode, const uint8 *querycode, const PQParams *params, + const float *pqDistanceTable, float *pqDistance) +{ + return g_pq_func.GetPQDistance(basecode, querycode, params, pqDistanceTable, pqDistance); +} \ No newline at end of file diff --git a/src/gausskernel/storage/access/datavec/ivfbuild.cpp b/src/gausskernel/storage/access/datavec/ivfbuild.cpp index 44e7c096ea..d24a30e8b1 100644 --- a/src/gausskernel/storage/access/datavec/ivfbuild.cpp +++ b/src/gausskernel/storage/access/datavec/ivfbuild.cpp @@ -47,6 +47,170 @@ #define PARALLEL_KEY_IVFFLAT_CENTERS UINT64CONST(0xA000000000000003) #define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000004) +/* + * Create PQ-related pages + */ +static void CreatePQPages(IvfflatBuildState *buildstate, ForkNumber fNum) +{ + uint16 nblks; + Relation index = buildstate->index; + ForkNumber forkNum = fNum; + Buffer buf; + Page page; + uint16 pqTableNblk; + uint16 pqPreComputeTableNblk; + + IvfGetPQInfoFromMetaPage(index, &pqTableNblk, NULL, &pqPreComputeTableNblk, NULL); + + /* create pq table page */ + for (uint16 i = 0; i < pqTableNblk; i++) { + buf = IvfflatNewBuffer(index, forkNum); + page = BufferGetPage(buf); + IvfflatInitPage(buf, page); + MarkBufferDirty(buf); + UnlockReleaseBuffer(buf); + } + + /* create pq distance table page */ + for (uint16 i = 0; i < pqPreComputeTableNblk; i++) { + buf = IvfflatNewBuffer(index, forkNum); + page = BufferGetPage(buf); + IvfflatInitPage(buf, page); + MarkBufferDirty(buf); + UnlockReleaseBuffer(buf); + } +} + +/* + * Caculate Residual + */ +static void ComputeResidual(IvfflatBuildState *buildstate, Vector* sample, int list) +{ + Vector *vec = (Vector *)lfirst(buildstate->rlist->tail); + Vector *center = (Vector *)VectorArrayGet(buildstate->centers, list); + + if (buildstate->byResidual) { + for (int i = 0; i < buildstate->dimensions; i++) { + vec->x[i] = sample->x[i] -center->x[i]; + } + } else { + for (int i = 0; i < buildstate->dimensions; i++) { + vec->x[i] = sample->x[i]; + } + } +} + +/* + * Caculate square of L2 normalform + */ +static float ComputeNormL2sqr(float *x, int dsub) +{ + float res = 0.0f; + for (int i = 0; i < dsub; i++) { + res += x[i] * x[i]; + } + return res; +} + +static void ComputeInnerProdAndSum(IvfflatBuildState *buildstate, float * l2Norm, float *center, float * tab, int dsub) +{ + Size itemSize = MAXALIGN(buildstate->typeInfo->itemSize(dsub)); + const float MULTIPLIER = 2.0; + + for (int i = 0; i < buildstate->pqM; i++) { + for (int j = 0; j < buildstate->pqKsub; j++) { + float *x = DatumGetVector(buildstate->pqTable + ((i * buildstate->pqKsub + j) * itemSize))->x; + tab[i * buildstate->pqKsub + j] = VectorInnerProduct(dsub, x, center + i * dsub); + float *pretable = &tab[i * buildstate->pqKsub + j]; + VectorMadd(1, l2Norm + (i * buildstate->pqKsub + j), MULTIPLIER, pretable, pretable); + } + } +} + +/* + * Compute precalculated table + */ +static void ComputePreTable(IvfflatBuildState *buildstate) +{ + Size size = buildstate->pqKsub * buildstate->pqM * sizeof(float); + float *l2Norm = (float *)palloc0(size); + + int dsub = buildstate->dimensions / buildstate->pqM; + Size itemSize = MAXALIGN(buildstate->typeInfo->itemSize(dsub)); + + for (int m = 0; m < buildstate->pqM; m++) { + for (int j = 0; j < buildstate->pqKsub; j++) { + float *x = DatumGetVector(buildstate->pqTable + (m * buildstate->pqKsub + j) * itemSize)->x; + l2Norm[m * buildstate->pqKsub + j] = ComputeNormL2sqr(x, dsub); + } + } + + for (int n = 0; n < buildstate->lists; n++) { + float *tab = buildstate->preComputeTable + n * buildstate->pqM * buildstate->pqKsub; + Vector *center = (Vector *)VectorArrayGet(buildstate->centers, n); + ComputeInnerProdAndSum(buildstate, l2Norm, center->x, tab, dsub); + } + + pfree(l2Norm); +} + +/* + * Compute PQTable + */ +static void ComputeIvfPQ(IvfflatBuildState *buildstate) +{ + MemoryContext pqCtx = AllocSetContextCreate(CurrentMemoryContext, + "Ivfflat PQ temporary context", + ALLOCSET_DEFAULT_SIZES); + MemoryContext oldCtx = MemoryContextSwitchTo(pqCtx); + + IvfComputePQTable(buildstate->residuals, buildstate->params); + MemoryContextSwitchTo(oldCtx); + MemoryContextDelete(pqCtx); +} + +/* + * Get all sample vector or residual vector to vector array + */ +static void CopyResidaulFromList(IvfflatBuildState *buildstate) +{ + if (buildstate->rlist == NIL) { + ereport(ERROR, (errmsg("when enable_pq = on, at least one vector needs to be include"))); + } + + ListCell *lc; + buildstate->residuals = VectorArrayInit( + buildstate->rlist->length, + buildstate->dimensions, + buildstate->typeInfo->itemSize(buildstate->dimensions) + ); + + foreach (lc, buildstate->rlist) { + Vector *vec = (Vector *)lfirst(lc); + Datum value = PointerGetDatum(vec); + value = PointerGetDatum(PG_DETOAST_DATUM(value)); + VectorArraySet(buildstate->residuals, buildstate->residuals->length, DatumGetPointer(value)); + buildstate->residuals->length++; + } + list_free_deep(buildstate->rlist); +} + +/* + * Init PQParam + */ +PQParams *InitIVFPQParamsInMemory(IvfflatBuildState *buildstate) +{ + PQParams *params = (PQParams*)palloc(sizeof(PQParams)); + params->pqM = buildstate->pqM; + params->pqKsub = buildstate->pqKsub; + params->funcType = getIVFPQfunctionType(buildstate->procinfo, buildstate->normprocinfo); + params->dim = buildstate->dimensions; + params->pqMode = IVF_PQMODE_DEFAULT; + params->subItemSize = buildstate->typeInfo->itemSize(buildstate->dimensions / buildstate->pqM); + params->pqTable = buildstate->pqTable; + return params; +} + /* * Add sample */ @@ -167,6 +331,14 @@ static void AddTupleToSort(Relation index, ItemPointer tid, Datum *values, Ivffl } } + Vector* residual = NULL; + if (buildstate->enablePQ) { + ComputeResidual(buildstate, DatumGetVector(value), closestCenter); + if (buildstate->byResidual) { + residual = (Vector *)lfirst(buildstate->rlist->tail); + } + } + #ifdef IVFFLAT_KMEANS_DEBUG buildstate->inertia += minDistance; buildstate->listSums[closestCenter] += minDistance; @@ -181,6 +353,8 @@ static void AddTupleToSort(Relation index, ItemPointer tid, Datum *values, Ivffl slot->tts_isnull[1] = false; slot->tts_values[2] = value; slot->tts_isnull[2] = false; + slot->tts_values[3] = residual == NULL ? NULL : PointerGetDatum(residual); + slot->tts_isnull[3] = residual == NULL ? true : false; ExecStoreVirtualTuple(slot); /* @@ -210,6 +384,9 @@ static void BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values, return; } + Vector *vec = InitVector(buildstate->dimensions); + buildstate->rlist = lappend(buildstate->rlist, vec); + /* Use memory context since detoast can allocate */ oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); @@ -262,6 +439,7 @@ static void InsertTuples(Relation index, IvfflatBuildState *buildstate, ForkNumb GenericXLogState *state; BlockNumber startPage; BlockNumber insertPage; + Size pqcodesSize = buildstate->pqcodeSize; /* Can take a while, so ensure we can interrupt */ /* Needs to be called when no buffer locks are held */ @@ -276,9 +454,22 @@ static void InsertTuples(Relation index, IvfflatBuildState *buildstate, ForkNumb while (list == i) { /* Check for free space */ Size itemsz = MAXALIGN(IndexTupleSize(itup)); - if (PageGetFreeSpace(page) < itemsz) + if (PageGetFreeSpace(page) < itemsz + MAXALIGN(pqcodesSize)) IvfflatAppendPage(index, &buf, &page, &state, forkNum); + if (buildstate->enablePQ) { + bool isnull; + Size codesize = buildstate->params->pqM * sizeof(uint8); + uint8 *pqcode = (uint8 *)palloc(codesize); + Datum datum = buildstate->byResidual ? heap_slot_getattr(slot, 4, &isnull) : index_getattr(itup, 1, tupdesc, &isnull); + + IvfComputeVectorPQCode(DatumGetVector(datum)->x, buildstate->params, pqcode); + ((PageHeader)page)->pd_upper -= MAXALIGN(pqcodesSize); + errno_t rc = memcpy_s( + ((char *)page) + ((PageHeader)page)->pd_upper, pqcodesSize, (char *)pqcode, pqcodesSize); + securec_check_c(rc, "\0", "\0"); + } + /* Add the item */ if (PageAddItem(page, (Item)itup, itemsz, InvalidOffsetNumber, false, false) == InvalidOffsetNumber) elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); @@ -338,10 +529,11 @@ static void InitBuildState(IvfflatBuildState *buildstate, Relation heap, Relatio elog(ERROR, "dimensions must be greater than one for this opclass"); /* Create tuple description for sorting */ - buildstate->tupdesc = CreateTemplateTupleDesc(3, false); + buildstate->tupdesc = CreateTemplateTupleDesc(4, false); TupleDescInitEntry(buildstate->tupdesc, (AttrNumber)1, "list", INT4OID, -1, 0); TupleDescInitEntry(buildstate->tupdesc, (AttrNumber)2, "tid", TIDOID, -1, 0); TupleDescInitEntry(buildstate->tupdesc, (AttrNumber)3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0); + TupleDescInitEntry(buildstate->tupdesc, (AttrNumber)4, "residual", VECTOROID, -1, 0); buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc); @@ -358,6 +550,46 @@ static void InitBuildState(IvfflatBuildState *buildstate, Relation heap, Relatio buildstate->listCounts = palloc0(sizeof(int) * buildstate->lists); #endif buildstate->ivfleader = NULL; + + buildstate->enablePQ = IvfGetEnablePQ(index); + if (buildstate->enablePQ && !buildstate->typeInfo->supportPQ) { + ereport(ERROR, (errmsg("this data type cannot support ivfpq."))); + } + if (buildstate->enablePQ && !g_instance.pq_inited) { + ereport(ERROR, (errmsg("this instance has not currently loaded the pq dynamic library."))); + } + + buildstate->pqM = IvfGetPqM(index); + buildstate->pqKsub = IvfGetPqKsub(index); + buildstate->byResidual = IvfGetByResidual(index); + buildstate->rlist = NIL; + buildstate->residuals = NULL; + + if (buildstate->enablePQ) { + Size subItemsize = buildstate->typeInfo->itemSize(buildstate->dimensions / buildstate->pqM); + subItemsize = MAXALIGN(subItemsize); + buildstate->pqTableSize = buildstate->pqM * buildstate->pqKsub * subItemsize; + buildstate->pqTable = (char*)palloc0(buildstate->pqTableSize); + buildstate->pqcodeSize = buildstate->pqM * sizeof(uint8); + buildstate->params = InitIVFPQParamsInMemory(buildstate); + + if (buildstate->byResidual && + (buildstate->params->funcType == IVF_PQ_DIS_L2 || buildstate->params->funcType == IVF_PQ_DIS_COSINE)) { + buildstate->preComputeTableSize = buildstate->lists * buildstate->pqM * buildstate->pqKsub; + buildstate->preComputeTable = (float*)palloc0(buildstate->preComputeTableSize * sizeof(float)); + } else { + buildstate->preComputeTableSize = 0; + buildstate->preComputeTable = NULL; + } + } else { + buildstate->pqTable = NULL; + buildstate->pqTableSize = 0; + buildstate->pqcodeSize = 0; + buildstate->params = NULL; + buildstate->preComputeTableSize = 0; + buildstate->preComputeTable = NULL; + } + buildstate->pqDistanceTable = NULL; } /* @@ -366,6 +598,9 @@ static void InitBuildState(IvfflatBuildState *buildstate, Relation heap, Relatio static void FreeBuildState(IvfflatBuildState *buildstate) { VectorArrayFree(buildstate->centers); + if (buildstate->residuals) { + VectorArrayFree(buildstate->residuals); + } pfree(buildstate->listInfo); #ifdef IVFFLAT_KMEANS_DEBUG @@ -417,7 +652,7 @@ static void ComputeCenters(IvfflatBuildState *buildstate) /* * Create the metapage */ -static void CreateMetaPage(Relation index, int dimensions, int lists, ForkNumber forkNum) +static void CreateMetaPage(Relation index, IvfflatBuildState *buildstate, ForkNumber forkNum) { Buffer buf; Page page; @@ -431,8 +666,34 @@ static void CreateMetaPage(Relation index, int dimensions, int lists, ForkNumber metap = IvfflatPageGetMeta(page); metap->magicNumber = IVFFLAT_MAGIC_NUMBER; metap->version = IVFFLAT_VERSION; - metap->dimensions = dimensions; - metap->lists = lists; + metap->dimensions = buildstate->dimensions; + metap->lists = buildstate->lists; + + /* set PQ info */ + metap->enablePQ = buildstate->enablePQ; + metap->pqM = buildstate->pqM; + metap->byResidual = buildstate->byResidual; + metap->pqKsub = buildstate->pqKsub; + metap->pqcodeSize = buildstate->pqcodeSize; + metap->pqPreComputeTableSize = 0; + metap->pqPreComputeTableNblk = 0; + + if (buildstate->enablePQ) { + metap->pqTableSize = (uint32)buildstate->pqTableSize; + metap->pqTableNblk = (uint16)( + (metap->pqTableSize + IVF_PQTABLE_STORAGE_SIZE - 1) / IVF_PQTABLE_STORAGE_SIZE); + if (buildstate->byResidual && + (buildstate->params->funcType == IVF_PQ_DIS_L2 || buildstate->params->funcType == IVF_PQ_DIS_COSINE)) { + uint32 TableLen = buildstate->lists * buildstate->pqM * buildstate->pqKsub; + metap->pqPreComputeTableSize = (uint32)TableLen * sizeof(float); + metap->pqPreComputeTableNblk = (uint16)( + (metap->pqPreComputeTableSize + IVF_PQTABLE_STORAGE_SIZE - 1) / IVF_PQTABLE_STORAGE_SIZE); + } + } else { + metap->pqTableSize = 0; + metap->pqTableNblk = 0; + } + ((PageHeader)page)->pd_lower = ((char *)metap + sizeof(IvfflatMetaPageData)) - (char *)page; IvfflatCommitBuffer(buf, state); @@ -548,6 +809,8 @@ static double ParallelHeapScan(IvfflatBuildState *buildstate) buildstate->indtuples = ivfshared->indtuples; reltuples = ivfshared->reltuples; + buildstate->rlist =list_copy(ivfshared->rlist); + list_free(ivfshared->rlist); #ifdef IVFFLAT_KMEANS_DEBUG buildstate->inertia = ivfshared->inertia; #endif @@ -602,6 +865,20 @@ static void IvfflatParallelScanAndSort(IvfflatSpool *ivfspool, IvfflatShared *iv /* Record statistics */ SpinLockAcquire(&ivfshared->mutex); + + MemoryContext oldCtx = MemoryContextSwitchTo(ivfshared->tmpCtx); + ListCell *lc; + foreach (lc, buildstate.rlist) { + ListCell *lc2 = (ListCell *)palloc0(sizeof(ListCell)); + Vector *vec = InitVector(buildstate.dimensions); + int size = VECTOR_SIZE(buildstate.dimensions); + error_t rc = memcpy_s(vec, size, lc->data.ptr_value, size); + securec_check_c(rc, "\0", "\0"); + ivfshared->rlist = lappend(ivfshared->rlist, vec); + } + MemoryContextSwitchTo(oldCtx); + list_free_deep(buildstate.rlist); + ivfshared->nparticipantsdone++; ivfshared->reltuples += reltuples; ivfshared->indtuples += buildstate.indtuples; @@ -655,6 +932,8 @@ static void IvfflatParallelCleanup(const BgWorkerContext *bwc) Assert(ivfshared->sharedsort); SharedFileSetDeleteAll(&ivfshared->sharedsort->fileset); pfree_ext(ivfshared->sharedsort); + + MemoryContextDelete(ivfshared->tmpCtx); } static IvfflatShared *IvfflatParallelInitshared(IvfflatBuildState *buildstate, int workmem, int scantuplesortstates) @@ -697,6 +976,8 @@ static IvfflatShared *IvfflatParallelInitshared(IvfflatBuildState *buildstate, i securec_check(rc, "\0", "\0"); ivfshared->ivfcenters = (Vector *)ivfcenters; + ivfshared->tmpCtx = + AllocSetContextCreate(CurrentMemoryContext, "Ivfflat build temporary context", ALLOCSET_DEFAULT_SIZES); return ivfshared; } @@ -824,6 +1105,14 @@ static void CreateEntryPages(IvfflatBuildState *buildstate, ForkNumber forkNum) /* Sort */ IvfflatBench("sort tuples", tuplesort_performsort(buildstate->sortstate)); + /* Build PQTable by residusal */ + if (buildstate->enablePQ) { + CopyResidaulFromList(buildstate); + ComputeIvfPQ(buildstate); + if (buildstate->byResidual && + (buildstate->params->funcType == IVF_PQ_DIS_L2 || buildstate->params->funcType == IVF_PQ_DIS_COSINE)) + ComputePreTable(buildstate); + } /* Load */ IvfflatBench("load tuples", InsertTuples(buildstate->index, buildstate, forkNum)); @@ -848,11 +1137,18 @@ static void BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo, Ivff ComputeCenters(buildstate); /* Create pages */ - CreateMetaPage(index, buildstate->dimensions, buildstate->lists, forkNum); + CreateMetaPage(index, buildstate, forkNum); + + CreatePQPages(buildstate, forkNum); + CreateListPages(index, buildstate->centers, buildstate->dimensions, buildstate->lists, forkNum, &buildstate->listInfo); CreateEntryPages(buildstate, forkNum); + if (buildstate->enablePQ) { + IvfFlushPQInfo(buildstate); + } + /* Write WAL for initialization fork since GenericXLog functions do not */ if (forkNum == INIT_FORKNUM) LogNewpageRange(index, forkNum, 0, RelationGetNumberOfBlocksInFork(index, forkNum), true); diff --git a/src/gausskernel/storage/access/datavec/ivfflat.cpp b/src/gausskernel/storage/access/datavec/ivfflat.cpp index 2602feea20..0d165885f3 100644 --- a/src/gausskernel/storage/access/datavec/ivfflat.cpp +++ b/src/gausskernel/storage/access/datavec/ivfflat.cpp @@ -112,6 +112,10 @@ static bytea *ivfflatoptions_internal(Datum reloptions, bool validate) { static const relopt_parse_elt tab[] = { {"lists", RELOPT_TYPE_INT, offsetof(IvfflatOptions, lists)}, + {"enable_pq", RELOPT_TYPE_BOOL, offsetof(IvfflatOptions, enablePQ)}, + {"pq_m", RELOPT_TYPE_INT, offsetof(IvfflatOptions, pqM)}, + {"pq_ksub", RELOPT_TYPE_INT, offsetof(IvfflatOptions, pqKsub)}, + {"by_residual", RELOPT_TYPE_BOOL, offsetof(IvfflatOptions, byResidual)}, {"parallel_workers", RELOPT_TYPE_INT, offsetof(StdRdOptions, parallel_workers)}}; relopt_value *options; diff --git a/src/gausskernel/storage/access/datavec/ivfutils.cpp b/src/gausskernel/storage/access/datavec/ivfutils.cpp index f73a78f662..9961ab8168 100644 --- a/src/gausskernel/storage/access/datavec/ivfutils.cpp +++ b/src/gausskernel/storage/access/datavec/ivfutils.cpp @@ -214,6 +214,7 @@ const IvfflatTypeInfo *IvfflatGetTypeInfo(Relation index) if (procinfo == NULL) { static const IvfflatTypeInfo typeInfo = {.maxDimensions = IVFFLAT_MAX_DIM, + .supportPQ = true, .normalize = l2_normalize, .itemSize = VectorItemSize, .updateCenter = VectorUpdateCenter, @@ -229,6 +230,7 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_support); Datum ivfflat_halfvec_support(PG_FUNCTION_ARGS) { static const IvfflatTypeInfo typeInfo = {.maxDimensions = IVFFLAT_MAX_DIM * 2, + .supportPQ = false, .normalize = halfvec_l2_normalize, .itemSize = HalfvecItemSize, .updateCenter = HalfvecUpdateCenter, @@ -241,6 +243,7 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_bit_support); Datum ivfflat_bit_support(PG_FUNCTION_ARGS) { static const IvfflatTypeInfo typeInfo = {.maxDimensions = IVFFLAT_MAX_DIM * 32, + .supportPQ = false, .normalize = NULL, .itemSize = BitItemSize, .updateCenter = BitUpdateCenter, @@ -248,3 +251,162 @@ Datum ivfflat_bit_support(PG_FUNCTION_ARGS) PG_RETURN_POINTER(&typeInfo); }; + +int getIVFPQfunctionType(FmgrInfo *procinfo, FmgrInfo *normprocinfo) +{ + if (procinfo->fn_oid == 8431) { + return IVF_PQ_DIS_L2; + } else if (procinfo->fn_oid == 8434) { + if (normprocinfo == NULL) { + return IVF_PQ_DIS_IP; + } else { + return IVF_PQ_DIS_COSINE; + } + } else { + ereport(ERROR, (errmsg("current data type or distance type can't support ivfflatpq."))); + return -1; + } +} + +/* +* Get the info related to pqTable in metapage +*/ +void IvfGetPQInfoFromMetaPage(Relation index, uint16 *pqTableNblk, uint32 *pqTableSize, + uint16 *pqPreComputeTableNblk, uint32 *pqPreComputeTableSize) +{ + Buffer buf; + Page page; + IvfflatMetaPage metap; + + buf = ReadBuffer(index, IVFFLAT_METAPAGE_BLKNO); + LockBuffer(buf, BUFFER_LOCK_SHARE); + page = BufferGetPage(buf); + metap = IvfflatPageGetMeta(page); + + PG_TRY(); + { + if (unlikely(metap->magicNumber != IVFFLAT_MAGIC_NUMBER)) { + elog(ERROR, "ivfflat index is not valid"); + } + } + PG_CATCH(); + { + UnlockReleaseBuffer(buf); + PG_RE_THROW(); + } + PG_END_TRY(); + + if (pqTableNblk != NULL) { + *pqTableNblk = metap->pqTableNblk; + } + if (pqTableSize != NULL) { + *pqTableSize = metap->pqTableSize; + } + if (pqPreComputeTableNblk != NULL) { + *pqPreComputeTableNblk = metap->pqPreComputeTableNblk; + } + if (pqPreComputeTableSize != NULL) { + *pqPreComputeTableSize = metap->pqPreComputeTableSize; + } + + UnlockReleaseBuffer(buf); +} + +/* + * Get whether to enable PQ + */ +bool IvfGetEnablePQ(Relation index) +{ + IvfflatOptions *opts = (IvfflatOptions *)index->rd_options; + + if (opts) { + return opts->enablePQ; + } + + return GENERIC_DEFAULT_ENABLE_PQ; +} + +/* + * Get the number of subquantizer + */ +int IvfGetPqM(Relation index) +{ + IvfflatOptions *opts = (IvfflatOptions *)index->rd_options; + + if (opts) { + return opts->pqM; + } + + return GENERIC_DEFAULT_PQ_M; +} + +/* + * Get the number of centroids for each subquantizer + */ +int IvfGetPqKsub(Relation index) +{ + IvfflatOptions *opts = (IvfflatOptions *)index->rd_options; + + if (opts) { + return opts->pqKsub; + } + + return GENERIC_DEFAULT_PQ_KSUB; +} + +/* + * Get whether to use residual + */ +int IvfGetByResidual(Relation index) +{ + IvfflatOptions *opts = (IvfflatOptions *)index->rd_options; + + if (opts) { + return opts->byResidual; + } + + return IVFFLAT_DEFAULT_PQ_RESIDUAL; +} + +void IvfFlushPQInfoInternal(Relation index, char* table, BlockNumber startBlkno, uint16 nblks, uint32 totalSize) +{ + Buffer buf; + Page page; + uint32 curFlushSize; + for (uint16 i = 0; i < nblks; i++) { + curFlushSize = (i == nblks - 1) ? + (totalSize - i * IVF_PQTABLE_STORAGE_SIZE) : IVF_PQTABLE_STORAGE_SIZE; + buf = ReadBufferExtended(index, MAIN_FORKNUM, startBlkno + i, RBM_NORMAL, NULL); + LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); + page = BufferGetPage(buf); + errno_t err = memcpy_s(PageGetContents(page), curFlushSize, + table + i * IVF_PQTABLE_STORAGE_SIZE, curFlushSize); + securec_check(err, "\0", "\0"); + MarkBufferDirty(buf); + UnlockReleaseBuffer(buf); + } +} + +/* +* Flush PQ table into page during index building +*/ +void IvfFlushPQInfo(IvfflatBuildState *buildstate) +{ + Relation index = buildstate->index; + char* pqTable = buildstate->pqTable; + float* preComputeTable = buildstate->preComputeTable; + uint16 pqTableNblk; + uint32 pqTableSize; + uint16 pqPrecomputeTableNblk; + uint32 pqPrecomputeTableSize; + + IvfGetPQInfoFromMetaPage(index, &pqTableNblk, &pqTableSize, &pqPrecomputeTableNblk, &pqPrecomputeTableSize); + + /* Flush pq table */ + IvfFlushPQInfoInternal(index, pqTable, IVF_PQTABLE_START_BLKNO, pqTableNblk, pqTableSize); + if (buildstate->byResidual && buildstate->params->funcType == IVF_PQ_DIS_L2) { + /* Flush pq distance table */ + IvfFlushPQInfoInternal(index, (char*)preComputeTable, + IVF_PQTABLE_START_BLKNO + pqTableNblk, pqPrecomputeTableNblk, pqPrecomputeTableSize); + } +} diff --git a/src/gausskernel/storage/access/datavec/utils.cpp b/src/gausskernel/storage/access/datavec/utils.cpp index 6d1834f803..b877cf5718 100644 --- a/src/gausskernel/storage/access/datavec/utils.cpp +++ b/src/gausskernel/storage/access/datavec/utils.cpp @@ -29,6 +29,8 @@ #include "access/datavec/bitvec.h" #include "access/datavec/vector.h" +pq_func_t g_pq_func = {0}; + Size VectorItemSize(int dimensions) { return VECTOR_SIZE(dimensions); diff --git a/src/include/access/datavec/hnsw.h b/src/include/access/datavec/hnsw.h index 7968c226e6..69703b1c0f 100644 --- a/src/include/access/datavec/hnsw.h +++ b/src/include/access/datavec/hnsw.h @@ -74,13 +74,6 @@ #define HNSW_MIN_THRESHOLD 160 #define HNSW_MAX_THRESHOLD INT32_MAX #define HNSW_MAX_EF_SEARCH 1000000 -#define HNSW_DEFAULT_ENABLE_PQ false -#define HNSW_DEFAULT_PQ_M 8 -#define HNSW_MIN_PQ_M 1 -#define HNSW_MAX_PQ_M HNSW_MAX_DIM -#define HNSW_DEFAULT_PQ_KSUB 256 -#define HNSW_MIN_PQ_KSUB 1 -#define HNSW_MAX_PQ_KSUB 256 #define HNSW_PQMODE_ADC 1 #define HNSW_PQMODE_SDC 2 @@ -109,8 +102,8 @@ /* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */ #define PROGRESS_HNSW_PHASE_LOAD 2 -#define HNSWPQ_SUCCESS 0 -#define HNSWPQ_ERROR (-1) +#define PQ_SUCCESS 0 +#define PQ_ERROR (-1) #define HNSWPQ_MAX_PATH_LEN 4096 #ifndef MAX_PATH_LEN @@ -119,8 +112,8 @@ #define HNSWPQ_DEFAULT_TARGET_ROWS 300 -#define HNSWPQ_ENV_PATH "DATAVEC_HNSWPQ_LIB_PATH" -#define HNSWPQ_SO_NAME "libkvecturbo.so" +#define PQ_ENV_PATH "DATAVEC_HNSWPQ_LIB_PATH" +#define PQ_SO_NAME "libkvecturbo.so" #define HNSW_MAX_SIZE \ (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData)) @@ -385,15 +378,6 @@ typedef struct HnswTypeInfo { void (*checkValue)(Pointer v); } HnswTypeInfo; -typedef struct PQParams { - int pqM; - int pqKsub; - int funcType; - int dim; - size_t subItemSize; - char *pqTable; -} PQParams; - typedef struct HnswBuildState { /* Info */ Relation heap; @@ -621,18 +605,6 @@ typedef struct Candidate { uint8 heaptidsLength; } Candidate; -typedef struct st_hnswpq_func { - bool inited; - void *handle; - int (*ComputePQTable)(VectorArray samples, PQParams *params); - int (*ComputeVectorPQCode)(float *vector, const PQParams *params, uint8 *pqCode); - int (*GetPQDistanceTableSdc)(const PQParams *params, float *pqDistanceTable); - int (*GetPQDistanceTableAdc)(float *vector, const PQParams *params, float *pqDistanceTable); - int (*GetPQDistance)(const uint8 *basecode, const uint8 *querycode, const PQParams *params, - const float *pqDistanceTable, float *PQDistance); -} hnswpq_func_t; - - /* Methods */ int HnswGetM(Relation index); int HnswGetEfConstruction(Relation index); diff --git a/src/include/access/datavec/ivfflat.h b/src/include/access/datavec/ivfflat.h index 6949da8899..16163a857b 100644 --- a/src/include/access/datavec/ivfflat.h +++ b/src/include/access/datavec/ivfflat.h @@ -63,6 +63,7 @@ #define IVFFLAT_MIN_LISTS 1 #define IVFFLAT_MAX_LISTS 32768 #define IVFFLAT_DEFAULT_PROBES 1 +#define IVFFLAT_DEFAULT_PQ_RESIDUAL false /* Build phases */ /* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */ @@ -93,6 +94,19 @@ #define RandomDouble() (((double)random()) / MAX_RANDOM_VALUE) #define RandomInt() random() +#define IVF_PQMODE_ADC 1 +#define IVF_PQMODE_SDC 2 +#define IVF_PQMODE_DEFAULT IVF_PQMODE_ADC +#define IVF_PQ_DIS_L2 1 +#define IVF_PQ_DIS_IP 2 +#define IVF_PQ_DIS_COSINE 3 + +/* Preserved page numbers */ +#define IVF_METAPAGE_BLKNO 0 +#define IVF_HEAD_BLKNO 1 /* first element page */ +#define IVF_PQTABLE_START_BLKNO 1 /* pqtable start page */ +#define IVF_PQTABLE_STORAGE_SIZE (uint16)(6 * 1024) /* pqtable storage size in each page */ + typedef struct ListInfo { BlockNumber blkno; OffsetNumber offno; @@ -102,6 +116,10 @@ typedef struct ListInfo { typedef struct IvfflatOptions { int32 vl_len_; /* varlena header (do not touch directly!) */ int lists; /* number of lists */ + bool enablePQ; + int pqM; + int pqKsub; + bool byResidual; /* whether to quantify by residual */ } IvfflatOptions; typedef struct IvfflatSpool { @@ -126,8 +144,12 @@ typedef struct IvfflatShared { Sharedsort *sharedsort; Vector *ivfcenters; + List *rlist; int workmem; + /* Memory */ + MemoryContext tmpCtx; + #ifdef IVFFLAT_KMEANS_DEBUG double inertia; #endif @@ -144,6 +166,7 @@ typedef struct IvfflatLeader { typedef struct IvfflatTypeInfo { int maxDimensions; + bool supportPQ; Datum (*normalize)(PG_FUNCTION_ARGS); Size (*itemSize)(int dimensions); void (*updateCenter)(Pointer v, int dimensions, const float *x); @@ -173,7 +196,9 @@ typedef struct IvfflatBuildState { /* Variables */ VectorArray samples; + VectorArray residuals; VectorArray centers; + List *rlist; ListInfo *listInfo; #ifdef IVFFLAT_KMEANS_DEBUG @@ -197,6 +222,19 @@ typedef struct IvfflatBuildState { /* Parallel builds */ IvfflatLeader *ivfleader; + + /* PQ info */ + bool enablePQ; + int pqM; + int pqKsub; + bool byResidual = false ; + char *pqTable; + Size pqTableSize; + float *pqDistanceTable; + uint16 pqcodeSize; + PQParams *params; + float *preComputeTable; + uint64 preComputeTableSize; } IvfflatBuildState; typedef struct IvfflatMetaPageData { @@ -204,6 +242,17 @@ typedef struct IvfflatMetaPageData { uint32 version; uint16 dimensions; uint16 lists; + + /* PQ info */ + bool enablePQ; + bool byResidual; + uint16 pqM; + uint16 pqKsub; + uint16 pqcodeSize; + uint32 pqTableSize; + uint16 pqTableNblk; + uint64 pqPreComputeTableSize; + uint16 pqPreComputeTableNblk; } IvfflatMetaPageData; typedef IvfflatMetaPageData *IvfflatMetaPage; @@ -273,6 +322,25 @@ void IvfflatInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLo PGDLLEXPORT void IvfflatParallelBuildMain(const BgWorkerContext *bwc); const IvfflatTypeInfo *IvfflatGetTypeInfo(Relation index); +bool IvfGetEnablePQ(Relation index); +int IvfGetPqM(Relation index); +int IvfGetPqKsub(Relation index); +int IvfGetByResidual(Relation index); + + +void IvfGetPQInfoFromMetaPage(Relation index, uint16 *pqTableNblk, uint32 *pqTableSize, + uint16 *pqPreComputeTableNblk, uint32 *pqPreComputeTableSize); +int getIVFPQfunctionType(FmgrInfo *procinfo, FmgrInfo *normprocinfo); +void IvfFlushPQInfoInternal(Relation index, char* table, BlockNumber startBlkno, uint16 nblks, uint32 totalSize); +void IvfFlushPQInfo(IvfflatBuildState *buildstate); + +int IvfComputePQTable(VectorArray samples, PQParams *params); +int IvfComputeVectorPQCode(float *vector, const PQParams *params, uint8 *pqCode); +int IvfGetPQDistanceTableAdc(float *vector, const PQParams *params, float *pqDistanceTable); +int IvfGetPQDistance(const uint8 *basecode, const uint8 *querycode, const PQParams *params, + const float *pqDistanceTable, float *pqDistance); + + Datum ivfflathandler(PG_FUNCTION_ARGS); Datum ivfflatbuild(PG_FUNCTION_ARGS); Datum ivfflatbuildempty(PG_FUNCTION_ARGS); diff --git a/src/include/access/datavec/utils.h b/src/include/access/datavec/utils.h index 7620b2e915..88f51ae236 100644 --- a/src/include/access/datavec/utils.h +++ b/src/include/access/datavec/utils.h @@ -4,6 +4,14 @@ #include "fmgr/fmgr_comp.h" #include +#define GENERIC_DEFAULT_ENABLE_PQ false +#define GENERIC_DEFAULT_PQ_M 8 +#define GENERIC_MIN_PQ_M 1 +#define GENERIC_MAX_PQ_M HNSW_MAX_DIM +#define GENERIC_DEFAULT_PQ_KSUB 256 +#define GENERIC_MIN_PQ_KSUB 1 +#define GENERIC_MAX_PQ_KSUB 256 + typedef struct VectorArrayData { int length; int maxlen; @@ -12,10 +20,32 @@ typedef struct VectorArrayData { char *items; } VectorArrayData; +typedef struct PQParams { + int pqM; + int pqKsub; + int funcType; + int dim; + int pqMode; + size_t subItemSize; + char *pqTable; +} PQParams; + #define VECTOR_ARRAY_SIZE(_length, _size) (sizeof(VectorArrayData) + (_length) * MAXALIGN(_size)) typedef VectorArrayData * VectorArray; +typedef struct st_pq_func { + bool inited; + void *handle; + int (*ComputePQTable)(VectorArray samples, PQParams *params); + int (*ComputeVectorPQCode)(float *vector, const PQParams *params, uint8 *pqCode); + int (*GetPQDistanceTableSdc)(const PQParams *params, float *pqDistanceTable); + int (*GetPQDistanceTableAdc)(float *vector, const PQParams *params, float *pqDistanceTable); + int (*GetPQDistance)(const uint8 *basecode, const uint8 *querycode, const PQParams *params, + const float *pqDistanceTable, float *pqDistance); +} pq_func_t; +extern pq_func_t g_pq_func; + static inline Pointer VectorArrayGet(VectorArray arr, int offset) { return ((char *) arr->items) + (offset * arr->itemsize); @@ -39,6 +69,6 @@ void BitSumCenter(Pointer v, float *x); VectorArray VectorArrayInit(int maxlen, int dimensions, Size itemsize); void VectorArrayFree(VectorArray arr); -int HNSWPQInit(); -void HNSWPQUinit(); +int PQInit(); +void PQUinit(); #endif \ No newline at end of file diff --git a/src/include/access/datavec/vector.h b/src/include/access/datavec/vector.h index 750b34f24d..01796d5eb8 100644 --- a/src/include/access/datavec/vector.h +++ b/src/include/access/datavec/vector.h @@ -46,9 +46,11 @@ typedef struct Vector { #endif VECTOR_TARGET_CLONES float VectorL2SquaredDistance(int dim, float *ax, float *bx); +VECTOR_TARGET_CLONES float VectorInnerProduct(int dim, float *ax, float *bx); Vector *InitVector(int dim); void PrintVector(char *msg, Vector *vector); int vector_cmp_internal(Vector *a, Vector *b); +void VectorMadd(size_t n, const float *ax, float bf, const float *bx, float *cx); void LogNewpageRange(Relation rel, ForkNumber forknum, BlockNumber startblk, BlockNumber endblk, bool page_std); int PlanCreateIndexWorkers(Relation heapRelation, IndexInfo *indexInfo); diff --git a/src/include/knl/knl_guc/knl_instance_attr_storage.h b/src/include/knl/knl_guc/knl_instance_attr_storage.h index e2218425b2..741a31089e 100755 --- a/src/include/knl/knl_guc/knl_instance_attr_storage.h +++ b/src/include/knl/knl_guc/knl_instance_attr_storage.h @@ -253,7 +253,7 @@ typedef struct knl_instance_attr_storage { int uwal_truncate_interval; bool uwal_async_append_switch; - bool enable_hnswpq; + bool enable_pq; int num_slru_buffers[SLRU_BUFFER_KIND]; char* num_slru_buffers_str; char* lmemfabric_client_path; diff --git a/src/include/knl/knl_instance.h b/src/include/knl/knl_instance.h index 9052e49a29..f94dbba362 100755 --- a/src/include/knl/knl_instance.h +++ b/src/include/knl/knl_instance.h @@ -1552,7 +1552,7 @@ typedef struct knl_instance_context { knl_g_listen_context listen_cxt; knl_g_datadir_context datadir_cxt; knl_g_dms_context dms_cxt; - bool hnswpq_inited; + bool pq_inited; #ifdef USE_SPQ knl_g_spq_context spq_cxt; #endif -- Gitee From 94392bfa88b42c24e54ac6045de6818e079662d6 Mon Sep 17 00:00:00 2001 From: wangjingyuan <1577039175@qq.com> Date: Mon, 10 Feb 2025 21:49:12 +0800 Subject: [PATCH 2/2] test --- src/bin/gs_guc/cluster_guc.conf | 1 + src/common/backend/utils/adt/vector.cpp | 160 +++++++++++ src/common/backend/utils/misc/guc/guc_sql.cpp | 13 + .../process/postmaster/postmaster.cpp | 2 +- .../process/threadpool/knl_session.cpp | 1 + .../storage/access/common/reloptions.cpp | 2 +- .../storage/access/datavec/hnswutils.cpp | 2 +- .../storage/access/datavec/ivfinsert.cpp | 89 +++++- .../storage/access/datavec/ivfscan.cpp | 253 +++++++++++++++++- .../storage/access/datavec/ivfutils.cpp | 173 +++++++++++- .../storage/access/datavec/ivfvacuum.cpp | 5 +- src/include/access/datavec/hnsw.h | 2 - src/include/access/datavec/ivfflat.h | 35 ++- src/include/access/datavec/vector.h | 2 + src/include/knl/knl_session.h | 1 + src/include/utils/rel.h | 2 +- 16 files changed, 727 insertions(+), 16 deletions(-) diff --git a/src/bin/gs_guc/cluster_guc.conf b/src/bin/gs_guc/cluster_guc.conf index 7908da000d..513a2e7a4a 100755 --- a/src/bin/gs_guc/cluster_guc.conf +++ b/src/bin/gs_guc/cluster_guc.conf @@ -849,6 +849,7 @@ enable_gazelle_performance_mode|bool|0,0|NULL|NULL| enable_aggr_coerce_type|bool|0,0|NULL|NULL| enable_default_local_index|bool|0,0|NULL|NULL| enable_pq|bool|0,0|NULL|NULL| +ivfpq_kreorder|int|0,2147483647|NULL|NULL| [cmserver] log_dir|string|0,0|NULL|NULL| log_file_size|int|0,2047|MB|NULL| diff --git a/src/common/backend/utils/adt/vector.cpp b/src/common/backend/utils/adt/vector.cpp index 0ec9734b14..ceee5faf5d 100644 --- a/src/common/backend/utils/adt/vector.cpp +++ b/src/common/backend/utils/adt/vector.cpp @@ -1565,6 +1565,166 @@ void VectorMadd(size_t n, const float *ax, float bf, const float *bx, float *cx) } #endif +#ifdef __aarch64__ +struct ElementOpL2 { + static float op(float x, float y) { + float tmp = x - y; + return tmp * tmp; + } + + static float32x4_t op(float32x4_t x, float32x4_t y) { + float32x4_t tmp = vsubq_f32(x, y); + return vmulq_f32(tmp, tmp); + } +}; + +struct ElementOpIP { + static float op(float x, float y) { + return x * y; + } + + static float32x4_t op(float32x4_t x, float32x4_t y) { + return vsubq_f32(x, y); + } +}; + +template +void VectorOpNYD4(size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ + float32x4_t x0 = vld1q_f32(x); + float *y; + __builtin_prefetch(pqTable, 0, 3); + + size_t i; + for (i = 0; i < ny; i++) { + y = DatumGetVector(pqTable + (offset + i) * subSize)->x; + float32x4_t accu = ElementOp::op(x0, vld1q_f32(y)); + accu = vaddq_f32(accu, accu); + accu = vaddq_f32(accu, accu); + dis[i] = vgetq_lane_f32(accu, 0); + } +} + +template +void VectorOpNYD8(size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ + /* neon support 128 bit, float is 32 bit, 4 float one batch*/ + int batch = 4; + float32x4_t x0 = vld1q_f32(x); + float32x4_t x1 = vld1q_f32(x + batch); + float *y; + __builtin_prefetch(pqTable, 0, 3); + + size_t i; + for (i = 0; i < ny; i++) { + y = DatumGetVector(pqTable + (offset + i) * subSize)->x; + float32x4_t accu = ElementOp::op(x0, vld1q_f32(y)); + y += batch; + accu = vaddq_f32(accu, ElementOp::op(x1, vld1q_f32(y))); + accu = vaddq_f32(accu, accu); + accu = vaddq_f32(accu, accu); + dis[i] = vgetq_lane_f32(accu, 0); + } +} + +template +void VectorOpNYD16(size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ + /* neon support 128 bit, float is 32 bit, 4 float one batch*/ + int batch = 4; + float32x4_t x0 = vld1q_f32(x); + float32x4_t x1 = vld1q_f32(x + batch); + float32x4_t x2 = vld1q_f32(x + batch * 2); + float32x4_t x3 = vld1q_f32(x + batch * 3); + float *y; + __builtin_prefetch(pqTable, 0, 3); + + size_t i; + for (i = 0; i < ny; i++) { + y = DatumGetVector(pqTable + (offset + i) * subSize)->x; + float32x4_t accu = ElementOp::op(x0, vld1q_f32(y)); + y += batch; + accu = vaddq_f32(accu, ElementOp::op(x1, vld1q_f32(y))); + y += batch; + accu = vaddq_f32(accu, ElementOp::op(x2, vld1q_f32(y))); + y += batch; + accu = vaddq_f32(accu, ElementOp::op(x3, vld1q_f32(y))); + accu = vaddq_f32(accu, accu); + accu = vaddq_f32(accu, accu); + dis[i] = vgetq_lane_f32(accu, 0); + } +} +#endif + +void VectorL2SquaredDistanceNYRef(size_t d, size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ + float *y; + for (size_t i = 0; i < ny; i++) { + y = DatumGetVector(pqTable + (offset + i) * subSize)->x; + dis[i] = VectorL2SquaredDistance(d, x, y); + y += d; + } +} + +#ifdef __aarch64__ +void VectorL2SquaredDistanceNY(size_t d, size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ +#define DISPATCH(dval) \ + case dval: \ + VectorOpNYD##dval(ny, x, pqTable, subSize, offset, dis); \ + return; + + switch (d) { + DISPATCH(4) + DISPATCH(8) + DISPATCH(16) + default: + VectorL2SquaredDistanceNYRef(d, ny, x, pqTable, subSize, offset, dis); + return; + } +#undef DISPATCH +} +#else +void VectorL2SquaredDistanceNY(size_t d, size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ + VectorL2SquaredDistanceNYRef(d, ny, x, pqTable, subSize, offset, dis); +} +#endif + +void VectorInnerProductNYRef(size_t d, size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ + float *y; + for (size_t i = 0; i < ny; i++) { + y = DatumGetVector(pqTable + (offset + i) * subSize)->x; + dis[i] = VectorInnerProduct(d, x, y); + } +} + +#ifdef __aarch64__ +void VectorInnerProductNY(size_t d, size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ +#define DISPATCH(dval) \ + case dval: \ + VectorOpNYD##dval(ny, x, pqTable, subSize, offset, dis); \ + return; + + switch (d) { + DISPATCH(4) + DISPATCH(8) + DISPATCH(16) + default: + VectorInnerProductNYRef(d, ny, x, pqTable, subSize, offset, dis); + return; + } +#undef DISPATCH +} +#else +void VectorInnerProductNY(size_t d, size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis) +{ + VectorInnerProductNYRef(d, ny, x, pqTable, subSize, offset, dis); +} +#endif + /* * WAL-log a range of blocks in a relation. * diff --git a/src/common/backend/utils/misc/guc/guc_sql.cpp b/src/common/backend/utils/misc/guc/guc_sql.cpp index 3c4a2e092d..831a6032c4 100755 --- a/src/common/backend/utils/misc/guc/guc_sql.cpp +++ b/src/common/backend/utils/misc/guc/guc_sql.cpp @@ -2645,6 +2645,19 @@ static void InitSqlConfigureNamesInt() NULL, NULL, NULL}, + {{"ivfpq_kreorder", + PGC_USERSET, + NODE_ALL, + QUERY_TUNING_OTHER, + gettext_noop("Sets the number of samples that need to be reordered after IVFPQ."), + NULL}, + &u_sess->datavec_ctx.ivfpq_kreorder, + 0, + 0, + INT_MAX, + NULL, + NULL, + NULL}, #endif /* End-of-list marker */ {{NULL, diff --git a/src/gausskernel/process/postmaster/postmaster.cpp b/src/gausskernel/process/postmaster/postmaster.cpp index 38f42017d1..eaa908af04 100644 --- a/src/gausskernel/process/postmaster/postmaster.cpp +++ b/src/gausskernel/process/postmaster/postmaster.cpp @@ -3151,7 +3151,7 @@ int PostmasterMain(int argc, char* argv[]) } } - /* init datavec hnswpq */ + /* init datavec pq */ if (g_instance.attr.attr_storage.enable_pq) { int ret = PQInit(); if (ret != 0) { diff --git a/src/gausskernel/process/threadpool/knl_session.cpp b/src/gausskernel/process/threadpool/knl_session.cpp index 92434f8926..f764cdb96c 100755 --- a/src/gausskernel/process/threadpool/knl_session.cpp +++ b/src/gausskernel/process/threadpool/knl_session.cpp @@ -1499,6 +1499,7 @@ static void knl_u_datavec_init(knl_u_datavec_context* datavec_cxt) datavec_cxt->hnsw_ef_search = 0; datavec_cxt->hnsw_earlystop_threshold = 0; datavec_cxt->ivfflat_probes = 0; + datavec_cxt->ivfpq_kreorder = 0; } #ifdef ENABLE_HTAP diff --git a/src/gausskernel/storage/access/common/reloptions.cpp b/src/gausskernel/storage/access/common/reloptions.cpp index c6e5eb3b69..6f07c0601d 100644 --- a/src/gausskernel/storage/access/common/reloptions.cpp +++ b/src/gausskernel/storage/access/common/reloptions.cpp @@ -126,7 +126,7 @@ static relopt_bool boolRelOpts[] = { false}, {{"deduplication", "Enables \"deduplication\" feature for btree index", RELOPT_KIND_BTREE}, false}, {{"enable_pq", "Whether to enable PQ", RELOPT_KIND_HNSW | RELOPT_KIND_IVFFLAT }, GENERIC_DEFAULT_ENABLE_PQ }, - {{"by_residual", "Whether to use residual during IVFPQ", RELOPT_KIND_IVFFLAT}, IVFFLAT_DEFAULT_PQ_RESIDUAL}, + {{"by_residual", "Whether to use residual during IVFPQ", RELOPT_KIND_IVFFLAT}, IVFPQ_DEFAULT_RESIDUAL}, /* list terminator */ {{NULL}}}; diff --git a/src/gausskernel/storage/access/datavec/hnswutils.cpp b/src/gausskernel/storage/access/datavec/hnswutils.cpp index 75c2498349..d1ce3b9997 100644 --- a/src/gausskernel/storage/access/datavec/hnswutils.cpp +++ b/src/gausskernel/storage/access/datavec/hnswutils.cpp @@ -1522,7 +1522,7 @@ void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entry } /* -* Get the info related to pqTable in append metapage +* Get the info related to pqTable in metapage */ void HnswGetPQInfoFromMetaPage(Relation index, uint16 *pqTableNblk, uint32 *pqTableSize, uint16 *pqDisTableNblk, uint32 *pqDisTableSize) diff --git a/src/gausskernel/storage/access/datavec/ivfinsert.cpp b/src/gausskernel/storage/access/datavec/ivfinsert.cpp index 152c97c119..84444ccbae 100644 --- a/src/gausskernel/storage/access/datavec/ivfinsert.cpp +++ b/src/gausskernel/storage/access/datavec/ivfinsert.cpp @@ -36,7 +36,10 @@ static void FindInsertPage(Relation index, Datum *values, BlockNumber *insertPage, ListInfo *listInfo) { double minDistance = DBL_MAX; - BlockNumber nextblkno = IVFFLAT_HEAD_BLKNO; + uint16 pqTableNblk; + uint16 pqDisTableNblk; + IvfGetPQInfoFromMetaPage(index, &pqTableNblk, NULL, &pqDisTableNblk, NULL); + BlockNumber nextblkno = IVFPQTABLE_START_BLKNO + pqTableNblk + pqDisTableNblk; FmgrInfo *procinfo; Oid collation; @@ -79,6 +82,58 @@ static void FindInsertPage(Relation index, Datum *values, BlockNumber *insertPag } } +static void InitPQParamsOnDisk(Relation index, PQParams *params, int dim, bool *enablePQ, + bool *byResidual, uint16 *pqcodeSize) +{ + Buffer buf; + Page page; + IvfflatMetaPage metap; + const IvfflatTypeInfo *typeInfo = IvfflatGetTypeInfo(index); + + buf = ReadBuffer(index, IVFFLAT_METAPAGE_BLKNO); + LockBuffer(buf, BUFFER_LOCK_SHARE); + page = BufferGetPage(buf); + metap = IvfflatPageGetMeta(page); + + PG_TRY(); + { + if (unlikely(metap->magicNumber != IVFFLAT_MAGIC_NUMBER)) { + elog(ERROR, "ivfflat index is not valid"); + } + } + PG_CATCH(); + { + UnlockReleaseBuffer(buf); + PG_RE_THROW(); + } + PG_END_TRY(); + + *enablePQ = metap->enablePQ; + params->pqM = metap->pqM; + params->pqKsub = metap->pqKsub; + *byResidual = metap->byResidual; + *pqcodeSize = metap->pqcodeSize; + UnlockReleaseBuffer(buf); + + if (*enablePQ) { + FmgrInfo *procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC); + FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC); + params->funcType = getIVFPQfunctionType(procinfo, normprocinfo); + params->dim = dim; + params->subItemSize = typeInfo->itemSize(dim / params->pqM); + + /* Now save pqTable in the relcache entry. */ + if (index->pqTable == NULL) { + MemoryContext oldcxt = MemoryContextSwitchTo(index->rd_indexcxt); + index->pqTable = IVFPQLoadPQtable(index); + (void)MemoryContextSwitchTo(oldcxt); + } + params->pqTable = index->pqTable; + } else { + params->pqTable = NULL; + } +} + /* * Insert a tuple into the index */ @@ -95,6 +150,11 @@ static void InsertTuple(Relation index, Datum *values, const bool *isnull, ItemP BlockNumber insertPage = InvalidBlockNumber; ListInfo listInfo; BlockNumber originalInsertPage; + PQParams params; + bool enablePQ; + bool byResidual; + uint16 pqcodeSize; + int dim = TupleDescAttr(index->rd_att, 0)->atttypmod; /* Detoast once for all calls */ value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); @@ -114,6 +174,8 @@ static void InsertTuple(Relation index, Datum *values, const bool *isnull, ItemP /* Ensure index is valid */ IvfflatGetMetaPageInfo(index, NULL, NULL); + InitPQParamsOnDisk(index, ¶ms, dim, &enablePQ, &byResidual, &pqcodeSize); + /* Find the insert page - sets the page and list info */ FindInsertPage(index, values, &insertPage, &listInfo); Assert(BlockNumberIsValid(insertPage)); @@ -135,7 +197,7 @@ static void InsertTuple(Relation index, Datum *values, const bool *isnull, ItemP state = GenericXLogStart(index); page = GenericXLogRegisterBuffer(state, buf, 0); - if (PageGetFreeSpace(page) >= itemsz) { + if (PageGetFreeSpace(page) >= itemsz + MAXALIGN(pqcodeSize)) { break; } @@ -177,6 +239,29 @@ static void InsertTuple(Relation index, Datum *values, const bool *isnull, ItemP } } + if (enablePQ) { + uint8 *pqcode = (uint8 *)palloc(pqcodeSize); + float *vec = ((Vector *)value)->x; + if (byResidual) { + float *resVec = (float *)palloc(dim * sizeof(float)); + Buffer cbuf = ReadBuffer(index, listInfo.blkno); + LockBuffer(buf, BUFFER_LOCK_SHARE); + Page cpage = BufferGetPage(cbuf); + IvfflatList list = (IvfflatList)PageGetItem(cpage, PageGetItemId(cpage, listInfo.offno)); + + for (int i = 0; i < dim; i++) { + resVec[i] = vec[i] - list->center.x[i]; + } + vec = resVec; + UnlockReleaseBuffer(cbuf); + IvfComputeVectorPQCode(vec, ¶ms, pqcode); + ((PageHeader)page)->pd_upper -= MAXALIGN(pqcodeSize); + errno_t rc = memcpy_s( + ((char *)page) + ((PageHeader)page)->pd_upper, pqcodeSize, (char *)pqcode, pqcodeSize); + securec_check_c(rc, "\0", "\0"); + } + } + /* Add to next offset */ if (PageAddItem(page, (Item)itup, itemsz, InvalidOffsetNumber, false, false) == InvalidOffsetNumber) elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); diff --git a/src/gausskernel/storage/access/datavec/ivfscan.cpp b/src/gausskernel/storage/access/datavec/ivfscan.cpp index 5e2fb4b96e..3569b9804d 100644 --- a/src/gausskernel/storage/access/datavec/ivfscan.cpp +++ b/src/gausskernel/storage/access/datavec/ivfscan.cpp @@ -53,9 +53,13 @@ static int CompareLists(const pairingheap_node *a, const pairingheap_node *b, vo static void GetScanLists(IndexScanDesc scan, Datum value) { IvfflatScanOpaque so = (IvfflatScanOpaque)scan->opaque; - BlockNumber nextblkno = IVFFLAT_HEAD_BLKNO; + uint16 pqTableNblk; + uint16 pqDisTableNblk; + IvfGetPQInfoFromMetaPage(scan->indexRelation, &pqTableNblk, NULL, &pqDisTableNblk, NULL); + BlockNumber nextblkno = IVFPQTABLE_START_BLKNO + pqTableNblk + pqDisTableNblk; int listCount = 0; double maxDistance = DBL_MAX; + int listId = 0; /* Search all list pages */ while (BlockNumberIsValid(nextblkno)) { @@ -82,6 +86,7 @@ static void GetScanLists(IndexScanDesc scan, Datum value) scanlist = &so->lists[listCount]; scanlist->startPage = list->startPage; scanlist->distance = distance; + scanlist->key = listId; listCount++; /* Add to heap */ @@ -99,11 +104,13 @@ static void GetScanLists(IndexScanDesc scan, Datum value) /* Reuse */ scanlist->startPage = list->startPage; scanlist->distance = distance; + scanlist->key = listId; pairingheap_add(so->listQueue, &scanlist->ph_node); /* Update max distance */ maxDistance = ((IvfflatScanList *)pairingheap_first(so->listQueue))->distance; } + listId++; } nextblkno = IvfflatPageGetOpaque(cpage)->nextblkno; @@ -203,6 +210,219 @@ static void GetScanItems(IndexScanDesc scan, Datum value) tuplesort_performsort(so->sortstate); } +/* + * Compare candidate distances + */ +static inline int CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg) +{ + if (((const IvfpqPairingHeapNode *)a)->distance < ((const IvfpqPairingHeapNode *)b)->distance) { + return -1; + } + if (((const IvfpqPairingHeapNode *)a)->distance > ((const IvfpqPairingHeapNode *)b)->distance) { + return 1; + } + + return 0; +} + +/* + * Compare candidate blocknumber + */ +static inline int CompareBlknoCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg) +{ + if (((const IvfpqPairingHeapNode *)a)->indexBlk < ((const IvfpqPairingHeapNode *)b)->indexBlk) { + return -1; + } + if (((const IvfpqPairingHeapNode *)a)->indexBlk > ((const IvfpqPairingHeapNode *)b)->indexBlk) { + return 1; + } + + return 0; +} + +/* + * Get items PQ + */ +static void GetScanItemsPQ(IndexScanDesc scan, Datum value, float *simTable) +{ + IvfflatScanOpaque so = (IvfflatScanOpaque)scan->opaque; + TupleDesc tupdesc = RelationGetDescr(scan->indexRelation); + double tuples = 0; + TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc); + Relation index = scan->indexRelation; + int pqM = so->pqM; + int pqKsub = so->pqKsub; + int kreorder = so->kreorder; + bool l2CosResidual = so->funcType != IVFPQ_DIS_IP && so->byResidual; + pairingheap *reOrderCandidate = pairingheap_allocate(CompareFurthestCandidates, NULL); + int canLen = 0; + + /* + * Reuse same set of shared buffers for scan + * + * See postgres/src/backend/storage/buffer/README for description + */ + BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD); + + /* Search closest probes lists */ + int listCount = 0; + while (!pairingheap_is_empty(so->listQueue)) { + IvfflatScanList *scanlist = (IvfflatScanList *)pairingheap_remove_first(so->listQueue); + double dis0 = so->byResidual ? scanlist->distance : 0; + BlockNumber searchPage = scanlist->startPage; + int key = scanlist->key; + float *simTable2; + /* Search all entry pages for list */ + bool isEmptyList = false; + bool isFirstPage = true; + + if (l2CosResidual) { + /* L2 or Cosine */ + float *preComputeDisTable = (float *)index->pqDistanceTable + key * pqM * pqKsub; + simTable2 = (float *)palloc(pqM * pqKsub * sizeof(float)); + VectorMadd(pqM * pqKsub, preComputeDisTable, -2.0, simTable, simTable2); + } + + while (BlockNumberIsValid(searchPage)) { + Buffer buf; + Page page; + OffsetNumber maxoffno; + + buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas); + LockBuffer(buf, BUFFER_LOCK_SHARE); + page = BufferGetPage(buf); + maxoffno = PageGetMaxOffsetNumber(page); + + isEmptyList = (isFirstPage && maxoffno <= 0 && !BlockNumberIsValid(IvfflatPageGetOpaque(page)->nextblkno)); + isFirstPage = false; + if (isEmptyList) { + UnlockReleaseBuffer(buf); + break; + } + + for (OffsetNumber offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) { + IndexTuple itup; + Datum datum; + bool isnull; + uint8 *code; + double distance; + double maxDistance = DBL_MAX; + + ItemId itemid = PageGetItemId(page, offno); + + itup = (IndexTuple)PageGetItem(page, itemid); + datum = index_getattr(itup, 1, tupdesc, &isnull); + code = LoadPQCode(itup); + if (l2CosResidual) { + distance = GetPQDistance(simTable2, code, dis0, pqM, pqKsub, false); + } else { + distance = GetPQDistance(simTable, code, dis0, pqM, pqKsub, so->funcType == IVFPQ_DIS_IP); + } + + if (kreorder == 0) { + /* + * Add virtual tuple + * + * Use procinfo from the index instead of scan key for + * performance + */ + ExecClearTuple(slot); + slot->tts_values[0] = Float8GetDatum(distance); + slot->tts_isnull[0] = false; + slot->tts_values[1] = PointerGetDatum(&itup->t_tid); + slot->tts_isnull[1] = false; + ExecStoreVirtualTuple(slot); + + tuplesort_puttupleslot(so->sortstate, slot); + } else { + /* need reorder, add to pairingheap */ + if (canLen < kreorder) { + IvfpqPairingHeapNode *e = IvfpqCreatePairingHeapNode(distance, &itup->t_tid, searchPage, offno); + pairingheap_add(reOrderCandidate, &e->ph_node); + canLen++; + if (canLen == kreorder) { + maxDistance = ((IvfpqPairingHeapNode *)pairingheap_first(reOrderCandidate))->distance; + } + } else if (distance < maxDistance) { + IvfpqPairingHeapNode *e = (IvfpqPairingHeapNode *)pairingheap_remove_first(reOrderCandidate); + e->distance = distance; + e->heapTid = &itup->t_tid; + e->indexBlk = searchPage; + e->indexOff = offno; + pairingheap_add(reOrderCandidate, &e->ph_node); + maxDistance = ((IvfpqPairingHeapNode *)pairingheap_first(reOrderCandidate))->distance; + } + } + tuples++; + } + + searchPage = IvfflatPageGetOpaque(page)->nextblkno; + + UnlockReleaseBuffer(buf); + } + + if (!isEmptyList) { + ++listCount; + if (listCount >= so->probes) { + break; + } + } + } + + FreeAccessStrategy(bas); + + if (tuples < 100) + ereport(DEBUG1, + (errmsg("index scan found few tuples"), errdetail("Index may have been created with little data."), + errhint("Recreate the index and possibly decrease lists."))); + + if (kreorder != 0) { + pairingheap *blkOrderCandidate = pairingheap_allocate(CompareBlknoCandidates, NULL); + BlockNumber blkno = InvalidBlockNumber; + Buffer buf; + Page page; + + while (!pairingheap_is_empty(reOrderCandidate)) { + pairingheap_add(blkOrderCandidate, pairingheap_remove_first(reOrderCandidate)); + } + + while (!pairingheap_is_empty(blkOrderCandidate)) { + bool isnull; + IvfpqPairingHeapNode *node = (IvfpqPairingHeapNode *)pairingheap_remove_first(blkOrderCandidate); + + if (blkno != node->indexBlk) { + if (BlockNumberIsValid(blkno)) { + UnlockReleaseBuffer(buf); + } + blkno = node->indexBlk; + buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, node->indexBlk, RBM_NORMAL, bas); + LockBuffer(buf, BUFFER_LOCK_SHARE); + page = BufferGetPage(buf); + } + + ItemId itemid = PageGetItemId(page, node->indexOff); + IndexTuple itup = (IndexTuple)PageGetItem(page, itemid); + Datum datum = index_getattr(itup, 1, tupdesc, &isnull); + + /* Add virtual tuple */ + ExecClearTuple(slot); + slot->tts_values[0] = so->distfunc(so->procinfo, so->collation, datum, value); + slot->tts_isnull[0] = false; + slot->tts_values[1] = PointerGetDatum(node->heapTid); + slot->tts_isnull[1] = false; + ExecStoreVirtualTuple(slot); + + tuplesort_puttupleslot(so->sortstate, slot); + } + + if (BlockNumberIsValid(blkno)) { + UnlockReleaseBuffer(buf); + } + } + + tuplesort_performsort(so->sortstate); +} + /* * Zero distance */ @@ -271,6 +491,7 @@ IndexScanDesc ivfflatbeginscan_internal(Relation index, int nkeys, int norderbys so->listCount = lists; so->probes = probes; so->dimensions = dimensions; + so->kreorder = u_sess->datavec_ctx.ivfpq_kreorder; /* Set support functions */ so->procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC); @@ -290,6 +511,9 @@ IndexScanDesc ivfflatbeginscan_internal(Relation index, int nkeys, int norderbys so->listQueue = pairingheap_allocate(CompareLists, scan); + GetPQInfoOnDisk(so, index); + so->pqCtx = AllocSetContextCreate(CurrentMemoryContext, "IVFPQ scan temporary context", ALLOCSET_DEFAULT_SIZES); + scan->opaque = so; return scan; @@ -345,8 +569,20 @@ bool ivfflatgettuple_internal(IndexScanDesc scan, ScanDirection dir) elog(ERROR, "non-MVCC snapshots are not supported with ivfflat"); value = GetScanValue(scan); + IvfflatBench("GetScanLists", GetScanLists(scan, value)); - IvfflatBench("GetScanItems", GetScanItems(scan, value)); + + if (so->enablePQ) { + MemoryContext oldCxt = MemoryContextSwitchTo(so->pqCtx); + + float *simTable = (float *)palloc0(so->pqM * so->pqKsub * sizeof(float)); + IvfpqComputeQueryRelTables(so, scan->indexRelation, value, simTable); + IvfflatBench("GetScanItemsPQ", GetScanItemsPQ(scan, value, simTable)); + + MemoryContextSwitchTo(oldCxt); + } else { + IvfflatBench("GetScanItems", GetScanItems(scan, value)); + } so->first = false; /* Clean up if we allocated a new value */ @@ -367,7 +603,17 @@ bool ivfflatgettuple_internal(IndexScanDesc scan, ScanDirection dir) so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, u_sess->attr.attr_memory.work_mem, NULL, false); Datum value = GetScanValue(scan); - IvfflatBench("GetScanItems", GetScanItems(scan, value)); + if (so->enablePQ) { + MemoryContext oldCxt = MemoryContextSwitchTo(so->pqCtx); + + float *simTable = (float *)palloc0(so->pqM * so->pqKsub * sizeof(float)); + IvfpqComputeQueryRelTables(so, scan->indexRelation, value, simTable); + IvfflatBench("GetScanItemsPQ", GetScanItemsPQ(scan, value, simTable)); + + MemoryContextSwitchTo(oldCxt); + } else { + IvfflatBench("GetScanItems", GetScanItems(scan, value)); + } isDone = tuplesort_gettupleslot(so->sortstate, true, so->slot, NULL); /* Clean up if we allocated a new value */ @@ -394,6 +640,7 @@ void ivfflatendscan_internal(IndexScanDesc scan) { IvfflatScanOpaque so = (IvfflatScanOpaque)scan->opaque; + MemoryContextDelete(so->pqCtx); pairingheap_free(so->listQueue); tuplesort_end(so->sortstate); diff --git a/src/gausskernel/storage/access/datavec/ivfutils.cpp b/src/gausskernel/storage/access/datavec/ivfutils.cpp index 9961ab8168..937d805575 100644 --- a/src/gausskernel/storage/access/datavec/ivfutils.cpp +++ b/src/gausskernel/storage/access/datavec/ivfutils.cpp @@ -205,6 +205,175 @@ void IvfflatUpdateList(Relation index, ListInfo listInfo, BlockNumber insertPage } } +char* IVFPQLoadPQtable(Relation index) +{ + Buffer buf; + Page page; + uint16 nblks; + uint32 curFlushSize; + uint32 pqTableSize; + char* pqTable; + + IvfGetPQInfoFromMetaPage(index, &nblks, &pqTableSize, NULL, NULL); + pqTable = (char*)palloc0(pqTableSize); + + for (uint16 i = 0; i < nblks; i++) { + curFlushSize = (i == nblks - 1) ? (pqTableSize - i * IVFPQTABLE_STORAGE_SIZE) : IVFPQTABLE_STORAGE_SIZE; + buf = ReadBuffer(index, IVFPQTABLE_START_BLKNO + i); + LockBuffer(buf, BUFFER_LOCK_SHARE); + page = BufferGetPage(buf); + errno_t err = memcpy_s(pqTable + i * IVFPQTABLE_STORAGE_SIZE, curFlushSize, + PageGetContents(page), curFlushSize); + securec_check(err, "\0", "\0"); + UnlockReleaseBuffer(buf); + } + return pqTable; +} + +float* IVFPQLoadPQDisTable(Relation index) +{ + Buffer buf; + Page page; + uint16 pqTableNblk; + uint16 nblks; + uint32 curFlushSize; + uint32 pqDisTableSize; + float* disTable; + + IvfGetPQInfoFromMetaPage(index, &pqTableNblk, NULL, &nblks, &pqDisTableSize); + disTable = (float*)palloc0(pqDisTableSize); + + BlockNumber startBlkno = IVFPQTABLE_START_BLKNO + pqTableNblk; + for (uint16 i = 0; i < nblks; i++) { + curFlushSize = (i == nblks - 1) ? (pqDisTableSize - i * IVFPQTABLE_STORAGE_SIZE) : IVFPQTABLE_STORAGE_SIZE; + buf = ReadBuffer(index, startBlkno + i); + LockBuffer(buf, BUFFER_LOCK_SHARE); + page = BufferGetPage(buf); + errno_t err = memcpy_s((char*)disTable + i * IVFPQTABLE_STORAGE_SIZE, curFlushSize, + PageGetContents(page), curFlushSize); + securec_check(err, "\0", "\0"); + UnlockReleaseBuffer(buf); + } + return disTable; +} + +/* + * Get Ivfflat PQ info + */ +void GetPQInfoOnDisk(IvfflatScanOpaque so, Relation index) +{ + Buffer buf; + Page page; + IvfflatMetaPage metap; + + buf = ReadBuffer(index, IVFFLAT_METAPAGE_BLKNO); + LockBuffer(buf, BUFFER_LOCK_SHARE); + page = BufferGetPage(buf); + metap = IvfflatPageGetMeta(page); + + PG_TRY(); + { + if (unlikely(metap->magicNumber != IVFFLAT_MAGIC_NUMBER)) { + elog(ERROR, "ivfflat index is not valid"); + } + } + PG_CATCH(); + { + UnlockReleaseBuffer(buf); + PG_RE_THROW(); + } + PG_END_TRY(); + + so->enablePQ = metap->enablePQ; + so->pqM = metap->pqM; + so->pqKsub = metap->pqKsub; + so->byResidual = metap->byResidual; + UnlockReleaseBuffer(buf); + + if (so->enablePQ) { + so->funcType = getIVFPQfunctionType(so->procinfo, so->normprocinfo); + /* Now save pqTable and pqDistanceTable in the relcache entry. */ + if (index->pqTable == NULL) { + MemoryContext oldcxt = MemoryContextSwitchTo(index->rd_indexcxt); + index->pqTable = IVFPQLoadPQtable(index); + (void)MemoryContextSwitchTo(oldcxt); + } + if (index->pqDistanceTable == NULL && so->byResidual && so->funcType != IVFPQ_DIS_IP) { + MemoryContext oldcxt = MemoryContextSwitchTo(index->rd_indexcxt); + index->pqDistanceTable = IVFPQLoadPQDisTable(index); + (void)MemoryContextSwitchTo(oldcxt); + } + } +} + +void IvfpqComputeQueryRelTablesInternal(IvfflatScanOpaque so, float *q, char *pqTable, bool innerPro, float *simTable) +{ + int pqM = so->pqM; + int pqKsub = so->pqKsub; + int dim = so->dimensions; + int dsub = dim / pqM; + Size subSize = MAXALIGN(so->typeInfo->itemSize(dsub)); + + for (int m = 0; m < pqM; m++) { + int offset = m * pqKsub; + float *qsubVector = q + m * dsub; + float *dis = simTable + offset; + /* one-to-many computation */ + if (innerPro) { + /* negate when GetPQDistance */ + VectorInnerProductNY(dsub, pqKsub, qsubVector, pqTable, subSize, offset, dis); + } else { + VectorL2SquaredDistanceNY(dsub, pqKsub, qsubVector, pqTable, subSize, offset, dis); + } + } +} + +/* + * Precompute some tables specific to query q, r is cluster center of PQ. + */ +void IvfpqComputeQueryRelTables(IvfflatScanOpaque so, Relation index, Datum q, float *simTable) +{ + if (so->funcType == IVFPQ_DIS_IP) { + /* compute q*r */ + IvfpqComputeQueryRelTablesInternal(so, DatumGetVector(q)->x, index->pqTable, true, simTable); + } else { + /* funcType is cosine or l2 */ + if (so->byResidual) { + /* compute q*r */ + IvfpqComputeQueryRelTablesInternal(so, DatumGetVector(q)->x, index->pqTable, true, simTable); + } else { + /* compute (q-r)^2 */ + IvfpqComputeQueryRelTablesInternal(so, DatumGetVector(q)->x, index->pqTable, false, simTable); + } + } +} + +uint8 *LoadPQCode(IndexTuple itup) +{ + return (uint8 *)((char *)itup + MAXALIGN(IndexTupleSize(itup))); +} + +float GetPQDistance(float *pqDistanceTable, uint8 *code, double dis0, int pqM, int pqKsub, bool innerPro) +{ + float resDistance = dis0; + for (int i = 0; i < pqM; i++) { + int offset = i * pqKsub + code[i]; + resDistance += pqDistanceTable[offset]; + } + return innerPro ? (0 - resDistance) : resDistance; +} + +IvfpqPairingHeapNode * IvfpqCreatePairingHeapNode(float distance, ItemPointer heapTid, + BlockNumber indexBlk, OffsetNumber indexOff) +{ + IvfpqPairingHeapNode *n = (IvfpqPairingHeapNode *)palloc(sizeof(IvfpqPairingHeapNode)); + n->distance = distance; + n->heapTid = heapTid; + n->indexBlk = indexBlk; + n->indexOff = indexOff; + return n; +} + /* * Get type info */ @@ -263,7 +432,7 @@ int getIVFPQfunctionType(FmgrInfo *procinfo, FmgrInfo *normprocinfo) return IVF_PQ_DIS_COSINE; } } else { - ereport(ERROR, (errmsg("current data type or distance type can't support ivfflatpq."))); + ereport(ERROR, (errmsg("current data type or distance type can't support IVFPQ."))); return -1; } } @@ -365,7 +534,7 @@ int IvfGetByResidual(Relation index) return opts->byResidual; } - return IVFFLAT_DEFAULT_PQ_RESIDUAL; + return IVFPQ_DEFAULT_RESIDUAL; } void IvfFlushPQInfoInternal(Relation index, char* table, BlockNumber startBlkno, uint16 nblks, uint32 totalSize) diff --git a/src/gausskernel/storage/access/datavec/ivfvacuum.cpp b/src/gausskernel/storage/access/datavec/ivfvacuum.cpp index 2ff912f4ac..ef8084c37d 100644 --- a/src/gausskernel/storage/access/datavec/ivfvacuum.cpp +++ b/src/gausskernel/storage/access/datavec/ivfvacuum.cpp @@ -33,8 +33,11 @@ IndexBulkDeleteResult *ivfflatbulkdelete_internal(IndexVacuumInfo *info, IndexBulkDeleteResult *stats, IndexBulkDeleteCallback callback, void *callbackState) { + uint16 pqTableNblk; + uint16 pqDisTableNblk; Relation index = info->index; - BlockNumber blkno = IVFFLAT_HEAD_BLKNO; + IvfGetPQInfoFromMetaPage(index, &pqTableNblk, NULL, &pqDisTableNblk, NULL); + BlockNumber blkno = IVFPQTABLE_START_BLKNO + pqTableNblk + pqDisTableNblk; BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD); if (stats == NULL) diff --git a/src/include/access/datavec/hnsw.h b/src/include/access/datavec/hnsw.h index 69703b1c0f..b8ca613295 100644 --- a/src/include/access/datavec/hnsw.h +++ b/src/include/access/datavec/hnsw.h @@ -653,8 +653,6 @@ bool HnswDelete(Relation index, Datum *values, const bool *isnull, ItemPointer h void HnswUpdateAppendMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber eleInsertPage, BlockNumber neiInsertPage, ForkNumber forkNum, bool building); void FlushPQInfo(HnswBuildState *buildstate); -char *LoadPQtable(Relation index); -float *LoadPQDisTable(Relation index); void HnswGetPQInfoFromMetaPage(Relation index, uint16 *pqTableNblk, uint32 *pqTableSize, uint16 *pqDisTableNblk, uint32 *pqDisTableSize); diff --git a/src/include/access/datavec/ivfflat.h b/src/include/access/datavec/ivfflat.h index 16163a857b..c414739d9e 100644 --- a/src/include/access/datavec/ivfflat.h +++ b/src/include/access/datavec/ivfflat.h @@ -56,14 +56,20 @@ /* Preserved page numbers */ #define IVFFLAT_METAPAGE_BLKNO 0 -#define IVFFLAT_HEAD_BLKNO 1 /* first list page */ +#define IVFPQTABLE_START_BLKNO 1 /* first list page of pqtable start page */ /* IVFFlat parameters */ #define IVFFLAT_DEFAULT_LISTS 100 #define IVFFLAT_MIN_LISTS 1 #define IVFFLAT_MAX_LISTS 32768 #define IVFFLAT_DEFAULT_PROBES 1 -#define IVFFLAT_DEFAULT_PQ_RESIDUAL false + +/* IVFPQ parameters */ +#define IVFPQ_DEFAULT_RESIDUAL false +#define IVFPQ_DIS_L2 1 +#define IVFPQ_DIS_IP 2 +#define IVFPQ_DIS_COSINE 3 +#define IVFPQTABLE_STORAGE_SIZE (uint16)(6 * 1024) /* pqtable storage size in each page */ /* Build phases */ /* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */ @@ -277,6 +283,7 @@ typedef struct IvfflatScanList { pairingheap_node ph_node; BlockNumber startPage; double distance; + int key; } IvfflatScanList; typedef struct IvfflatScanOpaqueData { @@ -298,6 +305,15 @@ typedef struct IvfflatScanOpaqueData { Oid collation; Datum (*distfunc)(FmgrInfo *flinfo, Oid collation, Datum arg1, Datum arg2); + /* PQ info */ + bool enablePQ; + int pqM; + int pqKsub; + int funcType; + bool byResidual; + int kreorder; + MemoryContext pqCtx; + /* Lists */ pairingheap *listQueue; IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */ @@ -305,6 +321,14 @@ typedef struct IvfflatScanOpaqueData { typedef IvfflatScanOpaqueData *IvfflatScanOpaque; +typedef struct IvfpqPairingHeapNode { + pairingheap_node ph_node; + double distance; + ItemPointer heapTid; + BlockNumber indexBlk; + OffsetNumber indexOff; +} IvfpqPairingHeapNode; + /* Methods */ void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, const IvfflatTypeInfo *typeInfo); FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum); @@ -340,6 +364,13 @@ int IvfGetPQDistanceTableAdc(float *vector, const PQParams *params, float *pqDis int IvfGetPQDistance(const uint8 *basecode, const uint8 *querycode, const PQParams *params, const float *pqDistanceTable, float *pqDistance); +void GetPQInfoOnDisk(IvfflatScanOpaque so, Relation index); +void IvfpqComputeQueryRelTables(IvfflatScanOpaque so, Relation index, Datum q, float *simTable); +uint8 *LoadPQCode(IndexTuple itup); +float GetPQDistance(float *pqDistanceTable, uint8 *code, double dis0, int pqM, int pqKsub, bool innerPro); +IvfpqPairingHeapNode * IvfpqCreatePairingHeapNode(float distance, ItemPointer heapTid, + BlockNumber indexBlk, OffsetNumber indexOff); +char* IVFPQLoadPQtable(Relation index); Datum ivfflathandler(PG_FUNCTION_ARGS); Datum ivfflatbuild(PG_FUNCTION_ARGS); diff --git a/src/include/access/datavec/vector.h b/src/include/access/datavec/vector.h index 01796d5eb8..c2fcb2bf66 100644 --- a/src/include/access/datavec/vector.h +++ b/src/include/access/datavec/vector.h @@ -51,6 +51,8 @@ Vector *InitVector(int dim); void PrintVector(char *msg, Vector *vector); int vector_cmp_internal(Vector *a, Vector *b); void VectorMadd(size_t n, const float *ax, float bf, const float *bx, float *cx); +void VectorL2SquaredDistanceNY(size_t d, size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis); +void VectorInnerProductNY(size_t d, size_t ny, float *x, char *pqTable, Size subSize, int offset, float *dis); void LogNewpageRange(Relation rel, ForkNumber forknum, BlockNumber startblk, BlockNumber endblk, bool page_std); int PlanCreateIndexWorkers(Relation heapRelation, IndexInfo *indexInfo); diff --git a/src/include/knl/knl_session.h b/src/include/knl/knl_session.h index 6e4eb6871c..8c6787d3ef 100644 --- a/src/include/knl/knl_session.h +++ b/src/include/knl/knl_session.h @@ -3047,6 +3047,7 @@ typedef struct knl_u_datavec_context { int hnsw_ef_search; int hnsw_earlystop_threshold; int ivfflat_probes; + int ivfpq_kreorder; } knl_u_datavec_context; #ifdef ENABLE_HTAP diff --git a/src/include/utils/rel.h b/src/include/utils/rel.h index 87968ae96b..6c4b26a90c 100644 --- a/src/include/utils/rel.h +++ b/src/include/utils/rel.h @@ -301,7 +301,7 @@ typedef struct RelationData { /* used only for gsc, keep it preserved if you modify the rel, otherwise set it null */ struct LocalRelationEntry *entry; - /* used only for datavec hnswpq */ + /* used only for datavec pq */ char *pqTable; float *pqDistanceTable; } RelationData; -- Gitee