Skip to content

Commit 6c0c1ba

Browse files
mitchwadairmitchwadair
committed
fix vector index diffing, tests
1 parent aed4387 commit 6c0c1ba

File tree

7 files changed

+113
-28
lines changed

7 files changed

+113
-28
lines changed

‎drizzle-kit/src/serializer/singlestoreSchema.ts‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ const vectorIndex = object({
2323
M: number().optional(),
2424
ef: number().optional(),
2525
efConstruction: number().optional(),
26-
});
26+
}).strict();
2727

2828
const column = object({
2929
name: string(),
@@ -182,13 +182,13 @@ export const SingleStoreSquasher = {
182182
column,
183183
indexType,
184184
metricType: metricType ? metricType : undefined,
185-
nlist: nlist ? nlist : undefined,
186-
nprobe: nprobe ? nprobe : undefined,
187-
nbits: nbits ? nbits : undefined,
188-
m: m ? m : undefined,
189-
M: M ? M : undefined,
190-
ef: ef ? ef : undefined,
191-
efConstruction: efConstruction ? efConstruction : undefined,
185+
nlist: nlist ? parseInt(nlist) : undefined,
186+
nprobe: nprobe ? parseInt(nprobe) : undefined,
187+
nbits: nbits ? parseInt(nbits) : undefined,
188+
m: m ? parseInt(m) : undefined,
189+
M: M ? parseInt(M) : undefined,
190+
ef: ef ? parseInt(ef) : undefined,
191+
efConstruction: efConstruction ? parseInt(efConstruction) : undefined,
192192
};
193193
return vectorIndex.parse(deconstructed);
194194
},

‎drizzle-kit/src/serializer/singlestoreSerializer.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { RowDataPacket } from 'mysql2/promise';
1111
import { withStyle } from '../cli/validations/outputs';
1212
import { IntrospectStage, IntrospectStatus } from '../cli/views';
1313

14+
import { VectorIndexType, VectorMetricType } from 'drizzle-orm/singlestore-core/indexes/vector';
1415
import { CasingType } from 'src/cli/validations/common';
1516
import type { DB } from '../utils';
1617
import {
@@ -728,6 +729,8 @@ export const fromDatabase = async (
728729
const constraintName = idxRow['INDEX_NAME'];
729730
const columnName: string = idxRow['COLUMN_NAME'];
730731
const isUnique = idxRow['NON_UNIQUE'] === 0;
732+
const idxType = idxRow['INDEX_TYPE'];
733+
const idxOptions = idxRow['INDEX_OPTIONS'];
731734

732735
const tableInResult = result[tableName];
733736
if (typeof tableInResult === 'undefined') continue;
@@ -752,6 +755,22 @@ export const fromDatabase = async (
752755
columns: [columnName],
753756
};
754757
}
758+
} else {
759+
if (idxType === 'VECTOR') {
760+
const { index_type: indexType, metric_type: metricType, ...restOpts } = JSON.parse(idxOptions) as
761+
& Omit<VectorIndex, 'name' | 'column' | 'indexType' | 'metricType'>
762+
& { index_type: VectorIndexType; metric_type: VectorMetricType };
763+
const vIdx: VectorIndex = {
764+
name: constraintName,
765+
column: columnName,
766+
indexType,
767+
metricType,
768+
...restOpts,
769+
};
770+
tableInResult.vectorIndexes[constraintName] = vIdx;
771+
} else {
772+
// TODO non-vector indexes
773+
}
755774
}
756775
}
757776

‎drizzle-kit/src/snapshotsDiffer.ts‎

Lines changed: 26 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2724,27 +2724,39 @@ export const applySingleStoreSnapshotsDiff = async (
27242724
for (const tableName in json1.tables) {
27252725
const table = json1.tables[tableName];
27262726
for (const indexName in table.indexes) {
2727-
const index = SingleStoreSquasher.unsquashIdx(table.indexes[indexName]);
2728-
if (index.isUnique) {
2729-
table.uniqueConstraints[indexName] = SingleStoreSquasher.squashUnique({
2730-
name: index.name,
2731-
columns: index.columns,
2732-
});
2733-
delete json1.tables[tableName].indexes[index.name];
2727+
try {
2728+
// skip vector indexes
2729+
SingleStoreSquasher.unsquashVectorIdx(table.indexes[indexName]);
2730+
continue;
2731+
} catch {
2732+
const index = SingleStoreSquasher.unsquashIdx(table.indexes[indexName]);
2733+
if (index.isUnique) {
2734+
table.uniqueConstraints[indexName] = SingleStoreSquasher.squashUnique({
2735+
name: index.name,
2736+
columns: index.columns,
2737+
});
2738+
delete json1.tables[tableName].indexes[index.name];
2739+
}
27342740
}
27352741
}
27362742
}
27372743

27382744
for (const tableName in json2.tables) {
27392745
const table = json2.tables[tableName];
27402746
for (const indexName in table.indexes) {
2741-
const index = SingleStoreSquasher.unsquashIdx(table.indexes[indexName]);
2742-
if (index.isUnique) {
2743-
table.uniqueConstraints[indexName] = SingleStoreSquasher.squashUnique({
2744-
name: index.name,
2745-
columns: index.columns,
2746-
});
2747-
delete json2.tables[tableName].indexes[index.name];
2747+
try {
2748+
// skip vector indexes
2749+
SingleStoreSquasher.unsquashVectorIdx(table.indexes[indexName]);
2750+
continue;
2751+
} catch {
2752+
const index = SingleStoreSquasher.unsquashIdx(table.indexes[indexName]);
2753+
if (index.isUnique) {
2754+
table.uniqueConstraints[indexName] = SingleStoreSquasher.squashUnique({
2755+
name: index.name,
2756+
columns: index.columns,
2757+
});
2758+
delete json2.tables[tableName].indexes[index.name];
2759+
}
27482760
}
27492761
}
27502762
}

‎drizzle-kit/src/sqlgenerator.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3526,7 +3526,9 @@ export class CreateSingleStoreIndexConvertor extends Convertor {
35263526

35273527
return `CREATE ${indexPart} \`${name}\` ON \`${statement.tableName}\` (${uniqueString});`;
35283528
} else {
3529-
const { name, column, indexType, ...restIndexOptions } = SingleStoreSquasher.unsquashVectorIdx(statement.data);
3529+
const { name, column, indexType, metricType, ...restIndexOptions } = SingleStoreSquasher.unsquashVectorIdx(
3530+
statement.data,
3531+
);
35303532

35313533
const columnString = statement.internal?.indexes
35323534
? statement.internal?.indexes[name]?.columns[column]?.isExpression
@@ -3537,6 +3539,9 @@ export class CreateSingleStoreIndexConvertor extends Convertor {
35373539
const indexOptions: Record<string, string | number> = {
35383540
index_type: indexType,
35393541
};
3542+
if (metricType) {
3543+
indexOptions['metric_type'] = metricType;
3544+
}
35403545
for (const [key, value] of Object.entries(restIndexOptions)) {
35413546
indexOptions[key] = value;
35423547
}

‎drizzle-kit/tests/push/singlestore.test.ts‎

Lines changed: 52 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,7 @@ const singlestoreSuite: DialectSuite = {
279279
await context.client.query(st);
280280
}
281281
},
282+
// TODO non-vector indexes
282283
addBasicIndexes: async function(context?: any): Promise<void> {
283284
const schema1 = {
284285
vectorIndex: singlestoreTable('vector_table', {
@@ -329,7 +330,48 @@ const singlestoreSuite: DialectSuite = {
329330
await context.client.query(st);
330331
}
331332
},
332-
changeIndexFields: function(context?: any): Promise<void> {
333+
changeIndexFields: async function(context?: any): Promise<void> {
334+
const schema1 = {
335+
vectorIndex: singlestoreTable('vector_table', {
336+
v: vector('v', { dimensions: 10 }),
337+
}, (table) => [
338+
vectorIndex('vector_index').on(table.v),
339+
]),
340+
};
341+
342+
const schema2 = {
343+
vectorIndex: singlestoreTable('vector_table', {
344+
v: vector('v', { dimensions: 10 }),
345+
}, (table) => [
346+
vectorIndex('vector_index', 'IVF_PQ').on(table.v).metricType('EUCLIDEAN_DISTANCE').nbits(16),
347+
]),
348+
};
349+
350+
const { statements, sqlStatements } = await diffTestSchemasPushSingleStore(
351+
context.client as Connection,
352+
schema1,
353+
schema2,
354+
[],
355+
'drizzle',
356+
false,
357+
);
358+
expect(statements.length).toBe(2);
359+
expect(sqlStatements).toStrictEqual([
360+
`DROP INDEX \`vector_index\` ON \`vector_table\`;`,
361+
`ALTER TABLE \`vector_table\` ADD VECTOR INDEX \`vector_index\` (\`v\`) INDEX_OPTIONS '{"index_type":"IVF_PQ","metric_type":"EUCLIDEAN_DISTANCE","nbits":16}';`,
362+
]);
363+
364+
const { sqlStatements: dropStatements } = await diffTestSchemasSingleStore(
365+
schema2,
366+
{},
367+
[],
368+
false,
369+
);
370+
371+
for (const st of dropStatements) {
372+
await context.client.query(st);
373+
}
374+
333375
return {} as any;
334376
},
335377
dropIndex: async function(context?: any): Promise<void> {
@@ -355,9 +397,16 @@ const singlestoreSuite: DialectSuite = {
355397
'drizzle',
356398
false,
357399
);
358-
console.log(statements);
359-
console.log(sqlStatements);
360400
expect(statements.length).toBe(1);
401+
expect(statements[0]).toStrictEqual({
402+
type: 'drop_index',
403+
tableName: 'vector_table',
404+
data: 'vector_index;v;AUTO;;;;;;;;',
405+
schema: '',
406+
});
407+
expect(sqlStatements[0]).toBe(
408+
`DROP INDEX \`vector_index\` ON \`vector_table\`;`,
409+
);
361410

362411
const { sqlStatements: dropStatements } = await diffTestSchemasSingleStore(
363412
schema2,

‎drizzle-kit/tests/schemaDiffer.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1657,7 +1657,7 @@ export const diffTestSchemasPushSingleStore = async (
16571657
await client.query(st);
16581658
}
16591659

1660-
// do introspect into PgSchemaInternal
1660+
// do introspect into SingleStoreSchemaInternal
16611661
const introspectedSchema = await fromSingleStoreDatabase(
16621662
{
16631663
query: async (sql: string, params?: any[]) => {

‎drizzle-orm/src/singlestore-core/indexes/vector.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { IndexColumn } from '~/singlestore-core/indexes.ts';
33
import { SingleStoreTable } from '~/singlestore-core/table.ts';
44

55
export type VectorIndexType = 'AUTO' | 'FLAT' | 'IVF_FLAT' | 'IVF_PQ' | 'IVF_PQFS' | 'HNSW_FLAT' | 'HNSW_PQ';
6-
type VectorMetricType = 'EUCLIDEAN_DISTANCE' | 'DOT_PRODUCT';
6+
export type VectorMetricType = 'EUCLIDEAN_DISTANCE' | 'DOT_PRODUCT';
77

88
// type VectorIndexConfigExtensions<T extends VectorIndexType> = T extends 'IVF_FLAT' ? { nlist?: number; nprobe?: number }
99
// : T extends 'IVF_PQ' ? { nlist?: number; m?: number; nbits?: number; nprobe?: number }

0 commit comments

Comments
 (0)