From 2055b04d73df7ae3751324f384dbf1183744921e Mon Sep 17 00:00:00 2001 From: Kristy Lee <18691669+kristyelee@users.noreply.github.com> Date: Wed, 4 Dec 2024 16:45:27 -0800 Subject: [PATCH] Partition Removal --- .../src/main/java/com/linkedin/davinci/DaVinciBackend.java | 6 ++++-- .../com/linkedin/venice/endToEnd/DaVinciClientTest.java | 1 - 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java index 2acc34a0ad..6266b9efc6 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java @@ -475,8 +475,10 @@ private synchronized void bootstrap() { AbstractStorageEngine storageEngine = storageService.getStorageEngine(versionTopic); aggVersionedStorageEngineStats.setStorageEngine(versionTopic, storageEngine); StoreBackend storeBackend = getStoreOrThrow(storeName); - ComplementSet persistedPartitionIds = ComplementSet.newSet(storageEngine.getPersistedPartitionIds()); - storeBackend.unsubscribe(persistedPartitionIds); + ComplementSet subscription = ComplementSet.newSet(storeBackend.getSubscription()); + ComplementSet unassignedPartitionSet = ComplementSet.newSet(storageEngine.getPersistedPartitionIds()); + unassignedPartitionSet.removeAll(subscription); + storeBackend.unsubscribe(unassignedPartitionSet); storeBackend.subscribe(ComplementSet.newSet(partitions), Optional.of(version)); }); } diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java index 6f8171d169..a08ac9cfbc 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java @@ -1288,7 +1288,6 @@ public void testBootstrapSubscription(DaVinciConfig daVinciConfig) throws Except try (DaVinciClient client = ServiceFactory .getGenericAvroDaVinciClient(storeName, cluster.getZk().getAddress(), daVinciConfig, backendConfig)) { - // We only subscribe to 1/3 of the partitions so some data will not be present locally. client.subscribe(Collections.singleton(0)).get(); assertThrows(() -> client.batchGet(keySet).get()); }