[kudu] 09/09: [spark-tools] Fix DistributedDataGenerator num-tasks

[gr...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

granthenke pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kudu.git

commit d51889f70beaac138ece3a3ea00633d1165038d3
Author: Grant Henke <gr...@apache.org>
AuthorDate: Wed Jan 9 20:48:28 2019 -0600

    [spark-tools] Fix DistributedDataGenerator num-tasks
    
    This patch fixes the num-tasks argument. Previously
    Spark was not creating the correct number of tasks
    to match the passed argument.
    
    Change-Id: I7fa560e71b2f84a75002e9e776011e1a11c5a1ff
    Reviewed-on: http://gerrit.cloudera.org:8080/12211
    Tested-by: Kudu Jenkins
    Reviewed-by: Adar Dembo <ad...@cloudera.com>
---
 .../scala/org/apache/kudu/spark/tools/DistributedDataGenerator.scala  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/java/kudu-spark-tools/src/main/scala/org/apache/kudu/spark/tools/DistributedDataGenerator.scala b/java/kudu-spark-tools/src/main/scala/org/apache/kudu/spark/tools/DistributedDataGenerator.scala
index 2ad4629..ee0954f 100644
--- a/java/kudu-spark-tools/src/main/scala/org/apache/kudu/spark/tools/DistributedDataGenerator.scala
+++ b/java/kudu-spark-tools/src/main/scala/org/apache/kudu/spark/tools/DistributedDataGenerator.scala
@@ -148,8 +148,8 @@ object DistributedDataGenerator {
     val sc = ss.sparkContext
     val context = new KuduContext(options.masterAddresses, sc)
     val metrics = GeneratorMetrics(sc)
-    sc.parallelize(0 until options.numTasks)
-      .foreach(taskNum => generateRows(context, options, taskNum, metrics))
+    sc.parallelize(0 until options.numTasks, numSlices = options.numTasks)
+      .foreachPartition(taskNum => generateRows(context, options, taskNum.next(), metrics))
     log.info(s"Rows written: ${metrics.rowsWritten.value}")
     log.info(s"Collisions: ${metrics.collisions.value}")
   }