query_engine.test.ts pass with pipelines

wu-hui · wu-hui · commit 667c398fd625 · 2024-11-19T15:50:24.000-05:00
diff --git a/packages/firestore/src/core/expressions.ts b/packages/firestore/src/core/expressions.ts
@@ -79,7 +79,12 @@ import {
   Field,
   Constant
 } from '../lite-api/expressions';
-import { FieldPath } from '../model/path';
+import {
+  CREATE_TIME_NAME,
+  DOCUMENT_KEY_NAME,
+  FieldPath,
+  UPDATE_TIME_NAME
+} from '../model/path';
 import {
   FALSE_VALUE,
   getVectorValue,
@@ -99,6 +104,7 @@ import {
 } from '../model/values';
 
 import { RE2JS } from 're2js';
+import { toName, toTimestamp, toVersion } from '../remote/serializer';
 
 export interface EvaluableExpr {
   evaluate(
@@ -246,6 +252,27 @@ export class CoreField implements EvaluableExpr {
     context: EvaluationContext,
     input: PipelineInputOutput
   ): Value | undefined {
+    if (this.expr.fieldName() === DOCUMENT_KEY_NAME) {
+      return {
+        referenceValue: toName(context.userDataReader.serializer, input.key)
+      };
+    }
+    if (this.expr.fieldName() === UPDATE_TIME_NAME) {
+      return {
+        timestampValue: toVersion(
+          context.userDataReader.serializer,
+          input.version
+        )
+      };
+    }
+    if (this.expr.fieldName() === CREATE_TIME_NAME) {
+      return {
+        timestampValue: toVersion(
+          context.userDataReader.serializer,
+          input.createTime
+        )
+      };
+    }
     return (
       input.data.field(FieldPath.fromServerFormat(this.expr.fieldName())) ??
       undefined
@@ -936,17 +963,17 @@ export class CoreArrayContainsAny implements EvaluableExpr {
     context: EvaluationContext,
     input: PipelineInputOutput
   ): Value | undefined {
-    const evaluated = toEvaluable(this.expr.array).evaluate(context, input);
-    if (evaluated === undefined || !isArray(evaluated)) {
+    const evaluatedExpr = toEvaluable(this.expr.array).evaluate(context, input);
+    if (evaluatedExpr === undefined || !isArray(evaluatedExpr)) {
       return undefined;
     }
 
-    const elements = this.expr.values.map(val =>
+    const candidates = this.expr.values.map(val =>
       toEvaluable(val).evaluate(context, input)
     );
 
-    for (const element of elements) {
-      for (const val of evaluated.arrayValue.values ?? []) {
+    for (const element of candidates) {
+      for (const val of evaluatedExpr.arrayValue.values ?? []) {
         if (element !== undefined && valueEquals(val, element!)) {
           return TRUE_VALUE;
         }
diff --git a/packages/firestore/src/core/pipeline-util.ts b/packages/firestore/src/core/pipeline-util.ts
@@ -64,14 +64,26 @@ import {
   Where
 } from '../lite-api/stage';
 import { Pipeline } from '../lite-api/pipeline';
-import { canonifyQuery, Query, queryEquals, stringifyQuery } from './query';
+import {
+  canonifyQuery,
+  isCollectionGroupQuery,
+  isDocumentQuery,
+  LimitType,
+  Query,
+  queryEquals,
+  queryNormalizedOrderBy,
+  stringifyQuery
+} from './query';
 import {
   canonifyTarget,
   Target,
   targetEquals,
   targetIsPipelineTarget
 } from './target';
 import { ResourcePath } from '../model/path';
+import { Firestore } from '../api/database';
+import { doc } from '../lite-api/reference';
+import { Direction } from './order_by';
 
 /* eslint @typescript-eslint/no-explicit-any: 0 */
 
@@ -222,34 +234,37 @@ export function toPipelineFilterCondition(
       const value = f.value;
       switch (f.op) {
         case Operator.LESS_THAN:
-          return and(field.exists(), field.lt(value));
+          return and(field.exists(), field.lt(Constant._fromProto(value)));
         case Operator.LESS_THAN_OR_EQUAL:
-          return and(field.exists(), field.lte(value));
+          return and(field.exists(), field.lte(Constant._fromProto(value)));
         case Operator.GREATER_THAN:
-          return and(field.exists(), field.gt(value));
+          return and(field.exists(), field.gt(Constant._fromProto(value)));
         case Operator.GREATER_THAN_OR_EQUAL:
-          return and(field.exists(), field.gte(value));
+          return and(field.exists(), field.gte(Constant._fromProto(value)));
         case Operator.EQUAL:
-          return and(field.exists(), field.eq(value));
+          return and(field.exists(), field.eq(Constant._fromProto(value)));
         case Operator.NOT_EQUAL:
-          return and(field.exists(), field.neq(value));
+          return and(field.exists(), field.neq(Constant._fromProto(value)));
         case Operator.ARRAY_CONTAINS:
-          return and(field.exists(), field.arrayContains(value));
+          return and(
+            field.exists(),
+            field.arrayContains(Constant._fromProto(value))
+          );
         case Operator.IN: {
           const values = value?.arrayValue?.values?.map((val: any) =>
-            Constant.of(val)
+            Constant._fromProto(val)
           );
           return and(field.exists(), field.in(...values!));
         }
         case Operator.ARRAY_CONTAINS_ANY: {
           const values = value?.arrayValue?.values?.map((val: any) =>
-            Constant.of(val)
+            Constant._fromProto(val)
           );
-          return and(field.exists(), field.arrayContainsAny(values!));
+          return and(field.exists(), field.arrayContainsAny(...values!));
         }
         case Operator.NOT_IN: {
           const values = value?.arrayValue?.values?.map((val: any) =>
-            Constant.of(val)
+            Constant._fromProto(val)
           );
           return and(field.exists(), not(field.in(...values!)));
         }
@@ -279,6 +294,56 @@ export function toPipelineFilterCondition(
   throw new Error(`Failed to convert filter to pipeline conditions: ${f}`);
 }
 
+export function toPipeline(query: Query, db: Firestore): Pipeline {
+  let pipeline: Pipeline;
+  if (isCollectionGroupQuery(query)) {
+    pipeline = db.pipeline().collectionGroup(query.collectionGroup!);
+  } else if (isDocumentQuery(query)) {
+    pipeline = db.pipeline().documents([doc(db, query.path.canonicalString())]);
+  } else {
+    pipeline = db.pipeline().collection(query.path.canonicalString());
+  }
+
+  // filters
+  for (const filter of query.filters) {
+    pipeline = pipeline.where(toPipelineFilterCondition(filter));
+  }
+
+  // orders
+  const orders = queryNormalizedOrderBy(query);
+  const existsConditions = orders.map(order =>
+    Field.of(order.field.canonicalString()).exists()
+  );
+  if (existsConditions.length > 1) {
+    pipeline = pipeline.where(
+      and(existsConditions[0], ...existsConditions.slice(1))
+    );
+  } else {
+    pipeline = pipeline.where(existsConditions[0]);
+  }
+
+  pipeline = pipeline.sort(
+    ...orders.map(order =>
+      order.dir === Direction.ASCENDING
+        ? Field.of(order.field.canonicalString()).ascending()
+        : Field.of(order.field.canonicalString()).descending()
+    )
+  );
+
+  // cursors and limits
+  if (query.startAt !== null || query.endAt !== null) {
+    throw new Error('Cursors are not supported yet.');
+  }
+  if (query.limitType === LimitType.Last) {
+    throw new Error('Limit to last are not supported yet.');
+  }
+  if (query.limit !== null) {
+    pipeline = pipeline.limit(query.limit);
+  }
+
+  return pipeline;
+}
+
 function canonifyExpr(expr: Expr): string {
   if (expr instanceof Field) {
     return `fld(${expr.fieldName()})`;
@@ -534,6 +599,6 @@ export function targetOrPipelineEqual(
 
 export function pipelineHasRanges(pipeline: Pipeline): boolean {
   return pipeline.stages.some(
-    stage => stage.name === Limit.name || stage.name === Offset.name
+    stage => stage instanceof Limit || stage instanceof Offset
   );
 }
diff --git a/packages/firestore/src/core/pipeline_run.ts b/packages/firestore/src/core/pipeline_run.ts
@@ -17,6 +17,9 @@ import {
   CollectionSource,
   DatabaseSource,
   DocumentsSource,
+  Exists,
+  exists,
+  Field,
   Limit,
   Offset,
   Ordering,
@@ -36,6 +39,7 @@ import { toEvaluable } from './expressions';
 import { UserDataReader } from '../lite-api/user_data_reader';
 import { Query, queryMatches, queryMatchesAllDocuments } from './query';
 import { isPipeline, QueryOrPipeline } from './pipeline-util';
+import { DOCUMENT_KEY_NAME } from '../model/path';
 
 export type PipelineInputOutput = MutableDocument;
 
@@ -78,8 +82,23 @@ export function queryOrPipelineMatches(
 }
 
 export function pipelineMatchesAllDocuments(pipeline: Pipeline): boolean {
-  // TODO(pipeline): implement properly.
-  return false;
+  for (const stage of pipeline.stages) {
+    if (stage instanceof Limit || stage instanceof Offset) {
+      return false;
+    }
+    if (stage instanceof Where) {
+      if (
+        stage.condition instanceof Exists &&
+        stage.condition.expr instanceof Field &&
+        stage.condition.expr.fieldName() === DOCUMENT_KEY_NAME
+      ) {
+        continue;
+      }
+      return false;
+    }
+  }
+
+  return true;
 }
 
 function evaluate(
@@ -178,8 +197,9 @@ function evaluateCollection(
 ): Array<PipelineInputOutput> {
   return inputs.filter(input => {
     return (
+      input.isFoundDocument() &&
       `/${input.key.getCollectionPath().canonicalString()}` ===
-      coll.collectionPath
+        coll.collectionPath
     );
   });
 }
@@ -191,7 +211,10 @@ function evaluateCollectionGroup(
 ): Array<PipelineInputOutput> {
   // return those records in input whose collection id is stage.collectionId
   return input.filter(input => {
-    return input.key.getCollectionPath().lastSegment() === stage.collectionId;
+    return (
+      input.isFoundDocument() &&
+      input.key.getCollectionPath().lastSegment() === stage.collectionId
+    );
   });
 }
 
@@ -200,7 +223,7 @@ function evaluateDatabase(
   stage: DatabaseSource,
   input: Array<PipelineInputOutput>
 ): Array<PipelineInputOutput> {
-  return input;
+  return input.filter(input => input.isFoundDocument());
 }
 
 function evaluateDocuments(
@@ -209,7 +232,10 @@ function evaluateDocuments(
   input: Array<PipelineInputOutput>
 ): Array<PipelineInputOutput> {
   return input.filter(input => {
-    return stage.docPaths.includes(input.key.path.canonicalString());
+    return (
+      input.isFoundDocument() &&
+      stage.docPaths.includes(input.key.path.canonicalString())
+    );
   });
 }
 
diff --git a/packages/firestore/src/lite-api/expressions.ts b/packages/firestore/src/lite-api/expressions.ts
@@ -2150,6 +2150,16 @@ export class Constant extends Expr {
     }
   }
 
+  /**
+   * @private
+   * @internal
+   */
+  static _fromProto(value: ProtoValue): Constant {
+    const result = new Constant(value);
+    result._protoValue = value;
+    return result;
+  }
+
   /**
    * @private
    * @internal
@@ -2179,6 +2189,10 @@ export class Constant extends Expr {
    * @internal
    */
   _readUserData(dataReader: UserDataReader): void {
+    if (!!this._protoValue) {
+      return;
+    }
+
     const context = dataReader.createContext(
       UserDataSource.Argument,
       'Constant.of'
diff --git a/packages/firestore/src/lite-api/user_data_reader.ts b/packages/firestore/src/lite-api/user_data_reader.ts
@@ -75,6 +75,8 @@ import {
 } from './reference';
 import { Timestamp } from './timestamp';
 import { VectorValue } from './vector_value';
+import { isFirestoreValue } from '../core/pipeline-util';
+import { Constant } from './expressions';
 
 const RESERVED_FIELD_REGEX = /^__.*__$/;
 
@@ -331,7 +333,7 @@ class ParseContextImpl implements ParseContext {
  * classes.
  */
 export class UserDataReader {
-  private readonly serializer: JsonProtoSerializer;
+  readonly serializer: JsonProtoSerializer;
 
   constructor(
     private readonly databaseId: DatabaseId,
@@ -797,6 +799,10 @@ export function parseData(
   // from firestore-exp.
   input = getModularInstance(input);
 
+  if (input instanceof Constant) {
+    return input._getValue();
+  }
+
   if (looksLikeJsonObject(input)) {
     validatePlainObject('Unsupported field value:', context, input);
     return parseObject(input, context);
diff --git a/packages/firestore/src/local/query_engine.ts b/packages/firestore/src/local/query_engine.ts
@@ -57,7 +57,10 @@ import {
   stringifyQueryOrPipeline
 } from '../core/pipeline-util';
 import * as querystring from 'node:querystring';
-import { pipelineMatchesAllDocuments } from '../core/pipeline_run';
+import {
+  pipelineMatches,
+  pipelineMatchesAllDocuments
+} from '../core/pipeline_run';
 import { compareByKey } from '../model/document_comparator';
 
 const DEFAULT_INDEX_AUTO_CREATION_MIN_COLLECTION_SIZE = 100;
@@ -428,6 +431,7 @@ export class QueryEngine {
       // TODO(pipeline): the order here does not actually matter, not until we implement
       // refill logic for pipelines as well.
       queryResults = new SortedSet<Document>(compareByKey);
+      matcher = doc => pipelineMatches(query, doc as MutableDocument);
     } else {
       // Sort the documents and re-apply the query filter since previously
       // matching documents do not necessarily still match the query.
diff --git a/packages/firestore/src/model/path.ts b/packages/firestore/src/model/path.ts
@@ -19,6 +19,8 @@ import { debugAssert, fail } from '../util/assert';
 import { Code, FirestoreError } from '../util/error';
 
 export const DOCUMENT_KEY_NAME = '__name__';
+export const UPDATE_TIME_NAME = '__update_time__';
+export const CREATE_TIME_NAME = '__create_time__';
 
 /**
  * Path represents an ordered sequence of string segments.
diff --git a/packages/firestore/test/unit/local/query_engine.test.ts b/packages/firestore/test/unit/local/query_engine.test.ts
