Implement indexing for bson types (#331)

Author: milaGGL

packages/firestore/src/core/target.ts

@@ -28,6 +28,8 @@ import {
   INTERNAL_MAX_VALUE,
   INTERNAL_MIN_VALUE,
   lowerBoundCompare,
+  MAX_KEY_VALUE,
+  MIN_KEY_VALUE,
   upperBoundCompare,
   valuesGetLowerBound,
   valuesGetUpperBound
@@ -387,7 +389,7 @@ function targetGetAscendingBound(
         break;
       case Operator.NOT_EQUAL:
       case Operator.NOT_IN:
-        filterValue = INTERNAL_MIN_VALUE;
+        filterValue = MIN_KEY_VALUE;
         break;
       default:
       // Remaining filters cannot be used as lower bounds.
@@ -462,7 +464,7 @@ function targetGetDescendingBound(
         break;
       case Operator.NOT_EQUAL:
       case Operator.NOT_IN:
-        filterValue = INTERNAL_MAX_VALUE;
+        filterValue = MAX_KEY_VALUE;
         break;
       default:
       // Remaining filters cannot be used as upper bounds.

packages/firestore/src/index/firestore_index_value_writer.ts

@@ -22,32 +22,45 @@ import {
   normalizeTimestamp
 } from '../model/normalize';
 import {
-  isVectorValue,
   VECTOR_MAP_VECTORS_KEY,
-  isMaxValue
+  detectSpecialMapType,
+  RESERVED_BSON_TIMESTAMP_KEY,
+  RESERVED_REGEX_KEY,
+  RESERVED_BSON_OBJECT_ID_KEY,
+  RESERVED_BSON_BINARY_KEY,
+  SpecialMapValueType,
+  RESERVED_REGEX_PATTERN_KEY,
+  RESERVED_REGEX_OPTIONS_KEY,
+  RESERVED_INT32_KEY
 } from '../model/values';
 import { ArrayValue, MapValue, Value } from '../protos/firestore_proto_api';
 import { fail } from '../util/assert';
 import { isNegativeZero } from '../util/types';
 
 import { DirectionalIndexByteEncoder } from './directional_index_byte_encoder';
 
-// Note: This code is copied from the backend. Code that is not used by
-// Firestore was removed.
+// Note: This file is copied from the backend. Code that is not used by
+// Firestore was removed. Code that has different behavior was modified.
 
 const INDEX_TYPE_NULL = 5;
+const INDEX_TYPE_MIN_KEY = 7;
 const INDEX_TYPE_BOOLEAN = 10;
 const INDEX_TYPE_NAN = 13;
 const INDEX_TYPE_NUMBER = 15;
 const INDEX_TYPE_TIMESTAMP = 20;
+const INDEX_TYPE_BSON_TIMESTAMP = 22;
 const INDEX_TYPE_STRING = 25;
 const INDEX_TYPE_BLOB = 30;
+const INDEX_TYPE_BSON_BINARY = 31;
 const INDEX_TYPE_REFERENCE = 37;
+const INDEX_TYPE_BSON_OBJECT_ID = 43;
 const INDEX_TYPE_GEOPOINT = 45;
+const INDEX_TYPE_REGEX = 47;
 const INDEX_TYPE_ARRAY = 50;
 const INDEX_TYPE_VECTOR = 53;
 const INDEX_TYPE_MAP = 55;
 const INDEX_TYPE_REFERENCE_SEGMENT = 60;
+const INDEX_TYPE_MAX_VALUE = 999;
 
 // A terminator that indicates that a truncatable value was not truncated.
 // This must be smaller than all other type labels.
@@ -124,11 +137,30 @@ export class FirestoreIndexValueWriter {
       encoder.writeNumber(geoPoint.latitude || 0);
       encoder.writeNumber(geoPoint.longitude || 0);
     } else if ('mapValue' in indexValue) {
-      // TODO(Mila/BSON): add bson types for indexing
-      if (isMaxValue(indexValue)) {
+      const type = detectSpecialMapType(indexValue);
+      if (type === SpecialMapValueType.INTERNAL_MAX) {
         this.writeValueTypeLabel(encoder, Number.MAX_SAFE_INTEGER);
-      } else if (isVectorValue(indexValue)) {
+      } else if (type === SpecialMapValueType.VECTOR) {
         this.writeIndexVector(indexValue.mapValue!, encoder);
+      } else if (type === SpecialMapValueType.MAX_KEY) {
+        this.writeValueTypeLabel(encoder, INDEX_TYPE_MAX_VALUE);
+      } else if (type === SpecialMapValueType.MIN_KEY) {
+        this.writeValueTypeLabel(encoder, INDEX_TYPE_MIN_KEY);
+      } else if (type === SpecialMapValueType.BSON_BINARY) {
+        this.writeIndexBsonBinaryData(indexValue.mapValue!, encoder);
+      } else if (type === SpecialMapValueType.REGEX) {
+        this.writeIndexRegex(indexValue.mapValue!, encoder);
+      } else if (type === SpecialMapValueType.BSON_TIMESTAMP) {
+        this.writeIndexBsonTimestamp(indexValue.mapValue!, encoder);
+      } else if (type === SpecialMapValueType.BSON_OBJECT_ID) {
+        this.writeIndexBsonObjectId(indexValue.mapValue!, encoder);
+      } else if (type === SpecialMapValueType.INT32) {
+        this.writeValueTypeLabel(encoder, INDEX_TYPE_NUMBER);
+        encoder.writeNumber(
+          normalizeNumber(
+            indexValue.mapValue!.fields![RESERVED_INT32_KEY]!.integerValue!
+          )
+        );
       } else {
         this.writeIndexMap(indexValue.mapValue!, encoder);
         this.writeTruncationMarker(encoder);
@@ -202,7 +234,10 @@ export class FirestoreIndexValueWriter {
     encoder: DirectionalIndexByteEncoder
   ): void {
     this.writeValueTypeLabel(encoder, INDEX_TYPE_REFERENCE);
-    const path = DocumentKey.fromName(referenceValue).path;
+    const segments: string[] = referenceValue
+      .split('/')
+      .filter(segment => segment.length > 0);
+    const path = DocumentKey.fromSegments(segments.slice(5)).path;
     path.forEach(segment => {
       this.writeValueTypeLabel(encoder, INDEX_TYPE_REFERENCE_SEGMENT);
       this.writeUnlabeledIndexString(segment, encoder);
@@ -222,4 +257,55 @@ export class FirestoreIndexValueWriter {
     // references, arrays and maps).
     encoder.writeNumber(NOT_TRUNCATED);
   }
+
+  private writeIndexBsonTimestamp(
+    mapValue: MapValue,
+    encoder: DirectionalIndexByteEncoder
+  ): void {
+    this.writeValueTypeLabel(encoder, INDEX_TYPE_BSON_TIMESTAMP);
+    const fields = mapValue.fields || {};
+    if (fields) {
+      // The JS SDK encodes BSON timestamps differently than the backend.
+      // This is due to the limitation of `number` in JS which handles up to 53-bit precision.
+      this.writeIndexMap(
+        fields[RESERVED_BSON_TIMESTAMP_KEY].mapValue!,
+        encoder
+      );
+    }
+  }
+
+  private writeIndexBsonObjectId(
+    mapValue: MapValue,
+    encoder: DirectionalIndexByteEncoder
+  ): void {
+    this.writeValueTypeLabel(encoder, INDEX_TYPE_BSON_OBJECT_ID);
+    const fields = mapValue.fields || {};
+    const oid = fields[RESERVED_BSON_OBJECT_ID_KEY]?.stringValue || '';
+    encoder.writeBytes(normalizeByteString(oid));
+  }
+
+  private writeIndexBsonBinaryData(
+    mapValue: MapValue,
+    encoder: DirectionalIndexByteEncoder
+  ): void {
+    this.writeValueTypeLabel(encoder, INDEX_TYPE_BSON_BINARY);
+    const fields = mapValue.fields || {};
+    const binary = fields[RESERVED_BSON_BINARY_KEY]?.bytesValue || '';
+    encoder.writeBytes(normalizeByteString(binary));
+    this.writeTruncationMarker(encoder);
+  }
+
+  private writeIndexRegex(
+    mapValue: MapValue,
+    encoder: DirectionalIndexByteEncoder
+  ): void {
+    this.writeValueTypeLabel(encoder, INDEX_TYPE_REGEX);
+    const fields = mapValue.fields || {};
+    const regex = fields[RESERVED_REGEX_KEY]?.mapValue?.fields || {};
+    if (regex) {
+      encoder.writeString(regex[RESERVED_REGEX_PATTERN_KEY]?.stringValue || '');
+      encoder.writeString(regex[RESERVED_REGEX_OPTIONS_KEY]?.stringValue || '');
+    }
+    this.writeTruncationMarker(encoder);
+  }
 }

packages/firestore/src/lite-api/query.ts

@@ -811,6 +811,8 @@ export function newQueryFilter(
   value: unknown
 ): FieldFilter {
   let fieldValue: ProtoValue;
+  validateQueryOperator(value, op);
+
   if (fieldPath.isKeyField()) {
     if (op === Operator.ARRAY_CONTAINS || op === Operator.ARRAY_CONTAINS_ANY) {
       throw new FirestoreError(
@@ -1064,6 +1066,31 @@ function validateDisjunctiveFilterElements(
   }
 }
 
+/**
+ * Validates the input string as a field comparison operator.
+ */
+export function validateQueryOperator(
+  value: unknown,
+  operator: Operator
+): void {
+  if (
+    typeof value === 'number' &&
+    isNaN(value) &&
+    operator !== '==' &&
+    operator !== '!='
+  ) {
+    throw new Error(
+      "Invalid query. You can only perform '==' and '!=' comparisons on NaN."
+    );
+  }
+
+  if (value === null && operator !== '==' && operator !== '!=') {
+    throw new Error(
+      "Invalid query. You can only perform '==' and '!=' comparisons on Null."
+    );
+  }
+}
+
 /**
  * Given an operator, returns the set of operators that cannot be used with it.
  *

packages/firestore/src/lite-api/user_data_writer.ts

@@ -83,9 +83,6 @@ export abstract class AbstractUserDataWriter {
   ): unknown {
     switch (typeOrder(value)) {
       case TypeOrder.NullValue:
-        if ('mapValue' in value) {
-          return minKey();
-        }
         return null;
       case TypeOrder.BooleanValue:
         return value.booleanValue!;
@@ -122,6 +119,8 @@ export abstract class AbstractUserDataWriter {
         return this.convertToBsonTimestampValue(value.mapValue!);
       case TypeOrder.MaxKeyValue:
         return maxKey();
+      case TypeOrder.MinKeyValue:
+        return minKey();
       default:
         throw fail('Invalid value type: ' + JSON.stringify(value));
     }

packages/firestore/src/model/type_order.ts

@@ -27,24 +27,25 @@ export const enum TypeOrder {
   // server timestamps and `MAX_VALUE` inside the SDK.
   // NULL and MIN_KEY sort the same.
   NullValue = 0,
-  MinKeyValue = 0,
-  BooleanValue = 1,
-  NumberValue = 2,
-  TimestampValue = 3,
-  // TODO(Mila/BSON): which should come first considering indexes?
-  BsonTimestampValue = 4,
-  ServerTimestampValue = 5,
-  StringValue = 6,
-  BlobValue = 7,
-  BsonBinaryValue = 8,
-  RefValue = 9,
-  BsonObjectIdValue = 10,
-  GeoPointValue = 11,
-  RegexValue = 12,
-  ArrayValue = 13,
-  VectorValue = 14,
-  ObjectValue = 15,
-  // TODO(Mila/BSON):should MaxKeyValue and MaxValue combined? how would this affect indexes?
-  MaxKeyValue = 16,
+  MinKeyValue = 1,
+  BooleanValue = 2,
+  // Note: all numbers (32-bit int, 64-bit int, 64-bit double, 128-bit decimal,
+  // etc.) are sorted together numerically. The `numberEquals` function
+  // distinguishes between different number types and compares them accordingly.
+  NumberValue = 3,
+  TimestampValue = 4,
+  BsonTimestampValue = 5,
+  ServerTimestampValue = 6,
+  StringValue = 7,
+  BlobValue = 8,
+  BsonBinaryValue = 9,
+  RefValue = 10,
+  BsonObjectIdValue = 11,
+  GeoPointValue = 12,
+  RegexValue = 13,
+  ArrayValue = 14,
+  VectorValue = 15,
+  ObjectValue = 16,
+  MaxKeyValue = 17,
   MaxValue = 9007199254740991 // Number.MAX_SAFE_INTEGER
 }