feat: Vector Search

dev/src/aggregate.ts

@@ -1,4 +1,4 @@
-/*!
+/**
  * Copyright 2023 Google LLC. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");

dev/src/collection-group.ts

@@ -93,7 +93,7 @@ export class CollectionGroup<
       // Partition queries require explicit ordering by __name__.
       const queryWithDefaultOrder = this.orderBy(FieldPath.documentId());
       const request: api.IPartitionQueryRequest =
-        queryWithDefaultOrder.toProto();
+        queryWithDefaultOrder._toProto();
 
       // Since we are always returning an extra partition (with an empty endBefore
       // cursor), we reduce the desired partition count by one.

dev/src/convert.ts

@@ -20,6 +20,7 @@ import {ApiMapValue, ProtobufJsValue} from './types';
 import {validateObject} from './validate';
 
 import api = google.firestore.v1;
+import {RESERVED_MAP_KEY, RESERVED_MAP_KEY_VECTOR_VALUE} from './map-type';
 
 /*!
  * @module firestore/convert
@@ -112,53 +113,72 @@ function bytesFromJson(bytesValue: string | Uint8Array): Uint8Array {
  * @return The string value for 'valueType'.
  */
 export function detectValueType(proto: ProtobufJsValue): string {
+  let valueType: string | undefined;
+
   if (proto.valueType) {
-    return proto.valueType;
-  }
+    valueType = proto.valueType;
+  } else {
+    const detectedValues: string[] = [];
 
-  const detectedValues: string[] = [];
+    if (proto.stringValue !== undefined) {
+      detectedValues.push('stringValue');
+    }
+    if (proto.booleanValue !== undefined) {
+      detectedValues.push('booleanValue');
+    }
+    if (proto.integerValue !== undefined) {
+      detectedValues.push('integerValue');
+    }
+    if (proto.doubleValue !== undefined) {
+      detectedValues.push('doubleValue');
+    }
+    if (proto.timestampValue !== undefined) {
+      detectedValues.push('timestampValue');
+    }
+    if (proto.referenceValue !== undefined) {
+      detectedValues.push('referenceValue');
+    }
+    if (proto.arrayValue !== undefined) {
+      detectedValues.push('arrayValue');
+    }
+    if (proto.nullValue !== undefined) {
+      detectedValues.push('nullValue');
+    }
+    if (proto.mapValue !== undefined) {
+      detectedValues.push('mapValue');
+    }
+    if (proto.geoPointValue !== undefined) {
+      detectedValues.push('geoPointValue');
+    }
+    if (proto.bytesValue !== undefined) {
+      detectedValues.push('bytesValue');
+    }
 
-  if (proto.stringValue !== undefined) {
-    detectedValues.push('stringValue');
-  }
-  if (proto.booleanValue !== undefined) {
-    detectedValues.push('booleanValue');
-  }
-  if (proto.integerValue !== undefined) {
-    detectedValues.push('integerValue');
-  }
-  if (proto.doubleValue !== undefined) {
-    detectedValues.push('doubleValue');
-  }
-  if (proto.timestampValue !== undefined) {
-    detectedValues.push('timestampValue');
-  }
-  if (proto.referenceValue !== undefined) {
-    detectedValues.push('referenceValue');
-  }
-  if (proto.arrayValue !== undefined) {
-    detectedValues.push('arrayValue');
-  }
-  if (proto.nullValue !== undefined) {
-    detectedValues.push('nullValue');
-  }
-  if (proto.mapValue !== undefined) {
-    detectedValues.push('mapValue');
-  }
-  if (proto.geoPointValue !== undefined) {
-    detectedValues.push('geoPointValue');
-  }
-  if (proto.bytesValue !== undefined) {
-    detectedValues.push('bytesValue');
+    if (detectedValues.length !== 1) {
+      throw new Error(
+        `Unable to infer type value from '${JSON.stringify(proto)}'.`
+      );
+    }
+
+    valueType = detectedValues[0];
   }
 
-  if (detectedValues.length !== 1) {
-    throw new Error(
-      `Unable to infer type value from '${JSON.stringify(proto)}'.`
-    );
+  // Special handling of mapValues used to represent other data types
+  if (valueType === 'mapValue') {
+    const fields = proto.mapValue?.fields;
+    if (fields) {
+      const props = Object.keys(fields);
+      if (
+        props.indexOf(RESERVED_MAP_KEY) !== -1 &&
+        detectValueType(fields[RESERVED_MAP_KEY]) === 'stringValue' &&
+        fields[RESERVED_MAP_KEY].stringValue === RESERVED_MAP_KEY_VECTOR_VALUE
+      ) {
+        valueType = 'vectorValue';
+      }
+    }
   }
 
-  return detectedValues[0];
+  return valueType;
 }
 
 /**
@@ -240,7 +260,8 @@ export function valueFromJson(fieldValue: api.IValue): api.IValue {
         },
       };
     }
-    case 'mapValue': {
+    case 'mapValue':
+    case 'vectorValue': {
       const mapValue: ApiMapValue = {};
       const fields = fieldValue.mapValue!.fields;
       if (fields) {

dev/src/field-value.ts

@@ -22,6 +22,7 @@ import * as proto from '../protos/firestore_v1_proto_api';
 
 import {FieldPath} from './path';
 import {Serializer, validateUserInput} from './serializer';
+import {isPrimitiveArrayEqual} from './util';
 import {
   invalidArgumentMessage,
   validateMinNumberOfArguments,
@@ -30,6 +31,58 @@ import {
 
 import api = proto.google.firestore.v1;
 
+/**
+ * Represent a vector type in Firestore documents.
+ * Create an instance with {@link FieldValue.vector}.
+ *
+ * @class VectorValue
+ */
+export class VectorValue implements firestore.VectorValue {
+  private readonly _values: number[];
+
+  /**
+   * @private
+   * @internal
+   */
+  constructor(values: number[] | undefined) {
+    // Making a copy of the parameter.
+    this._values = (values || []).map(n => n);
+  }
+
+  /**
+   * Returns a copy of the raw number array form of the vector.
+   */
+  public toArray(): number[] {
+    return this._values.map(n => n);
+  }
+
+  /**
+   * @private
+   * @internal
+   */
+  _toProto(serializer: Serializer): api.IValue {
+    return serializer.encodeVector(this._values);
+  }
+
+  /**
+   * @private
+   * @internal
+   */
+  static _fromProto(valueArray: api.IValue): VectorValue {
+    const values = valueArray.arrayValue?.values?.map(v => {
+      return v.doubleValue!;
+    });
+    return new VectorValue(values);
+  }
+
+  /**
+   * Returns `true` if the two VectorValue has the same raw number arrays, returns `false` otherwise.
+   */
+  isEqual(other: VectorValue): boolean {
+    return isPrimitiveArrayEqual(this._values, other._values);
+  }
+}
+
 /**
  * Sentinel values that can be used when writing documents with set(), create()
  * or update().
@@ -40,6 +93,17 @@ export class FieldValue implements firestore.FieldValue {
   /** @private */
   constructor() {}
 
+  /**
+   * Creates a new `VectorValue` constructed with a copy of the given array of numbers.
+   *
+   * @param values - Create a `VectorValue` instance with a copy of this array of numbers.
+   *
+   * @returns A new `VectorValue` constructed with a copy of the given array of numbers.
+   */
+  static vector(values?: number[]): VectorValue {
+    return new VectorValue(values);
+  }
+
   /**
    * Returns a sentinel for use with update() or set() with {merge:true} to mark
    * a field for deletion.
@@ -220,7 +284,7 @@ export class FieldValue implements firestore.FieldValue {
  * An internal interface shared by all field transforms.
  *
  * A 'FieldTransform` subclass should implement '.includeInDocumentMask',
- * '.includeInDocumentTransform' and 'toProto' (if '.includeInDocumentTransform'
+ * '.includeInDocumentTransform' and '_toProto' (if '.includeInDocumentTransform'
  * is 'true').
  *
  * @private

dev/src/index.ts

@@ -91,12 +91,17 @@ export {
   QuerySnapshot,
   Query,
 } from './reference';
-export type {AggregateQuery, AggregateQuerySnapshot} from './reference';
+export type {
+  AggregateQuery,
+  AggregateQuerySnapshot,
+  VectorQuery,
+  VectorQuerySnapshot,
+} from './reference';
 export {BulkWriter} from './bulk-writer';
 export type {BulkWriterError} from './bulk-writer';
 export type {BundleBuilder} from './bundle';
 export {DocumentSnapshot, QueryDocumentSnapshot} from './document';
-export {FieldValue} from './field-value';
+export {FieldValue, VectorValue} from './field-value';
 export {Filter} from './filter';
 export {WriteBatch, WriteResult} from './write-batch';
 export {Transaction} from './transaction';

dev/src/map-type.ts

@@ -0,0 +1,19 @@
+/*!
+ * Copyright 2024 Google LLC. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export const RESERVED_MAP_KEY = '__type__';
+export const RESERVED_MAP_KEY_VECTOR_VALUE = '__vector__';
+export const VECTOR_MAP_VECTORS_KEY = 'value';

dev/src/order.ts

@@ -34,7 +34,8 @@ enum TypeOrder {
   REF = 6,
   GEO_POINT = 7,
   ARRAY = 8,
-  OBJECT = 9,
+  VECTOR = 9,
+  OBJECT = 10,
 }
 
 /*!
@@ -67,6 +68,8 @@ function typeOrder(val: api.IValue): TypeOrder {
       return TypeOrder.REF;
     case 'mapValue':
       return TypeOrder.OBJECT;
+    case 'vectorValue':
+      return TypeOrder.VECTOR;
     default:
       throw new Error('Unexpected value type: ' + valueType);
   }
@@ -225,6 +228,26 @@ function compareObjects(left: ApiMapValue, right: ApiMapValue): number {
   return primitiveComparator(leftKeys.length, rightKeys.length);
 }
 
+/*!
+ * @private
+ * @internal
+ */
+function compareVectors(left: ApiMapValue, right: ApiMapValue): number {
+  // The vector is a map, but only vector value is compared.
+  const leftArray = left?.['value']?.arrayValue?.values ?? [];
+  const rightArray = right?.['value']?.arrayValue?.values ?? [];
+
+  const lengthCompare = primitiveComparator(
+    leftArray.length,
+    rightArray.length
+  );
+  if (lengthCompare !== 0) {
+    return lengthCompare;
+  }
+
+  return compareArrays(leftArray, rightArray);
+}
+
 /*!
  * @private
  * @internal
@@ -267,6 +290,11 @@ export function compare(left: api.IValue, right: api.IValue): number {
         left.mapValue!.fields || {},
         right.mapValue!.fields || {}
       );
+    case TypeOrder.VECTOR:
+      return compareVectors(
+        left.mapValue!.fields || {},
+        right.mapValue!.fields || {}
+      );
     default:
       throw new Error(`Encountered unknown type order: ${leftType}`);
   }