opusdev/vector-similarity-api

1import type { Document } from '../bson';2import * as BSON from '../bson';3import type { Collection } from '../collection';4import { MongoInvalidArgumentError } from '../error';5import type { DeleteStatement } from '../operations/delete';6import type { UpdateStatement } from '../operations/update';7import { Batch, BatchType, BulkOperationBase, type BulkWriteOptions } from './common';8 9/** @public */10export class OrderedBulkOperation extends BulkOperationBase {11  /** @internal */12  constructor(collection: Collection, options: BulkWriteOptions) {13    super(collection, options, true);14  }15 16  addToOperationsList(17    batchType: BatchType,18    document: Document | UpdateStatement | DeleteStatement19  ): this {20    // Get the bsonSize21    const bsonSize = BSON.calculateObjectSize(document, {22      checkKeys: false,23      // Since we don't know what the user selected for BSON options here,24      // err on the safe side, and check the size with ignoreUndefined: false.25      ignoreUndefined: false26    } as any);27 28    // Throw error if the doc is bigger than the max BSON size29    if (bsonSize >= this.s.maxBsonObjectSize)30      // TODO(NODE-3483): Change this to MongoBSONError31      throw new MongoInvalidArgumentError(32        `Document is larger than the maximum size ${this.s.maxBsonObjectSize}`33      );34 35    // Create a new batch object if we don't have a current one36    if (this.s.currentBatch == null) {37      this.s.currentBatch = new Batch(batchType, this.s.currentIndex);38    }39 40    const maxKeySize = this.s.maxKeySize;41 42    // Check if we need to create a new batch43    if (44      // New batch if we exceed the max batch op size45      this.s.currentBatchSize + 1 >= this.s.maxWriteBatchSize ||46      // New batch if we exceed the maxBatchSizeBytes. Only matters if batch already has a doc,47      // since we can't sent an empty batch48      (this.s.currentBatchSize > 0 &&49        this.s.currentBatchSizeBytes + maxKeySize + bsonSize >= this.s.maxBatchSizeBytes) ||50      // New batch if the new op does not have the same op type as the current batch51      this.s.currentBatch.batchType !== batchType52    ) {53      // Save the batch to the execution stack54      this.s.batches.push(this.s.currentBatch);55 56      // Create a new batch57      this.s.currentBatch = new Batch(batchType, this.s.currentIndex);58 59      // Reset the current size trackers60      this.s.currentBatchSize = 0;61      this.s.currentBatchSizeBytes = 0;62    }63 64    if (batchType === BatchType.INSERT) {65      this.s.bulkResult.insertedIds.push({66        index: this.s.currentIndex,67        _id: (document as Document)._id68      });69    }70 71    // We have an array of documents72    if (Array.isArray(document)) {73      throw new MongoInvalidArgumentError('Operation passed in cannot be an Array');74    }75 76    this.s.currentBatch.originalIndexes.push(this.s.currentIndex);77    this.s.currentBatch.operations.push(document);78    this.s.currentBatchSize += 1;79    this.s.currentBatchSizeBytes += maxKeySize + bsonSize;80    this.s.currentIndex += 1;81    return this;82  }83}84