feat: Vector Search

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,74 @@ 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;
 }
 
 /**
@@ -199,7 +221,7 @@ export function valueFromJson(fieldValue: api.IValue): api.IValue {
         },
       };
     }
-    case 'mapValue': {
+    case 'mapValue':{
       const mapValue: ApiMapValue = {};
       const fields = fieldValue.mapValue!.fields;
       if (fields) {

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,23 @@ 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
@@ -264,8 +284,13 @@ export function compare(left: api.IValue, right: api.IValue): number {
       );
     case TypeOrder.OBJECT:
       return compareObjects(
-        left.mapValue!.fields || {},
-        right.mapValue!.fields || {}
+          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}`);

dev/src/serializer.ts

@@ -18,17 +18,18 @@ import {DocumentData} from '@google-cloud/firestore';
 
 import * as proto from '../protos/firestore_v1_proto_api';
 
-import {detectValueType} from './convert';
 import {DeleteTransform, FieldTransform, VectorValue} from './field-value';
 import {GeoPoint} from './geo-point';
 import {DocumentReference, Firestore} from './index';
 import {FieldPath, QualifiedResourcePath} from './path';
 import {Timestamp} from './timestamp';
-import {ApiMapValue, ValidationOptions} from './types';
+import {ApiMapValue, ProtobufJsValue, ValidationOptions} from './types';
 import {isEmpty, isObject, isPlainObject} from './util';
 import {customObjectMessage, invalidArgumentMessage} from './validate';
 
 import api = proto.google.firestore.v1;
+import {detectValueType} from "./convert";
+import {RESERVED_MAP_KEY, RESERVED_MAP_KEY_VECTOR_VALUE, VECTOR_MAP_VECTORS_KEY} from "./map-type";
 
 /**
  * The maximum depth of a Firestore object.
@@ -38,10 +39,6 @@ import api = proto.google.firestore.v1;
  */
 const MAX_DEPTH = 20;
 
-const RESERVED_MAP_KEY = '__type__';
-const RESERVED_MAP_KEY_VECTOR_VALUE = '__vector__';
-const VECTOR_MAP_VECTORS_KEY = 'value';
-
 /**
  * An interface for Firestore types that can be serialized to Protobuf.
  *
@@ -298,24 +295,19 @@ export class Serializer {
       case 'mapValue': {
         const fields = proto.mapValue!.fields;
         if (fields) {
-          const props = Object.keys(fields);
-          if (
-            props.indexOf(RESERVED_MAP_KEY) !== -1 &&
-            this.decodeValue(fields[RESERVED_MAP_KEY]) ===
-              RESERVED_MAP_KEY_VECTOR_VALUE
-          ) {
-            return VectorValue._fromProto(fields[VECTOR_MAP_VECTORS_KEY]);
-          } else {
-            const obj: DocumentData = {};
-            for (const prop of Object.keys(fields)) {
-              obj[prop] = this.decodeValue(fields[prop]);
-            }
-            return obj;
+          const obj: DocumentData = {};
+          for (const prop of Object.keys(fields)) {
+            obj[prop] = this.decodeValue(fields[prop]);
           }
+          return obj;
         } else {
           return {};
         }
       }
+      case 'vectorValue': {
+        const fields = proto.mapValue!.fields!;
+        return VectorValue._fromProto(fields[VECTOR_MAP_VECTORS_KEY]);
+      }
       case 'geoPointValue': {
         return GeoPoint.fromProto(proto.geoPointValue!);
       }

dev/system-test/firestore.ts

@@ -3005,6 +3005,44 @@ describe('Query class', () => {
 
       unsubscribe();
     });
+
+    it.only('orders vector field correctly', async () => {
+      await randomCol.add({"embedding": {"HELLO": "WORLD"}});
+      await randomCol.add({"embedding": {"hello": "world"}});
+      await randomCol.add({"embedding": FieldValue.vector([1, 2, 3])});
+      await randomCol.add({"embedding": FieldValue.vector([1, 2])});
+      await randomCol.add({"embedding": FieldValue.vector([2, 2])});
+      await randomCol.add({"embedding": FieldValue.vector([100, 2, 3, 4, 5])});
+      await randomCol.add({"embedding": FieldValue.vector([1, 2, 3, 4, 5])});
+      await randomCol.add({"embedding": FieldValue.vector([1, 2, 100, 4, 5])});
+      await randomCol.add({"embedding": FieldValue.vector([1, 2, 3, 4])});
+
+      const orderedQuery = randomCol.orderBy("embedding");
+
+      const unsubscribe = orderedQuery.onSnapshot(
+          snapshot => {
+            currentDeferred.resolve(snapshot);
+          },
+          err => {
+            currentDeferred.reject!(err);
+          }
+      );
+
+      const watchSnapshot = await waitForSnapshot();
+      unsubscribe();
+      const getSnapshot = await orderedQuery.get();
+
+      console.log("---- watch snapshot ----")
+      watchSnapshot.docs.forEach(ds => console.log(ds.get("embedding")));
+
+      console.log("---- get snapshot ----")
+      getSnapshot.docs.forEach(ds => console.log(ds.get("embedding")));
+
+      snapshotsEqual(watchSnapshot, {
+        docs: getSnapshot.docs,
+        docChanges: getSnapshot.docChanges()
+      });
+    });
   });
 
   (process.env.FIRESTORE_EMULATOR_HOST === undefined