liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-02 10:40:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

Moved type lookup to support functions - #527

Browse Source
This commit is contained in:
Andrew Kane
2024-04-23 13:02:47 -07:00
parent d1b83991af
commit 0da6213a60
8 changed files with 151 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/hnsw.c
View File

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

4
src/hnsw.h
View File

@@ -23,6 +23,7 @@
#define HNSW_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2
#define HNSW_NORMALIZE_PROC 3
#define HNSW_TYPE_SUPPORT_PROC 4
#define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953
@@ -62,7 +63,8 @@ typedef enum HnswType
HNSW_TYPE_VECTOR,
HNSW_TYPE_HALFVEC,
HNSW_TYPE_BIT,
HNSW_TYPE_SPARSEVEC
HNSW_TYPE_SPARSEVEC,
HNSW_TYPE_UNSUPPORTED
} HnswType;
/* Build phases */

52
src/hnswutils.c
View File

@@ -13,7 +13,6 @@
#include "utils/datum.h"
#include "utils/memdebug.h"
#include "utils/rel.h"
#include "utils/syscache.h"
#if PG_VERSION_NUM >= 130000
#include "common/hashfn.h"
@@ -159,32 +158,17 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
HnswType
HnswGetType(Relation index)
{
FmgrInfo *procinfo = HnswOptionalProcInfo(index, HNSW_TYPE_SUPPORT_PROC);
Oid typid = TupleDescAttr(index->rd_att, 0)->atttypid;
HeapTuple tuple;
Form_pg_type type;
HnswType result;
if (typid == BITOID)
return HNSW_TYPE_BIT;
if (procinfo == NULL)
return HNSW_TYPE_VECTOR;
tuple = SearchSysCache1(TYPEOID, ObjectIdGetDatum(typid));
if (!HeapTupleIsValid(tuple))
elog(ERROR, "cache lookup failed for type %u", typid);
result = (HnswType) DatumGetInt32(FunctionCall1(procinfo, ObjectIdGetDatum(typid)));
type = (Form_pg_type) GETSTRUCT(tuple);
if (strcmp(NameStr(type->typname), "vector") == 0)
result = HNSW_TYPE_VECTOR;
else if (strcmp(NameStr(type->typname), "halfvec") == 0)
result = HNSW_TYPE_HALFVEC;
else if (strcmp(NameStr(type->typname), "sparsevec") == 0)
result = HNSW_TYPE_SPARSEVEC;
else
{
ReleaseSysCache(tuple);
if (result == HNSW_TYPE_UNSUPPORTED)
elog(ERROR, "type not supported for hnsw index");
}
ReleaseSysCache(tuple);
return result;
}
@@ -1318,3 +1302,29 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
ep = w;
}
}
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_hnsw_support);
Datum
halfvec_hnsw_support(PG_FUNCTION_ARGS)
{
PG_RETURN_INT32(HNSW_TYPE_HALFVEC);
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(bit_hnsw_support);
Datum
bit_hnsw_support(PG_FUNCTION_ARGS)
{
Oid typid = PG_GETARG_OID(0);
if (typid == BITOID)
PG_RETURN_INT32(HNSW_TYPE_BIT);
else
PG_RETURN_INT32(HNSW_TYPE_UNSUPPORTED);
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(sparsevec_hnsw_support);
Datum
sparsevec_hnsw_support(PG_FUNCTION_ARGS)
{
PG_RETURN_INT32(HNSW_TYPE_SPARSEVEC);
};

2
src/ivfflat.c
View File

@@ -188,7 +188,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0;
amroutine->amsupport = 5;
amroutine->amsupport = 6;
#if PG_VERSION_NUM >= 130000
amroutine->amoptsprocnum = 0;
#endif

4
src/ivfflat.h
View File

@@ -29,6 +29,7 @@
#define IVFFLAT_KMEANS_DISTANCE_PROC 3
#define IVFFLAT_KMEANS_NORM_PROC 4
#define IVFFLAT_NORMALIZE_PROC 5
#define IVFFLAT_TYPE_SUPPORT_PROC 6
#define IVFFLAT_VERSION 1
#define IVFFLAT_MAGIC_NUMBER 0x14FF1A7
@@ -48,7 +49,8 @@ typedef enum IvfflatType
{
IVFFLAT_TYPE_VECTOR,
IVFFLAT_TYPE_HALFVEC,
IVFFLAT_TYPE_BIT
IVFFLAT_TYPE_BIT,
IVFFLAT_TYPE_UNSUPPORTED
} IvfflatType;
/* Build phases */

43
src/ivfutils.c
View File

@@ -5,7 +5,6 @@
#include "fmgr.h"
#include "ivfflat.h"
#include "storage/bufmgr.h"
#include "utils/syscache.h"
/*
* Allocate a vector array
@@ -68,30 +67,17 @@ IvfflatOptionalProcInfo(Relation index, uint16 procnum)
IvfflatType
IvfflatGetType(Relation index)
{
FmgrInfo *procinfo = IvfflatOptionalProcInfo(index, IVFFLAT_TYPE_SUPPORT_PROC);
Oid typid = TupleDescAttr(index->rd_att, 0)->atttypid;
HeapTuple tuple;
Form_pg_type type;
IvfflatType result;
if (typid == BITOID)
return IVFFLAT_TYPE_BIT;
if (procinfo == NULL)
return IVFFLAT_TYPE_VECTOR;
tuple = SearchSysCache1(TYPEOID, ObjectIdGetDatum(typid));
if (!HeapTupleIsValid(tuple))
elog(ERROR, "cache lookup failed for type %u", typid);
result = (IvfflatType) DatumGetInt32(FunctionCall1(procinfo, ObjectIdGetDatum(typid)));
type = (Form_pg_type) GETSTRUCT(tuple);
if (strcmp(NameStr(type->typname), "vector") == 0)
result = IVFFLAT_TYPE_VECTOR;
else if (strcmp(NameStr(type->typname), "halfvec") == 0)
result = IVFFLAT_TYPE_HALFVEC;
else
{
ReleaseSysCache(tuple);
if (result == IVFFLAT_TYPE_UNSUPPORTED)
elog(ERROR, "type not supported for ivfflat index");
}
ReleaseSysCache(tuple);
return result;
}
@@ -259,3 +245,22 @@ IvfflatUpdateList(Relation index, ListInfo listInfo,
UnlockReleaseBuffer(buf);
}
}
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_ivfflat_support);
Datum
halfvec_ivfflat_support(PG_FUNCTION_ARGS)
{
PG_RETURN_INT32(IVFFLAT_TYPE_HALFVEC);
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(bit_ivfflat_support);
Datum
bit_ivfflat_support(PG_FUNCTION_ARGS)
{
Oid typid = PG_GETARG_OID(0);
if (typid == BITOID)
PG_RETURN_INT32(IVFFLAT_TYPE_BIT);
else
PG_RETURN_INT32(IVFFLAT_TYPE_UNSUPPORTED);
};