Skip to content

Commit 5dcb708

Browse files
sivukhinsivukhin
authored
Merge pull request #1606 from tursodatabase/vector-search-accounting
account reads and writes from ANN search
2 parents a77d45d + 18dd30e commit 5dcb708

File tree

10 files changed

+265
-109
lines changed

10 files changed

+265
-109
lines changed

‎libsql-ffi/bundled/SQLite3MultipleCiphers/src/sqlite3.c‎

Lines changed: 88 additions & 36 deletions
Large diffs are not rendered by default.

‎libsql-ffi/bundled/src/sqlite3.c‎

Lines changed: 88 additions & 36 deletions
Large diffs are not rendered by default.

‎libsql-sqlite3/benchmark/benchtest.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ int main(int argc, char* argv[]) {
138138
printf(" reads : %.2f (avg.), %lld (total)\n", total_reads * 1.0 / total_queries, total_reads);
139139
}
140140
if( total_writes > 0 ){
141-
printf(" writes: %.2f (avg.), %lld (total)\n", total_writes * 1.0 / total_writes, total_reads);
141+
printf(" writes: %.2f (avg.), %lld (total)\n", total_writes * 1.0 / total_queries, total_reads);
142142
}
143143
fflush(stdout);
144144

‎libsql-sqlite3/src/vdbe.c‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -815,12 +815,12 @@ static u32 saturating_add(u32 lhs, u32 rhs) {
815815
return (u32)MIN(0xFFFFFFFF, sum);
816816
}
817817

818-
static void inc_row_read(Vdbe *p, int count) {
818+
void libsql_inc_row_read(Vdbe *p, int count) {
819819
u32 *read = &p->aLibsqlCounter[LIBSQL_STMTSTATUS_ROWS_READ - LIBSQL_STMTSTATUS_BASE];
820820
*read = saturating_add(*read, count);
821821
}
822822

823-
static void inc_row_written(Vdbe *p, int count) {
823+
void libsql_inc_row_written(Vdbe *p, int count) {
824824
u32 *write = &p->aLibsqlCounter[LIBSQL_STMTSTATUS_ROWS_WRITTEN - LIBSQL_STMTSTATUS_BASE];
825825
*write = saturating_add(*write, count);
826826
}
@@ -3740,7 +3740,7 @@ case OP_Count: { /* out2 */
37403740
nEntry = 0; /* Not needed. Only used to silence a warning. */
37413741
i64 nPages = 0;
37423742
rc = sqlite3BtreeCount(db, pCrsr, &nEntry, &nPages);
3743-
inc_row_read(p, nPages);
3743+
libsql_inc_row_read(p, nPages);
37443744
if( rc ) goto abort_due_to_error;
37453745
}
37463746
pOut = out2Prerelease(p, pOp);
@@ -4942,7 +4942,7 @@ case OP_SeekGT: { /* jump, in3, group, ncycle */
49424942
goto seek_not_found;
49434943
}
49444944
}
4945-
inc_row_read(p, 1);
4945+
libsql_inc_row_read(p, 1);
49464946
#ifdef SQLITE_TEST
49474947
sqlite3_search_count++;
49484948
#endif
@@ -5512,7 +5512,7 @@ case OP_NotExists: /* jump, in3, ncycle */
55125512
pC->deferredMoveto = 0;
55135513
VdbeBranchTaken(res!=0,2);
55145514
pC->seekResult = res;
5515-
inc_row_read(p, 1);
5515+
libsql_inc_row_read(p, 1);
55165516
if( res!=0 ){
55175517
assert( rc==SQLITE_OK );
55185518
if( pOp->p2==0 ){
@@ -5814,7 +5814,7 @@ case OP_Insert: {
58145814
#endif
58155815

58165816
assert( (pOp->p5 & OPFLAG_LASTROWID)==0 || (pOp->p5 & OPFLAG_NCHANGE)!=0 );
5817-
if (!pC->isEphemeral) inc_row_written(p, 1);
5817+
if (!pC->isEphemeral) libsql_inc_row_written(p, 1);
58185818

58195819
if( pOp->p5 & OPFLAG_NCHANGE ){
58205820
p->nChange++;
@@ -6009,7 +6009,7 @@ case OP_Delete: {
60096009

60106010
/* Invoke the update-hook if required. */
60116011
if( opflags & OPFLAG_NCHANGE ){
6012-
if (!pC->isEphemeral) inc_row_written(p, 1);
6012+
if (!pC->isEphemeral) libsql_inc_row_written(p, 1);
60136013
p->nChange++;
60146014
if( db->xUpdateCallback && ALWAYS(pTab!=0) && HasRowid(pTab) ){
60156015
db->xUpdateCallback(db->pUpdateArg, SQLITE_DELETE, zDb, pTab->zName,
@@ -6297,7 +6297,7 @@ case OP_Last: { /* jump, ncycle */
62976297
pC->deferredMoveto = 0;
62986298
pC->cacheStatus = CACHE_STALE;
62996299
if( rc ) goto abort_due_to_error;
6300-
inc_row_read(p, 1);
6300+
libsql_inc_row_read(p, 1);
63016301
if( pOp->p2>0 ){
63026302
VdbeBranchTaken(res!=0,2);
63036303
if( res ) goto jump_to_p2;
@@ -6407,7 +6407,7 @@ case OP_Rewind: { /* jump, ncycle */
64076407
}
64086408
if( rc ) goto abort_due_to_error;
64096409
pC->nullRow = (u8)res;
6410-
inc_row_read(p, 1);
6410+
libsql_inc_row_read(p, 1);
64116411
if( pOp->p2>0 ){
64126412
VdbeBranchTaken(res!=0,2);
64136413
if( res ) goto jump_to_p2;
@@ -6513,7 +6513,7 @@ case OP_Next: /* jump, ncycle */
65136513
if( rc==SQLITE_OK ){
65146514
pC->nullRow = 0;
65156515
p->aCounter[pOp->p5]++;
6516-
inc_row_read(p, 1);
6516+
libsql_inc_row_read(p, 1);
65176517
#ifdef SQLITE_TEST
65186518
sqlite3_search_count++;
65196519
#endif
@@ -6565,7 +6565,7 @@ case OP_IdxInsert: { /* in2 */
65656565
pIn2 = &aMem[pOp->p2];
65666566
assert( (pIn2->flags & MEM_Blob) || (pOp->p5 & OPFLAG_PREFORMAT) );
65676567
if( pOp->p5 & OPFLAG_NCHANGE ) p->nChange++;
6568-
if (!pC->isEphemeral) inc_row_written(p, 1);
6568+
if (!pC->isEphemeral) libsql_inc_row_written(p, 1);
65696569
#ifndef SQLITE_OMIT_VECTOR
65706570
if( isVectorCursor(pC) ) {
65716571
UnpackedRecord idxKeyStatic;
@@ -7017,7 +7017,7 @@ case OP_Clear: {
70177017
rc = sqlite3BtreeClearTable(db->aDb[pOp->p2].pBt, (u32)pOp->p1, &nChange);
70187018
if( pOp->p3 ){
70197019
p->nChange += nChange;
7020-
inc_row_written(p, nChange);
7020+
libsql_inc_row_written(p, nChange);
70217021
if( pOp->p3>0 ){
70227022
assert( memIsValid(&aMem[pOp->p3]) );
70237023
memAboutToChange(p, &aMem[pOp->p3]);
@@ -8330,6 +8330,11 @@ case OP_VOpen: { /* ncycle */
83308330
if( pCur ){
83318331
pCur->uc.pVCur = pVCur;
83328332
pVtab->nRef++;
8333+
#ifndef SQLITE_OMIT_VECTOR
8334+
if( sqlite3StrICmp(pOp->p4.pVtab->pMod->zName, VECTOR_INDEX_VTAB_NAME) == 0 ){
8335+
pCur->isTracked = 1;
8336+
}
8337+
#endif
83338338
}else{
83348339
assert( db->mallocFailed );
83358340
pModule->xClose(pVCur);
@@ -8604,7 +8609,7 @@ case OP_VNext: { /* jump, ncycle */
86048609
rc = pModule->xNext(pCur->uc.pVCur);
86058610
sqlite3VtabImportErrmsg(p, pVtab);
86068611
if( rc ) goto abort_due_to_error;
8607-
inc_row_read(p, 1);
8612+
libsql_inc_row_read(p, 1);
86088613
res = pModule->xEof(pCur->uc.pVCur);
86098614
VdbeBranchTaken(!res,2);
86108615
if( !res ){
@@ -8726,7 +8731,7 @@ case OP_VUpdate: {
87268731
p->errorAction = ((pOp->p5==OE_Replace) ? OE_Abort : pOp->p5);
87278732
}
87288733
}else{
8729-
inc_row_written(p, 1);
8734+
libsql_inc_row_written(p, 1);
87308735
p->nChange++;
87318736
}
87328737
if( rc ) goto abort_due_to_error;

‎libsql-sqlite3/src/vdbeInt.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ struct VdbeCursor {
9898
Bool isEphemeral:1; /* True for an ephemeral table */
9999
Bool useRandomRowid:1; /* Generate new record numbers semi-randomly */
100100
Bool isOrdered:1; /* True if the table is not BTREE_UNORDERED */
101+
Bool isTracked:1; /* True if cursor created for virtual table which track reads/writes */
101102
Bool noReuse:1; /* OpenEphemeral may not reuse this cursor */
102103
Bool colCache:1; /* pCache pointer is initialized and non-NULL */
103104
u16 seekHit; /* See the OP_SeekHit and OP_IfNoHope opcodes */
@@ -529,6 +530,8 @@ struct Vdbe {
529530
#endif
530531
};
531532

533+
void libsql_inc_row_read(Vdbe *p, int count);
534+
void libsql_inc_row_written(Vdbe *p, int count);
532535
/*
533536
** The following are allowed values for Vdbe.eVdbeState
534537
*/

‎libsql-sqlite3/src/vdbeaux.c‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2713,6 +2713,13 @@ void sqlite3VdbeMakeReady(
27132713
void sqlite3VdbeFreeCursor(Vdbe *p, VdbeCursor *pCx){
27142714
if( pCx ) sqlite3VdbeFreeCursorNN(p,pCx);
27152715
}
2716+
2717+
struct sqlite3_vtab_cursor_tracked {
2718+
sqlite3_vtab_cursor base; /* Base class - must be first */
2719+
int nReads; /* Number of row read from the storage backing virtual table */
2720+
int nWrites; /* Number of row written to the storage backing virtual table */
2721+
};
2722+
27162723
static SQLITE_NOINLINE void freeCursorWithCache(Vdbe *p, VdbeCursor *pCx){
27172724
VdbeTxtBlbCache *pCache = pCx->pCache;
27182725
assert( pCx->colCache );
@@ -2742,17 +2749,26 @@ void sqlite3VdbeFreeCursorNN(Vdbe *p, VdbeCursor *pCx){
27422749
}
27432750
#ifndef SQLITE_OMIT_VIRTUALTABLE
27442751
case CURTYPE_VTAB: {
2752+
struct sqlite3_vtab_cursor_tracked *pTracked;
27452753
sqlite3_vtab_cursor *pVCur = pCx->uc.pVCur;
27462754
const sqlite3_module *pModule = pVCur->pVtab->pModule;
27472755
assert( pVCur->pVtab->nRef>0 );
2756+
if( pCx->isTracked ){
2757+
pTracked = (struct sqlite3_vtab_cursor_tracked*)pVCur;
2758+
libsql_inc_row_read(p, pTracked->nReads);
2759+
libsql_inc_row_written(p, pTracked->nWrites);
2760+
}
27482761
pVCur->pVtab->nRef--;
27492762
pModule->xClose(pVCur);
27502763
break;
27512764
}
27522765
#endif
27532766
#ifndef SQLITE_OMIT_VECTOR
27542767
case CURTYPE_VECTOR_IDX: {
2755-
vectorIndexCursorClose(p->db, pCx->uc.pVecIdx);
2768+
int nReads, nWrites;
2769+
vectorIndexCursorClose(p->db, pCx->uc.pVecIdx, &nReads, &nWrites);
2770+
libsql_inc_row_read(p, nReads);
2771+
libsql_inc_row_written(p, nWrites);
27562772
break;
27572773
}
27582774
#endif

‎libsql-sqlite3/src/vectorIndex.c‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -950,7 +950,16 @@ int vectorIndexCreate(Parse *pParse, const Index *pIdx, const char *zDbSName, co
950950
return CREATE_OK;
951951
}
952952

953-
int vectorIndexSearch(sqlite3 *db, const char* zDbSName, int argc, sqlite3_value **argv, VectorOutRows *pRows, char **pzErrMsg) {
953+
int vectorIndexSearch(
954+
sqlite3 *db,
955+
const char* zDbSName,
956+
int argc,
957+
sqlite3_value **argv,
958+
VectorOutRows *pRows,
959+
int *nReads,
960+
int *nWrites,
961+
char **pzErrMsg
962+
) {
954963
int type, dims, k, rc;
955964
const char *zIdxName;
956965
const char *zErrMsg;
@@ -1028,6 +1037,8 @@ int vectorIndexSearch(sqlite3 *db, const char* zDbSName, int argc, sqlite3_value
10281037
rc = diskAnnSearch(pDiskAnn, pVector, k, &pKey, pRows, pzErrMsg);
10291038
out:
10301039
if( pDiskAnn != NULL ){
1040+
*nReads += pDiskAnn->nReads;
1041+
*nWrites += pDiskAnn->nWrites;
10311042
diskAnnCloseIndex(pDiskAnn);
10321043
}
10331044
if( pVector != NULL ){
@@ -1107,7 +1118,10 @@ int vectorIndexCursorInit(
11071118
return SQLITE_OK;
11081119
}
11091120

1110-
void vectorIndexCursorClose(sqlite3 *db, VectorIdxCursor *pCursor){
1121+
void vectorIndexCursorClose(sqlite3 *db, VectorIdxCursor *pCursor, int *nReads, int *nWrites){
1122+
*nReads = pCursor->pIndex->nReads;
1123+
*nWrites = pCursor->pIndex->nWrites;
1124+
11111125
diskAnnCloseIndex(pCursor->pIndex);
11121126
sqlite3DbFree(db, pCursor);
11131127
}

‎libsql-sqlite3/src/vectorIndexInt.h‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@ struct DiskAnnIndex {
3030
float pruningAlpha; /* Alpha parameter for edge pruning during INSERT operation */
3131
int insertL; /* Max size of candidate set (L) visited during INSERT operation */
3232
int searchL; /* Max size of candidate set (L) visited during SEARCH operation (can be overriden from query in future) */
33+
34+
int nReads;
35+
int nWrites;
3336
};
3437

3538
/*
@@ -55,8 +58,8 @@ struct BlobSpot {
5558

5659
/* BlobSpot operations */
5760
int blobSpotCreate(const DiskAnnIndex *pIndex, BlobSpot **ppBlobSpot, u64 nRowid, int nBufferSize, int isWritable);
58-
int blobSpotReload(const DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, u64 nRowid, int nBufferSize);
59-
int blobSpotFlush(BlobSpot *pBlobSpot);
61+
int blobSpotReload(DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, u64 nRowid, int nBufferSize);
62+
int blobSpotFlush(DiskAnnIndex *pIndex, BlobSpot *pBlobSpot);
6063
void blobSpotFree(BlobSpot *pBlobSpot);
6164

6265
/*
@@ -220,9 +223,9 @@ int diskAnnClearIndex(sqlite3 *, const char *, const char *);
220223
int diskAnnDropIndex(sqlite3 *, const char *, const char *);
221224
int diskAnnOpenIndex(sqlite3 *, const char *, const char *, const VectorIdxParams *, DiskAnnIndex **);
222225
void diskAnnCloseIndex(DiskAnnIndex *);
223-
int diskAnnInsert(const DiskAnnIndex *, const VectorInRow *, char **);
224-
int diskAnnDelete(const DiskAnnIndex *, const VectorInRow *, char **);
225-
int diskAnnSearch(const DiskAnnIndex *, const Vector *, int, const VectorIdxKey *, VectorOutRows *, char **);
226+
int diskAnnInsert(DiskAnnIndex *, const VectorInRow *, char **);
227+
int diskAnnDelete(DiskAnnIndex *, const VectorInRow *, char **);
228+
int diskAnnSearch(DiskAnnIndex *, const Vector *, int, const VectorIdxKey *, VectorOutRows *, char **);
226229

227230
typedef struct VectorIdxCursor VectorIdxCursor;
228231

@@ -235,9 +238,9 @@ int vectorIdxParseColumnType(const char *, int *, int *, const char **);
235238
int vectorIndexCreate(Parse*, const Index*, const char *, const IdList*);
236239
int vectorIndexClear(sqlite3 *, const char *, const char *);
237240
int vectorIndexDrop(sqlite3 *, const char *, const char *);
238-
int vectorIndexSearch(sqlite3 *, const char *, int, sqlite3_value **, VectorOutRows *, char **);
241+
int vectorIndexSearch(sqlite3 *, const char *, int, sqlite3_value **, VectorOutRows *, int *, int *, char **);
239242
int vectorIndexCursorInit(sqlite3 *, const char *, const char *, VectorIdxCursor **);
240-
void vectorIndexCursorClose(sqlite3 *, VectorIdxCursor *);
243+
void vectorIndexCursorClose(sqlite3 *, VectorIdxCursor *, int *, int *);
241244
int vectorIndexInsert(VectorIdxCursor *, const UnpackedRecord *, char **);
242245
int vectorIndexDelete(VectorIdxCursor *, const UnpackedRecord *, char **);
243246

‎libsql-sqlite3/src/vectordiskann.c‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ int blobSpotCreate(const DiskAnnIndex *pIndex, BlobSpot **ppBlobSpot, u64 nRowid
200200
return rc;
201201
}
202202

203-
int blobSpotReload(const DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, u64 nRowid, int nBufferSize) {
203+
int blobSpotReload(DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, u64 nRowid, int nBufferSize) {
204204
int rc;
205205

206206
DiskAnnTrace(("blob spot reload: rowid=%lld\n", nRowid));
@@ -242,6 +242,7 @@ int blobSpotReload(const DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, u64 nRowid,
242242
if( rc != SQLITE_OK ){
243243
goto abort;
244244
}
245+
pIndex->nReads++;
245246
pBlobSpot->isInitialized = 1;
246247
return SQLITE_OK;
247248

@@ -251,8 +252,13 @@ int blobSpotReload(const DiskAnnIndex *pIndex, BlobSpot *pBlobSpot, u64 nRowid,
251252
return rc;
252253
}
253254

254-
int blobSpotFlush(BlobSpot *pBlobSpot) {
255-
return sqlite3_blob_write(pBlobSpot->pBlob, pBlobSpot->pBuffer, pBlobSpot->nBufferSize, 0);
255+
int blobSpotFlush(DiskAnnIndex* pIndex, BlobSpot *pBlobSpot) {
256+
int rc = sqlite3_blob_write(pBlobSpot->pBlob, pBlobSpot->pBuffer, pBlobSpot->nBufferSize, 0);
257+
if( rc != SQLITE_OK ){
258+
return rc;
259+
}
260+
pIndex->nWrites++;
261+
return rc;
256262
}
257263

258264
void blobSpotFree(BlobSpot *pBlobSpot) {
@@ -1085,7 +1091,7 @@ static void diskAnnPruneEdges(const DiskAnnIndex *pIndex, BlobSpot *pNodeBlob, i
10851091
}
10861092

10871093
// main search routine - called from both SEARCH and INSERT operation
1088-
static int diskAnnSearchInternal(const DiskAnnIndex *pIndex, DiskAnnSearchCtx *pCtx, u64 nStartRowid, char **pzErrMsg){
1094+
static int diskAnnSearchInternal(DiskAnnIndex *pIndex, DiskAnnSearchCtx *pCtx, u64 nStartRowid, char **pzErrMsg){
10891095
DiskAnnTrace(("diskAnnSearchInternal: ready to search: rootId=%lld\n", nStartRowid));
10901096
DiskAnnNode *start = NULL;
10911097
// in case of SEARCH operation (blobMode == DISKANN_BLOB_READONLY) we don't need to preserve all node blobs in the memory
@@ -1218,7 +1224,7 @@ static int diskAnnSearchInternal(const DiskAnnIndex *pIndex, DiskAnnSearchCtx *p
12181224

12191225
// search k nearest neighbours for pVector in the pIndex (with pKey primary key structure) and put result in the pRows output
12201226
int diskAnnSearch(
1221-
const DiskAnnIndex *pIndex,
1227+
DiskAnnIndex *pIndex,
12221228
const Vector *pVector,
12231229
int k,
12241230
const VectorIdxKey *pKey,
@@ -1290,7 +1296,7 @@ int diskAnnSearch(
12901296

12911297
// insert pVectorInRow in the pIndex
12921298
int diskAnnInsert(
1293-
const DiskAnnIndex *pIndex,
1299+
DiskAnnIndex *pIndex,
12941300
const VectorInRow *pVectorInRow,
12951301
char **pzErrMsg
12961302
){
@@ -1377,7 +1383,7 @@ int diskAnnInsert(
13771383
nodeBinReplaceEdge(pIndex, pVisited->pBlobSpot, iReplace, nNewRowid, pVectorInRow->pVector);
13781384
diskAnnPruneEdges(pIndex, pVisited->pBlobSpot, iReplace);
13791385

1380-
rc = blobSpotFlush(pVisited->pBlobSpot);
1386+
rc = blobSpotFlush(pIndex, pVisited->pBlobSpot);
13811387
if( rc != SQLITE_OK ){
13821388
*pzErrMsg = sqlite3_mprintf("vector index(insert): failed to flush blob");
13831389
goto out;
@@ -1387,7 +1393,7 @@ int diskAnnInsert(
13871393
rc = SQLITE_OK;
13881394
out:
13891395
if( rc == SQLITE_OK ){
1390-
rc = blobSpotFlush(pBlobSpot);
1396+
rc = blobSpotFlush(pIndex, pBlobSpot);
13911397
if( rc != SQLITE_OK ){
13921398
*pzErrMsg = sqlite3_mprintf("vector index(insert): failed to flush blob");
13931399
}
@@ -1401,7 +1407,7 @@ int diskAnnInsert(
14011407

14021408
// delete pInRow from pIndex
14031409
int diskAnnDelete(
1404-
const DiskAnnIndex *pIndex,
1410+
DiskAnnIndex *pIndex,
14051411
const VectorInRow *pInRow,
14061412
char **pzErrMsg
14071413
){
@@ -1450,7 +1456,7 @@ int diskAnnDelete(
14501456
continue;
14511457
}
14521458
nodeBinDeleteEdge(pIndex, pEdgeBlob, iDelete);
1453-
rc = blobSpotFlush(pEdgeBlob);
1459+
rc = blobSpotFlush(pIndex, pEdgeBlob);
14541460
if( rc != SQLITE_OK ){
14551461
*pzErrMsg = sqlite3_mprintf("vector index(delete): failed to flush blob for edge row");
14561462
goto out;
@@ -1510,6 +1516,8 @@ int diskAnnOpenIndex(
15101516
pIndex->pruningAlpha = vectorIdxParamsGetF64(pParams, VECTOR_PRUNING_ALPHA_PARAM_ID);
15111517
pIndex->insertL = vectorIdxParamsGetU64(pParams, VECTOR_INSERT_L_PARAM_ID);
15121518
pIndex->searchL = vectorIdxParamsGetU64(pParams, VECTOR_SEARCH_L_PARAM_ID);
1519+
pIndex->nReads = 0;
1520+
pIndex->nWrites = 0;
15131521
if( pIndex->nDistanceFunc == 0 ||
15141522
pIndex->nBlockSize == 0 ||
15151523
pIndex->nNodeVectorType == 0 ||

0 commit comments

Comments
 (0)