PQ增加码本训练日志打印和返回值报错

This commit is contained in:
wangjingyuan8
2025-07-21 19:47:21 +08:00
parent ba52145aa2
commit eb7bdab66c
2 changed files with 21 additions and 6 deletions

View File

@ -145,16 +145,17 @@ PQParams *InitPQParamsInMemory(HnswBuildState *buildstate)
return params;
}
static void ComputeHnswPQ(HnswBuildState *buildstate)
static int ComputeHnswPQ(HnswBuildState *buildstate)
{
MemoryContext pqCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw PQ temporary context",
ALLOCSET_DEFAULT_SIZES);
MemoryContext oldCtx = MemoryContextSwitchTo(pqCtx);
ComputePQTable(buildstate->samples, buildstate->params);
int res = ComputePQTable(buildstate->samples, buildstate->params);
MemoryContextSwitchTo(oldCtx);
MemoryContextDelete(pqCtx);
return res;
}
BlockNumber BlockSamplerGetBlock(BlockSampler bs)
@ -417,6 +418,7 @@ static void BuildPQtable(HnswBuildState *buildstate)
PG_TRY();
{
/* Sample rows */
ereport(LOG, (errmsg("HNSWPQ start sample rows.")));
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions,
buildstate->typeInfo->itemSize(buildstate->dimensions));
}
@ -435,8 +437,14 @@ static void BuildPQtable(HnswBuildState *buildstate)
errhint("Drop the index until the table has more data.")));
}
}
ComputeHnswPQ(buildstate);
ereport(LOG, (errmsg("HNSWPQ start to train codebook.")));
int success = ComputeHnswPQ(buildstate);
VectorArrayFree(buildstate->samples);
if (success == -1) {
ereport(ERROR, (errmsg("HNSWPQ training codebook is failed.")));
} else {
ereport(LOG, (errmsg("HNSWPQ finish to train codebook.")));
}
}

View File

@ -156,16 +156,17 @@ static void ComputePreTable(IvfflatBuildState *buildstate)
/*
* Compute PQTable
*/
static void ComputeIvfPQ(IvfflatBuildState *buildstate)
static int ComputeIvfPQ(IvfflatBuildState *buildstate)
{
MemoryContext pqCtx = AllocSetContextCreate(CurrentMemoryContext,
"Ivfflat PQ temporary context",
ALLOCSET_DEFAULT_SIZES);
MemoryContext oldCtx = MemoryContextSwitchTo(pqCtx);
IvfComputePQTable(buildstate->residuals, buildstate->params);
int res = IvfComputePQTable(buildstate->residuals, buildstate->params);
MemoryContextSwitchTo(oldCtx);
MemoryContextDelete(pqCtx);
return res;
}
/*
@ -1120,7 +1121,13 @@ static void CreateEntryPages(IvfflatBuildState *buildstate, ForkNumber forkNum)
/* Build PQTable by residusal */
if (buildstate->enablePQ) {
CopyResidaulFromList(buildstate);
ComputeIvfPQ(buildstate);
ereport(LOG, (errmsg("IVFPQ start to train codebook.")));
int success = ComputeIvfPQ(buildstate);
if (success == -1) {
ereport(ERROR, (errmsg("IVFPQ training codebook is failed.")));
} else {
ereport(LOG, (errmsg("IVFPQ finish to train codebook.")));
}
if (buildstate->byResidual &&
(buildstate->params->funcType == IVF_PQ_DIS_L2 || buildstate->params->funcType == IVF_PQ_DIS_COSINE))
ComputePreTable(buildstate);