diff --git a/src/main/java/org/apache/flink/benchmark/AdaptivePartitionThroughputBenchmark.java b/src/main/java/org/apache/flink/benchmark/AdaptivePartitionThroughputBenchmark.java
index 7426032..1b205fa 100644
--- a/src/main/java/org/apache/flink/benchmark/AdaptivePartitionThroughputBenchmark.java
+++ b/src/main/java/org/apache/flink/benchmark/AdaptivePartitionThroughputBenchmark.java
@@ -76,9 +76,9 @@ public void mapRebalanceMapRescaleMapSink(InputBenchmarkFlinkEnvironmentContext
         DataStreamSource<Long> source = sourceType.source(env, RECORDS_PER_INVOCATION);
         source.map(new MultiplyByTwo())
                 .rebalance()
-                .map(new SkewableHigherMultiply(1, 0, 5))
+                .map(new SkewableHigherMultiply(1, 100, 0, 5))
                 .rescale()
-                .map(new SkewableHigherMultiply(1, 0, 5))
+                .map(new SkewableHigherMultiply(1, 100, 0, 5))
                 .sinkTo(new DiscardingSink<>());
 
         env.execute();
diff --git a/src/main/java/org/apache/flink/benchmark/functions/SkewableHigherMultiply.java b/src/main/java/org/apache/flink/benchmark/functions/SkewableHigherMultiply.java
index f975694..aa8cab6 100644
--- a/src/main/java/org/apache/flink/benchmark/functions/SkewableHigherMultiply.java
+++ b/src/main/java/org/apache/flink/benchmark/functions/SkewableHigherMultiply.java
@@ -26,12 +26,17 @@
 
 public class SkewableHigherMultiply extends RichMapFunction<Long, Long> {
 
-    private final int additionalCostMillisPerRecord;
+    private final int additionalCostMillisPerBatch;
+    private final int recordsPerBatch;
+    private int numRecordsForCurrentBatch;
     private final Set<Integer> expectedSkewedSubTasksIndexes;
 
     public SkewableHigherMultiply(
-            int additionalCostMillisPerRecord, Integer... expectedSkewedSubTasksIndexes) {
-        this.additionalCostMillisPerRecord = additionalCostMillisPerRecord;
+            int additionalCostMillisPerBatch,
+            int recordsPerBatch,
+            Integer... expectedSkewedSubTasksIndexes) {
+        this.additionalCostMillisPerBatch = additionalCostMillisPerBatch;
+        this.recordsPerBatch = recordsPerBatch;
         this.expectedSkewedSubTasksIndexes =
                 new HashSet<>(Arrays.asList(expectedSkewedSubTasksIndexes));
     }
@@ -39,12 +44,15 @@ public SkewableHigherMultiply(
     @Override
     public Long map(Long value) throws Exception {
         double base = value * 2.0d;
-        if (additionalCostMillisPerRecord <= 0) {
+        if (additionalCostMillisPerBatch <= 0) {
             return (long) base;
         } else {
             if (expectedSkewedSubTasksIndexes.contains(
                     getRuntimeContext().getTaskInfo().getIndexOfThisSubtask())) {
-                Thread.sleep(additionalCostMillisPerRecord);
+                if (++numRecordsForCurrentBatch == recordsPerBatch) {
+                    numRecordsForCurrentBatch = 0;
+                    Thread.sleep(additionalCostMillisPerBatch);
+                }
             }
             return (long) base;
         }