ソースを参照

set k at startup

mertalev 1 年間 前
コミット
5f287449bf

+ 3 - 0
server/src/infra/database.config.ts

@@ -39,6 +39,9 @@ export async function initDataSource(): Promise<DataSource> {
   if (hasVectorExtension) {
     const dataSourceVector = await new DataSource(databaseConfigVector).initialize();
     await dataSourceVector.runMigrations();
+
+    await dataSourceVector.query(`SET vectors.enable_prefilter = on`);
+
     await dataSourceVector.destroy();
   }
 

+ 1 - 0
server/src/infra/entities/exif.entity.ts

@@ -98,6 +98,7 @@ export class ExifEntity {
   @Column({
     type: 'tsvector',
     generatedType: 'STORED',
+    select: false,
     asExpression: `TO_TSVECTOR('english',
                          COALESCE(make, '') || ' ' ||
                          COALESCE(model, '') || ' ' ||

+ 0 - 1
server/src/infra/migrations/vector/1699746301742-AddCLIPEmbeddingIndex.ts

@@ -7,7 +7,6 @@ export class AddCLIPEmbeddingIndex1699746301742 implements MigrationInterface {
     await queryRunner.query(`
         CREATE INDEX IF NOT EXISTS clip_index ON smart_search
         USING vectors (embedding cosine_ops) WITH (options = $$
-        capacity = 2097152
         [indexing.hnsw]
         m = 16
         ef_construction = 300

+ 0 - 1
server/src/infra/migrations/vector/1699746444644-AddFaceEmbeddingIndex.ts

@@ -7,7 +7,6 @@ export class AddFaceEmbeddingIndex1699746444644 implements MigrationInterface {
     await queryRunner.query(`
         CREATE INDEX IF NOT EXISTS face_index ON asset_faces
         USING vectors (embedding cosine_ops) WITH (options = $$
-        capacity = 2097152
         [indexing.hnsw]
         m = 16
         ef_construction = 300

+ 9 - 4
server/src/infra/repositories/smart-info.repository.ts

@@ -20,16 +20,22 @@ export class SmartInfoRepository implements ISmartInfoRepository {
   }
 
   async searchByEmbedding({ ownerId, embedding, numResults }: EmbeddingSearch): Promise<AssetEntity[]> {
-    let results: AssetEntity[] = await this.assetRepository.createQueryBuilder('a')
+    const query: string = this.assetRepository.createQueryBuilder('a')
     .innerJoin('a.smartSearch', 's')
     .where('a.ownerId = :ownerId')
     .leftJoinAndSelect('a.exifInfo', 'e')
     .orderBy('s.embedding <=> :embedding')
     .setParameters({ embedding: asVector(embedding), ownerId })
     .limit(numResults)
-    .getMany();
+    .getSql();
 
-    return results;
+    const queryWithK = `
+      BEGIN;
+      SET LOCAL vectors.k = ${numResults};
+      ${query};
+      COMMIT;
+    `
+    return this.assetRepository.create(await this.assetRepository.manager.query(queryWithK));
   }
 
   async upsert(smartInfo: Partial<SmartInfoEntity>, embedding?: Embedding): Promise<void> {
@@ -70,7 +76,6 @@ export class SmartInfoRepository implements ISmartInfoRepository {
 
         CREATE INDEX clip_index ON smart_search
         USING vectors (embedding dot_ops) WITH (options = $$
-        capacity = 2097152
         [indexing.hnsw]
         m = 16
         ef_construction = 300