database/openGauss-server
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
828674b29ce3c58fa6e234ce07c8ea14d8ff27c2
openGauss-server/src/common/backend/utils/adt/vector.cpp
jiwenke 828674b29c 新增hnsw & ivfflat reloptions
2025-04-22 14:40:28 +08:00

1507 lines
37 KiB
C++
Raw Blame History

/*
* Copyright (c) 2024 Huawei Technologies Co.,Ltd.
*
* openGauss is licensed under Mulan PSL v2.
* You can use this software according to the terms and conditions of the Mulan PSL v2.
* You may obtain a copy of Mulan PSL v2 at:
*
* http://license.coscl.org.cn/MulanPSL2
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
* See the Mulan PSL v2 for more details.
* -------------------------------------------------------------------------
*
* vector.cpp
*
* IDENTIFICATION
* src/common/backend/utils/adt/vector.cpp
*
* -------------------------------------------------------------------------
*/
#include "postgres.h"
#include <cmath>
#ifdef __aarch64__
#include <arm_neon.h>
#else
#include <immintrin.h>
#endif
#include "access/datavec/bitvec.h"
#include "catalog/pg_type.h"
#include "fmgr.h"
#include "access/datavec/halfutils.h"
#include "access/datavec/halfvec.h"
#include "access/datavec/hnsw.h"
#include "access/datavec/ivfflat.h"
#include "lib/stringinfo.h"
#include "libpq/pqformat.h"
#include "port.h" /* for strtof() */
#include "access/datavec/shortest_dec.h"
#include "access/datavec/sparsevec.h"
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/lsyscache.h"
#include "utils/numeric.h"
#include "commands/extension.h"
#include "knl/knl_session.h"
#include "access/datavec/vector.h"
#define TYPALIGN_DOUBLE 'd'
#define TYPALIGN_INT 'i'
#define STATE_DIMS(x) (ARR_DIMS(x)[0] - 1)
#define CreateStateDatums(dim) palloc(sizeof(Datum) * ((dim) + 1))
#if defined(USE_TARGET_CLONES) && !defined(__FMA__)
#define VECTOR_TARGET_CLONES __attribute__((target_clones("default", "fma")))
#else
#define VECTOR_TARGET_CLONES
#endif
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
/*
* Ensure same dimensions
*/
static inline void CheckDims(Vector *a, Vector *b)
{
if (a->dim != b->dim)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION), errmsg("different vector dimensions %d and %d", a->dim, b->dim)));
}
/*
* Ensure expected dimensions
*/
static inline void CheckExpectedDim(int32 typmod, int dim)
{
if (typmod != -1 && typmod != dim)
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("expected %d dimensions, not %d", typmod, dim)));
}
/*
* Ensure valid dimensions
*/
static inline void CheckDim(int dim)
{
if (dim < 1)
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("vector must have at least 1 dimension")));
if (dim > VECTOR_MAX_DIM)
ereport(ERROR, (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("vector cannot have more than %d dimensions", VECTOR_MAX_DIM)));
}
/*
* Ensure finite element
*/
static inline void CheckElement(float value)
{
if (isnan(value))
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("NaN not allowed in vector")));
if (isinf(value))
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("infinite value not allowed in vector")));
}
/*
* Allocate and initialize a new vector
*/
Vector *InitVector(int dim)
{
Vector *result;
int size;
size = VECTOR_SIZE(dim);
result = (Vector *)palloc0(size);
SET_VARSIZE(result, size);
result->dim = dim;
return result;
}
/*
* Check for whitespace, since array_isspace() is static
*/
static inline bool VectorIsspace(char ch)
{
if (ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r' || ch == '\v' || ch == '\f') {
return true;
}
return false;
}
/*
* Check state array
*/
static float8 *CheckStateArray(ArrayType *statearray, const char *caller)
{
if (ARR_NDIM(statearray) != 1 || ARR_DIMS(statearray)[0] < 1 || ARR_HASNULL(statearray))
elog(ERROR, "%s: expected state array", caller);
return (float8 *)ARR_DATA_PTR(statearray);
}
static pg_noinline void float_overflow_error(void)
{
ereport(ERROR, (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE), errmsg("value out of range: overflow")));
}
static pg_noinline void float_underflow_error(void)
{
ereport(ERROR, (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE), errmsg("value out of range: underflow")));
}
/*
* Convert textual representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_in);
Datum vector_in(PG_FUNCTION_ARGS)
{
char *lit = PG_GETARG_CSTRING(0);
int32 typmod = PG_GETARG_INT32(2);
float x[VECTOR_MAX_DIM];
int dim = 0;
char *pt = lit;
Vector *result;
while (VectorIsspace(*pt)) {
pt++;
}
if (*pt != '[')
ereport(ERROR, (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type vector: \"%s\"", lit),
errdetail("Vector contents must start with \"[\".")));
pt++;
while (VectorIsspace(*pt)) {
pt++;
}
if (*pt == ']') {
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("vector must have at least 1 dimension")));
}
for (;;) {
float val;
char *stringEnd;
if (dim == VECTOR_MAX_DIM)
ereport(ERROR, (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("vector cannot have more than %d dimensions", VECTOR_MAX_DIM)));
while (VectorIsspace(*pt)) {
pt++;
}
/* Check for empty string like float4in */
if (*pt == '\0')
ereport(ERROR, (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type vector: \"%s\"", lit)));
errno = 0;
/* Use strtof like float4in to avoid a double-rounding problem */
/* Postgres sets LC_NUMERIC to C on startup */
val = strtof(pt, &stringEnd);
if (stringEnd == pt) {
ereport(ERROR, (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type vector: \"%s\"", lit)));
}
/* Check for range error like float4in */
if (errno == ERANGE && isinf(val))
ereport(ERROR, (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("\"%s\" is out of range for type vector", pnstrdup(pt, stringEnd - pt))));
CheckElement(val);
x[dim++] = val;
pt = stringEnd;
while (VectorIsspace(*pt)) {
pt++;
}
if (*pt == ',') {
pt++;
} else if (*pt == ']') {
pt++;
break;
} else {
ereport(ERROR, (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type vector: \"%s\"", lit)));
}
}
/* Only whitespace is allowed after the closing brace */
while (VectorIsspace(*pt)) {
pt++;
}
if (*pt != '\0')
ereport(ERROR, (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type vector: \"%s\"", lit),
errdetail("Junk after closing right brace.")));
CheckDim(dim);
CheckExpectedDim(typmod, dim);
result = InitVector(dim);
for (int i = 0; i < dim; i++) {
result->x[i] = x[i];
}
PG_RETURN_POINTER(result);
}
#define AppendChar(ptr, c) (*(ptr)++ = (c))
#define AppendFloat(ptr, f) ((ptr) += FloatToShortestDecimalBufn((f), (ptr)))
/*
* Convert internal representation to textual representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_out);
Datum vector_out(PG_FUNCTION_ARGS)
{
Vector *vector = PG_GETARG_VECTOR_P(0);
int dim = vector->dim;
char *buf;
char *ptr;
/*
* Need:
*
* dim * (FLOAT_SHORTEST_DECIMAL_LEN - 1) bytes for
* FloatToShortestDecimalBufn
*
* dim - 1 bytes for separator
*
* 3 bytes for [, ], and \0
*/
buf = (char *)palloc(FLOAT_SHORTEST_DECIMAL_LEN * dim + 2);
ptr = buf;
AppendChar(ptr, '[');
for (int i = 0; i < dim; i++) {
if (i > 0) {
AppendChar(ptr, ',');
}
AppendFloat(ptr, vector->x[i]);
}
AppendChar(ptr, ']');
*ptr = '\0';
PG_FREE_IF_COPY(vector, 0);
PG_RETURN_CSTRING(buf);
}
/*
* Print vector - useful for debugging
*/
void PrintVector(char *msg, Vector *vector)
{
char *out = DatumGetPointer(DirectFunctionCall1(vector_out, PointerGetDatum(vector)));
elog(INFO, "%s = %s", msg, out);
pfree(out);
}
/*
* Convert type modifier
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_typmod_in);
Datum vector_typmod_in(PG_FUNCTION_ARGS)
{
ArrayType *ta = PG_GETARG_ARRAYTYPE_P(0);
int32 *tl;
int n;
tl = ArrayGetIntegerTypmods(ta, &n);
if (n != 1) {
ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE), errmsg("invalid type modifier")));
}
if (*tl < 1) {
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE), errmsg("dimensions for type vector must be at least 1")));
}
if (*tl > VECTOR_MAX_DIM) {
ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("dimensions for type vector cannot exceed %d", VECTOR_MAX_DIM)));
}
PG_RETURN_INT32(*tl);
}
/*
* Convert external binary representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_recv);
Datum vector_recv(PG_FUNCTION_ARGS)
{
StringInfo buf = (StringInfo)PG_GETARG_POINTER(0);
int32 typmod = PG_GETARG_INT32(2);
Vector *result;
int16 dim;
int16 unused;
dim = pq_getmsgint(buf, sizeof(int16));
unused = pq_getmsgint(buf, sizeof(int16));
CheckDim(dim);
CheckExpectedDim(typmod, dim);
if (unused != 0) {
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("expected unused to be 0, not %d", unused)));
}
result = InitVector(dim);
for (int i = 0; i < dim; i++) {
result->x[i] = pq_getmsgfloat4(buf);
CheckElement(result->x[i]);
}
PG_RETURN_POINTER(result);
}
/*
* Convert internal representation to the external binary representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_send);
Datum vector_send(PG_FUNCTION_ARGS)
{
Vector *vec = PG_GETARG_VECTOR_P(0);
StringInfoData buf;
pq_begintypsend(&buf);
pq_sendint(&buf, vec->dim, sizeof(int16));
pq_sendint(&buf, vec->unused, sizeof(int16));
for (int i = 0; i < vec->dim; i++)
pq_sendfloat4(&buf, vec->x[i]);
PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
}
/*
* Convert vector to vector
* This is needed to check the type modifier
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector);
Datum vector(PG_FUNCTION_ARGS)
{
Vector *vec = PG_GETARG_VECTOR_P(0);
int32 typmod = PG_GETARG_INT32(1);
CheckExpectedDim(typmod, vec->dim);
PG_RETURN_POINTER(vec);
}
/*
* Convert array to vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(array_to_vector);
Datum array_to_vector(PG_FUNCTION_ARGS)
{
ArrayType *array = PG_GETARG_ARRAYTYPE_P(0);
int32 typmod = PG_GETARG_INT32(1);
Vector *result;
int16 typlen;
bool typbyval;
char typalign;
Datum *elemsp;
int nelemsp;
if (ARR_NDIM(array) > 1) {
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("array must be 1-D")));
}
if (ARR_HASNULL(array) && array_contains_nulls(array)) {
ereport(ERROR, (errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED), errmsg("array must not contain nulls")));
}
get_typlenbyvalalign(ARR_ELEMTYPE(array), &typlen, &typbyval, &typalign);
deconstruct_array(array, ARR_ELEMTYPE(array), typlen, typbyval, typalign, &elemsp, NULL, &nelemsp);
CheckDim(nelemsp);
CheckExpectedDim(typmod, nelemsp);
result = InitVector(nelemsp);
if (ARR_ELEMTYPE(array) == INT4OID) {
for (int i = 0; i < nelemsp; i++)
result->x[i] = DatumGetInt32(elemsp[i]);
} else if (ARR_ELEMTYPE(array) == FLOAT8OID) {
for (int i = 0; i < nelemsp; i++)
result->x[i] = DatumGetFloat8(elemsp[i]);
} else if (ARR_ELEMTYPE(array) == FLOAT4OID) {
for (int i = 0; i < nelemsp; i++)
result->x[i] = DatumGetFloat4(elemsp[i]);
} else if (ARR_ELEMTYPE(array) == NUMERICOID) {
for (int i = 0; i < nelemsp; i++)
result->x[i] = DatumGetFloat4(DirectFunctionCall1(numeric_float4, elemsp[i]));
} else {
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("unsupported array type")));
}
/*
* Free allocation from deconstruct_array. Do not free individual elements
* when pass-by-reference since they point to original array.
*/
pfree(elemsp);
/* Check elements */
for (int i = 0; i < result->dim; i++) {
CheckElement(result->x[i]);
}
PG_RETURN_POINTER(result);
}
/*
* Convert vector to float4[]
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_to_float4);
Datum vector_to_float4(PG_FUNCTION_ARGS)
{
Vector *vec = PG_GETARG_VECTOR_P(0);
Datum *datums;
ArrayType *result;
datums = (Datum *)palloc(sizeof(Datum) * vec->dim);
for (int i = 0; i < vec->dim; i++) {
datums[i] = Float4GetDatum(vec->x[i]);
}
/* Use TYPALIGN_INT for float4 */
result = construct_array(datums, vec->dim, FLOAT4OID, sizeof(float4), true, TYPALIGN_INT);
pfree(datums);
PG_RETURN_POINTER(result);
}
/*
* Convert half vector to vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_to_vector);
Datum halfvec_to_vector(PG_FUNCTION_ARGS)
{
HalfVector *vec = PG_GETARG_HALFVEC_P(0);
int32 typmod = PG_GETARG_INT32(1);
Vector *result;
CheckDim(vec->dim);
CheckExpectedDim(typmod, vec->dim);
result = InitVector(vec->dim);
for (int i = 0; i < vec->dim; i++) {
result->x[i] = HalfToFloat4(vec->x[i]);
}
PG_RETURN_POINTER(result);
}
inline void prefetch_L1(const void *address)
{
#if defined(__SSE2__)
_mm_prefetch((const char*)address, _MM_HINT_T0);
#elif defined(__aarch64__)
asm volatile("prfm PLDL1KEEP, [%0]" : : "r"(address));
#else
__builtin_prefetch(address, 0, 3); // L3 cache
#endif
}
#ifdef __aarch64__
VECTOR_TARGET_CLONES static float
VectorL2SquaredDistance(int size, float *a, float *b)
{
float32x4_t sum = vdupq_n_f32(0.0f);
int prefetch_len = 8;
int batch_num = 4;
int i;
for (i = 0; i < size - (size % batch_num); i += batch_num) {
prefetch_L1(a + i + prefetch_len);
prefetch_L1(b + i + prefetch_len);
float32x4_t va = vld1q_f32(a + i);
float32x4_t vb = vld1q_f32(b + i);
float32x4_t diff = vsubq_f32(va, vb);
sum = vfmaq_f32(sum, diff, diff);
}
float scalar_sum = 0.0f;
if (size % batch_num > 0) {
int remaining_size = size % batch_num;
for (int j = 0; j < remaining_size; ++j) {
prefetch_L1(a + i + j);
prefetch_L1(b + i + j);
float value_a = a[i + j];
float value_b = b[i + j];
float diff = value_a - value_b;
float sq_diff = diff * diff;
scalar_sum += sq_diff;
}
}
sum = vpaddq_f32(sum, sum);
sum = vpaddq_f32(sum, sum);
float res = vgetq_lane_f32(sum, 0);
return res + scalar_sum;
}
#elif defined(__x86_64__)
static inline __m128 masked_read(int d, const float *x)
{
__attribute__((__aligned__(16))) float buf[4];
Assert(0 <= d && d < 4); /* reads 0 <= d < 4 floats as __m128 */
memset((void*)buf, 0, sizeof(buf));
switch (d) {
case 3:
buf[2] = x[2];
case 2:
buf[1] = x[1];
case 1:
buf[0] = x[0];
default:
break;
}
return _mm_load_ps(buf);
}
VECTOR_TARGET_CLONES static float
VectorL2SquaredDistance(int dim, float *ax, float *bx)
{
float* x = (float*)ax;
float* y = (float*)bx;
size_t d = (size_t)dim;
int batch_num1 = 8;
int batch_num2 = 4;
__m256 msum1 = _mm256_setzero_ps();
while (d >= batch_num1) {
__m256 mx = _mm256_loadu_ps(x);
x += batch_num1;
__m256 my = _mm256_loadu_ps(y);
y += batch_num1;
const __m256 a_m_b1 = mx - my;
msum1 += a_m_b1 * a_m_b1;
d -= batch_num1;
}
__m128 msum2 = _mm256_extractf128_ps(msum1, 1);
msum2 += _mm256_extractf128_ps(msum1, 0);
if (d >= batch_num2) {
__m128 mx = _mm_loadu_ps(x);
x += batch_num2;
__m128 my = _mm_loadu_ps(y);
y += batch_num2;
const __m128 a_m_b1 = mx - my;
msum2 += a_m_b1 * a_m_b1;
d -= batch_num2;
}
if (d > 0) {
__m128 mx = masked_read(d, x);
__m128 my = masked_read(d, y);
__m128 a_m_b1 = mx - my;
msum2 += a_m_b1 * a_m_b1;
}
msum2 = _mm_hadd_ps(msum2, msum2);
msum2 = _mm_hadd_ps(msum2, msum2);
return _mm_cvtss_f32(msum2);
}
#else
VECTOR_TARGET_CLONES static float VectorL2SquaredDistance(int dim, float *ax, float *bx)
{
float distance = 0.0;
/* Auto-vectorized */
for (int i = 0; i < dim; i++) {
float diff = ax[i] - bx[i];
distance += diff * diff;
}
return distance;
}
#endif
/*
* Get the L2 distance between vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(l2_distance);
Datum l2_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8(sqrt((double)VectorL2SquaredDistance(a->dim, a->x, b->x)));
}
/*
* Get the L2 squared distance between vectors
* This saves a sqrt calculation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_l2_squared_distance);
Datum vector_l2_squared_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8((double)VectorL2SquaredDistance(a->dim, a->x, b->x));
}
#ifdef __aarch64__
VECTOR_TARGET_CLONES static float
VectorInnerProduct(int dim, float *ax, float *bx)
{
float dis = 0.0f;
float32x4_t sum = vdupq_n_f32(0.0f);
float *pta = ax;
float *ptb = bx;
int i = 0;
int prefetch_len = 8;
int batch_num = 4;
for (; i + batch_num <= dim; i += batch_num) {
prefetch_L1(pta + prefetch_len);
prefetch_L1(ptb + prefetch_len);
float32x4_t sub_a = vld1q_f32(pta);
float32x4_t sub_b = vld1q_f32(ptb);
sum = vmlaq_f32(sum, sub_a, sub_b);
pta += batch_num;
ptb += batch_num;
}
dis = vaddvq_f32(sum);
for (; i < dim; ++i) {
dis += ax[i] * bx[i];
}
return dis;
}
#else
VECTOR_TARGET_CLONES static float VectorInnerProduct(int dim, float *ax, float *bx)
{
float distance = 0.0;
/* Auto-vectorized */
for (int i = 0; i < dim; i++) {
distance += ax[i] * bx[i];
}
return distance;
}
#endif
/*
* Get the inner product of two vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(inner_product);
Datum inner_product(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8((double)VectorInnerProduct(a->dim, a->x, b->x));
}
/*
* Get the negative inner product of two vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_negative_inner_product);
Datum vector_negative_inner_product(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8((double)-VectorInnerProduct(a->dim, a->x, b->x));
}
VECTOR_TARGET_CLONES static double VectorCosineSimilarity(int dim, float *ax, float *bx)
{
float similarity = 0.0;
float norma = 0.0;
float normb = 0.0;
/* Auto-vectorized */
for (int i = 0; i < dim; i++) {
similarity += ax[i] * bx[i];
norma += ax[i] * ax[i];
normb += bx[i] * bx[i];
}
/* Use sqrt(a * b) over sqrt(a) * sqrt(b) */
return static_cast<double>(similarity) / sqrt(static_cast<double>(norma) * static_cast<double>(normb));
}
/*
* Get the cosine distance between two vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(cosine_distance);
Datum cosine_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
double similarity;
CheckDims(a, b);
similarity = VectorCosineSimilarity(a->dim, a->x, b->x);
#ifdef _MSC_VER
/* /fp:fast may not propagate NaN */
if (isnan(similarity)) {
PG_RETURN_FLOAT8(NAN);
}
#endif
/* Keep in range */
if (similarity > 1) {
similarity = 1.0;
} else if (similarity < -1) {
similarity = -1.0;
}
PG_RETURN_FLOAT8(1.0 - similarity);
}
/*
* Get the distance for spherical k-means
* Currently uses angular distance since needs to satisfy triangle inequality
* Assumes inputs are unit vectors (skips norm)
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_spherical_distance);
Datum vector_spherical_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
double distance;
CheckDims(a, b);
distance = (double)VectorInnerProduct(a->dim, a->x, b->x);
/* Prevent NaN with acos with loss of precision */
if (distance > 1) {
distance = 1;
} else if (distance < -1) {
distance = -1;
}
PG_RETURN_FLOAT8(acos(distance) / M_PI);
}
/* Does not require FMA, but keep logic simple */
VECTOR_TARGET_CLONES static float VectorL1Distance(int dim, float *ax, float *bx)
{
float distance = 0.0;
/* Auto-vectorized */
for (int i = 0; i < dim; i++) {
distance += fabsf(ax[i] - bx[i]);
}
return distance;
}
/*
* Get the L1 distance between two vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(l1_distance);
Datum l1_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8((double)VectorL1Distance(a->dim, a->x, b->x));
}
/*
* Get the dimensions of a vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_dims);
Datum vector_dims(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
PG_RETURN_INT32(a->dim);
}
/*
* Get the L2 norm of a vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_norm);
Datum vector_norm(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
float *ax = a->x;
double norm = 0.0;
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++) {
norm += (double)ax[i] * (double)ax[i];
}
PG_RETURN_FLOAT8(sqrt(norm));
}
/*
* Normalize a vector with the L2 norm
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(l2_normalize);
Datum l2_normalize(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
float *ax = a->x;
double norm = 0;
Vector *result;
float *rx;
result = InitVector(a->dim);
rx = result->x;
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++) {
norm += (double)ax[i] * (double)ax[i];
}
norm = sqrt(norm);
/* Return zero vector for zero norm */
if (norm > 0) {
for (int i = 0; i < a->dim; i++) {
rx[i] = ax[i] / norm;
}
/* Check for overflow */
for (int i = 0; i < a->dim; i++) {
if (isinf(rx[i])) {
float_overflow_error();
}
}
}
PG_RETURN_POINTER(result);
}
/*
* Add vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_add);
Datum vector_add(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
Vector *result;
float *rx;
CheckDims(a, b);
result = InitVector(a->dim);
rx = result->x;
/* Auto-vectorized */
for (int i = 0, imax = a->dim; i < imax; i++) {
rx[i] = ax[i] + bx[i];
}
/* Check for overflow */
for (int i = 0, imax = a->dim; i < imax; i++) {
if (isinf(rx[i])) {
float_overflow_error();
}
}
PG_RETURN_POINTER(result);
}
/*
* Subtract vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_sub);
Datum vector_sub(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
Vector *result;
float *rx;
CheckDims(a, b);
result = InitVector(a->dim);
rx = result->x;
/* Auto-vectorized */
for (int i = 0, imax = a->dim; i < imax; i++) {
rx[i] = ax[i] - bx[i];
}
/* Check for overflow */
for (int i = 0, imax = a->dim; i < imax; i++) {
if (isinf(rx[i])) {
float_overflow_error();
}
}
PG_RETURN_POINTER(result);
}
/*
* Multiply vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_mul);
Datum vector_mul(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
Vector *result;
float *rx;
CheckDims(a, b);
result = InitVector(a->dim);
rx = result->x;
/* Auto-vectorized */
for (int i = 0, imax = a->dim; i < imax; i++) {
rx[i] = ax[i] * bx[i];
}
/* Check for overflow and underflow */
for (int i = 0, imax = a->dim; i < imax; i++) {
if (isinf(rx[i])) {
float_overflow_error();
}
if (rx[i] == 0 && !(ax[i] == 0 || bx[i] == 0)) {
float_underflow_error();
}
}
PG_RETURN_POINTER(result);
}
/*
* Concatenate 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);
}
/*
* Quantize a vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(binary_quantize);
Datum binary_quantize(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
float *ax = a->x;
VarBit *result = InitBitVector(a->dim);
unsigned char *rx = VARBITS(result);
for (int i = 0; i < a->dim; i++) {
rx[i / 8] |= (ax[i] > 0) << (7 - (i % 8));
}
PG_RETURN_VARBIT_P(result);
}
/*
* Get a subvector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(subvector);
Datum subvector(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
int32 start = PG_GETARG_INT32(1);
int32 count = PG_GETARG_INT32(2);
int32 end;
float *ax = a->x;
Vector *result;
int dim;
if (count < 1) {
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("vector must have at least 1 dimension")));
}
/*
* Check if (start + count > a->dim), avoiding integer overflow. a->dim
* and count are both positive, so a->dim - count won't overflow.
*/
if (start > a->dim - count) {
end = a->dim + 1;
} else {
end = start + count;
}
/* Indexing starts at 1, like substring */
if (start < 1) {
start = 1;
} else if (start > a->dim) {
ereport(ERROR, (errcode(ERRCODE_DATA_EXCEPTION), errmsg("vector must have at least 1 dimension")));
}
dim = end - start;
CheckDim(dim);
result = InitVector(dim);
for (int i = 0; i < dim; i++) {
result->x[i] = ax[start - 1 + i];
}
PG_RETURN_POINTER(result);
}
/*
* Internal helper to compare vectors
*/
int vector_cmp_internal(Vector *a, Vector *b)
{
int dim = Min(a->dim, b->dim);
/* Check values before dimensions to be consistent with Postgres arrays */
for (int i = 0; i < dim; i++) {
if (a->x[i] < b->x[i]) {
return -1;
}
if (a->x[i] > b->x[i]) {
return 1;
}
}
if (a->dim < b->dim) {
return -1;
}
if (a->dim > b->dim) {
return 1;
}
return 0;
}
/*
* Less than
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_lt);
Datum vector_lt(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) < 0);
}
/*
* Less than or equal
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_le);
Datum vector_le(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) <= 0);
}
/*
* Equal
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_eq);
Datum vector_eq(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) == 0);
}
/*
* Not equal
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_ne);
Datum vector_ne(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) != 0);
}
/*
* Greater than or equal
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_ge);
Datum vector_ge(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) >= 0);
}
/*
* Greater than
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_gt);
Datum vector_gt(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) > 0);
}
/*
* Compare vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_cmp);
Datum vector_cmp(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_INT32(vector_cmp_internal(a, b));
}
/*
* Accumulate vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_accum);
Datum vector_accum(PG_FUNCTION_ARGS)
{
ArrayType *statearray = PG_GETARG_ARRAYTYPE_P(0);
Vector *newval = PG_GETARG_VECTOR_P(1);
float8 *statevalues;
int16 dim;
bool newarr;
float8 n;
Datum *statedatums;
float *x = newval->x;
ArrayType *result;
/* Check array before using */
statevalues = CheckStateArray(statearray, "vector_accum");
dim = STATE_DIMS(statearray);
newarr = dim == 0;
if (newarr) {
dim = newval->dim;
} else {
CheckExpectedDim(dim, newval->dim);
}
n = statevalues[0] + 1.0;
statedatums = (Datum *)CreateStateDatums(dim);
statedatums[0] = Float8GetDatum(n);
if (newarr) {
for (int i = 0; i < dim; i++) {
statedatums[i + 1] = Float8GetDatum((double)x[i]);
}
} else {
for (int i = 0; i < dim; i++) {
double v = statevalues[i + 1] + x[i];
/* Check for overflow */
if (isinf(v)) {
float_overflow_error();
}
statedatums[i + 1] = Float8GetDatum(v);
}
}
/* Use float8 array like float4_accum */
result = construct_array(statedatums, dim + 1, FLOAT8OID, sizeof(float8), FLOAT8PASSBYVAL, TYPALIGN_DOUBLE);
pfree(statedatums);
PG_RETURN_ARRAYTYPE_P(result);
}
/*
* Combine vectors or half vectors (also used for halfvec_combine)
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_combine);
Datum vector_combine(PG_FUNCTION_ARGS)
{
/* Must also update parameters of halfvec_combine if modifying */
ArrayType *statearray1 = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *statearray2 = PG_GETARG_ARRAYTYPE_P(1);
float8 *statevalues1;
float8 *statevalues2;
float8 n;
float8 n1;
float8 n2;
int16 dim;
Datum *statedatums;
ArrayType *result;
/* Check arrays before using */
statevalues1 = CheckStateArray(statearray1, "vector_combine");
statevalues2 = CheckStateArray(statearray2, "vector_combine");
n1 = statevalues1[0];
n2 = statevalues2[0];
if (n1 == 0.0) {
n = n2;
dim = STATE_DIMS(statearray2);
statedatums = (Datum *)CreateStateDatums(dim);
for (int i = 1; i <= dim; i++)
statedatums[i] = Float8GetDatum(statevalues2[i]);
} else if (n2 == 0.0) {
n = n1;
dim = STATE_DIMS(statearray1);
statedatums = (Datum *)CreateStateDatums(dim);
for (int i = 1; i <= dim; i++)
statedatums[i] = Float8GetDatum(statevalues1[i]);
} else {
n = n1 + n2;
dim = STATE_DIMS(statearray1);
CheckExpectedDim(dim, STATE_DIMS(statearray2));
statedatums = (Datum *)CreateStateDatums(dim);
for (int i = 1; i <= dim; i++) {
double v = statevalues1[i] + statevalues2[i];
/* Check for overflow */
if (isinf(v))
float_overflow_error();
statedatums[i] = Float8GetDatum(v);
}
}
statedatums[0] = Float8GetDatum(n);
result = construct_array(statedatums, dim + 1, FLOAT8OID, sizeof(float8), FLOAT8PASSBYVAL, TYPALIGN_DOUBLE);
pfree(statedatums);
PG_RETURN_ARRAYTYPE_P(result);
}
/*
* Average vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_avg);
Datum vector_avg(PG_FUNCTION_ARGS)
{
ArrayType *statearray = PG_GETARG_ARRAYTYPE_P(0);
float8 *statevalues;
float8 n;
uint16 dim;
Vector *result;
/* Check array before using */
statevalues = CheckStateArray(statearray, "vector_avg");
n = statevalues[0];
/* SQL defines AVG of no values to be NULL */
if (n == 0.0) {
PG_RETURN_NULL();
}
/* Create vector */
dim = STATE_DIMS(statearray);
CheckDim(dim);
result = InitVector(dim);
for (int i = 0; i < dim; i++) {
result->x[i] = statevalues[i + 1] / n;
CheckElement(result->x[i]);
}
PG_RETURN_POINTER(result);
}
/*
* Convert sparse vector to dense vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(sparsevec_to_vector);
Datum sparsevec_to_vector(PG_FUNCTION_ARGS)
{
SparseVector *svec = PG_GETARG_SPARSEVEC_P(0);
int32 typmod = PG_GETARG_INT32(1);
Vector *result;
int dim = svec->dim;
float *values = SPARSEVEC_VALUES(svec);
CheckDim(dim);
CheckExpectedDim(typmod, dim);
result = InitVector(dim);
for (int i = 0; i < svec->nnz; i++) {
result->x[svec->indices[i]] = values[i];
}
PG_RETURN_POINTER(result);
}
/*
* WAL-log a range of blocks in a relation.
*
* An image of all pages with block numbers 'startblk' <= X < 'endblk' is
* written to the WAL. If the range is large, this is done in multiple WAL
* records.
*
* If all page follows the standard page layout, with a PageHeader and unused
* space between pd_lower and pd_upper, set 'page_std' to true. That allows
* the unused space to be left out from the WAL records, making them smaller.
*
* NOTE: This function acquires exclusive-locks on the pages. Typically, this
* is used on a newly-built relation, and the caller is holding a
* AccessExclusiveLock on it, so no other backend can be accessing it at the
* same time. If that's not the case, you must ensure that this does not
* cause a deadlock through some other means.
*/
void LogNewpageRange(Relation rel, ForkNumber forknum, BlockNumber startblk, BlockNumber endblk, bool page_std)
{
int flags;
BlockNumber blkno;
flags = REGBUF_FORCE_IMAGE;
if (page_std) {
flags |= REGBUF_STANDARD;
}
/*
* Iterate over all the pages in the range. They are collected into
* batches of XLR_MAX_BLOCK_ID pages, and a single WAL-record is written
* for each batch.
*/
XLogEnsureRecordSpace(XLR_MAX_BLOCK_ID - 1, 0);
blkno = startblk;
while (blkno < endblk) {
Buffer bufpack[XLR_MAX_BLOCK_ID];
XLogRecPtr recptr;
int nbufs;
int i;
CHECK_FOR_INTERRUPTS();
/* Collect a batch of blocks. */
nbufs = 0;
while (nbufs < XLR_MAX_BLOCK_ID && blkno < endblk) {
Buffer buf = ReadBufferExtended(rel, forknum, blkno, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
/*
* Completely empty pages are not WAL-logged. Writing a WAL record
* would change the LSN, and we don't want that. We want the page
* to stay empty.
*/
if (!PageIsNew(BufferGetPage(buf))) {
bufpack[nbufs++] = buf;
} else {
UnlockReleaseBuffer(buf);
}
blkno++;
}
/* Nothing more to do if all remaining blocks were empty. */
if (nbufs == 0) {
break;
}
/* Write WAL record for this batch. */
XLogBeginInsert();
START_CRIT_SECTION();
for (i = 0; i < nbufs; i++) {
MarkBufferDirty(bufpack[i]);
XLogRegisterBuffer(i, bufpack[i], flags);
}
recptr = XLogInsert(RM_XLOG_ID, XLOG_FPI);
for (i = 0; i < nbufs; i++) {
PageSetLSN(BufferGetPage(bufpack[i]), recptr);
UnlockReleaseBuffer(bufpack[i]);
}
END_CRIT_SECTION();
}
}
int PlanCreateIndexWorkers(Relation heapRelation, IndexInfo *indexInfo)
{
int parallelWorkers = RelationGetParallelWorkers(heapRelation, 0);
int maxHashbucketIndexWorker = 32;
if (parallelWorkers != 0) {
parallelWorkers = Min(maxHashbucketIndexWorker, parallelWorkers);
}
if (indexInfo->ii_Concurrent && indexInfo->ii_ParallelWorkers > 0) {
ereport(NOTICE, (errmsg("switch off parallel mode when concurrently flag is set")));
parallelWorkers = 0;
}
if (heapRelation->rd_rel->relpersistence == RELPERSISTENCE_GLOBAL_TEMP && indexInfo->ii_ParallelWorkers > 0) {
ereport(NOTICE, (errmsg("switch off parallel mode for global temp table")));
parallelWorkers = 0;
}
/* disable parallel building index for system table */
if (IsCatalogRelation(heapRelation)) {
parallelWorkers = 0;
}
return parallelWorkers;
}
Reference in New Issue View Git Blame Copy Permalink