Skip to content

Commit

Permalink
Merge pull request #1810 from tursodatabase/vector-search-avoid-unali…
Browse files Browse the repository at this point in the history 
…gned-reads

vector-search: avoid unaligned reads
  • Loading branch information
@sivukhin
sivukhin authored Nov 24, 2024
2 parents e673973 + 27b3fc4 commit a77f422
Show file tree
Hide file tree
Showing 8 changed files with 188 additions and 104 deletions.
79 changes: 46 additions & 33 deletions libsql-ffi/bundled/SQLite3MultipleCiphers/src/sqlite3.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85673,8 +85673,9 @@ void blobSpotFree(BlobSpot *pBlobSpot);

/*
* Accessor for node binary format
* - v1 format is the following:
* [u64 nRowid] [u16 nEdges] [node vector] [edge vector] * nEdges [trash vector] * (nMaxEdges - nEdges) ([u64 legacyField] [u64 edgeId]) * nEdges
* - default format is the following:
* [u64 nRowid] [u16 nEdges] [6 byte padding] [node vector] [edge vector] * nEdges [trash vector] * (nMaxEdges - nEdges) ([u32 unused] [f32 distance] [u64 edgeId]) * nEdges
* Note, that 6 byte padding after nEdges required to align [node vector] by word boundary and avoid unaligned reads
* Note, that node vector and edge vector can have different representations (and edge vector can be smaller in size than node vector)
*/
int nodeEdgesMaxCount(const DiskAnnIndex *pIndex);
Expand Down Expand Up @@ -85713,9 +85714,11 @@ typedef u8 MetricType;
/*
* 1 - v1 version; node block format: [node meta] [node vector] [edge vectors] ... [ [u64 unused ] [u64 edge rowid] ] ...
* 2 - v2 version; node block format: [node meta] [node vector] [edge vectors] ... [ [u32 unused] [f32 distance] [u64 edge rowid] ] ...
* 3 - v3 version; node meta aligned to 8-byte boundary (instead of having u64 + u16 size - we round it up to u64 + u64)
*/
#define VECTOR_FORMAT_V1 1
#define VECTOR_FORMAT_DEFAULT 2
#define VECTOR_FORMAT_V2 2
#define VECTOR_FORMAT_DEFAULT 3

/* type of the vector index */
#define VECTOR_INDEX_TYPE_PARAM_ID 2
Expand Down Expand Up @@ -212727,8 +212730,6 @@ SQLITE_PRIVATE void sqlite3RegisterVectorFunctions(void){
*/
#define DISKANN_BLOCK_SIZE_SHIFT 9

#define VECTOR_NODE_METADATA_SIZE (sizeof(u64) + sizeof(u16))
#define VECTOR_EDGE_METADATA_SIZE (sizeof(u64) + sizeof(u64))

typedef struct VectorPair VectorPair;
typedef struct DiskAnnSearchCtx DiskAnnSearchCtx;
Expand Down Expand Up @@ -212951,46 +212952,58 @@ void blobSpotFree(BlobSpot *pBlobSpot) {
** Layout specific utilities
**************************************************************************/

int nodeEdgeOverhead(int nEdgeVectorSize){
return nEdgeVectorSize + VECTOR_EDGE_METADATA_SIZE;
int nodeMetadataSize(int nFormatVersion){
if( nFormatVersion <= VECTOR_FORMAT_V2 ){
return (sizeof(u64) + sizeof(u16));
}else{
return (sizeof(u64) + sizeof(u64));
}
}

int edgeMetadataSize(int nFormatVersion){
return (sizeof(u64) + sizeof(u64));
}

int nodeEdgeOverhead(int nFormatVersion, int nEdgeVectorSize){
return nEdgeVectorSize + edgeMetadataSize(nFormatVersion);
}

int nodeOverhead(int nNodeVectorSize){
return nNodeVectorSize + VECTOR_NODE_METADATA_SIZE;
int nodeOverhead(int nFormatVersion, int nNodeVectorSize){
return nNodeVectorSize + nodeMetadataSize(nFormatVersion);
}

int nodeEdgesMaxCount(const DiskAnnIndex *pIndex){
unsigned int nMaxEdges = (pIndex->nBlockSize - nodeOverhead(pIndex->nNodeVectorSize)) / nodeEdgeOverhead(pIndex->nEdgeVectorSize);
unsigned int nMaxEdges = (pIndex->nBlockSize - nodeOverhead(pIndex->nFormatVersion, pIndex->nNodeVectorSize)) / nodeEdgeOverhead(pIndex->nFormatVersion, pIndex->nEdgeVectorSize);
assert( nMaxEdges > 0);
return nMaxEdges;
}

int nodeEdgesMetadataOffset(const DiskAnnIndex *pIndex){
unsigned int offset;
unsigned int nMaxEdges = nodeEdgesMaxCount(pIndex);
offset = VECTOR_NODE_METADATA_SIZE + pIndex->nNodeVectorSize + nMaxEdges * pIndex->nEdgeVectorSize;
offset = nodeMetadataSize(pIndex->nFormatVersion) + pIndex->nNodeVectorSize + nMaxEdges * pIndex->nEdgeVectorSize;
assert( offset <= pIndex->nBlockSize );
return offset;
}

void nodeBinInit(const DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, u64 nRowid, Vector *pVector){
assert( VECTOR_NODE_METADATA_SIZE + pIndex->nNodeVectorSize <= pBlobSpot->nBufferSize );
assert( nodeMetadataSize(pIndex->nFormatVersion) + pIndex->nNodeVectorSize <= pBlobSpot->nBufferSize );

memset(pBlobSpot->pBuffer, 0, pBlobSpot->nBufferSize);
writeLE64(pBlobSpot->pBuffer, nRowid);
// neighbours count already zero after memset - no need to set it explicitly

vectorSerializeToBlob(pVector, pBlobSpot->pBuffer + VECTOR_NODE_METADATA_SIZE, pIndex->nNodeVectorSize);
vectorSerializeToBlob(pVector, pBlobSpot->pBuffer + nodeMetadataSize(pIndex->nFormatVersion), pIndex->nNodeVectorSize);
}

void nodeBinVector(const DiskAnnIndex *pIndex, const BlobSpot *pBlobSpot, Vector *pVector) {
assert( VECTOR_NODE_METADATA_SIZE + pIndex->nNodeVectorSize <= pBlobSpot->nBufferSize );
assert( nodeMetadataSize(pIndex->nFormatVersion) + pIndex->nNodeVectorSize <= pBlobSpot->nBufferSize );

vectorInitStatic(pVector, pIndex->nNodeVectorType, pIndex->nVectorDims, pBlobSpot->pBuffer + VECTOR_NODE_METADATA_SIZE);
vectorInitStatic(pVector, pIndex->nNodeVectorType, pIndex->nVectorDims, pBlobSpot->pBuffer + nodeMetadataSize(pIndex->nFormatVersion));
}

u16 nodeBinEdges(const DiskAnnIndex *pIndex, const BlobSpot *pBlobSpot) {
assert( VECTOR_NODE_METADATA_SIZE <= pBlobSpot->nBufferSize );
assert( nodeMetadataSize(pIndex->nFormatVersion) <= pBlobSpot->nBufferSize );

return readLE16(pBlobSpot->pBuffer + sizeof(u64));
}
Expand All @@ -213000,20 +213013,20 @@ void nodeBinEdge(const DiskAnnIndex *pIndex, const BlobSpot *pBlobSpot, int iEdg
int offset = nodeEdgesMetadataOffset(pIndex);

if( pRowid != NULL ){
assert( offset + (iEdge + 1) * VECTOR_EDGE_METADATA_SIZE <= pBlobSpot->nBufferSize );
*pRowid = readLE64(pBlobSpot->pBuffer + offset + iEdge * VECTOR_EDGE_METADATA_SIZE + sizeof(u64));
assert( offset + (iEdge + 1) * edgeMetadataSize(pIndex->nFormatVersion) <= pBlobSpot->nBufferSize );
*pRowid = readLE64(pBlobSpot->pBuffer + offset + iEdge * edgeMetadataSize(pIndex->nFormatVersion) + sizeof(u64));
}
if( pIndex->nFormatVersion != VECTOR_FORMAT_V1 && pDistance != NULL ){
distance = readLE32(pBlobSpot->pBuffer + offset + iEdge * VECTOR_EDGE_METADATA_SIZE + sizeof(u32));
distance = readLE32(pBlobSpot->pBuffer + offset + iEdge * edgeMetadataSize(pIndex->nFormatVersion) + sizeof(u32));
*pDistance = *((float*)&distance);
}
if( pVector != NULL ){
assert( VECTOR_NODE_METADATA_SIZE + pIndex->nNodeVectorSize + iEdge * pIndex->nEdgeVectorSize < offset );
assert( nodeMetadataSize(pIndex->nFormatVersion) + pIndex->nNodeVectorSize + iEdge * pIndex->nEdgeVectorSize < offset );
vectorInitStatic(
pVector,
pIndex->nEdgeVectorType,
pIndex->nVectorDims,
pBlobSpot->pBuffer + VECTOR_NODE_METADATA_SIZE + pIndex->nNodeVectorSize + iEdge * pIndex->nEdgeVectorSize
pBlobSpot->pBuffer + nodeMetadataSize(pIndex->nFormatVersion) + pIndex->nNodeVectorSize + iEdge * pIndex->nEdgeVectorSize
);
}
}
Expand Down Expand Up @@ -213050,11 +213063,11 @@ void nodeBinReplaceEdge(const DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, int iRe
nEdges++;
}

edgeVectorOffset = VECTOR_NODE_METADATA_SIZE + pIndex->nNodeVectorSize + iReplace * pIndex->nEdgeVectorSize;
edgeMetaOffset = nodeEdgesMetadataOffset(pIndex) + iReplace * VECTOR_EDGE_METADATA_SIZE;
edgeVectorOffset = nodeMetadataSize(pIndex->nFormatVersion) + pIndex->nNodeVectorSize + iReplace * pIndex->nEdgeVectorSize;
edgeMetaOffset = nodeEdgesMetadataOffset(pIndex) + iReplace * edgeMetadataSize(pIndex->nFormatVersion);

assert( edgeVectorOffset + pIndex->nEdgeVectorSize <= pBlobSpot->nBufferSize );
assert( edgeMetaOffset + VECTOR_EDGE_METADATA_SIZE <= pBlobSpot->nBufferSize );
assert( edgeMetaOffset + edgeMetadataSize(pIndex->nFormatVersion) <= pBlobSpot->nBufferSize );

vectorSerializeToBlob(pVector, pBlobSpot->pBuffer + edgeVectorOffset, pIndex->nEdgeVectorSize);
writeLE32(pBlobSpot->pBuffer + edgeMetaOffset + sizeof(u32), *((u32*)&distance));
Expand All @@ -213070,19 +213083,19 @@ void nodeBinDeleteEdge(const DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, int iDel

assert( 0 <= iDelete && iDelete < nEdges );

edgeVectorOffset = VECTOR_NODE_METADATA_SIZE + pIndex->nNodeVectorSize + iDelete * pIndex->nEdgeVectorSize;
lastVectorOffset = VECTOR_NODE_METADATA_SIZE + pIndex->nNodeVectorSize + (nEdges - 1) * pIndex->nEdgeVectorSize;
edgeMetaOffset = nodeEdgesMetadataOffset(pIndex) + iDelete * VECTOR_EDGE_METADATA_SIZE;
lastMetaOffset = nodeEdgesMetadataOffset(pIndex) + (nEdges - 1) * VECTOR_EDGE_METADATA_SIZE;
edgeVectorOffset = nodeMetadataSize(pIndex->nFormatVersion) + pIndex->nNodeVectorSize + iDelete * pIndex->nEdgeVectorSize;
lastVectorOffset = nodeMetadataSize(pIndex->nFormatVersion) + pIndex->nNodeVectorSize + (nEdges - 1) * pIndex->nEdgeVectorSize;
edgeMetaOffset = nodeEdgesMetadataOffset(pIndex) + iDelete * edgeMetadataSize(pIndex->nFormatVersion);
lastMetaOffset = nodeEdgesMetadataOffset(pIndex) + (nEdges - 1) * edgeMetadataSize(pIndex->nFormatVersion);

assert( edgeVectorOffset + pIndex->nEdgeVectorSize <= pBlobSpot->nBufferSize );
assert( lastVectorOffset + pIndex->nEdgeVectorSize <= pBlobSpot->nBufferSize );
assert( edgeMetaOffset + VECTOR_EDGE_METADATA_SIZE <= pBlobSpot->nBufferSize );
assert( lastMetaOffset + VECTOR_EDGE_METADATA_SIZE <= pBlobSpot->nBufferSize );
assert( edgeMetaOffset + edgeMetadataSize(pIndex->nFormatVersion) <= pBlobSpot->nBufferSize );
assert( lastMetaOffset + edgeMetadataSize(pIndex->nFormatVersion) <= pBlobSpot->nBufferSize );

if( edgeVectorOffset < lastVectorOffset ){
memmove(pBlobSpot->pBuffer + edgeVectorOffset, pBlobSpot->pBuffer + lastVectorOffset, pIndex->nEdgeVectorSize);
memmove(pBlobSpot->pBuffer + edgeMetaOffset, pBlobSpot->pBuffer + lastMetaOffset, VECTOR_EDGE_METADATA_SIZE);
memmove(pBlobSpot->pBuffer + edgeMetaOffset, pBlobSpot->pBuffer + lastMetaOffset, edgeMetadataSize(pIndex->nFormatVersion));
}

writeLE16(pBlobSpot->pBuffer + sizeof(u64), nEdges - 1);
Expand Down Expand Up @@ -213168,9 +213181,9 @@ int diskAnnCreateIndex(
if( maxNeighborsParam == 0 ){
// 3 D**(1/2) gives good recall values (90%+)
// we also want to keep disk overhead at moderate level - 50x of the disk size increase is the current upper bound
maxNeighborsParam = MIN(3 * ((int)(sqrt(dims)) + 1), (50 * nodeOverhead(vectorDataSize(type, dims))) / nodeEdgeOverhead(vectorDataSize(neighbours, dims)) + 1);
maxNeighborsParam = MIN(3 * ((int)(sqrt(dims)) + 1), (50 * nodeOverhead(VECTOR_FORMAT_DEFAULT, vectorDataSize(type, dims))) / nodeEdgeOverhead(VECTOR_FORMAT_DEFAULT, vectorDataSize(neighbours, dims)) + 1);
}
blockSizeBytes = nodeOverhead(vectorDataSize(type, dims)) + maxNeighborsParam * (u64)nodeEdgeOverhead(vectorDataSize(neighbours, dims));
blockSizeBytes = nodeOverhead(VECTOR_FORMAT_DEFAULT, vectorDataSize(type, dims)) + maxNeighborsParam * (u64)nodeEdgeOverhead(VECTOR_FORMAT_DEFAULT, vectorDataSize(neighbours, dims));
if( blockSizeBytes > DISKANN_MAX_BLOCK_SZ ){
return SQLITE_ERROR;
}
Expand Down
16 changes: 12 additions & 4 deletions libsql-ffi/bundled/bindings/bindgen.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -940,7 +940,7 @@ extern "C" {
extern "C" {
pub fn sqlite3_vmprintf(
arg1: *const ::std::os::raw::c_char,
arg2: va_list,
arg2: *mut __va_list_tag,
) -> *mut ::std::os::raw::c_char;
}
extern "C" {
Expand All @@ -956,7 +956,7 @@ extern "C" {
arg1: ::std::os::raw::c_int,
arg2: *mut ::std::os::raw::c_char,
arg3: *const ::std::os::raw::c_char,
arg4: va_list,
arg4: *mut __va_list_tag,
) -> *mut ::std::os::raw::c_char;
}
extern "C" {
Expand Down Expand Up @@ -2503,7 +2503,7 @@ extern "C" {
pub fn sqlite3_str_vappendf(
arg1: *mut sqlite3_str,
zFormat: *const ::std::os::raw::c_char,
arg2: va_list,
arg2: *mut __va_list_tag,
);
}
extern "C" {
Expand Down Expand Up @@ -3570,4 +3570,12 @@ extern "C" {
extern "C" {
pub static sqlite3_wal_manager: libsql_wal_manager;
}
pub type __builtin_va_list = *mut ::std::os::raw::c_char;
pub type __builtin_va_list = [__va_list_tag; 1usize];
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct __va_list_tag {
pub gp_offset: ::std::os::raw::c_uint,
pub fp_offset: ::std::os::raw::c_uint,
pub overflow_arg_area: *mut ::std::os::raw::c_void,
pub reg_save_area: *mut ::std::os::raw::c_void,
}
Loading

0 comments on commit a77f422

Please sign in to comment.