From 7f1d1b58641a127e1295019a3e4490c6ec1261cb Mon Sep 17 00:00:00 2001 From: swethakann Date: Tue, 26 Nov 2024 14:52:56 -0800 Subject: [PATCH] Add default terminate after max recall count (#790) Add default terminate after max recall count --- .../proto/yelp/nrtsearch/luceneserver.proto | 4 ++ .../server/handler/LiveSettingsHandler.java | 6 +++ .../server/index/ImmutableIndexState.java | 12 ++++++ .../nrtsearch/server/index/IndexState.java | 3 ++ .../server/search/TerminateAfterWrapper.java | 5 +++ .../search/collectors/DocCollector.java | 2 +- .../tools/cli/LiveSettingsV2Command.java | 9 ++++ .../server/grpc/StateBackendServerTest.java | 5 +++ .../server/index/ImmutableIndexStateTest.java | 20 +++++++++ .../search/collectors/DocCollectorTest.java | 42 ++++++++++++++++++- 10 files changed, 106 insertions(+), 2 deletions(-) diff --git a/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto b/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto index 4cbd7881e..01eda8566 100644 --- a/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto +++ b/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto @@ -456,6 +456,8 @@ message LiveSettingsRequest { int32 defaultSearchTimeoutCheckEvery = 13; // Terminate after value to use when not specified in the search request int32 defaultTerminateAfter = 14; + //Terminate after max recall count value to use when not specified in the search request. + int32 defaultTerminateAfterMaxRecallCount = 15; } // Response from Server to liveSettings @@ -1446,6 +1448,8 @@ message IndexLiveSettings { google.protobuf.BoolValue parallelFetchByField = 16; // The number of documents/fields per parallel fetch task, default: 50 google.protobuf.Int32Value parallelFetchChunkSize = 17; + // Terminate after max recall count value to use when not specified in the search request, or 0 for none, default: 0 + google.protobuf.Int32Value defaultTerminateAfterMaxRecallCount = 18; } // Index state diff --git a/src/main/java/com/yelp/nrtsearch/server/handler/LiveSettingsHandler.java b/src/main/java/com/yelp/nrtsearch/server/handler/LiveSettingsHandler.java index e99baba0f..c8d23e911 100644 --- a/src/main/java/com/yelp/nrtsearch/server/handler/LiveSettingsHandler.java +++ b/src/main/java/com/yelp/nrtsearch/server/handler/LiveSettingsHandler.java @@ -133,6 +133,12 @@ private LiveSettingsResponse handleAsLiveSettingsV2( .setValue(liveSettingsRequest.getDefaultTerminateAfter()) .build()); } + if (liveSettingsRequest.getDefaultTerminateAfterMaxRecallCount() >= 0) { + settingsBuilder.setDefaultTerminateAfterMaxRecallCount( + Int32Value.newBuilder() + .setValue(liveSettingsRequest.getDefaultTerminateAfterMaxRecallCount()) + .build()); + } try { updatedSettings = indexStateManager.updateLiveSettings(settingsBuilder.build(), false); } catch (IOException e) { diff --git a/src/main/java/com/yelp/nrtsearch/server/index/ImmutableIndexState.java b/src/main/java/com/yelp/nrtsearch/server/index/ImmutableIndexState.java index 533d78067..ee9d55e78 100644 --- a/src/main/java/com/yelp/nrtsearch/server/index/ImmutableIndexState.java +++ b/src/main/java/com/yelp/nrtsearch/server/index/ImmutableIndexState.java @@ -146,6 +146,7 @@ public class ImmutableIndexState extends IndexState { .setDefaultSearchTimeoutSec(DoubleValue.newBuilder().setValue(0).build()) .setDefaultSearchTimeoutCheckEvery(Int32Value.newBuilder().setValue(0).build()) .setDefaultTerminateAfter(Int32Value.newBuilder().setValue(0).build()) + .setDefaultTerminateAfterMaxRecallCount(Int32Value.newBuilder().setValue(0).build()) .setMaxMergePreCopyDurationSec(UInt64Value.newBuilder().setValue(0)) .setVerboseMetrics(BoolValue.newBuilder().setValue(false).build()) .setParallelFetchByField(BoolValue.newBuilder().setValue(false).build()) @@ -167,6 +168,7 @@ public class ImmutableIndexState extends IndexState { private final double defaultSearchTimeoutSec; private final int defaultSearchTimeoutCheckEvery; private final int defaultTerminateAfter; + private final int defaultTerminateAfterMaxRecallCount; private final long maxMergePreCopyDurationSec; private final boolean verboseMetrics; private final ParallelFetchConfig parallelFetchConfig; @@ -262,6 +264,8 @@ public ImmutableIndexState( defaultSearchTimeoutCheckEvery = mergedLiveSettingsWithLocal.getDefaultSearchTimeoutCheckEvery().getValue(); defaultTerminateAfter = mergedLiveSettingsWithLocal.getDefaultTerminateAfter().getValue(); + defaultTerminateAfterMaxRecallCount = + mergedLiveSettingsWithLocal.getDefaultTerminateAfterMaxRecallCount().getValue(); maxMergePreCopyDurationSec = mergedLiveSettingsWithLocal.getMaxMergePreCopyDurationSec().getValue(); verboseMetrics = mergedLiveSettingsWithLocal.getVerboseMetrics().getValue(); @@ -718,6 +722,11 @@ public int getDefaultTerminateAfter() { return defaultTerminateAfter; } + @Override + public int getDefaultTerminateAfterMaxRecallCount() { + return defaultTerminateAfterMaxRecallCount; + } + @Override public int getDefaultSearchTimeoutCheckEvery() { return defaultSearchTimeoutCheckEvery; @@ -809,6 +818,9 @@ static void validateLiveSettings(IndexLiveSettings liveSettings) { if (liveSettings.getDefaultTerminateAfter().getValue() < 0) { throw new IllegalArgumentException("defaultTerminateAfter must be >= 0"); } + if (liveSettings.getDefaultTerminateAfterMaxRecallCount().getValue() < 0) { + throw new IllegalArgumentException("defaultTerminateAfterMaxRecallCount must be >= 0"); + } if (liveSettings.getMaxMergePreCopyDurationSec().getValue() < 0) { throw new IllegalArgumentException("maxMergePreCopyDurationSec must be >= 0"); } diff --git a/src/main/java/com/yelp/nrtsearch/server/index/IndexState.java b/src/main/java/com/yelp/nrtsearch/server/index/IndexState.java index 8c07128c2..75452bd82 100644 --- a/src/main/java/com/yelp/nrtsearch/server/index/IndexState.java +++ b/src/main/java/com/yelp/nrtsearch/server/index/IndexState.java @@ -487,6 +487,9 @@ public abstract IndexWriterConfig getIndexWriterConfig( /** Get the default terminate after. */ public abstract int getDefaultTerminateAfter(); + /** Get the default terminate after max recall count. */ + public abstract int getDefaultTerminateAfterMaxRecallCount(); + /** Get the default search timeout check every. */ public abstract int getDefaultSearchTimeoutCheckEvery(); diff --git a/src/main/java/com/yelp/nrtsearch/server/search/TerminateAfterWrapper.java b/src/main/java/com/yelp/nrtsearch/server/search/TerminateAfterWrapper.java index b5fbac459..6de20cabd 100644 --- a/src/main/java/com/yelp/nrtsearch/server/search/TerminateAfterWrapper.java +++ b/src/main/java/com/yelp/nrtsearch/server/search/TerminateAfterWrapper.java @@ -97,6 +97,11 @@ public int getTerminateAfter() { return terminateAfter; } + /** Max documents to count beyond terminateAfter. */ + public int getTerminateAfterMaxRecallCount() { + return terminateAfterMaxRecallCount; + } + /** * {@link Collector} implementation that wraps another collector and terminates collection after a * certain global count of documents is reached. diff --git a/src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java b/src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java index d170f3999..133cb1b1b 100644 --- a/src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java +++ b/src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java @@ -179,7 +179,7 @@ CollectorManager wrap int terminateAfterMaxRecallCount = request.getTerminateAfterMaxRecallCount() > 0 ? request.getTerminateAfterMaxRecallCount() - : 0; + : indexState.getDefaultTerminateAfterMaxRecallCount(); if (terminateAfter > 0) { wrapped = new TerminateAfterWrapper<>( diff --git a/src/main/java/com/yelp/nrtsearch/tools/cli/LiveSettingsV2Command.java b/src/main/java/com/yelp/nrtsearch/tools/cli/LiveSettingsV2Command.java index 92814de2f..012c5ef93 100644 --- a/src/main/java/com/yelp/nrtsearch/tools/cli/LiveSettingsV2Command.java +++ b/src/main/java/com/yelp/nrtsearch/tools/cli/LiveSettingsV2Command.java @@ -106,6 +106,11 @@ public class LiveSettingsV2Command implements Callable { description = "Terminate after to use when not provided by the request") private Integer defaultTerminateAfter; + @CommandLine.Option( + names = {"--defaultTerminateAfterMaxRecallCount"}, + description = "Terminate after max recall count to use when not provided by the request") + private Integer defaultTerminateAfterMaxRecallCount; + @CommandLine.Option( names = {"--maxMergePreCopyDurationSec"}, description = "Maximum time allowed for merge precopy in seconds") @@ -192,6 +197,10 @@ public Integer call() throws Exception { liveSettingsBuilder.setDefaultTerminateAfter( Int32Value.newBuilder().setValue(defaultTerminateAfter).build()); } + if (defaultTerminateAfterMaxRecallCount != null) { + liveSettingsBuilder.setDefaultTerminateAfterMaxRecallCount( + Int32Value.newBuilder().setValue(defaultTerminateAfterMaxRecallCount).build()); + } if (maxMergePreCopyDurationSec != null) { liveSettingsBuilder.setMaxMergePreCopyDurationSec( UInt64Value.newBuilder().setValue(maxMergePreCopyDurationSec)); diff --git a/src/test/java/com/yelp/nrtsearch/server/grpc/StateBackendServerTest.java b/src/test/java/com/yelp/nrtsearch/server/grpc/StateBackendServerTest.java index 40aa1eb1e..7cca7192b 100644 --- a/src/test/java/com/yelp/nrtsearch/server/grpc/StateBackendServerTest.java +++ b/src/test/java/com/yelp/nrtsearch/server/grpc/StateBackendServerTest.java @@ -775,6 +775,8 @@ public void testSetIndexLiveSettings() throws IOException { IndexLiveSettings.newBuilder() .setDefaultTerminateAfter( Int32Value.newBuilder().setValue(1000).build()) + .setDefaultTerminateAfterMaxRecallCount( + Int32Value.newBuilder().setValue(1000).build()) .setSegmentsPerTier(Int32Value.newBuilder().setValue(4).build()) .setSliceMaxSegments(Int32Value.newBuilder().setValue(50).build()) .setDefaultSearchTimeoutSec( @@ -784,6 +786,7 @@ public void testSetIndexLiveSettings() throws IOException { IndexLiveSettings expectedSettings = ImmutableIndexState.DEFAULT_INDEX_LIVE_SETTINGS.toBuilder() .setDefaultTerminateAfter(Int32Value.newBuilder().setValue(1000).build()) + .setDefaultTerminateAfterMaxRecallCount(Int32Value.newBuilder().setValue(1000).build()) .setSegmentsPerTier(Int32Value.newBuilder().setValue(4).build()) .setSliceMaxSegments(Int32Value.newBuilder().setValue(50).build()) .setDefaultSearchTimeoutSec(DoubleValue.newBuilder().setValue(5.1).build()) @@ -1647,6 +1650,7 @@ public void testLiveSettingsV1All() throws IOException { .setDefaultSearchTimeoutSec(13.0) .setDefaultSearchTimeoutCheckEvery(500) .setDefaultTerminateAfter(5000) + .setDefaultTerminateAfterMaxRecallCount(6000) .build(); LiveSettingsResponse response = primaryClient.getBlockingStub().liveSettings(request); @@ -1665,6 +1669,7 @@ public void testLiveSettingsV1All() throws IOException { .setDefaultSearchTimeoutSec(DoubleValue.newBuilder().setValue(13.0).build()) .setDefaultSearchTimeoutCheckEvery(Int32Value.newBuilder().setValue(500).build()) .setDefaultTerminateAfter(Int32Value.newBuilder().setValue(5000).build()) + .setDefaultTerminateAfterMaxRecallCount(Int32Value.newBuilder().setValue(6000).build()) .setMaxMergePreCopyDurationSec(UInt64Value.newBuilder().setValue(0)) .setVerboseMetrics(BoolValue.newBuilder().setValue(false).build()) .setParallelFetchByField(BoolValue.newBuilder().setValue(false).build()) diff --git a/src/test/java/com/yelp/nrtsearch/server/index/ImmutableIndexStateTest.java b/src/test/java/com/yelp/nrtsearch/server/index/ImmutableIndexStateTest.java index c8a6e1019..becc25fc5 100644 --- a/src/test/java/com/yelp/nrtsearch/server/index/ImmutableIndexStateTest.java +++ b/src/test/java/com/yelp/nrtsearch/server/index/ImmutableIndexStateTest.java @@ -781,6 +781,26 @@ public void testDefaultTerminateAfter_invalid() throws IOException { assertLiveSettingException(expectedMsg, b -> b.setDefaultTerminateAfter(wrap(-1))); } + @Test + public void testDefaultTerminateAfterMaxRecallCount_default() throws IOException { + assertEquals(0, getIndexState(getEmptyState()).getDefaultTerminateAfterMaxRecallCount()); + } + + @Test + public void testDefaultTerminateAfterMaxRecallCount_set() throws IOException { + verifyIntLiveSetting( + 100, + ImmutableIndexState::getDefaultTerminateAfterMaxRecallCount, + b -> b.setDefaultTerminateAfterMaxRecallCount(wrap(100))); + } + + @Test + public void testDefaultTerminateAfterMaxRecallCount_invalid() throws IOException { + String expectedMsg = "defaultTerminateAfterMaxRecallCount must be >= 0"; + assertLiveSettingException( + expectedMsg, b -> b.setDefaultTerminateAfterMaxRecallCount(wrap(-1))); + } + @Test public void testMaxMergePreCopyDurationSec_default() throws IOException { assertEquals(0, getIndexState(getEmptyState()).getMaxMergePreCopyDurationSec()); diff --git a/src/test/java/com/yelp/nrtsearch/server/search/collectors/DocCollectorTest.java b/src/test/java/com/yelp/nrtsearch/server/search/collectors/DocCollectorTest.java index 7de62e096..09e24ba0a 100644 --- a/src/test/java/com/yelp/nrtsearch/server/search/collectors/DocCollectorTest.java +++ b/src/test/java/com/yelp/nrtsearch/server/search/collectors/DocCollectorTest.java @@ -212,12 +212,21 @@ public void testNumHitsToCollect() { @Test public void testHasTerminateAfterWrapper() { - SearchRequest request = SearchRequest.newBuilder().setTopHits(10).setTerminateAfter(5).build(); + SearchRequest request = + SearchRequest.newBuilder() + .setTopHits(10) + .setTerminateAfter(5) + .setTerminateAfterMaxRecallCount(10) + .build(); TestDocCollector docCollector = new TestDocCollector(request); assertTrue(docCollector.getManager() instanceof TestDocCollector.TestCollectorManager); assertTrue(docCollector.getWrappedManager() instanceof TerminateAfterWrapper); assertEquals( 5, ((TerminateAfterWrapper) docCollector.getWrappedManager()).getTerminateAfter()); + assertEquals( + 10, + ((TerminateAfterWrapper) docCollector.getWrappedManager()) + .getTerminateAfterMaxRecallCount()); } @Test @@ -246,6 +255,37 @@ public void testOverrideDefaultTerminateAfter() { 75, ((TerminateAfterWrapper) docCollector.getWrappedManager()).getTerminateAfter()); } + @Test + public void testUsesDefaultTerminateAfterMaxRecallCount() { + IndexState indexState = Mockito.mock(IndexState.class); + when(indexState.getDefaultTerminateAfter()).thenReturn(100); + when(indexState.getDefaultTerminateAfterMaxRecallCount()).thenReturn(1000); + + SearchRequest request = SearchRequest.newBuilder().setTopHits(10).build(); + TestDocCollector docCollector = new TestDocCollector(request, indexState); + assertEquals( + 1000, + ((TerminateAfterWrapper) docCollector.getWrappedManager()) + .getTerminateAfterMaxRecallCount()); + } + + @Test + public void testOverrideDefaultTerminateAfterMaxRecallCount() { + IndexState indexState = Mockito.mock(IndexState.class); + when(indexState.getDefaultTerminateAfter()).thenReturn(100); + when(indexState.getDefaultTerminateAfterMaxRecallCount()).thenReturn(1000); + + SearchRequest request = + SearchRequest.newBuilder().setTopHits(10).setTerminateAfterMaxRecallCount(75).build(); + TestDocCollector docCollector = new TestDocCollector(request, indexState); + assertTrue(docCollector.getManager() instanceof TestDocCollector.TestCollectorManager); + assertTrue(docCollector.getWrappedManager() instanceof TerminateAfterWrapper); + assertEquals( + 75, + ((TerminateAfterWrapper) docCollector.getWrappedManager()) + .getTerminateAfterMaxRecallCount()); + } + @Test public void testWithAllWrappers() { SearchRequest request =