Compare commits

..

1 Commits

Author SHA1 Message Date
Andrew Kane
ece6ed7a0d Added concatenate operator for vectors 2024-03-17 15:49:38 -07:00
7 changed files with 46 additions and 311 deletions

View File

@@ -1,6 +1,6 @@
## 0.7.0 (unreleased)
- Added subscript function for vectors
- Added concatenate operator for vectors
## 0.6.2 (unreleased)

View File

@@ -682,6 +682,7 @@ Operator | Description | Added
\+ | element-wise addition |
\- | element-wise subtraction |
\* | element-wise multiplication | 0.5.0
\|\| | concatenate | 0.7.0
<-> | Euclidean distance |
<#> | negative inner product |
<=> | cosine distance |

View File

@@ -1,7 +1,9 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "ALTER EXTENSION vector UPDATE TO '0.7.0'" to load this file. \quit
CREATE FUNCTION vector_subscript(internal) RETURNS internal
CREATE FUNCTION vector_concat(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
ALTER TYPE vector SET (SUBSCRIPT = vector_subscript);
CREATE OPERATOR || (
LEFTARG = vector, RIGHTARG = vector, PROCEDURE = vector_concat
);

View File

@@ -20,16 +20,12 @@ CREATE FUNCTION vector_recv(internal, oid, integer) RETURNS vector
CREATE FUNCTION vector_send(vector) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_subscript(internal) RETURNS internal
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE vector (
INPUT = vector_in,
OUTPUT = vector_out,
TYPMOD_IN = vector_typmod_in,
RECEIVE = vector_recv,
SEND = vector_send,
SUBSCRIPT = vector_subscript,
STORAGE = external
);
@@ -103,6 +99,9 @@ CREATE FUNCTION vector_avg(double precision[]) RETURNS vector
CREATE FUNCTION vector_combine(double precision[], double precision[]) RETURNS double precision[]
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_concat(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- aggregates
CREATE AGGREGATE avg(vector) (
@@ -228,6 +227,10 @@ CREATE OPERATOR > (
RESTRICT = scalargtsel, JOIN = scalargtjoinsel
);
CREATE OPERATOR || (
LEFTARG = vector, RIGHTARG = vector, PROCEDURE = vector_concat
);
-- access methods
CREATE FUNCTION ivfflathandler(internal) RETURNS index_am_handler

View File

@@ -4,18 +4,11 @@
#include "catalog/pg_type.h"
#include "common/shortest_dec.h"
#include "executor/execExpr.h"
#include "fmgr.h"
#include "hnsw.h"
#include "ivfflat.h"
#include "lib/stringinfo.h"
#include "libpq/pqformat.h"
#include "nodes/makefuncs.h"
#include "nodes/nodeFuncs.h"
#include "nodes/subscripting.h"
#include "parser/parse_coerce.h"
#include "parser/parse_expr.h"
#include "parser/parse_node.h"
#include "port.h" /* for strtof() */
#include "utils/array.h"
#include "utils/builtins.h"
@@ -425,180 +418,6 @@ vector_send(PG_FUNCTION_ARGS)
PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
}
/*
* Transform the subscript expressions
*/
static void
vector_subscript_transform(SubscriptingRef *sbsref, List *indirection, ParseState *pstate, bool isSlice, bool isAssignment)
{
A_Indices *ai;
Node *subexpr;
if (list_length(indirection) != 1)
ereport(ERROR,
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
errmsg("vector allows only one subscript"),
parser_errposition(pstate,
exprLocation((Node *) indirection))));
ai = linitial_node(A_Indices, indirection);
if (isSlice)
{
if (ai->lidx)
{
subexpr = transformExpr(pstate, ai->lidx, pstate->p_expr_kind);
/* If it's not int4 already, try to coerce */
subexpr = coerce_to_target_type(pstate,
subexpr, exprType(subexpr),
INT4OID, -1,
COERCION_ASSIGNMENT,
COERCE_IMPLICIT_CAST,
-1);
if (subexpr == NULL)
ereport(ERROR,
(errcode(ERRCODE_DATATYPE_MISMATCH),
errmsg("vector subscript must have type integer"),
parser_errposition(pstate, exprLocation(ai->lidx))));
}
else if (!ai->is_slice)
{
/* Make a constant 1 */
subexpr = (Node *) makeConst(INT4OID,
-1,
InvalidOid,
sizeof(int32),
Int32GetDatum(1),
false,
true); /* pass by value */
}
else
{
/* Slice with omitted lower bound, put NULL into the list */
subexpr = NULL;
}
sbsref->reflowerindexpr = list_make1(subexpr);
}
else
Assert(ai->lidx == NULL && !ai->is_slice);
if (ai->uidx)
{
subexpr = transformExpr(pstate, ai->uidx, pstate->p_expr_kind);
/* If it's not int4 already, try to coerce */
subexpr = coerce_to_target_type(pstate,
subexpr, exprType(subexpr),
INT4OID, -1,
COERCION_ASSIGNMENT,
COERCE_IMPLICIT_CAST,
-1);
if (subexpr == NULL)
ereport(ERROR,
(errcode(ERRCODE_DATATYPE_MISMATCH),
errmsg("array subscript must have type integer"),
parser_errposition(pstate, exprLocation(ai->uidx))));
}
else
{
/* Slice with omitted upper bound, put NULL into the list */
Assert(isSlice && ai->is_slice);
subexpr = NULL;
}
sbsref->refupperindexpr = list_make1(subexpr);
if (isSlice)
sbsref->refrestype = sbsref->refcontainertype;
else
sbsref->refrestype = FLOAT4OID;
}
/*
* Fetch a vector element
*/
static void
vector_subscript_fetch(ExprState *state, ExprEvalStep *op, ExprContext *econtext)
{
SubscriptingRefState *sbsrefstate = op->d.sbsref.state;
Vector *vec = DatumGetVector(*op->resvalue);
int index = DatumGetInt32(sbsrefstate->upperindex[0]);
if (index < 1 || index > vec->dim)
*op->resnull = true;
else
*op->resvalue = Float4GetDatum(vec->x[index - 1]);
}
/*
* Fetch a vector slice
*/
static void
vector_subscript_fetch_slice(ExprState *state, ExprEvalStep *op, ExprContext *econtext)
{
SubscriptingRefState *sbsrefstate = op->d.sbsref.state;
if (sbsrefstate->upperprovided[0] && sbsrefstate->upperindexnull[0])
*op->resnull = true;
else if (sbsrefstate->lowerprovided[0] && sbsrefstate->lowerindexnull[0])
*op->resnull = true;
else
{
Vector *vec = DatumGetVector(*op->resvalue);
int upperindex = sbsrefstate->upperprovided[0] ? DatumGetInt32(sbsrefstate->upperindex[0]) : vec->dim;
int lowerindex = sbsrefstate->lowerprovided[0] ? DatumGetInt32(sbsrefstate->lowerindex[0]) : 1;
int dim;
Vector *result;
if (upperindex > vec->dim)
upperindex = vec->dim;
if (lowerindex < 1)
lowerindex = 1;
dim = upperindex - lowerindex + 1;
CheckDim(dim);
result = InitVector(dim);
for (int i = 0; i < dim; i++)
result->x[i] = vec->x[lowerindex + i - 1];
*op->resvalue = PointerGetDatum(result);
}
}
/*
* Set up execution state for a vector subscript operation
*/
static void
vector_exec_setup(const SubscriptingRef *sbsref, SubscriptingRefState *sbsrefstate, SubscriptExecSteps *methods)
{
methods->sbs_check_subscripts = NULL;
if (sbsrefstate->numlower != 0)
methods->sbs_fetch = vector_subscript_fetch_slice;
else
methods->sbs_fetch = vector_subscript_fetch;
methods->sbs_assign = NULL;
methods->sbs_fetch_old = NULL;
}
/*
* Subscript handler
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_subscript);
Datum
vector_subscript(PG_FUNCTION_ARGS)
{
static const SubscriptRoutines sbsroutines = {
.transform = vector_subscript_transform,
.exec_setup = vector_exec_setup,
.fetch_strict = true, /* fetch returns NULL for NULL inputs */
.fetch_leakproof = true, /* fetch returns NULL for bad subscript */
.store_leakproof = false /* ... but assignment throws error */
};
PG_RETURN_POINTER(&sbsroutines);
}
/*
* Convert vector to vector
* This is needed to check the type modifier
@@ -1341,3 +1160,28 @@ vector_avg(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(result);
}
/*
* Concat vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_concat);
Datum
vector_concat(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
Vector *result;
int dim = a->dim + b->dim;
CheckDim(dim);
result = InitVector(dim);
for (int i = 0; i < a->dim; i++)
result->x[i] = a->x[i];
for (int i = 0; i < b->dim; i++)
result->x[i + a->dim] = b->x[i];
PG_RETURN_POINTER(result);
}

View File

@@ -24,112 +24,14 @@ SELECT '[1e37]'::vector * '[1e37]';
ERROR: value out of range: overflow
SELECT '[1e-37]'::vector * '[1e-37]';
ERROR: value out of range: underflow
SELECT ('[1,2,3]'::vector)[0];
vector
--------
SELECT '[1,2,3]'::vector || '[4,5]'::vector;
?column?
-------------
[1,2,3,4,5]
(1 row)
SELECT ('[1,2,3]'::vector)[1];
vector
--------
1
(1 row)
SELECT ('[1,2,3]'::vector)[2];
vector
--------
2
(1 row)
SELECT ('[1,2,3]'::vector)[3];
vector
--------
3
(1 row)
SELECT ('[1,2,3]'::vector)[4];
vector
--------
(1 row)
SELECT ('[1,2,3]'::vector)[1:1];
vector
--------
[1]
(1 row)
SELECT ('[1,2,3]'::vector)[1:2];
vector
--------
[1,2]
(1 row)
SELECT ('[1,2,3]'::vector)[2:4];
vector
--------
[2,3]
(1 row)
SELECT ('[1,2,3]'::vector)[-2:2];
vector
--------
[1,2]
(1 row)
SELECT ('[1,2,3]'::vector)[2:1];
ERROR: vector must have at least 1 dimension
SELECT ('[1,2,3]'::vector)[:];
vector
---------
[1,2,3]
(1 row)
SELECT ('[1,2,3]'::vector)[:2];
vector
--------
[1,2]
(1 row)
SELECT ('[1,2,3]'::vector)[2:];
vector
--------
[2,3]
(1 row)
SELECT ('[1,2,3]'::vector)[:4];
vector
---------
[1,2,3]
(1 row)
SELECT ('[1,2,3]'::vector)[-2:];
vector
---------
[1,2,3]
(1 row)
SELECT ('[1,2,3]'::vector)[NULL];
vector
--------
(1 row)
SELECT ('[1,2,3]'::vector)[NULL:2];
vector
--------
(1 row)
SELECT ('[1,2,3]'::vector)[2:NULL];
vector
--------
(1 row)
SELECT ('[1,2,3]'::vector)[1][1];
ERROR: vector allows only one subscript
SELECT array_fill(0, ARRAY[16000])::vector || '[1]'::vector;
ERROR: vector cannot have more than 16000 dimensions
SELECT '[1,2,3]'::vector = '[1,2,3]';
?column?
----------

View File

@@ -6,25 +6,8 @@ SELECT '[1,2,3]'::vector * '[4,5,6]';
SELECT '[1e37]'::vector * '[1e37]';
SELECT '[1e-37]'::vector * '[1e-37]';
SELECT ('[1,2,3]'::vector)[0];
SELECT ('[1,2,3]'::vector)[1];
SELECT ('[1,2,3]'::vector)[2];
SELECT ('[1,2,3]'::vector)[3];
SELECT ('[1,2,3]'::vector)[4];
SELECT ('[1,2,3]'::vector)[1:1];
SELECT ('[1,2,3]'::vector)[1:2];
SELECT ('[1,2,3]'::vector)[2:4];
SELECT ('[1,2,3]'::vector)[-2:2];
SELECT ('[1,2,3]'::vector)[2:1];
SELECT ('[1,2,3]'::vector)[:];
SELECT ('[1,2,3]'::vector)[:2];
SELECT ('[1,2,3]'::vector)[2:];
SELECT ('[1,2,3]'::vector)[:4];
SELECT ('[1,2,3]'::vector)[-2:];
SELECT ('[1,2,3]'::vector)[NULL];
SELECT ('[1,2,3]'::vector)[NULL:2];
SELECT ('[1,2,3]'::vector)[2:NULL];
SELECT ('[1,2,3]'::vector)[1][1];
SELECT '[1,2,3]'::vector || '[4,5]'::vector;
SELECT array_fill(0, ARRAY[16000])::vector || '[1]'::vector;
SELECT '[1,2,3]'::vector = '[1,2,3]';
SELECT '[1,2,3]'::vector = '[1,2]';