1bit quantized embeddings search: somehow working version

Author: sivukhin

libsql-sqlite3/src/vector.c

@@ -74,10 +74,11 @@ Vector *vectorAlloc(VectorType type, VectorDims dims){
 ** Note that the vector object points to the blob so if
 ** you free the blob, the vector becomes invalid.
 **/
-void vectorInitStatic(Vector *pVector, VectorType type, const unsigned char *pBlob, size_t nBlobSize){
-  pVector->type = type;
+void vectorInitStatic(Vector *pVector, VectorType type, VectorDims dims, void *pBlob){
   pVector->flags = VECTOR_FLAGS_STATIC;
-  vectorInitFromBlob(pVector, pBlob, nBlobSize);
+  pVector->type = type;
+  pVector->dims = dims;
+  pVector->data = pBlob;
 }
 
 /*
@@ -479,6 +480,31 @@ void vectorInitFromBlob(Vector *pVector, const unsigned char *pBlob, size_t nBlo
   }
 }
 
+void vectorConvert(const Vector *pFrom, Vector *pTo){
+  int i;
+  u8 *bitData;
+  float *floatData;
+
+  assert( pFrom->dims == pTo->dims );
+
+  if( pFrom->type == VECTOR_TYPE_FLOAT32 && pTo->type == VECTOR_TYPE_1BIT ){
+    floatData = pFrom->data;
+    bitData = pTo->data;
+    for(i = 0; i < pFrom->dims; i += 8){
+      bitData[i / 8] = 0;
+    }
+    for(i = 0; i < pFrom->dims; i++){
+      if( floatData[i] < 0 ){
+        bitData[i / 8] &= ~(1 << (i & 7));
+      }else{
+        bitData[i / 8] |= (1 << (i & 7));
+      }
+    }
+  }else{
+    assert(0);
+  }
+}
+
 /**************************************************************************
 ** SQL function implementations
 ****************************************************************************/

libsql-sqlite3/src/vector1bit.c

@@ -56,17 +56,16 @@ size_t vector1BitSerializeToBlob(
   unsigned char *pBlob,
   size_t nBlobSize
 ){
-  float *elems = pVector->data;
-  unsigned char *pPtr = pBlob;
-  size_t len = 0;
+  u8 *elems = pVector->data;
+  u8 *pPtr = pBlob;
   unsigned i;
 
   assert( pVector->type == VECTOR_TYPE_1BIT );
   assert( pVector->dims <= MAX_VECTOR_SZ );
   assert( nBlobSize >= (pVector->dims + 7) / 8 );
 
-  for(i = 0; i < pVector->dims; i++){
-    elems[i] = pPtr[i];
+  for(i = 0; i < (pVector->dims + 7) / 8; i++){
+    pPtr[i] = elems[i];
   }
   return (pVector->dims + 7) / 8;
 }
@@ -92,7 +91,7 @@ static int BitsCount[256] = {
 };
 
 int vector1BitDistanceHamming(const Vector *v1, const Vector *v2){
-  int sum = 0;
+  int diff = 0;
   u8 *e1 = v1->data;
   u8 *e2 = v2->data;
   int i;
@@ -101,10 +100,10 @@ int vector1BitDistanceHamming(const Vector *v1, const Vector *v2){
   assert( v1->type == VECTOR_TYPE_1BIT );
   assert( v2->type == VECTOR_TYPE_1BIT );
 
-  for(i = 0; i < v1->dims; i++){
-    sum += BitsCount[e1[i]&e2[i]];
+  for(i = 0; i < v1->dims; i += 8){
+    diff += BitsCount[e1[i/8] ^ e2[i/8]];
   }
-  return sum;
+  return diff;
 }
 
 #endif /* !defined(SQLITE_OMIT_VECTOR) */

libsql-sqlite3/src/vectorIndex.c

@@ -396,10 +396,10 @@ struct VectorParamName {
 };
 
 static struct VectorParamName VECTOR_PARAM_NAMES[] = {
-  { "type",               VECTOR_INDEX_TYPE_PARAM_ID,    0, "diskann", VECTOR_INDEX_TYPE_DISKANN },
-  { "metric",             VECTOR_METRIC_TYPE_PARAM_ID,   0, "cosine",  VECTOR_METRIC_TYPE_COS },
-  { "metric",             VECTOR_METRIC_TYPE_PARAM_ID,   0, "l2",      VECTOR_METRIC_TYPE_L2 },
-  { "compress_neighbors", VECTOR_METRIC_TYPE_PARAM_ID,   0, "1bit",    VECTOR_TYPE_1BIT },
+  { "type",               VECTOR_INDEX_TYPE_PARAM_ID,         0, "diskann", VECTOR_INDEX_TYPE_DISKANN },
+  { "metric",             VECTOR_METRIC_TYPE_PARAM_ID,        0, "cosine",  VECTOR_METRIC_TYPE_COS },
+  { "metric",             VECTOR_METRIC_TYPE_PARAM_ID,        0, "l2",      VECTOR_METRIC_TYPE_L2 },
+  { "compress_neighbors", VECTOR_COMPRESS_NEIGHBORS_PARAM_ID, 0, "1bit",    VECTOR_TYPE_1BIT },
   { "alpha",              VECTOR_PRUNING_ALPHA_PARAM_ID, 2, 0, 0 },
   { "search_l",           VECTOR_SEARCH_L_PARAM_ID,      1, 0, 0 },
   { "insert_l",           VECTOR_INSERT_L_PARAM_ID,      1, 0, 0 },

libsql-sqlite3/src/vectorInt.h

@@ -102,11 +102,13 @@ int vectorParseSqliteBlob   (sqlite3_value *, Vector *, char **);
 void vectorF32DeserializeFromBlob(Vector *, const unsigned char *, size_t);
 void vectorF64DeserializeFromBlob(Vector *, const unsigned char *, size_t);
 
-void vectorInitStatic(Vector *, VectorType, const unsigned char *, size_t);
+void vectorInitStatic(Vector *, VectorType, VectorDims, void *);
 void vectorInitFromBlob(Vector *, const unsigned char *, size_t);
 void vectorF32InitFromBlob(Vector *, const unsigned char *, size_t);
 void vectorF64InitFromBlob(Vector *, const unsigned char *, size_t);
 
+void vectorConvert(const Vector *, Vector *);
+
 /* Detect type and dimension of vector provided with first parameter of sqlite3_value * type */
 int detectVectorParameters(sqlite3_value *, int, int *, int *, char **);