feat: add support for RQ

Author: bevzzz

src/collections/config/index.ts

@@ -19,6 +19,7 @@ import {
   CollectionConfigUpdate,
   PQConfig,
   QuantizerConfig,
+  RQConfig,
   SQConfig,
   VectorIndexConfig,
   VectorIndexConfigDynamic,
@@ -40,15 +41,15 @@ const config = <T>(
         .withClassName(name)
         .withProperty(resolveProperty<any>(property, []))
         .do()
-        .then(() => {}),
+        .then(() => { }),
     addReference: (
       reference: ReferenceSingleTargetConfigCreate<any> | ReferenceMultiTargetConfigCreate<any>
     ) =>
       new PropertyCreator(connection)
         .withClassName(name)
         .withProperty(resolveReference<any>(reference))
         .do()
-        .then(() => {}),
+        .then(() => { }),
     addVector: async (vectors: VectorizersConfigAdd<T>) => {
       const supportsDynamicVectorIndex = await dbVersionSupport.supportsDynamicVectorIndex();
       const { vectorsConfig } = makeVectorsConfig(vectors, supportsDynamicVectorIndex);
@@ -72,7 +73,7 @@ const config = <T>(
         })
       );
     },
-    updateShards: async function (status: 'READY' | 'READONLY', names?: string | string[]) {
+    updateShards: async function(status: 'READY' | 'READONLY', names?: string | string[]) {
       let shardNames: string[];
       if (names === undefined) {
         shardNames = await this.getShards().then((shards) => shards.map((s) => s.name));
@@ -97,7 +98,7 @@ const config = <T>(
           )
         )
         .then((merged) => new ClassUpdater(connection).withClass(merged).do())
-        .then(() => {});
+        .then(() => { });
     },
   };
 };
@@ -192,6 +193,9 @@ export class Quantizer {
   static isSQ(config?: QuantizerConfig): config is SQConfig {
     return config?.type === 'sq';
   }
+  static isRQ(config?: QuantizerConfig): config is RQConfig {
+    return config?.type === 'rq';
+  }
 }
 
 export const configGuards = {

src/collections/config/integration.test.ts

@@ -12,9 +12,6 @@ import {
   VectorIndexConfigHNSW,
 } from './types/index.js';
 
-const fail = (msg: string) => {
-  throw new Error(msg);
-};
 
 describe('Testing of the collection.config namespace', () => {
   let client: WeaviateClient;
@@ -210,6 +207,30 @@ describe('Testing of the collection.config namespace', () => {
     expect(config.vectorizers.default.vectorizer.name).toEqual('none');
   });
 
+  it('should be able to get the config of a collection with hnsw-rq', async () => {
+    const collectionName = 'TestCollectionConfigGetHNSWPlusRQ';
+    const collection = await client.collections.create({
+      name: collectionName,
+      vectorizers: weaviate.configure.vectorizer.none({
+        vectorIndexConfig: weaviate.configure.vectorIndex.hnsw({
+          quantizer: weaviate.configure.vectorIndex.quantizer.rq(),
+        }),
+      }),
+    });
+    const config = await collection.config.get();
+
+    const vectorIndexConfig = config.vectorizers.default.indexConfig as VectorIndexConfigHNSW;
+    expect(config.name).toEqual(collectionName);
+    expect(config.generative).toBeUndefined();
+    expect(config.reranker).toBeUndefined();
+    expect(vectorIndexConfig).toBeDefined();
+    expect(vectorIndexConfig.quantizer).toBeDefined();
+    expect(vectorIndexConfig.quantizer?.type).toEqual('rq');
+    expect(config.vectorizers.default.indexType).toEqual('hnsw');
+    expect(config.vectorizers.default.properties).toBeUndefined();
+    expect(config.vectorizers.default.vectorizer.name).toEqual('none');
+  });
+
   it('should be able to get the config of a collection with hnsw-bq', async () => {
     const collectionName = 'TestCollectionConfigGetHNSWPlusBQ';
     const query = () =>
@@ -534,8 +555,8 @@ describe('Testing of the collection.config namespace', () => {
       .update({
         propertyDescriptions: supportsUpdatingPropertyDescriptions
           ? {
-              testProp: 'This is a test property',
-            }
+            testProp: 'This is a test property',
+          }
           : undefined,
         vectorizers: weaviate.reconfigure.vectorizer.update({
           vectorIndexConfig: weaviate.reconfigure.vectorIndex.hnsw({

src/collections/config/types/vectorIndex.ts

@@ -9,7 +9,7 @@ export type VectorIndexConfigHNSW = {
   filterStrategy: VectorIndexFilterStrategy;
   flatSearchCutoff: number;
   maxConnections: number;
-  quantizer: PQConfig | BQConfig | SQConfig | undefined;
+  quantizer: PQConfig | BQConfig | SQConfig | RQConfig | undefined;
   skip: boolean;
   vectorCacheMaxObjects: number;
   type: 'hnsw';
@@ -61,6 +61,12 @@ export type PQConfig = {
   type: 'pq';
 };
 
+export type RQConfig = {
+  bits?: number;
+  rescoreLimit?: number;
+  type: 'rq';
+}
+
 export type PQEncoderConfig = {
   type: PQEncoderType;
   distribution: PQEncoderDistribution;
@@ -77,4 +83,4 @@ export type VectorIndexFilterStrategy = 'sweeping' | 'acorn';
 
 export type VectorIndexConfig = VectorIndexConfigHNSW | VectorIndexConfigFlat | VectorIndexConfigDynamic;
 
-export type QuantizerConfig = PQConfig | BQConfig | SQConfig;
+export type QuantizerConfig = PQConfig | BQConfig | SQConfig | RQConfig;

src/collections/config/utils.ts

@@ -46,6 +46,7 @@ import {
   PQEncoderType,
   PropertyConfig,
   PropertyVectorizerConfig,
+  RQConfig,
   ReferenceConfig,
   ReplicationConfig,
   Reranker,
@@ -182,6 +183,16 @@ export const parseVectorIndex = (module: ModuleConfig<VectorIndexType, VectorInd
       },
     };
   }
+  if (QuantizerGuards.isRQCreate(quantizer)) {
+    const { type, ...quant } = quantizer;
+    return {
+      ...conf,
+      rq: {
+        ...quant,
+        enabled: true,
+      },
+    };
+  }
 };
 
 export const parseVectorizerConfig = (config?: VectorizerConfig): any => {
@@ -202,11 +213,11 @@ export const makeVectorsConfig = <TProperties extends Properties | undefined = u
   const vectorizersConfig = Array.isArray(configVectorizers)
     ? configVectorizers
     : [
-        {
-          ...configVectorizers,
-          name: configVectorizers.name || 'default',
-        },
-      ];
+      {
+        ...configVectorizers,
+        name: configVectorizers.name || 'default',
+      },
+    ];
   vectorizersConfig.forEach((v) => {
     if (v.vectorIndex.name === 'dynamic' && !supportsDynamicVectorIndex.supports) {
       throw new WeaviateUnsupportedFeatureError(supportsDynamicVectorIndex.message);
@@ -329,18 +340,18 @@ class ConfigMapping {
         vectorizer:
           v.vectorizer === 'none'
             ? {
-                name: 'none',
-                config: undefined,
-              }
+              name: 'none',
+              config: undefined,
+            }
             : {
-                name: v.vectorizer,
-                config: v.moduleConfig
-                  ? ({
-                      ...(v.moduleConfig[v.vectorizer] as any),
-                      vectorizeCollectionName: (v.moduleConfig[v.vectorizer] as any).vectorizeClassName,
-                    } as VectorizerConfig)
-                  : undefined,
-              },
+              name: v.vectorizer,
+              config: v.moduleConfig
+                ? ({
+                  ...(v.moduleConfig[v.vectorizer] as any),
+                  vectorizeCollectionName: (v.moduleConfig[v.vectorizer] as any).vectorizeClassName,
+                } as VectorizerConfig)
+                : undefined,
+            },
         indexConfig: ConfigMapping.vectorIndex(v.vectorIndexConfig, v.vectorIndexType),
         indexType: ConfigMapping.vectorIndexType(v.vectorIndexType),
       },
@@ -468,11 +479,13 @@ class ConfigMapping {
       throw new WeaviateDeserializationError(
         'Vector index vector cache max objects was not returned by Weaviate'
       );
-    let quantizer: PQConfig | BQConfig | SQConfig | undefined;
+    let quantizer: PQConfig | BQConfig | SQConfig | RQConfig | 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.rq) && v.rq.enabled === true) {
+      quantizer = ConfigMapping.rq(v.rq);
     } else if (exists<Record<string, any>>(v.sq) && v.sq.enabled === true) {
       quantizer = ConfigMapping.sq(v.sq);
     } else {
@@ -508,6 +521,19 @@ class ConfigMapping {
       type: 'bq',
     };
   }
+  static rq(v?: Record<string, unknown>): RQConfig | undefined {
+    if (v === undefined) throw new WeaviateDeserializationError('RQ was not returned by Weaviate');
+    if (!exists<boolean>(v.enabled))
+      throw new WeaviateDeserializationError('RQ enabled was not returned by Weaviate');
+    if (v.enabled === false) return undefined;
+    const bits = v.bits === undefined ? 6 : (v.bits as number);
+    const rescoreLimit = v.rescoreLimit === undefined ? 20 : (v.rescoreLimit as number);
+    return {
+      bits,
+      rescoreLimit,
+      type: 'rq',
+    };
+  }
   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))

src/collections/configure/parsing.ts

@@ -3,6 +3,8 @@ import {
   BQConfigUpdate,
   PQConfigCreate,
   PQConfigUpdate,
+  RQConfigCreate,
+  RQConfigUpdate,
   SQConfigCreate,
   SQConfigUpdate,
 } from './types/index.js';
@@ -13,7 +15,9 @@ type QuantizerConfig =
   | BQConfigCreate
   | BQConfigUpdate
   | SQConfigCreate
-  | SQConfigUpdate;
+  | SQConfigUpdate
+  | RQConfigCreate
+  | RQConfigUpdate;
 
 export class QuantizerGuards {
   static isPQCreate(config?: QuantizerConfig): config is PQConfigCreate {
@@ -34,6 +38,12 @@ export class QuantizerGuards {
   static isSQUpdate(config?: QuantizerConfig): config is SQConfigUpdate {
     return (config as SQConfigUpdate)?.type === 'sq';
   }
+  static isRQCreate(config?: QuantizerConfig): config is RQConfigCreate {
+    return (config as RQConfigCreate)?.type === 'rq';
+  }
+  static isRQUpdate(config?: QuantizerConfig): config is RQConfigUpdate {
+    return (config as RQConfigUpdate)?.type === 'rq';
+  }
 }
 
 export function parseWithDefault<D>(value: D | undefined, defaultValue: D): D {

src/collections/configure/types/vectorIndex.ts

@@ -4,6 +4,7 @@ import {
   PQConfig,
   PQEncoderDistribution,
   PQEncoderType,
+  RQConfig,
   SQConfig,
   VectorDistance,
   VectorIndexConfigDynamic,
@@ -17,6 +18,13 @@ export type QuantizerRecursivePartial<T> = {
   [P in keyof T]: P extends 'type' ? T[P] : RecursivePartial<T[P]> | undefined;
 };
 
+export type RQConfigCreate = QuantizerRecursivePartial<RQConfig>;
+
+export type RQConfigUpdate = {
+  rescoreLimit?: number;
+  type: 'rq';
+}
+
 export type PQConfigCreate = QuantizerRecursivePartial<PQConfig>;
 
 export type PQConfigUpdate = {
@@ -59,7 +67,7 @@ export type VectorIndexConfigHNSWUpdate = {
   ef?: number;
   filterStrategy?: VectorIndexFilterStrategy;
   flatSearchCutoff?: number;
-  quantizer?: PQConfigUpdate | BQConfigUpdate | SQConfigUpdate;
+  quantizer?: PQConfigUpdate | BQConfigUpdate | SQConfigUpdate | RQConfigUpdate;
   vectorCacheMaxObjects?: number;
 };
 
@@ -131,7 +139,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 | SQConfigCreate;
+  quantizer?: PQConfigCreate | BQConfigCreate | SQConfigCreate | RQConfigCreate;
   /** Whether to skip the index. Default is false. */
   skip?: boolean;
   /** The maximum number of objects to cache in the vector cache. Default is 1000000000000. */