in storage-reverse-image-search/functions/src/functions/backfill_task.ts [27:104]
export async function backfillEmbeddingsTaskHandler(data: any) {
const {id, objects} = data;
const taskRef = admin.firestore().doc(`${config.tasksDoc}/enqueues/${id}`);
await taskRef.update({
status: BackfillStatus.PROCESSING,
});
const datapoints: any = [];
const featureVectors = await getFeatureVectors(objects);
if (featureVectors) {
for (let i = 0; i < featureVectors.length; i++) {
datapoints.push({
id: objects[i],
embedding: featureVectors[i],
});
}
}
if (datapoints.length === 0) {
functions.logger.info('No datapoints found, skipping...');
return;
}
const outputShape = featureVectors[0].length;
const localFilePath = await utils.saveEmbeddingsToTmpFile(datapoints);
const destinationPath = `datapoints/${id}.json`;
functions.logger.info(
`Embeddings will be saved to ${destinationPath} 📝, uploading to the bucket...`
);
await utils.uploadToCloudStorage(localFilePath, destinationPath);
functions.logger.info(
`Embeddings uploaded to the bucket ${config.bucketName} in ${destinationPath} 🚀`
);
await taskRef.update({
status: 'DONE',
filePath: `gs://${config.bucketName}/${destinationPath}`,
});
await utils.deleteTempFiles();
const tasksDoc = await admin.firestore().doc(config.tasksDoc).get();
const {totalLength} = tasksDoc.data() as any;
let {processedLength} = tasksDoc.data() as any;
processedLength += objects.length;
await admin
.firestore()
.doc(config.tasksDoc)
.update({
processedLength: admin.firestore.FieldValue.increment(objects.length),
});
if (processedLength === totalLength) {
// Update the metadata doc with the output shape to be used in the index creation.
await admin.firestore().doc(config.metadataDoc).set(
{
outputShape,
},
{merge: true}
);
await admin.firestore().doc(config.tasksDoc).update({
status: BackfillStatus.DONE,
});
} else {
await _createNextTask(id);
}
}