diff --git a/handwritten/firestore/dev/src/pipelines/expression.ts b/handwritten/firestore/dev/src/pipelines/expression.ts index d6a770cff67..e7d931f43ca 100644 --- a/handwritten/firestore/dev/src/pipelines/expression.ts +++ b/handwritten/firestore/dev/src/pipelines/expression.ts @@ -23,11 +23,14 @@ import { fieldOrExpression, isFirestoreValue, isString, + toField, valueToDefaultExpr, vectorToExpr, } from './pipeline-util'; import {HasUserData, Serializer, validateUserInput} from '../serializer'; import {cast} from '../util'; +import {GeoPoint} from '../geo-point'; +import {OptionsUtil} from './options-util'; /** * @beta @@ -3100,6 +3103,85 @@ export abstract class Expression ]).asBoolean(); } + /** + * Evaluates if the result of this `expression` is between + * the `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ```typescript + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * field('tireWidth').between(constant(2.2), constant(2.4)) + * + * // This is functionally equivalent to + * and(field('tireWidth').greaterThanOrEqual(contant(2.2)), field('tireWidth').lessThanOrEqual(constant(2.4))) + * ``` + * + * @param lowerBound - An `Expression` that evaluates to the lower bound (inclusive) of the range. + * @param upperBound - An `Expression` that evaluates to the upper bound (inclusive) of the range. + */ + between(lowerBound: Expression, upperBound: Expression): BooleanExpression; + + /** + * Evaluates if the result of this `expression` is between + * the `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * field('tireWidth').between(2.2, 2.4) + * + * // This is functionally equivalent to + * and(field('tireWidth').greaterThanOrEqual(2.2), field('tireWidth').lessThanOrEqual(2.4)) + * ``` + * + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ + between(lowerBound: unknown, upperBound: unknown): BooleanExpression; + + between(lowerBound: unknown, upperBound: unknown): BooleanExpression { + return new FunctionExpression('between', [ + this, + valueToDefaultExpr(lowerBound), + valueToDefaultExpr(upperBound), + ]).asBoolean(); + } + + /** + * Evaluates to an HTML-formatted text snippet that renders terms matching + * the search query in `bold`. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param rquery Define the search query using the search DTS (TODO(search) link). + */ + snippet(rquery: string): Expression; + + /** + * Evaluates to an HTML-formatted text snippet that renders terms matching + * the search query in `bold`. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param options Define how snippeting behaves. + */ + snippet(options: firestore.Pipelines.SnippetOptions): Expression; + + snippet( + queryOrOptions: string | firestore.Pipelines.SnippetOptions, + ): Expression { + const options: firestore.Pipelines.SnippetOptions = isString(queryOrOptions) + ? {rquery: queryOrOptions} + : queryOrOptions; + const rquery = options.rquery; + const internalOptions = { + maxSnippetWidth: options.maxSnippetWidth, + maxSnippets: options.maxSnippets, + separator: options.separator, + }; + return new SnippetExpression([this, constant(rquery)], internalOptions); + } + // TODO(new-expression): Add new expression method definitions above this line /** @@ -3351,6 +3433,35 @@ export class Field readonly expressionType: firestore.Pipelines.ExpressionType = 'Field'; selectable = true as const; + /** + * Perform a full-text search on this field. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param rquery Define the search query using the rquery DTS. + */ + matches(rquery: string | Expression): BooleanExpression { + return new FunctionExpression('matches', [ + this, + valueToDefaultExpr(rquery), + ]).asBoolean(); + } + + /** + * Evaluates to the distance in meters between the location specified + * by this field and the query location. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param location - Compute distance to this GeoPoint. + */ + geoDistance(location: GeoPoint | Expression): Expression { + return new FunctionExpression('geo_distance', [ + this, + valueToDefaultExpr(location), + ]).asBoolean(); + } + /** * @beta * @internal @@ -3686,6 +3797,51 @@ export class FunctionExpression extends Expression { } } +/** + * SnippetExpression extends from FunctionExpression because it + * supports options and requires the options util. + */ +export class SnippetExpression extends FunctionExpression { + /** + * @private + * @internal + */ + get _optionsUtil(): OptionsUtil { + return new OptionsUtil({ + maxSnippetWidth: { + serverName: 'max_snippet_width', + }, + maxSnippets: { + serverName: 'max_snippets', + }, + separator: { + serverName: 'separator', + }, + }); + } + + /** + * @hideconstructor + */ + constructor( + params: Expression[], + private _options?: {}, + ) { + super('snippet', params); + } + + _toProto(serializer: Serializer): api.IValue { + const proto = super._toProto(serializer); + proto.functionValue!.options = this._optionsUtil.getOptionsProto( + serializer, + this._options ?? {}, + {}, + ); + + return proto; + } +} + /** * @beta * This class defines the base class for Firestore `Pipeline` functions, which can be evaluated within pipeline @@ -10258,6 +10414,201 @@ export function isType( return fieldOrExpression(fieldNameOrExpression).isType(type); } +/** + * Perform a full-text search on the specified field. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param searchField Search the specified field. + * @param rquery Define the search query using the search DTS. + */ +export function matches( + searchField: string | Field, + rquery: string | Expression, +): BooleanExpression { + return toField(searchField).matches(rquery); +} + +/** + * Perform a full-text search on the document. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param rquery Define the search query using the rquery DTS. + */ +export function documentMatches( + rquery: string | Expression, +): BooleanExpression { + return new FunctionExpression('document_matches', [ + valueToDefaultExpr(rquery), + ]).asBoolean(); +} + +/** + * Evaluates to the search score that reflects the topicality of the document + * to all of the text predicates (`queryMatch`) + * in the search query. If `SearchOptions.query` is not set or does not contain + * any text predicates, then this topicality score will always be `0`. + * + * @remarks This Expression can only be used within a `Search` stage. + */ +export function score(): Expression { + return new FunctionExpression('score', []).asBoolean(); +} + +/** + * Evaluates to an HTML-formatted text snippet that highlights terms matching + * the search query in `bold`. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param searchField Search the specified field for matching terms. + * @param rquery Define the search query using the search DTS (TODO(search) link). + */ +export function snippet( + searchField: string | Field, + rquery: string, +): Expression; + +/** + * Evaluates to an HTML-formatted text snippet that highlights terms matching + * the search query in `bold`. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param searchField Search the specified field for matching terms. + * @param options Define the search query using the search DTS (TODO(search) link). + */ +export function snippet( + searchField: string | Field, + options: firestore.Pipelines.SnippetOptions, +): Expression; +export function snippet( + field: string | Field, + queryOrOptions: string | firestore.Pipelines.SnippetOptions, +): Expression { + return toField(field).snippet( + isString(queryOrOptions) ? {rquery: queryOrOptions} : queryOrOptions, + ); +} + +/** + * Evaluates to the distance in meters between the location in the specified + * field and the query location. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param fieldName - Specifies the field in the document which contains + * the first GeoPoint for distance computation. + * @param location - Compute distance to this GeoPoint. + */ +export function geoDistance( + fieldName: string | Field, + location: GeoPoint | Expression, +): Expression { + return toField(fieldName).geoDistance(location); +} + +/** + * Evaluates if the value in the field specified by `fieldName` is between + * the evaluated values for `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * between('tireWidth', constant(2.2), constant(2.4)) + * + * // This is functionally equivalent to + * and(greaterThanOrEqual('tireWidth', constant(2.2)), lessThanOrEqual('tireWidth', constant(2.4))) + * ``` + * + * @param fieldName - Evaluate if the value stored in this field is between the lower and upper bounds. + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ +export function between( + fieldName: string, + lowerBound: Expression, + upperBound: Expression, +): BooleanExpression; + +/** + * Evaluates if the value in the field specified by `fieldName` is between + * the values for `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * between('tireWidth', 2.2, 2.4) + * + * // This is functionally equivalent to + * and(greaterThanOrEqual('tireWidth', 2.2), lessThanOrEqual('tireWidth', 2.4)) + * ``` + * + * @param fieldName - Evaluate if the value stored in this field is between the lower and upper bounds. + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ +export function between( + fieldName: string, + lowerBound: unknown, + upperBound: unknown, +): BooleanExpression; + +/** + * Evaluates if the result of the specified `expression` is between + * the results of `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * between(field('tireWidth'), constant(2.2), constant(2.4)) + * + * // This is functionally equivalent to + * and(greaterThanOrEqual(field('tireWidth'), constant(2.2)), lessThanOrEqual(field('tireWidth'), constant(2.4))) + * ``` + * + * @param expression - Evaluate if the result of this expression is between the lower and upper bounds. + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ +export function between( + expression: Expression, + lowerBound: Expression, + upperBound: Expression, +): BooleanExpression; + +/** + * Evaluates if the result of the specified `expression` is between + * the `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * between(field('tireWidth'), 2.2, 2.4) + * + * // This is functionally equivalent to + * and(greaterThanOrEqual(field('tireWidth'), 2.2), lessThanOrEqual(field('tireWidth'), 2.4)) + * ``` + * + * @param expression - Evaluate if the result of this expression is between the lower and upper bounds. + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ +export function between( + expression: Expression, + lowerBound: unknown, + upperBound: unknown, +): BooleanExpression; + +export function between( + expression: Expression | string, + lowerBound: unknown, + upperBound: unknown, +): BooleanExpression { + return fieldOrExpression(expression).between(lowerBound, upperBound); +} + // TODO(new-expression): Add new top-level expression function definitions above this line /** diff --git a/handwritten/firestore/dev/src/pipelines/pipeline-util.ts b/handwritten/firestore/dev/src/pipelines/pipeline-util.ts index 5c12257c23f..909da33e7fe 100644 --- a/handwritten/firestore/dev/src/pipelines/pipeline-util.ts +++ b/handwritten/firestore/dev/src/pipelines/pipeline-util.ts @@ -577,6 +577,7 @@ export function isSelectable( export function isOrdering(val: unknown): val is firestore.Pipelines.Ordering { const candidate = val as firestore.Pipelines.Ordering; return ( + val !== undefined && isExpr(candidate.expr) && (candidate.direction === 'ascending' || candidate.direction === 'descending') @@ -721,10 +722,19 @@ export function fieldOrSelectable(value: string | Selectable): Selectable { } } +/** + * @deprecated use selectablesToObject instead + */ export function selectablesToMap( selectables: (firestore.Pipelines.Selectable | string)[], ): Map { - const result = new Map(); + return new Map(Object.entries(selectablesToObject(selectables))); +} + +export function selectablesToObject( + selectables: (firestore.Pipelines.Selectable | string)[], +): Record { + const result: Record = {}; for (const selectable of selectables) { let alias: string; let expression: Expression; @@ -736,11 +746,11 @@ export function selectablesToMap( expression = selectable._expr as unknown as Expression; } - if (result.get(alias) !== undefined) { + if (result[alias] !== undefined) { throw new Error(`Duplicate alias or field '${alias}'`); } - result.set(alias, expression); + result[alias] = expression; } return result; } diff --git a/handwritten/firestore/dev/src/pipelines/pipelines.ts b/handwritten/firestore/dev/src/pipelines/pipelines.ts index 2844c29ca57..61af8375a1a 100644 --- a/handwritten/firestore/dev/src/pipelines/pipelines.ts +++ b/handwritten/firestore/dev/src/pipelines/pipelines.ts @@ -39,6 +39,7 @@ import { selectablesToMap, toField, vectorToExpr, + selectablesToObject, } from './pipeline-util'; import {DocumentReference} from '../reference/document-reference'; import {PipelineResponse} from '../reference/types'; @@ -59,6 +60,7 @@ import { constant, _mapValue, field, + documentMatches, } from './expression'; import { AddFields, @@ -95,6 +97,8 @@ import { InternalDocumentsStageOptions, InternalCollectionGroupStageOptions, InternalCollectionStageOptions, + Search, + InternalSearchStageOptions, } from './stage'; import {StructuredPipeline} from './structured-pipeline'; import Selectable = FirebaseFirestore.Pipelines.Selectable; @@ -1420,6 +1424,42 @@ export class Pipeline implements firestore.Pipelines.Pipeline { return this._addStage(new Unnest(internalOptions)); } + /** + * Add a search stage to the Pipeline. + * + * @remarks This must be the first stage of the pipeline. + * @remarks A limited set of expressions are supported in the search stage. + * + * @param options - An object that specifies required and optional parameters + * for the stage. + * @return A new `Pipeline` object with this stage appended to the stage list. + */ + search(options: firestore.Pipelines.SearchStageOptions): Pipeline { + // Convert user land convenience types to internal types + const normalizedQuery: firestore.Pipelines.BooleanExpression = isExpr( + options.query, + ) + ? options.query + : documentMatches(options.query); + const normalizedSelect: Record | undefined = + options.select ? selectablesToObject(options.select) : undefined; + const normalizedAddFields: Record | undefined = + options.addFields ? selectablesToObject(options.addFields) : undefined; + const normalizedSort: firestore.Pipelines.Ordering[] | undefined = + isOrdering(options.sort) ? [options.sort] : options.sort; + + const internalOptions: InternalSearchStageOptions = { + ...options, + query: normalizedQuery, + select: normalizedSelect, + addFields: normalizedAddFields, + sort: normalizedSort, + }; + + // Add stage to the pipeline + return this._addStage(new Search(internalOptions)); + } + /** * @beta * Sorts the documents from previous stages based on one or more `Ordering` criteria. diff --git a/handwritten/firestore/dev/src/pipelines/stage.ts b/handwritten/firestore/dev/src/pipelines/stage.ts index 9695624fbdd..b06b70e310a 100644 --- a/handwritten/firestore/dev/src/pipelines/stage.ts +++ b/handwritten/firestore/dev/src/pipelines/stage.ts @@ -645,6 +645,67 @@ export class Sort implements Stage { } } +export type InternalSearchStageOptions = Omit< + firestore.Pipelines.SearchStageOptions, + 'query' | 'sort' | 'select' | 'addFields' +> & { + query: firestore.Pipelines.BooleanExpression; + sort?: Array; + select?: Record; + addFields?: Record; +}; + +/** + * Search stage. + */ +export class Search implements Stage { + name = 'search'; + + constructor(private options: InternalSearchStageOptions) {} + + readonly optionsUtil = new OptionsUtil({ + query: { + serverName: 'query', + }, + limit: { + serverName: 'limit', + }, + retrievalDepth: { + serverName: 'retrieval_depth', + }, + sort: { + serverName: 'sort', + }, + addFields: { + serverName: 'add_fields', + }, + select: { + serverName: 'select', + }, + offset: { + serverName: 'offset', + }, + queryEnhancement: { + serverName: 'query_enhancement', + }, + languageCode: { + serverName: 'language_code', + }, + }); + + _toProto(serializer: Serializer): api.Pipeline.IStage { + return { + name: this.name, + args: [], + options: this.optionsUtil.getOptionsProto( + serializer, + this.options, + this.options.rawOptions, + ), + }; + } +} + /** * Raw stage. */ diff --git a/handwritten/firestore/dev/system-test/pipeline.ts b/handwritten/firestore/dev/system-test/pipeline.ts index 8377424c8d2..6363a10c300 100644 --- a/handwritten/firestore/dev/system-test/pipeline.ts +++ b/handwritten/firestore/dev/system-test/pipeline.ts @@ -173,53 +173,56 @@ import {getTestDb, getTestRoot} from './firestore'; import {Firestore as InternalFirestore} from '../src'; import {ServiceError} from 'google-gax'; +import {documentMatches, score} from '../src/pipelines/expression'; use(chaiAsPromised); const timestampDeltaMS = 3000; +let beginDocCreation = 0; +let endDocCreation = 0; -describe.skipClassic('Pipeline class', () => { - let firestore: Firestore; - let randomCol: CollectionReference; - let beginDocCreation = 0; - let endDocCreation = 0; - - async function testCollectionWithDocs(docs: { +async function testCollectionWithDocs( + collection: CollectionReference, + docs: { [id: string]: DocumentData; - }): Promise> { - beginDocCreation = new Date().valueOf(); - for (const id in docs) { - const ref = randomCol.doc(id); - await ref.set(docs[id]); - } - endDocCreation = new Date().valueOf(); - return randomCol; + }, +): Promise> { + beginDocCreation = new Date().valueOf(); + for (const id in docs) { + const ref = collection.doc(id); + await ref.set(docs[id]); } - - function expectResults(result: PipelineSnapshot, ...docs: string[]): void; - function expectResults( - result: PipelineSnapshot, - ...data: DocumentData[] - ): void; - function expectResults( - result: PipelineSnapshot, - ...data: DocumentData[] | string[] - ): void { - if (data.length > 0) { - if (typeof data[0] === 'string') { - const actualIds = result.results.map(result => result.id); - expect(actualIds).to.deep.equal(data); - } else { - result.results.forEach(r => { - expect(r.data()).to.deep.equal(data.shift()); - }); - } + endDocCreation = new Date().valueOf(); + return collection; +} + +function expectResults(result: PipelineSnapshot, ...docs: string[]): void; +function expectResults(result: PipelineSnapshot, ...data: DocumentData[]): void; +function expectResults( + result: PipelineSnapshot, + ...data: DocumentData[] | string[] +): void { + if (data.length > 0) { + if (typeof data[0] === 'string') { + const actualIds = result.results.map(result => result.id); + expect(actualIds).to.deep.equal(data); } else { - expect(result.results.length).to.equal(data.length); + result.results.forEach(r => { + expect(r.data()).to.deep.equal(data.shift()); + }); } + } else { + expect(result.results.length).to.equal(data.length); } +} + +describe.skipClassic('Pipeline class', () => { + let firestore: Firestore; + let randomCol: CollectionReference; - async function setupBookDocs(): Promise> { + async function setupBookDocs( + collection: CollectionReference, + ): Promise> { const bookDocs: {[id: string]: DocumentData} = { book1: { title: "The Hitchhiker's Guide to the Galaxy", @@ -329,12 +332,12 @@ describe.skipClassic('Pipeline class', () => { embedding: FieldValue.vector([1, 1, 1, 1, 1, 1, 1, 1, 1, 10]), }, }; - return testCollectionWithDocs(bookDocs); + return testCollectionWithDocs(collection, bookDocs); } before(async () => { randomCol = getTestRoot(); - await setupBookDocs(); + await setupBookDocs(randomCol); firestore = randomCol.firestore; }); @@ -5627,6 +5630,601 @@ describe.skipClassic('Pipeline class', () => { }); }); +// Search tests require a collection with an index, so the test setup and tear +// down is managed different from the rest of the Pipeline tests. To accomplish +// this, we break these tests into a separate describe +describe.skipClassic('Pipeline search', () => { + let firestore: Firestore; + let restaurantsCollection: CollectionReference; + + async function setupRestaurantDocs( + collection: CollectionReference, + ): Promise> { + const restaurantDocs: {[id: string]: DocumentData} = { + sunnySideUp: { + name: 'The Sunny Side Up', + description: + 'A cozy neighborhood diner serving classic breakfast favorites all day long, from fluffy pancakes to savory omelets.', + location: new GeoPoint(39.7541, -105.0002), + menu: '

Breakfast Classics

  • Denver Omelet - $12
  • Buttermilk Pancakes - $10
  • Steak and Eggs - $16

Sides

  • Hash Browns - $4
  • Thick-cut Bacon - $5
  • Drip Coffee - $2
', + average_price_per_person: 15, + }, + goldenWaffle: { + name: 'The Golden Waffle', + description: + 'Specializing exclusively in Belgian-style waffles. Open daily from 6:00 AM to 11:00 AM.', + location: new GeoPoint(39.7183, -104.9621), + menu: '

Signature Waffles

  • Strawberry Delight - $11
  • Chicken and Waffles - $14
  • Chocolate Chip Crunch - $10

Drinks

  • Fresh OJ - $4
  • Artisan Coffee - $3
', + average_price_per_person: 13, + }, + lotusBlossomThai: { + name: 'Lotus Blossom Thai', + description: + 'Authentic Thai cuisine featuring hand-crushed spices and traditional family recipes from the Chiang Mai region.', + location: new GeoPoint(39.7315, -104.9847), + menu: '

Appetizers

  • Spring Rolls - $7
  • Chicken Satay - $9

Main Course

  • Pad Thai - $15
  • Green Curry - $16
  • Drunken Noodles - $15
', + average_price_per_person: 22, + }, + mileHighCatch: { + name: 'Mile High Catch', + description: + 'Freshly sourced seafood offering a wide variety of Pacific fish and Atlantic shellfish in an upscale atmosphere.', + location: new GeoPoint(39.7401, -104.9903), + menu: '

From the Raw Bar

  • Oysters (Half Dozen) - $18
  • Lobster Cocktail - $22

Entrees

  • Pan-Seared Salmon - $28
  • King Crab Legs - $45
  • Fish and Chips - $19
', + average_price_per_person: 45, + }, + peakBurgers: { + name: 'Peak Burgers', + description: + 'Casual burger joint focused on locally sourced Colorado beef and hand-cut fries.', + location: new GeoPoint(39.7622, -105.0125), + menu: '

Burgers

  • The Peak Double - $12
  • Bison Burger - $15
  • Veggie Stack - $11

Sides

  • Truffle Fries - $6
  • Onion Rings - $5
', + average_price_per_person: 18, + }, + solTacos: { + name: 'El Sol Tacos', + description: + 'A vibrant street-side taco stand serving up quick, delicious, and traditional Mexican street food.', + location: new GeoPoint(39.6952, -105.0274), + menu: '

Tacos ($3.50 each)

  • Al Pastor
  • Carne Asada
  • Pollo Asado
  • Nopales (Cactus)

Beverages

  • Horchata - $4
  • Mexican Coke - $3
', + average_price_per_person: 12, + }, + eastsideTacos: { + name: 'Eastside Cantina', + description: + 'Authentic street tacos and hand-shaken margaritas on the vibrant east side of the city.', + location: new GeoPoint(39.735, -104.885), + menu: '

Tacos

  • Carnitas Tacos - $4
  • Barbacoa Tacos - $4.50
  • Shrimp Tacos - $5

Drinks

  • House Margarita - $9
  • Jarritos - $3
', + average_price_per_person: 18, + }, + eastsideChicken: { + name: 'Eastside Chicken', + description: 'Fried chicken to go - next to Eastside Cantina.', + location: new GeoPoint(39.735, -104.885), + menu: '

Fried Chicken

  • Drumstick - $4
  • Wings - $1
  • Sandwich - $9

Drinks

  • House Margarita - $9
  • Jarritos - $3
', + average_price_per_person: 12, + }, + }; + + // TODO(search) - Migrate this over to IndexTestHelper when search supports the equal filter. + // Remove any restaurant docs not in the expected set - perhaps these were + // set by another dev or test suite. This has potential to cause flakes in another concurrent + // run of these tests, if they have added new test docs. + const collectionSnapshot = await collection.get(); + const expectedDocIds = Object.keys(restaurantDocs); + const deletes = collectionSnapshot.docs + .filter(ds => expectedDocIds.indexOf(ds.id) < 0) + .map(ds => ds.ref.delete()); + await Promise.all(deletes); + + // Add/overwrite all restaurant docs + return testCollectionWithDocs(collection, restaurantDocs); + } + + // Search tests will use restaurant docs + before(async () => { + // TODO(search) - Migrate this over to IndexTestHelper when search supports the equal filter. + // Note: using a static collection of documents for every search test has an inherent risk + // of flakiness. Search requires an index on the collection, which is the reason we use a pre-defined + // collection. We cannot use the IndexTestHelper because that relies on an equality match to the testID + // field. Search currently does not support the equal expression. + firestore = getTestDb(); + restaurantsCollection = firestore.collection('SearchIntegrationTests'); + await setupRestaurantDocs(restaurantsCollection); + firestore = restaurantsCollection.firestore; + }); + + describe('search stage', () => { + describe('DISABLE query expansion', () => { + describe('query', () => { + it('all search features', async () => { + const queryLocation = new GeoPoint(0, 0); + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: and( + documentMatches('waffles'), + field('description').matches('breakfast'), + field('location').geoDistance(queryLocation).lessThan(1000), + field('avgPrice').between(10, 20), + ), + select: [ + field('title'), + field('menu'), + field('description'), + field('location').geoDistance(queryLocation).as('distance'), + ], + addFields: [score().as('searchScore')], + offset: 0, + retrievalDepth: 1000, + limit: 50, + sort: [field('location').geoDistance(queryLocation).ascending()], + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'goldenWaffle'); + }); + + it('search full document', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: documentMatches('waffles'), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'goldenWaffle'); + }); + + it('search a specific field', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('waffles'), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'goldenWaffle'); + }); + + it('geo near query', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('location') + .geoDistance(new GeoPoint(39.6985, -105.024)) + .lessThan(1000 /* m */), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'solTacos'); + }); + + it('conjunction of text search predicates', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: and( + field('menu').matches('waffles'), + field('description').matches('diner'), + ), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'goldenWaffle', 'sunnySideUp'); + }); + + it('conjunction of text search and geo near', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: and( + field('menu').matches('tacos'), + field('location') + .geoDistance(new GeoPoint(39.6985, -105.024)) + .lessThan(10_000 /* meters */), + ), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'solTacos'); + }); + + it('negate match', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('-waffles'), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults( + snapshot, + 'eastSideTacos', + 'solTacos', + 'peakBurgers', + 'mileHighCatch', + 'lotusBlossomThai', + 'sunnySideUp', + ); + }); + + it('rquery search the document with conjunction and disjunction', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: documentMatches('(waffles OR pancakes) AND coffee'), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'goldenWaffle', 'sunnySideUp'); + }); + + it('rquery as query param', async () => { + const ppl = firestore.pipeline().collection('restaurants').search({ + query: '(waffles OR pancakes) AND coffee', + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'goldenWaffle', 'sunnySideUp'); + }); + + it('rquery supports field paths', async () => { + const ppl = firestore.pipeline().collection('restaurants').search({ + query: + 'menu:(waffles OR pancakes) AND description:"breakfast all day"', + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'sunnySideUp'); + }); + + it('conjunction of rquery and expression', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: and( + documentMatches('tacos'), + field('average_price_per_person').between(8, 15), + ), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'solTacos'); + }); + }); + + describe('addFields', () => { + it('topicality score and snippet', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('waffles'), + addFields: [ + score().as('searchScore'), + field('menu').snippet('waffles').as('snippet'), + ], + queryEnhancement: 'disabled', + }) + .select('name', 'searchScore', 'snippet'); + + const snapshot = await ppl.execute(); + expect(snapshot.results.length).to.equal(1); + expect(snapshot.results[0].get('name')).to.equal('The Golden Waffle'); + expect(snapshot.results[0].get('searchScore')).to.be.greaterThan(0); + expect(snapshot.results[0].get('snippet')?.length).to.be.greaterThan( + 0, + ); + }); + }); + + describe('select', () => { + it('topicality score and snippet', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('waffles'), + select: [ + field('name'), + 'location', + score().as('searchScore'), + field('menu').snippet('waffles').as('snippet'), + ], + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expect(snapshot.results.length).to.equal(1); + expect(snapshot.results[0].get('name')).to.equal('The Golden Waffle'); + expect(snapshot.results[0].get('location')).to.equal( + new GeoPoint(39.7183, -104.9621), + ); + expect(snapshot.results[0].get('searchScore')).to.be.greaterThan(0); + expect(snapshot.results[0].get('snippet')?.length).to.be.greaterThan( + 0, + ); + expect(Object.keys(snapshot.results[0].data()).sort()).to.deep.equal([ + 'location', + 'name', + 'searchScore', + 'snippet', + ]); + }); + }); + + describe('sort', () => { + it('by topicality', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('tacos'), + sort: score().descending(), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'eastsideTacos', 'solTacos'); + }); + + it('by distance', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('tacos'), + sort: field('location') + .geoDistance(new GeoPoint(39.6985, -105.024)) + .ascending(), + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'solTacos', 'eastsideTacos'); + }); + + it('by multiple orderings', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('tacos OR chicken'), + sort: [ + field('location') + .geoDistance(new GeoPoint(39.6985, -105.024)) + .ascending(), + score().descending(), + ], + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults( + snapshot, + 'solTacos', + 'eastsideTacos', + 'eastsideChicken', + ); + }); + }); + + describe('limit', () => { + it('limits the number of documents returned', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: constant(true), + sort: field('location') + .geoDistance(new GeoPoint(39.6985, -105.024)) + .ascending(), + limit: 5, + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults( + snapshot, + 'solTacos', + 'lotusBlossomThai', + 'goldenWaffle', + ); + }); + + it('limits the number of documents scored', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches( + 'chicken OR tacos OR fish OR waffles', + ), + retrievalDepth: 6, + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults( + snapshot, + 'eastsideChicken', + 'eastsideTacos', + 'solTacos', + 'mileHighCatch', + ); + }); + }); + + describe('offset', () => { + it('skips N documents', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: constant(true), + limit: 2, + offset: 2, + queryEnhancement: 'disabled', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'eastsideChicken', 'eastsideTacos'); + }); + }); + }); + + describe('REQUIRE query expansion', () => { + it('search full document', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: documentMatches('waffles'), + queryEnhancement: 'required', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'goldenWaffle', 'sunnySideUp'); + }); + + it('search a specific field', async () => { + const ppl = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('waffles'), + queryEnhancement: 'required', + }); + + const snapshot = await ppl.execute(); + expectResults(snapshot, 'goldenWaffle', 'sunnySideUp'); + }); + }); + }); + + describe('snippet', () => { + it('snippet options', async () => { + const ppl1 = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('waffles'), + addFields: [ + field('menu') + .snippet({ + rquery: 'waffles', + maxSnippetWidth: 10, + }) + .as('snippet'), + ], + queryEnhancement: 'disabled', + }); + + const snapshot1 = await ppl1.execute(); + expect(snapshot1.results.length).to.equal(1); + expect(snapshot1.results[0].get('name')).to.equal('The Golden Waffle'); + expect(snapshot1.results[0].get('snippet')?.length).to.be.greaterThan(0); + + const ppl2 = firestore + .pipeline() + .collection('restaurants') + .search({ + query: field('menu').matches('waffles'), + addFields: [ + field('menu') + .snippet({ + rquery: 'waffles', + maxSnippetWidth: 1000, + }) + .as('snippet'), + ], + queryEnhancement: 'disabled', + }); + + const snapshot2 = await ppl2.execute(); + expect(snapshot2.results.length).to.equal(1); + expect(snapshot2.results[0].get('name')).to.equal('The Golden Waffle'); + expect(snapshot2.results[0].get('snippet')?.length).to.be.greaterThan(0); + + expect(snapshot2.results[0].get('snippet')?.length).to.be.greaterThan( + snapshot2.results[0].get('snippet')?.length, + ); + }); + + it('snippet on multiple fields', async () => { + // Get snippet from 1 field + const ppl1 = firestore + .pipeline() + .collection('restaurants') + .search({ + query: documentMatches('waffle'), + addFields: [ + field('menu') + .snippet({ + rquery: 'waffles', + maxSnippetWidth: 2000, + }) + .as('snippet'), + ], + queryEnhancement: 'disabled', + }); + + const snapshot1 = await ppl1.execute(); + expect(snapshot1.results.length).to.equal(1); + expect(snapshot1.results[0].get('name')).to.equal('The Golden Waffle'); + expect(snapshot1.results[0].get('snippet')?.length).to.be.greaterThan(0); + + // Get snippet from 2 fields + const ppl2 = firestore + .pipeline() + .collection('restaurants') + .search({ + query: documentMatches('waffle'), + addFields: [ + concat(field('menu'), field('description')) + .snippet({ + rquery: 'waffles', + maxSnippetWidth: 2000, + }) + .as('snippet'), + ], + queryEnhancement: 'disabled', + }); + + const snapshot2 = await ppl2.execute(); + expect(snapshot2.results.length).to.equal(1); + expect(snapshot2.results[0].get('name')).to.equal('The Golden Waffle'); + expect(snapshot2.results[0].get('snippet')?.length).to.be.greaterThan(0); + + // Expect snippet from 2 fields to be longer than snippet from one field + expect(snapshot2.results[0].get('snippet')?.length).to.be.greaterThan( + snapshot2.results[0].get('snippet')?.length, + ); + }); + }); +}); + // This is the Query integration tests from the lite API (no cache support) // with some additional test cases added for more complete coverage. // eslint-disable-next-line no-restricted-properties diff --git a/handwritten/firestore/dev/test/pipelines/pipeline.ts b/handwritten/firestore/dev/test/pipelines/pipeline.ts index e67514a3342..fa48b1c49ff 100644 --- a/handwritten/firestore/dev/test/pipelines/pipeline.ts +++ b/handwritten/firestore/dev/test/pipelines/pipeline.ts @@ -19,14 +19,9 @@ import {expect} from 'chai'; import * as sinon from 'sinon'; import {createInstance, stream} from '../util/helpers'; import {google} from '../../protos/firestore_v1_proto_api'; -import {Timestamp, Pipelines, Firestore} from '../../src'; +import {Timestamp} from '../../src'; import IExecutePipelineRequest = google.firestore.v1.IExecutePipelineRequest; import IExecutePipelineResponse = google.firestore.v1.IExecutePipelineResponse; -import Pipeline = Pipelines.Pipeline; -import field = Pipelines.field; -import sum = Pipelines.sum; -import descending = Pipelines.descending; -import IValue = google.firestore.v1.IValue; const FIRST_CALL = 0; const EXECUTE_PIPELINE_REQUEST = 0; @@ -185,245 +180,3 @@ describe('execute(Pipeline|PipelineExecuteOptions)', () => { ); }); }); - -describe('stage option serialization', () => { - // Default rawOptions - const rawOptions: Record = { - foo: 'bar1', - }; - // Default expected serialized options - const expectedSerializedOptions: Record = { - foo: { - stringValue: 'bar1', - }, - }; - - const testDefinitions: Array<{ - name: string; - pipeline: (firestore: Firestore) => Pipeline; - stageIndex?: number; - expectedOptions?: Record; - }> = [ - { - name: 'collection stage', - pipeline: firestore => - firestore.pipeline().collection({ - collection: 'foo', - rawOptions, - forceIndex: 'foo-index', - }), - expectedOptions: { - ...expectedSerializedOptions, - force_index: { - stringValue: 'foo-index', - }, - }, - }, - { - name: 'collection group stage', - pipeline: firestore => - firestore.pipeline().collectionGroup({ - collectionId: 'foo', - rawOptions, - forceIndex: 'bar-index', - }), - expectedOptions: { - ...expectedSerializedOptions, - force_index: { - stringValue: 'bar-index', - }, - }, - }, - { - name: 'documents stage', - pipeline: firestore => - firestore.pipeline().documents({ - docs: ['foo/bar'], - rawOptions, - }), - }, - { - name: 'database stage', - pipeline: firestore => - firestore.pipeline().database({ - rawOptions, - }), - }, - { - name: 'distinct stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .distinct({ - groups: ['foo'], - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'findNearest stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .findNearest({ - field: 'foo', - vectorValue: [0], - distanceMeasure: 'euclidean', - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'select stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .select({ - selections: ['foo'], - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'unnest stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .unnest({ - selectable: field('foo'), - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'addFields stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .addFields({ - fields: [field('foo')], - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'aggregate stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .aggregate({ - accumulators: [sum('foo').as('fooSum')], - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'limit stage', - pipeline: firestore => - firestore.pipeline().database().limit({ - limit: 1, - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'offset stage', - pipeline: firestore => - firestore.pipeline().database().offset({ - offset: 1, - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'removeFields stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .removeFields({ - fields: ['foo'], - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'replaceWith stage', - pipeline: firestore => - firestore.pipeline().database().replaceWith({ - map: 'foo', - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'sample stage', - pipeline: firestore => - firestore.pipeline().database().sample({ - documents: 100, - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'sample stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .sort({ - orderings: [descending('foo')], - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'union stage', - pipeline: firestore => - firestore.pipeline().database().union({ - other: firestore.pipeline().database(), - rawOptions, - }), - stageIndex: 1, - }, - { - name: 'where stage', - pipeline: firestore => - firestore - .pipeline() - .database() - .where({ - condition: field('foo').equal(1), - rawOptions, - }), - stageIndex: 1, - }, - ]; - - testDefinitions.forEach(testDefinition => { - it(testDefinition.name, async () => { - const spy = sinon.fake.returns(stream()); - const firestore = await createInstance({ - executePipeline: spy, - }); - - await testDefinition.pipeline(firestore).execute(); - - const expectedOptions = testDefinition.expectedOptions - ? testDefinition.expectedOptions - : expectedSerializedOptions; - - expect( - spy.args[FIRST_CALL][EXECUTE_PIPELINE_REQUEST]['structuredPipeline'][ - 'pipeline' - ]['stages'][testDefinition.stageIndex ?? 0]['options'], - ).to.deep.equal(expectedOptions); - }); - }); -}); diff --git a/handwritten/firestore/dev/test/pipelines/stage.ts b/handwritten/firestore/dev/test/pipelines/stage.ts new file mode 100644 index 00000000000..d3b4535574c --- /dev/null +++ b/handwritten/firestore/dev/test/pipelines/stage.ts @@ -0,0 +1,359 @@ +/** + * @license + * Copyright 2026 Google LLC + * + * 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. + */ + +import {expect} from 'chai'; +import * as sinon from 'sinon'; +import {createInstance, stream} from '../util/helpers'; +import {google} from '../../protos/firestore_v1_proto_api'; +import {Pipelines, Firestore} from '../../src'; +import Pipeline = Pipelines.Pipeline; +import field = Pipelines.field; +import sum = Pipelines.sum; +import descending = Pipelines.descending; +import constant = Pipelines.constant; +import IValue = google.firestore.v1.IValue; + +const FIRST_CALL = 0; +const EXECUTE_PIPELINE_REQUEST = 0; + +describe('stage option serialization', () => { + // Default rawOptions + const rawOptions: Record = { + foo: 'bar1', + }; + // Default expected serialized options + const expectedSerializedOptions: Record = { + foo: { + stringValue: 'bar1', + }, + }; + + const testDefinitions: Array<{ + name: string; + pipeline: (firestore: Firestore) => Pipeline; + stageIndex?: number; + expectedOptions?: Record; + }> = [ + { + name: 'collection stage', + pipeline: firestore => + firestore.pipeline().collection({ + collection: 'foo', + rawOptions, + forceIndex: 'foo-index', + }), + expectedOptions: { + ...expectedSerializedOptions, + force_index: { + stringValue: 'foo-index', + }, + }, + }, + { + name: 'collection group stage', + pipeline: firestore => + firestore.pipeline().collectionGroup({ + collectionId: 'foo', + rawOptions, + forceIndex: 'bar-index', + }), + expectedOptions: { + ...expectedSerializedOptions, + force_index: { + stringValue: 'bar-index', + }, + }, + }, + { + name: 'documents stage', + pipeline: firestore => + firestore.pipeline().documents({ + docs: ['foo/bar'], + rawOptions, + }), + }, + { + name: 'database stage', + pipeline: firestore => + firestore.pipeline().database({ + rawOptions, + }), + }, + { + name: 'distinct stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .distinct({ + groups: ['foo'], + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'findNearest stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .findNearest({ + field: 'foo', + vectorValue: [0], + distanceMeasure: 'euclidean', + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'select stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .select({ + selections: ['foo'], + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'unnest stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .unnest({ + selectable: field('foo'), + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'addFields stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .addFields({ + fields: [field('foo')], + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'aggregate stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .aggregate({ + accumulators: [sum('foo').as('fooSum')], + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'limit stage', + pipeline: firestore => + firestore.pipeline().database().limit({ + limit: 1, + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'offset stage', + pipeline: firestore => + firestore.pipeline().database().offset({ + offset: 1, + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'removeFields stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .removeFields({ + fields: ['foo'], + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'replaceWith stage', + pipeline: firestore => + firestore.pipeline().database().replaceWith({ + map: 'foo', + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'sample stage', + pipeline: firestore => + firestore.pipeline().database().sample({ + documents: 100, + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'sample stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .sort({ + orderings: [descending('foo')], + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'union stage', + pipeline: firestore => + firestore.pipeline().database().union({ + other: firestore.pipeline().database(), + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'where stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .where({ + condition: field('foo').equal(1), + rawOptions, + }), + stageIndex: 1, + }, + { + name: 'search stage', + pipeline: firestore => + firestore + .pipeline() + .database() + .search({ + query: 'foo', + limit: 1, + retrievalDepth: 2, + offset: 3, + queryEnhancement: 'required', + languageCode: 'en-US', + sort: [field('foo').ascending()], + addFields: [constant(true).as('bar')], + select: [field('id')], + rawOptions, + }), + stageIndex: 1, + expectedOptions: { + add_fields: { + mapValue: { + fields: { + bar: { + booleanValue: true, + }, + }, + }, + }, + foo: { + stringValue: 'bar1', + }, + language_code: { + stringValue: 'en-US', + }, + limit: { + integerValue: '1', + }, + offset: { + integerValue: '3', + }, + query: { + functionValue: { + args: [ + { + stringValue: 'foo', + }, + ], + name: 'document_matches', + }, + }, + query_enhancement: { + stringValue: 'required', + }, + retrieval_depth: { + integerValue: '2', + }, + select: { + mapValue: { + fields: { + id: { + fieldReferenceValue: 'id', + }, + }, + }, + }, + sort: { + arrayValue: { + values: [ + { + mapValue: { + fields: { + direction: { + stringValue: 'ascending', + }, + expression: { + fieldReferenceValue: 'foo', + }, + }, + }, + }, + ], + }, + }, + }, + }, + ]; + + testDefinitions.forEach(testDefinition => { + it(testDefinition.name, async () => { + const spy = sinon.fake.returns(stream()); + const firestore = await createInstance({ + executePipeline: spy, + }); + + await testDefinition.pipeline(firestore).execute(); + + const expectedOptions = testDefinition.expectedOptions + ? testDefinition.expectedOptions + : expectedSerializedOptions; + + expect( + spy.args[FIRST_CALL][EXECUTE_PIPELINE_REQUEST]['structuredPipeline'][ + 'pipeline' + ]['stages'][testDefinition.stageIndex ?? 0]['options'], + ).to.deep.equal(expectedOptions); + }); + }); +}); diff --git a/handwritten/firestore/types/firestore.d.ts b/handwritten/firestore/types/firestore.d.ts index 69a3554988e..af47b29f6cf 100644 --- a/handwritten/firestore/types/firestore.d.ts +++ b/handwritten/firestore/types/firestore.d.ts @@ -5603,6 +5603,65 @@ declare namespace FirebaseFirestore { */ isType(type: Type): BooleanExpression; + /** + * Evaluates if the result of this `expression` is between + * the `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * field('tireWidth').between(constant(2.2), constant(2.4)) + * + * // This is functionally equivalent to + * and(field('tireWidth').greaterThanOrEqual(contant(2.2)), field('tireWidth').lessThanOrEqual(constant(2.4))) + * ``` + * + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ + between( + lowerBound: Expression, + upperBound: Expression, + ): BooleanExpression; + + /** + * Evaluates if the result of this `expression` is between + * the `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * field('tireWidth').between(2.2, 2.4) + * + * // This is functionally equivalent to + * and(field('tireWidth').greaterThanOrEqual(2.2), field('tireWidth').lessThanOrEqual(2.4)) + * ``` + * + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ + between(lowerBound: unknown, upperBound: unknown): BooleanExpression; + + /** + * Evaluates to an HTML-formatted text snippet that renders terms matching + * the search query in `bold`. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param rquery Define the search query using the search DTS (TODO(search) link). + */ + snippet(rquery: string): Expression; + + /** + * Evaluates to an HTML-formatted text snippet that renders terms matching + * the search query in `bold`. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param options Define how snippeting behaves. + */ + snippet(options: SnippetOptions): Expression; + // TODO(new-expression): Add new expression method declarations above this line /** * @beta @@ -5831,6 +5890,26 @@ declare namespace FirebaseFirestore { * @returns The name of the field. */ get fieldName(): string; + + /** + * Perform a full-text search on this field. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param rquery Define the search query using the rquery DTS. + */ + matches(rquery: string | Expression): BooleanExpression; + + /** + * Evaluates to the distance in meters between the location specified + * by this field and the query location. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param location - Compute distance to this GeoPoint. + */ + geoDistance(location: GeoPoint | Expression): Expression; + /** * @beta * @internal @@ -11431,6 +11510,174 @@ declare namespace FirebaseFirestore { expression: Expression, type: Type, ): BooleanExpression; + /** + * Perform a full-text search on the specified field. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param searchField Search the specified field. + * @param rquery Define the search query using the search DTS. + */ + export function matches( + searchField: string | Field, + rquery: string | Expression, + ): BooleanExpression; + + /** + * Perform a full-text search on the document. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param rquery Define the search query using the rquery DTS. + */ + export function documentMatches( + rquery: string | Expression, + ): BooleanExpression; + + /** + * Evaluates to the search score that reflects the topicality of the document + * to all of the text predicates (`queryMatch`) + * in the search query. If `SearchOptions.query` is not set or does not contain + * any text predicates, then this topicality score will always be `0`. + * + * @remarks This Expression can only be used within a `Search` stage. + */ + export function score(): Expression; + + /** + * Evaluates to an HTML-formatted text snippet that highlights terms matching + * the search query in `bold`. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param searchField Search the specified field for matching terms. + * @param rquery Define the search query using the search DTS (TODO(search) link). + */ + export function snippet( + searchField: string | Field, + rquery: string, + ): Expression; + + /** + * Evaluates to an HTML-formatted text snippet that highlights terms matching + * the search query in `bold`. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param searchField Search the specified field for matching terms. + * @param options Define the search query using the search DTS (TODO(search) link). + */ + export function snippet( + searchField: string | Field, + options: SnippetOptions, + ): Expression; + + /** + * Evaluates to the distance in meters between the location in the specified + * field and the query location. + * + * @remarks This Expression can only be used within a `Search` stage. + * + * @param fieldName - Specifies the field in the document which contains + * the first GeoPoint for distance computation. + * @param location - Compute distance to this GeoPoint. + */ + export function geoDistance( + fieldName: string | Field, + location: GeoPoint | Expression, + ): Expression; + + /** + * Evaluates if the value in the field specified by `fieldName` is between + * the evaluated values for `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * between('tireWidth', constant(2.2), constant(2.4)) + * + * // This is functionally equivalent to + * and(greaterThanOrEqual('tireWidth', constant(2.2)), lessThanOrEqual('tireWidth', constant(2.4))) + * ``` + * + * @param fieldName - Evaluate if the value stored in this field is between the lower and upper bounds. + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ + export function between( + fieldName: string, + lowerBound: Expression, + upperBound: Expression, + ): BooleanExpression; + + /** + * Evaluates if the value in the field specified by `fieldName` is between + * the values for `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * between('tireWidth', 2.2, 2.4) + * + * // This is functionally equivalent to + * and(greaterThanOrEqual('tireWidth', 2.2), lessThanOrEqual('tireWidth', 2.4)) + * ``` + * + * @param fieldName - Evaluate if the value stored in this field is between the lower and upper bounds. + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ + export function between( + fieldName: string, + lowerBound: unknown, + upperBound: unknown, + ): BooleanExpression; + + /** + * Evaluates if the result of the specified `expression` is between + * the results of `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * between(field('tireWidth'), constant(2.2), constant(2.4)) + * + * // This is functionally equivalent to + * and(greaterThanOrEqual(field('tireWidth'), constant(2.2)), lessThanOrEqual(field('tireWidth'), constant(2.4))) + * ``` + * + * @param expression - Evaluate if the result of this expression is between the lower and upper bounds. + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ + export function between( + expression: Expression, + lowerBound: Expression, + upperBound: Expression, + ): BooleanExpression; + + /** + * Evaluates if the result of the specified `expression` is between + * the `lowerBound` (inclusive) and `upperBound` (inclusive). + * + * @example + * ``` + * // Evaluate if the 'tireWidth' is between 2.2 and 2.4 + * between(field('tireWidth'), 2.2, 2.4) + * + * // This is functionally equivalent to + * and(greaterThanOrEqual(field('tireWidth'), 2.2), lessThanOrEqual(field('tireWidth'), 2.4)) + * ``` + * + * @param expression - Evaluate if the result of this expression is between the lower and upper bounds. + * @param lowerBound - Lower bound (inclusive) of the range. + * @param upperBound - Upper bound (inclusive) of the range. + */ + export function between( + expression: Expression, + lowerBound: unknown, + upperBound: unknown, + ): BooleanExpression; // TODO(new-expression): Add new top-level expression function declarations above this line /** @@ -12901,6 +13148,126 @@ declare namespace FirebaseFirestore { orderings: Ordering[]; }; + /** + * Specifies if the `matches` and `snippet` expressions will enhance the user + * provided query to perform matching of synonyms, misspellings, lemmatization, + * stemming. + * + * required - search will fail if the query enhancement times out or if the query + * enhancement is not supported by the project's DRZ compliance + * requirements. + * preferred - search will fall back to the un-enhanced, user provided query, if + * the query enhancement fails. + */ + export type QueryEnhancement = 'disabled' | 'required' | 'preferred'; + + /** + * Options defining how a SearchStage is evaluated. See {@link @firebase/firestore/pipelines#Pipeline.(search)}. + */ + export type SearchStageOptions = StageOptions & { + /** + * Specifies the search query that will be used to query and score documents + * by the search stage. + * + * The query can be expressed as an `Expression`, which will be used to score + * and filter the results. Not all expressions supported by Pipelines + * are supported in the Search query. + * + * @example + * ``` + * db.pipeline().collection('restaurants').search({ + * query: or( + * documentContainsText("breakfast"), + * field('menu').containsText('waffle AND coffee') + * ) + * }) + * ``` + * + * The query can also be expressed as a string in the Search DSL: + * + * @example + * ``` + * db.pipeline().collection('restaurants').search({ + * query: 'menu:(waffle and coffee) OR breakfast' + * }) + * ``` + */ + query: BooleanExpression | string; + + /** + * The BCP-47 language code of text in the search query, such as, “en-US” or “sr-Latn” + */ + languageCode?: string; + + // TODO(search) add indexPartition after languageCode + + /** + * The maximum number of documents for the search stage to score. Documents + * will be processed in the pre-sort order specified by the search index. + */ + retrievalDepth?: number; + + /** + * Orderings specify how the input documents are sorted. + * One or more ordering are required. + */ + sort?: Ordering | Ordering[]; + + /** + * The number of documents to skip. + */ + offset?: number; + + /** + * The maximum number of documents to return from the Search stage. + */ + limit?: number; + + /** + * The fields to keep or add to each document, + * specified as an array of {@link @firebase/firestore/pipelines#Selectable}. + */ + select?: Array; + + /** + * The fields to add to each document, specified as a {@link @firebase/firestore/pipelines#Selectable}. + */ + addFields?: Selectable[]; + + /** + * Define the query expansion behavior used by full-text search expressions + * in this search stage. + */ + queryEnhancement?: QueryEnhancement; + }; + + /** + * Options defining how a snippet expression is evaluated. + */ + export type SnippetOptions = { + /** + * Define the search query using the search DTS. + */ + rquery: string; + + /** + * The maximum width of the string estimated for a variable width font. The + * unit is tenths of ems. The default is `160`. + */ + maxSnippetWidth?: number; + + /** + * The maximum number of non-contiguous pieces of text in the returned snippet. + * The default is `1`. + */ + maxSnippets?: number; + + /** + * The string to join the pieces. The default value is '\n' + */ + separator?: string; + }; + /** * @beta * Represents a field value within the explain statistics, which can be a primitive type (null, string, number, boolean)