face optimizations

This commit is contained in:
mertalev 2023-11-21 01:52:30 -05:00
parent 5cb3312589
commit c1175be1d0
No known key found for this signature in database
GPG key ID: 9181CD92C0A1C5E3
3 changed files with 13 additions and 4 deletions

View file

@ -14,7 +14,7 @@ export class AssetFaceEntity {
@Column({ nullable: true, type: 'uuid' }) @Column({ nullable: true, type: 'uuid' })
personId!: string | null; personId!: string | null;
@Column({ type: 'float4', array: true }) @Column({ type: 'float4', array: true, select: false })
embedding!: number[]; embedding!: number[];
@Column({ default: 0, type: 'int' }) @Column({ default: 0, type: 'int' })

View file

@ -13,6 +13,7 @@ export class SmartSearchEntity {
@Column({ @Column({
type: 'float4', type: 'float4',
array: true, array: true,
select: false,
}) })
embedding!: number[]; embedding!: number[];
} }

View file

@ -15,11 +15,17 @@ import { asVector } from '../infra.utils';
import { dataSource } from '..'; import { dataSource } from '..';
export class PersonRepository implements IPersonRepository { export class PersonRepository implements IPersonRepository {
private readonly faceColumns: string[];
constructor( constructor(
@InjectRepository(AssetEntity) private assetRepository: Repository<AssetEntity>, @InjectRepository(AssetEntity) private assetRepository: Repository<AssetEntity>,
@InjectRepository(PersonEntity) private personRepository: Repository<PersonEntity>, @InjectRepository(PersonEntity) private personRepository: Repository<PersonEntity>,
@InjectRepository(AssetFaceEntity) private assetFaceRepository: Repository<AssetFaceEntity>, @InjectRepository(AssetFaceEntity) private assetFaceRepository: Repository<AssetFaceEntity>,
) {} ) {
this.faceColumns = this.assetFaceRepository.manager.connection
.getMetadata(AssetFaceEntity).ownColumns
.map((column) => column.propertyName)
.filter((propertyName) => propertyName !== 'embedding');
}
/** /**
* Before reassigning faces, delete potential key violations * Before reassigning faces, delete potential key violations
@ -245,12 +251,14 @@ export class PersonRepository implements IPersonRepository {
async searchByEmbedding({ ownerId, embedding, numResults, maxDistance }: EmbeddingSearch): Promise<AssetFaceEntity[]> { async searchByEmbedding({ ownerId, embedding, numResults, maxDistance }: EmbeddingSearch): Promise<AssetFaceEntity[]> {
const cte = this.assetFaceRepository.createQueryBuilder('faces') const cte = this.assetFaceRepository.createQueryBuilder('faces')
.select('1 + (faces.embedding <=> :embedding)', 'distance') .select('1 + (faces.embedding <=> :embedding)', 'distance')
.leftJoinAndSelect('faces.asset', 'asset') .innerJoin('faces.asset', 'asset')
.where('asset.ownerId = :ownerId') .where('asset.ownerId = :ownerId')
.orderBy(`faces.embedding <=> :embedding`) .orderBy(`faces.embedding <=> :embedding`)
.setParameters({ownerId, embedding: asVector(embedding)}) .setParameters({ownerId, embedding: asVector(embedding)})
.limit(numResults); .limit(numResults);
this.faceColumns.forEach((col) => cte.addSelect(`faces.${col} AS "${col}"`));
const res = await dataSource.createQueryBuilder() const res = await dataSource.createQueryBuilder()
.select('res.*') .select('res.*')
.addCommonTableExpression(cte, 'cte') .addCommonTableExpression(cte, 'cte')