Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit a075c99

Browse files
feat(firestore): Added search stage support for languageCode, offset, limit, and retrievalDepth (#2373)
* feat(firestore): Added search stage support for languageCode, offset, limit, and retrievalDepth * chore: generate libraries at Mon May 4 19:18:21 UTC 2026 --------- Co-authored-by: cloud-java-bot <[email protected]>
1 parent 26c830a commit a075c99

3 files changed

Lines changed: 105 additions & 100 deletions

File tree

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/Search.java

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -148,32 +148,33 @@ public Search withSort(Ordering order, Ordering... additionalOrderings) {
148148
options.with("sort", Lists.transform(Arrays.asList(allOrderings), Ordering::toProto)));
149149
}
150150

151-
// TODO(search) enable with backend support
152-
// /** Specify the maximum number of documents to return from the Search stage. */
153-
// public Search withLimit(long limit) {
154-
// return new Search(options.with("limit", encodeValue(limit)));
155-
// }
151+
/**
152+
* Specify the maximum number of documents to return from the `search` stage. The `limit` is
153+
* applied after documents are scored and sorted.
154+
*/
155+
public Search withLimit(long limit) {
156+
return new Search(options.with("limit", encodeValue(limit)));
157+
}
156158

157159
/**
158-
* Specify the maximum number of documents to retrieve. Documents will be retrieved in the
159-
* pre-sort order specified by the search index.
160+
* Specify the maximum number of documents to retrieve from the search index. Documents will be
161+
* retrieved in the pre-sort order specified by the search index. The `retrievalDepth` is a limit
162+
* applied before documents are scored and sorted, which can reduce costs of expensive scoring and
163+
* sorting operations.
160164
*/
161165
public Search withRetrievalDepth(long retrievalDepth) {
162166
return new Search(options.with("retrieval_depth", encodeValue(retrievalDepth)));
163167
}
164168

165-
// TODO(search) enable with backend support
166-
// /** Specify the number of documents to skip. */
167-
// public Search withOffset(long offset) {
168-
// return new Search(options.with("offset", encodeValue(offset)));
169-
// }
169+
/** Specify the number of documents to skip from the beginning of the search result set. */
170+
public Search withOffset(long offset) {
171+
return new Search(options.with("offset", encodeValue(offset)));
172+
}
170173

171-
// TODO(search) enable with backend support
172-
// /** Specify the BCP-47 language code of text in the search query, such as, “en-US” or
173-
// “sr-Latn” */
174-
// public Search withLanguageCode(String value) {
175-
// return new Search(options.with("language_code", encodeValue(value)));
176-
// }
174+
/** Specify the BCP-47 language code of text in the search query, such as “en” or “sr”. */
175+
public Search withLanguageCode(String value) {
176+
return new Search(options.with("language_code", encodeValue(value)));
177+
}
177178

178179
// TODO(search) enable with backend support
179180
// /**

google-cloud-firestore/src/test/java/com/google/cloud/firestore/PipelineProtoTest.java

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,11 @@ public void testSearchStageProtoEncoding() {
4545
.collection("foo")
4646
.search(
4747
Search.withQuery("foo")
48-
// .withLimit(1)
48+
.withLimit(1)
4949
.withRetrievalDepth(2)
50-
// .withOffset(3)
50+
.withOffset(3)
5151
// .withQueryEnhancement(Search.QueryEnhancement.REQUIRED)
52-
// .withLanguageCode("en-US")
52+
.withLanguageCode("en-US")
5353
.withSort(field("foo").ascending())
5454
.withAddFields(constant(true).as("bar")));
5555
// .withSelect(field("id")));
@@ -73,19 +73,19 @@ public void testSearchStageProtoEncoding() {
7373
assertThat(query.getFunctionValue().getArgs(0).getStringValue()).isEqualTo("foo");
7474

7575
// limit
76-
// assertThat(optionsMap.get("limit").getIntegerValue()).isEqualTo(1L);
76+
assertThat(optionsMap.get("limit").getIntegerValue()).isEqualTo(1L);
7777

7878
// retrieval_depth
7979
assertThat(optionsMap.get("retrieval_depth").getIntegerValue()).isEqualTo(2L);
8080

8181
// offset
82-
// assertThat(optionsMap.get("offset").getIntegerValue()).isEqualTo(3L);
82+
assertThat(optionsMap.get("offset").getIntegerValue()).isEqualTo(3L);
8383

8484
// query_enhancement
8585
// assertThat(optionsMap.get("query_enhancement").getStringValue()).isEqualTo("required");
8686

8787
// language_code
88-
// assertThat(optionsMap.get("language_code").getStringValue()).isEqualTo("en-US");
88+
assertThat(optionsMap.get("language_code").getStringValue()).isEqualTo("en-US");
8989

9090
// select
9191
// Value select = optionsMap.get("select");

google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineSearchTest.java

Lines changed: 80 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,10 @@
3535
import java.util.List;
3636
import java.util.Map;
3737
import java.util.Objects;
38+
import java.util.concurrent.ExecutionException;
3839
import java.util.concurrent.TimeUnit;
3940
import java.util.stream.Collectors;
41+
import org.junit.Assert;
4042
import org.junit.BeforeClass;
4143
import org.junit.Test;
4244
import org.junit.runner.RunWith;
@@ -238,21 +240,28 @@ private void assertResultIds(Pipeline.Snapshot snapshot, String... ids) {
238240
// --- DISABLE query expansion ---
239241

240242
// query
241-
// TODO(search) enable with backend support
242-
// @Test
243-
// public void searchWithLanguageCode() throws Exception {
244-
// Pipeline pipeline =
245-
// firestore
246-
// .pipeline()
247-
// .collection(COLLECTION_NAME)
248-
// .search(
249-
// Search.withQuery("waffles")
250-
// .withLanguageCode("en")
251-
// .withQueryEnhancement(Search.QueryEnhancement.DISABLED));
252-
//
253-
// Pipeline.Snapshot snapshot = pipeline.execute().get();
254-
// assertResultIds(snapshot, "goldenWaffle");
255-
// }
243+
@Test
244+
public void searchWithLanguageCode() throws Exception {
245+
Pipeline pipeline =
246+
firestore
247+
.pipeline()
248+
.collection(COLLECTION_NAME)
249+
.search(Search.withQuery("waffles").withLanguageCode("en"));
250+
251+
Pipeline.Snapshot snapshot = pipeline.execute().get();
252+
assertResultIds(snapshot, "goldenWaffle");
253+
}
254+
255+
@Test
256+
public void searchWithInvalidLanguageCode() throws Exception {
257+
Pipeline pipeline =
258+
firestore
259+
.pipeline()
260+
.collection(COLLECTION_NAME)
261+
.search(Search.withQuery("waffles").withLanguageCode("unknown"));
262+
263+
Assert.assertThrows(ExecutionException.class, () -> pipeline.execute().get());
264+
}
256265

257266
@Test
258267
public void searchFullDocument() throws Exception {
@@ -584,69 +593,64 @@ public void sort_byDistance() throws Exception {
584593
// }
585594

586595
// limit
587-
// TODO(search) enable with backend support
588-
// @Test
589-
// public void limit_limitsTheNumberOfDocumentsReturned() throws Exception {
590-
// Pipeline pipeline =
591-
// firestore
592-
// .pipeline()
593-
// .collection(COLLECTION_NAME)
594-
// .search(
595-
// Search.withQuery(constant(true))
596-
// .withSort(
597-
// field("location").geoDistance(new GeoPoint(39.6985,
598-
// -105.024)).ascending())
599-
// .withLimit(5)
600-
// .withQueryEnhancement(Search.QueryEnhancement.DISABLED));
601-
//
602-
// Pipeline.Snapshot snapshot = pipeline.execute().get();
603-
// assertResultIds(snapshot, "solTacos", "lotusBlossomThai", "goldenWaffle");
604-
// }
596+
@Test
597+
public void limit_limitsTheNumberOfDocumentsReturned() throws Exception {
598+
Pipeline pipeline =
599+
firestore
600+
.pipeline()
601+
.collection(COLLECTION_NAME)
602+
.search(
603+
Search.withQuery(
604+
field("location")
605+
.geoDistance(new GeoPoint(39.6985, -105.024))
606+
.lessThanOrEqual(100000000))
607+
.withSort(geoDistance("location", new GeoPoint(39.6985, -105.024)).ascending())
608+
.withLimit(3));
605609

606-
// @Test
607-
// public void limit_limitsTheNumberOfDocumentsRetrieved() throws Exception {
608-
// Pipeline pipeline =
609-
// firestore
610-
// .pipeline()
611-
// .collection(COLLECTION_NAME)
612-
// .search(
613-
// Search.withQuery(documentMatches("chicken"))
614-
// .withRetrievalDepth(3));
615-
//// .withQueryEnhancement(Search.QueryEnhancement.DISABLED));
616-
//
617-
// Pipeline.Snapshot snapshot = pipeline.execute().get();
618-
// assertResultIds(snapshot, "eastsideChicken", "lotusBlossomThai", "goldenWaffle");
619-
//
620-
// pipeline =
621-
// firestore
622-
// .pipeline()
623-
// .collection(COLLECTION_NAME)
624-
// .search(
625-
// Search.withQuery(documentMatches("chicken")));
626-
//// .withQueryEnhancement(Search.QueryEnhancement.DISABLED));
627-
//
628-
// snapshot = pipeline.execute().get();
629-
// assertResultIds(snapshot, "eastsideChicken", "lotusBlossomThai", "goldenWaffle",
630-
// "eastsideCantina");
631-
// }
610+
Pipeline.Snapshot snapshot = pipeline.execute().get();
611+
assertResultIds(snapshot, "solTacos", "lotusBlossomThai", "mileHighCatch");
612+
}
632613

633-
// offset
634-
// TODO(search) enable with backend support
635-
// @Test
636-
// public void offset_skipsNDocuments() throws Exception {
637-
// Pipeline pipeline =
638-
// firestore
639-
// .pipeline()
640-
// .collection(COLLECTION_NAME)
641-
// .search(
642-
// Search.withQuery(constant(true))
643-
// .withLimit(2)
644-
// .withOffset(2)
645-
// .withQueryEnhancement(Search.QueryEnhancement.DISABLED));
646-
//
647-
// Pipeline.Snapshot snapshot = pipeline.execute().get();
648-
// assertResultIds(snapshot, "eastsideChicken", "eastsideTacos");
649-
// }
614+
@Test
615+
public void limit_limitsTheNumberOfDocumentsScoredViaRetrievalDepth() throws Exception {
616+
Pipeline pipeline =
617+
firestore
618+
.pipeline()
619+
.collection(COLLECTION_NAME)
620+
.search(
621+
Search.withQuery(documentMatches("taco"))
622+
.withAddFields(score().as("score"))
623+
.withSort(score().descending())
624+
.withRetrievalDepth(2));
625+
626+
Pipeline.Snapshot snapshot = pipeline.execute().get();
627+
assertResultIds(snapshot, "solTacos", "eastsideTacos");
628+
629+
pipeline =
630+
firestore
631+
.pipeline()
632+
.collection(COLLECTION_NAME)
633+
.search(
634+
Search.withQuery(documentMatches("taco"))
635+
.withAddFields(score().as("score"))
636+
.withSort(score().descending())
637+
.withRetrievalDepth(1));
638+
639+
snapshot = pipeline.execute().get();
640+
assertResultIds(snapshot, "eastsideTacos");
641+
}
642+
643+
@Test
644+
public void offset_skipsNDocuments() throws Exception {
645+
Pipeline pipeline =
646+
firestore
647+
.pipeline()
648+
.collection(COLLECTION_NAME)
649+
.search(Search.withQuery("chicken").withLimit(2).withOffset(2));
650+
651+
Pipeline.Snapshot snapshot = pipeline.execute().get();
652+
assertResultIds(snapshot, "goldenWaffle");
653+
}
650654

651655
// =========================================================================
652656
// Snippet

0 commit comments

Comments
 (0)