Merge pull request #166 from weaviate/1.26/support-sq-configuration

tsmith023 · web-flow · commit bcb9f5f42b8c · 2024-07-10T12:00:58.000+01:00
Add support for configuring and reconfiguring SQ
diff --git a/src/collections/config/classes.ts b/src/collections/config/classes.ts
@@ -1,4 +1,5 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
+import { WeaviateInvalidInputError } from '../../errors.js';
 import {
   WeaviateClass,
   WeaviateInvertedIndexConfig,
@@ -118,10 +119,14 @@ export class MergeWithExisting {
   ): WeaviateVectorIndexConfig {
     if (update === undefined) return current;
     if (
-      (QuantizerGuards.isBQUpdate(update.quantizer) && ((current?.bq as any) || {}).enabled) ||
-      (QuantizerGuards.isPQUpdate(update.quantizer) && ((current?.pq as any) || {}).enabled)
+      (QuantizerGuards.isBQUpdate(update.quantizer) &&
+        (((current?.pq as any) || {}).enabled || ((current?.sq as any) || {}).enabled)) ||
+      (QuantizerGuards.isPQUpdate(update.quantizer) &&
+        (((current?.bq as any) || {}).enabled || ((current?.sq as any) || {}).enabled)) ||
+      (QuantizerGuards.isSQUpdate(update.quantizer) &&
+        (((current?.pq as any) || {}).enabled || ((current?.bq as any) || {}).enabled))
     )
-      throw Error(`Cannot update the quantizer type of an enabled vector index.`);
+      throw new WeaviateInvalidInputError(`Cannot update the quantizer type of an enabled vector index.`);
     const { quantizer, ...rest } = update;
     const merged: WeaviateVectorIndexConfig = { ...current, ...rest };
     if (QuantizerGuards.isBQUpdate(quantizer)) {
@@ -132,6 +137,10 @@ export class MergeWithExisting {
       const { type, ...quant } = quantizer;
       merged.pq = { ...current!.pq!, ...quant, enabled: true };
     }
+    if (QuantizerGuards.isSQUpdate(quantizer)) {
+      const { type, ...quant } = quantizer;
+      merged.sq = { ...current!.sq!, ...quant, enabled: true };
+    }
     return merged;
   }
 }
diff --git a/src/collections/config/types/vectorIndex.ts b/src/collections/config/types/vectorIndex.ts
@@ -8,7 +8,7 @@ export type VectorIndexConfigHNSW = {
   ef: number;
   flatSearchCutoff: number;
   maxConnections: number;
-  quantizer: PQConfig | BQConfig | undefined;
+  quantizer: PQConfig | BQConfig | SQConfig | undefined;
   skip: boolean;
   vectorCacheMaxObjects: number;
   type: 'hnsw';
@@ -45,6 +45,12 @@ export type BQConfig = {
   type: 'bq';
 };
 
+export type SQConfig = {
+  rescoreLimit: number;
+  trainingLimit: number;
+  type: 'sq';
+};
+
 export type PQConfig = {
   bitCompression: boolean;
   centroids: number;
diff --git a/src/collections/config/unit.test.ts b/src/collections/config/unit.test.ts
@@ -43,6 +43,9 @@ describe('Unit testing of the MergeWithExisting class', () => {
         bq: {
           enabled: false,
         },
+        sq: {
+          enabled: false,
+        },
       },
       vectorIndexType: 'hnsw',
       vectorizer: {
@@ -55,7 +58,7 @@ describe('Unit testing of the MergeWithExisting class', () => {
   };
 
   it('should merge a full invertedIndexUpdate with existing schema', () => {
-    const merged = MergeWithExisting.invertedIndex(Object.assign({}, invertedIndex), {
+    const merged = MergeWithExisting.invertedIndex(JSON.parse(JSON.stringify(invertedIndex)), {
       bm25: {
         b: 0.9,
         k1: 1.4,
@@ -110,7 +113,7 @@ describe('Unit testing of the MergeWithExisting class', () => {
   };
 
   it('should merge a partial invertedIndexUpdate with existing schema', () => {
-    const merged = MergeWithExisting.invertedIndex(Object.assign({}, invertedIndex), {
+    const merged = MergeWithExisting.invertedIndex(JSON.parse(JSON.stringify(invertedIndex)), {
       bm25: {
         b: 0.9,
       },
@@ -134,7 +137,7 @@ describe('Unit testing of the MergeWithExisting class', () => {
   });
 
   it('should merge a no quantizer HNSW vectorIndexConfig with existing schema', () => {
-    const merged = MergeWithExisting.vectors(Object.assign({}, hnswVectorConfig), [
+    const merged = MergeWithExisting.vectors(JSON.parse(JSON.stringify(hnswVectorConfig)), [
       {
         name: 'name',
         vectorIndex: {
@@ -158,6 +161,7 @@ describe('Unit testing of the MergeWithExisting class', () => {
     expect(merged).toEqual({
       name: {
         vectorIndexConfig: {
+          ...hnswVectorConfig.name.vectorIndexConfig,
           skip: true,
           cleanupIntervalSeconds: 301,
           maxConnections: 65,
@@ -169,20 +173,6 @@ describe('Unit testing of the MergeWithExisting class', () => {
           vectorCacheMaxObjects: 1000000000001,
           flatSearchCutoff: 40001,
           distance: 'euclidean',
-          pq: {
-            enabled: false,
-            bitCompression: false,
-            segments: 0,
-            centroids: 256,
-            trainingLimit: 100000,
-            encoder: {
-              type: 'kmeans',
-              distribution: 'log-normal',
-            },
-          },
-          bq: {
-            enabled: false,
-          },
         },
         vectorIndexType: 'hnsw',
         vectorizer: {
@@ -196,7 +186,7 @@ describe('Unit testing of the MergeWithExisting class', () => {
   });
 
   it('should merge a PQ quantizer HNSW vectorIndexConfig with existing schema', () => {
-    const merged = MergeWithExisting.vectors(Object.assign({}, hnswVectorConfig), [
+    const merged = MergeWithExisting.vectors(JSON.parse(JSON.stringify(hnswVectorConfig)), [
       {
         name: 'name',
         vectorIndex: {
@@ -220,17 +210,7 @@ describe('Unit testing of the MergeWithExisting class', () => {
     expect(merged).toEqual({
       name: {
         vectorIndexConfig: {
-          skip: true,
-          cleanupIntervalSeconds: 301,
-          maxConnections: 65,
-          efConstruction: 129,
-          ef: -2,
-          dynamicEfMin: 101,
-          dynamicEfMax: 501,
-          dynamicEfFactor: 9,
-          vectorCacheMaxObjects: 1000000000001,
-          flatSearchCutoff: 40001,
-          distance: 'euclidean',
+          ...hnswVectorConfig.name.vectorIndexConfig,
           pq: {
             enabled: true,
             bitCompression: true,
@@ -242,9 +222,6 @@ describe('Unit testing of the MergeWithExisting class', () => {
               distribution: 'normal',
             },
           },
-          bq: {
-            enabled: false,
-          },
         },
         vectorIndexType: 'hnsw',
         vectorizer: {
@@ -258,7 +235,7 @@ describe('Unit testing of the MergeWithExisting class', () => {
   });
 
   it('should merge a BQ quantizer HNSW vectorIndexConfig with existing schema', () => {
-    const merged = MergeWithExisting.vectors(Object.assign({}, hnswVectorConfig), [
+    const merged = MergeWithExisting.vectors(JSON.parse(JSON.stringify(hnswVectorConfig)), [
       {
         name: 'name',
         vectorIndex: {
@@ -292,8 +269,45 @@ describe('Unit testing of the MergeWithExisting class', () => {
     });
   });
 
+  it('should merge a SQ quantizer HNSW vectorIndexConfig with existing schema', () => {
+    const merged = MergeWithExisting.vectors(JSON.parse(JSON.stringify(hnswVectorConfig)), [
+      {
+        name: 'name',
+        vectorIndex: {
+          name: 'hnsw',
+          config: {
+            quantizer: {
+              type: 'sq',
+              rescoreLimit: 1000,
+              trainingLimit: 10000,
+            },
+          },
+        },
+      },
+    ]);
+    expect(merged).toEqual({
+      name: {
+        vectorIndexConfig: {
+          ...hnswVectorConfig.name.vectorIndexConfig,
+          sq: {
+            enabled: true,
+            rescoreLimit: 1000,
+            trainingLimit: 10000,
+          },
+        },
+        vectorIndexType: 'hnsw',
+        vectorizer: {
+          'text2vec-contextionary': {
+            properties: ['name'],
+            vectorizeCollectionName: false,
+          },
+        },
+      },
+    });
+  });
+
   it('should merge a BQ quantizer Flat vectorIndexConfig with existing schema', () => {
-    const merged = MergeWithExisting.vectors(Object.assign({}, flatVectorConfig), [
+    const merged = MergeWithExisting.vectors(JSON.parse(JSON.stringify(flatVectorConfig)), [
       {
         name: 'name',
         vectorIndex: {
diff --git a/src/collections/config/utils.ts b/src/collections/config/utils.ts
@@ -37,6 +37,7 @@ import {
   ReplicationConfig,
   Reranker,
   RerankerConfig,
+  SQConfig,
   ShardingConfig,
   VectorConfig,
   VectorDistance,
@@ -356,11 +357,13 @@ class ConfigMapping {
       throw new WeaviateDeserializationError(
         'Vector index vector cache max objects was not returned by Weaviate'
       );
-    let quantizer: PQConfig | BQConfig | undefined;
+    let quantizer: PQConfig | BQConfig | SQConfig | undefined;
     if (exists<Record<string, any>>(v.pq) && v.pq.enabled === true) {
       quantizer = ConfigMapping.pq(v.pq);
     } else if (exists<Record<string, any>>(v.bq) && v.bq.enabled === true) {
       quantizer = ConfigMapping.bq(v.bq);
+    } else if (exists<Record<string, any>>(v.sq) && v.sq.enabled === true) {
+      quantizer = ConfigMapping.sq(v.sq);
     } else {
       quantizer = undefined;
     }
@@ -393,6 +396,19 @@ class ConfigMapping {
       type: 'bq',
     };
   }
+  static sq(v?: Record<string, unknown>): SQConfig | undefined {
+    if (v === undefined) throw new WeaviateDeserializationError('SQ was not returned by Weaviate');
+    if (!exists<boolean>(v.enabled))
+      throw new WeaviateDeserializationError('SQ enabled was not returned by Weaviate');
+    if (v.enabled === false) return undefined;
+    const rescoreLimit = v.rescoreLimit === undefined ? 1000 : (v.rescoreLimit as number);
+    const trainingLimit = v.trainingLimit === undefined ? 100000 : (v.trainingLimit as number);
+    return {
+      rescoreLimit,
+      trainingLimit,
+      type: 'sq',
+    };
+  }
   static vectorIndexFlat(v: WeaviateVectorIndexConfig): VectorIndexConfigFlat {
     if (v === undefined) throw new WeaviateDeserializationError('Vector index was not returned by Weaviate');
     if (!exists<number>(v.vectorCacheMaxObjects))
diff --git a/src/collections/configure/parsing.ts b/src/collections/configure/parsing.ts
@@ -1,6 +1,19 @@
-import { BQConfigCreate, BQConfigUpdate, PQConfigCreate, PQConfigUpdate } from './types/index.js';
+import {
+  BQConfigCreate,
+  BQConfigUpdate,
+  PQConfigCreate,
+  PQConfigUpdate,
+  SQConfigCreate,
+  SQConfigUpdate,
+} from './types/index.js';
 
-type QuantizerConfig = PQConfigCreate | PQConfigUpdate | BQConfigCreate | BQConfigUpdate;
+type QuantizerConfig =
+  | PQConfigCreate
+  | PQConfigUpdate
+  | BQConfigCreate
+  | BQConfigUpdate
+  | SQConfigCreate
+  | SQConfigUpdate;
 
 export class QuantizerGuards {
   static isPQCreate(config?: QuantizerConfig): config is PQConfigCreate {
@@ -15,6 +28,12 @@ export class QuantizerGuards {
   static isBQUpdate(config?: QuantizerConfig): config is BQConfigUpdate {
     return (config as BQConfigUpdate)?.type === 'bq';
   }
+  static isSQCreate(config?: QuantizerConfig): config is SQConfigCreate {
+    return (config as SQConfigCreate)?.type === 'sq';
+  }
+  static isSQUpdate(config?: QuantizerConfig): config is SQConfigUpdate {
+    return (config as SQConfigUpdate)?.type === 'sq';
+  }
 }
 
 export function parseWithDefault<D>(value: D | undefined, defaultValue: D): D {
diff --git a/src/collections/configure/types/vectorIndex.ts b/src/collections/configure/types/vectorIndex.ts
@@ -4,6 +4,7 @@ import {
   PQConfig,
   PQEncoderDistribution,
   PQEncoderType,
+  SQConfig,
   VectorDistance,
   VectorIndexConfigDynamic,
   VectorIndexConfigFlat,
@@ -36,6 +37,14 @@ export type BQConfigUpdate = {
   type: 'bq';
 };
 
+export type SQConfigCreate = QuantizerRecursivePartial<SQConfig>;
+
+export type SQConfigUpdate = {
+  rescoreLimit?: number;
+  trainingLimit?: number;
+  type: 'sq';
+};
+
 export type VectorIndexConfigHNSWCreate = RecursivePartial<VectorIndexConfigHNSW>;
 
 export type VectorIndexConfigDynamicCreate = RecursivePartial<VectorIndexConfigDynamic>;
@@ -48,7 +57,7 @@ export type VectorIndexConfigHNSWUpdate = {
   dynamicEfFactor?: number;
   ef?: number;
   flatSearchCutoff?: number;
-  quantizer?: PQConfigUpdate | BQConfigUpdate;
+  quantizer?: PQConfigUpdate | BQConfigUpdate | SQConfigUpdate;
   vectorCacheMaxObjects?: number;
 };
 
@@ -118,7 +127,7 @@ export type VectorIndexConfigHNSWCreateOptions = {
   /** The maximum number of connections. Default is 64. */
   maxConnections?: number;
   /** The quantizer configuration to use. Use `vectorIndex.quantizer.bq` or `vectorIndex.quantizer.pq` to make one. */
-  quantizer?: PQConfigCreate | BQConfigCreate;
+  quantizer?: PQConfigCreate | BQConfigCreate | SQConfigCreate;
   /** Whether to skip the index. Default is false. */
   skip?: boolean;
   /** The maximum number of objects to cache in the vector cache. Default is 1000000000000. */
diff --git a/src/collections/configure/unit.test.ts b/src/collections/configure/unit.test.ts
@@ -198,6 +198,25 @@ describe('Unit testing of the configure factory class', () => {
       },
     });
   });
+
+  it('should create an hnsw VectorIndexConfig type with SQ quantizer', () => {
+    const config = configure.vectorIndex.hnsw({
+      quantizer: configure.vectorIndex.quantizer.sq({
+        rescoreLimit: 100,
+        trainingLimit: 200,
+      }),
+    });
+    expect(config).toEqual<ModuleConfig<'hnsw', VectorIndexConfigHNSWCreate>>({
+      name: 'hnsw',
+      config: {
+        quantizer: {
+          rescoreLimit: 100,
+          trainingLimit: 200,
+          type: 'sq',
+        },
+      },
+    });
+  });
 });
 
 describe('Unit testing of the vectorizer factory class', () => {
diff --git a/src/collections/configure/vectorIndex.ts b/src/collections/configure/vectorIndex.ts
