Skip to content

Commit 69c3e71

Browse files
ankaneankane
committed
Added support functions for max dimensions for ivfflat
1 parent e815478 commit 69c3e71

6 files changed

Lines changed: 55 additions & 18 deletions

File tree

‎sql/vector--0.6.2--0.7.0.sql‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,12 @@ CREATE OPERATOR || (
2222
LEFTARG = vector, RIGHTARG = vector, PROCEDURE = vector_concat
2323
);
2424

25+
CREATE FUNCTION ivfflat_bit_max_dims(internal) RETURNS internal
26+
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
27+
28+
CREATE FUNCTION ivfflat_halfvec_max_dims(internal) RETURNS internal
29+
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
30+
2531
CREATE FUNCTION ivfflat_bit_support(internal) RETURNS internal
2632
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
2733

@@ -66,7 +72,8 @@ CREATE OPERATOR CLASS bit_hamming_ops
6672
OPERATOR 1 <~> (bit, bit) FOR ORDER BY float_ops,
6773
FUNCTION 1 hamming_distance(bit, bit),
6874
FUNCTION 3 hamming_distance(bit, bit),
69-
FUNCTION 6 ivfflat_bit_support(internal);
75+
FUNCTION 6 ivfflat_bit_max_dims(internal),
76+
FUNCTION 7 ivfflat_bit_support(internal);
7077

7178
CREATE OPERATOR CLASS bit_hamming_ops
7279
FOR TYPE bit USING hnsw AS
@@ -333,7 +340,8 @@ CREATE OPERATOR CLASS halfvec_l2_ops
333340
OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops,
334341
FUNCTION 1 halfvec_l2_squared_distance(halfvec, halfvec),
335342
FUNCTION 3 l2_distance(halfvec, halfvec),
336-
FUNCTION 6 ivfflat_halfvec_support(internal);
343+
FUNCTION 6 ivfflat_halfvec_max_dims(internal),
344+
FUNCTION 7 ivfflat_halfvec_support(internal);
337345

338346
CREATE OPERATOR CLASS halfvec_ip_ops
339347
FOR TYPE halfvec USING ivfflat AS
@@ -342,7 +350,8 @@ CREATE OPERATOR CLASS halfvec_ip_ops
342350
FUNCTION 3 halfvec_spherical_distance(halfvec, halfvec),
343351
FUNCTION 4 l2_norm(halfvec),
344352
FUNCTION 5 l2_normalize(halfvec),
345-
FUNCTION 6 ivfflat_halfvec_support(internal);
353+
FUNCTION 6 ivfflat_halfvec_max_dims(internal),
354+
FUNCTION 7 ivfflat_halfvec_support(internal);
346355

347356
CREATE OPERATOR CLASS halfvec_cosine_ops
348357
FOR TYPE halfvec USING ivfflat AS
@@ -352,7 +361,8 @@ CREATE OPERATOR CLASS halfvec_cosine_ops
352361
FUNCTION 3 halfvec_spherical_distance(halfvec, halfvec),
353362
FUNCTION 4 l2_norm(halfvec),
354363
FUNCTION 5 l2_normalize(halfvec),
355-
FUNCTION 6 ivfflat_halfvec_support(internal);
364+
FUNCTION 6 ivfflat_halfvec_max_dims(internal),
365+
FUNCTION 7 ivfflat_halfvec_support(internal);
356366

357367
CREATE OPERATOR CLASS halfvec_l2_ops
358368
FOR TYPE halfvec USING hnsw AS

‎sql/vector.sql‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,12 @@ COMMENT ON ACCESS METHOD hnsw IS 'hnsw index access method';
263263

264264
-- access method private functions
265265

266+
CREATE FUNCTION ivfflat_bit_max_dims(internal) RETURNS internal
267+
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
268+
269+
CREATE FUNCTION ivfflat_halfvec_max_dims(internal) RETURNS internal
270+
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
271+
266272
CREATE FUNCTION ivfflat_bit_support(internal) RETURNS internal
267273
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
268274

@@ -361,7 +367,8 @@ CREATE OPERATOR CLASS bit_hamming_ops
361367
OPERATOR 1 <~> (bit, bit) FOR ORDER BY float_ops,
362368
FUNCTION 1 hamming_distance(bit, bit),
363369
FUNCTION 3 hamming_distance(bit, bit),
364-
FUNCTION 6 ivfflat_bit_support(internal);
370+
FUNCTION 6 ivfflat_bit_max_dims(internal),
371+
FUNCTION 7 ivfflat_bit_support(internal);
365372

366373
CREATE OPERATOR CLASS bit_hamming_ops
367374
FOR TYPE bit USING hnsw AS
@@ -644,7 +651,8 @@ CREATE OPERATOR CLASS halfvec_l2_ops
644651
OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops,
645652
FUNCTION 1 halfvec_l2_squared_distance(halfvec, halfvec),
646653
FUNCTION 3 l2_distance(halfvec, halfvec),
647-
FUNCTION 6 ivfflat_halfvec_support(internal);
654+
FUNCTION 6 ivfflat_halfvec_max_dims(internal),
655+
FUNCTION 7 ivfflat_halfvec_support(internal);
648656

649657
CREATE OPERATOR CLASS halfvec_ip_ops
650658
FOR TYPE halfvec USING ivfflat AS
@@ -653,7 +661,8 @@ CREATE OPERATOR CLASS halfvec_ip_ops
653661
FUNCTION 3 halfvec_spherical_distance(halfvec, halfvec),
654662
FUNCTION 4 l2_norm(halfvec),
655663
FUNCTION 5 l2_normalize(halfvec),
656-
FUNCTION 6 ivfflat_halfvec_support(internal);
664+
FUNCTION 6 ivfflat_halfvec_max_dims(internal),
665+
FUNCTION 7 ivfflat_halfvec_support(internal);
657666

658667
CREATE OPERATOR CLASS halfvec_cosine_ops
659668
FOR TYPE halfvec USING ivfflat AS
@@ -663,7 +672,8 @@ CREATE OPERATOR CLASS halfvec_cosine_ops
663672
FUNCTION 3 halfvec_spherical_distance(halfvec, halfvec),
664673
FUNCTION 4 l2_norm(halfvec),
665674
FUNCTION 5 l2_normalize(halfvec),
666-
FUNCTION 6 ivfflat_halfvec_support(internal);
675+
FUNCTION 6 ivfflat_halfvec_max_dims(internal),
676+
FUNCTION 7 ivfflat_halfvec_support(internal);
667677

668678
CREATE OPERATOR CLASS halfvec_l2_ops
669679
FOR TYPE halfvec USING hnsw AS

‎src/ivfbuild.c‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -323,16 +323,14 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
323323
* Get max dimensions
324324
*/
325325
static int
326-
GetMaxDimensions(IvfflatType type)
326+
GetMaxDimensions(Relation index)
327327
{
328-
int maxDimensions = IVFFLAT_MAX_DIM;
328+
FmgrInfo *procinfo = IvfflatOptionalProcInfo(index, IVFFLAT_MAX_DIMS_PROC);
329329

330-
if (type == IVFFLAT_TYPE_HALFVEC)
331-
maxDimensions *= 2;
332-
else if (type == IVFFLAT_TYPE_BIT)
333-
maxDimensions *= 32;
330+
if (procinfo == NULL)
331+
return IVFFLAT_MAX_DIM;
334332

335-
return maxDimensions;
333+
return DatumGetInt32(FunctionCall1(procinfo, PointerGetDatum(NULL)));
336334
}
337335

338336
/*
@@ -367,7 +365,11 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
367365
buildstate->lists = IvfflatGetLists(index);
368366
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
369367

370-
maxDimensions = GetMaxDimensions(buildstate->type);
368+
/* Disallow varbit since require fixed dimensions */
369+
if (TupleDescAttr(index->rd_att, 0)->atttypid == VARBITOID)
370+
elog(ERROR, "type not supported for ivfflat index");
371+
372+
maxDimensions = GetMaxDimensions(index);
371373

372374
/* Require column to have dimensions to be indexed */
373375
if (buildstate->dimensions < 0)

‎src/ivfflat.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
188188
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
189189

190190
amroutine->amstrategies = 0;
191-
amroutine->amsupport = 6;
191+
amroutine->amsupport = 7;
192192
#if PG_VERSION_NUM >= 130000
193193
amroutine->amoptsprocnum = 0;
194194
#endif

‎src/ivfflat.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@
2929
#define IVFFLAT_KMEANS_DISTANCE_PROC 3
3030
#define IVFFLAT_KMEANS_NORM_PROC 4
3131
#define IVFFLAT_NORMALIZE_PROC 5
32-
#define IVFFLAT_TYPE_SUPPORT_PROC 6
32+
#define IVFFLAT_MAX_DIMS_PROC 6
33+
#define IVFFLAT_TYPE_SUPPORT_PROC 7
3334

3435
#define IVFFLAT_VERSION 1
3536
#define IVFFLAT_MAGIC_NUMBER 0x14FF1A7

‎src/ivfutils.c‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,20 @@ IvfflatUpdateList(Relation index, ListInfo listInfo,
246246
}
247247
}
248248

249+
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_max_dims);
250+
Datum
251+
ivfflat_halfvec_max_dims(PG_FUNCTION_ARGS)
252+
{
253+
PG_RETURN_INT32(IVFFLAT_MAX_DIM * 2);
254+
};
255+
256+
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_bit_max_dims);
257+
Datum
258+
ivfflat_bit_max_dims(PG_FUNCTION_ARGS)
259+
{
260+
PG_RETURN_INT32(IVFFLAT_MAX_DIM * 32);
261+
};
262+
249263
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_support);
250264
Datum
251265
ivfflat_halfvec_support(PG_FUNCTION_ARGS)

0 commit comments

Comments
 (0)