Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for configuring and reconfiguring SQ #166

Merged
merged 2 commits into from
Jul 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions src/collections/config/classes.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
/* eslint-disable @typescript-eslint/no-non-null-assertion */
import { WeaviateInvalidInputError } from '../../errors.js';
import {
WeaviateClass,
WeaviateInvertedIndexConfig,
Expand Down Expand Up @@ -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)) {
Expand All @@ -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;
}
}
8 changes: 7 additions & 1 deletion src/collections/config/types/vectorIndex.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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;
Expand Down
82 changes: 48 additions & 34 deletions src/collections/config/unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,9 @@ describe('Unit testing of the MergeWithExisting class', () => {
bq: {
enabled: false,
},
sq: {
enabled: false,
},
},
vectorIndexType: 'hnsw',
vectorizer: {
Expand All @@ -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,
Expand Down Expand Up @@ -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,
},
Expand All @@ -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: {
Expand All @@ -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,
Expand All @@ -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: {
Expand All @@ -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: {
Expand All @@ -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,
Expand All @@ -242,9 +222,6 @@ describe('Unit testing of the MergeWithExisting class', () => {
distribution: 'normal',
},
},
bq: {
enabled: false,
},
},
vectorIndexType: 'hnsw',
vectorizer: {
Expand All @@ -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: {
Expand Down Expand Up @@ -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: {
Expand Down
18 changes: 17 additions & 1 deletion src/collections/config/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ import {
ReplicationConfig,
Reranker,
RerankerConfig,
SQConfig,
ShardingConfig,
VectorConfig,
VectorDistance,
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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))
Expand Down
23 changes: 21 additions & 2 deletions src/collections/configure/parsing.ts
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -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 {
Expand Down
13 changes: 11 additions & 2 deletions src/collections/configure/types/vectorIndex.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
PQConfig,
PQEncoderDistribution,
PQEncoderType,
SQConfig,
VectorDistance,
VectorIndexConfigDynamic,
VectorIndexConfigFlat,
Expand Down Expand Up @@ -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>;
Expand All @@ -48,7 +57,7 @@ export type VectorIndexConfigHNSWUpdate = {
dynamicEfFactor?: number;
ef?: number;
flatSearchCutoff?: number;
quantizer?: PQConfigUpdate | BQConfigUpdate;
quantizer?: PQConfigUpdate | BQConfigUpdate | SQConfigUpdate;
vectorCacheMaxObjects?: number;
};

Expand Down Expand Up @@ -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. */
Expand Down
19 changes: 19 additions & 0 deletions src/collections/configure/unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
Loading
Loading