opusdev/vector-similarity-api

1import type { Document } from '../bson';2import { type Connection } from '../cmap/connection';3import { CursorResponse, MongoDBResponse } from '../cmap/wire_protocol/responses';4import type { Collection } from '../collection';5import { type AbstractCursorOptions } from '../cursor/abstract_cursor';6import { MongoCompatibilityError } from '../error';7import { type OneOrMore } from '../mongo_types';8import { isObject, maxWireVersion, type MongoDBNamespace } from '../utils';9import {10  type CollationOptions,11  CommandOperation,12  type CommandOperationOptions,13  type OperationParent14} from './command';15import { Aspect, defineAspects } from './operation';16 17const VALID_INDEX_OPTIONS = new Set([18  'background',19  'unique',20  'name',21  'partialFilterExpression',22  'sparse',23  'hidden',24  'expireAfterSeconds',25  'storageEngine',26  'collation',27  'version',28 29  // text indexes30  'weights',31  'default_language',32  'language_override',33  'textIndexVersion',34 35  // 2d-sphere indexes36  '2dsphereIndexVersion',37 38  // 2d indexes39  'bits',40  'min',41  'max',42 43  // geoHaystack Indexes44  'bucketSize',45 46  // wildcard indexes47  'wildcardProjection'48]);49 50/** @public */51export type IndexDirection =52  | -153  | 154  | '2d'55  | '2dsphere'56  | 'text'57  | 'geoHaystack'58  | 'hashed'59  | number;60 61function isIndexDirection(x: unknown): x is IndexDirection {62  return (63    typeof x === 'number' || x === '2d' || x === '2dsphere' || x === 'text' || x === 'geoHaystack'64  );65}66/** @public */67export type IndexSpecification = OneOrMore<68  | string69  | [string, IndexDirection]70  | { [key: string]: IndexDirection }71  | Map<string, IndexDirection>72>;73 74/** @public */75export interface IndexInformationOptions extends ListIndexesOptions {76  /**77   * When `true`, an array of index descriptions is returned.78   * When `false`, the driver returns an object that with keys corresponding to index names with values79   * corresponding to the entries of the indexes' key.80   *81   * For example, the given the following indexes:82   * ```83   * [ { name: 'a_1', key: { a: 1 } }, { name: 'b_1_c_1' , key: { b: 1, c: 1 } }]84   * ```85   *86   * When `full` is `true`, the above array is returned.  When `full` is `false`, the following is returned:87   * ```88   * {89   *   'a_1': [['a', 1]],90   *   'b_1_c_1': [['b', 1], ['c', 1]],91   * }92   * ```93   */94  full?: boolean;95}96 97/** @public */98export interface IndexDescription99  extends Pick<100    CreateIndexesOptions,101    | 'background'102    | 'unique'103    | 'partialFilterExpression'104    | 'sparse'105    | 'hidden'106    | 'expireAfterSeconds'107    | 'storageEngine'108    | 'version'109    | 'weights'110    | 'default_language'111    | 'language_override'112    | 'textIndexVersion'113    | '2dsphereIndexVersion'114    | 'bits'115    | 'min'116    | 'max'117    | 'bucketSize'118    | 'wildcardProjection'119  > {120  collation?: CollationOptions;121  name?: string;122  key: { [key: string]: IndexDirection } | Map<string, IndexDirection>;123}124 125/** @public */126export interface CreateIndexesOptions extends Omit<CommandOperationOptions, 'writeConcern'> {127  /** Creates the index in the background, yielding whenever possible. */128  background?: boolean;129  /** Creates an unique index. */130  unique?: boolean;131  /** Override the autogenerated index name (useful if the resulting name is larger than 128 bytes) */132  name?: string;133  /** Creates a partial index based on the given filter object (MongoDB 3.2 or higher) */134  partialFilterExpression?: Document;135  /** Creates a sparse index. */136  sparse?: boolean;137  /** Allows you to expire data on indexes applied to a data (MongoDB 2.2 or higher) */138  expireAfterSeconds?: number;139  /** Allows users to configure the storage engine on a per-index basis when creating an index. (MongoDB 3.0 or higher) */140  storageEngine?: Document;141  /** (MongoDB 4.4. or higher) Specifies how many data-bearing members of a replica set, including the primary, must complete the index builds successfully before the primary marks the indexes as ready. This option accepts the same values for the "w" field in a write concern plus "votingMembers", which indicates all voting data-bearing nodes. */142  commitQuorum?: number | string;143  /** Specifies the index version number, either 0 or 1. */144  version?: number;145  // text indexes146  weights?: Document;147  default_language?: string;148  language_override?: string;149  textIndexVersion?: number;150  // 2d-sphere indexes151  '2dsphereIndexVersion'?: number;152  // 2d indexes153  bits?: number;154  /** For geospatial indexes set the lower bound for the co-ordinates. */155  min?: number;156  /** For geospatial indexes set the high bound for the co-ordinates. */157  max?: number;158  // geoHaystack Indexes159  bucketSize?: number;160  // wildcard indexes161  wildcardProjection?: Document;162  /** Specifies that the index should exist on the target collection but should not be used by the query planner when executing operations. (MongoDB 4.4 or higher) */163  hidden?: boolean;164}165 166function isSingleIndexTuple(t: unknown): t is [string, IndexDirection] {167  return Array.isArray(t) && t.length === 2 && isIndexDirection(t[1]);168}169 170/**171 * Converts an `IndexSpecification`, which can be specified in multiple formats, into a172 * valid `key` for the createIndexes command.173 */174function constructIndexDescriptionMap(indexSpec: IndexSpecification): Map<string, IndexDirection> {175  const key: Map<string, IndexDirection> = new Map();176 177  const indexSpecs =178    !Array.isArray(indexSpec) || isSingleIndexTuple(indexSpec) ? [indexSpec] : indexSpec;179 180  // Iterate through array and handle different types181  for (const spec of indexSpecs) {182    if (typeof spec === 'string') {183      key.set(spec, 1);184    } else if (Array.isArray(spec)) {185      key.set(spec[0], spec[1] ?? 1);186    } else if (spec instanceof Map) {187      for (const [property, value] of spec) {188        key.set(property, value);189      }190    } else if (isObject(spec)) {191      for (const [property, value] of Object.entries(spec)) {192        key.set(property, value);193      }194    }195  }196 197  return key;198}199 200/**201 * Receives an index description and returns a modified index description which has had invalid options removed202 * from the description and has mapped the `version` option to the `v` option.203 */204function resolveIndexDescription(205  description: IndexDescription206): Omit<ResolvedIndexDescription, 'key'> {207  const validProvidedOptions = Object.entries(description).filter(([optionName]) =>208    VALID_INDEX_OPTIONS.has(optionName)209  );210 211  return Object.fromEntries(212    // we support the `version` option, but the `createIndexes` command expects it to be the `v`213    validProvidedOptions.map(([name, value]) => (name === 'version' ? ['v', value] : [name, value]))214  );215}216 217/**218 * @public219 * The index information returned by the listIndexes command. https://www.mongodb.com/docs/manual/reference/command/listIndexes/#mongodb-dbcommand-dbcmd.listIndexes220 */221export type IndexDescriptionInfo = Omit<IndexDescription, 'key' | 'version'> & {222  key: { [key: string]: IndexDirection };223  v?: IndexDescription['version'];224} & Document;225 226/** @public */227export type IndexDescriptionCompact = Record<string, [name: string, direction: IndexDirection][]>;228 229/**230 * @internal231 *232 * Internally, the driver represents index description keys with `Map`s to preserve key ordering.233 * We don't require users to specify maps, so we transform user provided descriptions into234 * "resolved" by converting the `key` into a JS `Map`, if it isn't already a map.235 *236 * Additionally, we support the `version` option, but the `createIndexes` command uses the field `v`237 * to specify the index version so we map the value of `version` to `v`, if provided.238 */239type ResolvedIndexDescription = Omit<IndexDescription, 'key' | 'version'> & {240  key: Map<string, IndexDirection>;241  v?: IndexDescription['version'];242};243 244/** @internal */245export class CreateIndexesOperation extends CommandOperation<string[]> {246  override SERVER_COMMAND_RESPONSE_TYPE = MongoDBResponse;247  override options: CreateIndexesOptions;248  collectionName: string;249  indexes: ReadonlyArray<ResolvedIndexDescription>;250 251  private constructor(252    parent: OperationParent,253    collectionName: string,254    indexes: IndexDescription[],255    options?: CreateIndexesOptions256  ) {257    super(parent, options);258 259    this.options = options ?? {};260    // collation is set on each index, it should not be defined at the root261    this.options.collation = undefined;262    this.collectionName = collectionName;263    this.indexes = indexes.map((userIndex: IndexDescription): ResolvedIndexDescription => {264      // Ensure the key is a Map to preserve index key ordering265      const key =266        userIndex.key instanceof Map ? userIndex.key : new Map(Object.entries(userIndex.key));267      const name = userIndex.name ?? Array.from(key).flat().join('_');268      const validIndexOptions = resolveIndexDescription(userIndex);269      return {270        ...validIndexOptions,271        name,272        key273      };274    });275    this.ns = parent.s.namespace;276  }277 278  static fromIndexDescriptionArray(279    parent: OperationParent,280    collectionName: string,281    indexes: IndexDescription[],282    options?: CreateIndexesOptions283  ): CreateIndexesOperation {284    return new CreateIndexesOperation(parent, collectionName, indexes, options);285  }286 287  static fromIndexSpecification(288    parent: OperationParent,289    collectionName: string,290    indexSpec: IndexSpecification,291    options: CreateIndexesOptions = {}292  ): CreateIndexesOperation {293    const key = constructIndexDescriptionMap(indexSpec);294    const description: IndexDescription = { ...options, key };295    return new CreateIndexesOperation(parent, collectionName, [description], options);296  }297 298  override get commandName() {299    return 'createIndexes';300  }301 302  override buildCommandDocument(connection: Connection): Document {303    const options = this.options;304    const indexes = this.indexes;305 306    const serverWireVersion = maxWireVersion(connection);307 308    const cmd: Document = { createIndexes: this.collectionName, indexes };309 310    if (options.commitQuorum != null) {311      if (serverWireVersion < 9) {312        throw new MongoCompatibilityError(313          'Option `commitQuorum` for `createIndexes` not supported on servers < 4.4'314        );315      }316      cmd.commitQuorum = options.commitQuorum;317    }318    return cmd;319  }320 321  override handleOk(_response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>): string[] {322    const indexNames = this.indexes.map(index => index.name || '');323    return indexNames;324  }325}326 327/** @public */328export type DropIndexesOptions = CommandOperationOptions;329 330/** @internal */331export class DropIndexOperation extends CommandOperation<Document> {332  override SERVER_COMMAND_RESPONSE_TYPE = MongoDBResponse;333  override options: DropIndexesOptions;334  collection: Collection;335  indexName: string;336 337  constructor(collection: Collection, indexName: string, options?: DropIndexesOptions) {338    super(collection, options);339 340    this.options = options ?? {};341    this.collection = collection;342    this.indexName = indexName;343    this.ns = collection.fullNamespace;344  }345 346  override get commandName() {347    return 'dropIndexes' as const;348  }349 350  override buildCommandDocument(_connection: Connection): Document {351    return { dropIndexes: this.collection.collectionName, index: this.indexName };352  }353}354 355/** @public */356export type ListIndexesOptions = AbstractCursorOptions & {357  /** @internal */358  omitMaxTimeMS?: boolean;359  /** @internal */360  rawData?: boolean;361};362 363/** @internal */364export class ListIndexesOperation extends CommandOperation<CursorResponse> {365  override SERVER_COMMAND_RESPONSE_TYPE = CursorResponse;366  /**367   * @remarks WriteConcern can still be present on the options because368   * we inherit options from the client/db/collection.  The369   * key must be present on the options in order to delete it.370   * This allows typescript to delete the key but will371   * not allow a writeConcern to be assigned as a property on options.372   */373  override options: ListIndexesOptions & { writeConcern?: never };374  collectionNamespace: MongoDBNamespace;375 376  constructor(collection: Collection, options?: ListIndexesOptions) {377    super(collection, options);378 379    this.options = { ...options };380    delete this.options.writeConcern;381    this.collectionNamespace = collection.s.namespace;382  }383 384  override get commandName() {385    return 'listIndexes' as const;386  }387 388  override buildCommandDocument(connection: Connection): Document {389    const serverWireVersion = maxWireVersion(connection);390 391    const cursor = this.options.batchSize ? { batchSize: this.options.batchSize } : {};392 393    const command: Document = { listIndexes: this.collectionNamespace.collection, cursor };394 395    // we check for undefined specifically here to allow falsy values396    // eslint-disable-next-line no-restricted-syntax397    if (serverWireVersion >= 9 && this.options.comment !== undefined) {398      command.comment = this.options.comment;399    }400 401    return command;402  }403 404  override handleOk(405    response: InstanceType<typeof this.SERVER_COMMAND_RESPONSE_TYPE>406  ): CursorResponse {407    return response;408  }409}410 411defineAspects(ListIndexesOperation, [412  Aspect.READ_OPERATION,413  Aspect.RETRYABLE,414  Aspect.CURSOR_CREATING,415  Aspect.SUPPORTS_RAW_DATA416]);417defineAspects(CreateIndexesOperation, [Aspect.WRITE_OPERATION, Aspect.SUPPORTS_RAW_DATA]);418defineAspects(DropIndexOperation, [Aspect.WRITE_OPERATION, Aspect.SUPPORTS_RAW_DATA]);419