From d2cb984482fc77cc99260055ee5b33b1a4c79f38 Mon Sep 17 00:00:00 2001 From: Akshai Sarma Date: Fri, 1 Feb 2019 16:17:18 -0800 Subject: [PATCH] Fixing partition removal bug --- .../yahoo/bullet/querying/QueryManager.java | 7 +++- .../bullet/querying/QueryManagerTest.java | 36 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/yahoo/bullet/querying/QueryManager.java b/src/main/java/com/yahoo/bullet/querying/QueryManager.java index 83db23ac..329ee96c 100644 --- a/src/main/java/com/yahoo/bullet/querying/QueryManager.java +++ b/src/main/java/com/yahoo/bullet/querying/QueryManager.java @@ -130,7 +130,12 @@ public Querier removeAndGetQuery(String id) { Query query = querier.getQuery(); Set keys = partitioner.getKeys(query); for (String key : keys) { - partitioning.get(key).remove(id); + Set partition = partitioning.get(key); + partition.remove(id); + if (partition.isEmpty()) { + log.debug("Partition: {} is empty. Removing...", key); + partitioning.remove(key); + } log.debug("Removed query: {} from partition: {}", id, key); } } diff --git a/src/test/java/com/yahoo/bullet/querying/QueryManagerTest.java b/src/test/java/com/yahoo/bullet/querying/QueryManagerTest.java index d2999291..512be6fa 100644 --- a/src/test/java/com/yahoo/bullet/querying/QueryManagerTest.java +++ b/src/test/java/com/yahoo/bullet/querying/QueryManagerTest.java @@ -372,4 +372,40 @@ public void testQuerySeeingStatistics() { Assert.assertEquals(stats.get(QueryManager.PartitionStat.ACTUAL_QUERIES_SEEN), 6L); Assert.assertEquals(stats.get(QueryManager.PartitionStat.EXPECTED_QUERIES_SEEN), 9L); } + + @Test + public void testPartitionRemoval() { + QueryManager manager = new QueryManager(getEqualityPartitionerConfig("A", "B")); + Query queryA = getQuery(ImmutablePair.of("A", "foo")); + Query queryB = getQuery(ImmutablePair.of("A", "foo"), ImmutablePair.of("B", "bar")); + Query queryC = getQuery(); + manager.addQuery("idA", getQuerier(queryA)); + manager.addQuery("idB", getQuerier(queryB)); + manager.addQuery("idC", getQuerier(queryC)); + manager.addQuery("idD", getQuerier(queryC)); + + Map stats = manager.getStats(); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.QUERY_COUNT), 4); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.PARTITION_COUNT), 3); + + manager.removeAndGetQuery("idA"); + stats = manager.getStats(); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.QUERY_COUNT), 3); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.PARTITION_COUNT), 2); + + manager.removeAndGetQuery("idB"); + stats = manager.getStats(); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.QUERY_COUNT), 2); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.PARTITION_COUNT), 1); + + manager.removeAndGetQuery("idD"); + stats = manager.getStats(); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.QUERY_COUNT), 1); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.PARTITION_COUNT), 1); + + manager.removeAndGetQuery("idC"); + stats = manager.getStats(); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.QUERY_COUNT), 0); + Assert.assertEquals(stats.get(QueryManager.PartitionStat.PARTITION_COUNT), 0); + } }