liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-12 23:56:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added support function for l2_normalize [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-04-22 18:36:47 -07:00
parent 2b77005610
commit f14c21748b
8 changed files with 22 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
sql/vector--0.6.2--0.7.0.sql
View File

@@ -342,7 +342,8 @@ CREATE OPERATOR CLASS halfvec_cosine_ops
FOR TYPE halfvec USING hnsw AS FOR TYPE halfvec USING hnsw AS
OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops,
FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec), FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec),
FUNCTION 2 l2_norm(halfvec); FUNCTION 2 l2_norm(halfvec),
FUNCTION 3 l2_normalize(halfvec);
CREATE OPERATOR CLASS halfvec_l1_ops CREATE OPERATOR CLASS halfvec_l1_ops
FOR TYPE halfvec USING hnsw AS FOR TYPE halfvec USING hnsw AS
@@ -529,7 +530,8 @@ CREATE OPERATOR CLASS sparsevec_cosine_ops
FOR TYPE sparsevec USING hnsw AS FOR TYPE sparsevec USING hnsw AS
OPERATOR 1 <=> (sparsevec, sparsevec) FOR ORDER BY float_ops, OPERATOR 1 <=> (sparsevec, sparsevec) FOR ORDER BY float_ops,
FUNCTION 1 sparsevec_negative_inner_product(sparsevec, sparsevec), FUNCTION 1 sparsevec_negative_inner_product(sparsevec, sparsevec),
FUNCTION 2 l2_norm(sparsevec); FUNCTION 2 l2_norm(sparsevec),
FUNCTION 3 l2_normalize(sparsevec);
CREATE OPERATOR CLASS sparsevec_l1_ops CREATE OPERATOR CLASS sparsevec_l1_ops
FOR TYPE sparsevec USING hnsw AS FOR TYPE sparsevec USING hnsw AS

6
sql/vector.sql
View File

@@ -651,7 +651,8 @@ CREATE OPERATOR CLASS halfvec_cosine_ops
FOR TYPE halfvec USING hnsw AS FOR TYPE halfvec USING hnsw AS
OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops, OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops,
FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec), FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec),
FUNCTION 2 l2_norm(halfvec); FUNCTION 2 l2_norm(halfvec),
FUNCTION 3 l2_normalize(halfvec);
CREATE OPERATOR CLASS halfvec_l1_ops CREATE OPERATOR CLASS halfvec_l1_ops
FOR TYPE halfvec USING hnsw AS FOR TYPE halfvec USING hnsw AS
@@ -852,7 +853,8 @@ CREATE OPERATOR CLASS sparsevec_cosine_ops
FOR TYPE sparsevec USING hnsw AS FOR TYPE sparsevec USING hnsw AS
OPERATOR 1 <=> (sparsevec, sparsevec) FOR ORDER BY float_ops, OPERATOR 1 <=> (sparsevec, sparsevec) FOR ORDER BY float_ops,
FUNCTION 1 sparsevec_negative_inner_product(sparsevec, sparsevec), FUNCTION 1 sparsevec_negative_inner_product(sparsevec, sparsevec),
FUNCTION 2 l2_norm(sparsevec); FUNCTION 2 l2_norm(sparsevec),
FUNCTION 3 l2_normalize(sparsevec);
CREATE OPERATOR CLASS sparsevec_l1_ops CREATE OPERATOR CLASS sparsevec_l1_ops
FOR TYPE sparsevec USING hnsw AS FOR TYPE sparsevec USING hnsw AS

2
src/hnsw.c
View File

@@ -194,7 +194,7 @@ hnswhandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine); IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0; amroutine->amstrategies = 0;
amroutine->amsupport = 2; amroutine->amsupport = 3;
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
amroutine->amoptsprocnum = 0; amroutine->amoptsprocnum = 0;
#endif #endif

5
src/hnsw.h
View File

@@ -22,6 +22,7 @@
/* Support functions */ /* Support functions */
#define HNSW_DISTANCE_PROC 1 #define HNSW_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2 #define HNSW_NORM_PROC 2
#define HNSW_NORMALIZE_PROC 3
#define HNSW_VERSION 1 #define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953 #define HNSW_MAGIC_NUMBER 0xA953A953
@@ -265,6 +266,7 @@ typedef struct HnswBuildState
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
FmgrInfo *normalizeprocinfo;
Oid collation; Oid collation;
/* Variables */ /* Variables */
@@ -341,6 +343,7 @@ typedef struct HnswScanOpaqueData
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
FmgrInfo *normalizeprocinfo;
Oid collation; Oid collation;
} HnswScanOpaqueData; } HnswScanOpaqueData;
@@ -377,7 +380,7 @@ int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index); int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
HnswType HnswGetType(Relation index); HnswType HnswGetType(Relation index);
Datum HnswNormValue(Datum value, HnswType type); Datum HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum value);
bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value); bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
void HnswCheckValue(Datum value, HnswType type); void HnswCheckValue(Datum value, HnswType type);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);

3
src/hnswbuild.c
View File

@@ -496,7 +496,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
if (!HnswCheckNorm(buildstate->normprocinfo, buildstate->collation, value)) if (!HnswCheckNorm(buildstate->normprocinfo, buildstate->collation, value))
return false; return false;
value = HnswNormValue(value, buildstate->type); value = HnswNormValue(buildstate->normalizeprocinfo, buildstate->collation, value);
} }
/* Get datum size */ /* Get datum size */
@@ -725,6 +725,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
/* Get support functions */ /* Get support functions */
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
buildstate->collation = index->rd_indcollation[0]; buildstate->collation = index->rd_indcollation[0];
InitGraph(&buildstate->graphData, NULL, maintenance_work_mem * 1024L); InitGraph(&buildstate->graphData, NULL, maintenance_work_mem * 1024L);

2
src/hnswinsert.c
View File

@@ -629,7 +629,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
if (!HnswCheckNorm(normprocinfo, collation, value)) if (!HnswCheckNorm(normprocinfo, collation, value))
return; return;
value = HnswNormValue(value, type); value = HnswNormValue(HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC), collation, value);
} }
HnswInsertTupleOnDisk(index, value, values, isnull, heap_tid, false); HnswInsertTupleOnDisk(index, value, values, isnull, heap_tid, false);

3
src/hnswscan.c
View File

@@ -61,7 +61,7 @@ GetScanValue(IndexScanDesc scan)
/* Fine if normalization fails */ /* Fine if normalization fails */
if (so->normprocinfo != NULL) if (so->normprocinfo != NULL)
value = HnswNormValue(value, HnswGetType(scan->indexRelation)); value = HnswNormValue(so->normalizeprocinfo, so->collation, value);
} }
return value; return value;
@@ -87,6 +87,7 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
/* Set support functions */ /* Set support functions */
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
so->normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
so->collation = index->rd_indcollation[0]; so->collation = index->rd_indcollation[0];
scan->opaque = so; scan->opaque = so;

13
src/hnswutils.c
View File

@@ -195,17 +195,12 @@ HnswGetType(Relation index)
* Normalize value * Normalize value
*/ */
Datum Datum
HnswNormValue(Datum value, HnswType type) HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum value)
{ {
/* TODO Remove type-specific code */ if (procinfo == NULL)
if (type == HNSW_TYPE_VECTOR)
return DirectFunctionCall1(l2_normalize, value); return DirectFunctionCall1(l2_normalize, value);
else if (type == HNSW_TYPE_HALFVEC)
return DirectFunctionCall1(halfvec_l2_normalize, value); return FunctionCall1Coll(procinfo, collation, value);
else if (type == HNSW_TYPE_SPARSEVEC)
return DirectFunctionCall1(sparsevec_l2_normalize, value);
else
elog(ERROR, "Unsupported type");
} }
/* /*