Use topic level throughput information when partition level information is unavailable.

datastream-server/src/main/java/com/linkedin/datastream/server/assignment/LoadBasedPartitionAssigner.java

@@ -95,7 +95,15 @@ public Map<String, Set<DatastreamTask>> assignPartitions(
         LoadBasedPartitionAssignmentStrategyConfig.DEFAULT_PARTITION_MESSAGES_IN_RATE, "");
     newPartitions.forEach((task, partitions) -> {
       int totalThroughput = partitions.stream()
-          .mapToInt(p -> partitionInfoMap.getOrDefault(p, defaultPartitionInfo).getBytesInKBRate())
+          .mapToInt(p ->  {
+            int index = p.lastIndexOf('-');
+            String topic = p;
+            if (index > -1) {
+              topic = p.substring(0, index);
+            }
+            PartitionThroughputInfo defaultValue = partitionInfoMap.getOrDefault(topic, defaultPartitionInfo);
+            return partitionInfoMap.getOrDefault(p, defaultValue).getBytesInKBRate();
+          })
           .sum();
       taskThroughputMap.put(task, totalThroughput);
     });

datastream-server/src/main/java/com/linkedin/datastream/server/assignment/LoadBasedTaskCountEstimator.java

@@ -65,8 +65,15 @@ public int getTaskCount(ClusterThroughputInfo throughputInfo, List<String> assig
         LoadBasedPartitionAssignmentStrategyConfig.DEFAULT_PARTITION_MESSAGES_IN_RATE, "");
     // total throughput in KB/sec
     int totalThroughput = allPartitions.stream()
-        .map(p -> throughputMap.getOrDefault(p, defaultThroughputInfo))
-        .mapToInt(PartitionThroughputInfo::getBytesInKBRate)
+        .mapToInt(p ->  {
+          int index = p.lastIndexOf('-');
+          String topic = p;
+          if (index > -1) {
+            topic = p.substring(0, index);
+          }
+          PartitionThroughputInfo defaultValue = throughputMap.getOrDefault(topic, defaultThroughputInfo);
+          return throughputMap.getOrDefault(p, defaultValue).getBytesInKBRate();
+        })
         .sum();
     LOG.info("Total throughput in all {} partitions: {}KB/sec", allPartitions.size(), totalThroughput);
 

datastream-server/src/test/java/com/linkedin/datastream/server/TestLoadBasedTaskCountEstimator.java

@@ -97,4 +97,15 @@ public void partitionsHaveDefaultWeightTest() {
     int taskCount = estimator.getTaskCount(throughputInfo, assignedPartitions, unassignedPartitions);
     Assert.assertTrue(taskCount > 0);
   }
+
+  @Test
+  public void throughputTaskEstimatorWithTopicLevelInformation() {
+    ClusterThroughputInfo throughputInfo = _provider.getThroughputInfo("fruit");
+    List<String> assignedPartitions = Collections.emptyList();
+    List<String> unassignedPartitions = Arrays.asList("apple-0", "apple-1", "apple-2", "banana-0");
+    LoadBasedTaskCountEstimator estimator = new LoadBasedTaskCountEstimator(TASK_CAPACITY_MBPS,
+        TASK_CAPACITY_UTILIZATION_PCT);
+    int taskCount = estimator.getTaskCount(throughputInfo, assignedPartitions, unassignedPartitions);
+    Assert.assertEquals(taskCount, 4);
+  }
 }

datastream-server/src/test/resources/partitionThroughput.json

@@ -227,6 +227,10 @@
     "donut" : {
       "BostonCreme-1" : "bytesInKB: 5000, msgIn:200",
       "BostonCreme-2" : "bytesInKB: 5000, msgIn:200"
+    },
+    "fruit" : {
+      "apple" : "bytesInKB: 10000, msgIn:300",
+      "apple-2" : "bytesInKb: 8000, msgIn:200"
     }
   }
 }